From 83c29382b5d4b27c8a23235b2f472987551fa572 Mon Sep 17 00:00:00 2001 From: justinchuby Date: Wed, 12 Aug 2026 17:08:36 +0000 Subject: [PATCH 001/151] Simplify composite inference metadata Keep component I/O only under pipeline.models and make pipeline.phases the sole source of lifecycle scheduling. Strategy stages now describe control structure and ordering without duplicating run_on, while bare decoder metadata retains model.io. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: justinchuby --- .../onnx_genai/inference_metadata.py | 18 +-------- .../onnx_genai/inference_metadata_test.py | 37 ++++--------------- 2 files changed, 8 insertions(+), 47 deletions(-) diff --git a/src/mobius/integrations/onnx_genai/inference_metadata.py b/src/mobius/integrations/onnx_genai/inference_metadata.py index a196c2d42..f16b655c1 100644 --- a/src/mobius/integrations/onnx_genai/inference_metadata.py +++ b/src/mobius/integrations/onnx_genai/inference_metadata.py @@ -1324,12 +1324,6 @@ def annotate_strategy(strategy: dict[str, Any]) -> None: 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]] return metadata @@ -1551,7 +1545,6 @@ def build_native_vlm_package_metadata( { "name": f"run_{name}", "strategy": strategy, - "run_on": phases[name]["run_on"], } ) @@ -1589,7 +1582,6 @@ 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), @@ -2015,7 +2007,7 @@ def build_diffusion_pipeline_metadata( "to": f"denoiser.{denoiser_sample_input}", }, ] - phases: dict[str, Any] = {} + phases: dict[str, Any] = {"denoiser": {"run_on": "every_step"}} if text_encoder_filename is not None: models["text_encoder"] = { @@ -2136,7 +2128,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 +2168,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", }, ] ) @@ -2279,12 +2268,10 @@ 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", }, ], }, @@ -2355,12 +2342,10 @@ def build_audio_codec_pipeline_metadata( { "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", }, ], }, @@ -2522,7 +2507,6 @@ def build_tts_pipeline_metadata( { "name": "generate_codes", "strategy": stage_strategy, - "run_on": "every_step", }, ], }, diff --git a/src/mobius/integrations/onnx_genai/inference_metadata_test.py b/src/mobius/integrations/onnx_genai/inference_metadata_test.py index 06aef903f..160c5878c 100644 --- a/src/mobius/integrations/onnx_genai/inference_metadata_test.py +++ b/src/mobius/integrations/onnx_genai/inference_metadata_test.py @@ -765,25 +765,7 @@ def test_gemma4_routes_all_embedding_outputs(self, tmp_path): "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 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 @@ -803,13 +785,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 = { @@ -948,7 +932,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 +1114,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", @@ -1844,7 +1828,6 @@ def test_vision_only_pipeline(self): "kind": "single_pass", "model": "vision_encoder", }, - "run_on": "prompt_only", }, { "name": "fuse_embeddings", @@ -1852,7 +1835,6 @@ def test_vision_only_pipeline(self): "kind": "single_pass", "model": "embedding", }, - "run_on": "prompt_only", }, { "name": "decode", @@ -1860,7 +1842,6 @@ def test_vision_only_pipeline(self): "kind": "autoregressive", "decoder": "decoder", }, - "run_on": "every_step", }, ], }, @@ -1919,22 +1900,18 @@ 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", }, ] From 7f58445fa7e4eb5696f7a44b37071feaad496fcc Mon Sep 17 00:00:00 2001 From: justinchuby Date: Wed, 12 Aug 2026 18:30:48 +0000 Subject: [PATCH 002/151] Add reusable ONNX generation policy components Introduce model-agnostic ONNX graphs for sampling, termination, solver, masked update, speculative acceptance, and token state math. Persist policy artifacts with ModelPackage and expose schema-compliant workflow component declarations without guessing unfinished workflow bindings. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: justinchuby --- src/mobius/__init__.py | 3 +- src/mobius/_model_package.py | 55 +++- src/mobius/_model_package_test.py | 10 + src/mobius/generation/__init__.py | 34 +++ src/mobius/generation/_policy_components.py | 280 ++++++++++++++++++ .../generation/_policy_components_test.py | 156 ++++++++++ .../integrations/onnx_genai/__init__.py | 2 + .../integrations/onnx_genai/auto_export.py | 2 + .../onnx_genai/inference_metadata.py | 54 ++++ .../onnx_genai/inference_metadata_test.py | 28 ++ 10 files changed, 621 insertions(+), 3 deletions(-) create mode 100644 src/mobius/generation/__init__.py create mode 100644 src/mobius/generation/_policy_components.py create mode 100644 src/mobius/generation/_policy_components_test.py diff --git a/src/mobius/__init__.py b/src/mobius/__init__.py index 9e3438edd..cf3237219 100644 --- a/src/mobius/__init__.py +++ b/src/mobius/__init__.py @@ -47,6 +47,7 @@ "components", "ep_capabilities", "ep_registry", + "generation", "get_build_dtype", "get_ep", "inspect_components", @@ -59,7 +60,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 ( diff --git a/src/mobius/_model_package.py b/src/mobius/_model_package.py index 7740e240a..3d94c8eef 100644 --- a/src/mobius/_model_package.py +++ b/src/mobius/_model_package.py @@ -33,6 +33,7 @@ import tqdm from mobius._optimizations import fold_initializers_after_weights +from mobius.generation import PolicyComponent from mobius.integrations._weight_loading import _assign_weight logger = logging.getLogger(__name__) @@ -50,9 +51,11 @@ def __init__( self, models: dict[str, ir.Model] | None = None, config: object | None = None, + policy_components: dict[str, PolicyComponent] | None = None, ) -> None: super().__init__(models or {}) self.config = config + self.policy_components = dict(policy_components or {}) def __repr__(self) -> str: names = ", ".join(repr(k) for k in self.data) @@ -78,6 +81,7 @@ def save( components: Callable[[str], bool] | None = None, progress_bar: bool = True, check_weights: bool = True, + include_policy_components: bool = True, ) -> None: """Save all component models to a directory. @@ -129,6 +133,8 @@ 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``. Raises: ValueError: If *external_data* is not ``"onnx"`` or @@ -178,6 +184,35 @@ def save( save_kwargs["max_workers"] = max_workers ir.save(model, path, **save_kwargs) + if include_policy_components: + self.save_policy_components(directory, check_weights=check_weights) + + 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 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" + ir.save(component.model, os.path.join(directory, relative_path)) + artifacts[name] = relative_path + return artifacts + @classmethod def load(cls, directory: str) -> ModelPackage: """Load all ``.onnx`` files from a directory into a package. @@ -203,13 +238,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 ------------------------------------------------ diff --git a/src/mobius/_model_package_test.py b/src/mobius/_model_package_test.py index 50f2b9780..831f09309 100644 --- a/src/mobius/_model_package_test.py +++ b/src/mobius/_model_package_test.py @@ -17,6 +17,7 @@ from mobius._configs import VisionConfig from mobius._model_package import ModelPackage, _make_progress_callback from mobius._testing import make_config +from mobius.generation import PolicyRole, build_greedy_sampler from mobius.models.base import CausalLMModel from mobius.models.gemma3 import Gemma3MultiModalModel from mobius.tasks import CausalLMTask, VisionLanguageTask @@ -365,6 +366,15 @@ 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"].role is PolicyRole.TOKEN_SAMPLER + class TestModelPackageApplyWeights: def test_single_component(self): diff --git a/src/mobius/generation/__init__.py b/src/mobius/generation/__init__.py new file mode 100644 index 000000000..392192d0a --- /dev/null +++ b/src/mobius/generation/__init__.py @@ -0,0 +1,34 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT License. + +"""Reusable ONNX generation-policy components.""" + +from __future__ import annotations + +from mobius.generation._policy_components import ( + PolicyCapabilities, + PolicyComponent, + PolicyRole, + attach_policy_components, + build_eos_termination, + build_euler_solver_step, + build_greedy_sampler, + build_masked_token_update, + build_seeded_categorical_sampler, + build_speculative_acceptance, + build_token_state_update, +) + +__all__ = [ + "PolicyComponent", + "PolicyCapabilities", + "PolicyRole", + "attach_policy_components", + "build_eos_termination", + "build_euler_solver_step", + "build_greedy_sampler", + "build_masked_token_update", + "build_seeded_categorical_sampler", + "build_speculative_acceptance", + "build_token_state_update", +] diff --git a/src/mobius/generation/_policy_components.py b/src/mobius/generation/_policy_components.py new file mode 100644 index 000000000..245180677 --- /dev/null +++ b/src/mobius/generation/_policy_components.py @@ -0,0 +1,280 @@ +# 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 + +from dataclasses import dataclass +from enum import StrEnum +from typing import Protocol + +import onnx_ir as ir +from onnxscript import GraphBuilder + +from mobius._constants import OPSET_VERSION + +_POLICY_ROLE_METADATA = "mobius.generation.policy_role" + + +class PolicyRole(StrEnum): + """Architecture-neutral role performed by a policy component.""" + + TOKEN_SAMPLER = "token_sampler" + TERMINATION = "termination" + SOLVER_STEP = "solver_step" + MASKED_UPDATE = "masked_update" + SPECULATIVE_ACCEPTANCE = "speculative_acceptance" + STATE_UPDATE = "state_update" + + +@dataclass(frozen=True) +class PolicyComponent: + """A named role and its executable ONNX model.""" + + role: PolicyRole + model: ir.Model + + def __post_init__(self) -> None: + self.model.graph.metadata_props[_POLICY_ROLE_METADATA] = self.role.value + + @classmethod + def from_model(cls, model: ir.Model) -> PolicyComponent: + """Restore a component from role metadata embedded in its ONNX graph.""" + role = model.graph.metadata_props.get(_POLICY_ROLE_METADATA) + if role is None: + raise ValueError("ONNX policy component is missing its Mobius policy role") + return cls(PolicyRole(role), model) + + +@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 + 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 = {"euler": build_euler_solver_step} + 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.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} + + +def _component(role: PolicyRole, graph: ir.Graph) -> PolicyComponent: + model = ir.Model(graph, ir_version=11) + model.producer_name = "mobius" + return PolicyComponent(role, model) + + +def _make_graph(name: str) -> tuple[ir.Graph, GraphBuilder]: + graph = ir.Graph( + [], + [], + nodes=[], + name=name, + opset_imports={"": OPSET_VERSION}, + ) + return graph, GraphBuilder(graph) + + +def build_greedy_sampler() -> PolicyComponent: + """Build ``logits -> token_ids`` greedy sampling over the final axis.""" + graph, builder = _make_graph("greedy_sampler") + logits = builder.input( + "logits", + dtype=ir.DataType.FLOAT, + shape=["batch", "vocabulary"], + ) + token_ids = builder.op.ArgMax(logits, axis=-1, keepdims=0) + builder.add_output(token_ids, "token_ids") + return _component(PolicyRole.TOKEN_SAMPLER, graph) + + +def build_seeded_categorical_sampler() -> PolicyComponent: + """Build deterministic categorical sampling with explicit seed and counter. + + The integer hash is counter based: the same ``(seed, counter, logits, + temperature)`` inputs always produce the same token. The updated counter is + an explicit output, so no random or hidden mutable state exists in the graph. + """ + 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, []) + seed = builder.input("seed", ir.DataType.INT64, []) + counter = builder.input("counter", ir.DataType.INT64, []) + + # A compact LCG-style integer hash. Constants remain below signed-int64 + # limits, and the prime modulus keeps the result in a precisely castable range. + multiplier = op.Constant(value_int=1_103_515_245) + stream_multiplier = op.Constant(value_int=12_345) + increment = op.Constant(value_int=1_013_904_223) + modulus = op.Constant(value_int=2_147_483_647) + hashed = op.Add( + op.Add(op.Mul(seed, multiplier), op.Mul(counter, stream_multiplier)), + increment, + ) + hashed = op.Mod(hashed, modulus, fmod=0) + uniform = op.Div( + op.Add(op.Cast(hashed, to=ir.DataType.FLOAT), op.Constant(value_float=0.5)), + op.Constant(value_float=2_147_483_647.0), + ) + + scaled_logits = op.Div(logits, temperature) + probabilities = op.Softmax(scaled_logits, axis=-1) + axis = op.Constant(value_int=-1) + cumulative = op.CumSum(probabilities, axis) + uniform = op.Unsqueeze(uniform, op.Constant(value_ints=[0])) + candidates = op.GreaterOrEqual(cumulative, uniform) + token_ids = op.ArgMax( + op.Cast(candidates, to=ir.DataType.INT64), + axis=-1, + keepdims=0, + ) + next_counter = op.Add(counter, op.Constant(value_int=1)) + builder.add_output(token_ids, "token_ids") + builder.add_output(next_counter, "next_counter") + return _component(PolicyRole.TOKEN_SAMPLER, graph) + + +def build_eos_termination() -> 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("token_ids", ir.DataType.INT64, ["batch"]) + eos_token_ids = builder.input("eos_token_ids", ir.DataType.INT64, ["num_eos"]) + tokens = op.Unsqueeze(token_ids, op.Constant(value_ints=[-1])) + eos = op.Unsqueeze(eos_token_ids, op.Constant(value_ints=[0])) + matches = op.Equal(tokens, eos) + match_count = op.ReduceSum( + op.Cast(matches, to=ir.DataType.INT64), + axes=[-1], + keepdims=0, + ) + terminated = op.Greater(match_count, op.Constant(value_int=0)) + builder.add_output(terminated, "terminated") + return _component(PolicyRole.TERMINATION, graph) + + +def build_euler_solver_step() -> PolicyComponent: + """Build the generic Euler update ``x_next = x + dx * (sigma_next-sigma)``.""" + graph, builder = _make_graph("euler_solver_step") + op = builder.op + sample = builder.input( + "sample", + ir.DataType.FLOAT, + ["batch", "channels", "height", "width"], + ) + derivative = builder.input( + "derivative", + ir.DataType.FLOAT, + ["batch", "channels", "height", "width"], + ) + sigma = builder.input("sigma", ir.DataType.FLOAT, []) + sigma_next = builder.input("sigma_next", ir.DataType.FLOAT, []) + next_sample = op.Add(sample, op.Mul(derivative, op.Sub(sigma_next, sigma))) + builder.add_output(next_sample, "next_sample") + return _component(PolicyRole.SOLVER_STEP, graph) + + +def build_masked_token_update() -> PolicyComponent: + """Build confidence-thresholded replacement for masked token positions.""" + 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"]) + confidence = builder.input("confidence", ir.DataType.FLOAT, ["batch", "sequence"]) + masked = builder.input("masked", ir.DataType.BOOL, ["batch", "sequence"]) + threshold = builder.input("threshold", ir.DataType.FLOAT, []) + accepted = op.And(masked, op.GreaterOrEqual(confidence, threshold)) + updated = op.Where(accepted, proposed, current) + remaining = op.And(masked, op.Not(accepted)) + builder.add_output(updated, "updated_tokens") + builder.add_output(remaining, "remaining_mask") + return _component(PolicyRole.MASKED_UPDATE, graph) + + +def build_speculative_acceptance() -> PolicyComponent: + """Build per-token speculative acceptance and accepted-prefix length.""" + graph, builder = _make_graph("speculative_acceptance") + op = builder.op + target_probability = builder.input( + "target_probability", ir.DataType.FLOAT, ["batch", "draft_sequence"] + ) + draft_probability = builder.input( + "draft_probability", ir.DataType.FLOAT, ["batch", "draft_sequence"] + ) + uniform = builder.input("uniform", ir.DataType.FLOAT, ["batch", "draft_sequence"]) + ratio = op.Div(target_probability, draft_probability) + probability = op.Min(ratio, op.Constant(value_float=1.0)) + accepted = op.LessOrEqual(uniform, probability) + 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, + ) + accepted_count = op.ReduceSum(prefix, axes=[-1], keepdims=0) + builder.add_output(accepted, "accepted") + builder.add_output(accepted_count, "accepted_count") + return _component(PolicyRole.SPECULATIVE_ACCEPTANCE, graph) + + +def build_token_state_update() -> PolicyComponent: + """Build explicit token-history append and sequence-length update math.""" + graph, builder = _make_graph("token_state_update") + op = builder.op + tokens = builder.input("tokens", ir.DataType.INT64, ["batch", "sequence"]) + next_token = builder.input("next_token", ir.DataType.INT64, ["batch"]) + sequence_length = builder.input("sequence_length", ir.DataType.INT64, []) + appended = op.Concat( + tokens, + op.Unsqueeze(next_token, op.Constant(value_ints=[-1])), + axis=-1, + ) + next_length = op.Add(sequence_length, op.Constant(value_int=1)) + builder.add_output(appended, "updated_tokens") + builder.add_output(next_length, "updated_sequence_length") + return _component(PolicyRole.STATE_UPDATE, 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..6d10eada3 --- /dev/null +++ b/src/mobius/generation/_policy_components_test.py @@ -0,0 +1,156 @@ +# 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 + +from mobius._model_package import ModelPackage +from mobius.generation import ( + PolicyCapabilities, + PolicyRole, + attach_policy_components, + build_eos_termination, + build_euler_solver_step, + build_greedy_sampler, + build_masked_token_update, + build_seeded_categorical_sampler, + build_speculative_acceptance, + build_token_state_update, +) + + +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 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_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), + "seed": np.array(7, np.int64), + "counter": np.array(11, np.int64), + } + first = _run(component, tmp_path, feeds) + second = _run(component, tmp_path, feeds) + np.testing.assert_array_equal(first[0], second[0]) + assert first[1] == 12 + + +def test_eos_termination_runtime(tmp_path): + (terminated,) = _run( + build_eos_termination(), + tmp_path, + { + "token_ids": np.array([2, 8, 9], np.int64), + "eos_token_ids": np.array([2, 9], np.int64), + }, + ) + np.testing.assert_array_equal(terminated, [True, False, True]) + + +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, + "sigma": np.array(1.5, np.float32), + "sigma_next": np.array(0.5, np.float32), + }, + ) + np.testing.assert_allclose(actual, sample - derivative) + + +def test_masked_update_runtime_parity(tmp_path): + 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), + "confidence": np.array([[0.9, 0.8, 0.2]], np.float32), + "masked": np.array([[False, True, True]]), + "threshold": np.array(0.5, np.float32), + }, + ) + np.testing.assert_array_equal(outputs[0], [[1, 5, 99]]) + np.testing.assert_array_equal(outputs[1], [[False, False, True]]) + + +def test_speculative_acceptance_prefix_runtime(tmp_path): + accepted, count = _run( + build_speculative_acceptance(), + tmp_path, + { + "target_probability": np.array([[0.9, 0.5, 0.1, 0.9]], np.float32), + "draft_probability": np.array([[0.8, 0.5, 0.8, 0.8]], np.float32), + "uniform": np.array([[0.5, 0.5, 0.5, 0.5]], np.float32), + }, + ) + np.testing.assert_array_equal(accepted, [[True, True, False, True]]) + np.testing.assert_array_equal(count, [2]) + + +def test_token_state_update_runtime(tmp_path): + tokens, length = _run( + build_token_state_update(), + tmp_path, + { + "tokens": np.array([[1, 2], [3, 4]], np.int64), + "next_token": np.array([5, 6], np.int64), + "sequence_length": np.array(2, np.int64), + }, + ) + np.testing.assert_array_equal(tokens, [[1, 2, 5], [3, 4, 6]]) + assert length == 3 + + +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.role for component in package.policy_components.values()} == { + PolicyRole.TOKEN_SAMPLER, + PolicyRole.TERMINATION, + PolicyRole.SOLVER_STEP, + PolicyRole.MASKED_UPDATE, + PolicyRole.SPECULATIVE_ACCEPTANCE, + PolicyRole.STATE_UPDATE, + } diff --git a/src/mobius/integrations/onnx_genai/__init__.py b/src/mobius/integrations/onnx_genai/__init__.py index ad19d0bd7..9024ae581 100644 --- a/src/mobius/integrations/onnx_genai/__init__.py +++ b/src/mobius/integrations/onnx_genai/__init__.py @@ -50,6 +50,7 @@ ) from mobius.integrations.onnx_genai.inference_metadata import ( SchedulerConfig, + add_policy_components_to_workflow, build_audio_codec_pipeline_metadata, build_diffusion_pipeline_metadata, build_language_diffusion_pipeline_metadata, @@ -68,6 +69,7 @@ "ComfyUIWorkflow", "ConversionResult", "SchedulerConfig", + "add_policy_components_to_workflow", "build_decoder_metadata", "build_diffusion_pipeline_metadata", "build_language_diffusion_pipeline_metadata", diff --git a/src/mobius/integrations/onnx_genai/auto_export.py b/src/mobius/integrations/onnx_genai/auto_export.py index 0d9f9f793..3aa00eded 100644 --- a/src/mobius/integrations/onnx_genai/auto_export.py +++ b/src/mobius/integrations/onnx_genai/auto_export.py @@ -25,6 +25,7 @@ from mobius.integrations.onnx_genai.inference_metadata import ( SchedulerConfig, add_explicit_package_io, + add_policy_components_to_workflow, load_diffusers_scheduler_config, write_audio_codec_pipeline_metadata, write_diffusion_pipeline_metadata, @@ -49,6 +50,7 @@ def _add_explicit_io_to_file(path: str, pkg: Any, config: Any) -> None: with open(path, encoding="utf-8") as handle: metadata = yaml.safe_load(handle) add_explicit_package_io(metadata, pkg, config) + add_policy_components_to_workflow(metadata, pkg) with open(path, "w", encoding="utf-8") as handle: yaml.safe_dump(metadata, handle, sort_keys=False) diff --git a/src/mobius/integrations/onnx_genai/inference_metadata.py b/src/mobius/integrations/onnx_genai/inference_metadata.py index f16b655c1..01add57a1 100644 --- a/src/mobius/integrations/onnx_genai/inference_metadata.py +++ b/src/mobius/integrations/onnx_genai/inference_metadata.py @@ -1327,6 +1327,60 @@ def annotate_strategy(strategy: dict[str, Any]) -> None: return metadata +def add_policy_components_to_workflow( + metadata: dict[str, Any], + pkg: Any, +) -> dict[str, Any]: + """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", {}) + for name, component in policy_components.items(): + model = component.model + components[name] = { + "implementation": { + "kind": "onnx", + "artifact": f"policies/{name}.onnx", + }, + "ports": { + "inputs": { + value.name: { + "dtype": _port(value).dtype, + "rank": _port(value).rank, + **( + {"shape": _shape_metadata(_port(value))} + if value.shape is not None + else {} + ), + } + for value in model.graph.inputs + }, + "outputs": { + value.name: { + "dtype": _port(value).dtype, + "rank": _port(value).rank, + **( + {"shape": _shape_metadata(_port(value))} + if value.shape is not None + else {} + ), + } + for value in model.graph.outputs + }, + }, + } + return metadata + + def _topological_order( names: Iterable[str], edges: list[dict[str, Any]], diff --git a/src/mobius/integrations/onnx_genai/inference_metadata_test.py b/src/mobius/integrations/onnx_genai/inference_metadata_test.py index 160c5878c..f677faab6 100644 --- a/src/mobius/integrations/onnx_genai/inference_metadata_test.py +++ b/src/mobius/integrations/onnx_genai/inference_metadata_test.py @@ -22,12 +22,14 @@ 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, _port, add_explicit_package_io, + add_policy_components_to_workflow, build_diffusion_pipeline_metadata, build_language_diffusion_pipeline_metadata, build_multimodal_pipeline_metadata, @@ -42,6 +44,32 @@ ) +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": {"ir_version": "1.0"}, + "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 set(component["ports"]["inputs"]) == {"logits"} + assert set(component["ports"]["outputs"]) == {"token_ids"} + 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 = [ From ffc1b0439689de59142b82dd9eba074dd436b4f6 Mon Sep 17 00:00:00 2001 From: justinchuby Date: Wed, 12 Aug 2026 18:46:00 +0000 Subject: [PATCH 003/151] Align policy components with workflow contract Implement the published policy roles, shapes, linear effects, and counter-based Threefry RNG. Emit a schema-valid SSA decoder workflow with explicit component invocations, loop-carried token and KV state, effect tokens, termination, and token emission. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: justinchuby --- src/mobius/generation/_policy_components.py | 304 +++++++++++---- .../generation/_policy_components_test.py | 48 +-- .../integrations/onnx_genai/__init__.py | 6 + .../integrations/onnx_genai/auto_export.py | 9 +- .../onnx_genai/auto_export_test.py | 26 +- .../onnx_genai/inference_metadata.py | 7 +- .../onnx_genai/inference_metadata_test.py | 10 +- .../onnx_genai/workflow_metadata.py | 367 ++++++++++++++++++ 8 files changed, 668 insertions(+), 109 deletions(-) create mode 100644 src/mobius/integrations/onnx_genai/workflow_metadata.py diff --git a/src/mobius/generation/_policy_components.py b/src/mobius/generation/_policy_components.py index 245180677..edf0e3ca0 100644 --- a/src/mobius/generation/_policy_components.py +++ b/src/mobius/generation/_policy_components.py @@ -10,6 +10,7 @@ from __future__ import annotations +import json from dataclasses import dataclass from enum import StrEnum from typing import Protocol @@ -20,16 +21,18 @@ from mobius._constants import OPSET_VERSION _POLICY_ROLE_METADATA = "mobius.generation.policy_role" +_POLICY_CONTRACT_METADATA = "mobius.generation.policy_contract" +_POLICY_EFFECTS_METADATA = "mobius.generation.policy_effects" class PolicyRole(StrEnum): """Architecture-neutral role performed by a policy component.""" TOKEN_SAMPLER = "token_sampler" - TERMINATION = "termination" + TERMINATION = "termination_predicate" SOLVER_STEP = "solver_step" MASKED_UPDATE = "masked_update" - SPECULATIVE_ACCEPTANCE = "speculative_acceptance" + SPECULATIVE_ACCEPTANCE = "speculative_verifier" STATE_UPDATE = "state_update" @@ -39,9 +42,13 @@ class PolicyComponent: role: PolicyRole model: ir.Model + contract: dict[str, object] + effects: tuple[str, ...] def __post_init__(self) -> None: self.model.graph.metadata_props[_POLICY_ROLE_METADATA] = self.role.value + 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: @@ -49,7 +56,9 @@ def from_model(cls, model: ir.Model) -> PolicyComponent: role = model.graph.metadata_props.get(_POLICY_ROLE_METADATA) if role is None: raise ValueError("ONNX policy component is missing its Mobius policy role") - return cls(PolicyRole(role), model) + contract = json.loads(model.graph.metadata_props[_POLICY_CONTRACT_METADATA]) + effects = tuple(json.loads(model.graph.metadata_props[_POLICY_EFFECTS_METADATA])) + return cls(PolicyRole(role), model, contract, effects) @dataclass(frozen=True) @@ -102,10 +111,15 @@ def attach_policy_components( return {name: f"policies/{name}.onnx" for name, _ in selected} -def _component(role: PolicyRole, graph: ir.Graph) -> PolicyComponent: +def _component( + role: PolicyRole, + graph: ir.Graph, + contract: dict[str, object], + *effects: str, +) -> PolicyComponent: model = ir.Model(graph, ir_version=11) model.producer_name = "mobius" - return PolicyComponent(role, model) + return PolicyComponent(role, model, contract, effects) def _make_graph(name: str) -> tuple[ir.Graph, GraphBuilder]: @@ -128,55 +142,116 @@ def build_greedy_sampler() -> PolicyComponent: shape=["batch", "vocabulary"], ) token_ids = builder.op.ArgMax(logits, axis=-1, keepdims=0) - builder.add_output(token_ids, "token_ids") - return _component(PolicyRole.TOKEN_SAMPLER, graph) + builder.add_output(token_ids, "token") + return _component( + PolicyRole.TOKEN_SAMPLER, + graph, + { + "role": "token_sampler", + "mode": "greedy", + "logits": "logits", + "token": "token", + "effect": "sample", + }, + "sample", + ) def build_seeded_categorical_sampler() -> PolicyComponent: - """Build deterministic categorical sampling with explicit seed and counter. + """Build deterministic categorical sampling with explicit seed and offset. - The integer hash is counter based: the same ``(seed, counter, logits, - temperature)`` inputs always produce the same token. The updated counter is + Threefry is counter based: the same ``(seed, offset, logits, temperature)`` + inputs always produce the same token. The updated offset is an explicit output, so no random or hidden mutable state exists in the graph. """ 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, []) - seed = builder.input("seed", ir.DataType.INT64, []) - counter = builder.input("counter", ir.DataType.INT64, []) - - # A compact LCG-style integer hash. Constants remain below signed-int64 - # limits, and the prime modulus keeps the result in a precisely castable range. - multiplier = op.Constant(value_int=1_103_515_245) - stream_multiplier = op.Constant(value_int=12_345) - increment = op.Constant(value_int=1_013_904_223) - modulus = op.Constant(value_int=2_147_483_647) - hashed = op.Add( - op.Add(op.Mul(seed, multiplier), op.Mul(counter, stream_multiplier)), - increment, + temperature = builder.input("temperature", ir.DataType.FLOAT, ["batch"]) + seed = builder.input("seed", ir.DataType.INT64, ["batch"]) + offset = builder.input("offset", ir.DataType.INT64, ["batch"]) + + # 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(offset, 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", ) - hashed = op.Mod(hashed, modulus, fmod=0) uniform = op.Div( - op.Add(op.Cast(hashed, to=ir.DataType.FLOAT), op.Constant(value_float=0.5)), - op.Constant(value_float=2_147_483_647.0), + 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) - scaled_logits = op.Div(logits, temperature) + scaled_logits = op.Div( + logits, + op.Unsqueeze(temperature, op.Constant(value_ints=[-1])), + ) probabilities = op.Softmax(scaled_logits, axis=-1) axis = op.Constant(value_int=-1) cumulative = op.CumSum(probabilities, axis) - uniform = op.Unsqueeze(uniform, op.Constant(value_ints=[0])) + 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, ) - next_counter = op.Add(counter, op.Constant(value_int=1)) - builder.add_output(token_ids, "token_ids") - builder.add_output(next_counter, "next_counter") - return _component(PolicyRole.TOKEN_SAMPLER, graph) + next_offset = op.Add(offset, op.Constant(value_int=1)) + builder.add_output(token_ids, "token") + builder.add_output(next_offset, "next_offset") + return _component( + PolicyRole.TOKEN_SAMPLER, + graph, + { + "role": "token_sampler", + "mode": "seeded_stochastic", + "logits": "logits", + "token": "token", + "temperature": "temperature", + "rng": { + "seed": "seed", + "offset": "offset", + "next_offset": "next_offset", + }, + "effect": "rng", + }, + "rng", + ) def build_eos_termination() -> PolicyComponent: @@ -184,18 +259,38 @@ def build_eos_termination() -> PolicyComponent: graph, builder = _make_graph("eos_termination") op = builder.op token_ids = builder.input("token_ids", ir.DataType.INT64, ["batch"]) - eos_token_ids = builder.input("eos_token_ids", ir.DataType.INT64, ["num_eos"]) + eos_ids = builder.input("eos_ids", ir.DataType.INT64, ["num_eos"]) + iteration = builder.input("iteration", ir.DataType.INT64, ["batch"]) + max_iterations = builder.input("max_iterations", ir.DataType.INT64, ["batch"]) tokens = op.Unsqueeze(token_ids, op.Constant(value_ints=[-1])) - eos = op.Unsqueeze(eos_token_ids, op.Constant(value_ints=[0])) + eos = op.Unsqueeze(eos_ids, op.Constant(value_ints=[0])) matches = op.Equal(tokens, eos) match_count = op.ReduceSum( op.Cast(matches, to=ir.DataType.INT64), axes=[-1], keepdims=0, ) - terminated = op.Greater(match_count, op.Constant(value_int=0)) - builder.add_output(terminated, "terminated") - return _component(PolicyRole.TERMINATION, graph) + 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, + ) + done = op.Or(hit_eos, hit_limit) + builder.add_output(done, "done") + return _component( + PolicyRole.TERMINATION, + graph, + { + "role": "termination_predicate", + "tokens": "token_ids", + "eos_ids": "eos_ids", + "iteration": "iteration", + "max_iterations": "max_iterations", + "done": "done", + "effect": "termination", + }, + "termination", + ) def build_euler_solver_step() -> PolicyComponent: @@ -212,44 +307,76 @@ def build_euler_solver_step() -> PolicyComponent: ir.DataType.FLOAT, ["batch", "channels", "height", "width"], ) - sigma = builder.input("sigma", ir.DataType.FLOAT, []) - sigma_next = builder.input("sigma_next", ir.DataType.FLOAT, []) - next_sample = op.Add(sample, op.Mul(derivative, op.Sub(sigma_next, sigma))) - builder.add_output(next_sample, "next_sample") - return _component(PolicyRole.SOLVER_STEP, graph) + 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.Unsqueeze(delta, op.Constant(value_ints=[1, 2, 3])) + next_sample = op.Add(sample, op.Mul(derivative, delta)) + builder.add_output(next_sample, "next_state") + return _component( + PolicyRole.SOLVER_STEP, + graph, + { + "role": "solver_step", + "state": "sample", + "estimate": "derivative", + "step": "step", + "schedule": "schedule", + "next_state": "next_state", + "effect": "solver", + }, + "solver", + ) def build_masked_token_update() -> PolicyComponent: - """Build confidence-thresholded replacement for masked token positions.""" + """Build deterministic replacement of currently masked token positions.""" 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"]) - confidence = builder.input("confidence", ir.DataType.FLOAT, ["batch", "sequence"]) masked = builder.input("masked", ir.DataType.BOOL, ["batch", "sequence"]) - threshold = builder.input("threshold", ir.DataType.FLOAT, []) - accepted = op.And(masked, op.GreaterOrEqual(confidence, threshold)) - updated = op.Where(accepted, proposed, current) - remaining = op.And(masked, op.Not(accepted)) - builder.add_output(updated, "updated_tokens") - builder.add_output(remaining, "remaining_mask") - return _component(PolicyRole.MASKED_UPDATE, graph) + step = builder.input("step", ir.DataType.INT64, ["batch"]) + updated = op.Where(masked, proposed, current) + # Consume the declared step without changing values; schedules that remask + # tokens can be expressed by a richer artifact with the same semantic ports. + updated = op.Add(updated, op.Unsqueeze(op.Mul(step, 0), op.Constant(value_ints=[-1]))) + remaining = op.ConstantOfShape(op.Shape(masked), value=ir.tensor([False])) + builder.add_output(updated, "next_state") + builder.add_output(remaining, "next_mask") + return _component( + PolicyRole.MASKED_UPDATE, + graph, + { + "role": "masked_update", + "state": "current_tokens", + "proposal": "proposed_tokens", + "mask": "masked", + "step": "step", + "next_state": "next_state", + "next_mask": "next_mask", + "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_probability = builder.input( - "target_probability", ir.DataType.FLOAT, ["batch", "draft_sequence"] + target_scores = builder.input( + "target_scores", ir.DataType.FLOAT, ["batch", "draft_sequence", "vocabulary"] ) - draft_probability = builder.input( - "draft_probability", ir.DataType.FLOAT, ["batch", "draft_sequence"] + proposed_tokens = builder.input( + "proposed_tokens", ir.DataType.INT64, ["batch", "draft_sequence"] ) - uniform = builder.input("uniform", ir.DataType.FLOAT, ["batch", "draft_sequence"]) - ratio = op.Div(target_probability, draft_probability) - probability = op.Min(ratio, op.Constant(value_float=1.0)) - accepted = op.LessOrEqual(uniform, probability) + 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( @@ -257,24 +384,55 @@ def build_speculative_acceptance() -> PolicyComponent: to=ir.DataType.INT64, ) accepted_count = op.ReduceSum(prefix, axes=[-1], keepdims=0) - builder.add_output(accepted, "accepted") - builder.add_output(accepted_count, "accepted_count") - return _component(PolicyRole.SPECULATIVE_ACCEPTANCE, graph) + accepted_tokens = op.Where( + op.Cast(prefix, to=ir.DataType.BOOL), + proposed_tokens, + op.ConstantOfShape( + op.Shape(proposed_tokens), + value=ir.tensor([0], dtype=ir.DataType.INT64), + ), + ) + draft_length = op.Shape(proposed_tokens, start=1, end=2) + done = op.Equal(accepted_count, draft_length) + builder.add_output(accepted_tokens, "accepted_tokens") + builder.add_output(accepted_count, "accepted_len") + builder.add_output(done, "done") + return _component( + PolicyRole.SPECULATIVE_ACCEPTANCE, + graph, + { + "role": "speculative_verifier", + "target_scores": "target_scores", + "proposed_tokens": "proposed_tokens", + "accepted_tokens": "accepted_tokens", + "accepted_len": "accepted_len", + "done": "done", + "effect": "verify", + }, + "verify", + ) def build_token_state_update() -> PolicyComponent: """Build explicit token-history append and sequence-length update math.""" graph, builder = _make_graph("token_state_update") op = builder.op - tokens = builder.input("tokens", ir.DataType.INT64, ["batch", "sequence"]) - next_token = builder.input("next_token", ir.DataType.INT64, ["batch"]) - sequence_length = builder.input("sequence_length", ir.DataType.INT64, []) - appended = op.Concat( - tokens, - op.Unsqueeze(next_token, op.Constant(value_ints=[-1])), - axis=-1, + current = builder.input("current", ir.DataType.INT64, ["batch", 1]) + update = builder.input("update", ir.DataType.INT64, ["batch"]) + next_state = op.Add( + op.Mul(current, op.Constant(value_int=0)), + op.Unsqueeze(update, op.Constant(value_ints=[-1])), + ) + builder.add_output(next_state, "next") + return _component( + PolicyRole.STATE_UPDATE, + graph, + { + "role": "state_update", + "current": "current", + "update": "update", + "next": "next", + "effect": "state", + }, + "state", ) - next_length = op.Add(sequence_length, op.Constant(value_int=1)) - builder.add_output(appended, "updated_tokens") - builder.add_output(next_length, "updated_sequence_length") - return _component(PolicyRole.STATE_UPDATE, graph) diff --git a/src/mobius/generation/_policy_components_test.py b/src/mobius/generation/_policy_components_test.py index 6d10eada3..4e3760056 100644 --- a/src/mobius/generation/_policy_components_test.py +++ b/src/mobius/generation/_policy_components_test.py @@ -42,14 +42,14 @@ 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), - "seed": np.array(7, np.int64), - "counter": np.array(11, np.int64), + "temperature": np.array([1.0], np.float32), + "seed": np.array([7], np.int64), + "offset": np.array([11], np.int64), } first = _run(component, tmp_path, feeds) second = _run(component, tmp_path, feeds) np.testing.assert_array_equal(first[0], second[0]) - assert first[1] == 12 + np.testing.assert_array_equal(first[1], [12]) def test_eos_termination_runtime(tmp_path): @@ -58,10 +58,12 @@ def test_eos_termination_runtime(tmp_path): tmp_path, { "token_ids": np.array([2, 8, 9], np.int64), - "eos_token_ids": np.array([2, 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, False, True]) + np.testing.assert_array_equal(terminated, [True, True, True]) def test_euler_solver_runtime_parity(tmp_path): @@ -73,8 +75,8 @@ def test_euler_solver_runtime_parity(tmp_path): { "sample": sample, "derivative": derivative, - "sigma": np.array(1.5, np.float32), - "sigma_next": np.array(0.5, np.float32), + "step": np.array([0], np.int64), + "schedule": np.array([1.5, 0.5], np.float32), }, ) np.testing.assert_allclose(actual, sample - derivative) @@ -87,41 +89,41 @@ def test_masked_update_runtime_parity(tmp_path): { "current_tokens": np.array([[1, 99, 99]], np.int64), "proposed_tokens": np.array([[4, 5, 6]], np.int64), - "confidence": np.array([[0.9, 0.8, 0.2]], np.float32), "masked": np.array([[False, True, True]]), - "threshold": np.array(0.5, np.float32), + "step": np.array([0], np.int64), }, ) - np.testing.assert_array_equal(outputs[0], [[1, 5, 99]]) - np.testing.assert_array_equal(outputs[1], [[False, False, True]]) + np.testing.assert_array_equal(outputs[0], [[1, 5, 6]]) + np.testing.assert_array_equal(outputs[1], [[False, False, False]]) def test_speculative_acceptance_prefix_runtime(tmp_path): - accepted, count = _run( + accepted_tokens, count, done = _run( build_speculative_acceptance(), tmp_path, { - "target_probability": np.array([[0.9, 0.5, 0.1, 0.9]], np.float32), - "draft_probability": np.array([[0.8, 0.5, 0.8, 0.8]], np.float32), - "uniform": np.array([[0.5, 0.5, 0.5, 0.5]], np.float32), + "target_scores": np.array( + [[[0, 1], [1, 0], [0, 1], [1, 0]]], + np.float32, + ), + "proposed_tokens": np.array([[1, 0, 0, 0]], np.int64), }, ) - np.testing.assert_array_equal(accepted, [[True, True, False, True]]) + np.testing.assert_array_equal(accepted_tokens, [[1, 0, 0, 0]]) np.testing.assert_array_equal(count, [2]) + np.testing.assert_array_equal(done, [False]) def test_token_state_update_runtime(tmp_path): - tokens, length = _run( + (next_state,) = _run( build_token_state_update(), tmp_path, { - "tokens": np.array([[1, 2], [3, 4]], np.int64), - "next_token": np.array([5, 6], np.int64), - "sequence_length": np.array(2, np.int64), + "current": np.array([[1], [3]], np.int64), + "update": np.array([5, 7], np.int64), }, ) - np.testing.assert_array_equal(tokens, [[1, 2, 5], [3, 4, 6]]) - assert length == 3 + np.testing.assert_array_equal(next_state, [[5], [7]]) def test_capability_driven_attachment_is_model_agnostic(): diff --git a/src/mobius/integrations/onnx_genai/__init__.py b/src/mobius/integrations/onnx_genai/__init__.py index 9024ae581..c50f13bc1 100644 --- a/src/mobius/integrations/onnx_genai/__init__.py +++ b/src/mobius/integrations/onnx_genai/__init__.py @@ -64,6 +64,10 @@ write_speech_to_text_pipeline_metadata, write_tts_pipeline_metadata, ) +from mobius.integrations.onnx_genai.workflow_metadata import ( + build_decoder_workflow_metadata, + write_decoder_workflow_metadata, +) __all__ = [ "ComfyUIWorkflow", @@ -71,6 +75,7 @@ "SchedulerConfig", "add_policy_components_to_workflow", "build_decoder_metadata", + "build_decoder_workflow_metadata", "build_diffusion_pipeline_metadata", "build_language_diffusion_pipeline_metadata", "build_audio_codec_pipeline_metadata", @@ -87,6 +92,7 @@ "translate_comfyui_workflow", "translate_comfyui_workflow_file", "write_decoder_metadata", + "write_decoder_workflow_metadata", "write_diffusion_pipeline_metadata", "write_audio_codec_pipeline_metadata", "write_multimodal_pipeline_metadata", diff --git a/src/mobius/integrations/onnx_genai/auto_export.py b/src/mobius/integrations/onnx_genai/auto_export.py index 3aa00eded..af81b5797 100644 --- a/src/mobius/integrations/onnx_genai/auto_export.py +++ b/src/mobius/integrations/onnx_genai/auto_export.py @@ -20,7 +20,6 @@ from mobius.integrations.onnx_genai.decoder_metadata import ( decoder_metadata_from_config, - write_decoder_metadata, ) from mobius.integrations.onnx_genai.inference_metadata import ( SchedulerConfig, @@ -33,6 +32,9 @@ write_speech_to_text_pipeline_metadata, write_tts_pipeline_metadata, ) +from mobius.integrations.onnx_genai.workflow_metadata import ( + write_decoder_workflow_metadata, +) _LOGGER = logging.getLogger(__name__) @@ -625,10 +627,7 @@ 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 - ) - _add_explicit_io_to_file(path, pkg, resolved_config) + path = write_decoder_workflow_metadata(pkg, output_dir, resolved_config) artifacts = {"inference_metadata": path} tokenizer_path = _write_hf_tokenizer(output_dir, source, revision=revision) if tokenizer_path is not None: diff --git a/src/mobius/integrations/onnx_genai/auto_export_test.py b/src/mobius/integrations/onnx_genai/auto_export_test.py index c3e9f075d..eecafad47 100644 --- a/src/mobius/integrations/onnx_genai/auto_export_test.py +++ b/src/mobius/integrations/onnx_genai/auto_export_test.py @@ -14,7 +14,9 @@ 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.inference_metadata_test import _decoder_model @dataclasses.dataclass @@ -44,12 +46,26 @@ class _MultimodalPkg(dict): config = _Cfg() +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 workflow["manifest"]["ir_version"] == "1.0" + assert workflow["components"]["token_sampler"]["policy"]["role"] == "token_sampler" + assert workflow["components"]["termination"]["policy"]["role"] == ("termination_predicate") + assert workflow["graph"]["kind"] == "loop" + assert (tmp_path / "policies" / "token_sampler.onnx").is_file() def test_dispatch_diffusion(tmp_path): @@ -580,7 +596,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 +605,6 @@ 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() diff --git a/src/mobius/integrations/onnx_genai/inference_metadata.py b/src/mobius/integrations/onnx_genai/inference_metadata.py index 01add57a1..1b2ecc172 100644 --- a/src/mobius/integrations/onnx_genai/inference_metadata.py +++ b/src/mobius/integrations/onnx_genai/inference_metadata.py @@ -1344,6 +1344,7 @@ def add_policy_components_to_workflow( if not isinstance(workflow, dict): return metadata components = workflow.setdefault("components", {}) + workflow_dtypes = {"fp16": "float16", "bf16": "bfloat16", "fp32": "float32"} for name, component in policy_components.items(): model = component.model components[name] = { @@ -1354,7 +1355,7 @@ def add_policy_components_to_workflow( "ports": { "inputs": { value.name: { - "dtype": _port(value).dtype, + "dtype": workflow_dtypes.get(_port(value).dtype, _port(value).dtype), "rank": _port(value).rank, **( {"shape": _shape_metadata(_port(value))} @@ -1366,7 +1367,7 @@ def add_policy_components_to_workflow( }, "outputs": { value.name: { - "dtype": _port(value).dtype, + "dtype": workflow_dtypes.get(_port(value).dtype, _port(value).dtype), "rank": _port(value).rank, **( {"shape": _shape_metadata(_port(value))} @@ -1377,6 +1378,8 @@ def add_policy_components_to_workflow( for value in model.graph.outputs }, }, + "policy": component.contract, + "effects": list(component.effects), } return metadata diff --git a/src/mobius/integrations/onnx_genai/inference_metadata_test.py b/src/mobius/integrations/onnx_genai/inference_metadata_test.py index f677faab6..db67e52ff 100644 --- a/src/mobius/integrations/onnx_genai/inference_metadata_test.py +++ b/src/mobius/integrations/onnx_genai/inference_metadata_test.py @@ -66,7 +66,15 @@ def test_workflow_policy_components_reference_saved_onnx_artifacts(tmp_path): "artifact": "policies/sample.onnx", } assert set(component["ports"]["inputs"]) == {"logits"} - assert set(component["ports"]["outputs"]) == {"token_ids"} + assert set(component["ports"]["outputs"]) == {"token"} + assert component["policy"] == { + "role": "token_sampler", + "mode": "greedy", + "logits": "logits", + "token": "token", + "effect": "sample", + } + assert component["effects"] == ["sample"] assert (tmp_path / component["implementation"]["artifact"]).is_file() 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..d2b4e2f91 --- /dev/null +++ b/src/mobius/integrations/onnx_genai/workflow_metadata.py @@ -0,0 +1,367 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT License. + +"""ONNX GenAI workflow-IR metadata production.""" + +from __future__ import annotations + +import os +from typing import Any + +import onnx_ir as ir +import yaml + +from mobius._constants import OPSET_VERSION +from mobius.generation import ( + PolicyCapabilities, + attach_policy_components, +) +from mobius.integrations.onnx_genai.inference_metadata import ( + _port, + _shape_metadata, + add_policy_components_to_workflow, +) + + +def _contract(value: ir.Value) -> dict[str, Any]: + port = _port(value) + dtype = {"fp16": "float16", "bf16": "bfloat16", "fp32": "float32"}.get( + port.dtype, port.dtype + ) + return { + "dtype": dtype, + "rank": port.rank, + "shape": _shape_metadata(port), + } + + +def _component(model: ir.Model, artifact: str) -> dict[str, Any]: + return { + "implementation": {"kind": "onnx", "artifact": artifact}, + "ports": { + "inputs": {value.name: _contract(value) for value in model.graph.inputs}, + "outputs": {value.name: _contract(value) for value in model.graph.outputs}, + }, + } + + +def _effect(consumes: str, produces: str) -> dict[str, str]: + return {"consumes": consumes, "produces": produces} + + +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, + "effects": effects or {}, + } + + +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") + decoder_name, decoder = next(iter(pkg.items())) + attach_policy_components( + pkg, + PolicyCapabilities( + sampler=sampler, + eos_termination=True, + token_state_update=True, + ), + ) + + inputs = list(decoder.graph.inputs) + outputs = list(decoder.graph.outputs) + token_input = 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" + ) + + workflow_inputs: dict[str, Any] = {} + setup_decoder_inputs: dict[str, str] = {} + body_decoder_inputs: dict[str, str] = {} + for value in inputs: + 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 + + batch_dimension = _shape_metadata(_port(token_input))[0] + batch_int = {"dtype": "int64", "rank": 1, "shape": [batch_dimension]} + workflow_inputs.update( + { + "request.max_iterations": { + "contract": batch_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": "application", "name": "eos_token_ids"}, + "required": True, + }, + "loop.iteration": { + "contract": batch_int, + "role": {"kind": "opaque"}, + "source": {"kind": "application", "name": "iteration"}, + "required": True, + }, + "loop.token_slot": { + "contract": { + "dtype": "int64", + "rank": 2, + "shape": [batch_dimension, 1], + }, + "role": {"kind": "opaque"}, + "source": {"kind": "application", "name": "token_slot"}, + "required": True, + }, + } + ) + + cache_pairs: list[tuple[ir.Value, ir.Value]] = [] + output_by_suffix = {value.name: value for value in outputs} + for value in inputs: + if value is token_input: + continue + candidates = [ + value.name.replace("past_key_values", "present"), + value.name.replace("past.", "present."), + ] + present = next( + (output_by_suffix.get(name) for name in candidates if name in output_by_suffix), + None, + ) + if present is not None: + cache_pairs.append((value, present)) + body_decoder_inputs[value.name] = f"state.{value.name}.body" + body_decoder_inputs[token_input.name] = "state.token.body" + + setup_decoder_outputs = {logits_output.name: "decoder.setup.logits"} + body_decoder_outputs = {logits_output.name: "decoder.body.logits"} + state: dict[str, Any] = { + "token": { + "contract": { + "dtype": "int64", + "rank": 2, + "shape": [batch_dimension, 1], + }, + "scope": "invocation", + "initializer": f"request.{token_input.name}", + "recurrence": {"kind": "invariant"}, + } + } + initial_effects = { + "sample": "sample.0", + "termination": "termination.0", + "state": "state.0", + "emit": "emit.0", + "state:token": "state:token.0", + } + carried = [ + { + "cell": "token", + "current": "token.setup", + "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"), + } + ] + for past, present in cache_pairs: + cell = f"cache_{len(carried)}" + 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 + state[cell] = { + "contract": _contract(past), + "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": 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"), + } + ) + + setup = { + "kind": "sequence", + "nodes": [ + _invoke(decoder_name, setup_decoder_inputs, setup_decoder_outputs), + _invoke( + "token_sampler", + {"logits": "decoder.setup.logits"}, + {"token": "sample.setup"}, + {"sample": _effect("sample.0", "sample.1")}, + ), + _invoke( + "token_state_update", + { + "current": "loop.token_slot", + "update": "sample.setup", + }, + {"next": "token.setup"}, + {"state": _effect("state.0", "state.1")}, + ), + ], + } + body = { + "kind": "sequence", + "nodes": [ + _invoke(decoder_name, body_decoder_inputs, body_decoder_outputs), + _invoke( + "token_sampler", + {"logits": "decoder.body.logits"}, + {"token": "sample.body"}, + {"sample": _effect("sample.1", "sample.2")}, + ), + _invoke( + "token_state_update", + {"current": "state.token.body", "update": "sample.body"}, + {"next": "token.body"}, + {"state": _effect("state.1", "state.2")}, + ), + _invoke( + "termination", + { + "token_ids": "sample.body", + "eos_ids": "package.eos_ids", + "iteration": "loop.iteration", + "max_iterations": "request.max_iterations", + }, + {"done": "loop.done"}, + {"termination": _effect("termination.0", "termination.1")}, + ), + { + "kind": "emit", + "value": "sample.body", + "output": "tokens", + "mode": "append", + "effect_name": "emit", + "effect": _effect("emit.0", "emit.1"), + }, + ], + } + + use_subfolders = len(pkg) > 1 + artifact = f"{decoder_name}/model.onnx" if use_subfolders else "model.onnx" + workflow = { + "manifest": { + "ir_version": "1.0", + "onnx_opsets": {"ai.onnx": OPSET_VERSION}, + "capabilities": [ + "workflow_ssa", + "linear_effects", + "nested_control_flow", + "typed_emit", + ], + }, + "inputs": workflow_inputs, + "outputs": { + "tokens": { + "contract": batch_int, + "role": "tokens", + "stage": "pre_adapter", + } + }, + "components": {decoder_name: _component(decoder, artifact)}, + "state": state, + "initial_effects": initial_effects, + "graph": { + "kind": "loop", + "setup": setup, + "body": body, + "condition": "loop.done", + "max_iterations": "request.max_iterations", + "carried": carried, + }, + } + metadata = {"schema_version": "1.0", "pipeline": {"workflow": workflow}} + add_policy_components_to_workflow(metadata, pkg) + return metadata + + +def write_decoder_workflow_metadata( + pkg: Any, + output_dir: str, + config: Any, +) -> str: + """Write decoder workflow metadata and policy artifacts.""" + os.makedirs(output_dir, exist_ok=True) + metadata = build_decoder_workflow_metadata(pkg, config) + pkg.save_policy_components(output_dir) + path = os.path.join(output_dir, "inference_metadata.yaml") + with open(path, "w", encoding="utf-8") as handle: + yaml.safe_dump(metadata, handle, sort_keys=False) + return path From d2e08107e82847a1b63d57196e0de6707559096e Mon Sep 17 00:00:00 2001 From: justinchuby Date: Wed, 12 Aug 2026 19:06:36 +0000 Subject: [PATCH 004/151] Migrate masked diffusion metadata to SSA workflow Replace the legacy masked-diffusion strategy metadata with the generic workflow contract. Thread token, mask, RNG, effect, and loop state through the denoiser and masked-update ONNX components, and emit final tokens through typed workflow output. Pin the PR #828 schema for exact validation and add runtime parity, metadata, dispatcher, and full-suite coverage. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: justinchuby --- src/mobius/generation/_policy_components.py | 25 +- .../generation/_policy_components_test.py | 4 + .../integrations/onnx_genai/__init__.py | 4 +- .../integrations/onnx_genai/auto_export.py | 33 + .../onnx_genai/auto_export_test.py | 33 +- .../onnx_genai/inference_metadata.py | 70 - .../onnx_genai/inference_metadata_test.py | 73 - .../onnx_genai/workflow_metadata.py | 261 + .../onnx_genai/workflow_metadata_test.py | 99 + src/mobius/models/llada.py | 10 +- src/mobius/models/llada_test.py | 44 +- src/mobius/tasks/_masked_diffusion.py | 9 +- tests/schemas/onnx_genai_4c3c4b6.schema.json | 6586 +++++++++++++++++ 13 files changed, 7070 insertions(+), 181 deletions(-) create mode 100644 src/mobius/integrations/onnx_genai/workflow_metadata_test.py create mode 100644 tests/schemas/onnx_genai_4c3c4b6.schema.json diff --git a/src/mobius/generation/_policy_components.py b/src/mobius/generation/_policy_components.py index edf0e3ca0..98edc0b2d 100644 --- a/src/mobius/generation/_policy_components.py +++ b/src/mobius/generation/_policy_components.py @@ -334,20 +334,38 @@ def build_euler_solver_step() -> PolicyComponent: def build_masked_token_update() -> PolicyComponent: - """Build deterministic replacement of currently masked token positions.""" + """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"]) masked = builder.input("masked", ir.DataType.BOOL, ["batch", "sequence"]) step = builder.input("step", ir.DataType.INT64, ["batch"]) + seed = builder.input("seed", ir.DataType.INT64, ["batch"]) + offset = builder.input("offset", ir.DataType.INT64, ["batch"]) updated = op.Where(masked, proposed, current) # Consume the declared step without changing values; schedules that remask # tokens can be expressed by a richer artifact with the same semantic ports. updated = op.Add(updated, op.Unsqueeze(op.Mul(step, 0), op.Constant(value_ints=[-1]))) + updated.shape = ir.Shape(["batch", "sequence"]) remaining = op.ConstantOfShape(op.Shape(masked), value=ir.tensor([False])) + 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"]) + next_offset = op.Add( + op.Add(offset, op.Constant(value_int=1)), + 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") return _component( PolicyRole.MASKED_UPDATE, graph, @@ -359,6 +377,11 @@ def build_masked_token_update() -> PolicyComponent: "step": "step", "next_state": "next_state", "next_mask": "next_mask", + "rng": { + "seed": "seed", + "offset": "offset", + "next_offset": "next_offset", + }, "effect": "update", }, "update", diff --git a/src/mobius/generation/_policy_components_test.py b/src/mobius/generation/_policy_components_test.py index 4e3760056..1f4f78754 100644 --- a/src/mobius/generation/_policy_components_test.py +++ b/src/mobius/generation/_policy_components_test.py @@ -91,10 +91,14 @@ def test_masked_update_runtime_parity(tmp_path): "proposed_tokens": np.array([[4, 5, 6]], np.int64), "masked": np.array([[False, True, True]]), "step": np.array([0], np.int64), + "seed": np.array([7], np.int64), + "offset": np.array([11], np.int64), }, ) np.testing.assert_array_equal(outputs[0], [[1, 5, 6]]) np.testing.assert_array_equal(outputs[1], [[False, False, False]]) + np.testing.assert_array_equal(outputs[2], [12]) + np.testing.assert_array_equal(outputs[3], [True]) def test_speculative_acceptance_prefix_runtime(tmp_path): diff --git a/src/mobius/integrations/onnx_genai/__init__.py b/src/mobius/integrations/onnx_genai/__init__.py index c50f13bc1..7f11b5e08 100644 --- a/src/mobius/integrations/onnx_genai/__init__.py +++ b/src/mobius/integrations/onnx_genai/__init__.py @@ -53,7 +53,6 @@ add_policy_components_to_workflow, build_audio_codec_pipeline_metadata, build_diffusion_pipeline_metadata, - build_language_diffusion_pipeline_metadata, build_multimodal_pipeline_metadata, build_speech_to_text_pipeline_metadata, build_tts_pipeline_metadata, @@ -66,7 +65,9 @@ ) from mobius.integrations.onnx_genai.workflow_metadata import ( build_decoder_workflow_metadata, + build_language_diffusion_pipeline_metadata, write_decoder_workflow_metadata, + write_language_diffusion_workflow_metadata, ) __all__ = [ @@ -93,6 +94,7 @@ "translate_comfyui_workflow_file", "write_decoder_metadata", "write_decoder_workflow_metadata", + "write_language_diffusion_workflow_metadata", "write_diffusion_pipeline_metadata", "write_audio_codec_pipeline_metadata", "write_multimodal_pipeline_metadata", diff --git a/src/mobius/integrations/onnx_genai/auto_export.py b/src/mobius/integrations/onnx_genai/auto_export.py index af81b5797..77d89eed8 100644 --- a/src/mobius/integrations/onnx_genai/auto_export.py +++ b/src/mobius/integrations/onnx_genai/auto_export.py @@ -16,6 +16,7 @@ import os from typing import Any +import onnx_ir as ir import yaml from mobius.integrations.onnx_genai.decoder_metadata import ( @@ -34,6 +35,7 @@ ) from mobius.integrations.onnx_genai.workflow_metadata import ( write_decoder_workflow_metadata, + write_language_diffusion_workflow_metadata, ) _LOGGER = logging.getLogger(__name__) @@ -218,6 +220,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()) @@ -470,6 +491,18 @@ def write_onnx_genai_config( ``scheduler`` / ``guidance_scale`` set the loop. """ os.makedirs(output_dir, exist_ok=True) + if _looks_like_language_diffusion(pkg): + path = write_language_diffusion_workflow_metadata( + pkg, + output_dir, + num_inference_steps=num_inference_steps, + ) + artifacts = {"inference_metadata": path} + tokenizer_path = _write_hf_tokenizer(output_dir, source) + if tokenizer_path is not None: + artifacts["tokenizer"] = tokenizer_path + return artifacts + if _looks_like_diffusion(pkg): is_qwen_image_edit = getattr(getattr(pkg, "config", None), "model_type", None) == ( "qwen_image_edit" diff --git a/src/mobius/integrations/onnx_genai/auto_export_test.py b/src/mobius/integrations/onnx_genai/auto_export_test.py index eecafad47..07cbd2803 100644 --- a/src/mobius/integrations/onnx_genai/auto_export_test.py +++ b/src/mobius/integrations/onnx_genai/auto_export_test.py @@ -16,7 +16,11 @@ 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.inference_metadata_test import _decoder_model +from mobius.integrations.onnx_genai.inference_metadata_test import ( + _decoder_model, + _model, + _value, +) @dataclasses.dataclass @@ -68,6 +72,33 @@ def test_dispatch_decoder(tmp_path): assert (tmp_path / "policies" / "token_sampler.onnx").is_file() +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()}) arts = write_onnx_genai_config( diff --git a/src/mobius/integrations/onnx_genai/inference_metadata.py b/src/mobius/integrations/onnx_genai/inference_metadata.py index 1b2ecc172..10b022fc7 100644 --- a/src/mobius/integrations/onnx_genai/inference_metadata.py +++ b/src/mobius/integrations/onnx_genai/inference_metadata.py @@ -1937,76 +1937,6 @@ 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``. - - Returns: - A dict with a top-level ``pipeline`` key, ready to serialize to - ``inference_metadata.yaml``. - """ - 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 - - 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} - - def build_diffusion_pipeline_metadata( *, num_inference_steps: int, diff --git a/src/mobius/integrations/onnx_genai/inference_metadata_test.py b/src/mobius/integrations/onnx_genai/inference_metadata_test.py index db67e52ff..6b2515281 100644 --- a/src/mobius/integrations/onnx_genai/inference_metadata_test.py +++ b/src/mobius/integrations/onnx_genai/inference_metadata_test.py @@ -31,7 +31,6 @@ add_explicit_package_io, 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, @@ -1749,78 +1748,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( diff --git a/src/mobius/integrations/onnx_genai/workflow_metadata.py b/src/mobius/integrations/onnx_genai/workflow_metadata.py index d2b4e2f91..7430f8ed3 100644 --- a/src/mobius/integrations/onnx_genai/workflow_metadata.py +++ b/src/mobius/integrations/onnx_genai/workflow_metadata.py @@ -352,6 +352,248 @@ def build_decoder_workflow_metadata( 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 = {"dtype": "int64", "rank": 1, "shape": [batch_dimension]} + 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": batch_int, + "role": { + "kind": "runtime", + "version": "1.0", + "role": "max_iterations", + }, + "source": {"kind": "request", "field": "max_iterations"}, + "required": False, + "default": num_inference_steps, + }, + "loop.iteration": { + "contract": batch_int, + "role": {"kind": "opaque"}, + "source": {"kind": "application", "name": "iteration"}, + "required": True, + }, + } + + 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, + offset: str, + prefix: str, + effect_in: str, + effect_out: str, + ) -> dict[str, Any]: + return _invoke( + "masked_update", + { + "current_tokens": tokens, + "proposed_tokens": f"{prefix}.proposal", + "masked": mask, + "step": "loop.iteration", + "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", + }, + {"update": _effect(effect_in, effect_out)}, + ) + + setup = { + "kind": "sequence", + "nodes": [ + denoiser_invoke("request.input_ids", "denoiser.setup"), + update_invoke( + "request.input_ids", + "request.mask", + "request.rng_offset", + "denoiser.setup", + "update.0", + "update.1", + ), + ], + } + body = { + "kind": "sequence", + "nodes": [ + denoiser_invoke("state.tokens.body", "denoiser.body"), + update_invoke( + "state.tokens.body", + "state.mask.body", + "state.rng_offset.body", + "denoiser.body", + "update.1", + "update.2", + ), + { + "kind": "emit", + "value": "denoiser.body.tokens", + "output": "tokens", + "mode": "replace", + "effect_name": "emit", + "effect": _effect("emit.0", "emit.1"), + }, + ], + } + + state_specs = { + "tokens": (token_contract, "request.input_ids", "denoiser.setup.tokens"), + "mask": (mask_contract, "request.mask", "denoiser.setup.mask"), + "rng_offset": (batch_int, "request.rng_offset", "denoiser.setup.rng_offset"), + } + 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": { + "ir_version": "1.0", + "onnx_opsets": {"ai.onnx": OPSET_VERSION}, + "capabilities": [ + "workflow_ssa", + "linear_effects", + "nested_control_flow", + "typed_emit", + ], + }, + "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.done", + "max_iterations": "request.max_iterations", + "carried": carried, + }, + } + metadata = {"schema_version": "1.0", "pipeline": {"workflow": workflow}} + add_policy_components_to_workflow(metadata, pkg) + return metadata + + def write_decoder_workflow_metadata( pkg: Any, output_dir: str, @@ -365,3 +607,22 @@ def write_decoder_workflow_metadata( with open(path, "w", encoding="utf-8") as handle: yaml.safe_dump(metadata, handle, sort_keys=False) 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) + path = os.path.join(output_dir, "inference_metadata.yaml") + with open(path, "w", encoding="utf-8") as handle: + yaml.safe_dump(metadata, handle, sort_keys=False) + 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..918cb009c --- /dev/null +++ b/src/mobius/integrations/onnx_genai/workflow_metadata_test.py @@ -0,0 +1,99 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT License. + +from __future__ import annotations + +import json +from pathlib import Path + +import jsonschema +import onnx_ir as ir +import pytest + +from mobius._model_package import ModelPackage +from mobius.integrations.onnx_genai.workflow_metadata import ( + build_language_diffusion_pipeline_metadata, +) + + +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 _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"]["recurrence"] == {"kind": "invariant"} + assert workflow["state"]["tokens"]["contract"]["shape"] == ["batch", "sequence"] + assert workflow["state"]["rng_offset"]["contract"]["shape"] == ["batch"] + assert workflow["components"]["masked_update"]["policy"] == { + "role": "masked_update", + "state": "current_tokens", + "proposal": "proposed_tokens", + "mask": "masked", + "step": "step", + "next_state": "next_state", + "next_mask": "next_mask", + "rng": { + "seed": "seed", + "offset": "offset", + "next_offset": "next_offset", + }, + "effect": "update", + } + + graph = workflow["graph"] + assert graph["kind"] == "loop" + assert graph["condition"] == "denoiser.body.done" + assert graph["max_iterations"] == "request.max_iterations" + assert [node["component"] for node in graph["setup"]["nodes"]] == [ + "model", + "masked_update", + ] + assert [node["kind"] for node in graph["body"]["nodes"]] == [ + "invoke", + "invoke", + "emit", + ] + assert graph["body"]["nodes"][-1]["mode"] == "replace" + + +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 test_language_diffusion_matches_pr_828_schema(): + schema_path = ( + Path(__file__).parents[4] / "tests" / "schemas" / "onnx_genai_4c3c4b6.schema.json" + ) + with schema_path.open(encoding="utf-8") as handle: + schema = json.load(handle) + metadata = build_language_diffusion_pipeline_metadata( + _masked_denoiser_package(), + num_inference_steps=8, + ) + jsonschema.validate(instance=metadata, schema=schema) 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..c1e9917d1 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,11 @@ 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["graph"]["kind"] == "loop" + assert workflow["components"]["masked_update"]["policy"]["role"] == "masked_update" 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/tests/schemas/onnx_genai_4c3c4b6.schema.json b/tests/schemas/onnx_genai_4c3c4b6.schema.json new file mode 100644 index 000000000..53a8c5b94 --- /dev/null +++ b/tests/schemas/onnx_genai_4c3c4b6.schema.json @@ -0,0 +1,6586 @@ +{ + "$defs": { + "AbsentInputKind": { + "description": "Supported absent-input fallback kinds.", + "oneOf": [ + { + "const": "zeros", + "description": "Materialize a zero-initialized tensor.", + "type": "string" + } + ] + }, + "AbsentInputSpec": { + "description": "Explicit tensor fallback for an absent optional graph input.", + "properties": { + "kind": { + "$ref": "#/$defs/AbsentInputKind", + "description": "Fallback materialization kind." + }, + "shape": { + "description": "Runtime-resolved shape of the fallback tensor.", + "items": { + "$ref": "#/$defs/TensorDimension" + }, + "type": "array" + } + }, + "required": [ + "kind", + "shape" + ], + "type": "object" + }, + "AcceptanceMethod": { + "description": "Speculative acceptance-rule vocabulary.", + "oneOf": [ + { + "enum": [ + "rejection_sampling", + "greedy", + "typical" + ], + "title": "Known standard value" + }, + { + "not": { + "enum": [ + "rejection_sampling", + "greedy", + "typical" + ] + }, + "title": "Forward-compatible extension value", + "type": "string" + } + ], + "type": "string" + }, + "AttentionConfig": { + "description": "Build-time attention architecture and dimensions.", + "properties": { + "fallback_behavior": { + "anyOf": [ + { + "$ref": "#/$defs/AttentionType" + }, + { + "type": "null" + } + ], + "description": "Compatible attention behavior for runtimes that do not recognize `type`." + }, + "head_dim": { + "description": "Per-head hidden dimension.", + "format": "uint", + "minimum": 1, + "type": [ + "integer", + "null" + ] + }, + "key_sequence_lengths": { + "anyOf": [ + { + "$ref": "#/$defs/KeySequenceLengthsSpec" + }, + { + "type": "null" + } + ], + "description": "Representation compatibility for the attention key-sequence lengths.\n\nAbsent means the canonical contiguous `int32 [batch_size]` representation\nis required." + }, + "num_attention_heads": { + "description": "Number of query/attention heads.", + "format": "uint", + "minimum": 1, + "type": [ + "integer", + "null" + ] + }, + "num_kv_heads": { + "description": "Number of key/value heads; required by runtimes that need explicit GQA dimensions.", + "format": "uint", + "minimum": 1, + "type": [ + "integer", + "null" + ] + }, + "sink_tokens": { + "description": "Number of leading \"attention sink\" tokens always retained alongside the\nsliding window (StreamingLLM). Only meaningful when `sliding_window` is\nset; `null` or `0` disables sink retention. These first tokens stabilize\nthe attention distribution and are never evicted by the window.", + "format": "uint", + "minimum": 0, + "type": [ + "integer", + "null" + ] + }, + "sliding_window": { + "description": "Sliding-window length in tokens, or null for full-context attention.", + "format": "uint", + "minimum": 1, + "type": [ + "integer", + "null" + ] + }, + "type": { + "$ref": "#/$defs/AttentionType", + "description": "Attention architecture.\n\nCanonical values include `multi_head`, `grouped_query`, and\n`multi_latent`; future values are allowed when paired with a usable\n`fallback_behavior`." + } + }, + "required": [ + "type" + ], + "type": "object" + }, + "AttentionType": { + "description": "Attention architecture vocabulary with an extension branch.", + "oneOf": [ + { + "enum": [ + "multi_head", + "multi_head_attention", + "grouped_query", + "group_query_attention", + "grouped_query_attention", + "gqa", + "multi_latent", + "multi_latent_attention", + "mla" + ], + "title": "Known standard value" + }, + { + "not": { + "enum": [ + "multi_head", + "multi_head_attention", + "grouped_query", + "group_query_attention", + "grouped_query_attention", + "gqa", + "multi_latent", + "multi_latent_attention", + "mla" + ] + }, + "title": "Forward-compatible extension value", + "type": "string" + } + ], + "type": "string" + }, + "BatchingContract": { + "additionalProperties": false, + "properties": { + "batch_axis": { + "format": "uint", + "minimum": 0, + "type": "integer" + }, + "continuous": { + "default": false, + "type": "boolean" + }, + "max_batch_size": { + "format": "uint", + "minimum": 0, + "type": [ + "integer", + "null" + ] + } + }, + "required": [ + "batch_axis" + ], + "type": "object" + }, + "ChunkedPrefillConfig": { + "description": "Runtime chunked-prefill preference.", + "properties": { + "chunk_size": { + "description": "Preferred number of prompt tokens processed in each prefill chunk.", + "format": "uint", + "minimum": 1, + "type": [ + "integer", + "null" + ] + } + }, + "type": "object" + }, + "ComponentImplementation": { + "oneOf": [ + { + "additionalProperties": false, + "properties": { + "artifact": { + "type": "string" + }, + "kind": { + "const": "onnx", + "type": "string" + } + }, + "required": [ + "kind", + "artifact" + ], + "type": "object" + }, + { + "additionalProperties": false, + "properties": { + "abi": { + "type": "string" + }, + "artifact": { + "type": [ + "string", + "null" + ] + }, + "custom_ops": { + "additionalProperties": { + "type": "string" + }, + "default": {}, + "type": "object" + }, + "kind": { + "const": "adapter", + "type": "string" + }, + "version": { + "type": "string" + } + }, + "required": [ + "kind", + "abi", + "version" + ], + "type": "object" + }, + { + "additionalProperties": false, + "properties": { + "kind": { + "const": "binding", + "type": "string" + } + }, + "required": [ + "kind" + ], + "type": "object" + } + ] + }, + "ComponentPorts": { + "additionalProperties": false, + "description": "Explicit input/output ports of one executable component.", + "properties": { + "inputs": { + "additionalProperties": { + "$ref": "#/$defs/TensorContract" + }, + "default": {}, + "type": "object" + }, + "outputs": { + "additionalProperties": { + "$ref": "#/$defs/TensorContract" + }, + "default": {}, + "type": "object" + } + }, + "type": "object" + }, + "ControlFlow": { + "description": "Generic package control-flow algebra.", + "oneOf": [ + { + "additionalProperties": false, + "properties": { + "kind": { + "const": "sequence", + "type": "string" + }, + "steps": { + "items": { + "$ref": "#/$defs/ControlFlow" + }, + "type": "array" + } + }, + "required": [ + "kind", + "steps" + ], + "type": "object" + }, + { + "additionalProperties": false, + "properties": { + "component": { + "type": "string" + }, + "kind": { + "const": "invoke", + "type": "string" + }, + "when": { + "anyOf": [ + { + "$ref": "#/$defs/Predicate" + }, + { + "type": "null" + } + ] + } + }, + "required": [ + "kind", + "component" + ], + "type": "object" + }, + { + "additionalProperties": false, + "properties": { + "body": { + "$ref": "#/$defs/ControlFlow" + }, + "carried": { + "default": [], + "items": { + "$ref": "#/$defs/LoopCarry" + }, + "type": "array" + }, + "kind": { + "const": "loop", + "type": "string" + }, + "step_program": { + "type": [ + "string", + "null" + ] + }, + "termination": { + "$ref": "#/$defs/Termination" + } + }, + "required": [ + "kind", + "body", + "termination" + ], + "type": "object" + }, + { + "additionalProperties": false, + "properties": { + "cases": { + "additionalProperties": { + "$ref": "#/$defs/ControlFlow" + }, + "type": "object" + }, + "default": { + "anyOf": [ + { + "$ref": "#/$defs/ControlFlow" + }, + { + "type": "null" + } + ] + }, + "kind": { + "const": "branch", + "type": "string" + }, + "predicate": { + "$ref": "#/$defs/Predicate" + } + }, + "required": [ + "kind", + "predicate", + "cases" + ], + "type": "object" + } + ] + }, + "DType": { + "description": "Scalar dtype vocabulary with common ONNX and runtime aliases.", + "oneOf": [ + { + "enum": [ + "float32", + "fp32", + "float16", + "fp16", + "half", + "bfloat16", + "bf16", + "float8_e4m3fn", + "fp8_e4m3fn", + "float8_e4m3", + "fp8_e4m3", + "float8_e5m2", + "fp8_e5m2", + "int8", + "uint8", + "int4", + "uint4" + ], + "title": "Known standard value" + }, + { + "not": { + "enum": [ + "float32", + "fp32", + "float16", + "fp16", + "half", + "bfloat16", + "bf16", + "float8_e4m3fn", + "fp8_e4m3fn", + "float8_e4m3", + "fp8_e4m3", + "float8_e5m2", + "fp8_e5m2", + "int8", + "uint8", + "int4", + "uint4" + ] + }, + "title": "Forward-compatible extension value", + "type": "string" + } + ], + "type": "string" + }, + "DataflowEdge": { + "description": "Directed connection between two pipeline component ports.", + "properties": { + "device_transfer": { + "description": "Whether the runtime must move the value between execution devices.", + "type": [ + "boolean", + "null" + ] + }, + "dtype": { + "anyOf": [ + { + "$ref": "#/$defs/TensorDType" + }, + { + "type": "null" + } + ], + "description": "Scalar or logical data type at the component boundary." + }, + "from": { + "description": "Source package input or endpoint in `component.output_name` form.", + "examples": [ + "encoder.hidden_states" + ], + "pattern": "^[^.]+(?:\\.[^.]+)?$", + "type": "string" + }, + "to": { + "description": "Destination package output or endpoint in `component.input_name` form.", + "examples": [ + "decoder.encoder_hidden_states" + ], + "pattern": "^[^.]+(?:\\.[^.]+)?$", + "type": "string" + } + }, + "required": [ + "from", + "to" + ], + "type": "object" + }, + "DeviceKind": { + "enum": [ + "cpu", + "cuda", + "direct_ml", + "core_ml", + "web_gpu", + "npu" + ], + "type": "string" + }, + "DevicePreference": { + "description": "Execution-device preference vocabulary.", + "oneOf": [ + { + "enum": [ + "auto", + "cpu", + "cuda", + "rocm", + "directml", + "coreml", + "webgpu", + "npu" + ], + "title": "Known standard value" + }, + { + "not": { + "enum": [ + "auto", + "cpu", + "cuda", + "rocm", + "directml", + "coreml", + "webgpu", + "npu" + ] + }, + "title": "Forward-compatible extension value", + "type": "string" + } + ], + "type": "string" + }, + "DraftConfig": { + "description": "Draft-token producer configuration.", + "properties": { + "depth": { + "description": "Self-speculative early-exit depth.", + "format": "uint", + "minimum": 1, + "type": [ + "integer", + "null" + ] + }, + "heads": { + "description": "Named draft-head layout or selection.", + "type": [ + "string", + "null" + ] + }, + "ngram": { + "description": "Runtime-specific n-gram or prompt-lookup configuration." + }, + "producer": { + "$ref": "#/$defs/DraftProducer", + "description": "Producer family: `draft_model`, `self_speculative`, `ngram`, or `extra_heads`." + }, + "session": { + "description": "Named runtime session or pipeline component used as the producer.", + "type": [ + "string", + "null" + ] + } + }, + "required": [ + "producer" + ], + "type": "object" + }, + "DraftProducer": { + "description": "Speculative draft-producer vocabulary.", + "oneOf": [ + { + "enum": [ + "draft_model", + "self_speculative", + "ngram", + "extra_heads" + ], + "title": "Known standard value" + }, + { + "not": { + "enum": [ + "draft_model", + "self_speculative", + "ngram", + "extra_heads" + ] + }, + "title": "Forward-compatible extension value", + "type": "string" + } + ], + "type": "string" + }, + "EffectTransition": { + "additionalProperties": false, + "properties": { + "consumes": { + "type": "string" + }, + "produces": { + "type": "string" + } + }, + "required": [ + "consumes", + "produces" + ], + "type": "object" + }, + "ForkPrecisionPolicy": { + "description": "KV fork-precision policy vocabulary.", + "oneOf": [ + { + "enum": [ + "inherit", + "highest", + "independent" + ], + "title": "Known standard value" + }, + { + "not": { + "enum": [ + "inherit", + "highest", + "independent" + ] + }, + "title": "Forward-compatible extension value", + "type": "string" + } + ], + "type": "string" + }, + "GenerationDefaults": { + "description": "Author-declared text-generation defaults (sampling and beam search).\n\nMirrors the `search` section of an onnxruntime-genai `genai_config.json`.\nEvery field is optional so only values the author declared are carried over.", + "properties": { + "diversity_penalty": { + "description": "Diversity penalty for diverse beam groups.", + "format": "float", + "type": [ + "number", + "null" + ] + }, + "do_sample": { + "description": "Whether to randomize sampling through `top_k`/`top_p` (else greedy).", + "type": [ + "boolean", + "null" + ] + }, + "early_stopping": { + "description": "Whether beam search stops once enough beams have finished.", + "type": [ + "boolean", + "null" + ] + }, + "length_penalty": { + "description": "Exponential length penalty used with beam search.", + "format": "float", + "type": [ + "number", + "null" + ] + }, + "max_length": { + "description": "Maximum final sequence length.", + "format": "uint", + "minimum": 0, + "type": [ + "integer", + "null" + ] + }, + "min_length": { + "description": "Minimum final sequence length.", + "format": "uint", + "minimum": 0, + "type": [ + "integer", + "null" + ] + }, + "no_repeat_ngram_size": { + "description": "Disallow repeating n-grams of this size (`0` = disabled).", + "format": "uint", + "minimum": 0, + "type": [ + "integer", + "null" + ] + }, + "num_beams": { + "description": "Number of beams for beam search (`1` = no beam search).", + "format": "uint", + "minimum": 0, + "type": [ + "integer", + "null" + ] + }, + "num_return_sequences": { + "description": "Number of sequences returned after search.", + "format": "uint", + "minimum": 0, + "type": [ + "integer", + "null" + ] + }, + "repetition_penalty": { + "description": "Penalty applied to already-generated tokens (`1.0` = no penalty).", + "format": "float", + "type": [ + "number", + "null" + ] + }, + "temperature": { + "description": "Softmax temperature applied before sampling.", + "format": "float", + "type": [ + "number", + "null" + ] + }, + "top_k": { + "description": "Number of highest-probability tokens kept for top-k filtering.", + "format": "uint", + "minimum": 0, + "type": [ + "integer", + "null" + ] + }, + "top_p": { + "description": "Nucleus (top-p) cumulative-probability threshold.", + "format": "float", + "type": [ + "number", + "null" + ] + } + }, + "type": "object" + }, + "HardwareRequirements": { + "description": "Model-side hardware requirements and distribution-matching hints.", + "properties": { + "beneficial_dtypes": { + "description": "Dtypes that improve performance or memory use but are not mandatory.", + "items": { + "$ref": "#/$defs/DType" + }, + "type": [ + "array", + "null" + ] + }, + "kv_cache_memory_per_1k_tokens_mb": { + "description": "Estimated KV-cache memory in MiB per 1,000 cached tokens.", + "format": "float", + "minimum": 0.0, + "type": [ + "number", + "null" + ] + }, + "min_memory_gb": { + "description": "Minimum aggregate accelerator or system memory in GiB.", + "format": "float", + "minimum": 0.0, + "type": [ + "number", + "null" + ] + }, + "min_tp_degree": { + "description": "Minimum useful tensor-parallel degree when tensor parallelism is selected.", + "format": "uint", + "minimum": 1, + "type": [ + "integer", + "null" + ] + }, + "required_dtypes": { + "description": "Dtypes the selected device or execution provider must support.", + "items": { + "$ref": "#/$defs/DType" + }, + "type": [ + "array", + "null" + ] + }, + "supports_tensor_parallel": { + "description": "Whether the model can be partitioned with tensor parallelism.", + "type": [ + "boolean", + "null" + ] + } + }, + "type": "object" + }, + "ImageCorrespondence": { + "description": "Prompt-placeholder to image correspondence vocabulary.", + "oneOf": [ + { + "enum": [ + "prompt_order", + "explicit_indices" + ], + "title": "Known standard value" + }, + { + "not": { + "enum": [ + "prompt_order", + "explicit_indices" + ] + }, + "title": "Forward-compatible extension value", + "type": "string" + } + ], + "type": "string" + }, + "ImageOutputBinding": { + "description": "One named tensor output produced by an image preprocessing program.\n\nThe output binds a generic content role to an ARBITRARY endpoint name with a\nDECLARED dtype. Neither the name nor the content role is inferred from a model\nidentity, and the dtype is always explicit rather than derived from the model.", + "properties": { + "content": { + "$ref": "#/$defs/ImageOutputContent", + "description": "Generic content role this tensor carries (pixels, coordinates, grid,\noriginal size, or validity mask) — never a model-family label." + }, + "dtype": { + "$ref": "#/$defs/TensorDType", + "description": "Declared output dtype. Always explicit; never inferred from the model." + }, + "name": { + "description": "Arbitrary pipeline endpoint name this tensor is bound to (model DATA).", + "examples": [ + "vision_encoder.pixel_values" + ], + "minLength": 1, + "type": "string" + }, + "optional": { + "description": "Whether the runtime may omit this output when a model does not need it.", + "type": [ + "boolean", + "null" + ] + }, + "pad_value": { + "description": "Optional sentinel/pad value for padded entries (e.g. `-1` coordinates).", + "format": "double", + "type": [ + "number", + "null" + ] + }, + "source": { + "description": "Named value produced by a transform.\n\nAbsent preserves the legacy content-derived binding behavior.", + "minLength": 1, + "type": [ + "string", + "null" + ] + } + }, + "required": [ + "name", + "content", + "dtype" + ], + "type": "object" + }, + "ImageOutputContent": { + "description": "Generic image-output content-role vocabulary.", + "oneOf": [ + { + "enum": [ + "pixels", + "patch_coordinates", + "grid_dimensions", + "original_size", + "transformed_size", + "validity_mask" + ], + "title": "Known standard value" + }, + { + "not": { + "enum": [ + "pixels", + "patch_coordinates", + "grid_dimensions", + "original_size", + "transformed_size", + "validity_mask" + ] + }, + "title": "Forward-compatible extension value", + "type": "string" + } + ], + "type": "string" + }, + "ImagePreprocessingProgram": { + "description": "Generic image preprocessing program: an ordered transform pipeline plus the\nnamed tensor outputs it emits.\n\nThe program is expressed entirely as parameterized, architecture-neutral\ndata. Transform operations are generic (decode, resize, rescale, normalize,\ntile, patchify, pad); outputs bind a produced tensor to an ARBITRARY pipeline\nendpoint name with a DECLARED dtype. A model may name an output\n`pixel_position_ids`, `image_grid_thw`, or anything else — that string is\ndata carried in the model's metadata, never a branch in the runtime.", + "properties": { + "outputs": { + "description": "Named tensor outputs the program emits, each bound to a pipeline endpoint.", + "items": { + "$ref": "#/$defs/ImageOutputBinding" + }, + "minItems": 1, + "type": "array" + }, + "transforms": { + "description": "Ordered list of generic transform operations applied to decoded pixels.", + "items": { + "$ref": "#/$defs/ImageTransform" + }, + "type": "array" + } + }, + "required": [ + "outputs" + ], + "type": "object" + }, + "ImageSizeSpec": { + "anyOf": [ + { + "description": "A single edge length applied to both dimensions.", + "format": "uint32", + "minimum": 0, + "type": "integer" + }, + { + "description": "Explicit width and height.", + "properties": { + "height": { + "description": "Target height in pixels.", + "format": "uint32", + "minimum": 0, + "type": "integer" + }, + "width": { + "description": "Target width in pixels.", + "format": "uint32", + "minimum": 0, + "type": "integer" + } + }, + "required": [ + "width", + "height" + ], + "type": "object" + } + ], + "description": "A square size or an explicit width/height for an image transform." + }, + "ImageTokenCountSource": { + "description": "Image token-count source vocabulary.", + "oneOf": [ + { + "enum": [ + "per_tile", + "per_patch", + "from_grid" + ], + "title": "Known standard value" + }, + { + "not": { + "enum": [ + "per_tile", + "per_patch", + "from_grid" + ] + }, + "title": "Forward-compatible extension value", + "type": "string" + } + ], + "type": "string" + }, + "ImageTransform": { + "description": "One generic image transform operation.\n\n`op` selects the operation from a generic vocabulary; the remaining fields\nare the parameters that operation reads (only the relevant ones are set).\nEvery parameter is model DATA — concrete sizes, patch sizes, means, and so on\nlive in a model's fixture, never as constants baked into this schema.", + "properties": { + "canvas_pad_value": { + "description": "RGB canvas fill value applied before dynamic tiling.", + "format": "double", + "type": [ + "number", + "null" + ] + }, + "channel_order": { + "description": "Flattened patch feature order (`channels_first` or `channels_last`).", + "minLength": 1, + "type": [ + "string", + "null" + ] + }, + "coordinate_order": { + "description": "Patch-coordinate component order (`yx` or `xy`).", + "minLength": 1, + "type": [ + "string", + "null" + ] + }, + "flatten": { + "description": "Whether `patchify` flattens each patch into a single feature vector.", + "type": [ + "boolean", + "null" + ] + }, + "include_thumbnail": { + "description": "Whether a `tile` operation also emits a global thumbnail tile.", + "type": [ + "boolean", + "null" + ] + }, + "inputs": { + "description": "Named values consumed by this transform.\n\nAbsent means the operation consumes the immediately preceding value.\nExplicit names allow branching programs without tensor-name heuristics.", + "items": { + "minLength": 1, + "type": "string" + }, + "minItems": 1, + "type": [ + "array", + "null" + ] + }, + "interpolation": { + "description": "Interpolation filter for a `resize` operation — generic string data.", + "minLength": 1, + "type": [ + "string", + "null" + ] + }, + "mask_patch_size": { + "description": "Pixel edge represented by one validity-mask cell.", + "format": "uint", + "minimum": 1, + "type": [ + "integer", + "null" + ] + }, + "max_patches": { + "description": "Maximum number of spatial patches for a patch-budget resize.", + "format": "uint", + "minimum": 1, + "type": [ + "integer", + "null" + ] + }, + "max_pixels": { + "description": "Maximum pixel area for an aspect-preserving `pixel_area` resize.", + "format": "uint", + "minimum": 1, + "type": [ + "integer", + "null" + ] + }, + "max_tiles": { + "description": "Maximum number of local tiles for a `tile` operation.", + "format": "uint", + "minimum": 1, + "type": [ + "integer", + "null" + ] + }, + "mean": { + "description": "Per-channel mean for a `normalize` operation (length is model data).", + "items": { + "format": "float", + "type": "number" + }, + "type": [ + "array", + "null" + ] + }, + "merge_size": { + "description": "Spatial patch-group edge controlling packed patch traversal order.", + "format": "uint", + "minimum": 1, + "type": [ + "integer", + "null" + ] + }, + "min_pixels": { + "description": "Minimum pixel area for an aspect-preserving `pixel_area` resize.", + "format": "uint", + "minimum": 1, + "type": [ + "integer", + "null" + ] + }, + "mode": { + "description": "Resize/crop mode (e.g. `pad`, `crop`, `stretch`) — generic string data.", + "minLength": 1, + "type": [ + "string", + "null" + ] + }, + "op": { + "$ref": "#/$defs/ImageTransformOp", + "description": "Generic operation selector (e.g. `resize`, `normalize`, `patchify`)." + }, + "outputs": { + "description": "Named values produced by this transform.\n\nThese names are processor-local data. Final graph bindings select them\nthrough `ImageOutputBinding::source`.", + "items": { + "minLength": 1, + "type": "string" + }, + "minItems": 1, + "type": [ + "array", + "null" + ] + }, + "pad_value": { + "description": "Fill value for a `pad` operation, or sentinel for padded coordinates.", + "format": "double", + "type": [ + "number", + "null" + ] + }, + "patch_size": { + "description": "Edge length of a square patch for a `patchify` operation.", + "format": "uint", + "minimum": 1, + "type": [ + "integer", + "null" + ] + }, + "pooling_kernel_size": { + "description": "Spatial pooling edge used when resolving a patch-budget resize.", + "format": "uint", + "minimum": 1, + "type": [ + "integer", + "null" + ] + }, + "scale": { + "description": "Scalar multiplier for a `rescale` operation.", + "format": "double", + "type": [ + "number", + "null" + ] + }, + "size": { + "anyOf": [ + { + "$ref": "#/$defs/ImageSizeSpec" + }, + { + "type": "null" + } + ], + "description": "Target size for a `resize` operation." + }, + "size_multiple": { + "description": "Required divisibility of both resized dimensions.", + "format": "uint", + "minimum": 1, + "type": [ + "integer", + "null" + ] + }, + "std": { + "description": "Per-channel standard deviation for a `normalize` operation.", + "items": { + "format": "float", + "type": "number" + }, + "type": [ + "array", + "null" + ] + }, + "target_length": { + "description": "Exact first-axis length produced by a `pad` operation.", + "format": "uint", + "minimum": 1, + "type": [ + "integer", + "null" + ] + }, + "temporal_patch_size": { + "description": "Number of identical temporal frames packed into each spatial patch.", + "format": "uint", + "minimum": 1, + "type": [ + "integer", + "null" + ] + }, + "thumbnail_interpolation": { + "description": "Interpolation filter used specifically for a global thumbnail.", + "minLength": 1, + "type": [ + "string", + "null" + ] + }, + "thumbnail_order": { + "anyOf": [ + { + "$ref": "#/$defs/ThumbnailOrder" + }, + { + "type": "null" + } + ], + "description": "Ordering of a global thumbnail relative to local tiles." + }, + "tile_size": { + "description": "Edge length of a square tile for a `tile` operation.", + "format": "uint", + "minimum": 1, + "type": [ + "integer", + "null" + ] + } + }, + "required": [ + "op" + ], + "type": "object" + }, + "ImageTransformOp": { + "description": "Generic image transform-operation vocabulary.", + "oneOf": [ + { + "enum": [ + "decode", + "decode_rgb", + "convert_rgb", + "resize", + "rescale", + "normalize", + "tile", + "flatten", + "patchify", + "pad", + "emit_original_size", + "emit_transformed_size", + "emit_validity_mask", + "emit_patch_coordinates", + "emit_grid_coordinates" + ], + "title": "Known standard value" + }, + { + "not": { + "enum": [ + "decode", + "decode_rgb", + "convert_rgb", + "resize", + "rescale", + "normalize", + "tile", + "flatten", + "patchify", + "pad", + "emit_original_size", + "emit_transformed_size", + "emit_validity_mask", + "emit_patch_coordinates", + "emit_grid_coordinates" + ] + }, + "title": "Forward-compatible extension value", + "type": "string" + } + ], + "type": "string" + }, + "KeySequenceLengthsSpec": { + "description": "Explicit compatibility rules for attention key-sequence-length metadata.", + "properties": { + "scalar_broadcast": { + "anyOf": [ + { + "$ref": "#/$defs/SequenceLengthScalarBroadcast" + }, + { + "type": "null" + } + ], + "description": "Optional scalar compatibility. `unit_batch` authorizes a contiguous\nrank-0 one-element `int32` tensor only when the attention batch is one." + } + }, + "type": "object" + }, + "KvAxisStrides": { + "description": "Symbolic element stride of each of the four logical KV axes.\n\nThe stride of an axis is the product of the runtime dimensions in its factor\nlist; an **empty** list means unit stride (the innermost, contiguous axis).\nThe innermost axis of every layout the converted kernels honor is\n`head_dim`, whose stride is `1` (empty), because the fp16 read vectorizes\n`head_dim` as `half2` and the fused write addresses it as `dst + d`.\n\nThe two historical layouts map onto this as:\n\n| axis | head-major BNSH | seq-major BSNH |\n|----------|------------------------|-----------------------|\n| batch | `kv_heads·seq·head_dim`| `seq·kv_heads·head_dim`|\n| head | `seq·head_dim` | `head_dim` |\n| seq | `head_dim` | `kv_heads·head_dim` |\n| head_dim | `1` | `1` |", + "properties": { + "batch": { + "default": [], + "description": "Factors of the batch-axis stride.", + "items": { + "$ref": "#/$defs/KvStrideDim" + }, + "type": "array" + }, + "head": { + "default": [], + "description": "Factors of the KV-head-axis stride.", + "items": { + "$ref": "#/$defs/KvStrideDim" + }, + "type": "array" + }, + "head_dim": { + "default": [], + "description": "Factors of the head-dim-axis stride. Unit (empty) for every honored\nlayout.", + "items": { + "$ref": "#/$defs/KvStrideDim" + }, + "type": "array" + }, + "seq": { + "default": [], + "description": "Factors of the sequence-axis (per-token) stride.", + "items": { + "$ref": "#/$defs/KvStrideDim" + }, + "type": "array" + } + }, + "type": "object" + }, + "KvCacheLayout": { + "anyOf": [ + { + "$ref": "#/$defs/KvNamedLayout", + "description": "A readable shorthand for a standard layout. Deserializes from the strings\n`\"head_major_bnsh\"` and `\"seq_major_bsnh\"`; expands to explicit strides\nvia [`KvCacheLayout::resolve_strides`]." + }, + { + "$ref": "#/$defs/KvStrideDescriptor", + "description": "A fully explicit stride descriptor for layouts the named forms cannot\nexpress." + } + ], + "description": "Physical memory layout of a backend's KV cache tensors, as a stride\ndescriptor.\n\nThis is a **per-backend capability**, not a cross-backend constant: the two\nbackends own their KV buffers independently and never read each other's KV\nbytes, so they may store the cache differently. The ONNX Runtime backend\nrequires head-major BNSH (`[batch, kv_heads, seq, head_dim]`) because ORT's\nGroupQueryAttention past/present is BNSH on every dispatch path (Flash,\ncuDNN SDPA, memory-efficient, XQA). The native backend additionally supports\nseq-major BSNH (`[batch, seq, kv_heads, head_dim]`), which makes each token's\nlive prefix contiguous across heads — shrinking the VMM granule floor by the\n`kv_heads` factor, removing growth-triggered graph re-capture (the append\nstride is sequence-length independent), and making page-level prefix sharing\n(#777) practical. Absent preserves the historical head-major behavior.\n\nLayout preference is per-EP and per-platform rather than a global constant,\nand a JIT backend compiles a specialized kernel per descriptor, so this is a\ndescriptor rather than a closed enum: a raw stride tuple is unreadable, so\nthe common cases are still nameable (`head_major_bnsh`, `seq_major_bsnh`)\nwhile an explicit [`KvStrideDescriptor`] expresses anything the named forms\ncannot (e.g. a token-major view).\n\nOn-device, the native backend selects the layout by stamping the `kv_layout`\nattribute (`0` = BNSH, `1` = BSNH) on its GroupQueryAttention nodes; the\nCUDA EP honors it on the fused fp16 single-token decode pair. Seq-major is\nonly enabled end-to-end once the prefill (flash) read is also converted, so\nthe two never disagree about how a shared cache is physically laid out." + }, + "KvCacheOperations": { + "description": "Operational guarantees for mutable KV-cache state.", + "properties": { + "checkpoint_serializable": { + "description": "Whether checkpoints can be serialized for suspend/resume or migration.", + "type": [ + "boolean", + "null" + ] + }, + "fork_precision_policy": { + "anyOf": [ + { + "$ref": "#/$defs/ForkPrecisionPolicy" + }, + { + "type": "null" + } + ], + "description": "Precision policy for a copy-on-write fork, such as `inherit` or `highest`." + }, + "rewind_safe": { + "description": "Whether truncating cache state to an earlier token position is correctness-preserving.", + "type": [ + "boolean", + "null" + ] + } + }, + "type": "object" + }, + "KvCacheSpec": { + "description": "KV-cache storage, precision tolerance, and operational guarantees.", + "properties": { + "native_dtype": { + "anyOf": [ + { + "$ref": "#/$defs/DType" + }, + { + "type": "null" + } + ], + "description": "Native KV scalar dtype produced by the model before optional compression." + }, + "operations": { + "anyOf": [ + { + "$ref": "#/$defs/KvCacheOperations" + }, + { + "type": "null" + } + ], + "description": "Cache mutation and persistence operations known to be safe for this model." + }, + "quantization_tolerance": { + "anyOf": [ + { + "$ref": "#/$defs/KvQuantTolerance" + }, + { + "type": "null" + } + ], + "description": "Independent precision tolerance for key and value tensors." + }, + "sensitive_layers": { + "description": "Layer indices that should retain high precision; negative indices count from the end.", + "items": { + "format": "int32", + "type": "integer" + }, + "type": [ + "array", + "null" + ] + } + }, + "type": "object" + }, + "KvComponentTolerance": { + "description": "Quantization tolerance for one KV-cache component.", + "properties": { + "default": { + "anyOf": [ + { + "$ref": "#/$defs/DType" + }, + { + "type": "null" + } + ], + "description": "Default minimum acceptable scalar dtype for this component." + }, + "per_layer": { + "description": "Layer-specific minimum-precision overrides.", + "items": { + "$ref": "#/$defs/LayerPrecisionOverride" + }, + "type": [ + "array", + "null" + ] + }, + "quantization_axis": { + "anyOf": [ + { + "$ref": "#/$defs/QuantizationAxis" + }, + { + "type": "null" + } + ], + "description": "Quantization scaling axis, such as `per_tensor`, `per_channel`, or `per_token`." + } + }, + "type": "object" + }, + "KvNamedLayout": { + "description": "The named, human-readable KV cache layouts.", + "oneOf": [ + { + "const": "head_major_bnsh", + "description": "Head-major BNSH `[batch, kv_heads, seq, head_dim]`. ORT-compatible; the\ndefault for both backends.", + "type": "string" + }, + { + "const": "seq_major_bsnh", + "description": "Seq-major BSNH `[batch, seq, kv_heads, head_dim]`. Native backend only.", + "type": "string" + } + ] + }, + "KvOwnership": { + "description": "Ownership model for a graph's KV cache inputs.", + "oneOf": [ + { + "const": "owned", + "description": "The graph consumes past KV and emits replacement/extended present KV.", + "type": "string" + }, + { + "const": "shared", + "description": "The graph reads references to KV owned and advanced by another decoder.", + "type": "string" + } + ] + }, + "KvPagingMode": { + "enum": [ + "none", + "paged" + ], + "type": "string" + }, + "KvQuantTolerance": { + "description": "Precision tolerance for key and value cache components.", + "properties": { + "key": { + "anyOf": [ + { + "$ref": "#/$defs/KvComponentTolerance" + }, + { + "type": "null" + } + ], + "description": "Key-cache precision tolerance." + }, + "value": { + "anyOf": [ + { + "$ref": "#/$defs/KvComponentTolerance" + }, + { + "type": "null" + } + ], + "description": "Value-cache precision tolerance." + } + }, + "type": "object" + }, + "KvServiceContract": { + "additionalProperties": false, + "properties": { + "allocation": { + "$ref": "#/$defs/SlotAllocationMode" + }, + "compaction": { + "default": false, + "type": "boolean" + }, + "paging": { + "$ref": "#/$defs/KvPagingMode" + } + }, + "required": [ + "paging", + "allocation" + ], + "type": "object" + }, + "KvStrideDescriptor": { + "description": "A fully explicit KV-cache stride descriptor.\n\nThis is the general form the two named layouts expand into, and the shape a\nfuture layout (e.g. token-major) is expressed in without adding an enum\nvariant. The `reservation_*` fields describe a binding that is a **view into\na larger reservation** rather than the owner of its whole buffer:\ntoken-major stores every layer's tokens in one reservation and hands each\n`(layer, side)` a sub-view, so its per-token (seq) stride is taken over the\nreservation's total token count and its data starts at a non-zero offset.\nBoth historical layouts are whole-buffer bindings: `offset == 0` and no\nreservation override.", + "properties": { + "reservation_offset_elements": { + "description": "Element offset of this binding's first element within the reservation it\nviews. `0` for a binding that owns its whole buffer — the only case the\nconverted kernels honor today.", + "format": "uint64", + "minimum": 0, + "type": "integer" + }, + "reservation_seq_slots": { + "description": "Sequence-axis extent, in token slots, of the reservation this binding\nviews when the reservation is larger than the binding's own\n`cache_capacity`. Absent means the binding spans its own capacity (a\nwhole-buffer binding). Present expresses a token-major view whose seq\nstride collapses the per-`(layer, side)` buffer boundary; not honored by\nthe converted path yet.", + "format": "uint64", + "minimum": 0, + "type": [ + "integer", + "null" + ] + }, + "strides": { + "$ref": "#/$defs/KvAxisStrides", + "description": "Symbolic stride of each logical axis." + } + }, + "required": [ + "strides" + ], + "type": "object" + }, + "KvStrideDim": { + "description": "A runtime KV-cache dimension that an axis stride can be a multiple of.\n\nAbsolute element strides are a serving-time property — they depend on the\n`cache_capacity` a runtime picks — so metadata cannot store them as numbers.\nA stride is therefore stored **symbolically**, as the (unordered) set of\nruntime dimensions it multiplies. The concrete element stride of an axis is\nthe product of the sizes of the dimensions in its factor list.", + "oneOf": [ + { + "const": "kv_heads", + "description": "Number of KV heads (`kv_heads` / `N`).", + "type": "string" + }, + { + "const": "seq_capacity", + "description": "Sequence capacity of the growing axis (`cache_capacity` / `S`).", + "type": "string" + }, + { + "const": "head_dim", + "description": "Per-token head width (`head_dim` / `H`).", + "type": "string" + } + ] + }, + "KvUpdateKind": { + "description": "Paired KV-cache update-semantics vocabulary.", + "oneOf": [ + { + "enum": [ + "append", + "shared_buffer" + ], + "title": "Known standard value" + }, + { + "not": { + "enum": [ + "append", + "shared_buffer" + ] + }, + "title": "Forward-compatible extension value", + "type": "string" + } + ], + "type": "string" + }, + "LayerPrecisionOverride": { + "description": "Minimum precision required by a set of model layers.", + "properties": { + "layers": { + "description": "Non-empty layer-index list; negative indices count from the final layer.", + "items": { + "format": "int32", + "type": "integer" + }, + "minItems": 1, + "type": "array" + }, + "min_precision": { + "$ref": "#/$defs/DType", + "description": "Minimum acceptable scalar dtype for the listed layers." + } + }, + "required": [ + "layers", + "min_precision" + ], + "type": "object" + }, + "LoopCarry": { + "additionalProperties": false, + "properties": { + "from": { + "type": "string" + }, + "state": { + "type": "string" + }, + "to": { + "type": "string" + } + }, + "required": [ + "state", + "from", + "to" + ], + "type": "object" + }, + "LoopStatePair": { + "description": "One fixed-shape loop-carried recurrent-state port pair.\n\nGeneric and architecture-neutral: the runtime zero/other-initializes `input`\non the first step, runs the graph, and copies `output` back into `input` for\nthe next step (`replace` update). This models any fixed recurrent tensor\n(convolution state, linear-attention recurrent state, and so on) without\nreferencing a model family. It is intentionally distinct from growing or\nshared-buffer KV cache, which is declared through `kv_inputs`/`kv_outputs`\nand `kv_update`.", + "properties": { + "init": { + "$ref": "#/$defs/StateInitKind", + "description": "How `input` is initialized before the first step (e.g. `zeros`)." + }, + "input": { + "description": "Graph input port that receives the carried state for this step.", + "minLength": 1, + "type": "string" + }, + "output": { + "description": "Graph output port that produces the next-step state.", + "minLength": 1, + "type": "string" + }, + "update": { + "$ref": "#/$defs/StateUpdateKind", + "description": "How `output` becomes the next step's `input` (fixed state uses `replace`)." + } + }, + "required": [ + "input", + "output", + "init", + "update" + ], + "type": "object" + }, + "MixtureOfExpertsSpec": { + "description": "Explicit sparse mixture-of-experts structure and graph representation.", + "properties": { + "activation": { + "description": "Expert FFN activation name, such as `silu`.", + "minLength": 1, + "type": "string" + }, + "expert_intermediate_size": { + "description": "Intermediate width of each routed expert FFN.", + "format": "uint", + "minimum": 1, + "type": "integer" + }, + "experts_per_token": { + "description": "Number of routed experts selected for each token.", + "format": "uint", + "minimum": 1, + "type": "integer" + }, + "representation": { + "$ref": "#/$defs/MoERepresentation", + "description": "Expert graph representation: `dense_fallback`, `moe`, or `qmoe`." + }, + "routed_expert_count": { + "description": "Number of independently routed experts.", + "format": "uint", + "minimum": 1, + "type": "integer" + }, + "router": { + "$ref": "#/$defs/MoERouterSpec", + "description": "Router scoring, selection, normalization, and scaling semantics." + }, + "shared_expert_count": { + "description": "Number of dense shared experts evaluated for every token.", + "format": "uint", + "minimum": 0, + "type": "integer" + }, + "shared_expert_intermediate_size": { + "description": "Total intermediate width of the always-on shared-expert FFN.", + "format": "uint", + "minimum": 0, + "type": "integer" + } + }, + "required": [ + "representation", + "routed_expert_count", + "shared_expert_count", + "experts_per_token", + "expert_intermediate_size", + "shared_expert_intermediate_size", + "activation", + "router" + ], + "type": "object" + }, + "MoEGroupScore": { + "description": "Group-scoring reduction vocabulary.", + "oneOf": [ + { + "enum": [ + "maximum", + "top_2_sum" + ], + "title": "Known standard value" + }, + { + "not": { + "enum": [ + "maximum", + "top_2_sum" + ] + }, + "title": "Forward-compatible extension value", + "type": "string" + } + ], + "type": "string" + }, + "MoERepresentation": { + "description": "Sparse expert graph representation vocabulary.", + "oneOf": [ + { + "enum": [ + "dense_fallback", + "moe", + "qmoe" + ], + "title": "Known standard value" + }, + { + "not": { + "enum": [ + "dense_fallback", + "moe", + "qmoe" + ] + }, + "title": "Forward-compatible extension value", + "type": "string" + } + ], + "type": "string" + }, + "MoERouterScoreFunction": { + "description": "Router score-operation vocabulary.", + "oneOf": [ + { + "enum": [ + "softmax", + "sigmoid" + ], + "title": "Known standard value" + }, + { + "not": { + "enum": [ + "softmax", + "sigmoid" + ] + }, + "title": "Forward-compatible extension value", + "type": "string" + } + ], + "type": "string" + }, + "MoERouterSelectionMethod": { + "description": "Router expert-selection vocabulary.", + "oneOf": [ + { + "enum": [ + "top_k", + "grouped_top_k", + "sparse_mixer" + ], + "title": "Known standard value" + }, + { + "not": { + "enum": [ + "top_k", + "grouped_top_k", + "sparse_mixer" + ] + }, + "title": "Forward-compatible extension value", + "type": "string" + } + ], + "type": "string" + }, + "MoERouterSpec": { + "allOf": [ + { + "if": { + "properties": { + "selection_method": { + "const": "grouped_top_k" + } + }, + "required": [ + "selection_method" + ] + }, + "then": { + "required": [ + "group_count", + "groups_per_token", + "group_score" + ] + } + } + ], + "description": "Explicit router semantics, kept separate from expert FFN execution.", + "properties": { + "group_count": { + "description": "Number of expert groups considered by grouped selection.", + "format": "uint", + "minimum": 1, + "type": [ + "integer", + "null" + ] + }, + "group_score": { + "anyOf": [ + { + "$ref": "#/$defs/MoEGroupScore" + }, + { + "type": "null" + } + ], + "description": "Reduction used to score a group before group TopK." + }, + "groups_per_token": { + "description": "Number of groups retained per token by grouped selection.", + "format": "uint", + "minimum": 1, + "type": [ + "integer", + "null" + ] + }, + "normalize_weights": { + "description": "Whether selected aggregation weights are normalized to sum to one.", + "type": "boolean" + }, + "scaling_factor": { + "description": "Multiplicative scale applied to final aggregation weights.", + "format": "float", + "minimum": 0.0, + "type": "number" + }, + "score_function": { + "$ref": "#/$defs/MoERouterScoreFunction", + "description": "Elementwise score operation applied to router logits." + }, + "selection_method": { + "$ref": "#/$defs/MoERouterSelectionMethod", + "description": "Expert selection operation applied to the scores." + } + }, + "required": [ + "score_function", + "selection_method", + "normalize_weights", + "scaling_factor" + ], + "type": "object" + }, + "ModelCapabilities": { + "description": "Model properties that are baked into the graph or advertised as configurable.", + "properties": { + "attention": { + "anyOf": [ + { + "$ref": "#/$defs/AttentionConfig" + }, + { + "type": "null" + } + ], + "description": "Attention architecture and dimensions." + }, + "io": { + "anyOf": [ + { + "$ref": "#/$defs/ModelIoSpec" + }, + { + "type": "null" + } + ], + "default": null, + "description": "Explicit graph I/O port bindings for the single-decoder LLM path.\n\nThe runtime binds decode-step inputs and outputs from the declared names.\nA port that is not declared is resolved ONLY from an unambiguous io-shape\nsignal; when the shape is ambiguous the runtime fails with an actionable\nerror naming the exact key to declare, and never guesses from a tensor\nname." + }, + "max_sequence_length": { + "description": "Maximum total sequence length, in tokens.", + "format": "uint", + "minimum": 1, + "type": [ + "integer", + "null" + ] + }, + "mixture_of_experts": { + "anyOf": [ + { + "$ref": "#/$defs/MixtureOfExpertsSpec" + }, + { + "type": "null" + } + ], + "description": "Explicit sparse mixture-of-experts graph and routing contract.\n\nThis describes graph structure, never a model family. Runtimes use the\ndeclared representation and dimensions instead of inferring them from\nnode names, initializer shapes, or architecture strings." + }, + "runtime_configurable": { + "anyOf": [ + { + "$ref": "#/$defs/RuntimeConfigurable" + }, + { + "type": "null" + } + ], + "description": "Features that a serving runtime may configure at load time." + }, + "speculative": { + "anyOf": [ + { + "$ref": "#/$defs/SpeculativeModelInfo" + }, + { + "type": "null" + } + ], + "description": "Built-in draft-head or self-speculative model properties." + }, + "vocab_size": { + "description": "Vocabulary size (rows of the token-embedding / logits table).", + "format": "uint", + "minimum": 1, + "type": [ + "integer", + "null" + ] + } + }, + "type": "object" + }, + "ModelIoSpec": { + "description": "Explicit binding of the graph ports the decode step reads and writes.\n\nEvery field is optional so a model package can declare only the ports its\ngraph exposes. A port left unset is resolved ONLY from an unambiguous\ndtype/shape signal; when the shape cannot disambiguate the port, the runtime\nfails with an actionable error naming the key to declare rather than\ninterpreting a tensor name. A declared port is always authoritative.", + "properties": { + "attention_mask_input": { + "description": "Attention-mask input, if the graph takes one.", + "minLength": 1, + "type": [ + "string", + "null" + ] + }, + "audio_features_input": { + "description": "Raw audio-feature prompt input for an encoder-decoder encoder graph\n(e.g. Whisper `audio_features`, a log-mel `[batch, mels, frames]`\ntensor). Declared on the encoder component; a text encoder-decoder uses\n`token_input` instead.", + "minLength": 1, + "type": [ + "string", + "null" + ] + }, + "cross_kv_inputs": { + "description": "Cross-attention past-KV cache inputs for an encoder-decoder decoder, in\nthe SAME order as `cross_kv_outputs`. These are the encoder-derived KV\ntensors, distinct from the self-attention `kv_inputs`.", + "items": { + "minLength": 1, + "type": "string" + }, + "type": [ + "array", + "null" + ] + }, + "cross_kv_outputs": { + "description": "Cross-attention present-KV cache outputs (produced by the encoder for an\nencoder-decoder model), paired positionally with `cross_kv_inputs`.", + "items": { + "minLength": 1, + "type": "string" + }, + "type": [ + "array", + "null" + ] + }, + "encoder_hidden_states_input": { + "description": "Encoder-hidden-states input for an encoder-decoder (cross-attention)\ndecoder graph (e.g. `encoder_hidden_states`).", + "minLength": 1, + "type": [ + "string", + "null" + ] + }, + "hidden_output": { + "description": "Per-token hidden-state output for embedding / VLM hidden extraction, if\nthe graph exposes a distinct hidden output (e.g. `last_hidden_state`).", + "minLength": 1, + "type": [ + "string", + "null" + ] + }, + "inputs_embeds_input": { + "description": "Pre-embedded / routed sequence input (e.g. `inputs_embeds`).\n\nMay be declared alongside `token_input` (see its documentation): a graph\nthat consumes both a raw token input and one or more routed sequence\ninputs is explicitly permitted.", + "minLength": 1, + "type": [ + "string", + "null" + ] + }, + "kv_inputs": { + "description": "Past-KV cache inputs, in the SAME order as `kv_outputs` (positional\npairing). Length must match `kv_outputs`.", + "items": { + "minLength": 1, + "type": "string" + }, + "type": [ + "array", + "null" + ] + }, + "kv_layout": { + "anyOf": [ + { + "$ref": "#/$defs/KvCacheLayout" + }, + { + "type": "null" + } + ], + "description": "Physical layout of this backend's KV cache tensors, as a stride\ndescriptor. Accepts a readable named layout (`head_major_bnsh` or\n`seq_major_bsnh`) or a fully explicit [`KvStrideDescriptor`]. This is a\nper-backend capability — each backend owns its KV buffers and never reads\nthe other's KV bytes — so the ORT backend stays head-major while the\nnative backend may declare seq-major. Absent preserves the historical\nhead-major (BNSH) behavior. See [`KvCacheLayout`]." + }, + "kv_outputs": { + "description": "Present-KV cache outputs, paired positionally with `kv_inputs`.", + "items": { + "minLength": 1, + "type": "string" + }, + "type": [ + "array", + "null" + ] + }, + "kv_ownership": { + "anyOf": [ + { + "$ref": "#/$defs/KvOwnership" + }, + { + "type": "null" + } + ], + "description": "Whether this graph owns past/present KV state or reads target-owned KV.\n\nAbsent preserves the historical `owned` behavior." + }, + "kv_update": { + "anyOf": [ + { + "$ref": "#/$defs/KvUpdateKind" + }, + { + "type": "null" + } + ], + "description": "How the paired `kv_inputs`/`kv_outputs` cache tensors evolve each step.\n\nThis declares GROWING/append versus fixed shared-buffer cache semantics\nexplicitly, and is deliberately kept separate from `state_pairs` (which\ndescribes fixed recurrent tensors that are wholly REPLACED). The KV pair\nlists are the authoritative sparse layer ports: the runtime binds exactly\nthe ports named in `kv_inputs`/`kv_outputs` and never expands them from a\ntotal layer count. Absent means the historical growing-cache default." + }, + "logits_output": { + "description": "Logits output.", + "minLength": 1, + "type": [ + "string", + "null" + ] + }, + "optional_inputs": { + "additionalProperties": { + "$ref": "#/$defs/OptionalInputSpec" + }, + "description": "Optional graph inputs and their explicit absent-value contracts, keyed by\nthe real ONNX input port name.", + "type": "object" + }, + "position_ids_input": { + "description": "Position-ids input, if the graph takes one.", + "minLength": 1, + "type": [ + "string", + "null" + ] + }, + "sequence_source": { + "anyOf": [ + { + "$ref": "#/$defs/SequenceInputKind" + }, + { + "type": "null" + } + ], + "description": "Which declared sequence port drives autoregressive execution.\n\nAbsent preserves the historical `token_ids` behavior. Declaring\n`inputs_embeds` requires `inputs_embeds_input`; declaring `token_ids`\nrequires `token_input`." + }, + "state_pairs": { + "description": "Fixed-shape loop-carried recurrent state ports, distinct from KV cache.\n\nEach pair binds an input port to its matching output port and declares\nhow the input is initialized and how the output feeds the next step\n(`replace` semantics for fixed recurrent tensors). These are neither KV\ncache nor fixed conditioning; the sparse set of state ports comes from\nthis declared list, never expanded from a layer count.", + "items": { + "$ref": "#/$defs/LoopStatePair" + }, + "minItems": 1, + "type": [ + "array", + "null" + ] + }, + "static_cache": { + "anyOf": [ + { + "$ref": "#/$defs/StaticCacheIoSpec" + }, + { + "type": "null" + } + ], + "description": "Explicit port binding for a fixed-buffer TensorScatter static KV cache.\n\nA static-cache decoder scatters each step's K/V into pre-allocated,\nfixed-length buffers via an integer write-index vector and a non-pad\nsequence-length vector, rather than growing/appending a cache. These\ncontrol ports are integer vectors and are therefore SHAPE-indistinguish-\nable from one another, so shape cannot disambiguate them: the ABI must be\ndeclared explicitly. When present, this spec is authoritative and the\nruntime binds exactly these ports. When absent, a graph that exposes the\nscatter ABI is REJECTED with an actionable error naming this key rather\nthan having its integer control ports guessed by name." + }, + "token_input": { + "description": "Token-id input (e.g. `input_ids`).\n\nA graph MAY declare this together with `inputs_embeds_input`: some fused\ndecoders consume a raw token stream AND a routed pre-embedded sequence in\nthe same forward pass. The two are not mutually exclusive; declaring both\nis a valid, explicit contract.", + "minLength": 1, + "type": [ + "string", + "null" + ] + } + }, + "type": "object" + }, + "MtpHiddenLayout": { + "description": "Layout of the target state consumed by an MTP sidecar.", + "oneOf": [ + { + "const": "BSH", + "description": "`[batch, sequence, hidden]` legacy layout.", + "type": "string" + }, + { + "const": "BSHC", + "description": "`[batch, sequence, hc_mult, hidden]` Mobius Hyper-Connection layout.", + "type": "string" + } + ] + }, + "MtpKvMode": { + "description": "Lifetime declared for an MTP sidecar's private KV state.", + "oneOf": [ + { + "const": "proposal_local", + "description": "Reset sidecar KV at every target verification iteration.", + "type": "string" + }, + { + "const": "accepted_prefix", + "description": "Retain only KV corresponding to the accepted draft prefix.", + "type": "string" + } + ] + }, + "MtpTargetInitializer": { + "description": "Exact target-model initializer reference used by an MTP sidecar.", + "properties": { + "name": { + "description": "Exact initializer name in the target ONNX graph.", + "type": "string" + }, + "source": { + "$ref": "#/$defs/MtpWeightSource", + "description": "Initializer ownership source. The Phase-1 contract requires\n`target_initializer`." + } + }, + "required": [ + "source", + "name" + ], + "type": "object" + }, + "MtpWeightSource": { + "description": "Ownership source for an MTP shared weight.", + "oneOf": [ + { + "const": "target_initializer", + "description": "Borrow the named initializer from the target model package.", + "type": "string" + } + ] + }, + "OptionalInputSpec": { + "description": "Presence and absent-value contract for one optional graph input.", + "properties": { + "absent": { + "$ref": "#/$defs/AbsentInputSpec", + "description": "Tensor value supplied when the presence key is absent." + }, + "presence": { + "description": "Opaque, non-empty request presence key; not a port or model name.", + "minLength": 1, + "type": "string" + } + }, + "required": [ + "presence", + "absent" + ], + "type": "object" + }, + "OutputStage": { + "enum": [ + "pre_adapter", + "post_adapter" + ], + "type": "string" + }, + "PerformanceHints": { + "description": "Publisher-provided speculative decoding performance guidance.", + "properties": { + "expected_acceptance_rate": { + "description": "Expected fraction of proposed tokens accepted, from 0.0 through 1.0.", + "format": "float", + "maximum": 1.0, + "minimum": 0.0, + "type": [ + "number", + "null" + ] + }, + "optimal_k": { + "description": "Recommended number of draft tokens per verification step.", + "format": "uint", + "minimum": 1, + "type": [ + "integer", + "null" + ] + } + }, + "type": "object" + }, + "PhaseConfig": { + "description": "Phase gate for one pipeline component.", + "properties": { + "run_on": { + "$ref": "#/$defs/PhaseRunOn", + "description": "Pipeline phase in which the component runs." + }, + "when_present": { + "description": "Opaque presence key required for this component to run.", + "minLength": 1, + "type": [ + "string", + "null" + ] + } + }, + "required": [ + "run_on" + ], + "type": "object" + }, + "PhaseRunOn": { + "description": "Pipeline phase gate.\n\nKnown values are enumerated while future strings remain valid.", + "oneOf": [ + { + "enum": [ + "prompt_only", + "every_step", + "always", + "final_only", + "on_demand" + ], + "title": "Known standard value" + }, + { + "not": { + "enum": [ + "prompt_only", + "every_step", + "always", + "final_only", + "on_demand" + ] + }, + "title": "Forward-compatible extension value", + "type": "string" + } + ], + "type": "string" + }, + "PipelineAudioConfig": { + "description": "Waveform contract for a pipeline stage that emits audio.\n\nArchitecture-neutral: the endpoint is an arbitrary `component.output` name\ncarried in the package's metadata, and the sample rate is a declared number.\nNeither is inferred from a model or vendor name.", + "properties": { + "channels": { + "description": "Number of interleaved channels in the waveform. Defaults to 1 (mono).", + "format": "uint16", + "maximum": 65535, + "minimum": 1, + "type": [ + "integer", + "null" + ] + }, + "output": { + "description": "Endpoint carrying the waveform, in `component.output` form.\n\nWhen absent, the runtime uses the sole output of the final-phase\ncomponent, which is unambiguous for the common single-vocoder shape.", + "type": [ + "string", + "null" + ] + }, + "sample_rate": { + "description": "Sample rate, in hertz, of the waveform the pipeline emits.", + "format": "uint32", + "minimum": 1, + "type": [ + "integer", + "null" + ] + } + }, + "type": "object" + }, + "PipelineComponentSpec": { + "description": "One executable ONNX model in a pipeline.", + "properties": { + "device_preference": { + "anyOf": [ + { + "$ref": "#/$defs/DevicePreference" + }, + { + "type": "null" + } + ], + "description": "Optional execution or device preference declared by the model package." + }, + "filename": { + "description": "Non-empty ONNX filename relative to the model package root.", + "examples": [ + "decoder.onnx" + ], + "minLength": 1, + "type": "string" + }, + "io": { + "anyOf": [ + { + "$ref": "#/$defs/ModelIoSpec" + }, + { + "type": "null" + } + ], + "default": null, + "description": "Explicit graph I/O port bindings for this pipeline component.\n\nThe runtime binds decode-step ports from the declared names. A port that\nis not declared is resolved ONLY from an unambiguous io-shape signal;\nwhen the shape is ambiguous the runtime fails with an actionable error\nnaming the key to declare, and never guesses from a tensor name." + }, + "ports": { + "$ref": "#/$defs/ComponentPorts", + "default": { + "inputs": {}, + "outputs": {} + }, + "description": "Typed graph inputs and outputs exposed by this component." + }, + "tokenizer": { + "description": "Tokenizer filename relative to the package root.\n\nIf absent, loaders may use a shared top-level `tokenizer.json`.", + "examples": [ + "tokenizer.json" + ], + "minLength": 1, + "type": [ + "string", + "null" + ] + }, + "type": { + "$ref": "#/$defs/PipelineRole", + "description": "Component role, for example `encoder`, `decoder`, `draft`, `denoiser`, or `vocoder`." + } + }, + "required": [ + "filename", + "type" + ], + "type": "object" + }, + "PipelineRole": { + "description": "Pipeline component-role vocabulary.", + "oneOf": [ + { + "enum": [ + "encoder", + "vision_encoder", + "audio_encoder", + "decoder", + "draft", + "denoiser", + "scheduler", + "vocoder", + "speech_synthesis" + ], + "title": "Known standard value" + }, + { + "not": { + "enum": [ + "encoder", + "vision_encoder", + "audio_encoder", + "decoder", + "draft", + "denoiser", + "scheduler", + "vocoder", + "speech_synthesis" + ] + }, + "title": "Forward-compatible extension value", + "type": "string" + } + ], + "type": "string" + }, + "PipelineSpec": { + "description": "Multi-model pipeline represented as a directed acyclic dataflow graph.", + "properties": { + "audio": { + "anyOf": [ + { + "$ref": "#/$defs/PipelineAudioConfig" + }, + { + "type": "null" + } + ], + "description": "Waveform contract for a pipeline whose final stage emits audio.\n\nPresent for text-to-speech and any other package that produces sound.\nThe sample rate is model DATA: a runtime cannot infer it from a tensor,\nand guessing it silently changes playback pitch and duration." + }, + "batching": { + "anyOf": [ + { + "$ref": "#/$defs/BatchingContract" + }, + { + "type": "null" + } + ], + "description": "Package batching contract." + }, + "control": { + "anyOf": [ + { + "$ref": "#/$defs/ControlFlow" + }, + { + "type": "null" + } + ], + "description": "Universal nested control-flow program." + }, + "dataflow": { + "description": "Directed tensor or data edges between component ports.", + "items": { + "$ref": "#/$defs/DataflowEdge" + }, + "type": "array" + }, + "inputs": { + "additionalProperties": { + "$ref": "#/$defs/TensorContract" + }, + "default": {}, + "description": "Typed inputs exposed by the complete package.", + "type": "object" + }, + "models": { + "additionalProperties": { + "$ref": "#/$defs/PipelineComponentSpec" + }, + "description": "Named model components in the pipeline DAG; at least one component is required.", + "minProperties": 1, + "type": "object" + }, + "outputs": { + "additionalProperties": { + "$ref": "#/$defs/TensorContract" + }, + "default": {}, + "description": "Typed outputs produced by the complete package.", + "type": "object" + }, + "phases": { + "additionalProperties": { + "$ref": "#/$defs/PhaseConfig" + }, + "default": {}, + "description": "Auxiliary-component lifecycle scheduling, keyed by component name.\n\nModels referenced directly by strategy control-flow fields (`decoder`,\n`model`, `denoiser`, `outer`, or `inner`) must not appear here. Every\nother model must have exactly one phase entry.", + "type": "object" + }, + "positions": { + "anyOf": [ + { + "$ref": "#/$defs/PositionProgram" + }, + { + "type": "null" + } + ], + "description": "Declared position-id generation and prefill→decode continuation program.\n\nGeneric and architecture-neutral: parameterized by rank, axis labels, and\nsection sizes so it expresses both ordinary rank-2 linear positions and\nrank-N multimodal coordinates as data — never a model-family branch." + }, + "postprocessing": { + "anyOf": [ + { + "$ref": "#/$defs/PostprocessingSpec" + }, + { + "type": "null" + } + ], + "description": "Declarative output materialization." + }, + "programs": { + "additionalProperties": { + "$ref": "#/$defs/Program" + }, + "default": {}, + "description": "Named data-only sampler, scheduler, solver, and tensor programs.", + "type": "object" + }, + "reducers": { + "additionalProperties": { + "$ref": "#/$defs/ReducerSpec" + }, + "default": {}, + "description": "Explicit fan-in reducers keyed by destination endpoint.", + "type": "object" + }, + "resources": { + "additionalProperties": { + "$ref": "#/$defs/ResourceContract" + }, + "default": {}, + "description": "Typed resource contracts for named components.", + "type": "object" + }, + "states": { + "additionalProperties": { + "$ref": "#/$defs/StateDeclaration" + }, + "default": {}, + "description": "General tensor state, including loop-carried and persistent session state.", + "type": "object" + }, + "strategy": { + "$ref": "#/$defs/PipelineStrategy", + "description": "Loop and execution strategy for the pipeline." + }, + "vision": { + "anyOf": [ + { + "$ref": "#/$defs/PipelineVisionConfig" + }, + { + "type": "null" + } + ], + "description": "Vision-language model token-expansion contract.\n\nWhen present, the engine uses these fields to replace each image\nplaceholder token in the prompt with the declared expanded image-token\nsequence before KV-cache allocation." + }, + "workflow": { + "anyOf": [ + { + "$ref": "#/$defs/WorkflowSpec" + }, + { + "type": "null" + } + ], + "description": "North-star component-centric SSA workflow." + } + }, + "type": "object" + }, + "PipelineStrategy": { + "description": "Parameterized execution strategy for a pipeline or composite stage.", + "properties": { + "batching": { + "description": "Runtime-specific batching parameters." + }, + "cfg_conditioning_input": { + "default": null, + "description": "Denoiser conditioning input port zeroed for the unconditional pass of\nclassifier-free guidance. Required when `guidance_scale` != 1.0.", + "type": [ + "string", + "null" + ] + }, + "decoder": { + "description": "Autoregressive decoder component name.", + "type": [ + "string", + "null" + ] + }, + "denoiser": { + "description": "Iterative or diffusion denoiser component name.", + "type": [ + "string", + "null" + ] + }, + "guidance_scale": { + "description": "Classifier-free guidance scale or equivalent strategy-specific multiplier.", + "format": "float", + "minimum": 0.0, + "type": [ + "number", + "null" + ] + }, + "inner": { + "default": null, + "description": "Inner autoregressive decoder for a `nested_autoregressive` stage.\n\nThe code_predictor: for each outer frame it runs a short inner AR loop of\n`num_code_groups` steps over the residual codebooks, seeded at inner step\n0 by the outer decoder's `last_hidden_state` (routed via a dataflow edge\n`outer.last_hidden_state -> inner.inputs_embeds`) and threading its own\nper-step code embedding on later steps.", + "type": [ + "string", + "null" + ] + }, + "inner_embedding_output": { + "description": "Inner decoder output port threaded across inner steps for a\n`nested_autoregressive` stage.\n\nEach inner step consumes the previous step's per-code embedding as its\n`inputs_embeds` seed; this names the inner decoder OUTPUT port that\nproduces that embedding. It is declared explicitly because the port is\nshape-indistinguishable from other float outputs — the runtime must not\ninfer it by tensor name. Absent on a nested stage ⇒ actionable error\nnaming `pipeline.strategy.inner_embedding_output`.", + "type": [ + "string", + "null" + ] + }, + "kind": { + "$ref": "#/$defs/PipelineStrategyKind", + "description": "Strategy family; determines which strategy-specific fields are meaningful." + }, + "kv_cache": { + "description": "Runtime-specific KV-cache strategy parameters." + }, + "max_tokens": { + "description": "Maximum number of tokens generated by an autoregressive stage.", + "format": "uint", + "minimum": 1, + "type": [ + "integer", + "null" + ] + }, + "model": { + "description": "Single-pass component name.", + "type": [ + "string", + "null" + ] + }, + "num_code_groups": { + "default": null, + "description": "Inner-loop depth (RVQ residual codebook count) for a\n`nested_autoregressive` stage: the number of code tokens collected per\nouter frame. Must be at least 1.", + "format": "uint", + "minimum": 1, + "type": [ + "integer", + "null" + ] + }, + "num_steps": { + "description": "Number of iterative or diffusion steps.", + "format": "uint", + "minimum": 1, + "type": [ + "integer", + "null" + ] + }, + "outer": { + "default": null, + "description": "Outer autoregressive decoder for a `nested_autoregressive` stage.\n\nThe multi-decoder TTS shape: one outer step is one\naudio frame. The outer decoder (talker) produces a per-frame\n`last_hidden_state` that seeds the inner loop (see `inner`).", + "type": [ + "string", + "null" + ] + }, + "pre_embedder": { + "anyOf": [ + { + "$ref": "#/$defs/PreEmbedderSpec" + }, + { + "type": "null" + } + ], + "description": "Optional pre-embedder component driving the outer decoder (talker) of a\n`nested_autoregressive` stage through `inputs_embeds` instead of\n`input_ids`.\n\nA codec-driven TTS talker is not driven by token ids: each step's\n`inputs_embeds` is materialized from the PREVIOUS frame's codes as\n`codec_sum(+ text_embed)` (where\n`codec_sum = codec_embed(code_0) + Σ_i cp_codec_weights[i][codes[i+1]]`).\nWhen this field names such a component (inputs\n`frame_codes [batch, num_code_groups]` int64 `[+ text_embed [batch, 1,\nhidden]]` → output `inputs_embeds [batch, 1, hidden]`), the runtime builds\nthe outer decoder's per-step `inputs_embeds` through it, keeping the engine\ngeneric. Requires a dataflow edge\n`{pre_embedder}.inputs_embeds -> {outer}.inputs_embeds`.\n\nWhen absent the outer loop is `input_ids`-driven (backward compatible).\n\nAll graph-specific port bindings (the pre-embedder's `frame_codes` /\noptional `text_embed` inputs and the output feeding the outer decoder)\nare declared explicitly in [`PreEmbedderSpec`]; the runtime never guesses\nthem by tensor name or dtype." + }, + "prefill_embedder": { + "anyOf": [ + { + "$ref": "#/$defs/PrefillEmbedderSpec" + }, + { + "type": "null" + } + ], + "description": "Optional prefill embedder component that supplies the outer decoder\n(talker) with its real frame-0 PREFILL sequence and the per-frame\ntrailing-text conditioning of a `nested_autoregressive` stage.\n\nThe talker is prefilled with a multi-position embedding\nsequence built from the tokenized prompt, and each subsequent frame is\nconditioned on one trailing-text embedding. This component materializes\nboth from `text_ids`: inputs `text_ids [batch, text_len]` int64 → outputs\n`prefill_embeds [batch, prefill_len, hidden]` float (fed DIRECTLY to the\ntalker's `inputs_embeds` on frame 0) and `trailing_text_embeds [batch,\ntrailing_len, hidden]` float (one vector consumed per outer frame `k >= 1`\nas the pre-embedder's `text_embed`). It runs once in the prompt phase\n(`run_on: prompt_only`); its `text_ids` input is auto-seeded from the\ntokenized prompt.\n\nOnly meaningful together with [`Self::pre_embedder`] (the frame-`k >= 1`\npath feeds the trailing-text vectors through it). When absent, frame 0\nuses a zero seed and every `text_embed` is zero (backward compatible).\n\nAll graph-specific port bindings (the prompt input plus the prefill and\ntrailing-text outputs) are declared explicitly in [`PrefillEmbedderSpec`];\nthe runtime never guesses them by tensor name or dtype." + }, + "scheduler": { + "description": "Scheduler identifier for iterative or diffusion execution.", + "type": [ + "string", + "null" + ] + }, + "scheduler_config": { + "anyOf": [ + { + "$ref": "#/$defs/SchedulerSpec" + }, + { + "type": "null" + } + ], + "description": "Optional diffusion scheduler applied to the denoiser's loop-carried\noutput (treating it as a noise prediction) each step." + }, + "speculative": { + "description": "Runtime-specific speculative execution parameters." + }, + "stages": { + "description": "Ordered child stages for a composite strategy.", + "items": { + "$ref": "#/$defs/PipelineStrategyStage" + }, + "type": "array" + }, + "start_step": { + "default": null, + "description": "First step index for a partial (img2img) denoise loop.\n\nWhen set, the iterative loop runs `start_step..num_steps` instead of the\nfull `0..num_steps`, and the seed (`denoiser` sample input) is expected to\nalready be the encoded image noised to `timesteps[start_step]`. Matches\ndiffusers' img2img `get_timesteps(num_steps, strength)` skip. Default 0.", + "format": "uint", + "minimum": 0, + "type": [ + "integer", + "null" + ] + }, + "state": { + "description": "Runtime-specific iterative state declaration." + }, + "stop_conditions": { + "description": "Runtime-specific stop-condition declarations.", + "items": true, + "type": [ + "array", + "null" + ] + }, + "timestep_input": { + "default": null, + "description": "Denoiser input port that receives the per-step timestep/sigma scalar.\n\nWhen set, the iterative loop feeds this input a rank-1 `float32` value\neach step (from `timesteps` when provided, otherwise the 0-based step\nindex), so a step-aware denoiser can condition on the current step.", + "type": [ + "string", + "null" + ] + }, + "timesteps": { + "default": null, + "description": "Explicit per-step timestep/sigma schedule for an iterative strategy.\n\nWhen present its length must equal `num_steps`; when absent the loop\nuses the 0-based step index. Requires `timestep_input` to have any effect.", + "items": { + "format": "float", + "type": "number" + }, + "type": [ + "array", + "null" + ] + } + }, + "required": [ + "kind" + ], + "type": "object" + }, + "PipelineStrategyKind": { + "description": "Pipeline execution strategy family.\n\nKnown values are enumerated while future strings remain valid.", + "oneOf": [ + { + "enum": [ + "autoregressive", + "iterative", + "diffusion_steps", + "diffusion-steps", + "single_pass", + "single-pass", + "composite", + "nested_autoregressive", + "nested-autoregressive" + ], + "title": "Known standard value" + }, + { + "not": { + "enum": [ + "autoregressive", + "iterative", + "diffusion_steps", + "diffusion-steps", + "single_pass", + "single-pass", + "composite", + "nested_autoregressive", + "nested-autoregressive" + ] + }, + "title": "Forward-compatible extension value", + "type": "string" + } + ], + "type": "string" + }, + "PipelineStrategyStage": { + "additionalProperties": false, + "description": "Named child stage of a composite pipeline strategy.", + "properties": { + "name": { + "description": "Non-empty stage name unique within its containing composite.", + "minLength": 1, + "type": "string" + }, + "strategy": { + "$ref": "#/$defs/PipelineStrategy", + "description": "Execution strategy for this stage." + } + }, + "required": [ + "name", + "strategy" + ], + "type": "object" + }, + "PipelineVisionConfig": { + "description": "Image placeholder token-expansion contract for encoder-free VLM pipelines.\n\nEvery field is optional and additive: legacy documents that declare only\n`image_placeholder_token_id` and `tokens_per_tile` keep working. The richer\nfields mirror the generic expansion the preprocessor already models\n(separate emitted image token, per-tile/per-patch count source, per-image\ncorrespondence, optional row/column separators, and thumbnail order). All of\nit is generic data — no field names or values reference a model family.", + "properties": { + "column_separator_token_id": { + "description": "Optional token ID emitted between columns within a grid row.", + "format": "int64", + "type": [ + "integer", + "null" + ] + }, + "correspondence_summary": { + "description": "Named preprocessing value containing explicit image correspondence data.", + "minLength": 1, + "type": [ + "string", + "null" + ] + }, + "image_correspondence": { + "anyOf": [ + { + "$ref": "#/$defs/ImageCorrespondence" + }, + { + "type": "null" + } + ], + "description": "How prompt placeholders correspond to input images.\n\n`prompt_order` pairs each placeholder with the next input image.\n`explicit_indices` reads correspondence from `correspondence_summary`." + }, + "image_placeholder_token_id": { + "description": "Token ID of the image placeholder in the tokenized prompt.\n\nThe engine replaces every occurrence of this token with the expanded\nimage token sequence before sequence-length and KV-cache sizing.", + "format": "int64", + "type": [ + "integer", + "null" + ] + }, + "image_token_id": { + "description": "Token ID emitted for each expanded image position.\n\nDistinct from `image_placeholder_token_id`: the placeholder marks WHERE\nto expand, while this is the token actually written into the expanded\nsequence. When absent, the placeholder token itself is repeated.", + "format": "int64", + "type": [ + "integer", + "null" + ] + }, + "placeholder_per_image": { + "description": "Whether each placeholder occurrence corresponds to one input image in\nprompt order. Absent means the historical one-placeholder-per-image rule.", + "type": [ + "boolean", + "null" + ] + }, + "row_separator_token_id": { + "description": "Optional token ID emitted between rows of a tiled image grid.", + "format": "int64", + "type": [ + "integer", + "null" + ] + }, + "thumbnail_order": { + "anyOf": [ + { + "$ref": "#/$defs/ThumbnailOrder" + }, + { + "type": "null" + } + ], + "description": "Order of the optional global thumbnail tile relative to the local grid." + }, + "token_count_source": { + "anyOf": [ + { + "$ref": "#/$defs/ImageTokenCountSource" + }, + { + "type": "null" + } + ], + "description": "Where the per-placeholder token count comes from (per tile, per patch, or\na declared grid). Generic selector, never a model-family branch." + }, + "token_count_summary": { + "description": "Named preprocessing value that supplies per-image counts or grid\ndimensions when `token_count_source` is data-derived.\n\nThis is an arbitrary processor output name. A runtime resolves the name\nfrom the declared preprocessing program; it never dispatches on familiar\ntensor names.", + "minLength": 1, + "type": [ + "string", + "null" + ] + }, + "tokens_per_patch": { + "description": "Number of image tokens each patch expands to, used when the count source\nis per patch. Declared data.", + "format": "uint", + "minimum": 1, + "type": [ + "integer", + "null" + ] + }, + "tokens_per_tile": { + "description": "Number of image tokens each tile expands to.\n\nThe total per-tile expansion is `tokens_per_tile * num_tiles`.", + "format": "uint", + "minimum": 1, + "type": [ + "integer", + "null" + ] + } + }, + "type": "object" + }, + "PolicyComponentContract": { + "description": "Stable semantic roles for ONNX policy-math components.\n\nFields map semantic roles to concrete ONNX port names. The corresponding\ntensor contracts live in [`WorkflowComponent::ports`].", + "oneOf": [ + { + "additionalProperties": false, + "properties": { + "effect": { + "type": "string" + }, + "logits": { + "type": "string" + }, + "mode": { + "$ref": "#/$defs/SamplingPolicyMode" + }, + "rng": { + "anyOf": [ + { + "$ref": "#/$defs/RngPortContract" + }, + { + "type": "null" + } + ] + }, + "role": { + "const": "token_sampler", + "type": "string" + }, + "temperature": { + "type": [ + "string", + "null" + ] + }, + "token": { + "type": "string" + }, + "top_k": { + "type": [ + "string", + "null" + ] + }, + "top_p": { + "type": [ + "string", + "null" + ] + } + }, + "required": [ + "role", + "mode", + "logits", + "token", + "effect" + ], + "type": "object" + }, + { + "additionalProperties": false, + "properties": { + "done": { + "type": "string" + }, + "effect": { + "type": "string" + }, + "eos_ids": { + "type": "string" + }, + "iteration": { + "type": "string" + }, + "max_iterations": { + "type": "string" + }, + "role": { + "const": "termination_predicate", + "type": "string" + }, + "tokens": { + "type": "string" + } + }, + "required": [ + "role", + "tokens", + "eos_ids", + "iteration", + "max_iterations", + "done", + "effect" + ], + "type": "object" + }, + { + "additionalProperties": false, + "properties": { + "effect": { + "type": "string" + }, + "estimate": { + "type": "string" + }, + "next_state": { + "type": "string" + }, + "role": { + "const": "solver_step", + "type": "string" + }, + "schedule": { + "type": "string" + }, + "state": { + "type": "string" + }, + "step": { + "type": "string" + } + }, + "required": [ + "role", + "state", + "estimate", + "step", + "schedule", + "next_state", + "effect" + ], + "type": "object" + }, + { + "additionalProperties": false, + "properties": { + "effect": { + "type": "string" + }, + "mask": { + "type": "string" + }, + "next_mask": { + "type": "string" + }, + "next_state": { + "type": "string" + }, + "proposal": { + "type": "string" + }, + "rng": { + "anyOf": [ + { + "$ref": "#/$defs/RngPortContract" + }, + { + "type": "null" + } + ] + }, + "role": { + "const": "masked_update", + "type": "string" + }, + "state": { + "type": "string" + }, + "step": { + "type": "string" + } + }, + "required": [ + "role", + "state", + "proposal", + "mask", + "step", + "next_state", + "next_mask", + "effect" + ], + "type": "object" + }, + { + "additionalProperties": false, + "properties": { + "accepted_len": { + "type": "string" + }, + "accepted_tokens": { + "type": "string" + }, + "done": { + "type": "string" + }, + "effect": { + "type": "string" + }, + "proposal_scores": { + "type": [ + "string", + "null" + ] + }, + "proposed_tokens": { + "type": "string" + }, + "rng": { + "anyOf": [ + { + "$ref": "#/$defs/RngPortContract" + }, + { + "type": "null" + } + ] + }, + "role": { + "const": "speculative_verifier", + "type": "string" + }, + "target_scores": { + "type": "string" + } + }, + "required": [ + "role", + "target_scores", + "proposed_tokens", + "accepted_tokens", + "accepted_len", + "done", + "effect" + ], + "type": "object" + }, + { + "additionalProperties": false, + "properties": { + "current": { + "type": "string" + }, + "effect": { + "type": "string" + }, + "next": { + "type": "string" + }, + "role": { + "const": "state_update", + "type": "string" + }, + "update": { + "type": "string" + } + }, + "required": [ + "role", + "current", + "update", + "next", + "effect" + ], + "type": "object" + } + ] + }, + "PositionContinuation": { + "description": "Prefill→decode position-continuation vocabulary.", + "oneOf": [ + { + "enum": [ + "linear_increment", + "carry_max", + "from_grid" + ], + "title": "Known standard value" + }, + { + "not": { + "enum": [ + "linear_increment", + "carry_max", + "from_grid" + ] + }, + "title": "Forward-compatible extension value", + "type": "string" + } + ], + "type": "string" + }, + "PositionGeneration": { + "description": "Position-value generation vocabulary.", + "oneOf": [ + { + "enum": [ + "linear", + "processor_coordinates" + ], + "title": "Known standard value" + }, + { + "not": { + "enum": [ + "linear", + "processor_coordinates" + ] + }, + "title": "Forward-compatible extension value", + "type": "string" + } + ], + "type": "string" + }, + "PositionProgram": { + "description": "Declared position-id program for a decoder graph.\n\nThe runtime constructs the position tensor from these declared parameters\ninstead of assuming a fixed rank-2 layout. `rank` 1 (with a single axis)\nexpresses ordinary linear positions; `rank` N expresses multi-axis\nmultimodal coordinates. Axis labels and section sizes are opaque DATA — the\nruntime never infers them from a model name.", + "properties": { + "axes": { + "description": "Optional coordinate-stream labels, one per stream (DATA).", + "items": { + "minLength": 1, + "type": "string" + }, + "type": [ + "array", + "null" + ] + }, + "continuation": { + "anyOf": [ + { + "$ref": "#/$defs/PositionContinuation" + }, + { + "type": "null" + } + ], + "description": "How positions continue from the prompt (prefill) into per-token decode." + }, + "dtype": { + "anyOf": [ + { + "$ref": "#/$defs/TensorDType" + }, + { + "type": "null" + } + ], + "description": "Declared dtype of the position tensor." + }, + "generation": { + "anyOf": [ + { + "$ref": "#/$defs/PositionGeneration" + }, + { + "type": "null" + } + ], + "description": "How the position values are generated for prefill.\n\n`linear` generates ordinary sequence positions. `processor_coordinates`\nconsumes the declared processor summaries to construct multi-axis\ncoordinates. Future generation programs remain extensible capability\nstrings rather than model-family branches." + }, + "input": { + "description": "Graph input port that receives the position ids (arbitrary name, DATA).", + "minLength": 1, + "type": "string" + }, + "processor_summaries": { + "description": "Optional processor-summary endpoints this program reads to compute\nmulti-axis coordinates (e.g. a declared grid-dimensions output). Each\nentry is an arbitrary endpoint name (DATA), never a model-family hint.", + "items": { + "minLength": 1, + "type": "string" + }, + "type": [ + "array", + "null" + ] + }, + "rank": { + "description": "Number of coordinate streams carried by the position tensor.\n\n`1` is an ordinary linear position stream; values `> 1` describe\nmulti-axis multimodal coordinates. The physical ONNX tensor rank is\ndeclared separately by `tensor_rank`.", + "format": "uint", + "minimum": 1, + "type": "integer" + }, + "sections": { + "description": "Optional section sizes for sectioned rotary position embeddings.\n\nOpaque list of per-section widths; their meaning is model DATA, not a\nruntime branch.", + "items": { + "format": "uint", + "minimum": 0, + "type": "integer" + }, + "type": [ + "array", + "null" + ] + }, + "tensor_rank": { + "description": "Physical ONNX tensor rank.\n\nRank 2 declares a conventional `[batch, sequence]` linear input. Higher\nranks declare an explicit coordinate axis in addition to batch/sequence\naxes. Absent preserves the legacy mapping (`rank == 1` means tensor rank\n2; otherwise tensor rank 3).", + "format": "uint", + "minimum": 2, + "type": [ + "integer", + "null" + ] + } + }, + "required": [ + "input", + "rank" + ], + "type": "object" + }, + "PostprocessingSpec": { + "additionalProperties": false, + "properties": { + "outputs": { + "additionalProperties": { + "type": "string" + }, + "default": {}, + "type": "object" + }, + "program": { + "$ref": "#/$defs/Program" + } + }, + "required": [ + "program" + ], + "type": "object" + }, + "PreEmbedderSpec": { + "description": "Structured binding for the optional pre-embedder that drives the outer\ndecoder (talker) of a `nested_autoregressive` stage via `inputs_embeds`.\n\nEvery graph-specific port the runtime touches is declared here, so the\nengine never infers a port by tensor name or dtype.", + "properties": { + "component": { + "description": "Declared model name of the pre-embedder component.", + "minLength": 1, + "type": "string" + }, + "frame_codes_input": { + "description": "Pre-embedder input port receiving the previous frame's codes\n(`int64 [batch, num_code_groups]`).", + "minLength": 1, + "type": "string" + }, + "text_embed_input": { + "description": "Optional pre-embedder input port receiving the per-frame trailing-text\nconditioning vector (`float [batch, 1, hidden]`). When absent, the\npre-embedder exposes no trailing-text input.", + "type": [ + "string", + "null" + ] + } + }, + "required": [ + "component", + "frame_codes_input" + ], + "type": "object" + }, + "Precision": { + "description": "Weight precision and quantization-recipe vocabulary.", + "oneOf": [ + { + "enum": [ + "float32", + "fp32", + "float16", + "fp16", + "bfloat16", + "bf16", + "float8_e4m3fn", + "float8_e5m2", + "int8", + "int4", + "int4_group128" + ], + "title": "Known standard value" + }, + { + "not": { + "enum": [ + "float32", + "fp32", + "float16", + "fp16", + "bfloat16", + "bf16", + "float8_e4m3fn", + "float8_e5m2", + "int8", + "int4", + "int4_group128" + ] + }, + "title": "Forward-compatible extension value", + "type": "string" + } + ], + "type": "string" + }, + "Predicate": { + "description": "Data-only predicates for branch and loop termination.", + "oneOf": [ + { + "additionalProperties": false, + "properties": { + "input": { + "type": "string" + }, + "op": { + "const": "present", + "type": "string" + } + }, + "required": [ + "op", + "input" + ], + "type": "object" + }, + { + "additionalProperties": false, + "properties": { + "op": { + "const": "bool", + "type": "string" + }, + "value": { + "type": "boolean" + } + }, + "required": [ + "op", + "value" + ], + "type": "object" + }, + { + "additionalProperties": false, + "properties": { + "op": { + "const": "not", + "type": "string" + }, + "value": { + "$ref": "#/$defs/Predicate" + } + }, + "required": [ + "op", + "value" + ], + "type": "object" + }, + { + "additionalProperties": false, + "properties": { + "op": { + "const": "all", + "type": "string" + }, + "values": { + "items": { + "$ref": "#/$defs/Predicate" + }, + "type": "array" + } + }, + "required": [ + "op", + "values" + ], + "type": "object" + }, + { + "additionalProperties": false, + "properties": { + "op": { + "const": "any", + "type": "string" + }, + "values": { + "items": { + "$ref": "#/$defs/Predicate" + }, + "type": "array" + } + }, + "required": [ + "op", + "values" + ], + "type": "object" + }, + { + "additionalProperties": false, + "properties": { + "left": { + "$ref": "#/$defs/ScalarExpr" + }, + "op": { + "const": "equal", + "type": "string" + }, + "right": { + "$ref": "#/$defs/ScalarExpr" + } + }, + "required": [ + "op", + "left", + "right" + ], + "type": "object" + }, + { + "additionalProperties": false, + "properties": { + "left": { + "$ref": "#/$defs/ScalarExpr" + }, + "op": { + "const": "less", + "type": "string" + }, + "right": { + "$ref": "#/$defs/ScalarExpr" + } + }, + "required": [ + "op", + "left", + "right" + ], + "type": "object" + }, + { + "additionalProperties": false, + "properties": { + "left": { + "$ref": "#/$defs/ScalarExpr" + }, + "op": { + "const": "less_equal", + "type": "string" + }, + "right": { + "$ref": "#/$defs/ScalarExpr" + } + }, + "required": [ + "op", + "left", + "right" + ], + "type": "object" + } + ] + }, + "PrefillEmbedderSpec": { + "description": "Structured binding for the optional prefill embedder that supplies the outer\ndecoder (talker) of a `nested_autoregressive` stage with its frame-0 PREFILL\nsequence and per-frame trailing-text conditioning.\n\nEvery graph-specific port the runtime touches is declared here, so the\nengine never infers a port by tensor name or dtype.", + "properties": { + "component": { + "description": "Declared model name of the (prompt-phase) prefill embedder component.", + "minLength": 1, + "type": "string" + }, + "prefill_output": { + "description": "Prefill-embedder output port carrying the talker's frame-0 multi-position\nPREFILL sequence (`float [batch, prefill_len, hidden]`), fed DIRECTLY to\nthe outer decoder's `inputs_embeds` on frame 0.", + "minLength": 1, + "type": "string" + }, + "prompt_input": { + "description": "Prefill-embedder input port receiving the tokenized prompt\n(`int64 [batch, text_len]`, e.g. `text_ids`).", + "minLength": 1, + "type": "string" + }, + "trailing_output": { + "description": "Prefill-embedder output port carrying the per-frame trailing-text vectors\n(`float [batch, trailing_len, hidden]`), one sliced per outer frame\n`k >= 1` into the pre-embedder's `text_embed`.", + "minLength": 1, + "type": "string" + } + }, + "required": [ + "component", + "prompt_input", + "prefill_output", + "trailing_output" + ], + "type": "object" + }, + "PreprocessingSpec": { + "description": "Declared, architecture-neutral input preprocessing programs.", + "properties": { + "image": { + "anyOf": [ + { + "$ref": "#/$defs/ImagePreprocessingProgram" + }, + { + "type": "null" + } + ], + "description": "Typed image preprocessing transform program and its named tensor outputs." + } + }, + "type": "object" + }, + "Program": { + "additionalProperties": false, + "description": "Generic tensor/scalar program executed between component invocations.", + "properties": { + "operations": { + "items": { + "$ref": "#/$defs/ProgramOperation" + }, + "type": "array" + } + }, + "required": [ + "operations" + ], + "type": "object" + }, + "ProgramOperation": { + "oneOf": [ + { + "additionalProperties": false, + "properties": { + "from": { + "type": "string" + }, + "op": { + "const": "copy", + "type": "string" + }, + "to": { + "type": "string" + } + }, + "required": [ + "op", + "from", + "to" + ], + "type": "object" + }, + { + "additionalProperties": false, + "properties": { + "dtype": { + "$ref": "#/$defs/TensorDType" + }, + "input": { + "type": "string" + }, + "op": { + "const": "cast", + "type": "string" + }, + "output": { + "type": "string" + } + }, + "required": [ + "op", + "input", + "output", + "dtype" + ], + "type": "object" + }, + { + "additionalProperties": false, + "properties": { + "logits": { + "type": "string" + }, + "method": { + "$ref": "#/$defs/SamplingMethod" + }, + "op": { + "const": "sample", + "type": "string" + }, + "output": { + "type": "string" + } + }, + "required": [ + "op", + "logits", + "output", + "method" + ], + "type": "object" + }, + { + "additionalProperties": false, + "properties": { + "estimate": { + "type": "string" + }, + "op": { + "const": "solver_step", + "type": "string" + }, + "output": { + "type": "string" + }, + "solver": { + "$ref": "#/$defs/SolverSpec" + }, + "state": { + "type": "string" + } + }, + "required": [ + "op", + "estimate", + "state", + "output", + "solver" + ], + "type": "object" + } + ] + }, + "ProposalTopology": { + "description": "Speculative proposal-topology vocabulary.", + "oneOf": [ + { + "enum": [ + "linear", + "tree" + ], + "title": "Known standard value" + }, + { + "not": { + "enum": [ + "linear", + "tree" + ] + }, + "title": "Forward-compatible extension value", + "type": "string" + } + ], + "type": "string" + }, + "ProposalType": { + "description": "Speculator proposal architecture.\n\nKnown spellings are enumerated in the generated schema while unknown\nstrings remain valid to preserve forward compatibility.", + "oneOf": [ + { + "enum": [ + "eagle", + "eagle3", + "eagle-3", + "peagle", + "p-eagle", + "mtp", + "dflash", + "d-flash", + "shared_kv", + "shared-kv" + ], + "title": "Known standard value" + }, + { + "not": { + "enum": [ + "eagle", + "eagle3", + "eagle-3", + "peagle", + "p-eagle", + "mtp", + "dflash", + "d-flash", + "shared_kv", + "shared-kv" + ] + }, + "title": "Forward-compatible extension value", + "type": "string" + } + ], + "type": "string" + }, + "QuantizationAxis": { + "description": "Quantization scaling-axis vocabulary.", + "oneOf": [ + { + "enum": [ + "per_tensor", + "per_channel", + "per_token", + "per_head" + ], + "title": "Known standard value" + }, + { + "not": { + "enum": [ + "per_tensor", + "per_channel", + "per_token", + "per_head" + ] + }, + "title": "Forward-compatible extension value", + "type": "string" + } + ], + "type": "string" + }, + "QuantizationIntent": { + "description": "Runtime-independent model-weight quantization intent.", + "properties": { + "default": { + "anyOf": [ + { + "$ref": "#/$defs/Precision" + }, + { + "type": "null" + } + ], + "description": "Default precision or quantization recipe for model weights." + }, + "overrides": { + "description": "Layer- or component-specific precision overrides.", + "items": { + "$ref": "#/$defs/QuantizationOverride" + }, + "type": [ + "array", + "null" + ] + } + }, + "type": "object" + }, + "QuantizationOverride": { + "description": "Precision override for selected layers or a named graph component.", + "properties": { + "component": { + "description": "Logical component path, for example `attention.qk` or `lm_head`.", + "minLength": 1, + "type": [ + "string", + "null" + ] + }, + "layers": { + "description": "Layer indices to which the override applies; negative indices count from the end.", + "items": { + "format": "int32", + "type": "integer" + }, + "type": [ + "array", + "null" + ] + }, + "precision": { + "$ref": "#/$defs/Precision", + "description": "Required precision or quantization recipe." + } + }, + "required": [ + "precision" + ], + "type": "object" + }, + "ReducerKind": { + "description": "How multiple dataflow values are combined at one destination.", + "enum": [ + "first", + "last", + "sum", + "product", + "mean", + "min", + "max", + "concat", + "stack" + ], + "type": "string" + }, + "ReducerSpec": { + "additionalProperties": false, + "properties": { + "axis": { + "format": "int64", + "type": [ + "integer", + "null" + ] + }, + "kind": { + "$ref": "#/$defs/ReducerKind" + } + }, + "required": [ + "kind" + ], + "type": "object" + }, + "ResourceContract": { + "additionalProperties": false, + "properties": { + "device": { + "anyOf": [ + { + "$ref": "#/$defs/DeviceKind" + }, + { + "type": "null" + } + ] + }, + "device_index": { + "format": "uint32", + "minimum": 0, + "type": [ + "integer", + "null" + ] + }, + "memory_bytes": { + "format": "uint64", + "minimum": 0, + "type": [ + "integer", + "null" + ] + } + }, + "type": "object" + }, + "RngPortContract": { + "additionalProperties": false, + "description": "Counter-based RNG state. Producers should use Philox or Threefry inside ONNX.", + "properties": { + "next_offset": { + "type": "string" + }, + "offset": { + "type": "string" + }, + "seed": { + "type": "string" + } + }, + "required": [ + "seed", + "offset", + "next_offset" + ], + "type": "object" + }, + "RuntimeConfigurable": { + "description": "Features whose concrete settings may be selected by the runtime.", + "properties": { + "chunked_prefill": { + "anyOf": [ + { + "$ref": "#/$defs/ChunkedPrefillConfig" + }, + { + "type": "null" + } + ], + "description": "Chunked-prefill support and preferred chunk size." + }, + "continuous_batching": { + "description": "Whether continuous batching may be enabled.", + "type": [ + "boolean", + "null" + ] + }, + "kv_cache": { + "anyOf": [ + { + "$ref": "#/$defs/RuntimeKvConfig" + }, + { + "type": "null" + } + ], + "description": "Supported runtime-selectable KV-cache dtypes." + }, + "prefix_cache": { + "description": "Whether prefix caching may be enabled.", + "type": [ + "boolean", + "null" + ] + } + }, + "type": "object" + }, + "RuntimeInputRole": { + "enum": [ + "prompt_text", + "prompt_tokens", + "media", + "max_iterations", + "max_output_tokens", + "seed", + "sampling_temperature", + "sampling_top_k", + "sampling_top_p", + "constraint", + "session_id" + ], + "type": "string" + }, + "RuntimeKvConfig": { + "description": "Runtime-selectable KV-cache representations.", + "properties": { + "dtype": { + "description": "Non-empty list of supported KV-cache scalar dtypes, in preference order.", + "items": { + "$ref": "#/$defs/DType" + }, + "minItems": 1, + "type": "array" + } + }, + "required": [ + "dtype" + ], + "type": "object" + }, + "SamplingMethod": { + "oneOf": [ + { + "additionalProperties": false, + "properties": { + "kind": { + "const": "greedy", + "type": "string" + } + }, + "required": [ + "kind" + ], + "type": "object" + }, + { + "additionalProperties": false, + "properties": { + "kind": { + "const": "categorical", + "type": "string" + }, + "temperature": { + "format": "float", + "type": "number" + }, + "top_k": { + "format": "uint", + "minimum": 0, + "type": [ + "integer", + "null" + ] + }, + "top_p": { + "format": "float", + "type": [ + "number", + "null" + ] + } + }, + "required": [ + "kind", + "temperature" + ], + "type": "object" + } + ] + }, + "SamplingPolicyMode": { + "enum": [ + "greedy", + "seeded_stochastic" + ], + "type": "string" + }, + "ScalarExpr": { + "oneOf": [ + { + "additionalProperties": false, + "properties": { + "kind": { + "const": "literal", + "type": "string" + }, + "value": { + "$ref": "#/$defs/ScalarValue" + } + }, + "required": [ + "kind", + "value" + ], + "type": "object" + }, + { + "additionalProperties": false, + "properties": { + "kind": { + "const": "value", + "type": "string" + }, + "source": { + "type": "string" + } + }, + "required": [ + "kind", + "source" + ], + "type": "object" + }, + { + "additionalProperties": false, + "properties": { + "kind": { + "const": "iteration", + "type": "string" + } + }, + "required": [ + "kind" + ], + "type": "object" + } + ] + }, + "ScalarValue": { + "anyOf": [ + { + "type": "boolean" + }, + { + "format": "int64", + "type": "integer" + }, + { + "format": "double", + "type": "number" + }, + { + "type": "string" + } + ] + }, + "SchedulerSpec": { + "description": "Diffusion scheduler configuration for an iterative strategy.\n\nThe runtime treats the denoiser's loop-carried output as a noise prediction\n(or, for `flow_matching`, as a vector field and, for `masked_diffusion`, as\ntoken logits) and applies one scheduler step per iteration. Supported\n`kind`s: `ddpm`, `ddim`, `euler`, `euler_ancestral`, `dpmpp_2m`,\n`flow_matching`, and `masked_diffusion`.", + "properties": { + "beta_end": { + "description": "Linear beta-schedule end (default 0.012).", + "format": "float", + "minimum": 0.0, + "type": [ + "number", + "null" + ] + }, + "beta_schedule": { + "description": "Beta schedule shape: `\"linear\"` (default) or `\"scaled_linear\"` (Stable\nDiffusion).", + "type": [ + "string", + "null" + ] + }, + "beta_start": { + "description": "Linear beta-schedule start (default 0.00085).", + "format": "float", + "minimum": 0.0, + "type": [ + "number", + "null" + ] + }, + "block_length": { + "description": "Semi-autoregressive block length for a `masked_diffusion` scheduler, in\ntokens. When set (and smaller than the masked generation region), each\nstep only commits tokens inside the current left-to-right block, matching\nLLaDA's semi-autoregressive remasking. Defaults to a single block\nspanning the whole masked region.", + "format": "uint", + "minimum": 0, + "type": [ + "integer", + "null" + ] + }, + "kind": { + "description": "Scheduler algorithm: `\"ddpm\"`, `\"ddim\"`, `\"euler\"`,\n`\"euler_ancestral\"`, `\"dpmpp_2m\"`, `\"flow_matching\"`, or\n`\"masked_diffusion\"`.", + "type": "string" + }, + "mask_token_id": { + "description": "Mask token id for a `masked_diffusion` (language-diffusion) scheduler:\neach step commits the highest-confidence still-masked positions.", + "format": "int64", + "type": [ + "integer", + "null" + ] + }, + "num_train_timesteps": { + "description": "Training timesteps the noise schedule was defined over (default 1000).", + "format": "uint", + "minimum": 2, + "type": [ + "integer", + "null" + ] + }, + "prediction_type": { + "description": "Model output parameterization: `\"epsilon\"` (default, noise prediction),\n`\"v_prediction\"` (velocity; SD 2.x, SDXL refiner, many fine-tunes), or\n`\"sample\"`/`\"x0\"` (the model predicts the clean sample directly). All\nbuilt-in diffusion schedulers (`ddpm`, `ddim`, `euler`,\n`euler_ancestral`, `dpmpp_2m`) support every parameterization.\n`flow_matching` instead consumes the model's velocity/vector-field output\ndirectly and accepts an omitted value or `\"flow\"`/`\"velocity\"`.", + "type": [ + "string", + "null" + ] + }, + "remasking": { + "description": "Unmasking strategy for a `masked_diffusion` scheduler:\n * `\"low_confidence\"` (default) — LLaDA: each step commits the\n highest-confidence still-masked positions (confidence-ranked). Best\n for LLaDA checkpoints, but greedy/confidence-ranked decoding of other\n masked-diffusion LMs (e.g. MDLM) collapses into repetitive text.\n * `\"random\"` — MDLM-style ancestral: each still-masked position unmasks\n independently with the schedule probability `1/(steps_remaining)`,\n sampling its token from the model's categorical distribution (use\n `temperature: 1.0` for a true categorical sample). This per-position\n stochastic unmasking avoids the degenerate loops that confidence\n ranking produces. The mask token is never emitted.", + "type": [ + "string", + "null" + ] + }, + "shift": { + "description": "Static timestep shift for `flow_matching` (default `1.0`). The base\nrectified-flow sigma `s` is transformed to\n`shift * s / (1 + (shift - 1) * s)`.", + "format": "float", + "minimum": 0.0, + "type": [ + "number", + "null" + ] + }, + "temperature": { + "description": "Sampling temperature for a `masked_diffusion` scheduler. `0` (default)\nselects each masked position's argmax token deterministically; a positive\nvalue applies Gumbel noise (`logits.exp() / (-log u)^temperature`) before\nthe argmax, matching LLaDA's `add_gumbel_noise`. Confidence used for\nremasking is always the clean-softmax probability of the chosen token.", + "format": "float", + "type": [ + "number", + "null" + ] + }, + "use_exponential_sigmas": { + "description": "Use the exponential sigma spacing (`exp(linspace(log σ_max, log σ_min))`)\ninstead of linspace. Applies to `euler`/`dpmpp_2m`. Mutually exclusive\nwith `use_karras_sigmas` (Karras takes precedence).", + "type": [ + "boolean", + "null" + ] + }, + "use_karras_sigmas": { + "description": "Use the Karras (arXiv:2206.00364, rho=7) sigma spacing instead of the\ndefault linspace spacing. Applies to sigma-space schedulers (`euler`,\n`dpmpp_2m`); the most popular ComfyUI scheduler for those samplers.", + "type": [ + "boolean", + "null" + ] + } + }, + "required": [ + "kind" + ], + "type": "object" + }, + "SemanticInputRole": { + "oneOf": [ + { + "additionalProperties": false, + "properties": { + "kind": { + "const": "runtime", + "type": "string" + }, + "role": { + "$ref": "#/$defs/RuntimeInputRole" + }, + "version": { + "type": "string" + } + }, + "required": [ + "kind", + "version", + "role" + ], + "type": "object" + }, + { + "additionalProperties": false, + "properties": { + "kind": { + "const": "opaque", + "type": "string" + } + }, + "required": [ + "kind" + ], + "type": "object" + } + ] + }, + "SequenceInputKind": { + "description": "Primary autoregressive sequence source for a decoder or proposer graph.", + "oneOf": [ + { + "const": "token_ids", + "description": "Integer token ids supplied through `token_input`.", + "type": "string" + }, + { + "const": "inputs_embeds", + "description": "Precomputed floating-point embeddings supplied through\n`inputs_embeds_input`.", + "type": "string" + } + ] + }, + "SequenceLengthScalarBroadcast": { + "description": "Permitted scalar compatibility for attention key-sequence lengths.", + "oneOf": [ + { + "const": "unit_batch", + "description": "Interpret one rank-0 value as the canonical one-element vector only for\nan attention batch of exactly one.", + "type": "string" + } + ] + }, + "ServingServiceContract": { + "additionalProperties": false, + "properties": { + "accepted_len": { + "type": [ + "string", + "null" + ] + }, + "active": { + "type": "string" + }, + "done": { + "type": "string" + }, + "kv_service": { + "$ref": "#/$defs/KvServiceContract" + }, + "slot_ids": { + "type": "string" + } + }, + "required": [ + "active", + "done", + "slot_ids", + "kv_service" + ], + "type": "object" + }, + "SessionLeaseContract": { + "additionalProperties": false, + "properties": { + "optimistic_metadata_version": { + "default": false, + "type": "boolean" + }, + "policy": { + "$ref": "#/$defs/SessionMutationPolicy", + "default": "exclusive" + }, + "ttl_seconds": { + "format": "uint64", + "minimum": 0, + "type": [ + "integer", + "null" + ] + } + }, + "type": "object" + }, + "SessionMutationPolicy": { + "enum": [ + "exclusive", + "copy_on_write" + ], + "type": "string" + }, + "ShapeRecurrence": { + "oneOf": [ + { + "additionalProperties": false, + "properties": { + "kind": { + "const": "invariant", + "type": "string" + } + }, + "required": [ + "kind" + ], + "type": "object" + }, + { + "additionalProperties": false, + "properties": { + "axis": { + "format": "uint", + "minimum": 0, + "type": "integer" + }, + "increment": { + "type": "string" + }, + "kind": { + "const": "growing", + "type": "string" + }, + "max": { + "type": "string" + } + }, + "required": [ + "kind", + "axis", + "increment", + "max" + ], + "type": "object" + } + ] + }, + "SharedKvGroup": { + "description": "One shared-KV binding group for a shared-KV proposer.\n\nA `shared_kv` proposer graph exposes `shared_kv..key` and\n`shared_kv..value` inputs bound to slices of the target model's paged\nKV cache. `target_layers` lists the target KV layer indices feeding this\nslice.", + "properties": { + "key_input": { + "description": "Proposer input receiving this group's shared key cache.", + "type": [ + "string", + "null" + ] + }, + "name": { + "description": "Assistant input prefix, e.g. `sliding_attention` or `full_attention`.", + "type": "string" + }, + "target_key_input": { + "description": "Target decoder past-KV input whose current key cache is referenced.", + "type": [ + "string", + "null" + ] + }, + "target_layers": { + "default": [], + "description": "Target KV layer indices whose cache feeds this shared-KV slice.", + "items": { + "format": "uint", + "minimum": 0, + "type": "integer" + }, + "type": "array" + }, + "target_value_input": { + "description": "Target decoder past-KV input whose current value cache is referenced.", + "type": [ + "string", + "null" + ] + }, + "value_input": { + "description": "Proposer input receiving this group's shared value cache.", + "type": [ + "string", + "null" + ] + } + }, + "required": [ + "name" + ], + "type": "object" + }, + "SlotAllocationMode": { + "enum": [ + "static", + "runtime" + ], + "type": "string" + }, + "SolverSpec": { + "additionalProperties": false, + "properties": { + "algorithm": { + "type": "string" + }, + "parameters": { + "additionalProperties": { + "format": "double", + "type": "number" + }, + "default": {}, + "type": "object" + }, + "schedule": { + "default": [], + "items": { + "format": "double", + "type": "number" + }, + "type": "array" + } + }, + "required": [ + "algorithm" + ], + "type": "object" + }, + "SpecialTokens": { + "description": "Special / control token ids declared by a model author.\n\nEvery field is optional; `eos_token_id` is normalized to a list because\nonnxruntime-genai accepts either a scalar or an array for it.", + "properties": { + "bos_token_id": { + "description": "Beginning-of-stream token id.", + "format": "int64", + "type": [ + "integer", + "null" + ] + }, + "decoder_start_token_id": { + "description": "Token an encoder-decoder model starts decoding with, when not `bos`.", + "format": "int64", + "type": [ + "integer", + "null" + ] + }, + "eos_token_id": { + "description": "End-of-stream token ids (one or more).", + "items": { + "format": "int64", + "type": "integer" + }, + "type": [ + "array", + "null" + ] + }, + "image_token_id": { + "description": "Image placeholder token id (VLMs).", + "format": "int64", + "type": [ + "integer", + "null" + ] + }, + "pad_token_id": { + "description": "Padding token id.", + "format": "int64", + "type": [ + "integer", + "null" + ] + }, + "sep_token_id": { + "description": "Separator token id.", + "format": "int64", + "type": [ + "integer", + "null" + ] + }, + "video_token_id": { + "description": "Video placeholder token id (VLMs).", + "format": "int64", + "type": [ + "integer", + "null" + ] + }, + "vision_start_token_id": { + "description": "Vision-segment start token id (VLMs).", + "format": "int64", + "type": [ + "integer", + "null" + ] + } + }, + "type": "object" + }, + "SpeculativeModelInfo": { + "description": "Build-time support for self-contained speculative decoding.", + "properties": { + "has_draft_heads": { + "description": "Whether the exported graph contains Medusa/EAGLE/MTP-style draft heads.", + "type": [ + "boolean", + "null" + ] + }, + "self_speculative_depth": { + "description": "Early-exit layer depth usable for self-speculation.", + "format": "uint", + "minimum": 1, + "type": [ + "integer", + "null" + ] + } + }, + "type": "object" + }, + "SpeculatorConfig": { + "allOf": [ + { + "not": { + "required": [ + "num_speculative_tokens", + "tokens_per_step" + ] + } + } + ], + "description": "Configuration published with a standalone speculative proposer model.", + "oneOf": [ + { + "not": { + "required": [ + "method" + ] + }, + "required": [ + "proposal_type" + ] + }, + { + "not": { + "required": [ + "proposal_type" + ] + }, + "required": [ + "method" + ] + } + ], + "properties": { + "backbone_hidden_size": { + "default": null, + "description": "Target backbone hidden size `H` shared with the proposer.\n\nFor `shared_kv`, `inputs_embeds` is `[B, q, 2*H]` and\n`projected_state` is `[B, q, H]`.", + "format": "uint", + "minimum": 1, + "type": [ + "integer", + "null" + ] + }, + "embedding": { + "anyOf": [ + { + "$ref": "#/$defs/MtpTargetInitializer" + }, + { + "type": "null" + } + ], + "description": "Target embedding initializer shared with the MTP sidecar." + }, + "hc_mult": { + "default": null, + "description": "Number of Hyper-Connection lanes `C`.", + "format": "uint", + "minimum": 1, + "type": [ + "integer", + "null" + ] + }, + "input_embedding": { + "default": null, + "description": "Relative path (from the model directory) to the target model's raw\ninput-token embedding table, as a little-endian f32 matrix in\n`[vocab_size, backbone_hidden_size]` order.\n\nThe `shared_kv` proposer builds each step's `inputs_embeds` as\n`concat(target_input_embedding(last_token), hidden)`, so it must be able\nto look up the target's input embedding of the last drafted/accepted\ntoken. Required for the `shared_kv` proposer.", + "type": [ + "string", + "null" + ] + }, + "io": { + "anyOf": [ + { + "$ref": "#/$defs/ModelIoSpec" + }, + { + "type": "null" + } + ], + "default": null, + "description": "Explicit proposer graph execution contract.\n\nThis uses the same architecture-neutral I/O vocabulary as a target\ndecoder. `sequence_source` selects token ids versus embeddings,\n`kv_ownership` selects private past/present state versus references to\ntarget-owned cache, and the output fields assign semantic roles." + }, + "kv_mode": { + "anyOf": [ + { + "$ref": "#/$defs/MtpKvMode" + }, + { + "type": "null" + } + ], + "description": "Lifetime of the sidecar's KV state.\n\nDefaults to `proposal_local`." + }, + "lm_head": { + "anyOf": [ + { + "$ref": "#/$defs/MtpTargetInitializer" + }, + { + "type": "null" + } + ], + "description": "Target LM-head initializer shared with the MTP sidecar." + }, + "logits_output": { + "default": null, + "description": "Name of the proposer's draft-distribution output.\n\nDefaults to `logits` for `shared_kv`.", + "type": [ + "string", + "null" + ] + }, + "method": { + "allOf": [ + { + "$ref": "#/$defs/ProposalType", + "description": "Proposal architecture used by the speculator.\n\nThe deprecated `method` alias is accepted on input." + } + ], + "deprecated": true, + "description": "Deprecated alias for `proposal_type`." + }, + "model": { + "default": null, + "description": "Relative path (from the model directory) to the proposer ONNX model.\n\nUsed by the `shared_kv` proposer to locate the\nproposer graph. Optional for forward compatibility with proposer\nfamilies that do not ship a standalone model file.", + "type": [ + "string", + "null" + ] + }, + "mtp_hidden_output": { + "default": null, + "description": "Sidecar output projected through the shared target LM head.\n\nDefaults to `mtp_hidden`.", + "type": [ + "string", + "null" + ] + }, + "mtp_state_output": { + "default": null, + "description": "Sidecar recurrent Hyper-Connection state output.\n\nDefaults to `mtp_state`.", + "type": [ + "string", + "null" + ] + }, + "num_speculative_tokens": { + "default": 4, + "description": "Maximum number of tokens proposed per verifier step; defaults to 4.\n\nThe deprecated `tokens_per_step` alias is accepted on input.", + "format": "uint", + "minimum": 1, + "type": "integer" + }, + "projected_state_output": { + "default": null, + "description": "Name of the proposer output threaded forward between steps.\n\nDefaults to `projected_state` for `shared_kv`.", + "type": [ + "string", + "null" + ] + }, + "proposal_type": { + "$ref": "#/$defs/ProposalType", + "description": "Proposal architecture used by the speculator.\n\nThe deprecated `method` alias is accepted on input." + }, + "shared_kv": { + "description": "Shared-KV binding groups consumed by the proposer.\n\nEach group names an assistant input prefix\n(`shared_kv..{key,value}`) and the target KV layer indices whose\ncache feeds that slice. Empty for proposers that own their KV cache.", + "items": { + "$ref": "#/$defs/SharedKvGroup" + }, + "type": "array" + }, + "target_hidden_layout": { + "anyOf": [ + { + "$ref": "#/$defs/MtpHiddenLayout" + }, + { + "type": "null" + } + ], + "description": "Layout of `target_hidden_output`.\n\nMobius MTP sidecars use `BSHC`: batch, sequence, Hyper-Connection lane,\nhidden." + }, + "target_hidden_output": { + "default": null, + "description": "Target decoder output carrying the recurrent MTP seed.\n\nDefaults to `hidden_states` for `mtp`.", + "type": [ + "string", + "null" + ] + }, + "target_hidden_size": { + "default": null, + "description": "Target hidden width `H`.", + "format": "uint", + "minimum": 1, + "type": [ + "integer", + "null" + ] + }, + "tokens_per_step": { + "allOf": [ + { + "default": 4, + "description": "Maximum number of tokens proposed per verifier step; defaults to 4.\n\nThe deprecated `tokens_per_step` alias is accepted on input.", + "format": "uint", + "minimum": 1, + "type": "integer" + } + ], + "deprecated": true, + "description": "Deprecated alias for `num_speculative_tokens`." + }, + "verifier": { + "anyOf": [ + { + "$ref": "#/$defs/SpeculatorVerifier" + }, + { + "type": "null" + } + ], + "description": "Identity of the verifier model against which this proposer was trained." + }, + "vocab_size": { + "default": null, + "description": "Vocabulary size of the proposer's own `logits` output.", + "format": "uint", + "minimum": 1, + "type": [ + "integer", + "null" + ] + } + }, + "required": [], + "type": "object" + }, + "SpeculatorVerifier": { + "description": "Verifier identity embedded in a speculator package.", + "properties": { + "architectures": { + "default": [], + "description": "Verifier architecture class names, in preference order.", + "items": { + "type": "string" + }, + "type": "array" + }, + "name_or_path": { + "description": "HuggingFace-style verifier repository name or local model path.", + "type": [ + "string", + "null" + ] + } + }, + "type": "object" + }, + "StateDeclaration": { + "additionalProperties": false, + "properties": { + "init": { + "$ref": "#/$defs/StateInit" + }, + "scope": { + "$ref": "#/$defs/StateScope" + }, + "type": { + "$ref": "#/$defs/TensorContract" + }, + "update": { + "$ref": "#/$defs/StateUpdate" + } + }, + "required": [ + "type", + "init", + "update", + "scope" + ], + "type": "object" + }, + "StateInit": { + "oneOf": [ + { + "additionalProperties": false, + "properties": { + "kind": { + "const": "zeros", + "type": "string" + } + }, + "required": [ + "kind" + ], + "type": "object" + }, + { + "additionalProperties": false, + "properties": { + "kind": { + "const": "ones", + "type": "string" + } + }, + "required": [ + "kind" + ], + "type": "object" + }, + { + "additionalProperties": false, + "properties": { + "kind": { + "const": "input", + "type": "string" + }, + "source": { + "type": "string" + } + }, + "required": [ + "kind", + "source" + ], + "type": "object" + }, + { + "additionalProperties": false, + "properties": { + "kind": { + "const": "value", + "type": "string" + }, + "source": { + "type": "string" + } + }, + "required": [ + "kind", + "source" + ], + "type": "object" + }, + { + "additionalProperties": false, + "properties": { + "kind": { + "const": "scalar", + "type": "string" + }, + "value": { + "$ref": "#/$defs/ScalarValue" + } + }, + "required": [ + "kind", + "value" + ], + "type": "object" + } + ] + }, + "StateInitKind": { + "description": "Loop-carried state initialization vocabulary.", + "oneOf": [ + { + "enum": [ + "zeros" + ], + "title": "Known standard value" + }, + { + "not": { + "enum": [ + "zeros" + ] + }, + "title": "Forward-compatible extension value", + "type": "string" + } + ], + "type": "string" + }, + "StateScope": { + "enum": [ + "invocation", + "loop", + "request", + "session" + ], + "type": "string" + }, + "StateUpdate": { + "oneOf": [ + { + "additionalProperties": false, + "properties": { + "kind": { + "const": "replace", + "type": "string" + } + }, + "required": [ + "kind" + ], + "type": "object" + }, + { + "additionalProperties": false, + "properties": { + "axis": { + "format": "int64", + "type": "integer" + }, + "kind": { + "const": "append", + "type": "string" + } + }, + "required": [ + "kind", + "axis" + ], + "type": "object" + }, + { + "additionalProperties": false, + "properties": { + "axis": { + "format": "int64", + "type": "integer" + }, + "indices": { + "type": "string" + }, + "kind": { + "const": "scatter", + "type": "string" + } + }, + "required": [ + "kind", + "axis", + "indices" + ], + "type": "object" + }, + { + "additionalProperties": false, + "properties": { + "kind": { + "const": "accumulate", + "type": "string" + } + }, + "required": [ + "kind" + ], + "type": "object" + } + ] + }, + "StateUpdateKind": { + "description": "Loop-carried state update-semantics vocabulary.", + "oneOf": [ + { + "enum": [ + "replace" + ], + "title": "Known standard value" + }, + { + "not": { + "enum": [ + "replace" + ] + }, + "title": "Forward-compatible extension value", + "type": "string" + } + ], + "type": "string" + }, + "StaticCacheIoSpec": { + "description": "Explicit port ABI for a fixed-buffer TensorScatter static KV cache.\n\nDescribes GRAPH STRUCTURE, never a model family. The four per-layer cache\nlists pair positionally per layer and must all have the same length: index\n`i` in each list is layer `i`'s key/value input and updated key/value output.", + "properties": { + "key_cache_inputs": { + "description": "Per-layer static key-cache buffer inputs, positional per layer. Length\nmust equal `value_cache_inputs`, `key_cache_outputs`, and\n`value_cache_outputs`.", + "items": { + "minLength": 1, + "type": "string" + }, + "minItems": 1, + "type": "array" + }, + "key_cache_outputs": { + "description": "Per-layer updated key-cache outputs, paired positionally with\n`key_cache_inputs`.", + "items": { + "minLength": 1, + "type": "string" + }, + "minItems": 1, + "type": "array" + }, + "kv_sequence_length_input": { + "description": "Input port carrying the non-pad KV sequence length (`int` vector).\nShape-indistinguishable from `write_indices_input`, so it too must be\nnamed explicitly.", + "minLength": 1, + "type": "string" + }, + "value_cache_inputs": { + "description": "Per-layer static value-cache buffer inputs, paired positionally with\n`key_cache_inputs`.", + "items": { + "minLength": 1, + "type": "string" + }, + "minItems": 1, + "type": "array" + }, + "value_cache_outputs": { + "description": "Per-layer updated value-cache outputs, paired positionally with\n`value_cache_inputs`.", + "items": { + "minLength": 1, + "type": "string" + }, + "minItems": 1, + "type": "array" + }, + "write_indices_input": { + "description": "Input port carrying the per-token scatter write positions\n(`int` vector). Shape-indistinguishable from other integer control\ninputs, so it must be named explicitly.", + "minLength": 1, + "type": "string" + } + }, + "required": [ + "write_indices_input", + "kv_sequence_length_input", + "key_cache_inputs", + "value_cache_inputs", + "key_cache_outputs", + "value_cache_outputs" + ], + "type": "object" + }, + "StrategyKind": { + "description": "Generic inference-strategy vocabulary.", + "oneOf": [ + { + "enum": [ + "speculative" + ], + "title": "Known standard value" + }, + { + "not": { + "enum": [ + "speculative" + ] + }, + "title": "Forward-compatible extension value", + "type": "string" + } + ], + "type": "string" + }, + "StrategySpec": { + "description": "Generic inference strategy declaration.", + "properties": { + "acceptance": { + "anyOf": [ + { + "$ref": "#/$defs/AcceptanceMethod" + }, + { + "type": "null" + } + ], + "description": "Draft-token acceptance rule." + }, + "draft": { + "anyOf": [ + { + "$ref": "#/$defs/DraftConfig" + }, + { + "type": "null" + } + ], + "description": "Draft-token producer configuration for speculative decoding." + }, + "kind": { + "$ref": "#/$defs/StrategyKind", + "description": "Strategy vocabulary entry, such as `speculative`." + }, + "performance_hints": { + "anyOf": [ + { + "$ref": "#/$defs/PerformanceHints" + }, + { + "type": "null" + } + ], + "description": "Model-publisher performance guidance." + }, + "tokens_per_step": { + "description": "Number of draft tokens attempted per verification step.", + "format": "uint", + "minimum": 1, + "type": [ + "integer", + "null" + ] + }, + "topology": { + "anyOf": [ + { + "$ref": "#/$defs/ProposalTopology" + }, + { + "type": "null" + } + ], + "description": "Proposal topology, such as `linear` or `tree`." + }, + "verify": { + "anyOf": [ + { + "$ref": "#/$defs/VerifyConfig" + }, + { + "type": "null" + } + ], + "description": "Verification configuration for speculative decoding." + } + }, + "required": [ + "kind" + ], + "type": "object" + }, + "StructuredOutputFormat": { + "description": "Structured-output constraint-format vocabulary.", + "oneOf": [ + { + "enum": [ + "json_schema", + "regex", + "context_free_grammar", + "choice" + ], + "title": "Known standard value" + }, + { + "not": { + "enum": [ + "json_schema", + "regex", + "context_free_grammar", + "choice" + ] + }, + "title": "Forward-compatible extension value", + "type": "string" + } + ], + "type": "string" + }, + "StructuredOutputSpec": { + "description": "Structured-output capabilities and model formatting conventions.", + "properties": { + "stop_sequences": { + "description": "Literal token sequences that terminate a structured response.", + "items": { + "type": "string" + }, + "type": [ + "array", + "null" + ] + }, + "supported_formats": { + "description": "Supported constraint formats, such as JSON Schema, regular expressions, or CFGs.", + "items": { + "$ref": "#/$defs/StructuredOutputFormat" + }, + "type": [ + "array", + "null" + ] + }, + "training_format": { + "description": "Format in which the model was trained to emit tool calls or structured values.", + "type": [ + "string", + "null" + ] + } + }, + "type": "object" + }, + "TensorContract": { + "additionalProperties": false, + "description": "Typed tensor contract used at package and component boundaries.", + "properties": { + "dtype": { + "$ref": "#/$defs/TensorDType" + }, + "optional": { + "default": false, + "type": "boolean" + }, + "rank": { + "format": "uint", + "minimum": 0, + "type": "integer" + }, + "shape": { + "items": { + "$ref": "#/$defs/TensorDimension" + }, + "type": [ + "array", + "null" + ] + } + }, + "required": [ + "dtype", + "rank" + ], + "type": "object" + }, + "TensorDType": { + "description": "Tensor-boundary dtype vocabulary, including non-numeric pipeline values.", + "oneOf": [ + { + "enum": [ + "float32", + "fp32", + "float16", + "fp16", + "bfloat16", + "bf16", + "float8_e4m3fn", + "float8_e5m2", + "int64", + "int32", + "int8", + "uint8", + "bool", + "string" + ], + "title": "Known standard value" + }, + { + "not": { + "enum": [ + "float32", + "fp32", + "float16", + "fp16", + "bfloat16", + "bf16", + "float8_e4m3fn", + "float8_e5m2", + "int64", + "int32", + "int8", + "uint8", + "bool", + "string" + ] + }, + "title": "Forward-compatible extension value", + "type": "string" + } + ], + "type": "string" + }, + "TensorDimension": { + "anyOf": [ + { + "description": "A fixed, non-negative dimension.", + "format": "int64", + "minimum": 0, + "type": "integer" + }, + { + "description": "A runtime shape symbol.", + "minLength": 1, + "type": "string" + } + ], + "description": "One fixed or runtime-resolved tensor-shape dimension." + }, + "Termination": { + "oneOf": [ + { + "additionalProperties": false, + "properties": { + "count": { + "format": "uint", + "minimum": 0, + "type": "integer" + }, + "kind": { + "const": "iterations", + "type": "string" + }, + "start": { + "default": 0, + "format": "uint", + "minimum": 0, + "type": "integer" + } + }, + "required": [ + "kind", + "count" + ], + "type": "object" + }, + { + "additionalProperties": false, + "properties": { + "condition": { + "$ref": "#/$defs/Predicate" + }, + "kind": { + "const": "predicate", + "type": "string" + }, + "max_iterations": { + "format": "uint", + "minimum": 0, + "type": "integer" + } + }, + "required": [ + "kind", + "condition", + "max_iterations" + ], + "type": "object" + } + ] + }, + "ThumbnailOrder": { + "description": "Optional-thumbnail ordering vocabulary.", + "oneOf": [ + { + "enum": [ + "none", + "prepend", + "append" + ], + "title": "Known standard value" + }, + { + "not": { + "enum": [ + "none", + "prepend", + "append" + ] + }, + "title": "Forward-compatible extension value", + "type": "string" + } + ], + "type": "string" + }, + "VerificationMethod": { + "description": "Speculative verification-method vocabulary.", + "oneOf": [ + { + "enum": [ + "single_forward" + ], + "title": "Known standard value" + }, + { + "not": { + "enum": [ + "single_forward" + ] + }, + "title": "Forward-compatible extension value", + "type": "string" + } + ], + "type": "string" + }, + "VerifyConfig": { + "description": "Draft-token verification configuration.", + "properties": { + "method": { + "anyOf": [ + { + "$ref": "#/$defs/VerificationMethod" + }, + { + "type": "null" + } + ], + "description": "Verification method, such as `single_forward`." + }, + "session": { + "description": "Named verifier session or pipeline component.", + "type": [ + "string", + "null" + ] + } + }, + "type": "object" + }, + "WorkflowBatchingContract": { + "oneOf": [ + { + "additionalProperties": false, + "properties": { + "kind": { + "const": "none", + "type": "string" + } + }, + "required": [ + "kind" + ], + "type": "object" + }, + { + "additionalProperties": false, + "properties": { + "axis": { + "format": "uint", + "minimum": 0, + "type": "integer" + }, + "kind": { + "const": "stack", + "type": "string" + } + }, + "required": [ + "kind", + "axis" + ], + "type": "object" + }, + { + "additionalProperties": false, + "properties": { + "kind": { + "const": "ragged", + "type": "string" + }, + "offsets": { + "type": "string" + } + }, + "required": [ + "kind", + "offsets" + ], + "type": "object" + } + ] + }, + "WorkflowComponent": { + "additionalProperties": false, + "properties": { + "effects": { + "default": [], + "items": { + "type": "string" + }, + "type": "array" + }, + "implementation": { + "$ref": "#/$defs/ComponentImplementation" + }, + "policy": { + "anyOf": [ + { + "$ref": "#/$defs/PolicyComponentContract" + }, + { + "type": "null" + } + ] + }, + "ports": { + "$ref": "#/$defs/ComponentPorts" + }, + "resources": { + "anyOf": [ + { + "$ref": "#/$defs/WorkflowResourceContract" + }, + { + "type": "null" + } + ] + } + }, + "required": [ + "implementation", + "ports" + ], + "type": "object" + }, + "WorkflowEmitMode": { + "enum": [ + "replace", + "append", + "event" + ], + "type": "string" + }, + "WorkflowInput": { + "additionalProperties": false, + "properties": { + "contract": { + "$ref": "#/$defs/TensorContract" + }, + "default": { + "anyOf": [ + { + "$ref": "#/$defs/ScalarValue" + }, + { + "type": "null" + } + ] + }, + "required": { + "default": false, + "type": "boolean" + }, + "role": { + "$ref": "#/$defs/SemanticInputRole" + }, + "source": { + "$ref": "#/$defs/WorkflowInputSource" + } + }, + "required": [ + "contract", + "role", + "source" + ], + "type": "object" + }, + "WorkflowInputSource": { + "oneOf": [ + { + "additionalProperties": false, + "properties": { + "field": { + "$ref": "#/$defs/RuntimeInputRole" + }, + "kind": { + "const": "request", + "type": "string" + } + }, + "required": [ + "kind", + "field" + ], + "type": "object" + }, + { + "additionalProperties": false, + "properties": { + "kind": { + "const": "application", + "type": "string" + }, + "name": { + "type": "string" + } + }, + "required": [ + "kind", + "name" + ], + "type": "object" + }, + { + "additionalProperties": false, + "properties": { + "kind": { + "const": "literal", + "type": "string" + } + }, + "required": [ + "kind" + ], + "type": "object" + }, + { + "additionalProperties": false, + "properties": { + "kind": { + "const": "artifact", + "type": "string" + }, + "path": { + "type": "string" + } + }, + "required": [ + "kind", + "path" + ], + "type": "object" + } + ] + }, + "WorkflowLoopCarry": { + "additionalProperties": false, + "properties": { + "body_input": { + "type": "string" + }, + "body_output": { + "type": "string" + }, + "cell": { + "type": "string" + }, + "current": { + "type": "string" + }, + "next": { + "type": "string" + }, + "read_effect": { + "$ref": "#/$defs/EffectTransition" + }, + "write_effect": { + "$ref": "#/$defs/EffectTransition" + } + }, + "required": [ + "cell", + "current", + "body_input", + "body_output", + "next", + "read_effect", + "write_effect" + ], + "type": "object" + }, + "WorkflowManifest": { + "additionalProperties": false, + "properties": { + "adapter_abis": { + "additionalProperties": { + "type": "string" + }, + "default": {}, + "type": "object" + }, + "capabilities": { + "default": [], + "items": { + "type": "string" + }, + "type": "array", + "uniqueItems": true + }, + "custom_op_versions": { + "additionalProperties": { + "type": "string" + }, + "default": {}, + "type": "object" + }, + "ir_version": { + "type": "string" + }, + "onnx_opsets": { + "additionalProperties": { + "format": "uint32", + "minimum": 0, + "type": "integer" + }, + "default": {}, + "type": "object" + } + }, + "required": [ + "ir_version" + ], + "type": "object" + }, + "WorkflowMemoryClass": { + "enum": [ + "default", + "device", + "host", + "pinned" + ], + "type": "string" + }, + "WorkflowNode": { + "oneOf": [ + { + "additionalProperties": false, + "properties": { + "kind": { + "const": "sequence", + "type": "string" + }, + "nodes": { + "items": { + "$ref": "#/$defs/WorkflowNode" + }, + "type": "array" + } + }, + "required": [ + "kind", + "nodes" + ], + "type": "object" + }, + { + "additionalProperties": false, + "properties": { + "component": { + "type": "string" + }, + "effects": { + "additionalProperties": { + "$ref": "#/$defs/EffectTransition" + }, + "default": {}, + "type": "object" + }, + "inputs": { + "additionalProperties": { + "type": "string" + }, + "default": {}, + "type": "object" + }, + "kind": { + "const": "invoke", + "type": "string" + }, + "outputs": { + "additionalProperties": { + "type": "string" + }, + "default": {}, + "type": "object" + } + }, + "required": [ + "kind", + "component" + ], + "type": "object" + }, + { + "additionalProperties": false, + "properties": { + "body": { + "$ref": "#/$defs/WorkflowNode" + }, + "carried": { + "default": [], + "items": { + "$ref": "#/$defs/WorkflowLoopCarry" + }, + "type": "array" + }, + "condition": { + "type": "string" + }, + "kind": { + "const": "loop", + "type": "string" + }, + "max_iterations": { + "type": "string" + }, + "setup": { + "$ref": "#/$defs/WorkflowNode" + } + }, + "required": [ + "kind", + "setup", + "body", + "condition", + "max_iterations" + ], + "type": "object" + }, + { + "additionalProperties": false, + "properties": { + "cases": { + "additionalProperties": { + "$ref": "#/$defs/WorkflowNode" + }, + "type": "object" + }, + "default": { + "anyOf": [ + { + "$ref": "#/$defs/WorkflowNode" + }, + { + "type": "null" + } + ] + }, + "kind": { + "const": "branch", + "type": "string" + }, + "predicate": { + "type": "string" + } + }, + "required": [ + "kind", + "predicate", + "cases" + ], + "type": "object" + }, + { + "additionalProperties": false, + "properties": { + "effect": { + "$ref": "#/$defs/EffectTransition" + }, + "effect_name": { + "type": "string" + }, + "kind": { + "const": "emit", + "type": "string" + }, + "mode": { + "$ref": "#/$defs/WorkflowEmitMode" + }, + "output": { + "type": "string" + }, + "value": { + "type": "string" + } + }, + "required": [ + "kind", + "value", + "output", + "mode", + "effect_name", + "effect" + ], + "type": "object" + }, + { + "additionalProperties": false, + "properties": { + "device": { + "$ref": "#/$defs/DeviceKind" + }, + "input": { + "type": "string" + }, + "kind": { + "const": "transfer", + "type": "string" + }, + "output": { + "type": "string" + } + }, + "required": [ + "kind", + "input", + "output", + "device" + ], + "type": "object" + } + ] + }, + "WorkflowOutput": { + "additionalProperties": false, + "properties": { + "contract": { + "$ref": "#/$defs/TensorContract" + }, + "role": { + "$ref": "#/$defs/WorkflowOutputRole" + }, + "stage": { + "$ref": "#/$defs/OutputStage" + } + }, + "required": [ + "contract", + "role", + "stage" + ], + "type": "object" + }, + "WorkflowOutputRole": { + "enum": [ + "tokens", + "text", + "image", + "audio", + "tensor", + "event" + ], + "type": "string" + }, + "WorkflowResourceContract": { + "additionalProperties": false, + "properties": { + "allowed_devices": { + "default": [], + "items": { + "$ref": "#/$defs/DeviceKind" + }, + "type": "array" + }, + "batching": { + "$ref": "#/$defs/WorkflowBatchingContract" + }, + "concurrency": { + "format": "uint", + "minimum": 0, + "type": [ + "integer", + "null" + ] + }, + "memory_class": { + "$ref": "#/$defs/WorkflowMemoryClass" + }, + "preferred_device": { + "anyOf": [ + { + "$ref": "#/$defs/DeviceKind" + }, + { + "type": "null" + } + ] + } + }, + "required": [ + "memory_class", + "batching" + ], + "type": "object" + }, + "WorkflowSpec": { + "additionalProperties": false, + "description": "Sound, component-centric workflow IR. Tensor math lives in invoked components.", + "properties": { + "components": { + "additionalProperties": { + "$ref": "#/$defs/WorkflowComponent" + }, + "type": "object" + }, + "graph": { + "$ref": "#/$defs/WorkflowNode" + }, + "initial_effects": { + "additionalProperties": { + "type": "string" + }, + "default": {}, + "type": "object" + }, + "inputs": { + "additionalProperties": { + "$ref": "#/$defs/WorkflowInput" + }, + "default": {}, + "type": "object" + }, + "manifest": { + "$ref": "#/$defs/WorkflowManifest" + }, + "outputs": { + "additionalProperties": { + "$ref": "#/$defs/WorkflowOutput" + }, + "default": {}, + "type": "object" + }, + "serving": { + "anyOf": [ + { + "$ref": "#/$defs/ServingServiceContract" + }, + { + "type": "null" + } + ] + }, + "state": { + "additionalProperties": { + "$ref": "#/$defs/WorkflowStateCell" + }, + "default": {}, + "type": "object" + } + }, + "required": [ + "manifest", + "components", + "graph" + ], + "type": "object" + }, + "WorkflowStateCell": { + "additionalProperties": false, + "properties": { + "contract": { + "$ref": "#/$defs/TensorContract" + }, + "initializer": { + "type": "string" + }, + "recurrence": { + "$ref": "#/$defs/ShapeRecurrence" + }, + "scope": { + "$ref": "#/$defs/WorkflowStateScope" + }, + "session": { + "anyOf": [ + { + "$ref": "#/$defs/SessionLeaseContract" + }, + { + "type": "null" + } + ] + } + }, + "required": [ + "contract", + "scope", + "initializer", + "recurrence" + ], + "type": "object" + }, + "WorkflowStateScope": { + "enum": [ + "invocation", + "session" + ], + "type": "string" + } + }, + "$id": "https://github.com/onnx/onnx/issues/8184", + "$schema": "https://json-schema.org/draft/2020-12/schema", + "allOf": [ + { + "not": { + "required": [ + "speculative", + "speculator_config" + ] + } + }, + { + "not": { + "properties": { + "model": { + "required": [ + "io" + ] + } + }, + "required": [ + "pipeline", + "model" + ] + } + } + ], + "description": "Portable, runtime-agnostic inference metadata for ONNX generative models. All top-level sections are optional, and unknown fields are permitted for forward-compatible schema evolution.", + "properties": { + "generation": { + "anyOf": [ + { + "$ref": "#/$defs/GenerationDefaults" + }, + { + "type": "null" + } + ], + "description": "Author-declared text-generation / search defaults.\n\nPopulated from an onnxruntime-genai `genai_config.json` `search` block.\nEvery field is optional; readers treat an absent value as \"use the\nruntime default\"." + }, + "hardware_requirements": { + "anyOf": [ + { + "$ref": "#/$defs/HardwareRequirements" + }, + { + "type": "null" + } + ], + "description": "Minimum and beneficial hardware capabilities used for distribution matching." + }, + "kv_cache": { + "anyOf": [ + { + "$ref": "#/$defs/KvCacheSpec" + }, + { + "type": "null" + } + ], + "description": "KV-cache storage, quantization tolerance, and operational semantics." + }, + "model": { + "anyOf": [ + { + "$ref": "#/$defs/ModelCapabilities" + }, + { + "type": "null" + } + ], + "description": "Build-time model properties and runtime-configurable capabilities." + }, + "pipeline": { + "anyOf": [ + { + "$ref": "#/$defs/PipelineSpec" + }, + { + "type": "null" + } + ], + "description": "Declarative multi-model pipeline and its dataflow graph." + }, + "preprocessing": { + "anyOf": [ + { + "$ref": "#/$defs/PreprocessingSpec" + }, + { + "type": "null" + } + ], + "description": "Declared, architecture-neutral input preprocessing programs.\n\nCarries the typed multimodal preprocessing contract (currently the image\ntransform program and its named tensor outputs). Every operation and\noutput is generic, parameterized data — never a model family, vendor\nstring, or baked-in shape. Absent means the model declares no native\npreprocessing program and a runtime must obtain it elsewhere or fail." + }, + "quantization": { + "anyOf": [ + { + "$ref": "#/$defs/QuantizationIntent" + }, + { + "type": "null" + } + ], + "description": "Model weight quantization intent, independent of the packed representation." + }, + "required_capabilities": { + "default": [], + "description": "Capability identifiers that a runtime MUST support or refuse to load the model.", + "examples": [ + [ + "kv_cache", + "grouped_query_attention" + ] + ], + "items": { + "minLength": 1, + "type": "string" + }, + "type": "array" + }, + "schema_version": { + "description": "Schema version of this inference-metadata document, e.g. `\"v1\"`.\n\nAbsent means the initial `\"v1\"` contract (readers default to `v1`).\nBump this only for breaking schema changes; additive fields keep the\nsame major version and rely on the forward-compatible \"ignore unknown\nfields\" rule.", + "type": [ + "string", + "null" + ] + }, + "speculative": { + "anyOf": [ + { + "$ref": "#/$defs/SpeculatorConfig" + }, + { + "type": "null" + } + ], + "description": "Standalone speculative proposer declaration.\n\nThis is the preferred native source for speculator discovery;\nHuggingFace `config.json` is a compatibility fallback. The deprecated\n`speculator_config` alias is accepted on input." + }, + "speculator_config": { + "allOf": [ + { + "anyOf": [ + { + "$ref": "#/$defs/SpeculatorConfig" + }, + { + "type": "null" + } + ], + "description": "Standalone speculative proposer declaration.\n\nThis is the preferred native source for speculator discovery;\nHuggingFace `config.json` is a compatibility fallback. The deprecated\n`speculator_config` alias is accepted on input." + } + ], + "deprecated": true, + "description": "Deprecated alias for `speculative`." + }, + "strategy": { + "anyOf": [ + { + "$ref": "#/$defs/StrategySpec" + }, + { + "type": "null" + } + ], + "description": "Generic inference strategy, including speculative decoding." + }, + "structured_output": { + "anyOf": [ + { + "$ref": "#/$defs/StructuredOutputSpec" + }, + { + "type": "null" + } + ], + "description": "Structured-output formats and model training conventions." + }, + "tokens": { + "anyOf": [ + { + "$ref": "#/$defs/SpecialTokens" + }, + { + "type": "null" + } + ], + "description": "Special / control token ids declared by the model author.\n\nPopulated from the model-level token id fields of a `genai_config.json`." + } + }, + "title": "ONNX Inference Metadata", + "type": "object" +} From 6d265ccce9876e61df62e455822bfe2895cd96f0 Mon Sep 17 00:00:00 2001 From: justinchuby Date: Wed, 12 Aug 2026 19:06:30 +0000 Subject: [PATCH 005/151] Migrate codec metadata to typed workflow SSA Replace the codec strategy/phases producer with explicit typed component invokes, linear effects, codes SSA, and post-adapter audio emission. Remove superseded codec and TTS legacy producers, and fail Qwen3-TTS precisely because the generic nested loop contract does not expose the induction SSA needed for code-predictor steps. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: justinchuby --- .../integrations/onnx_genai/__init__.py | 23 +- .../integrations/onnx_genai/auto_export.py | 110 ++------- .../onnx_genai/auto_export_test.py | 73 ++---- .../codec_workflow_metadata_test.py | 109 ++++++++ .../onnx_genai/inference_metadata.py | 232 ------------------ .../onnx_genai/inference_metadata_test.py | 87 ------- .../onnx_genai/workflow_metadata.py | 158 +++++++++++- 7 files changed, 311 insertions(+), 481 deletions(-) create mode 100644 src/mobius/integrations/onnx_genai/codec_workflow_metadata_test.py diff --git a/src/mobius/integrations/onnx_genai/__init__.py b/src/mobius/integrations/onnx_genai/__init__.py index 7f11b5e08..2e6120753 100644 --- a/src/mobius/integrations/onnx_genai/__init__.py +++ b/src/mobius/integrations/onnx_genai/__init__.py @@ -15,9 +15,10 @@ 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. +* **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. @@ -51,23 +52,23 @@ from mobius.integrations.onnx_genai.inference_metadata import ( SchedulerConfig, add_policy_components_to_workflow, - build_audio_codec_pipeline_metadata, build_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_language_diffusion_pipeline_metadata, + build_tts_workflow_metadata, + write_audio_codec_workflow_metadata, write_decoder_workflow_metadata, write_language_diffusion_workflow_metadata, + write_tts_workflow_metadata, ) __all__ = [ @@ -79,11 +80,11 @@ "build_decoder_workflow_metadata", "build_diffusion_pipeline_metadata", "build_language_diffusion_pipeline_metadata", - "build_audio_codec_pipeline_metadata", + "build_audio_codec_workflow_metadata", "build_multimodal_pipeline_metadata", "build_pipeline_metadata_for_workflow", "build_speech_to_text_pipeline_metadata", - "build_tts_pipeline_metadata", + "build_tts_workflow_metadata", "convert_comfyui_workflow", "decoder_metadata_from_config", "moe_metadata_from_config", @@ -96,9 +97,9 @@ "write_decoder_workflow_metadata", "write_language_diffusion_workflow_metadata", "write_diffusion_pipeline_metadata", - "write_audio_codec_pipeline_metadata", + "write_audio_codec_workflow_metadata", "write_multimodal_pipeline_metadata", "write_speech_to_text_pipeline_metadata", - "write_tts_pipeline_metadata", + "write_tts_workflow_metadata", "write_onnx_genai_config", ] diff --git a/src/mobius/integrations/onnx_genai/auto_export.py b/src/mobius/integrations/onnx_genai/auto_export.py index 77d89eed8..880f7441c 100644 --- a/src/mobius/integrations/onnx_genai/auto_export.py +++ b/src/mobius/integrations/onnx_genai/auto_export.py @@ -27,15 +27,15 @@ add_explicit_package_io, add_policy_components_to_workflow, 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 ( + write_audio_codec_workflow_metadata, write_decoder_workflow_metadata, write_language_diffusion_workflow_metadata, + write_tts_workflow_metadata, ) _LOGGER = logging.getLogger(__name__) @@ -303,33 +303,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: @@ -341,9 +316,8 @@ 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()) @@ -352,44 +326,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") @@ -472,7 +410,7 @@ def write_onnx_genai_config( Pipeline shape Structural signal (detector) Emitted ``strategy`` ===================== ============================================ ================================= 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) @@ -545,9 +483,7 @@ 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} resolved_config = config if config is not None else getattr(pkg, "config", None) @@ -615,34 +551,20 @@ def write_onnx_genai_config( 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 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 diff --git a/src/mobius/integrations/onnx_genai/auto_export_test.py b/src/mobius/integrations/onnx_genai/auto_export_test.py index 07cbd2803..e7a39a666 100644 --- a/src/mobius/integrations/onnx_genai/auto_export_test.py +++ b/src/mobius/integrations/onnx_genai/auto_export_test.py @@ -489,39 +489,28 @@ def test_dispatch_speech_to_text_routes_encoder_mask(tmp_path): 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["graph"]["nodes"][0]["outputs"] == {"codes": "codec.codes"} + assert workflow["graph"]["nodes"][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): @@ -534,7 +523,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)) @@ -553,9 +542,6 @@ 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( { "talker": _FakeModel(["inputs_embeds"], ["logits", "last_hidden_state"]), @@ -567,31 +553,8 @@ def test_dispatch_multi_decoder_tts_with_pre_embedder(tmp_path): "embedding": _FakeModel(["text_ids"]), } ) - 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"] + with pytest.raises(NotImplementedError, match="nested-loop induction SSA value"): + write_onnx_genai_config(pkg, str(tmp_path)) def test_unrecognized_multi_component_package_fails_loudly(tmp_path): 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..c1ac61649 --- /dev/null +++ b/src/mobius/integrations/onnx_genai/codec_workflow_metadata_test.py @@ -0,0 +1,109 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT License. + +from __future__ import annotations + +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, +) + + +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_effects_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"], + } + assert workflow["components"]["encoder"]["ports"]["outputs"]["codes"] == { + "dtype": "int64", + "rank": 3, + "shape": ["batch", 16, "frames"], + } + assert workflow["components"]["encoder"]["effects"] == ["codec_encode"] + assert workflow["components"]["decoder"]["effects"] == ["codec_decode"] + + encode, decode, emit = workflow["graph"]["nodes"] + assert encode["outputs"] == {"codes": "codec.codes"} + assert decode["inputs"] == {"codes": "codec.codes"} + assert emit == { + "kind": "emit", + "value": "codec.waveform", + "output": "waveform", + "mode": "replace", + "effect_name": "audio_emit", + "effect": {"consumes": "audio_emit.0", "produces": "audio_emit.1"}, + } + 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) + + +def test_tts_reports_missing_nested_loop_induction_value(): + package = { + "talker": object(), + "code_predictor": object(), + "talker_step_embedder": object(), + "talker_prefill_embedder": object(), + } + with pytest.raises(NotImplementedError, match="code_predictor.step_index"): + build_tts_workflow_metadata(package, object()) diff --git a/src/mobius/integrations/onnx_genai/inference_metadata.py b/src/mobius/integrations/onnx_genai/inference_metadata.py index 10b022fc7..8433b4003 100644 --- a/src/mobius/integrations/onnx_genai/inference_metadata.py +++ b/src/mobius/integrations/onnx_genai/inference_metadata.py @@ -2285,238 +2285,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"}, - }, - { - "name": "decode_waveform", - "strategy": {"kind": "single_pass", "model": "decoder"}, - }, - ], - }, - "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, - }, - ], - }, - "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 6b2515281..e425f44d1 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 @@ -33,13 +32,11 @@ build_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, ) @@ -851,7 +848,6 @@ def test_gemma4_routes_all_embedding_outputs(self, tmp_path): "pad_value": -1, }, ] - broken = copy.deepcopy(metadata) broken["pipeline"]["dataflow"] = [ edge @@ -1877,86 +1873,3 @@ def test_vision_and_audio_pipeline(self): "strategy": {"kind": "autoregressive", "decoder": "decoder"}, }, ] - - -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/workflow_metadata.py b/src/mobius/integrations/onnx_genai/workflow_metadata.py index 7430f8ed3..f606e0197 100644 --- a/src/mobius/integrations/onnx_genai/workflow_metadata.py +++ b/src/mobius/integrations/onnx_genai/workflow_metadata.py @@ -35,14 +35,22 @@ def _contract(value: ir.Value) -> dict[str, Any]: } -def _component(model: ir.Model, artifact: str) -> dict[str, Any]: - return { +def _component( + model: ir.Model, + artifact: str, + *, + effects: tuple[str, ...] = (), +) -> dict[str, Any]: + component = { "implementation": {"kind": "onnx", "artifact": artifact}, "ports": { "inputs": {value.name: _contract(value) for value in model.graph.inputs}, "outputs": {value.name: _contract(value) for value in model.graph.outputs}, }, } + if effects: + component["effects"] = list(effects) + return component def _effect(consumes: str, produces: str) -> dict[str, str]: @@ -64,6 +72,152 @@ def _invoke( } +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": { + "ir_version": "1.0", + "onnx_opsets": {"ai.onnx": OPSET_VERSION}, + "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": 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) + path = os.path.join(output_dir, "inference_metadata.yaml") + with open(path, "w", encoding="utf-8") as handle: + yaml.safe_dump(metadata, handle, sort_keys=False) + return path + + +def build_tts_workflow_metadata(pkg: Any, config: Any) -> dict[str, Any]: + """Reject TTS until the generic nested-loop contract exposes induction SSA. + + Qwen3-TTS needs the inner loop index as the code predictor ``step_index`` and + to select the next code embedding. The workflow ``loop`` node at producer + commit 4c3c4b6 only accepts a condition and maximum value; it defines no + iteration SSA value. Consequently the existing prefill and per-frame + embedder artifacts can be invoked, but the code-predictor loop cannot be + expressed without host preprocessing or a model-specific counter component. + """ + required = { + "talker", + "code_predictor", + "talker_step_embedder", + } + missing = sorted(required.difference(pkg.keys())) + if missing: + raise ValueError(f"TTS workflow is missing required components: {missing}") + del config + raise NotImplementedError( + "generic TTS workflow requires a nested-loop induction SSA value: " + "ONNX GenAI workflow Loop at producer commit 4c3c4b6 exposes neither an " + "iteration output nor fixed-loop index, so code_predictor.step_index, " + "position_ids, and per-group code embedding selection cannot be wired " + "from the existing prefill/step embedder artifacts without host " + "preprocessing" + ) + + +def write_tts_workflow_metadata(pkg: Any, output_dir: str, config: Any) -> str: + """Build TTS workflow metadata, failing precisely on the producer defect.""" + metadata = build_tts_workflow_metadata(pkg, config) + os.makedirs(output_dir, exist_ok=True) + path = os.path.join(output_dir, "inference_metadata.yaml") + with open(path, "w", encoding="utf-8") as handle: + yaml.safe_dump(metadata, handle, sort_keys=False) + return path + + def build_decoder_workflow_metadata( pkg: Any, config: Any, From 66ab136f6c821473e051e50128b8f1eb58a9ab8f Mon Sep 17 00:00:00 2001 From: justinchuby Date: Wed, 12 Aug 2026 19:12:23 +0000 Subject: [PATCH 006/151] Resolve combined workflow migration tests Preserve masked-diffusion and codec workflow coverage after integrating both migrations, and make the TTS blocker assertion lint-safe. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: justinchuby --- .../integrations/onnx_genai/codec_workflow_metadata_test.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/mobius/integrations/onnx_genai/codec_workflow_metadata_test.py b/src/mobius/integrations/onnx_genai/codec_workflow_metadata_test.py index c1ac61649..917080209 100644 --- a/src/mobius/integrations/onnx_genai/codec_workflow_metadata_test.py +++ b/src/mobius/integrations/onnx_genai/codec_workflow_metadata_test.py @@ -105,5 +105,5 @@ def test_tts_reports_missing_nested_loop_induction_value(): "talker_step_embedder": object(), "talker_prefill_embedder": object(), } - with pytest.raises(NotImplementedError, match="code_predictor.step_index"): + with pytest.raises(NotImplementedError, match=r"code_predictor\.step_index"): build_tts_workflow_metadata(package, object()) From 5b6a3ff6f6bd2127af632c33778691bacaf88200 Mon Sep 17 00:00:00 2001 From: justinchuby Date: Wed, 12 Aug 2026 19:47:25 +0000 Subject: [PATCH 007/151] Fix workflow sampling and loop correctness Slice decoder logits to the last token before sampling, invert termination into continue predicates, emit and check the prefill token, declare growing KV state, progressively commit masked tokens with threaded RNG state, and remove strategy-owned legacy phases. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: justinchuby --- src/mobius/generation/__init__.py | 4 + src/mobius/generation/_policy_components.py | 57 ++++++++- .../generation/_policy_components_test.py | 39 +++++- .../onnx_genai/inference_metadata.py | 14 +-- .../onnx_genai/inference_metadata_test.py | 16 +-- .../onnx_genai/workflow_metadata.py | 114 +++++++++++++++--- .../onnx_genai/workflow_metadata_test.py | 3 +- 7 files changed, 194 insertions(+), 53 deletions(-) diff --git a/src/mobius/generation/__init__.py b/src/mobius/generation/__init__.py index 392192d0a..9646793cd 100644 --- a/src/mobius/generation/__init__.py +++ b/src/mobius/generation/__init__.py @@ -10,9 +10,11 @@ PolicyComponent, PolicyRole, attach_policy_components, + build_boolean_not, build_eos_termination, build_euler_solver_step, build_greedy_sampler, + build_last_token_logits, build_masked_token_update, build_seeded_categorical_sampler, build_speculative_acceptance, @@ -24,9 +26,11 @@ "PolicyCapabilities", "PolicyRole", "attach_policy_components", + "build_boolean_not", "build_eos_termination", "build_euler_solver_step", "build_greedy_sampler", + "build_last_token_logits", "build_masked_token_update", "build_seeded_categorical_sampler", "build_speculative_acceptance", diff --git a/src/mobius/generation/_policy_components.py b/src/mobius/generation/_policy_components.py index 98edc0b2d..53545f8cb 100644 --- a/src/mobius/generation/_policy_components.py +++ b/src/mobius/generation/_policy_components.py @@ -34,6 +34,7 @@ class PolicyRole(StrEnum): MASKED_UPDATE = "masked_update" SPECULATIVE_ACCEPTANCE = "speculative_verifier" STATE_UPDATE = "state_update" + AUXILIARY = "auxiliary" @dataclass(frozen=True) @@ -157,6 +158,30 @@ def build_greedy_sampler() -> PolicyComponent: ) +def build_last_token_logits() -> PolicyComponent: + """Build ``[B,T,V] -> [B,V]`` selection for decoder sampling.""" + graph, builder = _make_graph("last_token_logits") + logits = builder.input( + "logits", + dtype=ir.DataType.FLOAT, + shape=["batch", "sequence", "vocabulary"], + ) + selected = builder.op.Gather(logits, builder.op.Constant(value_int=-1), axis=1) + selected.shape = ir.Shape(["batch", "vocabulary"]) + builder.add_output(selected, "last_logits") + return _component(PolicyRole.AUXILIARY, graph, {}) + + +def build_boolean_not() -> PolicyComponent: + """Build an explicit ``continue = Not(done)`` predicate transform.""" + graph, builder = _make_graph("boolean_not") + done = builder.input("done", dtype=ir.DataType.BOOL, shape=["batch"]) + continued = builder.op.Not(done) + continued.shape = ir.Shape(["batch"]) + builder.add_output(continued, "continue") + return _component(PolicyRole.AUXILIARY, graph, {}) + + def build_seeded_categorical_sampler() -> PolicyComponent: """Build deterministic categorical sampling with explicit seed and offset. @@ -343,12 +368,32 @@ def build_masked_token_update() -> PolicyComponent: step = builder.input("step", ir.DataType.INT64, ["batch"]) seed = builder.input("seed", ir.DataType.INT64, ["batch"]) offset = builder.input("offset", ir.DataType.INT64, ["batch"]) - updated = op.Where(masked, proposed, current) - # Consume the declared step without changing values; schedules that remask - # tokens can be expressed by a richer artifact with the same semantic ports. - updated = op.Add(updated, op.Unsqueeze(op.Mul(step, 0), op.Constant(value_ints=[-1]))) + 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( + seed, + op.Constant(value_ints=[-1]), + ), + ) + bucket = op.Mod(stream, op.Constant(value_int=8), fmod=0) + scheduled = op.Less( + bucket, + op.Unsqueeze( + op.Min(op.Add(step, op.Constant(value_int=1)), op.Constant(value_int=8)), + op.Constant(value_ints=[-1]), + ), + ) + committed = op.And(masked, scheduled) + updated = op.Where(committed, proposed, current) updated.shape = ir.Shape(["batch", "sequence"]) - remaining = op.ConstantOfShape(op.Shape(masked), value=ir.tensor([False])) + remaining = op.And(masked, op.Not(committed)) remaining.shape = ir.Shape(["batch", "sequence"]) remaining_count = op.ReduceSum( op.Cast(remaining, to=ir.DataType.INT64), @@ -358,7 +403,7 @@ def build_masked_token_update() -> PolicyComponent: done = op.Equal(remaining_count, op.Constant(value_int=0)) done.shape = ir.Shape(["batch"]) next_offset = op.Add( - op.Add(offset, op.Constant(value_int=1)), + 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"]) diff --git a/src/mobius/generation/_policy_components_test.py b/src/mobius/generation/_policy_components_test.py index 1f4f78754..116ccf7e0 100644 --- a/src/mobius/generation/_policy_components_test.py +++ b/src/mobius/generation/_policy_components_test.py @@ -12,9 +12,11 @@ PolicyCapabilities, PolicyRole, attach_policy_components, + build_boolean_not, build_eos_termination, build_euler_solver_step, build_greedy_sampler, + build_last_token_logits, build_masked_token_update, build_seeded_categorical_sampler, build_speculative_acceptance, @@ -38,6 +40,19 @@ def test_greedy_sampler_runtime(tmp_path): np.testing.assert_array_equal(tokens, [1, 2]) +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, :]) + + (continued,) = _run( + build_boolean_not(), + tmp_path, + {"done": np.array([True, False])}, + ) + np.testing.assert_array_equal(continued, [False, True]) + + def test_seeded_sampler_is_counter_based_and_reproducible(tmp_path): component = build_seeded_categorical_sampler() feeds = { @@ -95,10 +110,26 @@ def test_masked_update_runtime_parity(tmp_path): "offset": np.array([11], np.int64), }, ) - np.testing.assert_array_equal(outputs[0], [[1, 5, 6]]) - np.testing.assert_array_equal(outputs[1], [[False, False, False]]) - np.testing.assert_array_equal(outputs[2], [12]) - np.testing.assert_array_equal(outputs[3], [True]) + np.testing.assert_array_equal(outputs[0], [[1, 5, 99]]) + np.testing.assert_array_equal(outputs[1], [[False, False, True]]) + np.testing.assert_array_equal(outputs[2], [14]) + np.testing.assert_array_equal(outputs[3], [False]) + + final = _run( + build_masked_token_update(), + tmp_path, + { + "current_tokens": outputs[0], + "proposed_tokens": np.array([[4, 5, 6]], np.int64), + "masked": outputs[1], + "step": outputs[2], + "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]) def test_speculative_acceptance_prefix_runtime(tmp_path): diff --git a/src/mobius/integrations/onnx_genai/inference_metadata.py b/src/mobius/integrations/onnx_genai/inference_metadata.py index 8433b4003..498d78d0c 100644 --- a/src/mobius/integrations/onnx_genai/inference_metadata.py +++ b/src/mobius/integrations/onnx_genai/inference_metadata.py @@ -1347,7 +1347,7 @@ def add_policy_components_to_workflow( workflow_dtypes = {"fp16": "float16", "bf16": "bfloat16", "fp32": "float32"} for name, component in policy_components.items(): model = component.model - components[name] = { + declaration = { "implementation": { "kind": "onnx", "artifact": f"policies/{name}.onnx", @@ -1378,9 +1378,11 @@ def add_policy_components_to_workflow( for value in model.graph.outputs }, }, - "policy": component.contract, "effects": list(component.effects), } + if component.contract: + declaration["policy"] = component.contract + components[name] = declaration return metadata @@ -1649,7 +1651,6 @@ def build_native_vlm_package_metadata( "models": models, "dataflow": dataflow, "strategy": {"kind": "composite", "stages": stages}, - "phases": phases, "vision": vision_config, } if positions is not None: @@ -1994,7 +1995,7 @@ def build_diffusion_pipeline_metadata( "to": f"denoiser.{denoiser_sample_input}", }, ] - phases: dict[str, Any] = {"denoiser": {"run_on": "every_step"}} + phases: dict[str, Any] = {} if text_encoder_filename is not None: models["text_encoder"] = { @@ -2170,7 +2171,6 @@ def add_encoder( "models": models, "dataflow": dataflow, "strategy": {"kind": "composite", "stages": stages}, - "phases": phases, } return metadata @@ -2262,10 +2262,6 @@ def build_speech_to_text_pipeline_metadata( }, ], }, - "phases": { - "encoder": {"run_on": "prompt_only"}, - "decoder": {"run_on": "every_step"}, - }, } return metadata diff --git a/src/mobius/integrations/onnx_genai/inference_metadata_test.py b/src/mobius/integrations/onnx_genai/inference_metadata_test.py index e425f44d1..357a5b05d 100644 --- a/src/mobius/integrations/onnx_genai/inference_metadata_test.py +++ b/src/mobius/integrations/onnx_genai/inference_metadata_test.py @@ -789,15 +789,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", - } - assert metadata["pipeline"]["phases"]["embedding"] == {"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 @@ -1804,11 +1797,6 @@ def test_vision_only_pipeline(self): }, ], }, - "phases": { - "vision_encoder": {"run_on": "prompt_only"}, - "embedding": {"run_on": "prompt_only"}, - "decoder": {"run_on": "every_step"}, - }, } } diff --git a/src/mobius/integrations/onnx_genai/workflow_metadata.py b/src/mobius/integrations/onnx_genai/workflow_metadata.py index f606e0197..82ae5efcf 100644 --- a/src/mobius/integrations/onnx_genai/workflow_metadata.py +++ b/src/mobius/integrations/onnx_genai/workflow_metadata.py @@ -15,6 +15,8 @@ from mobius.generation import ( PolicyCapabilities, attach_policy_components, + build_boolean_not, + build_last_token_logits, ) from mobius.integrations.onnx_genai.inference_metadata import ( _port, @@ -236,6 +238,8 @@ def build_decoder_workflow_metadata( token_state_update=True, ), ) + pkg.add_policy_component("last_token_logits", build_last_token_logits()) + pkg.add_policy_component("continue_predicate", build_boolean_not()) inputs = list(decoder.graph.inputs) outputs = list(decoder.graph.outputs) @@ -296,6 +300,10 @@ def build_decoder_workflow_metadata( batch_dimension = _shape_metadata(_port(token_input))[0] batch_int = {"dtype": "int64", "rank": 1, "shape": [batch_dimension]} + 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": { @@ -311,14 +319,16 @@ def build_decoder_workflow_metadata( "package.eos_ids": { "contract": {"dtype": "int64", "rank": 1, "shape": ["E"]}, "role": {"kind": "opaque"}, - "source": {"kind": "application", "name": "eos_token_ids"}, + "source": {"kind": "literal"}, "required": True, + "default": eos_token_id, }, "loop.iteration": { "contract": batch_int, "role": {"kind": "opaque"}, - "source": {"kind": "application", "name": "iteration"}, - "required": True, + "source": {"kind": "literal"}, + "required": False, + "default": 0, }, "loop.token_slot": { "contract": { @@ -327,8 +337,23 @@ def build_decoder_workflow_metadata( "shape": [batch_dimension, 1], }, "role": {"kind": "opaque"}, - "source": {"kind": "application", "name": "token_slot"}, - "required": True, + "source": {"kind": "literal"}, + "required": False, + "default": 0, + }, + "package.one_token": { + "contract": batch_int, + "role": {"kind": "opaque"}, + "source": {"kind": "literal"}, + "required": False, + "default": 1, + }, + "package.max_context": { + "contract": batch_int, + "role": {"kind": "opaque"}, + "source": {"kind": "literal"}, + "required": False, + "default": int(getattr(config, "max_position_embeddings", 4096)), }, } ) @@ -393,7 +418,19 @@ def build_decoder_workflow_metadata( "contract": _contract(past), "scope": "invocation", "initializer": setup_value, - "recurrence": {"kind": "invariant"}, + "recurrence": { + "kind": "growing", + "axis": next( + ( + index + for index, dimension in enumerate(_contract(past)["shape"]) + if "sequence" in str(dimension) + ), + 2, + ), + "increment": "package.one_token", + "max": "package.max_context", + }, } effect_name = f"state:{cell}" initial_effects[effect_name] = f"{effect_name}.0" @@ -414,11 +451,40 @@ def build_decoder_workflow_metadata( "nodes": [ _invoke(decoder_name, setup_decoder_inputs, setup_decoder_outputs), _invoke( - "token_sampler", + "last_token_logits", {"logits": "decoder.setup.logits"}, + {"last_logits": "decoder.setup.last_logits"}, + ), + _invoke( + "token_sampler", + {"logits": "decoder.setup.last_logits"}, {"token": "sample.setup"}, {"sample": _effect("sample.0", "sample.1")}, ), + _invoke( + "termination", + { + "token_ids": "sample.setup", + "eos_ids": "package.eos_ids", + "iteration": "loop.iteration", + "max_iterations": "request.max_iterations", + }, + {"done": "setup.done"}, + {"termination": _effect("termination.0", "termination.1")}, + ), + _invoke( + "continue_predicate", + {"done": "setup.done"}, + {"continue": "setup.continue"}, + ), + { + "kind": "emit", + "value": "sample.setup", + "output": "tokens", + "mode": "append", + "effect_name": "emit", + "effect": _effect("emit.0", "emit.1"), + }, _invoke( "token_state_update", { @@ -435,8 +501,13 @@ def build_decoder_workflow_metadata( "nodes": [ _invoke(decoder_name, body_decoder_inputs, body_decoder_outputs), _invoke( - "token_sampler", + "last_token_logits", {"logits": "decoder.body.logits"}, + {"last_logits": "decoder.body.last_logits"}, + ), + _invoke( + "token_sampler", + {"logits": "decoder.body.last_logits"}, {"token": "sample.body"}, {"sample": _effect("sample.1", "sample.2")}, ), @@ -455,7 +526,12 @@ def build_decoder_workflow_metadata( "max_iterations": "request.max_iterations", }, {"done": "loop.done"}, - {"termination": _effect("termination.0", "termination.1")}, + {"termination": _effect("termination.1", "termination.2")}, + ), + _invoke( + "continue_predicate", + {"done": "loop.done"}, + {"continue": "loop.continue"}, ), { "kind": "emit", @@ -463,7 +539,7 @@ def build_decoder_workflow_metadata( "output": "tokens", "mode": "append", "effect_name": "emit", - "effect": _effect("emit.0", "emit.1"), + "effect": _effect("emit.1", "emit.2"), }, ], } @@ -496,7 +572,7 @@ def build_decoder_workflow_metadata( "kind": "loop", "setup": setup, "body": body, - "condition": "loop.done", + "condition": "loop.continue", "max_iterations": "request.max_iterations", "carried": carried, }, @@ -546,6 +622,7 @@ def build_language_diffusion_pipeline_metadata( ) attach_policy_components(pkg, PolicyCapabilities(masked_update=True)) + pkg.add_policy_component("continue_predicate", build_boolean_not()) token_contract = _contract(token_input) mask_contract = { @@ -601,12 +678,6 @@ def build_language_diffusion_pipeline_metadata( "required": False, "default": num_inference_steps, }, - "loop.iteration": { - "contract": batch_int, - "role": {"kind": "opaque"}, - "source": {"kind": "application", "name": "iteration"}, - "required": True, - }, } def denoiser_invoke(tokens: str, prefix: str) -> dict[str, Any]: @@ -633,7 +704,7 @@ def update_invoke( "current_tokens": tokens, "proposed_tokens": f"{prefix}.proposal", "masked": mask, - "step": "loop.iteration", + "step": offset, "seed": "request.seed", "offset": offset, }, @@ -672,6 +743,11 @@ def update_invoke( "update.1", "update.2", ), + _invoke( + "continue_predicate", + {"done": "denoiser.body.done"}, + {"continue": "denoiser.body.continue"}, + ), { "kind": "emit", "value": "denoiser.body.tokens", @@ -738,7 +814,7 @@ def update_invoke( "kind": "loop", "setup": setup, "body": body, - "condition": "denoiser.body.done", + "condition": "denoiser.body.continue", "max_iterations": "request.max_iterations", "carried": carried, }, diff --git a/src/mobius/integrations/onnx_genai/workflow_metadata_test.py b/src/mobius/integrations/onnx_genai/workflow_metadata_test.py index 918cb009c..ef2a8d5d5 100644 --- a/src/mobius/integrations/onnx_genai/workflow_metadata_test.py +++ b/src/mobius/integrations/onnx_genai/workflow_metadata_test.py @@ -64,13 +64,14 @@ def test_language_diffusion_uses_exclusive_ssa_workflow(): graph = workflow["graph"] assert graph["kind"] == "loop" - assert graph["condition"] == "denoiser.body.done" + assert graph["condition"] == "denoiser.body.continue" assert graph["max_iterations"] == "request.max_iterations" assert [node["component"] for node in graph["setup"]["nodes"]] == [ "model", "masked_update", ] assert [node["kind"] for node in graph["body"]["nodes"]] == [ + "invoke", "invoke", "invoke", "emit", From 144dcfa1a284f358f4fafa9ab5f703f5a865aa5f Mon Sep 17 00:00:00 2001 From: justinchuby Date: Wed, 12 Aug 2026 20:18:07 +0000 Subject: [PATCH 008/151] Make generation workflows request executable Generate decoder masks, positions, empty caches, and loop counters from prompt and package literals so normal requests need no derived application tensors. Restructure decoder emission around one sampled token per loop iteration and carry logits, KV state, masks, positions, and iteration explicitly. Rank masked-diffusion proposals by model confidence with scheduled progressive commits and seeded counter tie-breaking. Add ORT multi-step parity coverage and refresh the PR #828 schema snapshot. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: justinchuby --- src/mobius/generation/__init__.py | 8 + src/mobius/generation/_policy_components.py | 208 +++++++++- .../generation/_policy_components_test.py | 182 +++++++- .../onnx_genai/auto_export_test.py | 12 + .../onnx_genai/workflow_metadata.py | 392 +++++++++++++----- .../onnx_genai/workflow_metadata_test.py | 9 +- tests/schemas/onnx_genai_4c3c4b6.schema.json | 64 +++ 7 files changed, 763 insertions(+), 112 deletions(-) diff --git a/src/mobius/generation/__init__.py b/src/mobius/generation/__init__.py index 9646793cd..b70afae59 100644 --- a/src/mobius/generation/__init__.py +++ b/src/mobius/generation/__init__.py @@ -11,11 +11,15 @@ PolicyRole, attach_policy_components, build_boolean_not, + build_decoder_state_initializer, + build_decoder_step_update, build_eos_termination, build_euler_solver_step, build_greedy_sampler, + build_integer_increment, build_last_token_logits, build_masked_token_update, + build_model_token_cast, build_seeded_categorical_sampler, build_speculative_acceptance, build_token_state_update, @@ -27,11 +31,15 @@ "PolicyRole", "attach_policy_components", "build_boolean_not", + "build_decoder_state_initializer", + "build_decoder_step_update", "build_eos_termination", "build_euler_solver_step", "build_greedy_sampler", + "build_integer_increment", "build_last_token_logits", "build_masked_token_update", + "build_model_token_cast", "build_seeded_categorical_sampler", "build_speculative_acceptance", "build_token_state_update", diff --git a/src/mobius/generation/_policy_components.py b/src/mobius/generation/_policy_components.py index 53545f8cb..85cea1e29 100644 --- a/src/mobius/generation/_policy_components.py +++ b/src/mobius/generation/_policy_components.py @@ -182,6 +182,161 @@ def build_boolean_not() -> PolicyComponent: return _component(PolicyRole.AUXILIARY, graph, {}) +def build_integer_increment() -> PolicyComponent: + """Build an explicit per-batch loop-counter increment.""" + graph, builder = _make_graph("integer_increment") + value = builder.input("value", dtype=ir.DataType.INT64, shape=["batch"]) + next_value = builder.op.Add(value, builder.op.Constant(value_int=1)) + next_value.shape = ir.Shape(["batch"]) + builder.add_output(next_value, "next_value") + return _component(PolicyRole.AUXILIARY, 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(PolicyRole.AUXILIARY, graph, {}) + + +def build_decoder_state_initializer( + decoder: ir.Model, + *, + token_input: str, + attention_mask_input: str, + position_ids_input: str | None, + cache_inputs: list[str], +) -> PolicyComponent: + """Build prompt-derived mask, position, token-slot, and empty-cache tensors.""" + graph, builder = _make_graph("decoder_state_initializer") + op = builder.op + decoder_inputs = {value.name: value for value in decoder.graph.inputs} + token_value = decoder_inputs[token_input] + prompt = builder.input( + "prompt_tokens", + dtype=token_value.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) + + attention_value = decoder_inputs[attention_mask_input] + attention = op.Cast( + op.ConstantOfShape(prompt_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", "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] + 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) + positions.shape = position_value.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]) + 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) + 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") + + 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: + 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 dimension {dimension_text!r}" + ) + cache_shape = op.Concat(*shape_parts, axis=0) + zero = 0.0 if value.dtype.is_floating_point else 0 + empty = op.ConstantOfShape( + cache_shape, + value=ir.tensor([zero], dtype=value.dtype), + ) + empty.shape = value.shape + builder.add_output(empty, name) + + return _component(PolicyRole.AUXILIARY, graph, {}) + + +def build_decoder_step_update( + *, + attention_dtype: ir.DataType, + position_dtype: ir.DataType | None, +) -> PolicyComponent: + """Build one-token attention-mask append and position increment.""" + graph, builder = _make_graph("decoder_step_update") + op = builder.op + attention = builder.input( + "attention_mask", + dtype=attention_dtype, + shape=["batch", "context"], + ) + 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: + position = builder.input( + "position_ids", + dtype=position_dtype, + shape=["batch", 1], + ) + next_position = op.Add(position, op.CastLike(op.Constant(value_int=1), position)) + next_position.shape = ir.Shape(["batch", 1]) + builder.add_output(next_position, "next_position_ids") + return _component(PolicyRole.AUXILIARY, graph, {}) + + def build_seeded_categorical_sampler() -> PolicyComponent: """Build deterministic categorical sampling with explicit seed and offset. @@ -364,8 +519,14 @@ def build_masked_token_update() -> PolicyComponent: 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) @@ -378,19 +539,50 @@ def build_masked_token_update() -> PolicyComponent: stream = op.Add( positions, op.Unsqueeze( - seed, + op.Add(op.Add(seed, offset), op.Mul(step, op.Constant(value_int=17))), op.Constant(value_ints=[-1]), ), ) - bucket = op.Mod(stream, op.Constant(value_int=8), fmod=0) - scheduled = op.Less( - bucket, - op.Unsqueeze( - op.Min(op.Add(step, op.Constant(value_int=1)), op.Constant(value_int=8)), - 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]))), ) - committed = op.And(masked, scheduled) updated = op.Where(committed, proposed, current) updated.shape = ir.Shape(["batch", "sequence"]) remaining = op.And(masked, op.Not(committed)) diff --git a/src/mobius/generation/_policy_components_test.py b/src/mobius/generation/_policy_components_test.py index 116ccf7e0..961c75ede 100644 --- a/src/mobius/generation/_policy_components_test.py +++ b/src/mobius/generation/_policy_components_test.py @@ -13,15 +13,19 @@ PolicyRole, attach_policy_components, build_boolean_not, + build_decoder_state_initializer, + build_decoder_step_update, build_eos_termination, build_euler_solver_step, build_greedy_sampler, build_last_token_logits, build_masked_token_update, + build_model_token_cast, build_seeded_categorical_sampler, build_speculative_acceptance, build_token_state_update, ) +from mobius.generation._policy_components import _make_graph def _run(component, tmp_path, feeds): @@ -31,6 +35,13 @@ def _run(component, tmp_path, feeds): 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_greedy_sampler_runtime(tmp_path): (tokens,) = _run( build_greedy_sampler(), @@ -53,6 +64,164 @@ def test_last_token_logits_and_continue_predicate_runtime(tmp_path): np.testing.assert_array_equal(continued, [False, True]) +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_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,) = _run( + build_eos_termination(), + tmp_path, + { + "token_ids": 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 = { @@ -98,20 +267,25 @@ def test_euler_solver_runtime_parity(tmp_path): 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, 5, 99]]) - np.testing.assert_array_equal(outputs[1], [[False, False, True]]) + 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]) @@ -121,8 +295,10 @@ def test_masked_update_runtime_parity(tmp_path): { "current_tokens": outputs[0], "proposed_tokens": np.array([[4, 5, 6]], np.int64), + "logits": logits, "masked": outputs[1], - "step": outputs[2], + "step": np.array([1], np.int64), + "total_steps": np.array([2], np.int64), "seed": np.array([7], np.int64), "offset": outputs[2], }, diff --git a/src/mobius/integrations/onnx_genai/auto_export_test.py b/src/mobius/integrations/onnx_genai/auto_export_test.py index e7a39a666..ee800473b 100644 --- a/src/mobius/integrations/onnx_genai/auto_export_test.py +++ b/src/mobius/integrations/onnx_genai/auto_export_test.py @@ -69,6 +69,18 @@ def test_dispatch_decoder(tmp_path): assert workflow["components"]["token_sampler"]["policy"]["role"] == "token_sampler" assert workflow["components"]["termination"]["policy"]["role"] == ("termination_predicate") assert workflow["graph"]["kind"] == "loop" + assert all( + value["source"]["kind"] != "application" for value in workflow["inputs"].values() + ) + assert [node["component"] for node in workflow["graph"]["setup"]["nodes"]] == [ + "decoder_state_initializer", + "model", + ] + body = workflow["graph"]["body"]["nodes"] + assert [node["kind"] for node in body].count("emit") == 1 + assert next(node for node in body if node["kind"] == "emit")["value"] == "sample.body" + assert workflow["state"]["iteration"]["initializer"] == "package.zero_iteration" + assert workflow["state"]["token"]["initializer"] == "initializer.token_slot" assert (tmp_path / "policies" / "token_sampler.onnx").is_file() diff --git a/src/mobius/integrations/onnx_genai/workflow_metadata.py b/src/mobius/integrations/onnx_genai/workflow_metadata.py index 82ae5efcf..816b3d873 100644 --- a/src/mobius/integrations/onnx_genai/workflow_metadata.py +++ b/src/mobius/integrations/onnx_genai/workflow_metadata.py @@ -16,7 +16,11 @@ PolicyCapabilities, attach_policy_components, build_boolean_not, + build_decoder_state_initializer, + build_decoder_step_update, + build_integer_increment, build_last_token_logits, + build_model_token_cast, ) from mobius.integrations.onnx_genai.inference_metadata import ( _port, @@ -240,6 +244,7 @@ def build_decoder_workflow_metadata( ) pkg.add_policy_component("last_token_logits", build_last_token_logits()) pkg.add_policy_component("continue_predicate", build_boolean_not()) + pkg.add_policy_component("iteration_increment", build_integer_increment()) inputs = list(decoder.graph.inputs) outputs = list(decoder.graph.outputs) @@ -247,11 +252,21 @@ def build_decoder_workflow_metadata( ( value for value in inputs - if value.dtype in {ir.DataType.INT32, ir.DataType.INT64} + 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 ), - None, + 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( ( @@ -274,10 +289,87 @@ def build_decoder_workflow_metadata( "decoder workflow requires rank-2 token input and rank-3 logits output" ) + output_by_suffix = {value.name: value for value in outputs} + cache_pairs: list[tuple[ir.Value, ir.Value]] = [] + for value in inputs: + candidates = [ + value.name.replace("past_key_values", "present"), + value.name.replace("past.", "present."), + ] + present = next( + (output_by_suffix.get(name) for name in candidates 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} + 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 "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 "position" in value.name + ), + next((value for value in integer_rank2 if value is not attention_input), None), + ) + if attention_input is None: + raise ValueError( + "standard decoder workflow requires a derived rank-2 attention-mask input" + ) + derived_names = cache_names | {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}") + pkg.add_policy_component( + "decoder_state_initializer", + build_decoder_state_initializer( + decoder, + token_input=token_input.name, + 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), + ), + ) + 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, + ), + ) + 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 = { @@ -323,24 +415,13 @@ def build_decoder_workflow_metadata( "required": True, "default": eos_token_id, }, - "loop.iteration": { + "package.zero_iteration": { "contract": batch_int, "role": {"kind": "opaque"}, "source": {"kind": "literal"}, "required": False, "default": 0, }, - "loop.token_slot": { - "contract": { - "dtype": "int64", - "rank": 2, - "shape": [batch_dimension, 1], - }, - "role": {"kind": "opaque"}, - "source": {"kind": "literal"}, - "required": False, - "default": 0, - }, "package.one_token": { "contract": batch_int, "role": {"kind": "opaque"}, @@ -358,23 +439,20 @@ def build_decoder_workflow_metadata( } ) - cache_pairs: list[tuple[ir.Value, ir.Value]] = [] - output_by_suffix = {value.name: value for value in outputs} for value in inputs: if value is token_input: continue - candidates = [ - value.name.replace("past_key_values", "present"), - value.name.replace("past.", "present."), - ] - present = next( - (output_by_suffix.get(name) for name in candidates if name in output_by_suffix), - None, - ) - if present is not None: - cache_pairs.append((value, present)) + if value.name in cache_names: body_decoder_inputs[value.name] = f"state.{value.name}.body" - body_decoder_inputs[token_input.name] = "state.token.body" + setup_decoder_inputs[value.name] = f"initializer.{value.name}" + setup_decoder_inputs[attention_input.name] = f"initializer.{attention_input.name}" + body_decoder_inputs[attention_input.name] = "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" + 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"} @@ -386,9 +464,21 @@ def build_decoder_workflow_metadata( "shape": [batch_dimension, 1], }, "scope": "invocation", - "initializer": f"request.{token_input.name}", + "initializer": "initializer.token_slot", "recurrence": {"kind": "invariant"}, - } + }, + "iteration": { + "contract": batch_int, + "scope": "invocation", + "initializer": "package.zero_iteration", + "recurrence": {"kind": "invariant"}, + }, + "logits": { + "contract": _contract(logits_output), + "scope": "invocation", + "initializer": "decoder.setup.logits", + "recurrence": {"kind": "invariant"}, + }, } initial_effects = { "sample": "sample.0", @@ -396,11 +486,13 @@ def build_decoder_workflow_metadata( "state": "state.0", "emit": "emit.0", "state:token": "state:token.0", + "state:iteration": "state:iteration.0", + "state:logits": "state:logits.0", } carried = [ { "cell": "token", - "current": "token.setup", + "current": "initializer.token_slot", "body_input": "state.token.body", "body_output": "token.body", "next": "token.final", @@ -408,6 +500,76 @@ def build_decoder_workflow_metadata( "write_effect": _effect("state:token.read", "state:token.1"), } ] + carried.extend( + [ + { + "cell": "iteration", + "current": "package.zero_iteration", + "body_input": "state.iteration.body", + "body_output": "iteration.body", + "next": "state.iteration.final", + "read_effect": _effect("state:iteration.0", "state:iteration.read"), + "write_effect": _effect("state:iteration.read", "state:iteration.1"), + }, + { + "cell": "logits", + "current": "decoder.setup.logits", + "body_input": "state.logits.body", + "body_output": "decoder.body.logits", + "next": "state.logits.final", + "read_effect": _effect("state:logits.0", "state:logits.read"), + "write_effect": _effect("state:logits.read", "state:logits.1"), + }, + ] + ) + decoder_state_specs = { + "attention_mask": ( + { + "dtype": _contract(attention_input)["dtype"], + "rank": 2, + "shape": [batch_dimension, "context"], + }, + "initializer.body_attention_mask", + "decoder_step.body_attention_mask", + { + "kind": "growing", + "axis": 1, + "increment": "package.one_token", + "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"), + } + ) for past, present in cache_pairs: cell = f"cache_{len(carried)}" setup_value = f"decoder.setup.{present.name}" @@ -449,84 +611,68 @@ def build_decoder_workflow_metadata( setup = { "kind": "sequence", "nodes": [ - _invoke(decoder_name, setup_decoder_inputs, setup_decoder_outputs), - _invoke( - "last_token_logits", - {"logits": "decoder.setup.logits"}, - {"last_logits": "decoder.setup.last_logits"}, - ), - _invoke( - "token_sampler", - {"logits": "decoder.setup.last_logits"}, - {"token": "sample.setup"}, - {"sample": _effect("sample.0", "sample.1")}, - ), _invoke( - "termination", - { - "token_ids": "sample.setup", - "eos_ids": "package.eos_ids", - "iteration": "loop.iteration", - "max_iterations": "request.max_iterations", - }, - {"done": "setup.done"}, - {"termination": _effect("termination.0", "termination.1")}, - ), - _invoke( - "continue_predicate", - {"done": "setup.done"}, - {"continue": "setup.continue"}, - ), - { - "kind": "emit", - "value": "sample.setup", - "output": "tokens", - "mode": "append", - "effect_name": "emit", - "effect": _effect("emit.0", "emit.1"), - }, - _invoke( - "token_state_update", + "decoder_state_initializer", + {"prompt_tokens": f"request.{token_input.name}"}, { - "current": "loop.token_slot", - "update": "sample.setup", + attention_input.name: f"initializer.{attention_input.name}", + "body_attention_mask": "initializer.body_attention_mask", + "token_slot": "initializer.token_slot", + **( + { + 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)}, }, - {"next": "token.setup"}, - {"state": _effect("state.0", "state.1")}, ), + _invoke(decoder_name, setup_decoder_inputs, setup_decoder_outputs), ], } body = { "kind": "sequence", "nodes": [ - _invoke(decoder_name, body_decoder_inputs, body_decoder_outputs), _invoke( "last_token_logits", - {"logits": "decoder.body.logits"}, + {"logits": "state.logits.body"}, {"last_logits": "decoder.body.last_logits"}, ), _invoke( "token_sampler", {"logits": "decoder.body.last_logits"}, {"token": "sample.body"}, - {"sample": _effect("sample.1", "sample.2")}, + {"sample": _effect("sample.0", "sample.1")}, ), _invoke( "token_state_update", {"current": "state.token.body", "update": "sample.body"}, {"next": "token.body"}, - {"state": _effect("state.1", "state.2")}, + {"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", { "token_ids": "sample.body", "eos_ids": "package.eos_ids", - "iteration": "loop.iteration", + "iteration": "state.iteration.body", "max_iterations": "request.max_iterations", }, {"done": "loop.done"}, - {"termination": _effect("termination.1", "termination.2")}, + {"termination": _effect("termination.0", "termination.1")}, ), _invoke( "continue_predicate", @@ -539,8 +685,33 @@ def build_decoder_workflow_metadata( "output": "tokens", "mode": "append", "effect_name": "emit", - "effect": _effect("emit.1", "emit.2"), + "effect": _effect("emit.0", "emit.1"), }, + _invoke( + "iteration_increment", + {"value": "state.iteration.body"}, + {"next_value": "iteration.body"}, + ), + _invoke(decoder_name, body_decoder_inputs, body_decoder_outputs), + _invoke( + "decoder_step_update", + { + "attention_mask": "state.attention_mask.body", + **( + {"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 {} + ), + }, + ), ], } @@ -623,6 +794,7 @@ def build_language_diffusion_pipeline_metadata( attach_policy_components(pkg, PolicyCapabilities(masked_update=True)) pkg.add_policy_component("continue_predicate", build_boolean_not()) + pkg.add_policy_component("iteration_increment", build_integer_increment()) token_contract = _contract(token_input) mask_contract = { @@ -678,6 +850,20 @@ def build_language_diffusion_pipeline_metadata( "required": False, "default": num_inference_steps, }, + "package.zero_iteration": { + "contract": batch_int, + "role": {"kind": "opaque"}, + "source": {"kind": "literal"}, + "required": False, + "default": 0, + }, + "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]: @@ -693,7 +879,10 @@ def denoiser_invoke(tokens: str, prefix: str) -> dict[str, Any]: def update_invoke( tokens: str, mask: str, + iteration: str, offset: str, + logits: str, + proposal: str, prefix: str, effect_in: str, effect_out: str, @@ -702,9 +891,11 @@ def update_invoke( "masked_update", { "current_tokens": tokens, - "proposed_tokens": f"{prefix}.proposal", + "proposed_tokens": proposal, + "logits": logits, "masked": mask, - "step": offset, + "step": iteration, + "total_steps": "package.num_steps", "seed": "request.seed", "offset": offset, }, @@ -719,35 +910,32 @@ def update_invoke( setup = { "kind": "sequence", - "nodes": [ - denoiser_invoke("request.input_ids", "denoiser.setup"), - update_invoke( - "request.input_ids", - "request.mask", - "request.rng_offset", - "denoiser.setup", - "update.0", - "update.1", - ), - ], + "nodes": [denoiser_invoke("request.input_ids", "denoiser.setup")], } body = { "kind": "sequence", "nodes": [ - denoiser_invoke("state.tokens.body", "denoiser.body"), update_invoke( "state.tokens.body", "state.mask.body", + "state.iteration.body", "state.rng_offset.body", + "state.logits.body", + "state.proposal.body", "denoiser.body", + "update.0", "update.1", - "update.2", ), _invoke( "continue_predicate", {"done": "denoiser.body.done"}, {"continue": "denoiser.body.continue"}, ), + _invoke( + "iteration_increment", + {"value": "state.iteration.body"}, + {"next_value": "denoiser.body.iteration"}, + ), { "kind": "emit", "value": "denoiser.body.tokens", @@ -756,13 +944,25 @@ def update_invoke( "effect_name": "emit", "effect": _effect("emit.0", "emit.1"), }, + denoiser_invoke("denoiser.body.tokens", "denoiser.body"), ], } state_specs = { - "tokens": (token_contract, "request.input_ids", "denoiser.setup.tokens"), - "mask": (mask_contract, "request.mask", "denoiser.setup.mask"), - "rng_offset": (batch_int, "request.rng_offset", "denoiser.setup.rng_offset"), + "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"), + "iteration": (batch_int, "package.zero_iteration", "package.zero_iteration"), + "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]] = [] diff --git a/src/mobius/integrations/onnx_genai/workflow_metadata_test.py b/src/mobius/integrations/onnx_genai/workflow_metadata_test.py index ef2a8d5d5..c9bbb2c98 100644 --- a/src/mobius/integrations/onnx_genai/workflow_metadata_test.py +++ b/src/mobius/integrations/onnx_genai/workflow_metadata_test.py @@ -66,17 +66,16 @@ def test_language_diffusion_uses_exclusive_ssa_workflow(): assert graph["kind"] == "loop" assert graph["condition"] == "denoiser.body.continue" assert graph["max_iterations"] == "request.max_iterations" - assert [node["component"] for node in graph["setup"]["nodes"]] == [ - "model", - "masked_update", - ] + assert [node["component"] for node in graph["setup"]["nodes"]] == ["model"] assert [node["kind"] for node in graph["body"]["nodes"]] == [ "invoke", "invoke", "invoke", "emit", + "invoke", ] - assert graph["body"]["nodes"][-1]["mode"] == "replace" + assert graph["body"]["nodes"][0]["inputs"]["total_steps"] == "package.num_steps" + assert graph["body"]["nodes"][-2]["mode"] == "replace" def test_language_diffusion_rejects_zero_steps(): diff --git a/tests/schemas/onnx_genai_4c3c4b6.schema.json b/tests/schemas/onnx_genai_4c3c4b6.schema.json index 53a8c5b94..16055a43b 100644 --- a/tests/schemas/onnx_genai_4c3c4b6.schema.json +++ b/tests/schemas/onnx_genai_4c3c4b6.schema.json @@ -5784,6 +5784,56 @@ } ] }, + "WorkflowBranchEffectMerge": { + "additionalProperties": false, + "properties": { + "cases": { + "additionalProperties": { + "type": "string" + }, + "type": "object" + }, + "default": { + "type": [ + "string", + "null" + ] + }, + "incoming": { + "type": "string" + }, + "produces": { + "type": "string" + } + }, + "required": [ + "incoming", + "cases", + "produces" + ], + "type": "object" + }, + "WorkflowBranchOutput": { + "additionalProperties": false, + "properties": { + "cases": { + "additionalProperties": { + "type": "string" + }, + "type": "object" + }, + "default": { + "type": [ + "string", + "null" + ] + } + }, + "required": [ + "cases" + ], + "type": "object" + }, "WorkflowComponent": { "additionalProperties": false, "properties": { @@ -6140,10 +6190,24 @@ } ] }, + "effects": { + "additionalProperties": { + "$ref": "#/$defs/WorkflowBranchEffectMerge" + }, + "default": {}, + "type": "object" + }, "kind": { "const": "branch", "type": "string" }, + "outputs": { + "additionalProperties": { + "$ref": "#/$defs/WorkflowBranchOutput" + }, + "default": {}, + "type": "object" + }, "predicate": { "type": "string" } From 99820aa326816962bef490f8c20e60a6f6d9f0ce Mon Sep 17 00:00:00 2001 From: justinchuby Date: Wed, 12 Aug 2026 20:27:32 +0000 Subject: [PATCH 009/151] Migrate diffusion metadata to workflow SSA Use lexical loop induction for denoiser timesteps, generated ONNX schedule and Euler solver components, explicit carried latent state, and final VAE image emission. Remove diffusion dispatch through the legacy iterative strategy. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: justinchuby --- src/mobius/generation/__init__.py | 4 + src/mobius/generation/_policy_components.py | 23 ++ .../integrations/onnx_genai/__init__.py | 4 + .../integrations/onnx_genai/auto_export.py | 8 +- .../onnx_genai/auto_export_test.py | 57 +++- .../onnx_genai/workflow_metadata.py | 299 ++++++++++++++++++ tests/schemas/onnx_genai_4c3c4b6.schema.json | 58 +++- 7 files changed, 431 insertions(+), 22 deletions(-) diff --git a/src/mobius/generation/__init__.py b/src/mobius/generation/__init__.py index b70afae59..4ed7e10d7 100644 --- a/src/mobius/generation/__init__.py +++ b/src/mobius/generation/__init__.py @@ -17,9 +17,11 @@ build_euler_solver_step, build_greedy_sampler, build_integer_increment, + build_iteration_cast, build_last_token_logits, build_masked_token_update, build_model_token_cast, + build_schedule_constant, build_seeded_categorical_sampler, build_speculative_acceptance, build_token_state_update, @@ -37,9 +39,11 @@ "build_euler_solver_step", "build_greedy_sampler", "build_integer_increment", + "build_iteration_cast", "build_last_token_logits", "build_masked_token_update", "build_model_token_cast", + "build_schedule_constant", "build_seeded_categorical_sampler", "build_speculative_acceptance", "build_token_state_update", diff --git a/src/mobius/generation/_policy_components.py b/src/mobius/generation/_policy_components.py index 85cea1e29..c8ad9e8fe 100644 --- a/src/mobius/generation/_policy_components.py +++ b/src/mobius/generation/_policy_components.py @@ -192,6 +192,29 @@ def build_integer_increment() -> PolicyComponent: return _component(PolicyRole.AUXILIARY, 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(PolicyRole.AUXILIARY, 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(PolicyRole.AUXILIARY, 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") diff --git a/src/mobius/integrations/onnx_genai/__init__.py b/src/mobius/integrations/onnx_genai/__init__.py index 2e6120753..ca8b02fb0 100644 --- a/src/mobius/integrations/onnx_genai/__init__.py +++ b/src/mobius/integrations/onnx_genai/__init__.py @@ -63,10 +63,12 @@ from mobius.integrations.onnx_genai.workflow_metadata import ( build_audio_codec_workflow_metadata, build_decoder_workflow_metadata, + build_diffusion_workflow_metadata, build_language_diffusion_pipeline_metadata, build_tts_workflow_metadata, write_audio_codec_workflow_metadata, write_decoder_workflow_metadata, + write_diffusion_workflow_metadata, write_language_diffusion_workflow_metadata, write_tts_workflow_metadata, ) @@ -78,6 +80,7 @@ "add_policy_components_to_workflow", "build_decoder_metadata", "build_decoder_workflow_metadata", + "build_diffusion_workflow_metadata", "build_diffusion_pipeline_metadata", "build_language_diffusion_pipeline_metadata", "build_audio_codec_workflow_metadata", @@ -95,6 +98,7 @@ "translate_comfyui_workflow_file", "write_decoder_metadata", "write_decoder_workflow_metadata", + "write_diffusion_workflow_metadata", "write_language_diffusion_workflow_metadata", "write_diffusion_pipeline_metadata", "write_audio_codec_workflow_metadata", diff --git a/src/mobius/integrations/onnx_genai/auto_export.py b/src/mobius/integrations/onnx_genai/auto_export.py index 880f7441c..14fb22eb7 100644 --- a/src/mobius/integrations/onnx_genai/auto_export.py +++ b/src/mobius/integrations/onnx_genai/auto_export.py @@ -27,13 +27,13 @@ add_explicit_package_io, add_policy_components_to_workflow, load_diffusers_scheduler_config, - write_diffusion_pipeline_metadata, write_multimodal_pipeline_metadata, write_speech_to_text_pipeline_metadata, ) from mobius.integrations.onnx_genai.workflow_metadata import ( write_audio_codec_workflow_metadata, write_decoder_workflow_metadata, + write_diffusion_workflow_metadata, write_language_diffusion_workflow_metadata, write_tts_workflow_metadata, ) @@ -464,12 +464,10 @@ def write_onnx_genai_config( # 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( + path = write_diffusion_workflow_metadata( + pkg, output_dir, num_inference_steps=num_inference_steps, - scheduler=scheduler, - guidance_scale=guidance_scale, - **kwargs, ) artifacts = {"inference_metadata": path} # Emit the CLIP tokenizer.json for text-conditioned pipelines so the diff --git a/src/mobius/integrations/onnx_genai/auto_export_test.py b/src/mobius/integrations/onnx_genai/auto_export_test.py index ee800473b..5f20559d4 100644 --- a/src/mobius/integrations/onnx_genai/auto_export_test.py +++ b/src/mobius/integrations/onnx_genai/auto_export_test.py @@ -46,6 +46,46 @@ 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() @@ -112,7 +152,7 @@ def test_dispatch_language_diffusion(tmp_path): def test_dispatch_diffusion(tmp_path): - pkg = _DiffusionPkg({"denoiser": object(), "vae": object()}) + pkg = _diffusion_package() arts = write_onnx_genai_config( pkg, str(tmp_path), @@ -121,8 +161,10 @@ 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["graph"]["nodes"][0]["iteration"]["value"] == "loop.iteration" + assert workflow["graph"]["nodes"][1]["component"] == "vae_decoder" + assert "strategy" not in meta["pipeline"] def test_single_diffusion_component_uses_flat_model_path(tmp_path): @@ -201,7 +243,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), @@ -226,7 +268,7 @@ 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), @@ -247,7 +289,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), @@ -256,7 +298,8 @@ 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"]["shape"] == [16] def test_dispatch_vision_multimodal_pipeline(tmp_path): diff --git a/src/mobius/integrations/onnx_genai/workflow_metadata.py b/src/mobius/integrations/onnx_genai/workflow_metadata.py index 816b3d873..ddc94774a 100644 --- a/src/mobius/integrations/onnx_genai/workflow_metadata.py +++ b/src/mobius/integrations/onnx_genai/workflow_metadata.py @@ -19,8 +19,10 @@ build_decoder_state_initializer, build_decoder_step_update, build_integer_increment, + build_iteration_cast, build_last_token_logits, build_model_token_cast, + build_schedule_constant, ) from mobius.integrations.onnx_genai.inference_metadata import ( _port, @@ -224,6 +226,303 @@ def write_tts_workflow_metadata(pkg: Any, output_dir: str, config: Any) -> str: 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 build_diffusion_workflow_metadata( + pkg: Any, + *, + num_inference_steps: int, + schedule: list[float] | None = None, +) -> dict[str, Any]: + """Build a fixed-schedule diffusion workflow with explicit latent state.""" + 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("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("Euler 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), + ) + + attach_policy_components(pkg, PolicyCapabilities(solver="euler")) + pkg.add_policy_component("continue_predicate", build_boolean_not()) + schedule_values = schedule or [ + 1.0 - index / num_inference_steps for index in range(num_inference_steps + 1) + ] + pkg.add_policy_component("diffusion_schedule", build_schedule_constant(schedule_values)) + if timestep_input.dtype != ir.DataType.INT64: + pkg.add_policy_component("iteration_cast", build_iteration_cast(timestep_input.dtype)) + + batch = _contract(sample_input)["shape"][0] + batch_int = {"dtype": "int64", "rank": 1, "shape": [batch]} + batch_bool = {"dtype": "bool", "rank": 1, "shape": [batch]} + inputs: dict[str, Any] = { + "request.latent": { + "contract": _contract(sample_input), + "role": {"kind": "opaque"}, + "source": {"kind": "application", "name": "latent"}, + "required": True, + }, + "request.max_iterations": { + "contract": batch_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"}) + ] + 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": _contract(value), + "role": ( + {"kind": "runtime", "version": "1.0", "role": "prompt_tokens"} + if index == 0 + else {"kind": "opaque"} + ), + "source": { + "kind": "request" if index == 0 else "application", + "field": "prompt_tokens" if index == 0 else None, + "name": value.name if index else None, + }, + "required": True, + } + inputs[name]["source"] = { + key: item for key, item in inputs[name]["source"].items() if item is not None + } + text_inputs[value.name] = name + conditioning_value = "conditioning.hidden_states" + setup_nodes.append( + _invoke( + text_name, + text_inputs, + {conditioning_output.name: conditioning_value}, + ) + ) + setup_nodes.append( + _invoke( + "continue_predicate", + {"done": "package.false"}, + {"continue": "setup.continue"}, + ) + ) + + denoiser_inputs = { + sample_input.name: "state.latent.body", + timestep_input.name: ( + "diffusion.timestep" + if timestep_input.dtype != ir.DataType.INT64 + else "loop.iteration" + ), + } + 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]] = [] + if timestep_input.dtype != ir.DataType.INT64: + body_nodes.append( + _invoke( + "iteration_cast", + {"iteration": "loop.iteration"}, + {"timestep": "diffusion.timestep"}, + ) + ) + body_nodes.extend( + [ + _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( + "continue_predicate", + {"done": "package.false"}, + {"continue": "loop.continue"}, + ), + ] + ) + latent_effect = "state:latent" + workflow = { + "manifest": { + "ir_version": "1.0", + "onnx_opsets": {"ai.onnx": OPSET_VERSION}, + "capabilities": [ + "workflow_ssa", + "linear_effects", + "nested_control_flow", + "typed_emit", + ], + }, + "inputs": inputs, + "outputs": { + "image": { + "contract": _contract(vae_output), + "role": "image", + "stage": "pre_adapter", + } + }, + "components": { + name: _component(model, _artifact(name, len(pkg))) for name, model in pkg.items() + }, + "state": { + "latent": { + "contract": _contract(sample_input), + "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( + vae_name, + {vae_input.name: "latent.final"}, + {vae_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": workflow}} + add_policy_components_to_workflow(metadata, pkg) + return metadata + + +def write_diffusion_workflow_metadata( + pkg: Any, + output_dir: str, + *, + num_inference_steps: int, +) -> str: + os.makedirs(output_dir, exist_ok=True) + metadata = build_diffusion_workflow_metadata(pkg, num_inference_steps=num_inference_steps) + path = os.path.join(output_dir, "inference_metadata.yaml") + with open(path, "w", encoding="utf-8") as handle: + yaml.safe_dump(metadata, handle, sort_keys=False) + return path + + def build_decoder_workflow_metadata( pkg: Any, config: Any, diff --git a/tests/schemas/onnx_genai_4c3c4b6.schema.json b/tests/schemas/onnx_genai_4c3c4b6.schema.json index 16055a43b..a8bbd429c 100644 --- a/tests/schemas/onnx_genai_4c3c4b6.schema.json +++ b/tests/schemas/onnx_genai_4c3c4b6.schema.json @@ -871,20 +871,31 @@ "type": "string" }, "ImageOutputBinding": { - "description": "One named tensor output produced by an image preprocessing program.\n\nThe output binds a generic content role to an ARBITRARY endpoint name with a\nDECLARED dtype. Neither the name nor the content role is inferred from a model\nidentity, and the dtype is always explicit rather than derived from the model.", + "description": "One named tensor output produced by an image preprocessing program.\n\nThe output binds a processor-local value to a typed workflow SSA name.\nNeither the name nor the content role is inferred from a model identity.", "properties": { "content": { "$ref": "#/$defs/ImageOutputContent", "description": "Generic content role this tensor carries (pixels, coordinates, grid,\noriginal size, or validity mask) — never a model-family label." }, + "contract": { + "anyOf": [ + { + "$ref": "#/$defs/TensorContract" + }, + { + "type": "null" + } + ], + "description": "Full workflow tensor contract. Required when `pipeline.workflow` is present." + }, "dtype": { "$ref": "#/$defs/TensorDType", "description": "Declared output dtype. Always explicit; never inferred from the model." }, "name": { - "description": "Arbitrary pipeline endpoint name this tensor is bound to (model DATA).", + "description": "Workflow SSA value produced by the preprocessing adapter invocation.", "examples": [ - "vision_encoder.pixel_values" + "image.pixel_values" ], "minLength": 1, "type": "string" @@ -905,15 +916,13 @@ ] }, "source": { - "description": "Named value produced by a transform.\n\nAbsent preserves the legacy content-derived binding behavior.", + "description": "Named processor-local value produced by a transform.", "minLength": 1, - "type": [ - "string", - "null" - ] + "type": "string" } }, "required": [ + "source", "name", "content", "dtype" @@ -952,10 +961,10 @@ "type": "string" }, "ImagePreprocessingProgram": { - "description": "Generic image preprocessing program: an ordered transform pipeline plus the\nnamed tensor outputs it emits.\n\nThe program is expressed entirely as parameterized, architecture-neutral\ndata. Transform operations are generic (decode, resize, rescale, normalize,\ntile, patchify, pad); outputs bind a produced tensor to an ARBITRARY pipeline\nendpoint name with a DECLARED dtype. A model may name an output\n`pixel_position_ids`, `image_grid_thw`, or anything else — that string is\ndata carried in the model's metadata, never a branch in the runtime.", + "description": "Generic image preprocessing program: an ordered transform pipeline plus the\nnamed workflow SSA tensor outputs it emits.\n\nThe program is expressed entirely as parameterized, architecture-neutral\ndata. Transform operations are generic (decode, resize, rescale, normalize,\ntile, patchify, pad). In workflow metadata, outputs are materialized by a\nmanifest-pinned preprocessing adapter invocation and bind processor-local\nvalues to typed SSA names. A package may name an output `pixel_position_ids`,\n`image_grid_thw`, or anything else without introducing runtime model-family\ndispatch.", "properties": { "outputs": { - "description": "Named tensor outputs the program emits, each bound to a pipeline endpoint.", + "description": "Named tensor outputs the program emits, each bound to a workflow SSA value.", "items": { "$ref": "#/$defs/ImageOutputBinding" }, @@ -6023,6 +6032,24 @@ ], "type": "object" }, + "WorkflowLoopIteration": { + "additionalProperties": false, + "properties": { + "contract": { + "$ref": "#/$defs/TensorContract", + "description": "`int64` scalar or rank-one broadcast contract." + }, + "value": { + "description": "SSA value containing the current zero-based iteration.", + "type": "string" + } + }, + "required": [ + "value", + "contract" + ], + "type": "object" + }, "WorkflowManifest": { "additionalProperties": false, "properties": { @@ -6151,6 +6178,17 @@ "condition": { "type": "string" }, + "iteration": { + "anyOf": [ + { + "$ref": "#/$defs/WorkflowLoopIteration" + }, + { + "type": "null" + } + ], + "description": "Optional zero-based induction value, scoped to this loop's body and condition." + }, "kind": { "const": "loop", "type": "string" From ba813b90f2e5ed77e7944167c456dc7243ed5a0b Mon Sep 17 00:00:00 2001 From: justinchuby Date: Wed, 12 Aug 2026 20:32:30 +0000 Subject: [PATCH 010/151] Migrate multimodal metadata to workflow SSA Invoke the published typed image preprocessing adapter, vision and optional audio encoders, embedding model, and autoregressive decoder through one generic SSA workflow. Carry decoder logits, tokens, masks, positions, and KV state explicitly and remove multimodal dispatch through the legacy composite strategy. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: justinchuby --- src/mobius/generation/_policy_components.py | 10 +- .../integrations/onnx_genai/__init__.py | 4 + .../integrations/onnx_genai/auto_export.py | 17 +- .../onnx_genai/auto_export_test.py | 198 +++---- .../onnx_genai/inference_metadata.py | 1 + .../onnx_genai/workflow_metadata.py | 548 ++++++++++++++++++ 6 files changed, 642 insertions(+), 136 deletions(-) diff --git a/src/mobius/generation/_policy_components.py b/src/mobius/generation/_policy_components.py index c8ad9e8fe..dbbcf4a32 100644 --- a/src/mobius/generation/_policy_components.py +++ b/src/mobius/generation/_policy_components.py @@ -228,7 +228,8 @@ def build_model_token_cast(dtype: ir.DataType) -> PolicyComponent: def build_decoder_state_initializer( decoder: ir.Model, *, - token_input: str, + token_input: str | None, + prompt_dtype: ir.DataType | None = None, attention_mask_input: str, position_ids_input: str | None, cache_inputs: list[str], @@ -237,10 +238,13 @@ def build_decoder_state_initializer( graph, builder = _make_graph("decoder_state_initializer") op = builder.op decoder_inputs = {value.name: value for value in decoder.graph.inputs} - token_value = decoder_inputs[token_input] + 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=token_value.dtype, + dtype=prompt_dtype, shape=["batch", "prompt_sequence"], ) prompt_shape = op.Shape(prompt) diff --git a/src/mobius/integrations/onnx_genai/__init__.py b/src/mobius/integrations/onnx_genai/__init__.py index ca8b02fb0..001c06eac 100644 --- a/src/mobius/integrations/onnx_genai/__init__.py +++ b/src/mobius/integrations/onnx_genai/__init__.py @@ -66,11 +66,13 @@ build_diffusion_workflow_metadata, build_language_diffusion_pipeline_metadata, build_tts_workflow_metadata, + build_vlm_workflow_metadata, write_audio_codec_workflow_metadata, write_decoder_workflow_metadata, write_diffusion_workflow_metadata, write_language_diffusion_workflow_metadata, write_tts_workflow_metadata, + write_vlm_workflow_metadata, ) __all__ = [ @@ -88,6 +90,7 @@ "build_pipeline_metadata_for_workflow", "build_speech_to_text_pipeline_metadata", "build_tts_workflow_metadata", + "build_vlm_workflow_metadata", "convert_comfyui_workflow", "decoder_metadata_from_config", "moe_metadata_from_config", @@ -105,5 +108,6 @@ "write_multimodal_pipeline_metadata", "write_speech_to_text_pipeline_metadata", "write_tts_workflow_metadata", + "write_vlm_workflow_metadata", "write_onnx_genai_config", ] diff --git a/src/mobius/integrations/onnx_genai/auto_export.py b/src/mobius/integrations/onnx_genai/auto_export.py index 14fb22eb7..7a0971775 100644 --- a/src/mobius/integrations/onnx_genai/auto_export.py +++ b/src/mobius/integrations/onnx_genai/auto_export.py @@ -27,7 +27,6 @@ add_explicit_package_io, add_policy_components_to_workflow, load_diffusers_scheduler_config, - write_multimodal_pipeline_metadata, write_speech_to_text_pipeline_metadata, ) from mobius.integrations.onnx_genai.workflow_metadata import ( @@ -36,6 +35,7 @@ write_diffusion_workflow_metadata, write_language_diffusion_workflow_metadata, write_tts_workflow_metadata, + write_vlm_workflow_metadata, ) _LOGGER = logging.getLogger(__name__) @@ -491,19 +491,12 @@ 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( + 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: diff --git a/src/mobius/integrations/onnx_genai/auto_export_test.py b/src/mobius/integrations/onnx_genai/auto_export_test.py index 5f20559d4..3ff9319be 100644 --- a/src/mobius/integrations/onnx_genai/auto_export_test.py +++ b/src/mobius/integrations/onnx_genai/auto_export_test.py @@ -90,6 +90,58 @@ 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( [], @@ -303,151 +355,55 @@ def test_dispatch_diffusion_auto_reads_scheduler_from_source(tmp_path): def test_dispatch_vision_multimodal_pipeline(tmp_path): - pkg = _MultimodalPkg( - { - "decoder": object(), - "vision_encoder": object(), - "embedding": object(), - } - ) + pkg = _vlm_package() artifacts = write_onnx_genai_config(pkg, str(tmp_path), kv_native_dtype="bf16") 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 pipeline["strategy"]["kind"] == "composite" + assert set(pipeline) == {"workflow"} + workflow = pipeline["workflow"] + assert workflow["manifest"]["adapter_abis"] == {"onnx-genai.image-preprocess": "1"} + assert workflow["graph"]["setup"]["nodes"][0]["component"] == "image_preprocess" + assert workflow["graph"]["setup"]["nodes"][1]["component"] == "vision_encoder" + assert workflow["graph"]["setup"]["nodes"][3]["component"] == "embedding" + assert workflow["graph"]["iteration"]["value"] == "loop.iteration" def test_dispatch_audio_only_multimodal_pipeline(tmp_path, monkeypatch): # 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(), - } - ) - audio_processor = tmp_path / "audio_processor.json" - audio_processor.write_text("{}") - calls: list[tuple[str | None, str | None]] = [] - - def fake_audio_processor(output_dir, source, *, revision=None): - calls.append((source, revision)) - return str(audio_processor) - - monkeypatch.setattr( - "mobius.integrations.onnx_genai.auto_export._write_hf_audio_processor", - fake_audio_processor, - ) - artifacts = write_onnx_genai_config( - pkg, - str(tmp_path), - source="zai-org/GLM-ASR-Nano-2512", - revision="pinned-revision", - ) - assert artifacts["audio_processor"] == str(audio_processor) - assert calls == [("zai-org/GLM-ASR-Nano-2512", "pinned-revision")] + 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["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", + setup = metadata["pipeline"]["workflow"]["graph"]["setup"]["nodes"] + 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_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: diff --git a/src/mobius/integrations/onnx_genai/inference_metadata.py b/src/mobius/integrations/onnx_genai/inference_metadata.py index 498d78d0c..68a3ec4f3 100644 --- a/src/mobius/integrations/onnx_genai/inference_metadata.py +++ b/src/mobius/integrations/onnx_genai/inference_metadata.py @@ -566,6 +566,7 @@ def _processor_values( "temporal_patch_size", "spatial_merge_size", "image_crop_size", + "size", ): value = getattr(vision, name, None) if value is None: diff --git a/src/mobius/integrations/onnx_genai/workflow_metadata.py b/src/mobius/integrations/onnx_genai/workflow_metadata.py index ddc94774a..4ee1570f6 100644 --- a/src/mobius/integrations/onnx_genai/workflow_metadata.py +++ b/src/mobius/integrations/onnx_genai/workflow_metadata.py @@ -28,6 +28,7 @@ _port, _shape_metadata, add_policy_components_to_workflow, + build_native_vlm_package_metadata, ) @@ -523,6 +524,553 @@ def write_diffusion_workflow_metadata( 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 _contract(value) == _contract(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 + + 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 ( + value.name.replace("past_key_values", "present"), + value.name.replace("past.", "present."), + ) + if name in decoder_outputs + ), + None, + ) + if present is not None: + cache_pairs.append((value, present)) + cache_names = {value.name for value, _ in cache_pairs} + rank2_integer = [ + value + for value in decoder.graph.inputs + if 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 = _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") + + 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") + 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["source"] = output["content"] + output["name"] = f"image.{port_name}" + adapter_outputs[port_name] = output["contract"] + preprocessing_values[port_name] = output["name"] + + 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()) + pkg.add_policy_component("continue_predicate", build_boolean_not()) + 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), + ), + ) + 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, + ), + ) + + batch = _contract(token_input)["shape"][0] + batch_int = {"dtype": "int64", "rank": 1, "shape": [batch]} + eos = getattr(config, "eos_token_id", 0) + if isinstance(eos, list): + eos = eos[0] if eos else 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": True, + }, + "request.max_iterations": { + "contract": batch_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": [1]}, + "role": {"kind": "opaque"}, + "source": {"kind": "literal"}, + "required": False, + "default": int(eos or 0), + }, + "package.max_context": { + "contract": batch_int, + "role": {"kind": "opaque"}, + "source": {"kind": "literal"}, + "required": False, + "default": int(getattr(config, "max_position_embeddings", 4096)), + }, + "package.one": { + "contract": batch_int, + "role": {"kind": "opaque"}, + "source": {"kind": "literal"}, + "required": False, + "default": 1, + }, + } + 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: "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" + 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"} + state: dict[str, Any] = { + "token": { + "contract": {"dtype": "int64", "rank": 2, "shape": [batch, 1]}, + "scope": "invocation", + "initializer": "initializer.token_slot", + "recurrence": {"kind": "invariant"}, + }, + "logits": { + "contract": _contract(logits_output), + "scope": "invocation", + "initializer": "decoder.setup.logits", + "recurrence": {"kind": "invariant"}, + }, + "attention_mask": { + "contract": { + "dtype": _contract(attention_input)["dtype"], + "rank": 2, + "shape": [batch, "context"], + }, + "scope": "invocation", + "initializer": "initializer.body_attention_mask", + "recurrence": { + "kind": "growing", + "axis": 1, + "increment": "package.one", + "max": "package.max_context", + }, + }, + } + state_specs = [ + ( + "token", + "initializer.token_slot", + "state.token.body", + "token.body", + "state.token.final", + ), + ( + "logits", + "decoder.setup.logits", + "state.logits.body", + "decoder.body.logits", + "state.logits.final", + ), + ( + "attention_mask", + "initializer.body_attention_mask", + "state.attention_mask.body", + "decoder_step.body_attention_mask", + "state.attention_mask.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}" + state[cell] = { + "contract": _contract(past), + "scope": "invocation", + "initializer": f"decoder.setup.{present.name}", + "recurrence": { + "kind": "growing", + "axis": next( + ( + axis + for axis, dimension in enumerate(_contract(past)["shape"]) + if "sequence" in str(dimension) + ), + 2, + ), + "increment": "package.one", + "max": "package.max_context", + }, + } + 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", + ) + ) + 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"), + } + ) + + setup = { + "kind": "sequence", + "nodes": [ + _invoke( + "image_preprocess", + {"encoded": "request.image"}, + dict(preprocessing_values), + ), + _invoke("vision_encoder", vision_invoke_inputs, vision_outputs), + *audio_setup_nodes, + _invoke( + "decoder_state_initializer", + {"prompt_tokens": "request.prompt_tokens"}, + { + attention_input.name: f"initializer.{attention_input.name}", + "body_attention_mask": "initializer.body_attention_mask", + "token_slot": "initializer.token_slot", + **( + { + 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( + "embedding", + embedding_setup_inputs, + {embedding_output.name: "embedding.setup.embeds"}, + ), + _invoke("decoder", setup_decoder_inputs, setup_decoder_outputs), + ], + } + body = { + "kind": "sequence", + "nodes": [ + _invoke( + "last_token_logits", + {"logits": "state.logits.body"}, + {"last_logits": "decoder.body.last_logits"}, + ), + _invoke( + "token_sampler", + {"logits": "decoder.body.last_logits"}, + {"token": "sample.body"}, + {"sample": _effect("sample.0", "sample.1")}, + ), + _invoke( + "termination", + { + "token_ids": "sample.body", + "eos_ids": "package.eos_ids", + "iteration": "loop.iteration", + "max_iterations": "request.max_iterations", + }, + {"done": "loop.done"}, + {"termination": _effect("termination.0", "termination.1")}, + ), + _invoke( + "continue_predicate", + {"done": "loop.done"}, + {"continue": "loop.continue"}, + ), + { + "kind": "emit", + "value": "sample.body", + "output": "tokens", + "mode": "append", + "effect_name": "emit", + "effect": _effect("emit.0", "emit.1"), + }, + _invoke( + "token_state_update", + {"current": "state.token.body", "update": "sample.body"}, + {"next": "token.body"}, + {"state": _effect("state.0", "state.1")}, + ), + _invoke( + "embedding", + embedding_body_inputs, + {embedding_output.name: "embedding.body.embeds"}, + ), + _invoke("decoder", body_decoder_inputs, body_decoder_outputs), + _invoke( + "decoder_step_update", + { + "attention_mask": "state.attention_mask.body", + **( + {"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 {} + ), + }, + ), + ], + } + 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": { + "ir_version": "1.0", + "onnx_opsets": {"ai.onnx": OPSET_VERSION}, + "adapter_abis": {"onnx-genai.image-preprocess": "1"}, + "capabilities": [ + "workflow_ssa", + "linear_effects", + "nested_control_flow", + "typed_emit", + ], + }, + "inputs": inputs, + "outputs": { + "tokens": {"contract": batch_int, "role": "tokens", "stage": "pre_adapter"} + }, + "components": components, + "state": state, + "initial_effects": initial_effects, + "graph": { + "kind": "loop", + "setup": setup, + "body": body, + "condition": "loop.continue", + "max_iterations": "request.max_iterations", + "iteration": {"value": "loop.iteration", "contract": batch_int}, + "carried": carried, + }, + } + metadata = { + "schema_version": "v1", + "preprocessing": preprocessing, + "pipeline": {"workflow": 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) + path = os.path.join(output_dir, "inference_metadata.yaml") + with open(path, "w", encoding="utf-8") as handle: + yaml.safe_dump(metadata, handle, sort_keys=False) + return path + + def build_decoder_workflow_metadata( pkg: Any, config: Any, From 62114b16da8b5f9f957674313e418b9fad5fbf8a Mon Sep 17 00:00:00 2001 From: justinchuby Date: Wed, 12 Aug 2026 20:36:48 +0000 Subject: [PATCH 011/151] Add nested TTS workflow generation Use lexical outer talker and inner code-group induction values, generated frame/history state components, greedy code sampling, and codec layout conversion before waveform emission. Replace the prior induction-contract blocker path. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: justinchuby --- src/mobius/generation/__init__.py | 8 + src/mobius/generation/_policy_components.py | 77 ++++ .../onnx_genai/auto_export_test.py | 46 +- .../codec_workflow_metadata_test.py | 71 ++- .../onnx_genai/workflow_metadata.py | 416 +++++++++++++++++- 5 files changed, 582 insertions(+), 36 deletions(-) diff --git a/src/mobius/generation/__init__.py b/src/mobius/generation/__init__.py index 4ed7e10d7..145221344 100644 --- a/src/mobius/generation/__init__.py +++ b/src/mobius/generation/__init__.py @@ -11,6 +11,9 @@ PolicyRole, attach_policy_components, build_boolean_not, + build_code_frame_update, + build_code_history_append, + build_codec_layout_transpose, build_decoder_state_initializer, build_decoder_step_update, build_eos_termination, @@ -25,6 +28,7 @@ build_seeded_categorical_sampler, build_speculative_acceptance, build_token_state_update, + build_tts_state_initializer, ) __all__ = [ @@ -33,6 +37,9 @@ "PolicyRole", "attach_policy_components", "build_boolean_not", + "build_code_frame_update", + "build_code_history_append", + "build_codec_layout_transpose", "build_decoder_state_initializer", "build_decoder_step_update", "build_eos_termination", @@ -47,4 +54,5 @@ "build_seeded_categorical_sampler", "build_speculative_acceptance", "build_token_state_update", + "build_tts_state_initializer", ] diff --git a/src/mobius/generation/_policy_components.py b/src/mobius/generation/_policy_components.py index dbbcf4a32..9747dfd17 100644 --- a/src/mobius/generation/_policy_components.py +++ b/src/mobius/generation/_policy_components.py @@ -215,6 +215,83 @@ def build_schedule_constant(values: list[float]) -> PolicyComponent: return _component(PolicyRole.AUXILIARY, 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)) + history = op.ConstantOfShape(history_shape, value=ir.tensor([0], dtype=ir.DataType.INT64)) + frame.shape = ir.Shape(["batch", num_code_groups]) + history.shape = ir.Shape(["batch", 0, num_code_groups]) + builder.add_output(frame, "frame_codes") + builder.add_output(history, "code_history") + return _component(PolicyRole.AUXILIARY, graph, {}) + + +def build_code_frame_update(num_code_groups: int) -> 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=["batch"]) + 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(PolicyRole.AUXILIARY, 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(PolicyRole.AUXILIARY, 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(PolicyRole.AUXILIARY, 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") diff --git a/src/mobius/integrations/onnx_genai/auto_export_test.py b/src/mobius/integrations/onnx_genai/auto_export_test.py index 3ff9319be..64b104698 100644 --- a/src/mobius/integrations/onnx_genai/auto_export_test.py +++ b/src/mobius/integrations/onnx_genai/auto_export_test.py @@ -553,19 +553,45 @@ class _TTSPkg(dict): def test_dispatch_multi_decoder_tts_with_pre_embedder(tmp_path): - 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(), ) - with pytest.raises(NotImplementedError, match="nested-loop induction SSA value"): - write_onnx_genai_config(pkg, str(tmp_path)) + artifacts = write_onnx_genai_config(pkg, str(tmp_path)) + with open(artifacts["inference_metadata"]) as handle: + workflow = yaml.safe_load(handle)["pipeline"]["workflow"] + outer = workflow["graph"]["nodes"][0] + assert outer["iteration"]["value"] == "talker.iteration" + assert outer["body"]["nodes"][2]["iteration"]["value"] == "code.iteration" def test_unrecognized_multi_component_package_fails_loudly(tmp_path): diff --git a/src/mobius/integrations/onnx_genai/codec_workflow_metadata_test.py b/src/mobius/integrations/onnx_genai/codec_workflow_metadata_test.py index 917080209..3818fc033 100644 --- a/src/mobius/integrations/onnx_genai/codec_workflow_metadata_test.py +++ b/src/mobius/integrations/onnx_genai/codec_workflow_metadata_test.py @@ -3,6 +3,7 @@ from __future__ import annotations +import dataclasses import json import os @@ -98,12 +99,64 @@ def test_codec_workflow_matches_producer_schema(): jsonschema.validate(build_audio_codec_workflow_metadata(_codec_package()), schema) -def test_tts_reports_missing_nested_loop_induction_value(): - package = { - "talker": object(), - "code_predictor": object(), - "talker_step_embedder": object(), - "talker_prefill_embedder": object(), - } - with pytest.raises(NotImplementedError, match=r"code_predictor\.step_index"): - build_tts_workflow_metadata(package, object()) +@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["graph"]["nodes"][0] + inner = outer["body"]["nodes"][2] + assert outer["iteration"]["value"] == "talker.iteration" + assert inner["iteration"]["value"] == "code.iteration" + assert inner["body"]["nodes"][0]["inputs"]["step_index"] == "code.iteration" + assert workflow["graph"]["nodes"][-2]["component"] == "codec" + assert workflow["outputs"]["waveform"]["stage"] == "post_adapter" diff --git a/src/mobius/integrations/onnx_genai/workflow_metadata.py b/src/mobius/integrations/onnx_genai/workflow_metadata.py index 4ee1570f6..f8e2c3b3f 100644 --- a/src/mobius/integrations/onnx_genai/workflow_metadata.py +++ b/src/mobius/integrations/onnx_genai/workflow_metadata.py @@ -16,6 +16,9 @@ PolicyCapabilities, attach_policy_components, build_boolean_not, + build_code_frame_update, + build_code_history_append, + build_codec_layout_transpose, build_decoder_state_initializer, build_decoder_step_update, build_integer_increment, @@ -23,6 +26,7 @@ build_last_token_logits, build_model_token_cast, build_schedule_constant, + build_tts_state_initializer, ) from mobius.integrations.onnx_genai.inference_metadata import ( _port, @@ -189,32 +193,410 @@ def write_audio_codec_workflow_metadata(pkg: Any, output_dir: str) -> str: def build_tts_workflow_metadata(pkg: Any, config: Any) -> dict[str, Any]: - """Reject TTS until the generic nested-loop contract exposes induction SSA. - - Qwen3-TTS needs the inner loop index as the code predictor ``step_index`` and - to select the next code embedding. The workflow ``loop`` node at producer - commit 4c3c4b6 only accepts a condition and maximum value; it defines no - iteration SSA value. Consequently the existing prefill and per-frame - embedder artifacts can be invoked, but the code-predictor loop cannot be - expressed without host preprocessing or a model-specific counter component. - """ + """Build nested talker/code-predictor loops with lexical induction SSA.""" 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}") - del config - raise NotImplementedError( - "generic TTS workflow requires a nested-loop induction SSA value: " - "ONNX GenAI workflow Loop at producer commit 4c3c4b6 exposes neither an " - "iteration output nor fixed-loop index, so code_predictor.step_index, " - "position_ids, and per-group code embedding selection cannot be wired " - "from the existing prefill/step embedder artifacts without host " - "preprocessing" + 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") + + 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)) + 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 = {"dtype": "int64", "rank": 1, "shape": [batch]} + batch_bool = {"dtype": "bool", "rank": 1, "shape": [batch]} + 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": batch_int, + "role": { + "kind": "runtime", + "version": "1.0", + "role": "max_output_tokens", + }, + "source": {"kind": "request", "field": "max_output_tokens"}, + "required": True, + }, + "package.code_groups": { + "contract": batch_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": batch_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 + + 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, + {predictor_logits.name: "code.logits"}, + ), + _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": batch_int}, + "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, + {step_output.name: "talker.step_embeds"}, + ), + _invoke( + "talker", + talker_body_inputs, + {talker_hidden.name: "talker.body.hidden"}, + ), + 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 = {prefill_output.name: "talker.prefill_embeds"} + if talker_hidden.name in {value.name for value in talker.graph.outputs}: + talker_setup_outputs = {talker_hidden.name: "talker.prefill.hidden"} + 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}, + {waveform_output.name: "tts.waveform"}, + ), + { + "kind": "emit", + "value": "tts.waveform", + "output": "waveform", + "mode": "replace", + "effect_name": "emit", + "effect": _effect("emit.0", "emit.1"), + }, + ] ) + workflow = { + "manifest": { + "ir_version": "1.0", + "onnx_opsets": {"ai.onnx": OPSET_VERSION}, + "capabilities": [ + "workflow_ssa", + "linear_effects", + "nested_control_flow", + "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": workflow}} + add_policy_components_to_workflow(metadata, pkg) + return metadata def write_tts_workflow_metadata(pkg: Any, output_dir: str, config: Any) -> str: From 3a40b0bad803cbf4be95ebabed15a0c2c2e2dfd0 Mon Sep 17 00:00:00 2001 From: justinchuby Date: Wed, 12 Aug 2026 20:40:59 +0000 Subject: [PATCH 012/151] Add speculative workflow branch joins Wire proposer and verifier components through the seeded acceptance policy, carry token, KV, and RNG state, and publish accepted/corrected token blocks through typed branch phi outputs and linear effect merges. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: justinchuby --- src/mobius/generation/__init__.py | 2 + src/mobius/generation/_policy_components.py | 25 ++ .../generation/_policy_components_test.py | 5 +- .../integrations/onnx_genai/__init__.py | 4 + .../integrations/onnx_genai/auto_export.py | 12 + .../onnx_genai/workflow_metadata.py | 355 ++++++++++++++++++ .../onnx_genai/workflow_metadata_test.py | 67 ++++ 7 files changed, 469 insertions(+), 1 deletion(-) diff --git a/src/mobius/generation/__init__.py b/src/mobius/generation/__init__.py index 145221344..518776ea8 100644 --- a/src/mobius/generation/__init__.py +++ b/src/mobius/generation/__init__.py @@ -27,6 +27,7 @@ build_schedule_constant, build_seeded_categorical_sampler, build_speculative_acceptance, + build_token_block_identity, build_token_state_update, build_tts_state_initializer, ) @@ -54,5 +55,6 @@ "build_seeded_categorical_sampler", "build_speculative_acceptance", "build_token_state_update", + "build_token_block_identity", "build_tts_state_initializer", ] diff --git a/src/mobius/generation/_policy_components.py b/src/mobius/generation/_policy_components.py index 9747dfd17..2d6ebb843 100644 --- a/src/mobius/generation/_policy_components.py +++ b/src/mobius/generation/_policy_components.py @@ -739,6 +739,8 @@ def build_speculative_acceptance() -> PolicyComponent: 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) @@ -756,11 +758,21 @@ def build_speculative_acceptance() -> PolicyComponent: value=ir.tensor([0], dtype=ir.DataType.INT64), ), ) + accepted_tokens.shape = ir.Shape(["batch", "draft_sequence"]) draft_length = op.Shape(proposed_tokens, start=1, end=2) done = op.Equal(accepted_count, draft_length) + 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"]) + 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") return _component( PolicyRole.SPECULATIVE_ACCEPTANCE, graph, @@ -771,12 +783,25 @@ def build_speculative_acceptance() -> PolicyComponent: "accepted_tokens": "accepted_tokens", "accepted_len": "accepted_len", "done": "done", + "rng": { + "seed": "seed", + "offset": "offset", + "next_offset": "next_offset", + }, "effect": "verify", }, "verify", ) +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(PolicyRole.AUXILIARY, graph, {}, "state") + + def build_token_state_update() -> PolicyComponent: """Build explicit token-history append and sequence-length update math.""" graph, builder = _make_graph("token_state_update") diff --git a/src/mobius/generation/_policy_components_test.py b/src/mobius/generation/_policy_components_test.py index 961c75ede..69a074d9e 100644 --- a/src/mobius/generation/_policy_components_test.py +++ b/src/mobius/generation/_policy_components_test.py @@ -309,7 +309,7 @@ def test_masked_update_runtime_parity(tmp_path): def test_speculative_acceptance_prefix_runtime(tmp_path): - accepted_tokens, count, done = _run( + accepted_tokens, count, done, next_offset = _run( build_speculative_acceptance(), tmp_path, { @@ -318,11 +318,14 @@ def test_speculative_acceptance_prefix_runtime(tmp_path): 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, 0, 0]]) np.testing.assert_array_equal(count, [2]) np.testing.assert_array_equal(done, [False]) + np.testing.assert_array_equal(next_offset, [12]) def test_token_state_update_runtime(tmp_path): diff --git a/src/mobius/integrations/onnx_genai/__init__.py b/src/mobius/integrations/onnx_genai/__init__.py index 001c06eac..7abcae6b1 100644 --- a/src/mobius/integrations/onnx_genai/__init__.py +++ b/src/mobius/integrations/onnx_genai/__init__.py @@ -65,12 +65,14 @@ build_decoder_workflow_metadata, build_diffusion_workflow_metadata, build_language_diffusion_pipeline_metadata, + build_speculative_workflow_metadata, build_tts_workflow_metadata, build_vlm_workflow_metadata, write_audio_codec_workflow_metadata, write_decoder_workflow_metadata, write_diffusion_workflow_metadata, write_language_diffusion_workflow_metadata, + write_speculative_workflow_metadata, write_tts_workflow_metadata, write_vlm_workflow_metadata, ) @@ -85,6 +87,7 @@ "build_diffusion_workflow_metadata", "build_diffusion_pipeline_metadata", "build_language_diffusion_pipeline_metadata", + "build_speculative_workflow_metadata", "build_audio_codec_workflow_metadata", "build_multimodal_pipeline_metadata", "build_pipeline_metadata_for_workflow", @@ -103,6 +106,7 @@ "write_decoder_workflow_metadata", "write_diffusion_workflow_metadata", "write_language_diffusion_workflow_metadata", + "write_speculative_workflow_metadata", "write_diffusion_pipeline_metadata", "write_audio_codec_workflow_metadata", "write_multimodal_pipeline_metadata", diff --git a/src/mobius/integrations/onnx_genai/auto_export.py b/src/mobius/integrations/onnx_genai/auto_export.py index 7a0971775..f03d19ee0 100644 --- a/src/mobius/integrations/onnx_genai/auto_export.py +++ b/src/mobius/integrations/onnx_genai/auto_export.py @@ -34,6 +34,7 @@ write_decoder_workflow_metadata, write_diffusion_workflow_metadata, write_language_diffusion_workflow_metadata, + write_speculative_workflow_metadata, write_tts_workflow_metadata, write_vlm_workflow_metadata, ) @@ -312,6 +313,13 @@ 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. @@ -484,6 +492,10 @@ def write_onnx_genai_config( path = write_audio_codec_workflow_metadata(pkg, output_dir) return {"inference_metadata": path} + if _looks_like_speculative(pkg): + path = write_speculative_workflow_metadata(pkg, output_dir) + return {"inference_metadata": path} + resolved_config = config if config is not None else getattr(pkg, "config", None) if resolved_config is None: raise ValueError( diff --git a/src/mobius/integrations/onnx_genai/workflow_metadata.py b/src/mobius/integrations/onnx_genai/workflow_metadata.py index f8e2c3b3f..5e76acddb 100644 --- a/src/mobius/integrations/onnx_genai/workflow_metadata.py +++ b/src/mobius/integrations/onnx_genai/workflow_metadata.py @@ -26,6 +26,7 @@ build_last_token_logits, build_model_token_cast, build_schedule_constant, + build_token_block_identity, build_tts_state_initializer, ) from mobius.integrations.onnx_genai.inference_metadata import ( @@ -1453,6 +1454,360 @@ def write_vlm_workflow_metadata( return path +def build_speculative_workflow_metadata(pkg: Any) -> dict[str, Any]: + """Build proposer/verifier workflow with branch phi and effect joins.""" + if not {"proposer", "verifier"} <= set(pkg.keys()): + raise ValueError("speculative workflow requires proposer and verifier") + proposer = pkg["proposer"] + verifier = pkg["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 + 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)) + pkg.add_policy_component("continue_predicate", build_boolean_not()) + pkg.add_policy_component("branch_state", build_token_block_identity()) + batch = _contract(proposer_input)["shape"][0] + batch_int = {"dtype": "int64", "rank": 1, "shape": [batch]} + batch_bool = {"dtype": "bool", "rank": 1, "shape": [batch]} + 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": batch_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.false": { + "contract": batch_bool, + "role": {"kind": "opaque"}, + "source": {"kind": "literal"}, + "required": False, + "default": False, + }, + } + + proposer_inputs = {proposer_input.name: "state.tokens.body"} + for value in proposer.graph.inputs: + if value is proposer_input: + 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" + branch = { + "kind": "branch", + "predicate": "acceptance.done", + "cases": { + "true": _invoke( + "branch_state", + {"tokens": "acceptance.tokens"}, + {"next_tokens": "branch.accepted"}, + {"state": _effect("branch.state.in", "branch.state.accepted")}, + ), + "false": _invoke( + "branch_state", + {"tokens": "proposal.tokens"}, + {"next_tokens": "branch.corrected"}, + {"state": _effect("branch.state.in", "branch.state.corrected")}, + ), + }, + "outputs": { + "tokens.next": { + "cases": { + "true": "branch.accepted", + "false": "branch.corrected", + } + } + }, + "effects": { + "state": { + "incoming": "branch.state.in", + "cases": { + "true": "branch.state.accepted", + "false": "branch.state.corrected", + }, + "produces": "branch.state.out", + } + }, + } + body_nodes = [ + _invoke("proposer", proposer_inputs, proposal_outputs), + _invoke("verifier", verifier_inputs, verifier_outputs), + _invoke( + "speculative_acceptance", + acceptance_inputs, + { + "accepted_tokens": "acceptance.tokens", + "accepted_len": "acceptance.length", + "done": "acceptance.done", + "next_offset": "rng_offset.body", + }, + {"verify": _effect("verify.0", "verify.1")}, + ), + branch, + _invoke( + "continue_predicate", + {"done": "package.false"}, + {"continue": "speculative.continue"}, + ), + { + "kind": "emit", + "value": "tokens.next", + "output": "tokens", + "mode": "append", + "effect_name": "emit", + "effect": _effect("emit.0", "emit.1"), + }, + ] + state = { + "tokens": { + "contract": _contract(proposer_input), + "scope": "invocation", + "initializer": "request.tokens", + "recurrence": {"kind": "invariant"}, + }, + "rng_offset": { + "contract": batch_int, + "scope": "invocation", + "initializer": "package.zero", + "recurrence": {"kind": "invariant"}, + }, + } + state_specs = [ + ( + "tokens", + "request.tokens", + "state.tokens.body", + "tokens.next", + "state.tokens.final", + ), + ( + "rng_offset", + "package.zero", + "state.rng_offset.body", + "rng_offset.body", + "state.rng_offset.final", + ), + ] + for index, (past, present) in enumerate(cache_pairs): + cell = f"cache_{index}" + initializer = f"request.verifier.{past.name}" + state[cell] = { + "contract": _contract(past), + "scope": "invocation", + "initializer": initializer, + "recurrence": { + "kind": "growing", + "axis": next( + ( + axis + for axis, dimension in enumerate(_contract(past)["shape"]) + if "sequence" in str(dimension) + ), + 2, + ), + "increment": "package.one", + "max": "request.max_iterations", + }, + } + state_specs.append( + ( + cell, + initializer, + f"state.{cell}.body", + f"verifier.{present.name}", + f"state.{cell}.final", + ) + ) + initial_effects = { + "verify": "verify.0", + "emit": "emit.0", + "state": "branch.state.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": { + "ir_version": "1.0", + "onnx_opsets": {"ai.onnx": OPSET_VERSION}, + "capabilities": [ + "workflow_ssa", + "linear_effects", + "nested_control_flow", + "typed_emit", + ], + }, + "inputs": inputs, + "outputs": { + "tokens": { + "contract": _contract(proposed_tokens), + "role": "tokens", + "stage": "pre_adapter", + } + }, + "components": { + name: _component(model, _artifact(name, len(pkg))) for name, model in pkg.items() + }, + "state": state, + "initial_effects": initial_effects, + "graph": { + "kind": "loop", + "setup": { + "kind": "sequence", + "nodes": [ + _invoke( + "continue_predicate", + {"done": "package.false"}, + {"continue": "speculative.setup.continue"}, + ) + ], + }, + "body": {"kind": "sequence", "nodes": body_nodes}, + "condition": "speculative.continue", + "max_iterations": "request.max_iterations", + "iteration": {"value": "speculative.iteration", "contract": batch_int}, + "carried": carried, + }, + } + metadata = {"schema_version": "v1", "pipeline": {"workflow": workflow}} + add_policy_components_to_workflow(metadata, pkg) + return metadata + + +def write_speculative_workflow_metadata(pkg: Any, output_dir: str) -> str: + os.makedirs(output_dir, exist_ok=True) + metadata = build_speculative_workflow_metadata(pkg) + path = os.path.join(output_dir, "inference_metadata.yaml") + with open(path, "w", encoding="utf-8") as handle: + yaml.safe_dump(metadata, handle, sort_keys=False) + return path + + def build_decoder_workflow_metadata( pkg: Any, config: Any, diff --git a/src/mobius/integrations/onnx_genai/workflow_metadata_test.py b/src/mobius/integrations/onnx_genai/workflow_metadata_test.py index c9bbb2c98..0ee2626f3 100644 --- a/src/mobius/integrations/onnx_genai/workflow_metadata_test.py +++ b/src/mobius/integrations/onnx_genai/workflow_metadata_test.py @@ -13,6 +13,7 @@ from mobius._model_package import ModelPackage from mobius.integrations.onnx_genai.workflow_metadata import ( build_language_diffusion_pipeline_metadata, + build_speculative_workflow_metadata, ) @@ -97,3 +98,69 @@ def test_language_diffusion_matches_pr_828_schema(): num_inference_steps=8, ) jsonschema.validate(instance=metadata, schema=schema) + + +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() -> ModelPackage: + proposer = _graph_model( + "proposer", + [_value("tokens", ir.DataType.INT64, ["batch", 4])], + [ + _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_speculative_workflow_uses_branch_phi_effect_join_and_rng(): + workflow = build_speculative_workflow_metadata(_speculative_package())["pipeline"][ + "workflow" + ] + body = workflow["graph"]["body"]["nodes"] + branch = body[3] + assert branch["kind"] == "branch" + assert branch["outputs"]["tokens.next"]["cases"] == { + "true": "branch.accepted", + "false": "branch.corrected", + } + assert branch["effects"]["state"]["produces"] == "branch.state.out" + acceptance = body[2] + assert acceptance["inputs"]["offset"] == "state.rng_offset.body" + assert acceptance["outputs"]["next_offset"] == "rng_offset.body" + assert any(item["cell"].startswith("cache_") for item in workflow["graph"]["carried"]) From 56e76f5fd43300aad0fd49ad482f3aa975137f84 Mon Sep 17 00:00:00 2001 From: justinchuby Date: Wed, 12 Aug 2026 20:47:50 +0000 Subject: [PATCH 013/151] Validate migrated workflows and Muse metadata Add structural packed-grid preprocessing for max-token processor configs, normalize symbolic VLM graph contracts, and declare loop-induction capabilities required by the North Star validator. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: justinchuby --- .../onnx_genai/inference_metadata.py | 36 +++++++++++++++++++ .../onnx_genai/inference_metadata_test.py | 30 ++++++++++++++++ .../onnx_genai/workflow_metadata.py | 30 +++++++++++++++- 3 files changed, 95 insertions(+), 1 deletion(-) diff --git a/src/mobius/integrations/onnx_genai/inference_metadata.py b/src/mobius/integrations/onnx_genai/inference_metadata.py index 68a3ec4f3..831d28f23 100644 --- a/src/mobius/integrations/onnx_genai/inference_metadata.py +++ b/src/mobius/integrations/onnx_genai/inference_metadata.py @@ -413,6 +413,41 @@ 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 +494,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, ) diff --git a/src/mobius/integrations/onnx_genai/inference_metadata_test.py b/src/mobius/integrations/onnx_genai/inference_metadata_test.py index 357a5b05d..64c4ef7a7 100644 --- a/src/mobius/integrations/onnx_genai/inference_metadata_test.py +++ b/src/mobius/integrations/onnx_genai/inference_metadata_test.py @@ -26,6 +26,7 @@ SchedulerConfig, _decoder_io, _input_source_map, + _match_max_token_grid, _port, add_explicit_package_io, add_policy_components_to_workflow, @@ -40,6 +41,35 @@ ) +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()) diff --git a/src/mobius/integrations/onnx_genai/workflow_metadata.py b/src/mobius/integrations/onnx_genai/workflow_metadata.py index 5e76acddb..988ce7c02 100644 --- a/src/mobius/integrations/onnx_genai/workflow_metadata.py +++ b/src/mobius/integrations/onnx_genai/workflow_metadata.py @@ -553,6 +553,7 @@ def bind_remaining( "workflow_ssa", "linear_effects", "nested_control_flow", + "loop_induction_values", "typed_emit", ], }, @@ -621,6 +622,22 @@ def _find_port(values: Any, *fragments: str) -> ir.Value | 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 build_diffusion_workflow_metadata( pkg: Any, *, @@ -819,6 +836,7 @@ def build_diffusion_workflow_metadata( "workflow_ssa", "linear_effects", "nested_control_flow", + "loop_induction_values", "typed_emit", ], }, @@ -945,7 +963,7 @@ def build_vlm_workflow_metadata( value for value in decoder.graph.inputs if embedding_output is not None - and _contract(value) == _contract(embedding_output) + and _contracts_compatible(value, embedding_output) ), None, ) @@ -959,6 +977,12 @@ def build_vlm_workflow_metadata( 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]] = [] @@ -975,6 +999,8 @@ def build_vlm_workflow_metadata( 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} rank2_integer = [ @@ -1410,6 +1436,7 @@ def build_vlm_workflow_metadata( "workflow_ssa", "linear_effects", "nested_control_flow", + "loop_induction_values", "typed_emit", ], }, @@ -1759,6 +1786,7 @@ def build_speculative_workflow_metadata(pkg: Any) -> dict[str, Any]: "workflow_ssa", "linear_effects", "nested_control_flow", + "loop_induction_values", "typed_emit", ], }, From 69153d56b3ac2d7f89d40f9d7424246f8778cf7f Mon Sep 17 00:00:00 2001 From: justinchuby Date: Wed, 12 Aug 2026 21:26:04 +0000 Subject: [PATCH 014/151] Fix workflow policy artifact and solver execution Persist generated policies for every composite workflow, materialize Euler timesteps and sigma scaling with model-native dtypes, correct speculative verifier substitutions, and support scalar TTS loop induction while binding complete component outputs. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: justinchuby --- src/mobius/generation/__init__.py | 4 + src/mobius/generation/_policy_components.py | 71 +++++++-- .../generation/_policy_components_test.py | 32 ++++- .../integrations/onnx_genai/auto_export.py | 64 ++++++++- .../onnx_genai/auto_export_test.py | 4 + .../onnx_genai/workflow_metadata.py | 136 ++++++++++++++---- .../onnx_genai/workflow_metadata_test.py | 7 + 7 files changed, 274 insertions(+), 44 deletions(-) diff --git a/src/mobius/generation/__init__.py b/src/mobius/generation/__init__.py index 518776ea8..51f107871 100644 --- a/src/mobius/generation/__init__.py +++ b/src/mobius/generation/__init__.py @@ -17,6 +17,7 @@ build_decoder_state_initializer, build_decoder_step_update, build_eos_termination, + build_euler_model_input, build_euler_solver_step, build_greedy_sampler, build_integer_increment, @@ -25,6 +26,7 @@ build_masked_token_update, build_model_token_cast, build_schedule_constant, + build_schedule_lookup, build_seeded_categorical_sampler, build_speculative_acceptance, build_token_block_identity, @@ -44,6 +46,7 @@ "build_decoder_state_initializer", "build_decoder_step_update", "build_eos_termination", + "build_euler_model_input", "build_euler_solver_step", "build_greedy_sampler", "build_integer_increment", @@ -52,6 +55,7 @@ "build_masked_token_update", "build_model_token_cast", "build_schedule_constant", + "build_schedule_lookup", "build_seeded_categorical_sampler", "build_speculative_acceptance", "build_token_state_update", diff --git a/src/mobius/generation/_policy_components.py b/src/mobius/generation/_policy_components.py index 2d6ebb843..282338a17 100644 --- a/src/mobius/generation/_policy_components.py +++ b/src/mobius/generation/_policy_components.py @@ -215,6 +215,18 @@ def build_schedule_constant(values: list[float]) -> PolicyComponent: return _component(PolicyRole.AUXILIARY, 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(PolicyRole.AUXILIARY, 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: @@ -236,7 +248,9 @@ def build_tts_state_initializer(num_code_groups: int) -> PolicyComponent: return _component(PolicyRole.AUXILIARY, graph, {}) -def build_code_frame_update(num_code_groups: int) -> PolicyComponent: +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 @@ -246,7 +260,13 @@ def build_code_frame_update(num_code_groups: int) -> PolicyComponent: shape=["batch", num_code_groups], ) token = builder.input("token", dtype=ir.DataType.INT64, shape=["batch"]) - index = builder.input("index", 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])), @@ -577,18 +597,39 @@ def build_eos_termination() -> PolicyComponent: ) -def build_euler_solver_step() -> PolicyComponent: +def build_euler_model_input(dtype: ir.DataType = ir.DataType.FLOAT) -> PolicyComponent: + """Scale a latent for the Euler denoiser input at the current sigma.""" + graph, builder = _make_graph("euler_model_input") + op = builder.op + sample = builder.input( + "sample", dtype, ["batch", "channels", "height", "width"] + ) + 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=[1, 2, 3])) + model_input = op.Div(sample, scale) + model_input.shape = sample.shape + builder.add_output(model_input, "model_input") + return _component(PolicyRole.AUXILIARY, graph, {}) + + +def build_euler_solver_step( + dtype: ir.DataType = ir.DataType.FLOAT, +) -> PolicyComponent: """Build the generic Euler update ``x_next = x + dx * (sigma_next-sigma)``.""" graph, builder = _make_graph("euler_solver_step") op = builder.op sample = builder.input( "sample", - ir.DataType.FLOAT, + dtype, ["batch", "channels", "height", "width"], ) derivative = builder.input( "derivative", - ir.DataType.FLOAT, + dtype, ["batch", "channels", "height", "width"], ) step = builder.input("step", ir.DataType.INT64, ["batch"]) @@ -598,6 +639,7 @@ def build_euler_solver_step() -> PolicyComponent: 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=[1, 2, 3])) next_sample = op.Add(sample, op.Mul(derivative, delta)) builder.add_output(next_sample, "next_state") @@ -750,17 +792,28 @@ def build_speculative_acceptance() -> PolicyComponent: to=ir.DataType.INT64, ) accepted_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.ConstantOfShape( - op.Shape(proposed_tokens), - value=ir.tensor([0], dtype=ir.DataType.INT64), - ), + 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(accepted_count, draft_length) + accepted_count = op.Min( + op.Add(accepted_count, op.Cast(op.Not(done), to=ir.DataType.INT64)), + draft_length, + ) next_offset = op.Add( offset, op.Squeeze(draft_length, op.Constant(value_ints=[0])), diff --git a/src/mobius/generation/_policy_components_test.py b/src/mobius/generation/_policy_components_test.py index 69a074d9e..f49774bf7 100644 --- a/src/mobius/generation/_policy_components_test.py +++ b/src/mobius/generation/_policy_components_test.py @@ -13,9 +13,11 @@ PolicyRole, attach_policy_components, build_boolean_not, + build_code_frame_update, build_decoder_state_initializer, build_decoder_step_update, build_eos_termination, + build_euler_model_input, build_euler_solver_step, build_greedy_sampler, build_last_token_logits, @@ -322,12 +324,38 @@ def test_speculative_acceptance_prefix_runtime(tmp_path): "offset": np.array([8], np.int64), }, ) - np.testing.assert_array_equal(accepted_tokens, [[1, 0, 0, 0]]) - np.testing.assert_array_equal(count, [2]) + 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]) +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_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(), diff --git a/src/mobius/integrations/onnx_genai/auto_export.py b/src/mobius/integrations/onnx_genai/auto_export.py index f03d19ee0..961281912 100644 --- a/src/mobius/integrations/onnx_genai/auto_export.py +++ b/src/mobius/integrations/onnx_genai/auto_export.py @@ -16,6 +16,7 @@ import os from typing import Any +import numpy as np import onnx_ir as ir import yaml @@ -41,6 +42,54 @@ _LOGGER = logging.getLogger(__name__) + +def _euler_schedule( + scheduler: SchedulerConfig, num_inference_steps: int +) -> tuple[list[float], list[float]]: + """Materialize diffusers-compatible Euler timesteps and sigma values.""" + if scheduler.kind != "euler" or scheduler.prediction_type != "epsilon": + raise ValueError( + "workflow diffusion currently supports deterministic Euler epsilon " + f"schedulers, got kind={scheduler.kind!r}, " + f"prediction_type={scheduler.prediction_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 " + f"{scheduler.beta_schedule!r}" + ) + training_sigmas = np.sqrt((1.0 - np.cumprod(1.0 - betas)) / np.cumprod(1.0 - betas)) + 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") @@ -468,14 +517,21 @@ 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 + if guidance_scale is not None and not np.isclose(guidance_scale, 1.0): + raise ValueError( + "workflow diffusion requires an explicit classifier-free guidance " + "component before guidance_scale can differ from 1.0" + ) + resolved_scheduler = scheduler or SchedulerConfig(kind="euler") + timesteps, sigma_schedule = _euler_schedule( + resolved_scheduler, num_inference_steps + ) path = write_diffusion_workflow_metadata( pkg, output_dir, num_inference_steps=num_inference_steps, + schedule=sigma_schedule, + timesteps=timesteps, ) artifacts = {"inference_metadata": path} # Emit the CLIP tokenizer.json for text-conditioned pipelines so the diff --git a/src/mobius/integrations/onnx_genai/auto_export_test.py b/src/mobius/integrations/onnx_genai/auto_export_test.py index 64b104698..d4a79c43d 100644 --- a/src/mobius/integrations/onnx_genai/auto_export_test.py +++ b/src/mobius/integrations/onnx_genai/auto_export_test.py @@ -217,6 +217,8 @@ def test_dispatch_diffusion(tmp_path): assert workflow["graph"]["nodes"][0]["iteration"]["value"] == "loop.iteration" assert workflow["graph"]["nodes"][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): @@ -369,6 +371,7 @@ def test_dispatch_vision_multimodal_pipeline(tmp_path): assert workflow["graph"]["setup"]["nodes"][1]["component"] == "vision_encoder" assert workflow["graph"]["setup"]["nodes"][3]["component"] == "embedding" assert workflow["graph"]["iteration"]["value"] == "loop.iteration" + assert (tmp_path / "policies" / "token_sampler.onnx").is_file() def test_dispatch_audio_only_multimodal_pipeline(tmp_path, monkeypatch): @@ -592,6 +595,7 @@ def test_dispatch_multi_decoder_tts_with_pre_embedder(tmp_path): outer = workflow["graph"]["nodes"][0] assert outer["iteration"]["value"] == "talker.iteration" assert outer["body"]["nodes"][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): diff --git a/src/mobius/integrations/onnx_genai/workflow_metadata.py b/src/mobius/integrations/onnx_genai/workflow_metadata.py index 988ce7c02..d180c6b87 100644 --- a/src/mobius/integrations/onnx_genai/workflow_metadata.py +++ b/src/mobius/integrations/onnx_genai/workflow_metadata.py @@ -21,11 +21,13 @@ build_codec_layout_transpose, build_decoder_state_initializer, build_decoder_step_update, + build_euler_model_input, + build_euler_solver_step, build_integer_increment, - build_iteration_cast, build_last_token_logits, build_model_token_cast, build_schedule_constant, + build_schedule_lookup, build_token_block_identity, build_tts_state_initializer, ) @@ -250,6 +252,13 @@ def build_tts_workflow_metadata(pkg: Any, config: Any) -> dict[str, Any]: 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), @@ -264,7 +273,10 @@ def build_tts_workflow_metadata(pkg: Any, config: Any) -> dict[str, Any]: 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)) + 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 @@ -337,6 +349,14 @@ def bind_remaining( 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, @@ -401,7 +421,11 @@ def bind_remaining( _invoke( "code_predictor", predictor_inputs, - {predictor_logits.name: "code.logits"}, + bind_outputs( + predictor.graph.outputs, + {predictor_logits.name: "code.logits"}, + "code_predictor.body", + ), ), _invoke( "token_sampler", @@ -427,7 +451,10 @@ def bind_remaining( }, "condition": "code.continue", "max_iterations": "package.code_groups", - "iteration": {"value": "code.iteration", "contract": batch_int}, + "iteration": { + "value": "code.iteration", + "contract": predictor_step_contract, + }, "carried": [ { "cell": "frame", @@ -446,12 +473,20 @@ def bind_remaining( _invoke( "talker_step_embedder", step_inputs, - {step_output.name: "talker.step_embeds"}, + bind_outputs( + step_embedder.graph.outputs, + {step_output.name: "talker.step_embeds"}, + "talker_step_embedder.body", + ), ), _invoke( "talker", talker_body_inputs, - {talker_hidden.name: "talker.body.hidden"}, + bind_outputs( + talker.graph.outputs, + {talker_hidden.name: "talker.body.hidden"}, + "talker.body", + ), ), inner_loop, _invoke( @@ -469,9 +504,17 @@ def bind_remaining( ), ], } - setup_outputs = {prefill_output.name: "talker.prefill_embeds"} + 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 = {talker_hidden.name: "talker.prefill.hidden"} + talker_setup_outputs = bind_outputs( + talker.graph.outputs, + {talker_hidden.name: "talker.prefill.hidden"}, + "talker.setup", + ) else: talker_setup_outputs = {} outer_loop = { @@ -533,7 +576,11 @@ def bind_remaining( _invoke( codec_name, {codec_input.name: codec_value}, - {waveform_output.name: "tts.waveform"}, + bind_outputs( + codec.graph.outputs, + {waveform_output.name: "tts.waveform"}, + "codec.final", + ), ), { "kind": "emit", @@ -602,9 +649,10 @@ def bind_remaining( def write_tts_workflow_metadata(pkg: Any, output_dir: str, config: Any) -> str: - """Build TTS workflow metadata, failing precisely on the producer defect.""" + """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) path = os.path.join(output_dir, "inference_metadata.yaml") with open(path, "w", encoding="utf-8") as handle: yaml.safe_dump(metadata, handle, sort_keys=False) @@ -643,6 +691,7 @@ def build_diffusion_workflow_metadata( *, num_inference_steps: int, schedule: list[float] | None = None, + timesteps: list[float] | None = None, ) -> dict[str, Any]: """Build a fixed-schedule diffusion workflow with explicit latent state.""" if num_inference_steps < 1: @@ -707,14 +756,23 @@ def build_diffusion_workflow_metadata( next(iter(text_encoder.graph.outputs), None), ) - attach_policy_components(pkg, PolicyCapabilities(solver="euler")) + attach_policy_components(pkg, PolicyCapabilities()) + pkg.add_policy_component( + "euler_model_input", build_euler_model_input(sample_input.dtype) + ) + pkg.add_policy_component("solver_step", build_euler_solver_step(sample_input.dtype)) pkg.add_policy_component("continue_predicate", build_boolean_not()) 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)) - if timestep_input.dtype != ir.DataType.INT64: - pkg.add_policy_component("iteration_cast", build_iteration_cast(timestep_input.dtype)) + pkg.add_policy_component("diffusion_timesteps", build_schedule_constant(timestep_values)) + pkg.add_policy_component("schedule_lookup", build_schedule_lookup(timestep_input.dtype)) batch = _contract(sample_input)["shape"][0] batch_int = {"dtype": "int64", "rank": 1, "shape": [batch]} @@ -742,7 +800,8 @@ def build_diffusion_workflow_metadata( }, } setup_nodes: list[dict[str, Any]] = [ - _invoke("diffusion_schedule", {}, {"schedule": "diffusion.schedule"}) + _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: @@ -784,24 +843,33 @@ def build_diffusion_workflow_metadata( ) denoiser_inputs = { - sample_input.name: "state.latent.body", - timestep_input.name: ( - "diffusion.timestep" - if timestep_input.dtype != ir.DataType.INT64 - else "loop.iteration" - ), + 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]] = [] - if timestep_input.dtype != ir.DataType.INT64: - body_nodes.append( - _invoke( - "iteration_cast", - {"iteration": "loop.iteration"}, - {"timestep": "diffusion.timestep"}, - ) + body_nodes.append( + _invoke( + "schedule_lookup", + { + "schedule": "diffusion.timesteps", + "step": "loop.iteration", + }, + {"timestep": "diffusion.timestep"}, + ) + ) + body_nodes.append( + _invoke( + "euler_model_input", + { + "sample": "state.latent.body", + "step": "loop.iteration", + "schedule": "diffusion.schedule", + }, + {"model_input": "diffusion.model_input"}, ) + ) body_nodes.extend( [ _invoke( @@ -916,9 +984,17 @@ def write_diffusion_workflow_metadata( output_dir: str, *, num_inference_steps: int, + schedule: list[float] | None = None, + timesteps: list[float] | None = None, ) -> str: os.makedirs(output_dir, exist_ok=True) - metadata = build_diffusion_workflow_metadata(pkg, num_inference_steps=num_inference_steps) + metadata = build_diffusion_workflow_metadata( + pkg, + num_inference_steps=num_inference_steps, + schedule=schedule, + timesteps=timesteps, + ) + pkg.save_policy_components(output_dir) path = os.path.join(output_dir, "inference_metadata.yaml") with open(path, "w", encoding="utf-8") as handle: yaml.safe_dump(metadata, handle, sort_keys=False) @@ -1475,6 +1551,7 @@ def write_vlm_workflow_metadata( ) -> str: os.makedirs(output_dir, exist_ok=True) metadata = build_vlm_workflow_metadata(pkg, config, source=source) + pkg.save_policy_components(output_dir) path = os.path.join(output_dir, "inference_metadata.yaml") with open(path, "w", encoding="utf-8") as handle: yaml.safe_dump(metadata, handle, sort_keys=False) @@ -1645,7 +1722,7 @@ def build_speculative_workflow_metadata(pkg: Any) -> dict[str, Any]: ), "false": _invoke( "branch_state", - {"tokens": "proposal.tokens"}, + {"tokens": "acceptance.tokens"}, {"next_tokens": "branch.corrected"}, {"state": _effect("branch.state.in", "branch.state.corrected")}, ), @@ -1830,6 +1907,7 @@ def build_speculative_workflow_metadata(pkg: Any) -> dict[str, Any]: def write_speculative_workflow_metadata(pkg: Any, output_dir: str) -> str: os.makedirs(output_dir, exist_ok=True) metadata = build_speculative_workflow_metadata(pkg) + pkg.save_policy_components(output_dir) path = os.path.join(output_dir, "inference_metadata.yaml") with open(path, "w", encoding="utf-8") as handle: yaml.safe_dump(metadata, handle, sort_keys=False) diff --git a/src/mobius/integrations/onnx_genai/workflow_metadata_test.py b/src/mobius/integrations/onnx_genai/workflow_metadata_test.py index 0ee2626f3..7b2d9e5b6 100644 --- a/src/mobius/integrations/onnx_genai/workflow_metadata_test.py +++ b/src/mobius/integrations/onnx_genai/workflow_metadata_test.py @@ -14,9 +14,16 @@ from mobius.integrations.onnx_genai.workflow_metadata import ( build_language_diffusion_pipeline_metadata, build_speculative_workflow_metadata, + write_speculative_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" / "branch_state.onnx").is_file() + + 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)) From b69bbfaac697a907eb022f11657d0568e4d6ab9d Mon Sep 17 00:00:00 2001 From: justinchuby Date: Wed, 12 Aug 2026 21:30:36 +0000 Subject: [PATCH 015/151] Carry normalized decoder logits through workflows Extract last-token logits after setup and body decoder invokes so the invariant loop cell remains rank two while full decoder outputs may change from prompt length to one token. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: justinchuby --- .../onnx_genai/auto_export_test.py | 13 ++++ .../onnx_genai/workflow_metadata.py | 62 +++++++++++++------ 2 files changed, 55 insertions(+), 20 deletions(-) diff --git a/src/mobius/integrations/onnx_genai/auto_export_test.py b/src/mobius/integrations/onnx_genai/auto_export_test.py index d4a79c43d..9683a34ba 100644 --- a/src/mobius/integrations/onnx_genai/auto_export_test.py +++ b/src/mobius/integrations/onnx_genai/auto_export_test.py @@ -167,12 +167,19 @@ def test_dispatch_decoder(tmp_path): assert [node["component"] for node in workflow["graph"]["setup"]["nodes"]] == [ "decoder_state_initializer", "model", + "last_token_logits", ] body = workflow["graph"]["body"]["nodes"] assert [node["kind"] for node in body].count("emit") == 1 assert next(node for node in body if node["kind"] == "emit")["value"] == "sample.body" assert workflow["state"]["iteration"]["initializer"] == "package.zero_iteration" assert workflow["state"]["token"]["initializer"] == "initializer.token_slot" + assert workflow["state"]["logits"] == { + "contract": {"dtype": "float32", "rank": 2, "shape": ["batch", 128]}, + "scope": "invocation", + "initializer": "decoder.setup.last_logits", + "recurrence": {"kind": "invariant"}, + } assert (tmp_path / "policies" / "token_sampler.onnx").is_file() @@ -371,6 +378,12 @@ def test_dispatch_vision_multimodal_pipeline(tmp_path): assert workflow["graph"]["setup"]["nodes"][1]["component"] == "vision_encoder" assert workflow["graph"]["setup"]["nodes"][3]["component"] == "embedding" assert workflow["graph"]["iteration"]["value"] == "loop.iteration" + assert workflow["state"]["logits"]["contract"] == { + "dtype": "float32", + "rank": 2, + "shape": ["batch", 128], + } + assert workflow["state"]["logits"]["initializer"] == "decoder.setup.last_logits" assert (tmp_path / "policies" / "token_sampler.onnx").is_file() diff --git a/src/mobius/integrations/onnx_genai/workflow_metadata.py b/src/mobius/integrations/onnx_genai/workflow_metadata.py index d180c6b87..c27fa8faf 100644 --- a/src/mobius/integrations/onnx_genai/workflow_metadata.py +++ b/src/mobius/integrations/onnx_genai/workflow_metadata.py @@ -1251,6 +1251,12 @@ def build_vlm_workflow_metadata( 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": logits_contract["dtype"], + "rank": 2, + "shape": [logits_contract["shape"][0], logits_contract["shape"][-1]], + } state: dict[str, Any] = { "token": { "contract": {"dtype": "int64", "rank": 2, "shape": [batch, 1]}, @@ -1259,9 +1265,9 @@ def build_vlm_workflow_metadata( "recurrence": {"kind": "invariant"}, }, "logits": { - "contract": _contract(logits_output), + "contract": last_logits_contract, "scope": "invocation", - "initializer": "decoder.setup.logits", + "initializer": "decoder.setup.last_logits", "recurrence": {"kind": "invariant"}, }, "attention_mask": { @@ -1290,9 +1296,9 @@ def build_vlm_workflow_metadata( ), ( "logits", - "decoder.setup.logits", + "decoder.setup.last_logits", "state.logits.body", - "decoder.body.logits", + "decoder.body.last_logits", "state.logits.final", ), ( @@ -1410,19 +1416,19 @@ def build_vlm_workflow_metadata( {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"}, + ), ], } body = { "kind": "sequence", "nodes": [ - _invoke( - "last_token_logits", - {"logits": "state.logits.body"}, - {"last_logits": "decoder.body.last_logits"}, - ), _invoke( "token_sampler", - {"logits": "decoder.body.last_logits"}, + {"logits": "state.logits.body"}, {"token": "sample.body"}, {"sample": _effect("sample.0", "sample.1")}, ), @@ -1462,6 +1468,11 @@ def build_vlm_workflow_metadata( {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"}, + ), _invoke( "decoder_step_update", { @@ -2146,6 +2157,12 @@ def build_decoder_workflow_metadata( 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": logits_contract["dtype"], + "rank": 2, + "shape": [logits_contract["shape"][0], logits_contract["shape"][-1]], + } state: dict[str, Any] = { "token": { "contract": { @@ -2164,9 +2181,9 @@ def build_decoder_workflow_metadata( "recurrence": {"kind": "invariant"}, }, "logits": { - "contract": _contract(logits_output), + "contract": last_logits_contract, "scope": "invocation", - "initializer": "decoder.setup.logits", + "initializer": "decoder.setup.last_logits", "recurrence": {"kind": "invariant"}, }, } @@ -2203,9 +2220,9 @@ def build_decoder_workflow_metadata( }, { "cell": "logits", - "current": "decoder.setup.logits", + "current": "decoder.setup.last_logits", "body_input": "state.logits.body", - "body_output": "decoder.body.logits", + "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"), @@ -2320,19 +2337,19 @@ def build_decoder_workflow_metadata( }, ), _invoke(decoder_name, setup_decoder_inputs, setup_decoder_outputs), + _invoke( + "last_token_logits", + {"logits": "decoder.setup.logits"}, + {"last_logits": "decoder.setup.last_logits"}, + ), ], } body = { "kind": "sequence", "nodes": [ - _invoke( - "last_token_logits", - {"logits": "state.logits.body"}, - {"last_logits": "decoder.body.last_logits"}, - ), _invoke( "token_sampler", - {"logits": "decoder.body.last_logits"}, + {"logits": "state.logits.body"}, {"token": "sample.body"}, {"sample": _effect("sample.0", "sample.1")}, ), @@ -2383,6 +2400,11 @@ def build_decoder_workflow_metadata( {"next_value": "iteration.body"}, ), _invoke(decoder_name, body_decoder_inputs, body_decoder_outputs), + _invoke( + "last_token_logits", + {"logits": "decoder.body.logits"}, + {"last_logits": "decoder.body.last_logits"}, + ), _invoke( "decoder_step_update", { From 8a77aeeed85c8ab4574838f7aea528479f062518 Mon Sep 17 00:00:00 2001 From: justinchuby Date: Wed, 12 Aug 2026 21:42:12 +0000 Subject: [PATCH 016/151] Add speculative accepted-state rollback policies Synchronize batched accepted prefixes, generate typed ONNX truncation components for tentative recurrent state, and join accepted versus corrected state values and effects through workflow branch phi outputs. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: justinchuby --- src/mobius/generation/__init__.py | 4 + src/mobius/generation/_policy_components.py | 69 +++++++ .../generation/_policy_components_test.py | 43 +++++ .../onnx_genai/workflow_metadata.py | 171 ++++++++++++++---- .../onnx_genai/workflow_metadata_test.py | 11 +- 5 files changed, 265 insertions(+), 33 deletions(-) diff --git a/src/mobius/generation/__init__.py b/src/mobius/generation/__init__.py index 51f107871..bd109c5dc 100644 --- a/src/mobius/generation/__init__.py +++ b/src/mobius/generation/__init__.py @@ -16,6 +16,7 @@ build_codec_layout_transpose, build_decoder_state_initializer, build_decoder_step_update, + build_effectful_identity, build_eos_termination, build_euler_model_input, build_euler_solver_step, @@ -29,6 +30,7 @@ build_schedule_lookup, build_seeded_categorical_sampler, build_speculative_acceptance, + build_speculative_state_rollback, build_token_block_identity, build_token_state_update, build_tts_state_initializer, @@ -48,6 +50,7 @@ "build_eos_termination", "build_euler_model_input", "build_euler_solver_step", + "build_effectful_identity", "build_greedy_sampler", "build_integer_increment", "build_iteration_cast", @@ -58,6 +61,7 @@ "build_schedule_lookup", "build_seeded_categorical_sampler", "build_speculative_acceptance", + "build_speculative_state_rollback", "build_token_state_update", "build_token_block_identity", "build_tts_state_initializer", diff --git a/src/mobius/generation/_policy_components.py b/src/mobius/generation/_policy_components.py index 282338a17..2b8a3662c 100644 --- a/src/mobius/generation/_policy_components.py +++ b/src/mobius/generation/_policy_components.py @@ -814,6 +814,27 @@ def build_speculative_acceptance() -> PolicyComponent: op.Add(accepted_count, op.Cast(op.Not(done), to=ir.DataType.INT64)), draft_length, ) + # Dense batched state has one physical sequence length. Synchronize to the + # shortest verified prefix so every row can share the same rollback point. + synchronized_len = op.ReduceMin(accepted_count, axes=[0], keepdims=1) + accepted_count = op.Expand(synchronized_len, op.Shape(accepted_count)) + synchronized_done = op.ReduceMin( + op.Cast(done, to=ir.DataType.INT64), axes=[0], keepdims=1 + ) + done = op.Expand( + op.Cast(synchronized_done, to=ir.DataType.BOOL), + op.Shape(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) next_offset = op.Add( offset, op.Squeeze(draft_length, op.Constant(value_ints=[0])), @@ -847,6 +868,54 @@ def build_speculative_acceptance() -> PolicyComponent: ) +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, ["batch"]) + past_len = op.Shape(past, start=sequence_axis, end=sequence_axis + 1) + synchronized_len = op.ReduceMin(accepted_len, axes=[0], keepdims=1) + end = op.Add(past_len, synchronized_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(PolicyRole.AUXILIARY, 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(PolicyRole.AUXILIARY, 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") diff --git a/src/mobius/generation/_policy_components_test.py b/src/mobius/generation/_policy_components_test.py index f49774bf7..d2c2d1cb3 100644 --- a/src/mobius/generation/_policy_components_test.py +++ b/src/mobius/generation/_policy_components_test.py @@ -25,6 +25,7 @@ build_model_token_cast, build_seeded_categorical_sampler, build_speculative_acceptance, + build_speculative_state_rollback, build_token_state_update, ) from mobius.generation._policy_components import _make_graph @@ -330,6 +331,48 @@ def test_speculative_acceptance_prefix_runtime(tmp_path): np.testing.assert_array_equal(next_offset, [12]) +def test_speculative_acceptance_synchronizes_batched_prefixes(tmp_path): + accepted_tokens, count, done, _ = _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, 2]) + np.testing.assert_array_equal(accepted_tokens, [[1, 0, 0, 0], [1, 0, 0, 0]]) + np.testing.assert_array_equal(done, [False, 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(), diff --git a/src/mobius/integrations/onnx_genai/workflow_metadata.py b/src/mobius/integrations/onnx_genai/workflow_metadata.py index c27fa8faf..41812cb23 100644 --- a/src/mobius/integrations/onnx_genai/workflow_metadata.py +++ b/src/mobius/integrations/onnx_genai/workflow_metadata.py @@ -21,6 +21,7 @@ build_codec_layout_transpose, build_decoder_state_initializer, build_decoder_step_update, + build_effectful_identity, build_euler_model_input, build_euler_solver_step, build_integer_increment, @@ -28,6 +29,7 @@ build_model_token_cast, build_schedule_constant, build_schedule_lookup, + build_speculative_state_rollback, build_token_block_identity, build_tts_state_initializer, ) @@ -1721,41 +1723,142 @@ def build_speculative_workflow_metadata(pkg: Any) -> dict[str, Any]: } if proposal_scores is not None: proposal_outputs[proposal_scores.name] = "proposal.scores" - branch = { - "kind": "branch", - "predicate": "acceptance.done", - "cases": { - "true": _invoke( - "branch_state", - {"tokens": "acceptance.tokens"}, - {"next_tokens": "branch.accepted"}, - {"state": _effect("branch.state.in", "branch.state.accepted")}, + rollback_nodes: list[dict[str, Any]] = [] + accepted_case_nodes = [ + _invoke( + "branch_state", + {"tokens": "acceptance.tokens"}, + {"next_tokens": "branch.accepted"}, + {"state": _effect("branch.state.in", "branch.state.accepted")}, + ) + ] + corrected_case_nodes = [ + _invoke( + "branch_state", + {"tokens": "acceptance.tokens"}, + {"next_tokens": "branch.corrected"}, + {"state": _effect("branch.state.in", "branch.state.corrected")}, + ) + ] + branch_outputs: dict[str, Any] = { + "tokens.next": { + "cases": { + "true": "branch.accepted", + "false": "branch.corrected", + } + } + } + branch_effects: dict[str, Any] = { + "state": { + "incoming": "branch.state.in", + "cases": { + "true": "branch.state.accepted", + "false": "branch.state.corrected", + }, + "produces": "branch.state.out", + } + } + for index, (past, present) in enumerate(cache_pairs): + cache_name = f"cache_{index}" + cache_contract = _contract(past) + sequence_axis = next( + ( + axis + for axis, dimension in enumerate(cache_contract["shape"]) + if "sequence" in str(dimension) ), - "false": _invoke( - "branch_state", - {"tokens": "acceptance.tokens"}, - {"next_tokens": "branch.corrected"}, - {"state": _effect("branch.state.in", "branch.state.corrected")}, + 2, + ) + rollback_name = f"rollback_{cache_name}" + publisher_name = f"publish_{cache_name}" + branch_effect = f"branch:{cache_name}" + pkg.add_policy_component( + rollback_name, + build_speculative_state_rollback( + past.dtype, + cache_contract["shape"], + sequence_axis=sequence_axis, + effect=rollback_name, ), - }, - "outputs": { - "tokens.next": { - "cases": { - "true": "branch.accepted", - "false": "branch.corrected", - } - } - }, - "effects": { - "state": { - "incoming": "branch.state.in", - "cases": { - "true": "branch.state.accepted", - "false": "branch.state.corrected", + ) + pkg.add_policy_component( + publisher_name, + build_effectful_identity( + publisher_name, + past.dtype, + [ + "branch_sequence" if axis == sequence_axis else dimension + for axis, dimension in enumerate(cache_contract["shape"]) + ], + effect=branch_effect, + ), + ) + rollback_nodes.append( + _invoke( + rollback_name, + { + "past_state": f"state.{cache_name}.body", + "tentative_state": f"verifier.{present.name}", + "accepted_len": "acceptance.length", + }, + {"corrected_state": f"rollback.{cache_name}"}, + { + rollback_name: _effect( + f"rollback.{cache_name}.0", + f"rollback.{cache_name}.1", + ) + }, + ) + ) + accepted_case_nodes.append( + _invoke( + publisher_name, + {"value": f"verifier.{present.name}"}, + {"next_value": f"branch.accepted.{cache_name}"}, + { + branch_effect: _effect( + f"branch.{cache_name}.in", + f"branch.{cache_name}.accepted", + ) + }, + ) + ) + corrected_case_nodes.append( + _invoke( + publisher_name, + {"value": f"rollback.{cache_name}"}, + {"next_value": f"branch.corrected.{cache_name}"}, + { + branch_effect: _effect( + f"branch.{cache_name}.in", + f"branch.{cache_name}.corrected", + ) }, - "produces": "branch.state.out", + ) + ) + branch_outputs[f"{cache_name}.next"] = { + "cases": { + "true": f"branch.accepted.{cache_name}", + "false": f"branch.corrected.{cache_name}", } + } + branch_effects[branch_effect] = { + "incoming": f"branch.{cache_name}.in", + "cases": { + "true": f"branch.{cache_name}.accepted", + "false": f"branch.{cache_name}.corrected", + }, + "produces": f"branch.{cache_name}.out", + } + branch = { + "kind": "branch", + "predicate": "acceptance.done", + "cases": { + "true": {"kind": "sequence", "nodes": accepted_case_nodes}, + "false": {"kind": "sequence", "nodes": corrected_case_nodes}, }, + "outputs": branch_outputs, + "effects": branch_effects, } body_nodes = [ _invoke("proposer", proposer_inputs, proposal_outputs), @@ -1771,6 +1874,7 @@ def build_speculative_workflow_metadata(pkg: Any) -> dict[str, Any]: }, {"verify": _effect("verify.0", "verify.1")}, ), + *rollback_nodes, branch, _invoke( "continue_predicate", @@ -1816,7 +1920,7 @@ def build_speculative_workflow_metadata(pkg: Any) -> dict[str, Any]: "state.rng_offset.final", ), ] - for index, (past, present) in enumerate(cache_pairs): + for index, (past, _present) in enumerate(cache_pairs): cell = f"cache_{index}" initializer = f"request.verifier.{past.name}" state[cell] = { @@ -1842,7 +1946,7 @@ def build_speculative_workflow_metadata(pkg: Any) -> dict[str, Any]: cell, initializer, f"state.{cell}.body", - f"verifier.{present.name}", + f"{cell}.next", f"state.{cell}.final", ) ) @@ -1851,6 +1955,9 @@ def build_speculative_workflow_metadata(pkg: Any) -> dict[str, Any]: "emit": "emit.0", "state": "branch.state.in", } + 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}" diff --git a/src/mobius/integrations/onnx_genai/workflow_metadata_test.py b/src/mobius/integrations/onnx_genai/workflow_metadata_test.py index 7b2d9e5b6..95113879b 100644 --- a/src/mobius/integrations/onnx_genai/workflow_metadata_test.py +++ b/src/mobius/integrations/onnx_genai/workflow_metadata_test.py @@ -160,7 +160,7 @@ def test_speculative_workflow_uses_branch_phi_effect_join_and_rng(): "workflow" ] body = workflow["graph"]["body"]["nodes"] - branch = body[3] + branch = next(node for node in body if node["kind"] == "branch") assert branch["kind"] == "branch" assert branch["outputs"]["tokens.next"]["cases"] == { "true": "branch.accepted", @@ -170,4 +170,13 @@ def test_speculative_workflow_uses_branch_phi_effect_join_and_rng(): acceptance = body[2] assert acceptance["inputs"]["offset"] == "state.rng_offset.body" assert acceptance["outputs"]["next_offset"] == "rng_offset.body" + rollback = next( + node for node in body if node.get("component") == "rollback_cache_0" + ) + assert rollback["inputs"]["accepted_len"] == "acceptance.length" + assert branch["outputs"]["cache_0.next"]["cases"] == { + "true": "branch.accepted.cache_0", + "false": "branch.corrected.cache_0", + } + assert branch["effects"]["branch:cache_0"]["produces"] == "branch.cache_0.out" assert any(item["cell"].startswith("cache_") for item in workflow["graph"]["carried"]) From 13b99548cd6f18d5843f11c8f2a5cd9dcbe50063 Mon Sep 17 00:00:00 2001 From: justinchuby Date: Wed, 12 Aug 2026 22:11:57 +0000 Subject: [PATCH 017/151] Migrate Qwen3 TTS to real nested workflow Export trained predictor transition graphs and generate typed talker and predictor state policies. Build the nested workflow with complete first-frame code prediction, KV recurrence, codec layout, and waveform emission. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: justinchuby --- src/mobius/generation/__init__.py | 6 + src/mobius/generation/_policy_components.py | 175 +++- .../codec_workflow_metadata_test.py | 41 + .../onnx_genai/workflow_metadata.py | 978 +++++++++++++++++- src/mobius/models/qwen3_tts_test.py | 47 + src/mobius/tasks/_tts.py | 101 ++ 6 files changed, 1331 insertions(+), 17 deletions(-) diff --git a/src/mobius/generation/__init__.py b/src/mobius/generation/__init__.py index bd109c5dc..c17b643f1 100644 --- a/src/mobius/generation/__init__.py +++ b/src/mobius/generation/__init__.py @@ -33,6 +33,9 @@ build_speculative_state_rollback, build_token_block_identity, build_token_state_update, + build_token_to_slot, + build_tts_decoder_state_initializer, + build_tts_decoder_step_update, build_tts_state_initializer, ) @@ -64,5 +67,8 @@ "build_speculative_state_rollback", "build_token_state_update", "build_token_block_identity", + "build_token_to_slot", + "build_tts_decoder_state_initializer", + "build_tts_decoder_step_update", "build_tts_state_initializer", ] diff --git a/src/mobius/generation/_policy_components.py b/src/mobius/generation/_policy_components.py index 2b8a3662c..40846598f 100644 --- a/src/mobius/generation/_policy_components.py +++ b/src/mobius/generation/_policy_components.py @@ -134,7 +134,7 @@ def _make_graph(name: str) -> tuple[ir.Graph, GraphBuilder]: return graph, GraphBuilder(graph) -def build_greedy_sampler() -> PolicyComponent: +def build_greedy_sampler(*, effect: str = "sample") -> PolicyComponent: """Build ``logits -> token_ids`` greedy sampling over the final axis.""" graph, builder = _make_graph("greedy_sampler") logits = builder.input( @@ -152,9 +152,9 @@ def build_greedy_sampler() -> PolicyComponent: "mode": "greedy", "logits": "logits", "token": "token", - "effect": "sample", + "effect": effect, }, - "sample", + effect, ) @@ -240,14 +240,162 @@ def build_tts_state_initializer(num_code_groups: int) -> PolicyComponent: 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(PolicyRole.AUXILIARY, 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(PolicyRole.AUXILIARY, 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(PolicyRole.AUXILIARY, graph, {}) + + def build_code_frame_update( num_code_groups: int, *, scalar_index: bool = False ) -> PolicyComponent: @@ -416,7 +564,8 @@ def build_decoder_state_initializer( shape_parts.append(op.Constant(value_ints=[dimension])) else: raise ValueError( - f"cache input {name!r} has unsupported symbolic dimension {dimension_text!r}" + f"cache input {name!r} has unsupported symbolic " + f"dimension {dimension_text!r}" ) cache_shape = op.Concat(*shape_parts, axis=0) zero = 0.0 if value.dtype.is_floating_point else 0 @@ -601,9 +750,7 @@ def build_euler_model_input(dtype: ir.DataType = ir.DataType.FLOAT) -> PolicyCom """Scale a latent for the Euler denoiser input at the current sigma.""" graph, builder = _make_graph("euler_model_input") op = builder.op - sample = builder.input( - "sample", dtype, ["batch", "channels", "height", "width"] - ) + sample = builder.input("sample", dtype, ["batch", "channels", "height", "width"]) step = builder.input("step", ir.DataType.INT64, ["batch"]) schedule = builder.input("schedule", ir.DataType.FLOAT, ["schedule_length"]) sigma = op.Gather(schedule, step, axis=0) @@ -818,9 +965,7 @@ def build_speculative_acceptance() -> PolicyComponent: # shortest verified prefix so every row can share the same rollback point. synchronized_len = op.ReduceMin(accepted_count, axes=[0], keepdims=1) accepted_count = op.Expand(synchronized_len, op.Shape(accepted_count)) - synchronized_done = op.ReduceMin( - op.Cast(done, to=ir.DataType.INT64), axes=[0], keepdims=1 - ) + synchronized_done = op.ReduceMin(op.Cast(done, to=ir.DataType.INT64), axes=[0], keepdims=1) done = op.Expand( op.Cast(synchronized_done, to=ir.DataType.BOOL), op.Shape(done), @@ -947,3 +1092,13 @@ def build_token_state_update() -> PolicyComponent: }, "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(PolicyRole.AUXILIARY, graph, {}) diff --git a/src/mobius/integrations/onnx_genai/codec_workflow_metadata_test.py b/src/mobius/integrations/onnx_genai/codec_workflow_metadata_test.py index 3818fc033..fd68df735 100644 --- a/src/mobius/integrations/onnx_genai/codec_workflow_metadata_test.py +++ b/src/mobius/integrations/onnx_genai/codec_workflow_metadata_test.py @@ -22,6 +22,9 @@ 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: @@ -160,3 +163,41 @@ def test_tts_uses_nested_lexical_loop_induction_and_codec(): assert inner["body"]["nodes"][0]["inputs"]["step_index"] == "code.iteration" assert workflow["graph"]["nodes"][-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"] == "growing" + assert workflow["state"]["predictor_cache_0"]["scope"] == "invocation" + assert ( + workflow["state"]["predictor_cache_0"]["recurrence"]["max"] + == "package.predictor_context_limit" + ) + outer = workflow["graph"]["nodes"][0] + setup_history = next( + node + for node in outer["setup"]["nodes"] + 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["body"]["nodes"] if node["kind"] == "loop") + assert inner["iteration"]["value"] == "code.iteration" + assert any( + node.get("component") == "code_predictor_step_embedder" + for node in inner["body"]["nodes"] + ) diff --git a/src/mobius/integrations/onnx_genai/workflow_metadata.py b/src/mobius/integrations/onnx_genai/workflow_metadata.py index 41812cb23..4e6dd8dbb 100644 --- a/src/mobius/integrations/onnx_genai/workflow_metadata.py +++ b/src/mobius/integrations/onnx_genai/workflow_metadata.py @@ -24,6 +24,7 @@ build_effectful_identity, build_euler_model_input, build_euler_solver_step, + build_greedy_sampler, build_integer_increment, build_last_token_logits, build_model_token_cast, @@ -31,6 +32,9 @@ build_schedule_lookup, build_speculative_state_rollback, build_token_block_identity, + build_token_to_slot, + build_tts_decoder_state_initializer, + build_tts_decoder_step_update, build_tts_state_initializer, ) from mobius.integrations.onnx_genai.inference_metadata import ( @@ -197,8 +201,970 @@ def write_audio_codec_workflow_metadata(pkg: Any, output_dir: str) -> str: 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 ( + past.name.replace("past_key_values", "present"), + past.name.replace("past.", "present."), + ) + if name in outputs + ), + None, + ) + if present is not None: + pairs.append((past, present)) + return pairs + + +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) + 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)) + 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 = {"dtype": "int64", "rank": 1, "shape": [batch]} + batch_bool = {"dtype": "bool", "rank": 1, "shape": [batch]} + 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": batch_int, + "role": {"kind": "runtime", "version": "1.0", "role": "max_output_tokens"}, + "source": {"kind": "request", "field": "max_output_tokens"}, + "required": True, + }, + "package.false": { + "contract": 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": batch_int, + "role": {"kind": "opaque"}, + "source": {"kind": "literal"}, + "required": False, + "default": num_groups - 2, + }, + "package.predictor_context_limit": { + "contract": {"dtype": "int64", "rank": 0, "shape": []}, + "role": {"kind": "opaque"}, + "source": {"kind": "literal"}, + "required": False, + "default": num_groups, + }, + "package.predictor_mask_limit": { + "contract": {"dtype": "int64", "rank": 0, "shape": []}, + "role": {"kind": "opaque"}, + "source": {"kind": "literal"}, + "required": False, + "default": num_groups + 1, + }, + "package.talker_context_limit": { + "contract": {"dtype": "int64", "rank": 0, "shape": []}, + "role": {"kind": "opaque"}, + "source": {"kind": "literal"}, + "required": False, + "default": int(getattr(config, "max_position_embeddings", 4096)), + }, + } + 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( + "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" + ), + }, + ] + 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( + "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_scalar", + "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_scalar", + "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_scalar", + "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"}, + }, + } + for index, (past, present) in enumerate(talker_caches): + state[f"talker_cache_{index}"] = { + "contract": _contract(past), + "scope": "invocation", + "initializer": f"talker.setup.{present.name}", + "recurrence": { + "kind": "growing", + "axis": 2, + "increment": "package.one_scalar", + "max": "package.talker_context_limit", + }, + } + for index, (past, present) in enumerate(predictor_caches): + state[f"predictor_cache_{index}"] = { + "contract": _contract(past), + "scope": "invocation", + "initializer": f"frame.predictor.{present.name}", + "recurrence": { + "kind": "growing", + "axis": 2, + "increment": "package.one_scalar", + "max": "package.predictor_context_limit", + }, + } + + outer_carried = [ + { + "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": { + "ir_version": "1.0", + "onnx_opsets": {"ai.onnx": OPSET_VERSION}, + "capabilities": [ + "workflow_ssa", + "linear_effects", + "nested_control_flow", + "loop_induction_values", + "typed_emit", + ], + }, + "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, + "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": 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", @@ -351,9 +1317,7 @@ def bind_remaining( result[value.name] = name return result - def bind_outputs( - values: Any, bound: dict[str, str], prefix: str - ) -> dict[str, str]: + 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}") @@ -759,9 +1723,7 @@ def build_diffusion_workflow_metadata( ) attach_policy_components(pkg, PolicyCapabilities()) - pkg.add_policy_component( - "euler_model_input", build_euler_model_input(sample_input.dtype) - ) + pkg.add_policy_component("euler_model_input", build_euler_model_input(sample_input.dtype)) pkg.add_policy_component("solver_step", build_euler_solver_step(sample_input.dtype)) pkg.add_policy_component("continue_predicate", build_boolean_not()) schedule_values = schedule or [ @@ -769,7 +1731,9 @@ def build_diffusion_workflow_metadata( ] 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") + 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)) diff --git a/src/mobius/models/qwen3_tts_test.py b/src/mobius/models/qwen3_tts_test.py index 4851867c9..7cc3eb301 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,49 @@ 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_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/tasks/_tts.py b/src/mobius/tasks/_tts.py index 331f2c59a..fb05403b1 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. """ @@ -86,6 +90,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 +103,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 +225,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 +314,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 +437,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") From 3de81570a3f4df7938f6d2edeee6d9d23adc48e9 Mon Sep 17 00:00:00 2001 From: justinchuby Date: Wed, 12 Aug 2026 22:31:58 +0000 Subject: [PATCH 018/151] Emit typed SSA image preprocessing programs Name every preprocessing transform value and add explicit derived size, mask, coordinate, and grid producers so VLM adapters satisfy the workflow contract and execute through vision, embedding, and decoder stages. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: justinchuby --- .../onnx_genai/workflow_metadata.py | 53 ++++++++++++++++++- .../onnx_genai/workflow_metadata_test.py | 50 +++++++++++++++-- 2 files changed, 99 insertions(+), 4 deletions(-) diff --git a/src/mobius/integrations/onnx_genai/workflow_metadata.py b/src/mobius/integrations/onnx_genai/workflow_metadata.py index 4e6dd8dbb..bf5274b68 100644 --- a/src/mobius/integrations/onnx_genai/workflow_metadata.py +++ b/src/mobius/integrations/onnx_genai/workflow_metadata.py @@ -79,6 +79,57 @@ def _effect(consumes: str, produces: str) -> dict[str, str]: return {"consumes": consumes, "produces": produces} +def _name_image_preprocessing_program(image: dict[str, Any]) -> None: + """Convert structural preprocessing transforms into explicit typed SSA values.""" + transforms = image["transforms"] + 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 _invoke( component: str, inputs: dict[str, str], @@ -2062,6 +2113,7 @@ def build_vlm_workflow_metadata( 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] = {} @@ -2073,7 +2125,6 @@ def build_vlm_workflow_metadata( raise ValueError(f"preprocessing output {endpoint!r} has no vision input") output["contract"] = _contract(vision_inputs[port_name]) output["dtype"] = output["contract"]["dtype"] - output["source"] = output["content"] output["name"] = f"image.{port_name}" adapter_outputs[port_name] = output["contract"] preprocessing_values[port_name] = output["name"] diff --git a/src/mobius/integrations/onnx_genai/workflow_metadata_test.py b/src/mobius/integrations/onnx_genai/workflow_metadata_test.py index 95113879b..42b7ecc0f 100644 --- a/src/mobius/integrations/onnx_genai/workflow_metadata_test.py +++ b/src/mobius/integrations/onnx_genai/workflow_metadata_test.py @@ -11,9 +11,15 @@ import pytest 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 ( build_language_diffusion_pipeline_metadata, build_speculative_workflow_metadata, + build_vlm_workflow_metadata, write_speculative_workflow_metadata, ) @@ -28,6 +34,46 @@ 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 _masked_denoiser_package() -> ModelPackage: input_ids = _value("input_ids", ir.DataType.INT64, ["batch", "sequence"]) logits = _value("logits", ir.DataType.FLOAT, ["batch", "sequence", 128]) @@ -170,9 +216,7 @@ def test_speculative_workflow_uses_branch_phi_effect_join_and_rng(): acceptance = body[2] assert acceptance["inputs"]["offset"] == "state.rng_offset.body" assert acceptance["outputs"]["next_offset"] == "rng_offset.body" - rollback = next( - node for node in body if node.get("component") == "rollback_cache_0" - ) + rollback = next(node for node in body if node.get("component") == "rollback_cache_0") assert rollback["inputs"]["accepted_len"] == "acceptance.length" assert branch["outputs"]["cache_0.next"]["cases"] == { "true": "branch.accepted.cache_0", From c6fa8fd98fd8d90247852d9bbe5fec068b8a1f45 Mon Sep 17 00:00:00 2001 From: justinchuby Date: Wed, 12 Aug 2026 23:05:17 +0000 Subject: [PATCH 019/151] Complete speculative prefix rollback workflow Emit synchronized accepted-prefix lengths, use bounded dynamic cache recurrence for verifier rollback, and expose the final-axis prefix contract required by generic valid-length emission. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: justinchuby --- src/mobius/generation/_policy_components.py | 18 ++++++---- .../generation/_policy_components_test.py | 8 +++-- .../onnx_genai/workflow_metadata.py | 35 +++++++++++++++---- .../onnx_genai/workflow_metadata_test.py | 18 +++++++++- 4 files changed, 62 insertions(+), 17 deletions(-) diff --git a/src/mobius/generation/_policy_components.py b/src/mobius/generation/_policy_components.py index 40846598f..c61af6104 100644 --- a/src/mobius/generation/_policy_components.py +++ b/src/mobius/generation/_policy_components.py @@ -965,11 +965,11 @@ def build_speculative_acceptance() -> PolicyComponent: # shortest verified prefix so every row can share the same rollback point. synchronized_len = op.ReduceMin(accepted_count, axes=[0], keepdims=1) accepted_count = op.Expand(synchronized_len, op.Shape(accepted_count)) - synchronized_done = op.ReduceMin(op.Cast(done, to=ir.DataType.INT64), axes=[0], keepdims=1) - done = op.Expand( - op.Cast(synchronized_done, to=ir.DataType.BOOL), - op.Shape(done), + synchronized_done = op.Cast( + op.ReduceMin(op.Cast(done, to=ir.DataType.INT64), axes=[0], keepdims=1), + to=ir.DataType.BOOL, ) + done = op.Expand(synchronized_done, op.Shape(done)) positions = op.Range( op.Constant(value_int=0), op.Squeeze(draft_length, op.Constant(value_ints=[0])), @@ -980,6 +980,7 @@ def build_speculative_acceptance() -> PolicyComponent: 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])), @@ -987,11 +988,15 @@ def build_speculative_acceptance() -> PolicyComponent: 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"]) + synchronized_len.shape = ir.Shape([1]) + synchronized_done.shape = ir.Shape([1]) 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(synchronized_len, "synchronized_len") + builder.add_output(synchronized_done, "synchronized_done") return _component( PolicyRole.SPECULATIVE_ACCEPTANCE, graph, @@ -1029,10 +1034,9 @@ def build_speculative_state_rollback( 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, ["batch"]) + accepted_len = builder.input("accepted_len", ir.DataType.INT64, [1]) past_len = op.Shape(past, start=sequence_axis, end=sequence_axis + 1) - synchronized_len = op.ReduceMin(accepted_len, axes=[0], keepdims=1) - end = op.Add(past_len, synchronized_len) + end = op.Add(past_len, accepted_len) corrected = op.Slice( tentative, op.Constant(value_ints=[0]), diff --git a/src/mobius/generation/_policy_components_test.py b/src/mobius/generation/_policy_components_test.py index d2c2d1cb3..8036cde29 100644 --- a/src/mobius/generation/_policy_components_test.py +++ b/src/mobius/generation/_policy_components_test.py @@ -312,7 +312,7 @@ def test_masked_update_runtime_parity(tmp_path): def test_speculative_acceptance_prefix_runtime(tmp_path): - accepted_tokens, count, done, next_offset = _run( + accepted_tokens, count, done, next_offset, synchronized_len, synchronized_done = _run( build_speculative_acceptance(), tmp_path, { @@ -329,10 +329,12 @@ def test_speculative_acceptance_prefix_runtime(tmp_path): 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(synchronized_len, [3]) + np.testing.assert_array_equal(synchronized_done, [False]) def test_speculative_acceptance_synchronizes_batched_prefixes(tmp_path): - accepted_tokens, count, done, _ = _run( + accepted_tokens, count, done, _, synchronized_len, synchronized_done = _run( build_speculative_acceptance(), tmp_path, { @@ -351,6 +353,8 @@ def test_speculative_acceptance_synchronizes_batched_prefixes(tmp_path): np.testing.assert_array_equal(count, [2, 2]) np.testing.assert_array_equal(accepted_tokens, [[1, 0, 0, 0], [1, 0, 0, 0]]) np.testing.assert_array_equal(done, [False, False]) + np.testing.assert_array_equal(synchronized_len, [2]) + np.testing.assert_array_equal(synchronized_done, [False]) def test_speculative_state_rollback_trims_tentative_cache(tmp_path): diff --git a/src/mobius/integrations/onnx_genai/workflow_metadata.py b/src/mobius/integrations/onnx_genai/workflow_metadata.py index bf5274b68..d1005acf9 100644 --- a/src/mobius/integrations/onnx_genai/workflow_metadata.py +++ b/src/mobius/integrations/onnx_genai/workflow_metadata.py @@ -2586,8 +2586,12 @@ def write_vlm_workflow_metadata( return path -def build_speculative_workflow_metadata(pkg: Any) -> dict[str, Any]: +def build_speculative_workflow_metadata( + pkg: Any, + config: Any | 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"] @@ -2668,6 +2672,13 @@ def build_speculative_workflow_metadata(pkg: Any) -> dict[str, Any]: "required": False, "default": 1, }, + "package.max_context": { + "contract": batch_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"}, @@ -2814,7 +2825,7 @@ def build_speculative_workflow_metadata(pkg: Any) -> dict[str, Any]: { "past_state": f"state.{cache_name}.body", "tentative_state": f"verifier.{present.name}", - "accepted_len": "acceptance.length", + "accepted_len": "acceptance.synchronized_length", }, {"corrected_state": f"rollback.{cache_name}"}, { @@ -2867,7 +2878,7 @@ def build_speculative_workflow_metadata(pkg: Any) -> dict[str, Any]: } branch = { "kind": "branch", - "predicate": "acceptance.done", + "predicate": "acceptance.synchronized_done", "cases": { "true": {"kind": "sequence", "nodes": accepted_case_nodes}, "false": {"kind": "sequence", "nodes": corrected_case_nodes}, @@ -2886,6 +2897,8 @@ def build_speculative_workflow_metadata(pkg: Any) -> dict[str, Any]: "accepted_len": "acceptance.length", "done": "acceptance.done", "next_offset": "rng_offset.body", + "synchronized_len": "acceptance.synchronized_length", + "synchronized_done": "acceptance.synchronized_done", }, {"verify": _effect("verify.0", "verify.1")}, ), @@ -2899,6 +2912,7 @@ def build_speculative_workflow_metadata(pkg: Any) -> dict[str, Any]: { "kind": "emit", "value": "tokens.next", + "valid_length": "acceptance.synchronized_length", "output": "tokens", "mode": "append", "effect_name": "emit", @@ -2943,7 +2957,7 @@ def build_speculative_workflow_metadata(pkg: Any) -> dict[str, Any]: "scope": "invocation", "initializer": initializer, "recurrence": { - "kind": "growing", + "kind": "bounded", "axis": next( ( axis @@ -2952,8 +2966,7 @@ def build_speculative_workflow_metadata(pkg: Any) -> dict[str, Any]: ), 2, ), - "increment": "package.one", - "max": "request.max_iterations", + "max": "package.max_context", }, } state_specs.append( @@ -2998,12 +3011,20 @@ def build_speculative_workflow_metadata(pkg: Any) -> dict[str, Any]: "nested_control_flow", "loop_induction_values", "typed_emit", + "emit_valid_length", + "bounded_state_recurrence", ], }, "inputs": inputs, "outputs": { "tokens": { - "contract": _contract(proposed_tokens), + "contract": { + **_contract(proposed_tokens), + "shape": [ + *_contract(proposed_tokens)["shape"][:-1], + "accepted_sequence", + ], + }, "role": "tokens", "stage": "pre_adapter", } diff --git a/src/mobius/integrations/onnx_genai/workflow_metadata_test.py b/src/mobius/integrations/onnx_genai/workflow_metadata_test.py index 42b7ecc0f..4cc652595 100644 --- a/src/mobius/integrations/onnx_genai/workflow_metadata_test.py +++ b/src/mobius/integrations/onnx_genai/workflow_metadata_test.py @@ -30,6 +30,22 @@ def test_speculative_writer_saves_policy_artifacts(tmp_path): assert (tmp_path / "policies" / "branch_state.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["graph"]["body"]["nodes"] if node["kind"] == "emit") + assert emit["valid_length"] == "acceptance.synchronized_length" + assert "emit_valid_length" in workflow["manifest"]["capabilities"] + 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["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)) @@ -217,7 +233,7 @@ def test_speculative_workflow_uses_branch_phi_effect_join_and_rng(): assert acceptance["inputs"]["offset"] == "state.rng_offset.body" assert acceptance["outputs"]["next_offset"] == "rng_offset.body" rollback = next(node for node in body if node.get("component") == "rollback_cache_0") - assert rollback["inputs"]["accepted_len"] == "acceptance.length" + assert rollback["inputs"]["accepted_len"] == "acceptance.synchronized_length" assert branch["outputs"]["cache_0.next"]["cases"] == { "true": "branch.accepted.cache_0", "false": "branch.corrected.cache_0", From 7fdfe86c7aa88b690509aa3b5fda68d2fe0523ac Mon Sep 17 00:00:00 2001 From: justinchuby Date: Wed, 12 Aug 2026 23:30:02 +0000 Subject: [PATCH 020/151] Harden workflow export contracts Keep corrected speculative tokens out of verifier KV rollback, reject unsupported KV dtype overrides, and require explicit unguided mode until workflow diffusion gains classifier-free guidance. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: justinchuby --- src/mobius/generation/_policy_components.py | 9 +++-- .../generation/_policy_components_test.py | 36 +++++++++++++++++-- .../integrations/onnx_genai/auto_export.py | 26 ++++++++++++++ .../onnx_genai/auto_export_test.py | 23 ++++++++++-- .../onnx_genai/workflow_metadata.py | 3 +- .../onnx_genai/workflow_metadata_test.py | 2 +- 6 files changed, 90 insertions(+), 9 deletions(-) diff --git a/src/mobius/generation/_policy_components.py b/src/mobius/generation/_policy_components.py index c61af6104..e16fc957e 100644 --- a/src/mobius/generation/_policy_components.py +++ b/src/mobius/generation/_policy_components.py @@ -938,7 +938,7 @@ def build_speculative_acceptance() -> PolicyComponent: op.Equal(rejection_count, op.Constant(value_int=0)), to=ir.DataType.INT64, ) - accepted_count = op.ReduceSum(prefix, axes=[-1], keepdims=0) + 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)), @@ -956,14 +956,15 @@ def build_speculative_acceptance() -> PolicyComponent: ) accepted_tokens.shape = ir.Shape(["batch", "draft_sequence"]) draft_length = op.Shape(proposed_tokens, start=1, end=2) - done = op.Equal(accepted_count, draft_length) + done = op.Equal(verified_count, draft_length) accepted_count = op.Min( - op.Add(accepted_count, op.Cast(op.Not(done), to=ir.DataType.INT64)), + op.Add(verified_count, op.Cast(op.Not(done), to=ir.DataType.INT64)), draft_length, ) # Dense batched state has one physical sequence length. Synchronize to the # shortest verified prefix so every row can share the same rollback point. synchronized_len = op.ReduceMin(accepted_count, axes=[0], keepdims=1) + rollback_len = op.ReduceMin(verified_count, axes=[0], keepdims=1) accepted_count = op.Expand(synchronized_len, op.Shape(accepted_count)) synchronized_done = op.Cast( op.ReduceMin(op.Cast(done, to=ir.DataType.INT64), axes=[0], keepdims=1), @@ -989,6 +990,7 @@ def build_speculative_acceptance() -> PolicyComponent: accepted_count.shape = ir.Shape(["batch"]) done.shape = ir.Shape(["batch"]) synchronized_len.shape = ir.Shape([1]) + rollback_len.shape = ir.Shape([1]) synchronized_done.shape = ir.Shape([1]) next_offset.shape = ir.Shape(["batch"]) builder.add_output(accepted_tokens, "accepted_tokens") @@ -997,6 +999,7 @@ def build_speculative_acceptance() -> PolicyComponent: builder.add_output(next_offset, "next_offset") builder.add_output(synchronized_len, "synchronized_len") builder.add_output(synchronized_done, "synchronized_done") + builder.add_output(rollback_len, "rollback_len") return _component( PolicyRole.SPECULATIVE_ACCEPTANCE, graph, diff --git a/src/mobius/generation/_policy_components_test.py b/src/mobius/generation/_policy_components_test.py index 8036cde29..1da87156f 100644 --- a/src/mobius/generation/_policy_components_test.py +++ b/src/mobius/generation/_policy_components_test.py @@ -312,7 +312,15 @@ def test_masked_update_runtime_parity(tmp_path): def test_speculative_acceptance_prefix_runtime(tmp_path): - accepted_tokens, count, done, next_offset, synchronized_len, synchronized_done = _run( + ( + accepted_tokens, + count, + done, + next_offset, + synchronized_len, + synchronized_done, + rollback_len, + ) = _run( build_speculative_acceptance(), tmp_path, { @@ -331,10 +339,33 @@ def test_speculative_acceptance_prefix_runtime(tmp_path): np.testing.assert_array_equal(next_offset, [12]) np.testing.assert_array_equal(synchronized_len, [3]) np.testing.assert_array_equal(synchronized_done, [False]) + np.testing.assert_array_equal(rollback_len, [2]) + (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_synchronizes_batched_prefixes(tmp_path): - accepted_tokens, count, done, _, synchronized_len, synchronized_done = _run( + ( + accepted_tokens, + count, + done, + _, + synchronized_len, + synchronized_done, + rollback_len, + ) = _run( build_speculative_acceptance(), tmp_path, { @@ -355,6 +386,7 @@ def test_speculative_acceptance_synchronizes_batched_prefixes(tmp_path): np.testing.assert_array_equal(done, [False, False]) np.testing.assert_array_equal(synchronized_len, [2]) np.testing.assert_array_equal(synchronized_done, [False]) + np.testing.assert_array_equal(rollback_len, [1]) def test_speculative_state_rollback_trims_tentative_cache(tmp_path): diff --git a/src/mobius/integrations/onnx_genai/auto_export.py b/src/mobius/integrations/onnx_genai/auto_export.py index 961281912..3dde138d7 100644 --- a/src/mobius/integrations/onnx_genai/auto_export.py +++ b/src/mobius/integrations/onnx_genai/auto_export.py @@ -517,6 +517,12 @@ def write_onnx_genai_config( derived = _diffusion_component_kwargs(pkg) for name, value in derived.items(): kwargs.setdefault(name, value) + if "text_encoder_filename" in kwargs and guidance_scale is None: + raise ValueError( + "text-conditioned workflow diffusion does not implement " + "classifier-free guidance; pass guidance_scale=1.0 explicitly " + "to request unguided generation" + ) if guidance_scale is not None and not np.isclose(guidance_scale, 1.0): raise ValueError( "workflow diffusion requires an explicit classifier-free guidance " @@ -549,6 +555,11 @@ def write_onnx_genai_config( return {"inference_metadata": path} 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) return {"inference_metadata": path} @@ -559,6 +570,11 @@ def write_onnx_genai_config( "or a package carrying `.config`)" ) if _looks_like_multimodal(pkg): + 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, @@ -615,6 +631,11 @@ def write_onnx_genai_config( # 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) " @@ -641,6 +662,11 @@ def write_onnx_genai_config( "Multi-decoder pipelines such as TTS require a dedicated emitter." ) + 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) artifacts = {"inference_metadata": path} tokenizer_path = _write_hf_tokenizer(output_dir, source, revision=revision) diff --git a/src/mobius/integrations/onnx_genai/auto_export_test.py b/src/mobius/integrations/onnx_genai/auto_export_test.py index 9683a34ba..51952b21c 100644 --- a/src/mobius/integrations/onnx_genai/auto_export_test.py +++ b/src/mobius/integrations/onnx_genai/auto_export_test.py @@ -312,6 +312,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" @@ -336,6 +337,7 @@ def _boom(*args, **kwargs): 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 @@ -365,7 +367,7 @@ def test_dispatch_diffusion_auto_reads_scheduler_from_source(tmp_path): def test_dispatch_vision_multimodal_pipeline(tmp_path): pkg = _vlm_package() - artifacts = write_onnx_genai_config(pkg, str(tmp_path), kv_native_dtype="bf16") + artifacts = write_onnx_genai_config(pkg, str(tmp_path)) with open(artifacts["inference_metadata"]) as handle: metadata = yaml.safe_load(handle) @@ -387,7 +389,24 @@ def test_dispatch_vision_multimodal_pipeline(tmp_path): 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_unguided_mode(tmp_path): + with pytest.raises(ValueError, match=r"pass guidance_scale=1\.0 explicitly"): + write_onnx_genai_config( + _diffusion_package(text=True), + str(tmp_path), + ) + + +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 = _vlm_package(audio=True) diff --git a/src/mobius/integrations/onnx_genai/workflow_metadata.py b/src/mobius/integrations/onnx_genai/workflow_metadata.py index d1005acf9..e2e9cb704 100644 --- a/src/mobius/integrations/onnx_genai/workflow_metadata.py +++ b/src/mobius/integrations/onnx_genai/workflow_metadata.py @@ -2825,7 +2825,7 @@ def build_speculative_workflow_metadata( { "past_state": f"state.{cache_name}.body", "tentative_state": f"verifier.{present.name}", - "accepted_len": "acceptance.synchronized_length", + "accepted_len": "acceptance.rollback_length", }, {"corrected_state": f"rollback.{cache_name}"}, { @@ -2899,6 +2899,7 @@ def build_speculative_workflow_metadata( "next_offset": "rng_offset.body", "synchronized_len": "acceptance.synchronized_length", "synchronized_done": "acceptance.synchronized_done", + "rollback_len": "acceptance.rollback_length", }, {"verify": _effect("verify.0", "verify.1")}, ), diff --git a/src/mobius/integrations/onnx_genai/workflow_metadata_test.py b/src/mobius/integrations/onnx_genai/workflow_metadata_test.py index 4cc652595..61ae3e0a2 100644 --- a/src/mobius/integrations/onnx_genai/workflow_metadata_test.py +++ b/src/mobius/integrations/onnx_genai/workflow_metadata_test.py @@ -233,7 +233,7 @@ def test_speculative_workflow_uses_branch_phi_effect_join_and_rng(): assert acceptance["inputs"]["offset"] == "state.rng_offset.body" assert acceptance["outputs"]["next_offset"] == "rng_offset.body" rollback = next(node for node in body if node.get("component") == "rollback_cache_0") - assert rollback["inputs"]["accepted_len"] == "acceptance.synchronized_length" + assert rollback["inputs"]["accepted_len"] == "acceptance.rollback_length" assert branch["outputs"]["cache_0.next"]["cases"] == { "true": "branch.accepted.cache_0", "false": "branch.corrected.cache_0", From e464ed9d3584313c9974fbf437f8ceb8c54507eb Mon Sep 17 00:00:00 2001 From: justinchuby Date: Wed, 12 Aug 2026 23:39:40 +0000 Subject: [PATCH 021/151] Add grammar and adaptive proposal policy components Generate schema-aligned grammar-guided sampling and packed advisory adaptive-K policy graphs. Preserve advisory state for invalid telemetry and cover adjacent probing, batching, and proposal metrics numerically.\n\nCo-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: justinchuby --- src/mobius/generation/__init__.py | 8 + src/mobius/generation/_policy_components.py | 321 ++++++++++++++++++ .../generation/_policy_components_test.py | 97 ++++++ 3 files changed, 426 insertions(+) diff --git a/src/mobius/generation/__init__.py b/src/mobius/generation/__init__.py index c17b643f1..881f7104f 100644 --- a/src/mobius/generation/__init__.py +++ b/src/mobius/generation/__init__.py @@ -10,6 +10,7 @@ PolicyComponent, PolicyRole, attach_policy_components, + build_adaptive_k_policy, build_boolean_not, build_code_frame_update, build_code_history_append, @@ -20,12 +21,15 @@ build_eos_termination, build_euler_model_input, build_euler_solver_step, + build_grammar_logits_processor, build_greedy_sampler, build_integer_increment, + build_integer_minimum, build_iteration_cast, build_last_token_logits, build_masked_token_update, build_model_token_cast, + build_proposal_metrics, build_schedule_constant, build_schedule_lookup, build_seeded_categorical_sampler, @@ -44,6 +48,7 @@ "PolicyCapabilities", "PolicyRole", "attach_policy_components", + "build_adaptive_k_policy", "build_boolean_not", "build_code_frame_update", "build_code_history_append", @@ -55,11 +60,14 @@ "build_euler_solver_step", "build_effectful_identity", "build_greedy_sampler", + "build_grammar_logits_processor", "build_integer_increment", + "build_integer_minimum", "build_iteration_cast", "build_last_token_logits", "build_masked_token_update", "build_model_token_cast", + "build_proposal_metrics", "build_schedule_constant", "build_schedule_lookup", "build_seeded_categorical_sampler", diff --git a/src/mobius/generation/_policy_components.py b/src/mobius/generation/_policy_components.py index e16fc957e..034f5de4e 100644 --- a/src/mobius/generation/_policy_components.py +++ b/src/mobius/generation/_policy_components.py @@ -33,6 +33,8 @@ class PolicyRole(StrEnum): SOLVER_STEP = "solver_step" MASKED_UPDATE = "masked_update" SPECULATIVE_ACCEPTANCE = "speculative_verifier" + GRAMMAR_GUIDANCE = "grammar_guidance" + ADAPTIVE_K = "adaptive_k" STATE_UPDATE = "state_update" AUXILIARY = "auxiliary" @@ -71,6 +73,8 @@ class PolicyCapabilities: 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 @@ -104,6 +108,15 @@ def attach_policy_components( 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())) @@ -192,6 +205,33 @@ def build_integer_increment() -> PolicyComponent: return _component(PolicyRole.AUXILIARY, 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(PolicyRole.AUXILIARY, 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(PolicyRole.AUXILIARY, 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") @@ -610,6 +650,287 @@ def build_decoder_step_update( return _component(PolicyRole.AUXILIARY, 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(PolicyRole.GRAMMAR_GUIDANCE, 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( + PolicyRole.ADAPTIVE_K, + 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 deterministic categorical sampling with explicit seed and offset. diff --git a/src/mobius/generation/_policy_components_test.py b/src/mobius/generation/_policy_components_test.py index 1da87156f..6de433545 100644 --- a/src/mobius/generation/_policy_components_test.py +++ b/src/mobius/generation/_policy_components_test.py @@ -12,6 +12,7 @@ PolicyCapabilities, PolicyRole, attach_policy_components, + build_adaptive_k_policy, build_boolean_not, build_code_frame_update, build_decoder_state_initializer, @@ -19,10 +20,13 @@ build_eos_termination, build_euler_model_input, build_euler_solver_step, + build_grammar_logits_processor, build_greedy_sampler, + build_integer_minimum, build_last_token_logits, build_masked_token_update, build_model_token_cast, + build_proposal_metrics, build_seeded_categorical_sampler, build_speculative_acceptance, build_speculative_state_rollback, @@ -45,6 +49,99 @@ def _run_model(model, tmp_path, feeds): 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]) + + +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(), From 91bc6c2f535997a5ccd8a3528d51e4d5bee990fa Mon Sep 17 00:00:00 2001 From: justinchuby Date: Thu, 13 Aug 2026 00:15:41 +0000 Subject: [PATCH 022/151] Migrate workflows to versioned guidance contracts Lower all generic workflows to the public v1 step IR and emit versioned component contracts from the latest ONNX GenAI schema. Wire reusable grammar clone/lookahead/commit actions and advisory adaptive proposal budgets into speculative generation, including synchronized valid-prefix emission and collision-free semantic state names.\n\nCo-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: justinchuby --- src/mobius/generation/__init__.py | 4 + src/mobius/generation/_policy_components.py | 36 +- .../generation/_policy_components_test.py | 9 +- .../integrations/onnx_genai/auto_export.py | 33 +- .../onnx_genai/auto_export_test.py | 36 +- .../codec_workflow_metadata_test.py | 25 +- .../onnx_genai/inference_metadata.py | 29 +- .../onnx_genai/inference_metadata_test.py | 26 +- .../onnx_genai/workflow_metadata.py | 580 +++- .../onnx_genai/workflow_metadata_test.py | 122 +- src/mobius/models/llada_test.py | 6 +- ...ma.json => onnx_genai_b90d949.schema.json} | 3045 ++--------------- 12 files changed, 1138 insertions(+), 2813 deletions(-) rename tests/schemas/{onnx_genai_4c3c4b6.schema.json => onnx_genai_b90d949.schema.json} (63%) diff --git a/src/mobius/generation/__init__.py b/src/mobius/generation/__init__.py index 881f7104f..715b0a006 100644 --- a/src/mobius/generation/__init__.py +++ b/src/mobius/generation/__init__.py @@ -11,6 +11,7 @@ PolicyRole, attach_policy_components, build_adaptive_k_policy, + build_batch_minimum, build_boolean_not, build_code_frame_update, build_code_history_append, @@ -33,6 +34,7 @@ build_schedule_constant, build_schedule_lookup, build_seeded_categorical_sampler, + build_sequence_length, build_speculative_acceptance, build_speculative_state_rollback, build_token_block_identity, @@ -49,6 +51,7 @@ "PolicyRole", "attach_policy_components", "build_adaptive_k_policy", + "build_batch_minimum", "build_boolean_not", "build_code_frame_update", "build_code_history_append", @@ -70,6 +73,7 @@ "build_proposal_metrics", "build_schedule_constant", "build_schedule_lookup", + "build_sequence_length", "build_seeded_categorical_sampler", "build_speculative_acceptance", "build_speculative_state_rollback", diff --git a/src/mobius/generation/_policy_components.py b/src/mobius/generation/_policy_components.py index 034f5de4e..7afdd65b9 100644 --- a/src/mobius/generation/_policy_components.py +++ b/src/mobius/generation/_policy_components.py @@ -34,7 +34,7 @@ class PolicyRole(StrEnum): MASKED_UPDATE = "masked_update" SPECULATIVE_ACCEPTANCE = "speculative_verifier" GRAMMAR_GUIDANCE = "grammar_guidance" - ADAPTIVE_K = "adaptive_k" + ADAPTIVE_K = "adaptive_proposal_budget" STATE_UPDATE = "state_update" AUXILIARY = "auxiliary" @@ -186,11 +186,15 @@ def build_last_token_logits() -> PolicyComponent: def build_boolean_not() -> PolicyComponent: - """Build an explicit ``continue = Not(done)`` predicate transform.""" + """Build one synchronized ``continue = Not(Any(done))`` predicate.""" graph, builder = _make_graph("boolean_not") done = builder.input("done", dtype=ir.DataType.BOOL, shape=["batch"]) - continued = builder.op.Not(done) - continued.shape = ir.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(PolicyRole.AUXILIARY, graph, {}) @@ -216,6 +220,16 @@ def build_integer_minimum() -> PolicyComponent: return _component(PolicyRole.AUXILIARY, 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(PolicyRole.AUXILIARY, graph, {}) + + def build_proposal_metrics() -> PolicyComponent: """Derive evaluated width and budget fullness from a dense proposal.""" graph, builder = _make_graph("proposal_metrics") @@ -232,6 +246,20 @@ def build_proposal_metrics() -> PolicyComponent: return _component(PolicyRole.AUXILIARY, 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(PolicyRole.AUXILIARY, 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") diff --git a/src/mobius/generation/_policy_components_test.py b/src/mobius/generation/_policy_components_test.py index 6de433545..48a9c663f 100644 --- a/src/mobius/generation/_policy_components_test.py +++ b/src/mobius/generation/_policy_components_test.py @@ -13,6 +13,7 @@ PolicyRole, attach_policy_components, build_adaptive_k_policy, + build_batch_minimum, build_boolean_not, build_code_frame_update, build_decoder_state_initializer, @@ -120,6 +121,12 @@ def test_speculative_guidance_length_and_budget_math(tmp_path): 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): @@ -161,7 +168,7 @@ def test_last_token_logits_and_continue_predicate_runtime(tmp_path): tmp_path, {"done": np.array([True, False])}, ) - np.testing.assert_array_equal(continued, [False, True]) + np.testing.assert_array_equal(continued, [False]) def test_decoder_state_initializer_and_step_update_runtime(tmp_path): diff --git a/src/mobius/integrations/onnx_genai/auto_export.py b/src/mobius/integrations/onnx_genai/auto_export.py index 3dde138d7..5ed95f882 100644 --- a/src/mobius/integrations/onnx_genai/auto_export.py +++ b/src/mobius/integrations/onnx_genai/auto_export.py @@ -58,12 +58,15 @@ def _euler_schedule( "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 + 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, @@ -73,8 +76,7 @@ def _euler_schedule( ) else: raise ValueError( - f"workflow diffusion does not support beta schedule " - f"{scheduler.beta_schedule!r}" + 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)) timesteps = np.linspace( @@ -90,6 +92,7 @@ def _euler_schedule( ) return timesteps.tolist(), [*sigmas.tolist(), 0.0] + _DENOISER_KEYS = ("denoiser", "transformer", "unet") @@ -454,7 +457,8 @@ def write_onnx_genai_config( scheduler: SchedulerConfig | None = None, 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. @@ -529,9 +533,7 @@ def write_onnx_genai_config( "component before guidance_scale can differ from 1.0" ) resolved_scheduler = scheduler or SchedulerConfig(kind="euler") - timesteps, sigma_schedule = _euler_schedule( - resolved_scheduler, num_inference_steps - ) + timesteps, sigma_schedule = _euler_schedule(resolved_scheduler, num_inference_steps) path = write_diffusion_workflow_metadata( pkg, output_dir, @@ -560,7 +562,12 @@ def write_onnx_genai_config( "workflow speculative export derives KV state dtype from ONNX ports; " "kv_native_dtype overrides are unsupported" ) - path = write_speculative_workflow_metadata(pkg, output_dir) + path = write_speculative_workflow_metadata( + pkg, + output_dir, + grammar_guidance=grammar_guidance, + adaptive_k_max=adaptive_k_max, + ) return {"inference_metadata": path} resolved_config = config if config is not None else getattr(pkg, "config", None) diff --git a/src/mobius/integrations/onnx_genai/auto_export_test.py b/src/mobius/integrations/onnx_genai/auto_export_test.py index 51952b21c..9ecefb4d0 100644 --- a/src/mobius/integrations/onnx_genai/auto_export_test.py +++ b/src/mobius/integrations/onnx_genai/auto_export_test.py @@ -158,18 +158,22 @@ def test_dispatch_decoder(tmp_path): meta = yaml.safe_load(handle) workflow = meta["pipeline"]["workflow"] assert workflow["manifest"]["ir_version"] == "1.0" - assert workflow["components"]["token_sampler"]["policy"]["role"] == "token_sampler" - assert workflow["components"]["termination"]["policy"]["role"] == ("termination_predicate") - assert workflow["graph"]["kind"] == "loop" + 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" assert all( value["source"]["kind"] != "application" for value in workflow["inputs"].values() ) - assert [node["component"] for node in workflow["graph"]["setup"]["nodes"]] == [ + assert [node["component"] for node in workflow["steps"][0]["setup"]] == [ "decoder_state_initializer", "model", "last_token_logits", ] - body = workflow["graph"]["body"]["nodes"] + 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"] == "sample.body" assert workflow["state"]["iteration"]["initializer"] == "package.zero_iteration" @@ -221,8 +225,8 @@ def test_dispatch_diffusion(tmp_path): with open(arts["inference_metadata"]) as handle: meta = yaml.safe_load(handle) workflow = meta["pipeline"]["workflow"] - assert workflow["graph"]["nodes"][0]["iteration"]["value"] == "loop.iteration" - assert workflow["graph"]["nodes"][1]["component"] == "vae_decoder" + 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() @@ -376,10 +380,10 @@ def test_dispatch_vision_multimodal_pipeline(tmp_path): assert set(pipeline) == {"workflow"} workflow = pipeline["workflow"] assert workflow["manifest"]["adapter_abis"] == {"onnx-genai.image-preprocess": "1"} - assert workflow["graph"]["setup"]["nodes"][0]["component"] == "image_preprocess" - assert workflow["graph"]["setup"]["nodes"][1]["component"] == "vision_encoder" - assert workflow["graph"]["setup"]["nodes"][3]["component"] == "embedding" - assert workflow["graph"]["iteration"]["value"] == "loop.iteration" + assert workflow["steps"][0]["setup"][0]["component"] == "image_preprocess" + assert workflow["steps"][0]["setup"][1]["component"] == "vision_encoder" + assert workflow["steps"][0]["setup"][3]["component"] == "embedding" + assert workflow["steps"][0]["iteration"]["value"] == "loop.iteration" assert workflow["state"]["logits"]["contract"] == { "dtype": "float32", "rank": 2, @@ -415,7 +419,7 @@ def test_dispatch_audio_only_multimodal_pipeline(tmp_path): with open(artifacts["inference_metadata"]) as handle: metadata = yaml.safe_load(handle) - setup = metadata["pipeline"]["workflow"]["graph"]["setup"]["nodes"] + setup = metadata["pipeline"]["workflow"]["steps"][0]["setup"] assert [node["component"] for node in setup[:3]] == [ "image_preprocess", "vision_encoder", @@ -554,8 +558,8 @@ def test_dispatch_audio_codec_pipeline(tmp_path): assert "model" not in metadata assert not {"models", "dataflow", "strategy", "phases"}.intersection(metadata["pipeline"]) workflow = metadata["pipeline"]["workflow"] - assert workflow["graph"]["nodes"][0]["outputs"] == {"codes": "codec.codes"} - assert workflow["graph"]["nodes"][1]["inputs"] == {"codes": "codec.codes"} + assert workflow["steps"][0]["outputs"] == {"codes": "codec.codes"} + assert workflow["steps"][1]["inputs"] == {"codes": "codec.codes"} assert workflow["outputs"]["waveform"]["stage"] == "post_adapter" @@ -624,9 +628,9 @@ def test_dispatch_multi_decoder_tts_with_pre_embedder(tmp_path): artifacts = write_onnx_genai_config(pkg, str(tmp_path)) with open(artifacts["inference_metadata"]) as handle: workflow = yaml.safe_load(handle)["pipeline"]["workflow"] - outer = workflow["graph"]["nodes"][0] + outer = workflow["steps"][0] assert outer["iteration"]["value"] == "talker.iteration" - assert outer["body"]["nodes"][2]["iteration"]["value"] == "code.iteration" + assert outer["steps"][2]["iteration"]["value"] == "code.iteration" assert (tmp_path / "policies" / "code_frame_update.onnx").is_file() diff --git a/src/mobius/integrations/onnx_genai/codec_workflow_metadata_test.py b/src/mobius/integrations/onnx_genai/codec_workflow_metadata_test.py index fd68df735..4950a41d9 100644 --- a/src/mobius/integrations/onnx_genai/codec_workflow_metadata_test.py +++ b/src/mobius/integrations/onnx_genai/codec_workflow_metadata_test.py @@ -41,7 +41,7 @@ def _codec_package() -> ModelPackage: return ModelPackage({"encoder": encoder, "decoder": decoder}) -def test_codec_workflow_has_typed_ssa_effects_and_audio_emit(): +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) @@ -60,7 +60,7 @@ def test_codec_workflow_has_typed_ssa_effects_and_audio_emit(): assert workflow["components"]["encoder"]["effects"] == ["codec_encode"] assert workflow["components"]["decoder"]["effects"] == ["codec_decode"] - encode, decode, emit = workflow["graph"]["nodes"] + encode, decode, emit = workflow["steps"] assert encode["outputs"] == {"codes": "codec.codes"} assert decode["inputs"] == {"codes": "codec.codes"} assert emit == { @@ -68,8 +68,6 @@ def test_codec_workflow_has_typed_ssa_effects_and_audio_emit(): "value": "codec.waveform", "output": "waveform", "mode": "replace", - "effect_name": "audio_emit", - "effect": {"consumes": "audio_emit.0", "produces": "audio_emit.1"}, } assert workflow["outputs"]["waveform"]["role"] == "audio" assert workflow["outputs"]["waveform"]["stage"] == "post_adapter" @@ -156,12 +154,12 @@ def test_tts_uses_nested_lexical_loop_induction_and_codec(): workflow = build_tts_workflow_metadata(_tts_package(), _TtsConfig())["pipeline"][ "workflow" ] - outer = workflow["graph"]["nodes"][0] - inner = outer["body"]["nodes"][2] + outer = workflow["steps"][0] + inner = outer["steps"][2] assert outer["iteration"]["value"] == "talker.iteration" assert inner["iteration"]["value"] == "code.iteration" - assert inner["body"]["nodes"][0]["inputs"]["step_index"] == "code.iteration" - assert workflow["graph"]["nodes"][-2]["component"] == "codec" + assert inner["steps"][0]["inputs"]["step_index"] == "code.iteration" + assert workflow["steps"][-2]["component"] == "codec" assert workflow["outputs"]["waveform"]["stage"] == "post_adapter" @@ -186,18 +184,15 @@ def test_real_qwen3_tts_workflow_carries_trained_transitions_and_kv_state(): workflow["state"]["predictor_cache_0"]["recurrence"]["max"] == "package.predictor_context_limit" ) - outer = workflow["graph"]["nodes"][0] + outer = workflow["steps"][0] setup_history = next( - node - for node in outer["setup"]["nodes"] - if node.get("component") == "code_history_append" + 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["body"]["nodes"] if node["kind"] == "loop") + inner = next(node for node in outer["steps"] if node["kind"] == "loop") assert inner["iteration"]["value"] == "code.iteration" assert any( - node.get("component") == "code_predictor_step_embedder" - for node in inner["body"]["nodes"] + node.get("component") == "code_predictor_step_embedder" for node in inner["steps"] ) diff --git a/src/mobius/integrations/onnx_genai/inference_metadata.py b/src/mobius/integrations/onnx_genai/inference_metadata.py index 831d28f23..c7235fe94 100644 --- a/src/mobius/integrations/onnx_genai/inference_metadata.py +++ b/src/mobius/integrations/onnx_genai/inference_metadata.py @@ -425,9 +425,7 @@ def _max_token_grid_transforms(config: Any, values: dict[str, Any]) -> list[dict return _area_grid_transforms(config, declared) -def _match_max_token_grid( - ports: list[_Port], values: dict[str, Any] -) -> _ImageProgram | None: +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) @@ -1382,6 +1380,29 @@ def add_policy_components_to_workflow( return metadata components = workflow.setdefault("components", {}) workflow_dtypes = {"fp16": "float16", "bf16": "bfloat16", "fp32": "float32"} + + def semantic_contract(contract: dict[str, Any]) -> dict[str, Any]: + role = str(contract["role"]).replace("_", "-") + bindings = { + key: value + for key, value in contract.items() + if key not in {"role", "mode", "effect", "rng", "state_class"} + 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": f"onnx-genai.{role}", + "version": "1", + "bindings": bindings, + } + if "mode" in contract: + declaration["parameters"] = {"mode": contract["mode"]} + return declaration + for name, component in policy_components.items(): model = component.model declaration = { @@ -1418,7 +1439,7 @@ def add_policy_components_to_workflow( "effects": list(component.effects), } if component.contract: - declaration["policy"] = component.contract + declaration["contract"] = semantic_contract(component.contract) components[name] = declaration return metadata diff --git a/src/mobius/integrations/onnx_genai/inference_metadata_test.py b/src/mobius/integrations/onnx_genai/inference_metadata_test.py index 64c4ef7a7..eb9f853b9 100644 --- a/src/mobius/integrations/onnx_genai/inference_metadata_test.py +++ b/src/mobius/integrations/onnx_genai/inference_metadata_test.py @@ -58,12 +58,15 @@ def test_max_token_packed_grid_derives_pixel_area_bounds(): 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, - }) + 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 @@ -93,12 +96,11 @@ def test_workflow_policy_components_reference_saved_onnx_artifacts(tmp_path): } assert set(component["ports"]["inputs"]) == {"logits"} assert set(component["ports"]["outputs"]) == {"token"} - assert component["policy"] == { - "role": "token_sampler", - "mode": "greedy", - "logits": "logits", - "token": "token", - "effect": "sample", + assert component["contract"] == { + "id": "onnx-genai.token-sampler", + "version": "1", + "bindings": {"logits": "logits", "token": "token"}, + "parameters": {"mode": "greedy"}, } assert component["effects"] == ["sample"] assert (tmp_path / component["implementation"]["artifact"]).is_file() diff --git a/src/mobius/integrations/onnx_genai/workflow_metadata.py b/src/mobius/integrations/onnx_genai/workflow_metadata.py index e2e9cb704..38494ca2a 100644 --- a/src/mobius/integrations/onnx_genai/workflow_metadata.py +++ b/src/mobius/integrations/onnx_genai/workflow_metadata.py @@ -15,6 +15,7 @@ from mobius.generation import ( PolicyCapabilities, attach_policy_components, + build_batch_minimum, build_boolean_not, build_code_frame_update, build_code_history_append, @@ -26,10 +27,13 @@ build_euler_solver_step, build_greedy_sampler, build_integer_increment, + build_integer_minimum, build_last_token_logits, build_model_token_cast, + build_proposal_metrics, build_schedule_constant, build_schedule_lookup, + build_sequence_length, build_speculative_state_rollback, build_token_block_identity, build_token_to_slot, @@ -75,10 +79,170 @@ def _component( return component +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"], + } + + 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]: + """Lower the producer's explicit-effect graph into the public v1 step IR.""" + graph = workflow.pop("graph") + workflow.pop("initial_effects", None) + substitutions: dict[str, str] = {} + 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]: + 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 "valid_length" in node: + result["valid_length"] = rewrite(node["valid_length"]) + 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": + 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)] + ) + result = { + "kind": "loop", + "setup": setup_steps, + "steps": body_steps, + "condition": rewrite(node["condition"]), + "max_iterations": rewrite(node["max_iterations"]), + "carried": [ + { + "cell": cell_aliases.get(carry["cell"], carry["cell"]), + "initial": rewrite(carry["current"]), + "next": rewrite(carry["body_output"]), + } + for carry in node.get("carried", []) + ], + } + 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] + return workflow + + def _name_image_preprocessing_program(image: dict[str, Any]) -> None: """Convert structural preprocessing transforms into explicit typed SSA values.""" transforms = image["transforms"] @@ -239,7 +403,7 @@ def build_audio_codec_workflow_metadata(pkg: Any) -> dict[str, Any]: ], }, } - return {"schema_version": "v1", "pipeline": {"workflow": workflow}} + return {"schema_version": "v1", "pipeline": {"workflow": _publish_workflow_v1(workflow)}} def write_audio_codec_workflow_metadata(pkg: Any, output_dir: str) -> str: @@ -1200,7 +1364,10 @@ def frame_generation_nodes(prefix: str, hidden: str, logits: str) -> list[dict[s ], }, } - metadata = {"schema_version": "v1", "pipeline": {"workflow": workflow}} + metadata = { + "schema_version": "v1", + "pipeline": {"workflow": _publish_workflow_v1(workflow)}, + } add_policy_components_to_workflow(metadata, pkg) return metadata @@ -1660,7 +1827,10 @@ def bind_outputs(values: Any, bound: dict[str, str], prefix: str) -> dict[str, s "initial_effects": initial_effects, "graph": {"kind": "sequence", "nodes": post_nodes}, } - metadata = {"schema_version": "v1", "pipeline": {"workflow": workflow}} + metadata = { + "schema_version": "v1", + "pipeline": {"workflow": _publish_workflow_v1(workflow)}, + } add_policy_components_to_workflow(metadata, pkg) return metadata @@ -1991,7 +2161,10 @@ def build_diffusion_workflow_metadata( ], }, } - metadata = {"schema_version": "v1", "pipeline": {"workflow": workflow}} + metadata = { + "schema_version": "v1", + "pipeline": {"workflow": _publish_workflow_v1(workflow)}, + } add_policy_components_to_workflow(metadata, pkg) return metadata @@ -2100,7 +2273,7 @@ def build_vlm_workflow_metadata( value for value in decoder.graph.inputs if value.name not in cache_names - and value.dtype in {ir.DataType.INT32, ir.DataType.INT64} + and value.dtype == ir.DataType.INT64 and value.shape is not None and len(value.shape) == 2 ] @@ -2564,7 +2737,7 @@ def build_vlm_workflow_metadata( metadata = { "schema_version": "v1", "preprocessing": preprocessing, - "pipeline": {"workflow": workflow}, + "pipeline": {"workflow": _publish_workflow_v1(workflow)}, } add_policy_components_to_workflow(metadata, pkg) return metadata @@ -2589,6 +2762,9 @@ def write_vlm_workflow_metadata( 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) @@ -2623,17 +2799,50 @@ def build_speculative_workflow_metadata( 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)) + attach_policy_components( + pkg, + PolicyCapabilities( + speculative_acceptance=True, + grammar_guidance=grammar_guidance, + adaptive_k_max=adaptive_k_max, + ), + ) pkg.add_policy_component("continue_predicate", build_boolean_not()) pkg.add_policy_component("branch_state", build_token_block_identity()) + if grammar_guidance: + pkg.add_policy_component("grammar_length", build_integer_minimum()) + pkg.add_policy_component("grammar_emit_length", build_batch_minimum()) + pkg.add_policy_component("grammar_rollback_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()) batch = _contract(proposer_input)["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]} + control_bool = {"dtype": "bool", "rank": 1, "shape": [1]} inputs: dict[str, Any] = { "request.tokens": { "contract": _contract(proposer_input), @@ -2649,7 +2858,7 @@ def build_speculative_workflow_metadata( "default": 0, }, "request.max_iterations": { - "contract": batch_int, + "contract": control_int, "role": { "kind": "runtime", "version": "1.0", @@ -2673,25 +2882,102 @@ def build_speculative_workflow_metadata( "default": 1, }, "package.max_context": { - "contract": batch_int, + "contract": control_int, "role": {"kind": "opaque"}, "source": {"kind": "literal"}, "required": False, "default": int(getattr(config, "max_position_embeddings", 4096)), }, "package.false": { - "contract": batch_bool, + "contract": control_bool, "role": {"kind": "opaque"}, "source": {"kind": "literal"}, "required": False, "default": False, }, } + 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), @@ -2749,6 +3035,157 @@ def build_speculative_workflow_metadata( } 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.synchronized_length" + cache_rollback_length = "acceptance.rollback_length" + grammar_post_nodes: list[dict[str, Any]] = [] + if grammar_guidance: + emit_length = "grammar.synchronized_length" + cache_rollback_length = "grammar.rollback_length" + grammar_post_nodes.extend( + [ + _invoke( + "grammar_length", + { + "left": "acceptance.length", + "right": "grammar.valid_length", + }, + {"minimum": "grammar.committed_length"}, + ), + _invoke( + "grammar_rollback_length", + { + "left": "acceptance.rollback_length", + "right": "grammar.valid_length", + }, + {"minimum": "grammar.rollback_length"}, + ), + _invoke( + "grammar_emit_length", + {"values": "grammar.committed_length"}, + {"minimum": "grammar.synchronized_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")}, + ) + ) rollback_nodes: list[dict[str, Any]] = [] accepted_case_nodes = [ _invoke( @@ -2825,7 +3262,7 @@ def build_speculative_workflow_metadata( { "past_state": f"state.{cache_name}.body", "tentative_state": f"verifier.{present.name}", - "accepted_len": "acceptance.rollback_length", + "accepted_len": cache_rollback_length, }, {"corrected_state": f"rollback.{cache_name}"}, { @@ -2888,6 +3325,8 @@ def build_speculative_workflow_metadata( } body_nodes = [ _invoke("proposer", proposer_inputs, proposal_outputs), + *proposal_measure_nodes, + *grammar_pre_nodes, _invoke("verifier", verifier_inputs, verifier_outputs), _invoke( "speculative_acceptance", @@ -2903,6 +3342,8 @@ def build_speculative_workflow_metadata( }, {"verify": _effect("verify.0", "verify.1")}, ), + *grammar_post_nodes, + *adaptive_nodes, *rollback_nodes, branch, _invoke( @@ -2913,22 +3354,35 @@ def build_speculative_workflow_metadata( { "kind": "emit", "value": "tokens.next", - "valid_length": "acceptance.synchronized_length", + "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", + "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"}, @@ -2950,11 +3404,62 @@ def build_speculative_workflow_metadata( "state.rng_offset.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", + ), + ] + ) for index, (past, _present) in enumerate(cache_pairs): cell = f"cache_{index}" initializer = f"request.verifier.{past.name}" state[cell] = { "contract": _contract(past), + "class": "semantic", "scope": "invocation", "initializer": initializer, "recurrence": { @@ -2984,6 +3489,10 @@ def build_speculative_workflow_metadata( "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" @@ -3028,7 +3537,7 @@ def build_speculative_workflow_metadata( }, "role": "tokens", "stage": "pre_adapter", - } + }, }, "components": { name: _component(model, _artifact(name, len(pkg))) for name, model in pkg.items() @@ -3054,14 +3563,41 @@ def build_speculative_workflow_metadata( "carried": carried, }, } - metadata = {"schema_version": "v1", "pipeline": {"workflow": workflow}} + 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"), + } + ) + 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) -> str: +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) + metadata = build_speculative_workflow_metadata( + pkg, + grammar_guidance=grammar_guidance, + adaptive_k_max=adaptive_k_max, + ) pkg.save_policy_components(output_dir) path = os.path.join(output_dir, "inference_metadata.yaml") with open(path, "w", encoding="utf-8") as handle: @@ -3604,7 +4140,10 @@ def build_decoder_workflow_metadata( "carried": carried, }, } - metadata = {"schema_version": "1.0", "pipeline": {"workflow": workflow}} + metadata = { + "schema_version": "1.0", + "pipeline": {"workflow": _publish_workflow_v1(workflow)}, + } add_policy_components_to_workflow(metadata, pkg) return metadata @@ -3875,7 +4414,10 @@ def update_invoke( "carried": carried, }, } - metadata = {"schema_version": "1.0", "pipeline": {"workflow": workflow}} + metadata = { + "schema_version": "1.0", + "pipeline": {"workflow": _publish_workflow_v1(workflow)}, + } add_policy_components_to_workflow(metadata, pkg) return metadata diff --git a/src/mobius/integrations/onnx_genai/workflow_metadata_test.py b/src/mobius/integrations/onnx_genai/workflow_metadata_test.py index 61ae3e0a2..ac3ac31b6 100644 --- a/src/mobius/integrations/onnx_genai/workflow_metadata_test.py +++ b/src/mobius/integrations/onnx_genai/workflow_metadata_test.py @@ -30,11 +30,23 @@ def test_speculative_writer_saves_policy_artifacts(tmp_path): assert (tmp_path / "policies" / "branch_state.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() + assert (tmp_path / "policies" / "grammar_emit_length.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["graph"]["body"]["nodes"] if node["kind"] == "emit") + emit = next(node for node in workflow["steps"][0]["steps"] if node["kind"] == "emit") assert emit["valid_length"] == "acceptance.synchronized_length" assert "emit_valid_length" in workflow["manifest"]["capabilities"] assert workflow["outputs"]["tokens"]["contract"]["shape"][-1] == "accepted_sequence" @@ -113,39 +125,42 @@ def test_language_diffusion_uses_exclusive_ssa_workflow(): assert set(pipeline) == {"workflow"} workflow = pipeline["workflow"] - assert workflow["state"]["tokens"]["recurrence"] == {"kind": "invariant"} - assert workflow["state"]["tokens"]["contract"]["shape"] == ["batch", "sequence"] + 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"]["policy"] == { - "role": "masked_update", - "state": "current_tokens", - "proposal": "proposed_tokens", - "mask": "masked", - "step": "step", - "next_state": "next_state", - "next_mask": "next_mask", - "rng": { + 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", }, - "effect": "update", } - graph = workflow["graph"] + graph = workflow["steps"][0] assert graph["kind"] == "loop" assert graph["condition"] == "denoiser.body.continue" assert graph["max_iterations"] == "request.max_iterations" - assert [node["component"] for node in graph["setup"]["nodes"]] == ["model"] - assert [node["kind"] for node in graph["body"]["nodes"]] == [ + assert [node["component"] for node in graph["setup"]] == ["model"] + assert [node["kind"] for node in graph["steps"]] == [ "invoke", "invoke", "invoke", "emit", "invoke", ] - assert graph["body"]["nodes"][0]["inputs"]["total_steps"] == "package.num_steps" - assert graph["body"]["nodes"][-2]["mode"] == "replace" + assert graph["steps"][0]["inputs"]["total_steps"] == "package.num_steps" + assert graph["steps"][-2]["mode"] == "replace" def test_language_diffusion_rejects_zero_steps(): @@ -158,7 +173,7 @@ def test_language_diffusion_rejects_zero_steps(): def test_language_diffusion_matches_pr_828_schema(): schema_path = ( - Path(__file__).parents[4] / "tests" / "schemas" / "onnx_genai_4c3c4b6.schema.json" + Path(__file__).parents[4] / "tests" / "schemas" / "onnx_genai_b90d949.schema.json" ) with schema_path.open(encoding="utf-8") as handle: schema = json.load(handle) @@ -186,10 +201,17 @@ def _graph_model( ) -def _speculative_package() -> ModelPackage: +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", - [_value("tokens", ir.DataType.INT64, ["batch", 4])], + proposer_inputs, [ _value("proposed_tokens", ir.DataType.INT64, ["batch", 4]), _value("proposal_scores", ir.DataType.FLOAT, ["batch", 4, 32]), @@ -217,20 +239,67 @@ def _speculative_package() -> ModelPackage: return ModelPackage({"proposer": proposer, "verifier": verifier}) -def test_speculative_workflow_uses_branch_phi_effect_join_and_rng(): +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" + proposer = next( + node for node in workflow["steps"][0]["steps"] if node.get("component") == "proposer" + ) + assert proposer["inputs"]["proposal_budget"] == "proposal_k" + schema_path = ( + Path(__file__).parents[4] / "tests" / "schemas" / "onnx_genai_b90d949.schema.json" + ) + with schema_path.open(encoding="utf-8") as handle: + jsonschema.validate(instance=metadata, schema=json.load(handle)) + + +def test_speculative_workflow_uses_branch_phi_and_rng(): workflow = build_speculative_workflow_metadata(_speculative_package())["pipeline"][ "workflow" ] - body = workflow["graph"]["body"]["nodes"] + body = workflow["steps"][0]["steps"] branch = next(node for node in body if node["kind"] == "branch") assert branch["kind"] == "branch" assert branch["outputs"]["tokens.next"]["cases"] == { "true": "branch.accepted", "false": "branch.corrected", } - assert branch["effects"]["state"]["produces"] == "branch.state.out" acceptance = body[2] - assert acceptance["inputs"]["offset"] == "state.rng_offset.body" + assert acceptance["inputs"]["offset"] == "rng_offset" assert acceptance["outputs"]["next_offset"] == "rng_offset.body" rollback = next(node for node in body if node.get("component") == "rollback_cache_0") assert rollback["inputs"]["accepted_len"] == "acceptance.rollback_length" @@ -238,5 +307,4 @@ def test_speculative_workflow_uses_branch_phi_effect_join_and_rng(): "true": "branch.accepted.cache_0", "false": "branch.corrected.cache_0", } - assert branch["effects"]["branch:cache_0"]["produces"] == "branch.cache_0.out" - assert any(item["cell"].startswith("cache_") for item in workflow["graph"]["carried"]) + assert any(item["cell"].startswith("cache_") for item in workflow["steps"][0]["carried"]) diff --git a/src/mobius/models/llada_test.py b/src/mobius/models/llada_test.py index c1e9917d1..89fe02641 100644 --- a/src/mobius/models/llada_test.py +++ b/src/mobius/models/llada_test.py @@ -261,5 +261,7 @@ def test_llada_export_signature_matches_masked_diffusion_metadata(): ) workflow = meta["pipeline"]["workflow"] assert "strategy" not in meta["pipeline"] - assert workflow["graph"]["kind"] == "loop" - assert workflow["components"]["masked_update"]["policy"]["role"] == "masked_update" + assert workflow["steps"][0]["kind"] == "loop" + assert workflow["components"]["masked_update"]["contract"]["id"] == ( + "onnx-genai.masked-update" + ) diff --git a/tests/schemas/onnx_genai_4c3c4b6.schema.json b/tests/schemas/onnx_genai_b90d949.schema.json similarity index 63% rename from tests/schemas/onnx_genai_4c3c4b6.schema.json rename to tests/schemas/onnx_genai_b90d949.schema.json index a8bbd429c..3c7d8a6e5 100644 --- a/tests/schemas/onnx_genai_4c3c4b6.schema.json +++ b/tests/schemas/onnx_genai_b90d949.schema.json @@ -173,32 +173,6 @@ ], "type": "string" }, - "BatchingContract": { - "additionalProperties": false, - "properties": { - "batch_axis": { - "format": "uint", - "minimum": 0, - "type": "integer" - }, - "continuous": { - "default": false, - "type": "boolean" - }, - "max_batch_size": { - "format": "uint", - "minimum": 0, - "type": [ - "integer", - "null" - ] - } - }, - "required": [ - "batch_axis" - ], - "type": "object" - }, "ChunkedPrefillConfig": { "description": "Runtime chunked-prefill preference.", "properties": { @@ -214,6 +188,39 @@ }, "type": "object" }, + "ComponentContract": { + "additionalProperties": false, + "properties": { + "bindings": { + "additionalProperties": { + "type": "string" + }, + "default": {}, + "description": "Semantic role to concrete component port name.", + "type": "object" + }, + "id": { + "description": "Versioned semantic capability identifier. It never selects execution behavior.", + "type": "string" + }, + "parameters": { + "additionalProperties": { + "$ref": "#/$defs/ScalarValue" + }, + "default": {}, + "description": "Contract parameters that are not tensor ports, such as adapter actions.", + "type": "object" + }, + "version": { + "type": "string" + } + }, + "required": [ + "id", + "version" + ], + "type": "object" + }, "ComponentImplementation": { "oneOf": [ { @@ -303,126 +310,6 @@ }, "type": "object" }, - "ControlFlow": { - "description": "Generic package control-flow algebra.", - "oneOf": [ - { - "additionalProperties": false, - "properties": { - "kind": { - "const": "sequence", - "type": "string" - }, - "steps": { - "items": { - "$ref": "#/$defs/ControlFlow" - }, - "type": "array" - } - }, - "required": [ - "kind", - "steps" - ], - "type": "object" - }, - { - "additionalProperties": false, - "properties": { - "component": { - "type": "string" - }, - "kind": { - "const": "invoke", - "type": "string" - }, - "when": { - "anyOf": [ - { - "$ref": "#/$defs/Predicate" - }, - { - "type": "null" - } - ] - } - }, - "required": [ - "kind", - "component" - ], - "type": "object" - }, - { - "additionalProperties": false, - "properties": { - "body": { - "$ref": "#/$defs/ControlFlow" - }, - "carried": { - "default": [], - "items": { - "$ref": "#/$defs/LoopCarry" - }, - "type": "array" - }, - "kind": { - "const": "loop", - "type": "string" - }, - "step_program": { - "type": [ - "string", - "null" - ] - }, - "termination": { - "$ref": "#/$defs/Termination" - } - }, - "required": [ - "kind", - "body", - "termination" - ], - "type": "object" - }, - { - "additionalProperties": false, - "properties": { - "cases": { - "additionalProperties": { - "$ref": "#/$defs/ControlFlow" - }, - "type": "object" - }, - "default": { - "anyOf": [ - { - "$ref": "#/$defs/ControlFlow" - }, - { - "type": "null" - } - ] - }, - "kind": { - "const": "branch", - "type": "string" - }, - "predicate": { - "$ref": "#/$defs/Predicate" - } - }, - "required": [ - "kind", - "predicate", - "cases" - ], - "type": "object" - } - ] - }, "DType": { "description": "Scalar dtype vocabulary with common ONNX and runtime aliases.", "oneOf": [ @@ -476,50 +363,6 @@ ], "type": "string" }, - "DataflowEdge": { - "description": "Directed connection between two pipeline component ports.", - "properties": { - "device_transfer": { - "description": "Whether the runtime must move the value between execution devices.", - "type": [ - "boolean", - "null" - ] - }, - "dtype": { - "anyOf": [ - { - "$ref": "#/$defs/TensorDType" - }, - { - "type": "null" - } - ], - "description": "Scalar or logical data type at the component boundary." - }, - "from": { - "description": "Source package input or endpoint in `component.output_name` form.", - "examples": [ - "encoder.hidden_states" - ], - "pattern": "^[^.]+(?:\\.[^.]+)?$", - "type": "string" - }, - "to": { - "description": "Destination package output or endpoint in `component.input_name` form.", - "examples": [ - "decoder.encoder_hidden_states" - ], - "pattern": "^[^.]+(?:\\.[^.]+)?$", - "type": "string" - } - }, - "required": [ - "from", - "to" - ], - "type": "object" - }, "DeviceKind": { "enum": [ "cpu", @@ -531,41 +374,6 @@ ], "type": "string" }, - "DevicePreference": { - "description": "Execution-device preference vocabulary.", - "oneOf": [ - { - "enum": [ - "auto", - "cpu", - "cuda", - "rocm", - "directml", - "coreml", - "webgpu", - "npu" - ], - "title": "Known standard value" - }, - { - "not": { - "enum": [ - "auto", - "cpu", - "cuda", - "rocm", - "directml", - "coreml", - "webgpu", - "npu" - ] - }, - "title": "Forward-compatible extension value", - "type": "string" - } - ], - "type": "string" - }, "DraftConfig": { "description": "Draft-token producer configuration.", "properties": { @@ -632,22 +440,6 @@ ], "type": "string" }, - "EffectTransition": { - "additionalProperties": false, - "properties": { - "consumes": { - "type": "string" - }, - "produces": { - "type": "string" - } - }, - "required": [ - "consumes", - "produces" - ], - "type": "object" - }, "ForkPrecisionPolicy": { "description": "KV fork-precision policy vocabulary.", "oneOf": [ @@ -847,29 +639,6 @@ }, "type": "object" }, - "ImageCorrespondence": { - "description": "Prompt-placeholder to image correspondence vocabulary.", - "oneOf": [ - { - "enum": [ - "prompt_order", - "explicit_indices" - ], - "title": "Known standard value" - }, - { - "not": { - "enum": [ - "prompt_order", - "explicit_indices" - ] - }, - "title": "Forward-compatible extension value", - "type": "string" - } - ], - "type": "string" - }, "ImageOutputBinding": { "description": "One named tensor output produced by an image preprocessing program.\n\nThe output binds a processor-local value to a typed workflow SSA name.\nNeither the name nor the content role is inferred from a model identity.", "properties": { @@ -1017,31 +786,6 @@ ], "description": "A square size or an explicit width/height for an image transform." }, - "ImageTokenCountSource": { - "description": "Image token-count source vocabulary.", - "oneOf": [ - { - "enum": [ - "per_tile", - "per_patch", - "from_grid" - ], - "title": "Known standard value" - }, - { - "not": { - "enum": [ - "per_tile", - "per_patch", - "from_grid" - ] - }, - "title": "Forward-compatible extension value", - "type": "string" - } - ], - "type": "string" - }, "ImageTransform": { "description": "One generic image transform operation.\n\n`op` selects the operation from a generic vocabulary; the remaining fields\nare the parameters that operation reads (only the relevant ones are set).\nEvery parameter is model DATA — concrete sizes, patch sizes, means, and so on\nlive in a model's fixture, never as constants baked into this schema.", "properties": { @@ -1724,26 +1468,6 @@ ], "type": "object" }, - "LoopCarry": { - "additionalProperties": false, - "properties": { - "from": { - "type": "string" - }, - "state": { - "type": "string" - }, - "to": { - "type": "string" - } - }, - "required": [ - "state", - "from", - "to" - ], - "type": "object" - }, "LoopStatePair": { "description": "One fixed-shape loop-carried recurrent-state port pair.\n\nGeneric and architecture-neutral: the runtime zero/other-initializes `input`\non the first step, runs the graph, and copies `output` back into `input` for\nthe next step (`replace` update). This models any fixed recurrent tensor\n(convolution state, linear-attention recurrent state, and so on) without\nreferencing a model family. It is intentionally distinct from growing or\nshared-buffer KV cache, which is declared through `kv_inputs`/`kv_outputs`\nand `kv_update`.", "properties": { @@ -2385,1172 +2109,17 @@ }, "type": "object" }, - "PhaseConfig": { - "description": "Phase gate for one pipeline component.", + "PipelineSpec": { + "additionalProperties": false, + "description": "Executable package described by the universal typed workflow IR.", "properties": { - "run_on": { - "$ref": "#/$defs/PhaseRunOn", - "description": "Pipeline phase in which the component runs." - }, - "when_present": { - "description": "Opaque presence key required for this component to run.", - "minLength": 1, - "type": [ - "string", - "null" - ] - } - }, - "required": [ - "run_on" - ], - "type": "object" - }, - "PhaseRunOn": { - "description": "Pipeline phase gate.\n\nKnown values are enumerated while future strings remain valid.", - "oneOf": [ - { - "enum": [ - "prompt_only", - "every_step", - "always", - "final_only", - "on_demand" - ], - "title": "Known standard value" - }, - { - "not": { - "enum": [ - "prompt_only", - "every_step", - "always", - "final_only", - "on_demand" - ] - }, - "title": "Forward-compatible extension value", - "type": "string" - } - ], - "type": "string" - }, - "PipelineAudioConfig": { - "description": "Waveform contract for a pipeline stage that emits audio.\n\nArchitecture-neutral: the endpoint is an arbitrary `component.output` name\ncarried in the package's metadata, and the sample rate is a declared number.\nNeither is inferred from a model or vendor name.", - "properties": { - "channels": { - "description": "Number of interleaved channels in the waveform. Defaults to 1 (mono).", - "format": "uint16", - "maximum": 65535, - "minimum": 1, - "type": [ - "integer", - "null" - ] - }, - "output": { - "description": "Endpoint carrying the waveform, in `component.output` form.\n\nWhen absent, the runtime uses the sole output of the final-phase\ncomponent, which is unambiguous for the common single-vocoder shape.", - "type": [ - "string", - "null" - ] - }, - "sample_rate": { - "description": "Sample rate, in hertz, of the waveform the pipeline emits.", - "format": "uint32", - "minimum": 1, - "type": [ - "integer", - "null" - ] - } - }, - "type": "object" - }, - "PipelineComponentSpec": { - "description": "One executable ONNX model in a pipeline.", - "properties": { - "device_preference": { - "anyOf": [ - { - "$ref": "#/$defs/DevicePreference" - }, - { - "type": "null" - } - ], - "description": "Optional execution or device preference declared by the model package." - }, - "filename": { - "description": "Non-empty ONNX filename relative to the model package root.", - "examples": [ - "decoder.onnx" - ], - "minLength": 1, - "type": "string" - }, - "io": { - "anyOf": [ - { - "$ref": "#/$defs/ModelIoSpec" - }, - { - "type": "null" - } - ], - "default": null, - "description": "Explicit graph I/O port bindings for this pipeline component.\n\nThe runtime binds decode-step ports from the declared names. A port that\nis not declared is resolved ONLY from an unambiguous io-shape signal;\nwhen the shape is ambiguous the runtime fails with an actionable error\nnaming the key to declare, and never guesses from a tensor name." - }, - "ports": { - "$ref": "#/$defs/ComponentPorts", - "default": { - "inputs": {}, - "outputs": {} - }, - "description": "Typed graph inputs and outputs exposed by this component." - }, - "tokenizer": { - "description": "Tokenizer filename relative to the package root.\n\nIf absent, loaders may use a shared top-level `tokenizer.json`.", - "examples": [ - "tokenizer.json" - ], - "minLength": 1, - "type": [ - "string", - "null" - ] - }, - "type": { - "$ref": "#/$defs/PipelineRole", - "description": "Component role, for example `encoder`, `decoder`, `draft`, `denoiser`, or `vocoder`." - } - }, - "required": [ - "filename", - "type" - ], - "type": "object" - }, - "PipelineRole": { - "description": "Pipeline component-role vocabulary.", - "oneOf": [ - { - "enum": [ - "encoder", - "vision_encoder", - "audio_encoder", - "decoder", - "draft", - "denoiser", - "scheduler", - "vocoder", - "speech_synthesis" - ], - "title": "Known standard value" - }, - { - "not": { - "enum": [ - "encoder", - "vision_encoder", - "audio_encoder", - "decoder", - "draft", - "denoiser", - "scheduler", - "vocoder", - "speech_synthesis" - ] - }, - "title": "Forward-compatible extension value", - "type": "string" - } - ], - "type": "string" - }, - "PipelineSpec": { - "description": "Multi-model pipeline represented as a directed acyclic dataflow graph.", - "properties": { - "audio": { - "anyOf": [ - { - "$ref": "#/$defs/PipelineAudioConfig" - }, - { - "type": "null" - } - ], - "description": "Waveform contract for a pipeline whose final stage emits audio.\n\nPresent for text-to-speech and any other package that produces sound.\nThe sample rate is model DATA: a runtime cannot infer it from a tensor,\nand guessing it silently changes playback pitch and duration." - }, - "batching": { - "anyOf": [ - { - "$ref": "#/$defs/BatchingContract" - }, - { - "type": "null" - } - ], - "description": "Package batching contract." - }, - "control": { - "anyOf": [ - { - "$ref": "#/$defs/ControlFlow" - }, - { - "type": "null" - } - ], - "description": "Universal nested control-flow program." - }, - "dataflow": { - "description": "Directed tensor or data edges between component ports.", - "items": { - "$ref": "#/$defs/DataflowEdge" - }, - "type": "array" - }, - "inputs": { - "additionalProperties": { - "$ref": "#/$defs/TensorContract" - }, - "default": {}, - "description": "Typed inputs exposed by the complete package.", - "type": "object" - }, - "models": { - "additionalProperties": { - "$ref": "#/$defs/PipelineComponentSpec" - }, - "description": "Named model components in the pipeline DAG; at least one component is required.", - "minProperties": 1, - "type": "object" - }, - "outputs": { - "additionalProperties": { - "$ref": "#/$defs/TensorContract" - }, - "default": {}, - "description": "Typed outputs produced by the complete package.", - "type": "object" - }, - "phases": { - "additionalProperties": { - "$ref": "#/$defs/PhaseConfig" - }, - "default": {}, - "description": "Auxiliary-component lifecycle scheduling, keyed by component name.\n\nModels referenced directly by strategy control-flow fields (`decoder`,\n`model`, `denoiser`, `outer`, or `inner`) must not appear here. Every\nother model must have exactly one phase entry.", - "type": "object" - }, - "positions": { - "anyOf": [ - { - "$ref": "#/$defs/PositionProgram" - }, - { - "type": "null" - } - ], - "description": "Declared position-id generation and prefill→decode continuation program.\n\nGeneric and architecture-neutral: parameterized by rank, axis labels, and\nsection sizes so it expresses both ordinary rank-2 linear positions and\nrank-N multimodal coordinates as data — never a model-family branch." - }, - "postprocessing": { - "anyOf": [ - { - "$ref": "#/$defs/PostprocessingSpec" - }, - { - "type": "null" - } - ], - "description": "Declarative output materialization." - }, - "programs": { - "additionalProperties": { - "$ref": "#/$defs/Program" - }, - "default": {}, - "description": "Named data-only sampler, scheduler, solver, and tensor programs.", - "type": "object" - }, - "reducers": { - "additionalProperties": { - "$ref": "#/$defs/ReducerSpec" - }, - "default": {}, - "description": "Explicit fan-in reducers keyed by destination endpoint.", - "type": "object" - }, - "resources": { - "additionalProperties": { - "$ref": "#/$defs/ResourceContract" - }, - "default": {}, - "description": "Typed resource contracts for named components.", - "type": "object" - }, - "states": { - "additionalProperties": { - "$ref": "#/$defs/StateDeclaration" - }, - "default": {}, - "description": "General tensor state, including loop-carried and persistent session state.", - "type": "object" - }, - "strategy": { - "$ref": "#/$defs/PipelineStrategy", - "description": "Loop and execution strategy for the pipeline." - }, - "vision": { - "anyOf": [ - { - "$ref": "#/$defs/PipelineVisionConfig" - }, - { - "type": "null" - } - ], - "description": "Vision-language model token-expansion contract.\n\nWhen present, the engine uses these fields to replace each image\nplaceholder token in the prompt with the declared expanded image-token\nsequence before KV-cache allocation." - }, - "workflow": { - "anyOf": [ - { - "$ref": "#/$defs/WorkflowSpec" - }, - { - "type": "null" - } - ], - "description": "North-star component-centric SSA workflow." - } - }, - "type": "object" - }, - "PipelineStrategy": { - "description": "Parameterized execution strategy for a pipeline or composite stage.", - "properties": { - "batching": { - "description": "Runtime-specific batching parameters." - }, - "cfg_conditioning_input": { - "default": null, - "description": "Denoiser conditioning input port zeroed for the unconditional pass of\nclassifier-free guidance. Required when `guidance_scale` != 1.0.", - "type": [ - "string", - "null" - ] - }, - "decoder": { - "description": "Autoregressive decoder component name.", - "type": [ - "string", - "null" - ] - }, - "denoiser": { - "description": "Iterative or diffusion denoiser component name.", - "type": [ - "string", - "null" - ] - }, - "guidance_scale": { - "description": "Classifier-free guidance scale or equivalent strategy-specific multiplier.", - "format": "float", - "minimum": 0.0, - "type": [ - "number", - "null" - ] - }, - "inner": { - "default": null, - "description": "Inner autoregressive decoder for a `nested_autoregressive` stage.\n\nThe code_predictor: for each outer frame it runs a short inner AR loop of\n`num_code_groups` steps over the residual codebooks, seeded at inner step\n0 by the outer decoder's `last_hidden_state` (routed via a dataflow edge\n`outer.last_hidden_state -> inner.inputs_embeds`) and threading its own\nper-step code embedding on later steps.", - "type": [ - "string", - "null" - ] - }, - "inner_embedding_output": { - "description": "Inner decoder output port threaded across inner steps for a\n`nested_autoregressive` stage.\n\nEach inner step consumes the previous step's per-code embedding as its\n`inputs_embeds` seed; this names the inner decoder OUTPUT port that\nproduces that embedding. It is declared explicitly because the port is\nshape-indistinguishable from other float outputs — the runtime must not\ninfer it by tensor name. Absent on a nested stage ⇒ actionable error\nnaming `pipeline.strategy.inner_embedding_output`.", - "type": [ - "string", - "null" - ] - }, - "kind": { - "$ref": "#/$defs/PipelineStrategyKind", - "description": "Strategy family; determines which strategy-specific fields are meaningful." - }, - "kv_cache": { - "description": "Runtime-specific KV-cache strategy parameters." - }, - "max_tokens": { - "description": "Maximum number of tokens generated by an autoregressive stage.", - "format": "uint", - "minimum": 1, - "type": [ - "integer", - "null" - ] - }, - "model": { - "description": "Single-pass component name.", - "type": [ - "string", - "null" - ] - }, - "num_code_groups": { - "default": null, - "description": "Inner-loop depth (RVQ residual codebook count) for a\n`nested_autoregressive` stage: the number of code tokens collected per\nouter frame. Must be at least 1.", - "format": "uint", - "minimum": 1, - "type": [ - "integer", - "null" - ] - }, - "num_steps": { - "description": "Number of iterative or diffusion steps.", - "format": "uint", - "minimum": 1, - "type": [ - "integer", - "null" - ] - }, - "outer": { - "default": null, - "description": "Outer autoregressive decoder for a `nested_autoregressive` stage.\n\nThe multi-decoder TTS shape: one outer step is one\naudio frame. The outer decoder (talker) produces a per-frame\n`last_hidden_state` that seeds the inner loop (see `inner`).", - "type": [ - "string", - "null" - ] - }, - "pre_embedder": { - "anyOf": [ - { - "$ref": "#/$defs/PreEmbedderSpec" - }, - { - "type": "null" - } - ], - "description": "Optional pre-embedder component driving the outer decoder (talker) of a\n`nested_autoregressive` stage through `inputs_embeds` instead of\n`input_ids`.\n\nA codec-driven TTS talker is not driven by token ids: each step's\n`inputs_embeds` is materialized from the PREVIOUS frame's codes as\n`codec_sum(+ text_embed)` (where\n`codec_sum = codec_embed(code_0) + Σ_i cp_codec_weights[i][codes[i+1]]`).\nWhen this field names such a component (inputs\n`frame_codes [batch, num_code_groups]` int64 `[+ text_embed [batch, 1,\nhidden]]` → output `inputs_embeds [batch, 1, hidden]`), the runtime builds\nthe outer decoder's per-step `inputs_embeds` through it, keeping the engine\ngeneric. Requires a dataflow edge\n`{pre_embedder}.inputs_embeds -> {outer}.inputs_embeds`.\n\nWhen absent the outer loop is `input_ids`-driven (backward compatible).\n\nAll graph-specific port bindings (the pre-embedder's `frame_codes` /\noptional `text_embed` inputs and the output feeding the outer decoder)\nare declared explicitly in [`PreEmbedderSpec`]; the runtime never guesses\nthem by tensor name or dtype." - }, - "prefill_embedder": { - "anyOf": [ - { - "$ref": "#/$defs/PrefillEmbedderSpec" - }, - { - "type": "null" - } - ], - "description": "Optional prefill embedder component that supplies the outer decoder\n(talker) with its real frame-0 PREFILL sequence and the per-frame\ntrailing-text conditioning of a `nested_autoregressive` stage.\n\nThe talker is prefilled with a multi-position embedding\nsequence built from the tokenized prompt, and each subsequent frame is\nconditioned on one trailing-text embedding. This component materializes\nboth from `text_ids`: inputs `text_ids [batch, text_len]` int64 → outputs\n`prefill_embeds [batch, prefill_len, hidden]` float (fed DIRECTLY to the\ntalker's `inputs_embeds` on frame 0) and `trailing_text_embeds [batch,\ntrailing_len, hidden]` float (one vector consumed per outer frame `k >= 1`\nas the pre-embedder's `text_embed`). It runs once in the prompt phase\n(`run_on: prompt_only`); its `text_ids` input is auto-seeded from the\ntokenized prompt.\n\nOnly meaningful together with [`Self::pre_embedder`] (the frame-`k >= 1`\npath feeds the trailing-text vectors through it). When absent, frame 0\nuses a zero seed and every `text_embed` is zero (backward compatible).\n\nAll graph-specific port bindings (the prompt input plus the prefill and\ntrailing-text outputs) are declared explicitly in [`PrefillEmbedderSpec`];\nthe runtime never guesses them by tensor name or dtype." - }, - "scheduler": { - "description": "Scheduler identifier for iterative or diffusion execution.", - "type": [ - "string", - "null" - ] - }, - "scheduler_config": { - "anyOf": [ - { - "$ref": "#/$defs/SchedulerSpec" - }, - { - "type": "null" - } - ], - "description": "Optional diffusion scheduler applied to the denoiser's loop-carried\noutput (treating it as a noise prediction) each step." - }, - "speculative": { - "description": "Runtime-specific speculative execution parameters." - }, - "stages": { - "description": "Ordered child stages for a composite strategy.", - "items": { - "$ref": "#/$defs/PipelineStrategyStage" - }, - "type": "array" - }, - "start_step": { - "default": null, - "description": "First step index for a partial (img2img) denoise loop.\n\nWhen set, the iterative loop runs `start_step..num_steps` instead of the\nfull `0..num_steps`, and the seed (`denoiser` sample input) is expected to\nalready be the encoded image noised to `timesteps[start_step]`. Matches\ndiffusers' img2img `get_timesteps(num_steps, strength)` skip. Default 0.", - "format": "uint", - "minimum": 0, - "type": [ - "integer", - "null" - ] - }, - "state": { - "description": "Runtime-specific iterative state declaration." - }, - "stop_conditions": { - "description": "Runtime-specific stop-condition declarations.", - "items": true, - "type": [ - "array", - "null" - ] - }, - "timestep_input": { - "default": null, - "description": "Denoiser input port that receives the per-step timestep/sigma scalar.\n\nWhen set, the iterative loop feeds this input a rank-1 `float32` value\neach step (from `timesteps` when provided, otherwise the 0-based step\nindex), so a step-aware denoiser can condition on the current step.", - "type": [ - "string", - "null" - ] - }, - "timesteps": { - "default": null, - "description": "Explicit per-step timestep/sigma schedule for an iterative strategy.\n\nWhen present its length must equal `num_steps`; when absent the loop\nuses the 0-based step index. Requires `timestep_input` to have any effect.", - "items": { - "format": "float", - "type": "number" - }, - "type": [ - "array", - "null" - ] - } - }, - "required": [ - "kind" - ], - "type": "object" - }, - "PipelineStrategyKind": { - "description": "Pipeline execution strategy family.\n\nKnown values are enumerated while future strings remain valid.", - "oneOf": [ - { - "enum": [ - "autoregressive", - "iterative", - "diffusion_steps", - "diffusion-steps", - "single_pass", - "single-pass", - "composite", - "nested_autoregressive", - "nested-autoregressive" - ], - "title": "Known standard value" - }, - { - "not": { - "enum": [ - "autoregressive", - "iterative", - "diffusion_steps", - "diffusion-steps", - "single_pass", - "single-pass", - "composite", - "nested_autoregressive", - "nested-autoregressive" - ] - }, - "title": "Forward-compatible extension value", - "type": "string" - } - ], - "type": "string" - }, - "PipelineStrategyStage": { - "additionalProperties": false, - "description": "Named child stage of a composite pipeline strategy.", - "properties": { - "name": { - "description": "Non-empty stage name unique within its containing composite.", - "minLength": 1, - "type": "string" - }, - "strategy": { - "$ref": "#/$defs/PipelineStrategy", - "description": "Execution strategy for this stage." - } - }, - "required": [ - "name", - "strategy" - ], - "type": "object" - }, - "PipelineVisionConfig": { - "description": "Image placeholder token-expansion contract for encoder-free VLM pipelines.\n\nEvery field is optional and additive: legacy documents that declare only\n`image_placeholder_token_id` and `tokens_per_tile` keep working. The richer\nfields mirror the generic expansion the preprocessor already models\n(separate emitted image token, per-tile/per-patch count source, per-image\ncorrespondence, optional row/column separators, and thumbnail order). All of\nit is generic data — no field names or values reference a model family.", - "properties": { - "column_separator_token_id": { - "description": "Optional token ID emitted between columns within a grid row.", - "format": "int64", - "type": [ - "integer", - "null" - ] - }, - "correspondence_summary": { - "description": "Named preprocessing value containing explicit image correspondence data.", - "minLength": 1, - "type": [ - "string", - "null" - ] - }, - "image_correspondence": { - "anyOf": [ - { - "$ref": "#/$defs/ImageCorrespondence" - }, - { - "type": "null" - } - ], - "description": "How prompt placeholders correspond to input images.\n\n`prompt_order` pairs each placeholder with the next input image.\n`explicit_indices` reads correspondence from `correspondence_summary`." - }, - "image_placeholder_token_id": { - "description": "Token ID of the image placeholder in the tokenized prompt.\n\nThe engine replaces every occurrence of this token with the expanded\nimage token sequence before sequence-length and KV-cache sizing.", - "format": "int64", - "type": [ - "integer", - "null" - ] - }, - "image_token_id": { - "description": "Token ID emitted for each expanded image position.\n\nDistinct from `image_placeholder_token_id`: the placeholder marks WHERE\nto expand, while this is the token actually written into the expanded\nsequence. When absent, the placeholder token itself is repeated.", - "format": "int64", - "type": [ - "integer", - "null" - ] - }, - "placeholder_per_image": { - "description": "Whether each placeholder occurrence corresponds to one input image in\nprompt order. Absent means the historical one-placeholder-per-image rule.", - "type": [ - "boolean", - "null" - ] - }, - "row_separator_token_id": { - "description": "Optional token ID emitted between rows of a tiled image grid.", - "format": "int64", - "type": [ - "integer", - "null" - ] - }, - "thumbnail_order": { - "anyOf": [ - { - "$ref": "#/$defs/ThumbnailOrder" - }, - { - "type": "null" - } - ], - "description": "Order of the optional global thumbnail tile relative to the local grid." - }, - "token_count_source": { - "anyOf": [ - { - "$ref": "#/$defs/ImageTokenCountSource" - }, - { - "type": "null" - } - ], - "description": "Where the per-placeholder token count comes from (per tile, per patch, or\na declared grid). Generic selector, never a model-family branch." - }, - "token_count_summary": { - "description": "Named preprocessing value that supplies per-image counts or grid\ndimensions when `token_count_source` is data-derived.\n\nThis is an arbitrary processor output name. A runtime resolves the name\nfrom the declared preprocessing program; it never dispatches on familiar\ntensor names.", - "minLength": 1, - "type": [ - "string", - "null" - ] - }, - "tokens_per_patch": { - "description": "Number of image tokens each patch expands to, used when the count source\nis per patch. Declared data.", - "format": "uint", - "minimum": 1, - "type": [ - "integer", - "null" - ] - }, - "tokens_per_tile": { - "description": "Number of image tokens each tile expands to.\n\nThe total per-tile expansion is `tokens_per_tile * num_tiles`.", - "format": "uint", - "minimum": 1, - "type": [ - "integer", - "null" - ] - } - }, - "type": "object" - }, - "PolicyComponentContract": { - "description": "Stable semantic roles for ONNX policy-math components.\n\nFields map semantic roles to concrete ONNX port names. The corresponding\ntensor contracts live in [`WorkflowComponent::ports`].", - "oneOf": [ - { - "additionalProperties": false, - "properties": { - "effect": { - "type": "string" - }, - "logits": { - "type": "string" - }, - "mode": { - "$ref": "#/$defs/SamplingPolicyMode" - }, - "rng": { - "anyOf": [ - { - "$ref": "#/$defs/RngPortContract" - }, - { - "type": "null" - } - ] - }, - "role": { - "const": "token_sampler", - "type": "string" - }, - "temperature": { - "type": [ - "string", - "null" - ] - }, - "token": { - "type": "string" - }, - "top_k": { - "type": [ - "string", - "null" - ] - }, - "top_p": { - "type": [ - "string", - "null" - ] - } - }, - "required": [ - "role", - "mode", - "logits", - "token", - "effect" - ], - "type": "object" - }, - { - "additionalProperties": false, - "properties": { - "done": { - "type": "string" - }, - "effect": { - "type": "string" - }, - "eos_ids": { - "type": "string" - }, - "iteration": { - "type": "string" - }, - "max_iterations": { - "type": "string" - }, - "role": { - "const": "termination_predicate", - "type": "string" - }, - "tokens": { - "type": "string" - } - }, - "required": [ - "role", - "tokens", - "eos_ids", - "iteration", - "max_iterations", - "done", - "effect" - ], - "type": "object" - }, - { - "additionalProperties": false, - "properties": { - "effect": { - "type": "string" - }, - "estimate": { - "type": "string" - }, - "next_state": { - "type": "string" - }, - "role": { - "const": "solver_step", - "type": "string" - }, - "schedule": { - "type": "string" - }, - "state": { - "type": "string" - }, - "step": { - "type": "string" - } - }, - "required": [ - "role", - "state", - "estimate", - "step", - "schedule", - "next_state", - "effect" - ], - "type": "object" - }, - { - "additionalProperties": false, - "properties": { - "effect": { - "type": "string" - }, - "mask": { - "type": "string" - }, - "next_mask": { - "type": "string" - }, - "next_state": { - "type": "string" - }, - "proposal": { - "type": "string" - }, - "rng": { - "anyOf": [ - { - "$ref": "#/$defs/RngPortContract" - }, - { - "type": "null" - } - ] - }, - "role": { - "const": "masked_update", - "type": "string" - }, - "state": { - "type": "string" - }, - "step": { - "type": "string" - } - }, - "required": [ - "role", - "state", - "proposal", - "mask", - "step", - "next_state", - "next_mask", - "effect" - ], - "type": "object" - }, - { - "additionalProperties": false, - "properties": { - "accepted_len": { - "type": "string" - }, - "accepted_tokens": { - "type": "string" - }, - "done": { - "type": "string" - }, - "effect": { - "type": "string" - }, - "proposal_scores": { - "type": [ - "string", - "null" - ] - }, - "proposed_tokens": { - "type": "string" - }, - "rng": { - "anyOf": [ - { - "$ref": "#/$defs/RngPortContract" - }, - { - "type": "null" - } - ] - }, - "role": { - "const": "speculative_verifier", - "type": "string" - }, - "target_scores": { - "type": "string" - } - }, - "required": [ - "role", - "target_scores", - "proposed_tokens", - "accepted_tokens", - "accepted_len", - "done", - "effect" - ], - "type": "object" - }, - { - "additionalProperties": false, - "properties": { - "current": { - "type": "string" - }, - "effect": { - "type": "string" - }, - "next": { - "type": "string" - }, - "role": { - "const": "state_update", - "type": "string" - }, - "update": { - "type": "string" - } - }, - "required": [ - "role", - "current", - "update", - "next", - "effect" - ], - "type": "object" - } - ] - }, - "PositionContinuation": { - "description": "Prefill→decode position-continuation vocabulary.", - "oneOf": [ - { - "enum": [ - "linear_increment", - "carry_max", - "from_grid" - ], - "title": "Known standard value" - }, - { - "not": { - "enum": [ - "linear_increment", - "carry_max", - "from_grid" - ] - }, - "title": "Forward-compatible extension value", - "type": "string" - } - ], - "type": "string" - }, - "PositionGeneration": { - "description": "Position-value generation vocabulary.", - "oneOf": [ - { - "enum": [ - "linear", - "processor_coordinates" - ], - "title": "Known standard value" - }, - { - "not": { - "enum": [ - "linear", - "processor_coordinates" - ] - }, - "title": "Forward-compatible extension value", - "type": "string" - } - ], - "type": "string" - }, - "PositionProgram": { - "description": "Declared position-id program for a decoder graph.\n\nThe runtime constructs the position tensor from these declared parameters\ninstead of assuming a fixed rank-2 layout. `rank` 1 (with a single axis)\nexpresses ordinary linear positions; `rank` N expresses multi-axis\nmultimodal coordinates. Axis labels and section sizes are opaque DATA — the\nruntime never infers them from a model name.", - "properties": { - "axes": { - "description": "Optional coordinate-stream labels, one per stream (DATA).", - "items": { - "minLength": 1, - "type": "string" - }, - "type": [ - "array", - "null" - ] - }, - "continuation": { - "anyOf": [ - { - "$ref": "#/$defs/PositionContinuation" - }, - { - "type": "null" - } - ], - "description": "How positions continue from the prompt (prefill) into per-token decode." - }, - "dtype": { - "anyOf": [ - { - "$ref": "#/$defs/TensorDType" - }, - { - "type": "null" - } - ], - "description": "Declared dtype of the position tensor." - }, - "generation": { - "anyOf": [ - { - "$ref": "#/$defs/PositionGeneration" - }, - { - "type": "null" - } - ], - "description": "How the position values are generated for prefill.\n\n`linear` generates ordinary sequence positions. `processor_coordinates`\nconsumes the declared processor summaries to construct multi-axis\ncoordinates. Future generation programs remain extensible capability\nstrings rather than model-family branches." - }, - "input": { - "description": "Graph input port that receives the position ids (arbitrary name, DATA).", - "minLength": 1, - "type": "string" - }, - "processor_summaries": { - "description": "Optional processor-summary endpoints this program reads to compute\nmulti-axis coordinates (e.g. a declared grid-dimensions output). Each\nentry is an arbitrary endpoint name (DATA), never a model-family hint.", - "items": { - "minLength": 1, - "type": "string" - }, - "type": [ - "array", - "null" - ] - }, - "rank": { - "description": "Number of coordinate streams carried by the position tensor.\n\n`1` is an ordinary linear position stream; values `> 1` describe\nmulti-axis multimodal coordinates. The physical ONNX tensor rank is\ndeclared separately by `tensor_rank`.", - "format": "uint", - "minimum": 1, - "type": "integer" - }, - "sections": { - "description": "Optional section sizes for sectioned rotary position embeddings.\n\nOpaque list of per-section widths; their meaning is model DATA, not a\nruntime branch.", - "items": { - "format": "uint", - "minimum": 0, - "type": "integer" - }, - "type": [ - "array", - "null" - ] - }, - "tensor_rank": { - "description": "Physical ONNX tensor rank.\n\nRank 2 declares a conventional `[batch, sequence]` linear input. Higher\nranks declare an explicit coordinate axis in addition to batch/sequence\naxes. Absent preserves the legacy mapping (`rank == 1` means tensor rank\n2; otherwise tensor rank 3).", - "format": "uint", - "minimum": 2, - "type": [ - "integer", - "null" - ] - } - }, - "required": [ - "input", - "rank" - ], - "type": "object" - }, - "PostprocessingSpec": { - "additionalProperties": false, - "properties": { - "outputs": { - "additionalProperties": { - "type": "string" - }, - "default": {}, - "type": "object" - }, - "program": { - "$ref": "#/$defs/Program" - } - }, - "required": [ - "program" - ], - "type": "object" - }, - "PreEmbedderSpec": { - "description": "Structured binding for the optional pre-embedder that drives the outer\ndecoder (talker) of a `nested_autoregressive` stage via `inputs_embeds`.\n\nEvery graph-specific port the runtime touches is declared here, so the\nengine never infers a port by tensor name or dtype.", - "properties": { - "component": { - "description": "Declared model name of the pre-embedder component.", - "minLength": 1, - "type": "string" - }, - "frame_codes_input": { - "description": "Pre-embedder input port receiving the previous frame's codes\n(`int64 [batch, num_code_groups]`).", - "minLength": 1, - "type": "string" - }, - "text_embed_input": { - "description": "Optional pre-embedder input port receiving the per-frame trailing-text\nconditioning vector (`float [batch, 1, hidden]`). When absent, the\npre-embedder exposes no trailing-text input.", - "type": [ - "string", - "null" - ] + "workflow": { + "$ref": "#/$defs/WorkflowSpec", + "description": "Required component-centric SSA workflow." } }, "required": [ - "component", - "frame_codes_input" + "workflow" ], "type": "object" }, @@ -3595,197 +2164,6 @@ ], "type": "string" }, - "Predicate": { - "description": "Data-only predicates for branch and loop termination.", - "oneOf": [ - { - "additionalProperties": false, - "properties": { - "input": { - "type": "string" - }, - "op": { - "const": "present", - "type": "string" - } - }, - "required": [ - "op", - "input" - ], - "type": "object" - }, - { - "additionalProperties": false, - "properties": { - "op": { - "const": "bool", - "type": "string" - }, - "value": { - "type": "boolean" - } - }, - "required": [ - "op", - "value" - ], - "type": "object" - }, - { - "additionalProperties": false, - "properties": { - "op": { - "const": "not", - "type": "string" - }, - "value": { - "$ref": "#/$defs/Predicate" - } - }, - "required": [ - "op", - "value" - ], - "type": "object" - }, - { - "additionalProperties": false, - "properties": { - "op": { - "const": "all", - "type": "string" - }, - "values": { - "items": { - "$ref": "#/$defs/Predicate" - }, - "type": "array" - } - }, - "required": [ - "op", - "values" - ], - "type": "object" - }, - { - "additionalProperties": false, - "properties": { - "op": { - "const": "any", - "type": "string" - }, - "values": { - "items": { - "$ref": "#/$defs/Predicate" - }, - "type": "array" - } - }, - "required": [ - "op", - "values" - ], - "type": "object" - }, - { - "additionalProperties": false, - "properties": { - "left": { - "$ref": "#/$defs/ScalarExpr" - }, - "op": { - "const": "equal", - "type": "string" - }, - "right": { - "$ref": "#/$defs/ScalarExpr" - } - }, - "required": [ - "op", - "left", - "right" - ], - "type": "object" - }, - { - "additionalProperties": false, - "properties": { - "left": { - "$ref": "#/$defs/ScalarExpr" - }, - "op": { - "const": "less", - "type": "string" - }, - "right": { - "$ref": "#/$defs/ScalarExpr" - } - }, - "required": [ - "op", - "left", - "right" - ], - "type": "object" - }, - { - "additionalProperties": false, - "properties": { - "left": { - "$ref": "#/$defs/ScalarExpr" - }, - "op": { - "const": "less_equal", - "type": "string" - }, - "right": { - "$ref": "#/$defs/ScalarExpr" - } - }, - "required": [ - "op", - "left", - "right" - ], - "type": "object" - } - ] - }, - "PrefillEmbedderSpec": { - "description": "Structured binding for the optional prefill embedder that supplies the outer\ndecoder (talker) of a `nested_autoregressive` stage with its frame-0 PREFILL\nsequence and per-frame trailing-text conditioning.\n\nEvery graph-specific port the runtime touches is declared here, so the\nengine never infers a port by tensor name or dtype.", - "properties": { - "component": { - "description": "Declared model name of the (prompt-phase) prefill embedder component.", - "minLength": 1, - "type": "string" - }, - "prefill_output": { - "description": "Prefill-embedder output port carrying the talker's frame-0 multi-position\nPREFILL sequence (`float [batch, prefill_len, hidden]`), fed DIRECTLY to\nthe outer decoder's `inputs_embeds` on frame 0.", - "minLength": 1, - "type": "string" - }, - "prompt_input": { - "description": "Prefill-embedder input port receiving the tokenized prompt\n(`int64 [batch, text_len]`, e.g. `text_ids`).", - "minLength": 1, - "type": "string" - }, - "trailing_output": { - "description": "Prefill-embedder output port carrying the per-frame trailing-text vectors\n(`float [batch, trailing_len, hidden]`), one sliced per outer frame\n`k >= 1` into the pre-embedder's `text_embed`.", - "minLength": 1, - "type": "string" - } - }, - "required": [ - "component", - "prompt_input", - "prefill_output", - "trailing_output" - ], - "type": "object" - }, "PreprocessingSpec": { "description": "Declared, architecture-neutral input preprocessing programs.", "properties": { @@ -3795,133 +2173,13 @@ "$ref": "#/$defs/ImagePreprocessingProgram" }, { - "type": "null" - } - ], - "description": "Typed image preprocessing transform program and its named tensor outputs." - } - }, - "type": "object" - }, - "Program": { - "additionalProperties": false, - "description": "Generic tensor/scalar program executed between component invocations.", - "properties": { - "operations": { - "items": { - "$ref": "#/$defs/ProgramOperation" - }, - "type": "array" - } - }, - "required": [ - "operations" - ], - "type": "object" - }, - "ProgramOperation": { - "oneOf": [ - { - "additionalProperties": false, - "properties": { - "from": { - "type": "string" - }, - "op": { - "const": "copy", - "type": "string" - }, - "to": { - "type": "string" - } - }, - "required": [ - "op", - "from", - "to" - ], - "type": "object" - }, - { - "additionalProperties": false, - "properties": { - "dtype": { - "$ref": "#/$defs/TensorDType" - }, - "input": { - "type": "string" - }, - "op": { - "const": "cast", - "type": "string" - }, - "output": { - "type": "string" - } - }, - "required": [ - "op", - "input", - "output", - "dtype" - ], - "type": "object" - }, - { - "additionalProperties": false, - "properties": { - "logits": { - "type": "string" - }, - "method": { - "$ref": "#/$defs/SamplingMethod" - }, - "op": { - "const": "sample", - "type": "string" - }, - "output": { - "type": "string" - } - }, - "required": [ - "op", - "logits", - "output", - "method" - ], - "type": "object" - }, - { - "additionalProperties": false, - "properties": { - "estimate": { - "type": "string" - }, - "op": { - "const": "solver_step", - "type": "string" - }, - "output": { - "type": "string" - }, - "solver": { - "$ref": "#/$defs/SolverSpec" - }, - "state": { - "type": "string" + "type": "null" } - }, - "required": [ - "op", - "estimate", - "state", - "output", - "solver" ], - "type": "object" + "description": "Typed image preprocessing transform program and its named tensor outputs." } - ] + }, + "type": "object" }, "ProposalTopology": { "description": "Speculative proposal-topology vocabulary.", @@ -4071,93 +2329,6 @@ ], "type": "object" }, - "ReducerKind": { - "description": "How multiple dataflow values are combined at one destination.", - "enum": [ - "first", - "last", - "sum", - "product", - "mean", - "min", - "max", - "concat", - "stack" - ], - "type": "string" - }, - "ReducerSpec": { - "additionalProperties": false, - "properties": { - "axis": { - "format": "int64", - "type": [ - "integer", - "null" - ] - }, - "kind": { - "$ref": "#/$defs/ReducerKind" - } - }, - "required": [ - "kind" - ], - "type": "object" - }, - "ResourceContract": { - "additionalProperties": false, - "properties": { - "device": { - "anyOf": [ - { - "$ref": "#/$defs/DeviceKind" - }, - { - "type": "null" - } - ] - }, - "device_index": { - "format": "uint32", - "minimum": 0, - "type": [ - "integer", - "null" - ] - }, - "memory_bytes": { - "format": "uint64", - "minimum": 0, - "type": [ - "integer", - "null" - ] - } - }, - "type": "object" - }, - "RngPortContract": { - "additionalProperties": false, - "description": "Counter-based RNG state. Producers should use Philox or Threefry inside ONNX.", - "properties": { - "next_offset": { - "type": "string" - }, - "offset": { - "type": "string" - }, - "seed": { - "type": "string" - } - }, - "required": [ - "seed", - "offset", - "next_offset" - ], - "type": "object" - }, "RuntimeConfigurable": { "description": "Features whose concrete settings may be selected by the runtime.", "properties": { @@ -4233,114 +2404,6 @@ ], "type": "object" }, - "SamplingMethod": { - "oneOf": [ - { - "additionalProperties": false, - "properties": { - "kind": { - "const": "greedy", - "type": "string" - } - }, - "required": [ - "kind" - ], - "type": "object" - }, - { - "additionalProperties": false, - "properties": { - "kind": { - "const": "categorical", - "type": "string" - }, - "temperature": { - "format": "float", - "type": "number" - }, - "top_k": { - "format": "uint", - "minimum": 0, - "type": [ - "integer", - "null" - ] - }, - "top_p": { - "format": "float", - "type": [ - "number", - "null" - ] - } - }, - "required": [ - "kind", - "temperature" - ], - "type": "object" - } - ] - }, - "SamplingPolicyMode": { - "enum": [ - "greedy", - "seeded_stochastic" - ], - "type": "string" - }, - "ScalarExpr": { - "oneOf": [ - { - "additionalProperties": false, - "properties": { - "kind": { - "const": "literal", - "type": "string" - }, - "value": { - "$ref": "#/$defs/ScalarValue" - } - }, - "required": [ - "kind", - "value" - ], - "type": "object" - }, - { - "additionalProperties": false, - "properties": { - "kind": { - "const": "value", - "type": "string" - }, - "source": { - "type": "string" - } - }, - "required": [ - "kind", - "source" - ], - "type": "object" - }, - { - "additionalProperties": false, - "properties": { - "kind": { - "const": "iteration", - "type": "string" - } - }, - "required": [ - "kind" - ], - "type": "object" - } - ] - }, "ScalarValue": { "anyOf": [ { @@ -4359,115 +2422,6 @@ } ] }, - "SchedulerSpec": { - "description": "Diffusion scheduler configuration for an iterative strategy.\n\nThe runtime treats the denoiser's loop-carried output as a noise prediction\n(or, for `flow_matching`, as a vector field and, for `masked_diffusion`, as\ntoken logits) and applies one scheduler step per iteration. Supported\n`kind`s: `ddpm`, `ddim`, `euler`, `euler_ancestral`, `dpmpp_2m`,\n`flow_matching`, and `masked_diffusion`.", - "properties": { - "beta_end": { - "description": "Linear beta-schedule end (default 0.012).", - "format": "float", - "minimum": 0.0, - "type": [ - "number", - "null" - ] - }, - "beta_schedule": { - "description": "Beta schedule shape: `\"linear\"` (default) or `\"scaled_linear\"` (Stable\nDiffusion).", - "type": [ - "string", - "null" - ] - }, - "beta_start": { - "description": "Linear beta-schedule start (default 0.00085).", - "format": "float", - "minimum": 0.0, - "type": [ - "number", - "null" - ] - }, - "block_length": { - "description": "Semi-autoregressive block length for a `masked_diffusion` scheduler, in\ntokens. When set (and smaller than the masked generation region), each\nstep only commits tokens inside the current left-to-right block, matching\nLLaDA's semi-autoregressive remasking. Defaults to a single block\nspanning the whole masked region.", - "format": "uint", - "minimum": 0, - "type": [ - "integer", - "null" - ] - }, - "kind": { - "description": "Scheduler algorithm: `\"ddpm\"`, `\"ddim\"`, `\"euler\"`,\n`\"euler_ancestral\"`, `\"dpmpp_2m\"`, `\"flow_matching\"`, or\n`\"masked_diffusion\"`.", - "type": "string" - }, - "mask_token_id": { - "description": "Mask token id for a `masked_diffusion` (language-diffusion) scheduler:\neach step commits the highest-confidence still-masked positions.", - "format": "int64", - "type": [ - "integer", - "null" - ] - }, - "num_train_timesteps": { - "description": "Training timesteps the noise schedule was defined over (default 1000).", - "format": "uint", - "minimum": 2, - "type": [ - "integer", - "null" - ] - }, - "prediction_type": { - "description": "Model output parameterization: `\"epsilon\"` (default, noise prediction),\n`\"v_prediction\"` (velocity; SD 2.x, SDXL refiner, many fine-tunes), or\n`\"sample\"`/`\"x0\"` (the model predicts the clean sample directly). All\nbuilt-in diffusion schedulers (`ddpm`, `ddim`, `euler`,\n`euler_ancestral`, `dpmpp_2m`) support every parameterization.\n`flow_matching` instead consumes the model's velocity/vector-field output\ndirectly and accepts an omitted value or `\"flow\"`/`\"velocity\"`.", - "type": [ - "string", - "null" - ] - }, - "remasking": { - "description": "Unmasking strategy for a `masked_diffusion` scheduler:\n * `\"low_confidence\"` (default) — LLaDA: each step commits the\n highest-confidence still-masked positions (confidence-ranked). Best\n for LLaDA checkpoints, but greedy/confidence-ranked decoding of other\n masked-diffusion LMs (e.g. MDLM) collapses into repetitive text.\n * `\"random\"` — MDLM-style ancestral: each still-masked position unmasks\n independently with the schedule probability `1/(steps_remaining)`,\n sampling its token from the model's categorical distribution (use\n `temperature: 1.0` for a true categorical sample). This per-position\n stochastic unmasking avoids the degenerate loops that confidence\n ranking produces. The mask token is never emitted.", - "type": [ - "string", - "null" - ] - }, - "shift": { - "description": "Static timestep shift for `flow_matching` (default `1.0`). The base\nrectified-flow sigma `s` is transformed to\n`shift * s / (1 + (shift - 1) * s)`.", - "format": "float", - "minimum": 0.0, - "type": [ - "number", - "null" - ] - }, - "temperature": { - "description": "Sampling temperature for a `masked_diffusion` scheduler. `0` (default)\nselects each masked position's argmax token deterministically; a positive\nvalue applies Gumbel noise (`logits.exp() / (-log u)^temperature`) before\nthe argmax, matching LLaDA's `add_gumbel_noise`. Confidence used for\nremasking is always the clean-softmax probability of the chosen token.", - "format": "float", - "type": [ - "number", - "null" - ] - }, - "use_exponential_sigmas": { - "description": "Use the exponential sigma spacing (`exp(linspace(log σ_max, log σ_min))`)\ninstead of linspace. Applies to `euler`/`dpmpp_2m`. Mutually exclusive\nwith `use_karras_sigmas` (Karras takes precedence).", - "type": [ - "boolean", - "null" - ] - }, - "use_karras_sigmas": { - "description": "Use the Karras (arXiv:2206.00364, rho=7) sigma spacing instead of the\ndefault linspace spacing. Applies to sigma-space schedulers (`euler`,\n`dpmpp_2m`); the most popular ComfyUI scheduler for those samplers.", - "type": [ - "boolean", - "null" - ] - } - }, - "required": [ - "kind" - ], - "type": "object" - }, "SemanticInputRole": { "oneOf": [ { @@ -4596,28 +2550,53 @@ "additionalProperties": false, "properties": { "kind": { - "const": "invariant", + "const": "invariant", + "type": "string" + } + }, + "required": [ + "kind" + ], + "type": "object" + }, + { + "additionalProperties": false, + "properties": { + "axis": { + "format": "uint", + "minimum": 0, + "type": "integer" + }, + "increment": { + "type": "string" + }, + "kind": { + "const": "growing", + "type": "string" + }, + "max": { "type": "string" } }, "required": [ - "kind" + "kind", + "axis", + "increment", + "max" ], "type": "object" }, { "additionalProperties": false, + "description": "The selected axis may grow or shrink between iterations, but never exceed `max`.", "properties": { "axis": { "format": "uint", "minimum": 0, "type": "integer" }, - "increment": { - "type": "string" - }, "kind": { - "const": "growing", + "const": "bounded", "type": "string" }, "max": { @@ -4627,7 +2606,6 @@ "required": [ "kind", "axis", - "increment", "max" ], "type": "object" @@ -4692,34 +2670,6 @@ ], "type": "string" }, - "SolverSpec": { - "additionalProperties": false, - "properties": { - "algorithm": { - "type": "string" - }, - "parameters": { - "additionalProperties": { - "format": "double", - "type": "number" - }, - "default": {}, - "type": "object" - }, - "schedule": { - "default": [], - "items": { - "format": "double", - "type": "number" - }, - "type": "array" - } - }, - "required": [ - "algorithm" - ], - "type": "object" - }, "SpecialTokens": { "description": "Special / control token ids declared by a model author.\n\nEvery field is optional; `eos_token_id` is normalized to a list because\nonnxruntime-genai accepts either a scalar or an array for it.", "properties": { @@ -5079,111 +3029,6 @@ }, "type": "object" }, - "StateDeclaration": { - "additionalProperties": false, - "properties": { - "init": { - "$ref": "#/$defs/StateInit" - }, - "scope": { - "$ref": "#/$defs/StateScope" - }, - "type": { - "$ref": "#/$defs/TensorContract" - }, - "update": { - "$ref": "#/$defs/StateUpdate" - } - }, - "required": [ - "type", - "init", - "update", - "scope" - ], - "type": "object" - }, - "StateInit": { - "oneOf": [ - { - "additionalProperties": false, - "properties": { - "kind": { - "const": "zeros", - "type": "string" - } - }, - "required": [ - "kind" - ], - "type": "object" - }, - { - "additionalProperties": false, - "properties": { - "kind": { - "const": "ones", - "type": "string" - } - }, - "required": [ - "kind" - ], - "type": "object" - }, - { - "additionalProperties": false, - "properties": { - "kind": { - "const": "input", - "type": "string" - }, - "source": { - "type": "string" - } - }, - "required": [ - "kind", - "source" - ], - "type": "object" - }, - { - "additionalProperties": false, - "properties": { - "kind": { - "const": "value", - "type": "string" - }, - "source": { - "type": "string" - } - }, - "required": [ - "kind", - "source" - ], - "type": "object" - }, - { - "additionalProperties": false, - "properties": { - "kind": { - "const": "scalar", - "type": "string" - }, - "value": { - "$ref": "#/$defs/ScalarValue" - } - }, - "required": [ - "kind", - "value" - ], - "type": "object" - } - ] - }, "StateInitKind": { "description": "Loop-carried state initialization vocabulary.", "oneOf": [ @@ -5205,85 +3050,6 @@ ], "type": "string" }, - "StateScope": { - "enum": [ - "invocation", - "loop", - "request", - "session" - ], - "type": "string" - }, - "StateUpdate": { - "oneOf": [ - { - "additionalProperties": false, - "properties": { - "kind": { - "const": "replace", - "type": "string" - } - }, - "required": [ - "kind" - ], - "type": "object" - }, - { - "additionalProperties": false, - "properties": { - "axis": { - "format": "int64", - "type": "integer" - }, - "kind": { - "const": "append", - "type": "string" - } - }, - "required": [ - "kind", - "axis" - ], - "type": "object" - }, - { - "additionalProperties": false, - "properties": { - "axis": { - "format": "int64", - "type": "integer" - }, - "indices": { - "type": "string" - }, - "kind": { - "const": "scatter", - "type": "string" - } - }, - "required": [ - "kind", - "axis", - "indices" - ], - "type": "object" - }, - { - "additionalProperties": false, - "properties": { - "kind": { - "const": "accumulate", - "type": "string" - } - }, - "required": [ - "kind" - ], - "type": "object" - } - ] - }, "StateUpdateKind": { "description": "Loop-carried state update-semantics vocabulary.", "oneOf": [ @@ -5599,76 +3365,24 @@ "title": "Forward-compatible extension value", "type": "string" } - ], - "type": "string" - }, - "TensorDimension": { - "anyOf": [ - { - "description": "A fixed, non-negative dimension.", - "format": "int64", - "minimum": 0, - "type": "integer" - }, - { - "description": "A runtime shape symbol.", - "minLength": 1, - "type": "string" - } - ], - "description": "One fixed or runtime-resolved tensor-shape dimension." - }, - "Termination": { - "oneOf": [ - { - "additionalProperties": false, - "properties": { - "count": { - "format": "uint", - "minimum": 0, - "type": "integer" - }, - "kind": { - "const": "iterations", - "type": "string" - }, - "start": { - "default": 0, - "format": "uint", - "minimum": 0, - "type": "integer" - } - }, - "required": [ - "kind", - "count" - ], - "type": "object" - }, - { - "additionalProperties": false, - "properties": { - "condition": { - "$ref": "#/$defs/Predicate" - }, - "kind": { - "const": "predicate", - "type": "string" - }, - "max_iterations": { - "format": "uint", - "minimum": 0, - "type": "integer" - } - }, - "required": [ - "kind", - "condition", - "max_iterations" - ], - "type": "object" + ], + "type": "string" + }, + "TensorDimension": { + "anyOf": [ + { + "description": "A fixed, non-negative dimension.", + "format": "int64", + "minimum": 0, + "type": "integer" + }, + { + "description": "A runtime shape symbol.", + "minLength": 1, + "type": "string" } - ] + ], + "description": "One fixed or runtime-resolved tensor-shape dimension." }, "ThumbnailOrder": { "description": "Optional-thumbnail ordering vocabulary.", @@ -5793,7 +3507,7 @@ } ] }, - "WorkflowBranchEffectMerge": { + "WorkflowBranchOutput": { "additionalProperties": false, "properties": { "cases": { @@ -5807,45 +3521,48 @@ "string", "null" ] - }, - "incoming": { - "type": "string" - }, - "produces": { - "type": "string" } }, "required": [ - "incoming", - "cases", - "produces" + "cases" ], "type": "object" }, - "WorkflowBranchOutput": { + "WorkflowCarry": { "additionalProperties": false, "properties": { - "cases": { - "additionalProperties": { - "type": "string" - }, - "type": "object" + "cell": { + "type": "string" }, - "default": { + "initial": { "type": [ "string", "null" ] + }, + "next": { + "type": "string" } }, "required": [ - "cases" + "cell", + "next" ], "type": "object" }, "WorkflowComponent": { "additionalProperties": false, "properties": { + "contract": { + "anyOf": [ + { + "$ref": "#/$defs/ComponentContract" + }, + { + "type": "null" + } + ] + }, "effects": { "default": [], "items": { @@ -5856,18 +3573,12 @@ "implementation": { "$ref": "#/$defs/ComponentImplementation" }, - "policy": { - "anyOf": [ - { - "$ref": "#/$defs/PolicyComponentContract" - }, - { - "type": "null" - } - ] - }, "ports": { - "$ref": "#/$defs/ComponentPorts" + "$ref": "#/$defs/ComponentPorts", + "default": { + "inputs": {}, + "outputs": {} + } }, "resources": { "anyOf": [ @@ -5881,8 +3592,7 @@ } }, "required": [ - "implementation", - "ports" + "implementation" ], "type": "object" }, @@ -5996,42 +3706,6 @@ } ] }, - "WorkflowLoopCarry": { - "additionalProperties": false, - "properties": { - "body_input": { - "type": "string" - }, - "body_output": { - "type": "string" - }, - "cell": { - "type": "string" - }, - "current": { - "type": "string" - }, - "next": { - "type": "string" - }, - "read_effect": { - "$ref": "#/$defs/EffectTransition" - }, - "write_effect": { - "$ref": "#/$defs/EffectTransition" - } - }, - "required": [ - "cell", - "current", - "body_input", - "body_output", - "next", - "read_effect", - "write_effect" - ], - "type": "object" - }, "WorkflowLoopIteration": { "additionalProperties": false, "properties": { @@ -6065,257 +3739,42 @@ "items": { "type": "string" }, - "type": "array", - "uniqueItems": true - }, - "custom_op_versions": { - "additionalProperties": { - "type": "string" - }, - "default": {}, - "type": "object" - }, - "ir_version": { - "type": "string" - }, - "onnx_opsets": { - "additionalProperties": { - "format": "uint32", - "minimum": 0, - "type": "integer" - }, - "default": {}, - "type": "object" - } - }, - "required": [ - "ir_version" - ], - "type": "object" - }, - "WorkflowMemoryClass": { - "enum": [ - "default", - "device", - "host", - "pinned" - ], - "type": "string" - }, - "WorkflowNode": { - "oneOf": [ - { - "additionalProperties": false, - "properties": { - "kind": { - "const": "sequence", - "type": "string" - }, - "nodes": { - "items": { - "$ref": "#/$defs/WorkflowNode" - }, - "type": "array" - } - }, - "required": [ - "kind", - "nodes" - ], - "type": "object" - }, - { - "additionalProperties": false, - "properties": { - "component": { - "type": "string" - }, - "effects": { - "additionalProperties": { - "$ref": "#/$defs/EffectTransition" - }, - "default": {}, - "type": "object" - }, - "inputs": { - "additionalProperties": { - "type": "string" - }, - "default": {}, - "type": "object" - }, - "kind": { - "const": "invoke", - "type": "string" - }, - "outputs": { - "additionalProperties": { - "type": "string" - }, - "default": {}, - "type": "object" - } - }, - "required": [ - "kind", - "component" - ], - "type": "object" - }, - { - "additionalProperties": false, - "properties": { - "body": { - "$ref": "#/$defs/WorkflowNode" - }, - "carried": { - "default": [], - "items": { - "$ref": "#/$defs/WorkflowLoopCarry" - }, - "type": "array" - }, - "condition": { - "type": "string" - }, - "iteration": { - "anyOf": [ - { - "$ref": "#/$defs/WorkflowLoopIteration" - }, - { - "type": "null" - } - ], - "description": "Optional zero-based induction value, scoped to this loop's body and condition." - }, - "kind": { - "const": "loop", - "type": "string" - }, - "max_iterations": { - "type": "string" - }, - "setup": { - "$ref": "#/$defs/WorkflowNode" - } - }, - "required": [ - "kind", - "setup", - "body", - "condition", - "max_iterations" - ], - "type": "object" - }, - { - "additionalProperties": false, - "properties": { - "cases": { - "additionalProperties": { - "$ref": "#/$defs/WorkflowNode" - }, - "type": "object" - }, - "default": { - "anyOf": [ - { - "$ref": "#/$defs/WorkflowNode" - }, - { - "type": "null" - } - ] - }, - "effects": { - "additionalProperties": { - "$ref": "#/$defs/WorkflowBranchEffectMerge" - }, - "default": {}, - "type": "object" - }, - "kind": { - "const": "branch", - "type": "string" - }, - "outputs": { - "additionalProperties": { - "$ref": "#/$defs/WorkflowBranchOutput" - }, - "default": {}, - "type": "object" - }, - "predicate": { - "type": "string" - } - }, - "required": [ - "kind", - "predicate", - "cases" - ], - "type": "object" + "type": "array", + "uniqueItems": true }, - { - "additionalProperties": false, - "properties": { - "effect": { - "$ref": "#/$defs/EffectTransition" - }, - "effect_name": { - "type": "string" - }, - "kind": { - "const": "emit", - "type": "string" - }, - "mode": { - "$ref": "#/$defs/WorkflowEmitMode" - }, - "output": { - "type": "string" - }, - "value": { - "type": "string" - } + "custom_op_versions": { + "additionalProperties": { + "type": "string" }, - "required": [ - "kind", - "value", - "output", - "mode", - "effect_name", - "effect" - ], + "default": {}, "type": "object" }, - { - "additionalProperties": false, - "properties": { - "device": { - "$ref": "#/$defs/DeviceKind" - }, - "input": { - "type": "string" - }, - "kind": { - "const": "transfer", - "type": "string" - }, - "output": { - "type": "string" - } + "ir_version": { + "type": "string" + }, + "onnx_opsets": { + "additionalProperties": { + "format": "uint32", + "minimum": 0, + "type": "integer" }, - "required": [ - "kind", - "input", - "output", - "device" - ], + "default": {}, "type": "object" } - ] + }, + "required": [ + "ir_version" + ], + "type": "object" + }, + "WorkflowMemoryClass": { + "enum": [ + "default", + "device", + "host", + "pinned" + ], + "type": "string" }, "WorkflowOutput": { "additionalProperties": false, @@ -6399,16 +3858,6 @@ }, "type": "object" }, - "graph": { - "$ref": "#/$defs/WorkflowNode" - }, - "initial_effects": { - "additionalProperties": { - "type": "string" - }, - "default": {}, - "type": "object" - }, "inputs": { "additionalProperties": { "$ref": "#/$defs/WorkflowInput" @@ -6442,18 +3891,28 @@ }, "default": {}, "type": "object" + }, + "steps": { + "items": { + "$ref": "#/$defs/WorkflowStep" + }, + "type": "array" } }, "required": [ "manifest", "components", - "graph" + "steps" ], "type": "object" }, "WorkflowStateCell": { "additionalProperties": false, "properties": { + "class": { + "$ref": "#/$defs/WorkflowStateClass", + "default": "semantic" + }, "contract": { "$ref": "#/$defs/TensorContract" }, @@ -6485,12 +3944,198 @@ ], "type": "object" }, + "WorkflowStateClass": { + "enum": [ + "semantic", + "advisory" + ], + "type": "string" + }, "WorkflowStateScope": { "enum": [ "invocation", "session" ], "type": "string" + }, + "WorkflowStep": { + "oneOf": [ + { + "additionalProperties": false, + "properties": { + "kind": { + "const": "sequence", + "type": "string" + }, + "steps": { + "items": { + "$ref": "#/$defs/WorkflowStep" + }, + "type": "array" + } + }, + "required": [ + "kind", + "steps" + ], + "type": "object" + }, + { + "additionalProperties": false, + "properties": { + "component": { + "type": "string" + }, + "inputs": { + "additionalProperties": { + "type": "string" + }, + "default": {}, + "type": "object" + }, + "kind": { + "const": "invoke", + "type": "string" + }, + "outputs": { + "additionalProperties": { + "type": "string" + }, + "default": {}, + "type": "object" + } + }, + "required": [ + "kind", + "component" + ], + "type": "object" + }, + { + "additionalProperties": false, + "properties": { + "carried": { + "default": [], + "items": { + "$ref": "#/$defs/WorkflowCarry" + }, + "type": "array" + }, + "condition": { + "type": "string" + }, + "iteration": { + "anyOf": [ + { + "$ref": "#/$defs/WorkflowLoopIteration" + }, + { + "type": "null" + } + ] + }, + "kind": { + "const": "loop", + "type": "string" + }, + "max_iterations": { + "type": "string" + }, + "setup": { + "default": [], + "items": { + "$ref": "#/$defs/WorkflowStep" + }, + "type": "array" + }, + "steps": { + "items": { + "$ref": "#/$defs/WorkflowStep" + }, + "type": "array" + } + }, + "required": [ + "kind", + "steps", + "condition", + "max_iterations" + ], + "type": "object" + }, + { + "additionalProperties": false, + "properties": { + "cases": { + "additionalProperties": { + "$ref": "#/$defs/WorkflowStep" + }, + "type": "object" + }, + "default": { + "anyOf": [ + { + "$ref": "#/$defs/WorkflowStep" + }, + { + "type": "null" + } + ] + }, + "kind": { + "const": "branch", + "type": "string" + }, + "outputs": { + "additionalProperties": { + "$ref": "#/$defs/WorkflowBranchOutput" + }, + "default": {}, + "type": "object" + }, + "predicate": { + "type": "string" + } + }, + "required": [ + "kind", + "predicate", + "cases" + ], + "type": "object" + }, + { + "additionalProperties": false, + "properties": { + "kind": { + "const": "emit", + "type": "string" + }, + "mode": { + "$ref": "#/$defs/WorkflowEmitMode" + }, + "output": { + "type": "string" + }, + "valid_length": { + "type": [ + "string", + "null" + ] + }, + "value": { + "type": "string" + } + }, + "required": [ + "kind", + "value", + "output", + "mode" + ], + "type": "object" + } + ] } }, "$id": "https://github.com/onnx/onnx/issues/8184", From 7279dff308884ab1d6f971bbf92152ea931d56fe Mon Sep 17 00:00:00 2001 From: justinchuby Date: Thu, 13 Aug 2026 00:29:06 +0000 Subject: [PATCH 023/151] Simplify public workflow metadata Emit concise structured steps with logical loop carries, infer pure ONNX port contracts from artifacts, and reserve explicit effects for stateful adapters. Document structural once/per-iteration semantics and the internal SSA/effect lowering model.\n\nCo-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: justinchuby --- docs/onnx-genai-workflows.md | 166 ++++++++++++++++++ src/mobius/generation/_policy_components.py | 5 +- .../onnx_genai/auto_export_test.py | 4 +- .../codec_workflow_metadata_test.py | 10 +- .../onnx_genai/inference_metadata.py | 29 --- .../onnx_genai/inference_metadata_test.py | 5 +- .../onnx_genai/workflow_metadata.py | 32 ++-- .../onnx_genai/workflow_metadata_test.py | 11 +- ...ma.json => onnx_genai_c553a16.schema.json} | 5 + 9 files changed, 205 insertions(+), 62 deletions(-) create mode 100644 docs/onnx-genai-workflows.md rename tests/schemas/{onnx_genai_b90d949.schema.json => onnx_genai_c553a16.schema.json} (99%) diff --git a/docs/onnx-genai-workflows.md b/docs/onnx-genai-workflows.md new file mode 100644 index 000000000..a468fdb12 --- /dev/null +++ b/docs/onnx-genai-workflows.md @@ -0,0 +1,166 @@ +# ONNX GenAI workflow metadata + +Mobius emits the concise public workflow source form. The only control primitives are +`sequence`/`steps`, `invoke`, `loop`, `branch`, and `emit`. + +## 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 } +``` + +## Compact examples + +### 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: invoke + component: initialize_latent + inputs: { noise: noise } + outputs: { latent: latent.initial } + - kind: loop + setup: [] + steps: + - kind: invoke + component: denoiser + inputs: { sample: latent, step: diffusion_step } + outputs: { estimate: estimate } + - kind: invoke + component: solver + inputs: { state: latent, estimate: estimate, step: diffusion_step } + outputs: { next_state: latent.next } + condition: continue + max_iterations: num_steps + iteration: + value: diffusion_step + contract: { dtype: int64, rank: 0, shape: [] } + carried: [{ cell: latent, initial: latent.initial, next: latent.next }] + - kind: invoke + component: vae_decoder + inputs: { latent: latent } + outputs: { image: image } + - kind: emit + value: image + output: image + mode: replace +``` + +Latent initialization and VAE decoding run once; denoiser and solver run per iteration. diff --git a/src/mobius/generation/_policy_components.py b/src/mobius/generation/_policy_components.py index 7afdd65b9..40ae4fd93 100644 --- a/src/mobius/generation/_policy_components.py +++ b/src/mobius/generation/_policy_components.py @@ -129,11 +129,12 @@ def _component( role: PolicyRole, graph: ir.Graph, contract: dict[str, object], - *effects: str, + *_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(role, model, contract, effects) + return PolicyComponent(role, model, contract, ()) def _make_graph(name: str) -> tuple[ir.Graph, GraphBuilder]: diff --git a/src/mobius/integrations/onnx_genai/auto_export_test.py b/src/mobius/integrations/onnx_genai/auto_export_test.py index 9ecefb4d0..a8b3dace3 100644 --- a/src/mobius/integrations/onnx_genai/auto_export_test.py +++ b/src/mobius/integrations/onnx_genai/auto_export_test.py @@ -366,7 +366,9 @@ def test_dispatch_diffusion_auto_reads_scheduler_from_source(tmp_path): with open(arts["inference_metadata"]) as handle: meta = yaml.safe_load(handle) components = meta["pipeline"]["workflow"]["components"] - assert components["diffusion_schedule"]["ports"]["outputs"]["schedule"]["shape"] == [16] + assert "ports" not in components["diffusion_schedule"] + schedule = ir.load(out / "policies" / "diffusion_schedule.onnx") + assert list(schedule.graph.outputs[0].shape) == [16] def test_dispatch_vision_multimodal_pipeline(tmp_path): diff --git a/src/mobius/integrations/onnx_genai/codec_workflow_metadata_test.py b/src/mobius/integrations/onnx_genai/codec_workflow_metadata_test.py index 4950a41d9..d47b4aa99 100644 --- a/src/mobius/integrations/onnx_genai/codec_workflow_metadata_test.py +++ b/src/mobius/integrations/onnx_genai/codec_workflow_metadata_test.py @@ -52,13 +52,9 @@ def test_codec_workflow_has_typed_ssa_and_audio_emit(): "rank": 3, "shape": ["batch", 1, "audio_samples"], } - assert workflow["components"]["encoder"]["ports"]["outputs"]["codes"] == { - "dtype": "int64", - "rank": 3, - "shape": ["batch", 16, "frames"], - } - assert workflow["components"]["encoder"]["effects"] == ["codec_encode"] - assert workflow["components"]["decoder"]["effects"] == ["codec_decode"] + assert "ports" not in workflow["components"]["encoder"] + 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"} diff --git a/src/mobius/integrations/onnx_genai/inference_metadata.py b/src/mobius/integrations/onnx_genai/inference_metadata.py index c7235fe94..5b7bdc736 100644 --- a/src/mobius/integrations/onnx_genai/inference_metadata.py +++ b/src/mobius/integrations/onnx_genai/inference_metadata.py @@ -1379,7 +1379,6 @@ def add_policy_components_to_workflow( if not isinstance(workflow, dict): return metadata components = workflow.setdefault("components", {}) - workflow_dtypes = {"fp16": "float16", "bf16": "bfloat16", "fp32": "float32"} def semantic_contract(contract: dict[str, Any]) -> dict[str, Any]: role = str(contract["role"]).replace("_", "-") @@ -1404,39 +1403,11 @@ def semantic_contract(contract: dict[str, Any]) -> dict[str, Any]: return declaration for name, component in policy_components.items(): - model = component.model declaration = { "implementation": { "kind": "onnx", "artifact": f"policies/{name}.onnx", }, - "ports": { - "inputs": { - value.name: { - "dtype": workflow_dtypes.get(_port(value).dtype, _port(value).dtype), - "rank": _port(value).rank, - **( - {"shape": _shape_metadata(_port(value))} - if value.shape is not None - else {} - ), - } - for value in model.graph.inputs - }, - "outputs": { - value.name: { - "dtype": workflow_dtypes.get(_port(value).dtype, _port(value).dtype), - "rank": _port(value).rank, - **( - {"shape": _shape_metadata(_port(value))} - if value.shape is not None - else {} - ), - } - for value in model.graph.outputs - }, - }, - "effects": list(component.effects), } if component.contract: declaration["contract"] = semantic_contract(component.contract) diff --git a/src/mobius/integrations/onnx_genai/inference_metadata_test.py b/src/mobius/integrations/onnx_genai/inference_metadata_test.py index eb9f853b9..4b750b2df 100644 --- a/src/mobius/integrations/onnx_genai/inference_metadata_test.py +++ b/src/mobius/integrations/onnx_genai/inference_metadata_test.py @@ -94,15 +94,14 @@ def test_workflow_policy_components_reference_saved_onnx_artifacts(tmp_path): "kind": "onnx", "artifact": "policies/sample.onnx", } - assert set(component["ports"]["inputs"]) == {"logits"} - assert set(component["ports"]["outputs"]) == {"token"} + assert "ports" not in component assert component["contract"] == { "id": "onnx-genai.token-sampler", "version": "1", "bindings": {"logits": "logits", "token": "token"}, "parameters": {"mode": "greedy"}, } - assert component["effects"] == ["sample"] + assert "effects" not in component assert (tmp_path / component["implementation"]["artifact"]).is_file() diff --git a/src/mobius/integrations/onnx_genai/workflow_metadata.py b/src/mobius/integrations/onnx_genai/workflow_metadata.py index 38494ca2a..6e0503935 100644 --- a/src/mobius/integrations/onnx_genai/workflow_metadata.py +++ b/src/mobius/integrations/onnx_genai/workflow_metadata.py @@ -67,16 +67,8 @@ def _component( *, effects: tuple[str, ...] = (), ) -> dict[str, Any]: - component = { - "implementation": {"kind": "onnx", "artifact": artifact}, - "ports": { - "inputs": {value.name: _contract(value) for value in model.graph.inputs}, - "outputs": {value.name: _contract(value) for value in model.graph.outputs}, - }, - } - if effects: - component["effects"] = list(effects) - return component + del model, effects + return {"implementation": {"kind": "onnx", "artifact": artifact}} def _grammar_adapter_component(action: str) -> dict[str, Any]: @@ -217,20 +209,24 @@ def convert(node: dict[str, Any]) -> dict[str, Any]: 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) result = { "kind": "loop", "setup": setup_steps, "steps": body_steps, "condition": rewrite(node["condition"]), "max_iterations": rewrite(node["max_iterations"]), - "carried": [ - { - "cell": cell_aliases.get(carry["cell"], carry["cell"]), - "initial": rewrite(carry["current"]), - "next": rewrite(carry["body_output"]), - } - for carry in node.get("carried", []) - ], + "carried": carried, } if "iteration" in node: result["iteration"] = node["iteration"] diff --git a/src/mobius/integrations/onnx_genai/workflow_metadata_test.py b/src/mobius/integrations/onnx_genai/workflow_metadata_test.py index ac3ac31b6..89361968d 100644 --- a/src/mobius/integrations/onnx_genai/workflow_metadata_test.py +++ b/src/mobius/integrations/onnx_genai/workflow_metadata_test.py @@ -173,7 +173,7 @@ def test_language_diffusion_rejects_zero_steps(): def test_language_diffusion_matches_pr_828_schema(): schema_path = ( - Path(__file__).parents[4] / "tests" / "schemas" / "onnx_genai_b90d949.schema.json" + Path(__file__).parents[4] / "tests" / "schemas" / "onnx_genai_c553a16.schema.json" ) with schema_path.open(encoding="utf-8") as handle: schema = json.load(handle) @@ -276,12 +276,19 @@ def test_speculative_grammar_and_adaptive_k_use_typed_state_contracts(): assert workflow["state"]["grammar"]["class"] == "semantic" assert workflow["state"]["proposal_k"]["class"] == "advisory" assert workflow["state"]["adaptive_estimates"]["class"] == "advisory" + assert all( + "ports" not in component and "effects" not in component + for component in workflow["components"].values() + if component["implementation"]["kind"] == "onnx" + ) + 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" + assert all("initial" not in carry for carry in workflow["steps"][0]["carried"]) schema_path = ( - Path(__file__).parents[4] / "tests" / "schemas" / "onnx_genai_b90d949.schema.json" + Path(__file__).parents[4] / "tests" / "schemas" / "onnx_genai_c553a16.schema.json" ) with schema_path.open(encoding="utf-8") as handle: jsonschema.validate(instance=metadata, schema=json.load(handle)) diff --git a/tests/schemas/onnx_genai_b90d949.schema.json b/tests/schemas/onnx_genai_c553a16.schema.json similarity index 99% rename from tests/schemas/onnx_genai_b90d949.schema.json rename to tests/schemas/onnx_genai_c553a16.schema.json index 3c7d8a6e5..4978380f4 100644 --- a/tests/schemas/onnx_genai_b90d949.schema.json +++ b/tests/schemas/onnx_genai_c553a16.schema.json @@ -3553,6 +3553,11 @@ "WorkflowComponent": { "additionalProperties": false, "properties": { + "application_overridable": { + "default": false, + "description": "Allow an application to select another declared component with the same\nversioned contract ABI for this invocation.", + "type": "boolean" + }, "contract": { "anyOf": [ { From ceabcaed72788ecbfaa7467c801d89598aed7114 Mon Sep 17 00:00:00 2001 From: justinchuby Date: Thu, 13 Aug 2026 00:46:06 +0000 Subject: [PATCH 024/151] Parameterize workflow token sampling Add request-driven temperature, top-k, top-p, grammar constraints, and counter RNG state to the generic stochastic sampler. Mark sampler components as versioned application overrides and document direct decoder KV recurrence without synthetic update artifacts. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: justinchuby --- docs/onnx-genai-workflows.md | 17 +++ src/mobius/generation/_policy_components.py | 119 ++++++++++++++++-- .../generation/_policy_components_test.py | 37 ++++++ .../integrations/onnx_genai/auto_export.py | 7 +- .../onnx_genai/auto_export_test.py | 39 ++++++ .../onnx_genai/inference_metadata.py | 2 + .../onnx_genai/workflow_metadata.py | 117 ++++++++++++++++- 7 files changed, 322 insertions(+), 16 deletions(-) diff --git a/docs/onnx-genai-workflows.md b/docs/onnx-genai-workflows.md index a468fdb12..29539d7e0 100644 --- a/docs/onnx-genai-workflows.md +++ b/docs/onnx-genai-workflows.md @@ -60,6 +60,23 @@ contract: parameters: { mode: greedy } ``` +### 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`, `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. 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. + ## Compact examples ### Decoder diff --git a/src/mobius/generation/_policy_components.py b/src/mobius/generation/_policy_components.py index 40ae4fd93..c898b5172 100644 --- a/src/mobius/generation/_policy_components.py +++ b/src/mobius/generation/_policy_components.py @@ -961,16 +961,19 @@ def ewma(old, sample): def build_seeded_categorical_sampler() -> PolicyComponent: - """Build deterministic categorical sampling with explicit seed and offset. + """Build request-parameterized categorical sampling with explicit RNG state. - Threefry is counter based: the same ``(seed, offset, logits, temperature)`` - inputs always produce the same token. The updated offset is - an explicit output, so no random or hidden mutable state exists in the graph. + Threefry is counter based: identical tensor inputs produce the same token. + Temperature, top-k, top-p, and grammar constraints 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"]) + temperature = builder.input("temperature", ir.DataType.FLOAT, [1]) + top_k = builder.input("top_k", ir.DataType.INT64, [1]) + top_p = builder.input("top_p", ir.DataType.FLOAT, [1]) + grammar_mask = builder.input("grammar_mask", ir.DataType.BOOL, ["batch", "vocabulary"]) seed = builder.input("seed", ir.DataType.INT64, ["batch"]) offset = builder.input("offset", ir.DataType.INT64, ["batch"]) @@ -1020,11 +1023,90 @@ def build_seeded_categorical_sampler() -> PolicyComponent: ) uniform = op.Cast(uniform, to=ir.DataType.FLOAT) - scaled_logits = op.Div( - logits, - op.Unsqueeze(temperature, op.Constant(value_ints=[-1])), + blocked = op.CastLike(op.Constant(value_float=-3.4028235e38), logits) + constrained_logits = op.Where(grammar_mask, logits, blocked) + safe_temperature = op.Max(temperature, op.Constant(value_float=1e-6)) + scaled_logits = op.Div(constrained_logits, safe_temperature) + + 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, + effective_k, + axis=-1, + largest=1, + sorted=0, + _outputs=2, + ) + top_k_mask = op.Greater( + op.ScatterElements( + op.ConstantOfShape( + op.Shape(logits), + value=ir.tensor([0], dtype=ir.DataType.INT64), + ), + top_indices, + op.ConstantOfShape( + op.Shape(top_indices), + value=ir.tensor([1], dtype=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.Clip( + top_p, + op.Constant(value_float=1e-6), + op.Constant(value_float=1.0), + ) + 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), + ), ) - probabilities = op.Softmax(scaled_logits, axis=-1) axis = op.Constant(value_int=-1) cumulative = op.CumSum(probabilities, axis) uniform = op.Unsqueeze(uniform, op.Constant(value_ints=[-1])) @@ -1034,6 +1116,16 @@ def build_seeded_categorical_sampler() -> PolicyComponent: axis=-1, keepdims=0, ) + has_allowed_token = op.ReduceMax( + op.Cast(grammar_mask, to=ir.DataType.INT64), + axes=[-1], + keepdims=0, + ) + token_ids = op.Where( + op.Greater(has_allowed_token, op.Constant(value_int=0)), + token_ids, + op.Constant(value_int=-1), + ) next_offset = op.Add(offset, op.Constant(value_int=1)) builder.add_output(token_ids, "token") builder.add_output(next_offset, "next_offset") @@ -1046,10 +1138,13 @@ def build_seeded_categorical_sampler() -> PolicyComponent: "logits": "logits", "token": "token", "temperature": "temperature", + "top_k": "top_k", + "top_p": "top_p", + "grammar_mask": "grammar_mask", "rng": { - "seed": "seed", - "offset": "offset", - "next_offset": "next_offset", + "rng_seed": "seed", + "rng_offset": "offset", + "rng_next_offset": "next_offset", }, "effect": "rng", }, diff --git a/src/mobius/generation/_policy_components_test.py b/src/mobius/generation/_policy_components_test.py index 48a9c663f..534b8c0c7 100644 --- a/src/mobius/generation/_policy_components_test.py +++ b/src/mobius/generation/_policy_components_test.py @@ -334,6 +334,9 @@ def test_seeded_sampler_is_counter_based_and_reproducible(tmp_path): 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), + "grammar_mask": np.array([[True, True, True, True]], np.bool_), "seed": np.array([7], np.int64), "offset": np.array([11], np.int64), } @@ -343,6 +346,40 @@ def test_seeded_sampler_is_counter_based_and_reproducible(tmp_path): np.testing.assert_array_equal(first[1], [12]) +def test_seeded_sampler_applies_request_top_k_and_grammar_mask(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), + "grammar_mask": np.array([[False, True, True, True]], np.bool_), + "seed": np.array([17], np.int64), + "offset": np.array([0], np.int64), + }, + ) + np.testing.assert_array_equal(token, [1]) + + +def test_seeded_sampler_rejects_empty_grammar_vocabulary(tmp_path): + (token, _) = _run( + build_seeded_categorical_sampler(), + tmp_path, + { + "logits": np.array([[1.0, 2.0, 3.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), + "grammar_mask": np.array([[False, False, False]], np.bool_), + "seed": np.array([1], np.int64), + "offset": np.array([0], np.int64), + }, + ) + np.testing.assert_array_equal(token, [-1]) + + def test_eos_termination_runtime(tmp_path): (terminated,) = _run( build_eos_termination(), diff --git a/src/mobius/integrations/onnx_genai/auto_export.py b/src/mobius/integrations/onnx_genai/auto_export.py index 5ed95f882..37152ac15 100644 --- a/src/mobius/integrations/onnx_genai/auto_export.py +++ b/src/mobius/integrations/onnx_genai/auto_export.py @@ -674,7 +674,12 @@ def write_onnx_genai_config( "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) + path = write_decoder_workflow_metadata( + pkg, + output_dir, + resolved_config, + sampler=str(getattr(resolved_config, "workflow_sampler", "greedy")), + ) artifacts = {"inference_metadata": path} tokenizer_path = _write_hf_tokenizer(output_dir, source, revision=revision) if tokenizer_path is not None: diff --git a/src/mobius/integrations/onnx_genai/auto_export_test.py b/src/mobius/integrations/onnx_genai/auto_export_test.py index a8b3dace3..c89d87ad4 100644 --- a/src/mobius/integrations/onnx_genai/auto_export_test.py +++ b/src/mobius/integrations/onnx_genai/auto_export_test.py @@ -21,6 +21,9 @@ _model, _value, ) +from mobius.integrations.onnx_genai.workflow_metadata import ( + build_decoder_workflow_metadata, +) @dataclasses.dataclass @@ -187,6 +190,42 @@ def test_dispatch_decoder(tmp_path): 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"]["bindings"] == { + "logits": "logits", + "token": "token", + "temperature": "temperature", + "top_k": "top_k", + "top_p": "top_p", + "grammar_mask": "grammar_mask", + "rng_seed": "seed", + "rng_offset": "offset", + "rng_next_offset": "next_offset", + } + 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", + "grammar_mask": "request.grammar_mask", + "seed": "request.seed", + "offset": "rng_offset", + } + assert workflow["state"]["rng_offset"]["class"] == "semantic" + assert workflow["state"]["rng_offset"]["initializer"] == "request.rng_offset" + assert not any("kv_update" in name for name in workflow["components"]) + + def test_dispatch_language_diffusion(tmp_path): package = ModelPackage( { diff --git a/src/mobius/integrations/onnx_genai/inference_metadata.py b/src/mobius/integrations/onnx_genai/inference_metadata.py index 5b7bdc736..f3e105f41 100644 --- a/src/mobius/integrations/onnx_genai/inference_metadata.py +++ b/src/mobius/integrations/onnx_genai/inference_metadata.py @@ -1411,6 +1411,8 @@ def semantic_contract(contract: dict[str, Any]) -> dict[str, Any]: } if component.contract: declaration["contract"] = semantic_contract(component.contract) + if component.contract.get("role") == "token_sampler": + declaration["application_overridable"] = True components[name] = declaration return metadata diff --git a/src/mobius/integrations/onnx_genai/workflow_metadata.py b/src/mobius/integrations/onnx_genai/workflow_metadata.py index 6e0503935..49fac988f 100644 --- a/src/mobius/integrations/onnx_genai/workflow_metadata.py +++ b/src/mobius/integrations/onnx_genai/workflow_metadata.py @@ -3815,6 +3815,78 @@ def build_decoder_workflow_metadata( }, } ) + stochastic_sampler = sampler != "greedy" + if stochastic_sampler: + workflow_inputs.update( + { + "request.temperature": { + "contract": {"dtype": "float32", "rank": 1, "shape": [1]}, + "role": { + "kind": "runtime", + "version": "1.0", + "role": "sampling_temperature", + }, + "source": { + "kind": "request", + "field": "sampling_temperature", + }, + "required": False, + "default": 1.0, + }, + "request.top_k": { + "contract": {"dtype": "int64", "rank": 1, "shape": [1]}, + "role": { + "kind": "runtime", + "version": "1.0", + "role": "sampling_top_k", + }, + "source": {"kind": "request", "field": "sampling_top_k"}, + "required": False, + "default": 0, + }, + "request.top_p": { + "contract": {"dtype": "float32", "rank": 1, "shape": [1]}, + "role": { + "kind": "runtime", + "version": "1.0", + "role": "sampling_top_p", + }, + "source": {"kind": "request", "field": "sampling_top_p"}, + "required": False, + "default": 1.0, + }, + "request.seed": { + "contract": batch_int, + "role": { + "kind": "runtime", + "version": "1.0", + "role": "seed", + }, + "source": {"kind": "request", "field": "seed"}, + "required": True, + }, + "request.grammar_mask": { + "contract": { + "dtype": "bool", + "rank": 2, + "shape": [ + batch_dimension, + _contract(logits_output)["shape"][-1], + ], + }, + "role": {"kind": "opaque"}, + "source": {"kind": "application", "name": "grammar_mask"}, + "required": True, + }, + "request.rng_offset": { + "contract": batch_int, + "role": {"kind": "opaque"}, + "source": {"kind": "application", "name": "rng_offset"}, + "required": False, + "default": 0, + }, + } + ) for value in inputs: if value is token_input: @@ -3905,6 +3977,26 @@ def build_decoder_workflow_metadata( }, ] ) + if stochastic_sampler: + state["rng_offset"] = { + "contract": batch_int, + "scope": "invocation", + "class": "semantic", + "initializer": "request.rng_offset", + "recurrence": {"kind": "invariant"}, + } + initial_effects["state:rng_offset"] = "state:rng_offset.0" + carried.append( + { + "cell": "rng_offset", + "current": "request.rng_offset", + "body_input": "state.rng_offset.body", + "body_output": "sample.next_offset", + "next": "state.rng_offset.final", + "read_effect": _effect("state:rng_offset.0", "state:rng_offset.read"), + "write_effect": _effect("state:rng_offset.read", "state:rng_offset.1"), + } + ) decoder_state_specs = { "attention_mask": ( { @@ -4025,8 +4117,25 @@ def build_decoder_workflow_metadata( "nodes": [ _invoke( "token_sampler", - {"logits": "state.logits.body"}, - {"token": "sample.body"}, + { + "logits": "state.logits.body", + **( + { + "temperature": "request.temperature", + "top_k": "request.top_k", + "top_p": "request.top_p", + "grammar_mask": "request.grammar_mask", + "seed": "request.seed", + "offset": "state.rng_offset.body", + } + if stochastic_sampler + else {} + ), + }, + { + "token": "sample.body", + **({"next_offset": "sample.next_offset"} if stochastic_sampler else {}), + }, {"sample": _effect("sample.0", "sample.1")}, ), _invoke( @@ -4422,10 +4531,12 @@ 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) + metadata = build_decoder_workflow_metadata(pkg, config, sampler=sampler) pkg.save_policy_components(output_dir) path = os.path.join(output_dir, "inference_metadata.yaml") with open(path, "w", encoding="utf-8") as handle: From 61e3fa7331e005fabd90d71b6bf2750de1957e0c Mon Sep 17 00:00:00 2001 From: justinchuby Date: Thu, 13 Aug 2026 00:55:31 +0000 Subject: [PATCH 025/151] Add capture-friendly min-p sampling Keep min-p as a request-provided tensor and filter temperature-scaled logits without changing sampler shapes. Document planner-derived execution islands, per-island CUDA Graph eligibility, diagnostics, and fallback while preserving modular workflow components. Refresh the generated ONNX GenAI schema fixture to b2157a2. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: justinchuby --- docs/onnx-genai-workflows.md | 42 ++++++++++++++++--- src/mobius/generation/_policy_components.py | 16 +++++++ .../generation/_policy_components_test.py | 22 ++++++++++ .../onnx_genai/auto_export_test.py | 2 + .../onnx_genai/workflow_metadata.py | 15 +++++++ .../onnx_genai/workflow_metadata_test.py | 4 +- ...ma.json => onnx_genai_b2157a2.schema.json} | 1 + 7 files changed, 95 insertions(+), 7 deletions(-) rename tests/schemas/{onnx_genai_c553a16.schema.json => onnx_genai_b2157a2.schema.json} (99%) diff --git a/docs/onnx-genai-workflows.md b/docs/onnx-genai-workflows.md index 29539d7e0..11005ac5f 100644 --- a/docs/onnx-genai-workflows.md +++ b/docs/onnx-genai-workflows.md @@ -60,6 +60,34 @@ contract: 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. + ### Decoder KV boundary Normal decode carries each decoder `present` tensor directly to the corresponding @@ -71,11 +99,15 @@ rollback, or format conversion. ### Request-parameterized sampling -The stochastic sampler ABI accepts `temperature`, `top_k`, `top_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. 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. +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. ## Compact examples diff --git a/src/mobius/generation/_policy_components.py b/src/mobius/generation/_policy_components.py index c898b5172..1f271386c 100644 --- a/src/mobius/generation/_policy_components.py +++ b/src/mobius/generation/_policy_components.py @@ -973,6 +973,7 @@ def build_seeded_categorical_sampler() -> PolicyComponent: temperature = builder.input("temperature", ir.DataType.FLOAT, [1]) top_k = builder.input("top_k", ir.DataType.INT64, [1]) top_p = builder.input("top_p", ir.DataType.FLOAT, [1]) + min_p = builder.input("min_p", ir.DataType.FLOAT, [1]) grammar_mask = builder.input("grammar_mask", ir.DataType.BOOL, ["batch", "vocabulary"]) seed = builder.input("seed", ir.DataType.INT64, ["batch"]) offset = builder.input("offset", ir.DataType.INT64, ["batch"]) @@ -1027,6 +1028,20 @@ def build_seeded_categorical_sampler() -> PolicyComponent: constrained_logits = op.Where(grammar_mask, logits, blocked) safe_temperature = op.Max(temperature, op.Constant(value_float=1e-6)) scaled_logits = op.Div(constrained_logits, safe_temperature) + safe_min_p = op.Clip( + min_p, + op.Constant(value_float=1e-20), + op.Constant(value_float=1.0), + ) + min_p_threshold = op.Add( + op.ReduceMax(scaled_logits, axes=[-1], keepdims=1), + op.Log(safe_min_p), + ) + min_p_mask = op.Or( + op.LessOrEqual(min_p, op.Constant(value_float=0.0)), + 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( @@ -1140,6 +1155,7 @@ def build_seeded_categorical_sampler() -> PolicyComponent: "temperature": "temperature", "top_k": "top_k", "top_p": "top_p", + "min_p": "min_p", "grammar_mask": "grammar_mask", "rng": { "rng_seed": "seed", diff --git a/src/mobius/generation/_policy_components_test.py b/src/mobius/generation/_policy_components_test.py index 534b8c0c7..a4ba5e80a 100644 --- a/src/mobius/generation/_policy_components_test.py +++ b/src/mobius/generation/_policy_components_test.py @@ -336,6 +336,7 @@ def test_seeded_sampler_is_counter_based_and_reproducible(tmp_path): "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), "grammar_mask": np.array([[True, True, True, True]], np.bool_), "seed": np.array([7], np.int64), "offset": np.array([11], np.int64), @@ -355,6 +356,7 @@ def test_seeded_sampler_applies_request_top_k_and_grammar_mask(tmp_path): "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), "grammar_mask": np.array([[False, True, True, True]], np.bool_), "seed": np.array([17], np.int64), "offset": np.array([0], np.int64), @@ -372,6 +374,7 @@ def test_seeded_sampler_rejects_empty_grammar_vocabulary(tmp_path): "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), "grammar_mask": np.array([[False, False, False]], np.bool_), "seed": np.array([1], np.int64), "offset": np.array([0], np.int64), @@ -380,6 +383,25 @@ def test_seeded_sampler_rejects_empty_grammar_vocabulary(tmp_path): np.testing.assert_array_equal(token, [-1]) +def test_seeded_sampler_applies_request_min_p_in_logit_space(tmp_path): + (token, next_offset) = _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), + "grammar_mask": np.array([[True, True, True]], np.bool_), + "seed": np.array([23], np.int64), + "offset": np.array([4], np.int64), + }, + ) + np.testing.assert_array_equal(token, [0]) + np.testing.assert_array_equal(next_offset, [5]) + + def test_eos_termination_runtime(tmp_path): (terminated,) = _run( build_eos_termination(), diff --git a/src/mobius/integrations/onnx_genai/auto_export_test.py b/src/mobius/integrations/onnx_genai/auto_export_test.py index c89d87ad4..1e0fbcd3e 100644 --- a/src/mobius/integrations/onnx_genai/auto_export_test.py +++ b/src/mobius/integrations/onnx_genai/auto_export_test.py @@ -202,6 +202,7 @@ def test_seeded_decoder_sampler_uses_request_controls_and_direct_kv_carry(): "temperature": "temperature", "top_k": "top_k", "top_p": "top_p", + "min_p": "min_p", "grammar_mask": "grammar_mask", "rng_seed": "seed", "rng_offset": "offset", @@ -217,6 +218,7 @@ def test_seeded_decoder_sampler_uses_request_controls_and_direct_kv_carry(): "temperature": "request.temperature", "top_k": "request.top_k", "top_p": "request.top_p", + "min_p": "request.min_p", "grammar_mask": "request.grammar_mask", "seed": "request.seed", "offset": "rng_offset", diff --git a/src/mobius/integrations/onnx_genai/workflow_metadata.py b/src/mobius/integrations/onnx_genai/workflow_metadata.py index 49fac988f..5a105a8e0 100644 --- a/src/mobius/integrations/onnx_genai/workflow_metadata.py +++ b/src/mobius/integrations/onnx_genai/workflow_metadata.py @@ -3855,6 +3855,20 @@ def build_decoder_workflow_metadata( "required": False, "default": 1.0, }, + "request.min_p": { + "contract": {"dtype": "float32", "rank": 1, "shape": [1]}, + "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": { @@ -4124,6 +4138,7 @@ def build_decoder_workflow_metadata( "temperature": "request.temperature", "top_k": "request.top_k", "top_p": "request.top_p", + "min_p": "request.min_p", "grammar_mask": "request.grammar_mask", "seed": "request.seed", "offset": "state.rng_offset.body", diff --git a/src/mobius/integrations/onnx_genai/workflow_metadata_test.py b/src/mobius/integrations/onnx_genai/workflow_metadata_test.py index 89361968d..edf1c6a7c 100644 --- a/src/mobius/integrations/onnx_genai/workflow_metadata_test.py +++ b/src/mobius/integrations/onnx_genai/workflow_metadata_test.py @@ -173,7 +173,7 @@ def test_language_diffusion_rejects_zero_steps(): def test_language_diffusion_matches_pr_828_schema(): schema_path = ( - Path(__file__).parents[4] / "tests" / "schemas" / "onnx_genai_c553a16.schema.json" + Path(__file__).parents[4] / "tests" / "schemas" / "onnx_genai_b2157a2.schema.json" ) with schema_path.open(encoding="utf-8") as handle: schema = json.load(handle) @@ -288,7 +288,7 @@ def test_speculative_grammar_and_adaptive_k_use_typed_state_contracts(): assert proposer["inputs"]["proposal_budget"] == "proposal_k" assert all("initial" not in carry for carry in workflow["steps"][0]["carried"]) schema_path = ( - Path(__file__).parents[4] / "tests" / "schemas" / "onnx_genai_c553a16.schema.json" + Path(__file__).parents[4] / "tests" / "schemas" / "onnx_genai_b2157a2.schema.json" ) with schema_path.open(encoding="utf-8") as handle: jsonschema.validate(instance=metadata, schema=json.load(handle)) diff --git a/tests/schemas/onnx_genai_c553a16.schema.json b/tests/schemas/onnx_genai_b2157a2.schema.json similarity index 99% rename from tests/schemas/onnx_genai_c553a16.schema.json rename to tests/schemas/onnx_genai_b2157a2.schema.json index 4978380f4..91f345fe7 100644 --- a/tests/schemas/onnx_genai_c553a16.schema.json +++ b/tests/schemas/onnx_genai_b2157a2.schema.json @@ -2382,6 +2382,7 @@ "sampling_temperature", "sampling_top_k", "sampling_top_p", + "sampling_min_p", "constraint", "session_id" ], From 50a2144e24031f071829b4709d2ae59ec8ee82ec Mon Sep 17 00:00:00 2001 From: justinchuby Date: Thu, 13 Aug 2026 01:04:44 +0000 Subject: [PATCH 026/151] Gate workflow performance against native Add strict paired-run identity validation and acceptance checks for throughput, TTFT, memory, transfers, synchronization, session/kernel boundaries, device residency, and required island capture/replay. Document required decoder, min-p, speculative, and grammar-boundary scenarios plus the measured upstream CUDA gaps that still block readiness. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: justinchuby --- docs/onnx-genai-performance-conformance.md | 103 ++++++++ docs/onnx-genai-workflows.md | 3 + scripts/compare_workflow_performance.py | 43 ++++ .../integrations/onnx_genai/performance.py | 221 ++++++++++++++++++ .../onnx_genai/performance_test.py | 141 +++++++++++ 5 files changed, 511 insertions(+) create mode 100644 docs/onnx-genai-performance-conformance.md create mode 100755 scripts/compare_workflow_performance.py create mode 100644 src/mobius/integrations/onnx_genai/performance.py create mode 100644 src/mobius/integrations/onnx_genai/performance_test.py diff --git a/docs/onnx-genai-performance-conformance.md b/docs/onnx-genai-performance-conformance.md new file mode 100644 index 000000000..5dff1217c --- /dev/null +++ b/docs/onnx-genai-performance-conformance.md @@ -0,0 +1,103 @@ +# 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. + +## 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 index 11005ac5f..af5a80e78 100644 --- a/docs/onnx-genai-workflows.md +++ b/docs/onnx-genai-workflows.md @@ -88,6 +88,9 @@ to ordinary island execution without changing workflow semantics. Diagnostics sh 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 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/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) From 045f20a00d2599b8486f06ce356fba9d49b28a4c Mon Sep 17 00:00:00 2001 From: justinchuby Date: Thu, 13 Aug 2026 01:07:46 +0000 Subject: [PATCH 027/151] Make policy contracts data driven Replace the closed policy-role enum with versioned namespaced contract IDs embedded in each ONNX artifact. Workflow declarations now derive contract identity and version from artifact metadata rather than a producer-side role registry. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: justinchuby --- src/mobius/_model_package_test.py | 4 +- src/mobius/generation/__init__.py | 2 - src/mobius/generation/_policy_components.py | 105 ++++++++---------- .../generation/_policy_components_test.py | 15 ++- .../onnx_genai/inference_metadata.py | 11 +- 5 files changed, 61 insertions(+), 76 deletions(-) diff --git a/src/mobius/_model_package_test.py b/src/mobius/_model_package_test.py index 831f09309..169b856c1 100644 --- a/src/mobius/_model_package_test.py +++ b/src/mobius/_model_package_test.py @@ -17,7 +17,7 @@ from mobius._configs import VisionConfig from mobius._model_package import ModelPackage, _make_progress_callback from mobius._testing import make_config -from mobius.generation import PolicyRole, build_greedy_sampler +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 @@ -373,7 +373,7 @@ def test_policy_components_roundtrip(self, tmp_path): assert (tmp_path / "policies" / "sample.onnx").exists() loaded = ModelPackage.load(str(tmp_path)) - assert loaded.policy_components["sample"].role is PolicyRole.TOKEN_SAMPLER + assert loaded.policy_components["sample"].contract_id == "onnx-genai.token-sampler@1" class TestModelPackageApplyWeights: diff --git a/src/mobius/generation/__init__.py b/src/mobius/generation/__init__.py index 715b0a006..bf6ca077c 100644 --- a/src/mobius/generation/__init__.py +++ b/src/mobius/generation/__init__.py @@ -8,7 +8,6 @@ from mobius.generation._policy_components import ( PolicyCapabilities, PolicyComponent, - PolicyRole, attach_policy_components, build_adaptive_k_policy, build_batch_minimum, @@ -48,7 +47,6 @@ __all__ = [ "PolicyComponent", "PolicyCapabilities", - "PolicyRole", "attach_policy_components", "build_adaptive_k_policy", "build_batch_minimum", diff --git a/src/mobius/generation/_policy_components.py b/src/mobius/generation/_policy_components.py index 1f271386c..c26870ce9 100644 --- a/src/mobius/generation/_policy_components.py +++ b/src/mobius/generation/_policy_components.py @@ -12,7 +12,6 @@ import json from dataclasses import dataclass -from enum import StrEnum from typing import Protocol import onnx_ir as ir @@ -20,48 +19,36 @@ from mobius._constants import OPSET_VERSION -_POLICY_ROLE_METADATA = "mobius.generation.policy_role" +_POLICY_CONTRACT_ID_METADATA = "mobius.generation.policy_contract_id" _POLICY_CONTRACT_METADATA = "mobius.generation.policy_contract" _POLICY_EFFECTS_METADATA = "mobius.generation.policy_effects" -class PolicyRole(StrEnum): - """Architecture-neutral role performed by a policy component.""" - - TOKEN_SAMPLER = "token_sampler" - TERMINATION = "termination_predicate" - SOLVER_STEP = "solver_step" - MASKED_UPDATE = "masked_update" - SPECULATIVE_ACCEPTANCE = "speculative_verifier" - GRAMMAR_GUIDANCE = "grammar_guidance" - ADAPTIVE_K = "adaptive_proposal_budget" - STATE_UPDATE = "state_update" - AUXILIARY = "auxiliary" - - @dataclass(frozen=True) class PolicyComponent: - """A named role and its executable ONNX model.""" + """A versioned semantic contract and its executable ONNX model.""" - role: PolicyRole + contract_id: str model: ir.Model contract: dict[str, object] effects: tuple[str, ...] def __post_init__(self) -> None: - self.model.graph.metadata_props[_POLICY_ROLE_METADATA] = self.role.value + 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 role metadata embedded in its ONNX graph.""" - role = model.graph.metadata_props.get(_POLICY_ROLE_METADATA) - if role is None: - raise ValueError("ONNX policy component is missing its Mobius policy role") + """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(PolicyRole(role), model, contract, effects) + return cls(contract_id, model, contract, effects) @dataclass(frozen=True) @@ -126,7 +113,7 @@ def attach_policy_components( def _component( - role: PolicyRole, + contract_id: str, graph: ir.Graph, contract: dict[str, object], *_effects: str, @@ -134,7 +121,7 @@ def _component( # 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(role, model, contract, ()) + return PolicyComponent(contract_id, model, contract, ()) def _make_graph(name: str) -> tuple[ir.Graph, GraphBuilder]: @@ -159,7 +146,7 @@ def build_greedy_sampler(*, effect: str = "sample") -> PolicyComponent: token_ids = builder.op.ArgMax(logits, axis=-1, keepdims=0) builder.add_output(token_ids, "token") return _component( - PolicyRole.TOKEN_SAMPLER, + "onnx-genai.token-sampler@1", graph, { "role": "token_sampler", @@ -183,7 +170,7 @@ def build_last_token_logits() -> PolicyComponent: selected = builder.op.Gather(logits, builder.op.Constant(value_int=-1), axis=1) selected.shape = ir.Shape(["batch", "vocabulary"]) builder.add_output(selected, "last_logits") - return _component(PolicyRole.AUXILIARY, graph, {}) + return _component("mobius.policy.auxiliary@1", graph, {}) def build_boolean_not() -> PolicyComponent: @@ -197,7 +184,7 @@ def build_boolean_not() -> PolicyComponent: continued = builder.op.Equal(any_done, builder.op.Constant(value_int=0)) continued.shape = ir.Shape([1]) builder.add_output(continued, "continue") - return _component(PolicyRole.AUXILIARY, graph, {}) + return _component("mobius.policy.auxiliary@1", graph, {}) def build_integer_increment() -> PolicyComponent: @@ -207,7 +194,7 @@ def build_integer_increment() -> PolicyComponent: next_value = builder.op.Add(value, builder.op.Constant(value_int=1)) next_value.shape = ir.Shape(["batch"]) builder.add_output(next_value, "next_value") - return _component(PolicyRole.AUXILIARY, graph, {}) + return _component("mobius.policy.auxiliary@1", graph, {}) def build_integer_minimum() -> PolicyComponent: @@ -218,7 +205,7 @@ def build_integer_minimum() -> PolicyComponent: minimum = builder.op.Min(left, right) minimum.shape = ir.Shape(["batch"]) builder.add_output(minimum, "minimum") - return _component(PolicyRole.AUXILIARY, graph, {}) + return _component("mobius.policy.auxiliary@1", graph, {}) def build_batch_minimum() -> PolicyComponent: @@ -228,7 +215,7 @@ def build_batch_minimum() -> PolicyComponent: minimum = builder.op.ReduceMin(values, keepdims=1) minimum.shape = ir.Shape([1]) builder.add_output(minimum, "minimum") - return _component(PolicyRole.AUXILIARY, graph, {}) + return _component("mobius.policy.auxiliary@1", graph, {}) def build_proposal_metrics() -> PolicyComponent: @@ -244,7 +231,7 @@ def build_proposal_metrics() -> PolicyComponent: filled.shape = ir.Shape(["batch"]) builder.add_output(length, "evaluated") builder.add_output(filled, "filled_proposal_budget") - return _component(PolicyRole.AUXILIARY, graph, {}) + return _component("mobius.policy.auxiliary@1", graph, {}) def build_sequence_length() -> PolicyComponent: @@ -258,7 +245,7 @@ def build_sequence_length() -> PolicyComponent: ) length.shape = ir.Shape(["batch"]) builder.add_output(length, "length") - return _component(PolicyRole.AUXILIARY, graph, {}) + return _component("mobius.policy.auxiliary@1", graph, {}) def build_iteration_cast(dtype: ir.DataType) -> PolicyComponent: @@ -268,7 +255,7 @@ def build_iteration_cast(dtype: ir.DataType) -> PolicyComponent: timestep = builder.op.Cast(iteration, to=dtype) timestep.shape = ir.Shape(["batch"]) builder.add_output(timestep, "timestep") - return _component(PolicyRole.AUXILIARY, graph, {}) + return _component("mobius.policy.auxiliary@1", graph, {}) def build_schedule_constant(values: list[float]) -> PolicyComponent: @@ -281,7 +268,7 @@ def build_schedule_constant(values: list[float]) -> PolicyComponent: ) schedule.shape = ir.Shape([len(values)]) builder.add_output(schedule, "schedule") - return _component(PolicyRole.AUXILIARY, graph, {}) + return _component("mobius.policy.auxiliary@1", graph, {}) def build_schedule_lookup(dtype: ir.DataType) -> PolicyComponent: @@ -293,7 +280,7 @@ def build_schedule_lookup(dtype: ir.DataType) -> PolicyComponent: timestep = op.Cast(op.Gather(schedule, step, axis=0), to=dtype) timestep.shape = ir.Shape(["batch"]) builder.add_output(timestep, "timestep") - return _component(PolicyRole.AUXILIARY, graph, {}) + return _component("mobius.policy.auxiliary@1", graph, {}) def build_tts_state_initializer(num_code_groups: int) -> PolicyComponent: @@ -320,7 +307,7 @@ def build_tts_state_initializer(num_code_groups: int) -> PolicyComponent: builder.add_output(frame, "frame_codes") builder.add_output(token_slot, "token_slot") builder.add_output(history, "code_history") - return _component(PolicyRole.AUXILIARY, graph, {}) + return _component("mobius.policy.auxiliary@1", graph, {}) def build_tts_decoder_state_initializer( @@ -429,7 +416,7 @@ def build_tts_decoder_state_initializer( ) empty.shape = value.shape builder.add_output(empty, name) - return _component(PolicyRole.AUXILIARY, graph, {}) + return _component("mobius.policy.auxiliary@1", graph, {}) def build_tts_decoder_step_update( @@ -462,7 +449,7 @@ def build_tts_decoder_step_update( next_position.shape = position.shape builder.add_output(next_attention, "next_attention_mask") builder.add_output(next_position, "next_position_ids") - return _component(PolicyRole.AUXILIARY, graph, {}) + return _component("mobius.policy.auxiliary@1", graph, {}) def build_code_frame_update( @@ -492,7 +479,7 @@ def build_code_frame_update( ) updated.shape = frame.shape builder.add_output(updated, "next_frame") - return _component(PolicyRole.AUXILIARY, graph, {}) + return _component("mobius.policy.auxiliary@1", graph, {}) def build_code_history_append(num_code_groups: int) -> PolicyComponent: @@ -512,7 +499,7 @@ def build_code_history_append(num_code_groups: int) -> PolicyComponent: ) next_history.shape = ir.Shape(["batch", "frames + 1", num_code_groups]) builder.add_output(next_history, "next_history") - return _component(PolicyRole.AUXILIARY, graph, {}) + return _component("mobius.policy.auxiliary@1", graph, {}) def build_codec_layout_transpose(num_code_groups: int) -> PolicyComponent: @@ -526,7 +513,7 @@ def build_codec_layout_transpose(num_code_groups: int) -> PolicyComponent: 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(PolicyRole.AUXILIARY, graph, {}) + return _component("mobius.policy.auxiliary@1", graph, {}) def build_model_token_cast(dtype: ir.DataType) -> PolicyComponent: @@ -536,7 +523,7 @@ def build_model_token_cast(dtype: ir.DataType) -> PolicyComponent: model_token = builder.op.Cast(token, to=dtype) model_token.shape = ir.Shape(["batch", 1]) builder.add_output(model_token, "model_token") - return _component(PolicyRole.AUXILIARY, graph, {}) + return _component("mobius.policy.auxiliary@1", graph, {}) def build_decoder_state_initializer( @@ -645,7 +632,7 @@ def build_decoder_state_initializer( empty.shape = value.shape builder.add_output(empty, name) - return _component(PolicyRole.AUXILIARY, graph, {}) + return _component("mobius.policy.auxiliary@1", graph, {}) def build_decoder_step_update( @@ -676,7 +663,7 @@ def build_decoder_step_update( next_position = op.Add(position, op.CastLike(op.Constant(value_int=1), position)) next_position.shape = ir.Shape(["batch", 1]) builder.add_output(next_position, "next_position_ids") - return _component(PolicyRole.AUXILIARY, graph, {}) + return _component("mobius.policy.auxiliary@1", graph, {}) def build_grammar_logits_processor() -> PolicyComponent: @@ -697,7 +684,7 @@ def build_grammar_logits_processor() -> PolicyComponent: ) token.shape = ir.Shape(["batch", 1]) builder.add_output(token, "token") - return _component(PolicyRole.GRAMMAR_GUIDANCE, graph, {}) + return _component("onnx-genai.grammar-guidance@1", graph, {}) def build_adaptive_k_policy(*, max_k: int = 16, min_k: int = 1) -> PolicyComponent: @@ -940,7 +927,7 @@ def ewma(old, sample): builder.add_output(next_k, "next_k") builder.add_output(next_estimates, "next_estimates") return _component( - PolicyRole.ADAPTIVE_K, + "onnx-genai.adaptive-proposal-budget@1", graph, { "role": "adaptive_proposal_budget", @@ -1145,7 +1132,7 @@ def build_seeded_categorical_sampler() -> PolicyComponent: builder.add_output(token_ids, "token") builder.add_output(next_offset, "next_offset") return _component( - PolicyRole.TOKEN_SAMPLER, + "onnx-genai.token-sampler@1", graph, { "role": "token_sampler", @@ -1192,7 +1179,7 @@ def build_eos_termination() -> PolicyComponent: done = op.Or(hit_eos, hit_limit) builder.add_output(done, "done") return _component( - PolicyRole.TERMINATION, + "onnx-genai.termination-predicate@1", graph, { "role": "termination_predicate", @@ -1221,7 +1208,7 @@ def build_euler_model_input(dtype: ir.DataType = ir.DataType.FLOAT) -> PolicyCom model_input = op.Div(sample, scale) model_input.shape = sample.shape builder.add_output(model_input, "model_input") - return _component(PolicyRole.AUXILIARY, graph, {}) + return _component("mobius.policy.auxiliary@1", graph, {}) def build_euler_solver_step( @@ -1252,7 +1239,7 @@ def build_euler_solver_step( next_sample = op.Add(sample, op.Mul(derivative, delta)) builder.add_output(next_sample, "next_state") return _component( - PolicyRole.SOLVER_STEP, + "onnx-genai.solver-step@1", graph, { "role": "solver_step", @@ -1358,7 +1345,7 @@ def build_masked_token_update() -> PolicyComponent: builder.add_output(next_offset, "next_offset") builder.add_output(done, "done") return _component( - PolicyRole.MASKED_UPDATE, + "onnx-genai.masked-update@1", graph, { "role": "masked_update", @@ -1462,7 +1449,7 @@ def build_speculative_acceptance() -> PolicyComponent: builder.add_output(synchronized_done, "synchronized_done") builder.add_output(rollback_len, "rollback_len") return _component( - PolicyRole.SPECULATIVE_ACCEPTANCE, + "onnx-genai.speculative-verifier@1", graph, { "role": "speculative_verifier", @@ -1512,7 +1499,7 @@ def build_speculative_state_rollback( corrected_shape[sequence_axis] = "accepted_sequence" corrected.shape = ir.Shape(corrected_shape) builder.add_output(corrected, "corrected_state") - return _component(PolicyRole.AUXILIARY, graph, {}, effect) + return _component("mobius.policy.auxiliary@1", graph, {}, effect) def build_effectful_identity( @@ -1526,7 +1513,7 @@ def build_effectful_identity( graph, builder = _make_graph(name) value = builder.input("value", dtype, shape) builder.add_output(builder.op.Identity(value), "next_value") - return _component(PolicyRole.AUXILIARY, graph, {}, effect) + return _component("mobius.policy.auxiliary@1", graph, {}, effect) def build_token_block_identity() -> PolicyComponent: @@ -1534,7 +1521,7 @@ def build_token_block_identity() -> PolicyComponent: 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(PolicyRole.AUXILIARY, graph, {}, "state") + return _component("mobius.policy.auxiliary@1", graph, {}, "state") def build_token_state_update() -> PolicyComponent: @@ -1549,7 +1536,7 @@ def build_token_state_update() -> PolicyComponent: ) builder.add_output(next_state, "next") return _component( - PolicyRole.STATE_UPDATE, + "onnx-genai.state-update@1", graph, { "role": "state_update", @@ -1569,4 +1556,4 @@ def build_token_to_slot() -> PolicyComponent: slot = builder.op.Unsqueeze(token, [-1]) slot.shape = ir.Shape(["batch", 1]) builder.add_output(slot, "slot") - return _component(PolicyRole.AUXILIARY, graph, {}) + 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 index a4ba5e80a..71c60099f 100644 --- a/src/mobius/generation/_policy_components_test.py +++ b/src/mobius/generation/_policy_components_test.py @@ -10,7 +10,6 @@ from mobius._model_package import ModelPackage from mobius.generation import ( PolicyCapabilities, - PolicyRole, attach_policy_components, build_adaptive_k_policy, build_batch_minimum, @@ -632,11 +631,11 @@ def test_capability_driven_attachment_is_model_agnostic(): "speculative_acceptance": "policies/speculative_acceptance.onnx", "token_state_update": "policies/token_state_update.onnx", } - assert {component.role for component in package.policy_components.values()} == { - PolicyRole.TOKEN_SAMPLER, - PolicyRole.TERMINATION, - PolicyRole.SOLVER_STEP, - PolicyRole.MASKED_UPDATE, - PolicyRole.SPECULATIVE_ACCEPTANCE, - PolicyRole.STATE_UPDATE, + 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", } diff --git a/src/mobius/integrations/onnx_genai/inference_metadata.py b/src/mobius/integrations/onnx_genai/inference_metadata.py index f3e105f41..15b15d727 100644 --- a/src/mobius/integrations/onnx_genai/inference_metadata.py +++ b/src/mobius/integrations/onnx_genai/inference_metadata.py @@ -1380,8 +1380,9 @@ def add_policy_components_to_workflow( return metadata components = workflow.setdefault("components", {}) - def semantic_contract(contract: dict[str, Any]) -> dict[str, Any]: - role = str(contract["role"]).replace("_", "-") + 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() @@ -1394,8 +1395,8 @@ def semantic_contract(contract: dict[str, Any]) -> dict[str, Any]: {key: value for key, value in rng.items() if isinstance(value, str)} ) declaration: dict[str, Any] = { - "id": f"onnx-genai.{role}", - "version": "1", + "id": contract_name, + "version": version, "bindings": bindings, } if "mode" in contract: @@ -1410,7 +1411,7 @@ def semantic_contract(contract: dict[str, Any]) -> dict[str, Any]: }, } if component.contract: - declaration["contract"] = semantic_contract(component.contract) + declaration["contract"] = semantic_contract(component) if component.contract.get("role") == "token_sampler": declaration["application_overridable"] = True components[name] = declaration From a201acda033f2b441117df6129758b339e0e4f20 Mon Sep 17 00:00:00 2001 From: justinchuby Date: Thu, 13 Aug 2026 01:11:26 +0000 Subject: [PATCH 028/151] Use lexical workflow loop iteration Remove generated iteration state and increment ONNX artifacts from decoder and masked-diffusion workflows. Public invokes now contain only semantic tensor mappings while the workflow serializer emits structured steps and logical carries without effect-token fields. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: justinchuby --- src/mobius/generation/__init__.py | 2 - src/mobius/generation/_policy_components.py | 10 --- .../onnx_genai/auto_export_test.py | 12 +++- .../onnx_genai/workflow_metadata.py | 61 ++++--------------- .../onnx_genai/workflow_metadata_test.py | 4 +- 5 files changed, 25 insertions(+), 64 deletions(-) diff --git a/src/mobius/generation/__init__.py b/src/mobius/generation/__init__.py index bf6ca077c..505dc29d6 100644 --- a/src/mobius/generation/__init__.py +++ b/src/mobius/generation/__init__.py @@ -23,7 +23,6 @@ build_euler_solver_step, build_grammar_logits_processor, build_greedy_sampler, - build_integer_increment, build_integer_minimum, build_iteration_cast, build_last_token_logits, @@ -62,7 +61,6 @@ "build_effectful_identity", "build_greedy_sampler", "build_grammar_logits_processor", - "build_integer_increment", "build_integer_minimum", "build_iteration_cast", "build_last_token_logits", diff --git a/src/mobius/generation/_policy_components.py b/src/mobius/generation/_policy_components.py index c26870ce9..be8e1e157 100644 --- a/src/mobius/generation/_policy_components.py +++ b/src/mobius/generation/_policy_components.py @@ -187,16 +187,6 @@ def build_boolean_not() -> PolicyComponent: return _component("mobius.policy.auxiliary@1", graph, {}) -def build_integer_increment() -> PolicyComponent: - """Build an explicit per-batch loop-counter increment.""" - graph, builder = _make_graph("integer_increment") - value = builder.input("value", dtype=ir.DataType.INT64, shape=["batch"]) - next_value = builder.op.Add(value, builder.op.Constant(value_int=1)) - next_value.shape = ir.Shape(["batch"]) - builder.add_output(next_value, "next_value") - 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") diff --git a/src/mobius/integrations/onnx_genai/auto_export_test.py b/src/mobius/integrations/onnx_genai/auto_export_test.py index 1e0fbcd3e..9560af08a 100644 --- a/src/mobius/integrations/onnx_genai/auto_export_test.py +++ b/src/mobius/integrations/onnx_genai/auto_export_test.py @@ -179,7 +179,17 @@ def test_dispatch_decoder(tmp_path): 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"] == "sample.body" - assert workflow["state"]["iteration"]["initializer"] == "package.zero_iteration" + 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]}, diff --git a/src/mobius/integrations/onnx_genai/workflow_metadata.py b/src/mobius/integrations/onnx_genai/workflow_metadata.py index 5a105a8e0..9bed96117 100644 --- a/src/mobius/integrations/onnx_genai/workflow_metadata.py +++ b/src/mobius/integrations/onnx_genai/workflow_metadata.py @@ -26,7 +26,6 @@ build_euler_model_input, build_euler_solver_step, build_greedy_sampler, - build_integer_increment, build_integer_minimum, build_last_token_logits, build_model_token_cast, @@ -123,7 +122,7 @@ def _effect(consumes: str, produces: str) -> dict[str, str]: def _publish_workflow_v1(workflow: dict[str, Any]) -> dict[str, Any]: - """Lower the producer's explicit-effect graph into the public v1 step IR.""" + """Publish structured steps and logical carries without compiler bookkeeping.""" graph = workflow.pop("graph") workflow.pop("initial_effects", None) substitutions: dict[str, str] = {} @@ -294,14 +293,13 @@ def _invoke( component: str, inputs: dict[str, str], outputs: dict[str, str], - effects: dict[str, dict[str, str]] | None = None, + _effects: dict[str, dict[str, str]] | None = None, ) -> dict[str, Any]: return { "kind": "invoke", "component": component, "inputs": inputs, "outputs": outputs, - "effects": effects or {}, } @@ -3621,7 +3619,6 @@ def build_decoder_workflow_metadata( ) pkg.add_policy_component("last_token_logits", build_last_token_logits()) pkg.add_policy_component("continue_predicate", build_boolean_not()) - pkg.add_policy_component("iteration_increment", build_integer_increment()) inputs = list(decoder.graph.inputs) outputs = list(decoder.graph.outputs) @@ -3792,13 +3789,6 @@ def build_decoder_workflow_metadata( "required": True, "default": eos_token_id, }, - "package.zero_iteration": { - "contract": batch_int, - "role": {"kind": "opaque"}, - "source": {"kind": "literal"}, - "required": False, - "default": 0, - }, "package.one_token": { "contract": batch_int, "role": {"kind": "opaque"}, @@ -3936,12 +3926,6 @@ def build_decoder_workflow_metadata( "initializer": "initializer.token_slot", "recurrence": {"kind": "invariant"}, }, - "iteration": { - "contract": batch_int, - "scope": "invocation", - "initializer": "package.zero_iteration", - "recurrence": {"kind": "invariant"}, - }, "logits": { "contract": last_logits_contract, "scope": "invocation", @@ -3955,7 +3939,6 @@ def build_decoder_workflow_metadata( "state": "state.0", "emit": "emit.0", "state:token": "state:token.0", - "state:iteration": "state:iteration.0", "state:logits": "state:logits.0", } carried = [ @@ -3971,15 +3954,6 @@ def build_decoder_workflow_metadata( ] carried.extend( [ - { - "cell": "iteration", - "current": "package.zero_iteration", - "body_input": "state.iteration.body", - "body_output": "iteration.body", - "next": "state.iteration.final", - "read_effect": _effect("state:iteration.0", "state:iteration.read"), - "write_effect": _effect("state:iteration.read", "state:iteration.1"), - }, { "cell": "logits", "current": "decoder.setup.last_logits", @@ -4175,7 +4149,7 @@ def build_decoder_workflow_metadata( { "token_ids": "sample.body", "eos_ids": "package.eos_ids", - "iteration": "state.iteration.body", + "iteration": "loop.iteration", "max_iterations": "request.max_iterations", }, {"done": "loop.done"}, @@ -4194,11 +4168,6 @@ def build_decoder_workflow_metadata( "effect_name": "emit", "effect": _effect("emit.0", "emit.1"), }, - _invoke( - "iteration_increment", - {"value": "state.iteration.body"}, - {"next_value": "iteration.body"}, - ), _invoke(decoder_name, body_decoder_inputs, body_decoder_outputs), _invoke( "last_token_logits", @@ -4257,6 +4226,10 @@ def build_decoder_workflow_metadata( "body": body, "condition": "loop.continue", "max_iterations": "request.max_iterations", + "iteration": { + "value": "loop.iteration", + "contract": {"dtype": "int64", "rank": 1, "shape": [1]}, + }, "carried": carried, }, } @@ -4309,7 +4282,6 @@ def build_language_diffusion_pipeline_metadata( attach_policy_components(pkg, PolicyCapabilities(masked_update=True)) pkg.add_policy_component("continue_predicate", build_boolean_not()) - pkg.add_policy_component("iteration_increment", build_integer_increment()) token_contract = _contract(token_input) mask_contract = { @@ -4365,13 +4337,6 @@ def build_language_diffusion_pipeline_metadata( "required": False, "default": num_inference_steps, }, - "package.zero_iteration": { - "contract": batch_int, - "role": {"kind": "opaque"}, - "source": {"kind": "literal"}, - "required": False, - "default": 0, - }, "package.num_steps": { "contract": batch_int, "role": {"kind": "opaque"}, @@ -4433,7 +4398,7 @@ def update_invoke( update_invoke( "state.tokens.body", "state.mask.body", - "state.iteration.body", + "loop.iteration", "state.rng_offset.body", "state.logits.body", "state.proposal.body", @@ -4446,11 +4411,6 @@ def update_invoke( {"done": "denoiser.body.done"}, {"continue": "denoiser.body.continue"}, ), - _invoke( - "iteration_increment", - {"value": "state.iteration.body"}, - {"next_value": "denoiser.body.iteration"}, - ), { "kind": "emit", "value": "denoiser.body.tokens", @@ -4467,7 +4427,6 @@ def update_invoke( "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"), - "iteration": (batch_int, "package.zero_iteration", "package.zero_iteration"), "logits": ( _contract(logits_output), "denoiser.setup.logits", @@ -4531,6 +4490,10 @@ def update_invoke( "body": body, "condition": "denoiser.body.continue", "max_iterations": "request.max_iterations", + "iteration": { + "value": "loop.iteration", + "contract": {"dtype": "int64", "rank": 1, "shape": [1]}, + }, "carried": carried, }, } diff --git a/src/mobius/integrations/onnx_genai/workflow_metadata_test.py b/src/mobius/integrations/onnx_genai/workflow_metadata_test.py index edf1c6a7c..162202312 100644 --- a/src/mobius/integrations/onnx_genai/workflow_metadata_test.py +++ b/src/mobius/integrations/onnx_genai/workflow_metadata_test.py @@ -153,14 +153,14 @@ def test_language_diffusion_uses_exclusive_ssa_workflow(): 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", "invoke", "invoke", "emit", "invoke", ] + assert graph["iteration"]["value"] == "loop.iteration" assert graph["steps"][0]["inputs"]["total_steps"] == "package.num_steps" - assert graph["steps"][-2]["mode"] == "replace" + assert graph["steps"][2]["mode"] == "replace" def test_language_diffusion_rejects_zero_steps(): From d8804826f1c1cb6ceb70ec70c21a10f71216d150 Mon Sep 17 00:00:00 2001 From: justinchuby Date: Thu, 13 Aug 2026 02:08:30 +0000 Subject: [PATCH 029/151] Align workflows with ONNX GenAI serving contracts Emit per-row speculative acceptance and rollback lengths, serving-aware KV aliases and logical lengths, and static control contracts across decoder, VLM, diffusion, and TTS workflows. Replace pinned JSON-schema checks with authoritative semantic package validation in CI. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: justinchuby --- .github/workflows/main.yml | 29 + src/mobius/generation/__init__.py | 12 +- src/mobius/generation/_policy_components.py | 47 +- .../generation/_policy_components_test.py | 32 +- .../onnx_genai/workflow_metadata.py | 810 ++- .../onnx_genai/workflow_metadata_test.py | 57 +- ...generate_onnx_genai_validation_packages.py | 84 + tests/schemas/onnx_genai_b2157a2.schema.json | 4339 ----------------- 8 files changed, 752 insertions(+), 4658 deletions(-) create mode 100644 tests/generate_onnx_genai_validation_packages.py delete mode 100644 tests/schemas/onnx_genai_b2157a2.schema.json diff --git a/.github/workflows/main.yml b/.github/workflows/main.yml index 671f259f4..63f399053 100644 --- a/.github/workflows/main.yml +++ b/.github/workflows/main.yml @@ -15,6 +15,35 @@ 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 + ref: 8215649100a0a27be15b04045fddde777c8248fc + 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 + pip install -e '.[testing]' + - name: Generate representative packages + run: PYTHONPATH=src python tests/generate_onnx_genai_validation_packages.py validation/packages + - name: Validate package semantics + run: | + for package in validation/packages/*; do + cargo run --quiet \ + --manifest-path validation/onnx-genai/Cargo.toml \ + -p onnx-genai-metadata --bin validate_metadata -- "$package" + done + lint: name: Lint runs-on: ubuntu-latest diff --git a/src/mobius/generation/__init__.py b/src/mobius/generation/__init__.py index 505dc29d6..57d4f0726 100644 --- a/src/mobius/generation/__init__.py +++ b/src/mobius/generation/__init__.py @@ -23,6 +23,7 @@ build_euler_solver_step, build_grammar_logits_processor, build_greedy_sampler, + build_integer_add, build_integer_minimum, build_iteration_cast, build_last_token_logits, @@ -44,8 +45,8 @@ ) __all__ = [ - "PolicyComponent", "PolicyCapabilities", + "PolicyComponent", "attach_policy_components", "build_adaptive_k_policy", "build_batch_minimum", @@ -55,12 +56,13 @@ "build_codec_layout_transpose", "build_decoder_state_initializer", "build_decoder_step_update", + "build_effectful_identity", "build_eos_termination", "build_euler_model_input", "build_euler_solver_step", - "build_effectful_identity", - "build_greedy_sampler", "build_grammar_logits_processor", + "build_greedy_sampler", + "build_integer_add", "build_integer_minimum", "build_iteration_cast", "build_last_token_logits", @@ -69,12 +71,12 @@ "build_proposal_metrics", "build_schedule_constant", "build_schedule_lookup", - "build_sequence_length", "build_seeded_categorical_sampler", + "build_sequence_length", "build_speculative_acceptance", "build_speculative_state_rollback", - "build_token_state_update", "build_token_block_identity", + "build_token_state_update", "build_token_to_slot", "build_tts_decoder_state_initializer", "build_tts_decoder_step_update", diff --git a/src/mobius/generation/_policy_components.py b/src/mobius/generation/_policy_components.py index be8e1e157..c1fcf4866 100644 --- a/src/mobius/generation/_policy_components.py +++ b/src/mobius/generation/_policy_components.py @@ -198,6 +198,17 @@ def build_integer_minimum() -> PolicyComponent: 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_batch_minimum() -> PolicyComponent: """Synchronize a per-batch integer length to one conservative scalar.""" graph, builder = _make_graph("batch_minimum") @@ -1167,7 +1178,13 @@ def build_eos_termination() -> PolicyComponent: max_iterations, ) done = op.Or(hit_eos, hit_limit) + continued = op.Equal( + op.ReduceMax(op.Cast(done, to=ir.DataType.INT64), keepdims=1), + op.Constant(value_int=0), + ) + continued.shape = ir.Shape([1]) builder.add_output(done, "done") + builder.add_output(continued, "continue") return _component( "onnx-genai.termination-predicate@1", graph, @@ -1178,6 +1195,7 @@ def build_eos_termination() -> PolicyComponent: "iteration": "iteration", "max_iterations": "max_iterations", "done": "done", + "continue": "continue", "effect": "termination", }, "termination", @@ -1325,6 +1343,11 @@ def build_masked_token_update() -> PolicyComponent: ) 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)), @@ -1334,6 +1357,7 @@ def build_masked_token_update() -> PolicyComponent: 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, @@ -1345,6 +1369,7 @@ def build_masked_token_update() -> PolicyComponent: "step": "step", "next_state": "next_state", "next_mask": "next_mask", + "continue": "continue", "rng": { "seed": "seed", "offset": "offset", @@ -1399,16 +1424,7 @@ def build_speculative_acceptance() -> PolicyComponent: op.Add(verified_count, op.Cast(op.Not(done), to=ir.DataType.INT64)), draft_length, ) - # Dense batched state has one physical sequence length. Synchronize to the - # shortest verified prefix so every row can share the same rollback point. - synchronized_len = op.ReduceMin(accepted_count, axes=[0], keepdims=1) - rollback_len = op.ReduceMin(verified_count, axes=[0], keepdims=1) - accepted_count = op.Expand(synchronized_len, op.Shape(accepted_count)) - synchronized_done = op.Cast( - op.ReduceMin(op.Cast(done, to=ir.DataType.INT64), axes=[0], keepdims=1), - to=ir.DataType.BOOL, - ) - done = op.Expand(synchronized_done, op.Shape(done)) + continued = op.Not(done) positions = op.Range( op.Constant(value_int=0), op.Squeeze(draft_length, op.Constant(value_ints=[0])), @@ -1427,17 +1443,15 @@ def build_speculative_acceptance() -> PolicyComponent: 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"]) - synchronized_len.shape = ir.Shape([1]) - rollback_len.shape = ir.Shape([1]) - synchronized_done.shape = ir.Shape([1]) + 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(synchronized_len, "synchronized_len") - builder.add_output(synchronized_done, "synchronized_done") - builder.add_output(rollback_len, "rollback_len") + builder.add_output(verified_count, "rollback_len") + builder.add_output(continued, "continue") return _component( "onnx-genai.speculative-verifier@1", graph, @@ -1448,6 +1462,7 @@ def build_speculative_acceptance() -> PolicyComponent: "accepted_tokens": "accepted_tokens", "accepted_len": "accepted_len", "done": "done", + "continue": "continue", "rng": { "seed": "seed", "offset": "offset", diff --git a/src/mobius/generation/_policy_components_test.py b/src/mobius/generation/_policy_components_test.py index 71c60099f..d5245953a 100644 --- a/src/mobius/generation/_policy_components_test.py +++ b/src/mobius/generation/_policy_components_test.py @@ -285,7 +285,7 @@ def test_decoder_policy_chain_generates_multiple_tokens_from_prompt_only(tmp_pat (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,) = _run( + done, _continue = _run( build_eos_termination(), tmp_path, { @@ -402,7 +402,7 @@ def test_seeded_sampler_applies_request_min_p_in_logit_space(tmp_path): def test_eos_termination_runtime(tmp_path): - (terminated,) = _run( + terminated, continued = _run( build_eos_termination(), tmp_path, { @@ -413,6 +413,7 @@ def test_eos_termination_runtime(tmp_path): }, ) np.testing.assert_array_equal(terminated, [True, True, True]) + np.testing.assert_array_equal(continued, [False]) def test_euler_solver_runtime_parity(tmp_path): @@ -453,6 +454,7 @@ def test_masked_update_runtime_parity(tmp_path): 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(), @@ -471,6 +473,7 @@ def test_masked_update_runtime_parity(tmp_path): 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): @@ -479,9 +482,8 @@ def test_speculative_acceptance_prefix_runtime(tmp_path): count, done, next_offset, - synchronized_len, - synchronized_done, rollback_len, + continued, ) = _run( build_speculative_acceptance(), tmp_path, @@ -499,9 +501,8 @@ def test_speculative_acceptance_prefix_runtime(tmp_path): 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(synchronized_len, [3]) - np.testing.assert_array_equal(synchronized_done, [False]) 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, @@ -518,15 +519,14 @@ def test_speculative_acceptance_prefix_runtime(tmp_path): assert corrected_cache.shape[2] == 4 -def test_speculative_acceptance_synchronizes_batched_prefixes(tmp_path): +def test_speculative_acceptance_preserves_per_row_prefixes(tmp_path): ( accepted_tokens, count, done, _, - synchronized_len, - synchronized_done, rollback_len, + continued, ) = _run( build_speculative_acceptance(), tmp_path, @@ -543,12 +543,14 @@ def test_speculative_acceptance_synchronizes_batched_prefixes(tmp_path): "offset": np.array([0, 0], np.int64), }, ) - np.testing.assert_array_equal(count, [2, 2]) - np.testing.assert_array_equal(accepted_tokens, [[1, 0, 0, 0], [1, 0, 0, 0]]) - np.testing.assert_array_equal(done, [False, False]) - np.testing.assert_array_equal(synchronized_len, [2]) - np.testing.assert_array_equal(synchronized_done, [False]) - np.testing.assert_array_equal(rollback_len, [1]) + 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): diff --git a/src/mobius/integrations/onnx_genai/workflow_metadata.py b/src/mobius/integrations/onnx_genai/workflow_metadata.py index 9bed96117..f930e9867 100644 --- a/src/mobius/integrations/onnx_genai/workflow_metadata.py +++ b/src/mobius/integrations/onnx_genai/workflow_metadata.py @@ -15,17 +15,16 @@ from mobius.generation import ( PolicyCapabilities, attach_policy_components, - build_batch_minimum, build_boolean_not, build_code_frame_update, build_code_history_append, build_codec_layout_transpose, build_decoder_state_initializer, build_decoder_step_update, - build_effectful_identity, build_euler_model_input, build_euler_solver_step, build_greedy_sampler, + build_integer_add, build_integer_minimum, build_last_token_logits, build_model_token_cast, @@ -33,8 +32,6 @@ build_schedule_constant, build_schedule_lookup, build_sequence_length, - build_speculative_state_rollback, - build_token_block_identity, build_token_to_slot, build_tts_decoder_state_initializer, build_tts_decoder_step_update, @@ -125,7 +122,12 @@ 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) 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", {}) @@ -162,6 +164,7 @@ def rewrite(value: Any) -> Any: return value def convert(node: dict[str, Any]) -> dict[str, Any]: + nonlocal loop_index kind = node["kind"] if kind == "sequence": return { @@ -196,6 +199,8 @@ def convert(node: dict[str, Any]) -> dict[str, Any]: 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 = ( @@ -219,11 +224,34 @@ def convert(node: dict[str, Any]) -> dict[str, Any]: 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, - "condition": rewrite(node["condition"]), + "continue_when": active_cell, "max_iterations": rewrite(node["max_iterations"]), "carried": carried, } @@ -568,6 +596,7 @@ def _build_real_tts_workflow_metadata(pkg: Any, config: Any) -> dict[str, Any]: batch = _contract(prompt)["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]} inputs = { "request.prompt_tokens": { "contract": _contract(prompt), @@ -576,7 +605,7 @@ def _build_real_tts_workflow_metadata(pkg: Any, config: Any) -> dict[str, Any]: "required": True, }, "request.max_iterations": { - "contract": batch_int, + "contract": control_int, "role": {"kind": "runtime", "version": "1.0", "role": "max_output_tokens"}, "source": {"kind": "request", "field": "max_output_tokens"}, "required": True, @@ -603,33 +632,40 @@ def _build_real_tts_workflow_metadata(pkg: Any, config: Any) -> dict[str, Any]: "default": 1, }, "package.remaining_groups": { - "contract": batch_int, + "contract": control_int, "role": {"kind": "opaque"}, "source": {"kind": "literal"}, "required": False, "default": num_groups - 2, }, "package.predictor_context_limit": { - "contract": {"dtype": "int64", "rank": 0, "shape": []}, + "contract": control_int, "role": {"kind": "opaque"}, "source": {"kind": "literal"}, "required": False, "default": num_groups, }, "package.predictor_mask_limit": { - "contract": {"dtype": "int64", "rank": 0, "shape": []}, + "contract": control_int, "role": {"kind": "opaque"}, "source": {"kind": "literal"}, "required": False, "default": num_groups + 1, }, "package.talker_context_limit": { - "contract": {"dtype": "int64", "rank": 0, "shape": []}, + "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, + }, } for iteration in range(num_groups - 2): inputs[f"package.setup_predictor_iteration_{iteration}"] = { @@ -1130,7 +1166,7 @@ def frame_generation_nodes(prefix: str, hidden: str, logits: str) -> list[dict[s "recurrence": { "kind": "growing", "axis": 1, - "increment": "package.one_scalar", + "increment": "package.one_control", "max": "package.talker_context_limit", }, }, @@ -1145,7 +1181,7 @@ def frame_generation_nodes(prefix: str, hidden: str, logits: str) -> list[dict[s "recurrence": { "kind": "growing", "axis": 1, - "increment": "package.one_scalar", + "increment": "package.one_control", "max": "package.talker_context_limit", }, }, @@ -1186,7 +1222,7 @@ def frame_generation_nodes(prefix: str, hidden: str, logits: str) -> list[dict[s "recurrence": { "kind": "growing", "axis": 1, - "increment": "package.one_scalar", + "increment": "package.one_control", "max": "package.predictor_mask_limit", }, }, @@ -1213,7 +1249,7 @@ def frame_generation_nodes(prefix: str, hidden: str, logits: str) -> list[dict[s "recurrence": { "kind": "growing", "axis": 2, - "increment": "package.one_scalar", + "increment": "package.one_control", "max": "package.talker_context_limit", }, } @@ -1225,7 +1261,7 @@ def frame_generation_nodes(prefix: str, hidden: str, logits: str) -> list[dict[s "recurrence": { "kind": "growing", "axis": 2, - "increment": "package.one_scalar", + "increment": "package.one_control", "max": "package.predictor_context_limit", }, } @@ -1470,6 +1506,7 @@ def build_tts_workflow_metadata(pkg: Any, config: Any) -> dict[str, Any]: batch = _contract(prompt_input)["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]} inputs: dict[str, Any] = { "request.prompt_tokens": { "contract": _contract(prompt_input), @@ -1478,7 +1515,7 @@ def build_tts_workflow_metadata(pkg: Any, config: Any) -> dict[str, Any]: "required": True, }, "request.max_iterations": { - "contract": batch_int, + "contract": control_int, "role": { "kind": "runtime", "version": "1.0", @@ -1488,7 +1525,7 @@ def build_tts_workflow_metadata(pkg: Any, config: Any) -> dict[str, Any]: "required": True, }, "package.code_groups": { - "contract": batch_int, + "contract": control_int, "role": {"kind": "opaque"}, "source": {"kind": "literal"}, "required": False, @@ -1502,7 +1539,7 @@ def build_tts_workflow_metadata(pkg: Any, config: Any) -> dict[str, Any]: "default": False, }, "package.one": { - "contract": batch_int, + "contract": control_int, "role": {"kind": "opaque"}, "source": {"kind": "literal"}, "required": False, @@ -1958,6 +1995,7 @@ def build_diffusion_workflow_metadata( batch = _contract(sample_input)["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]} inputs: dict[str, Any] = { "request.latent": { "contract": _contract(sample_input), @@ -1966,7 +2004,7 @@ def build_diffusion_workflow_metadata( "required": True, }, "request.max_iterations": { - "contract": batch_int, + "contract": control_int, "role": {"kind": "runtime", "version": "1.0", "role": "max_iterations"}, "source": {"kind": "request", "field": "max_iterations"}, "required": False, @@ -2305,7 +2343,6 @@ def build_vlm_workflow_metadata( ), ) pkg.add_policy_component("last_token_logits", build_last_token_logits()) - pkg.add_policy_component("continue_predicate", build_boolean_not()) pkg.add_policy_component( "decoder_state_initializer", build_decoder_state_initializer( @@ -2324,9 +2361,13 @@ def build_vlm_workflow_metadata( position_dtype=position_input.dtype if position_input is not None else None, ), ) + if cache_pairs: + pkg.add_policy_component("cache_length_update", build_integer_add()) batch = _contract(token_input)["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]} eos = getattr(config, "eos_token_id", 0) if isinstance(eos, list): eos = eos[0] if eos else 0 @@ -2344,7 +2385,7 @@ def build_vlm_workflow_metadata( "required": True, }, "request.max_iterations": { - "contract": batch_int, + "contract": control_int, "role": { "kind": "runtime", "version": "1.0", @@ -2361,19 +2402,47 @@ def build_vlm_workflow_metadata( "default": int(eos or 0), }, "package.max_context": { - "contract": batch_int, + "contract": control_int, "role": {"kind": "opaque"}, "source": {"kind": "literal"}, "required": False, "default": int(getattr(config, "max_position_embeddings", 4096)), }, "package.one": { - "contract": batch_int, + "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, + }, + "package.slot_ids": { + "contract": batch_int, + "role": {"kind": "opaque"}, + "source": {"kind": "literal"}, + "required": False, + "default": 0, + }, } vision_invoke_inputs = { name: preprocessing_values[name] @@ -2469,6 +2538,41 @@ def build_vlm_workflow_metadata( "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"}, + }, + "slot_ids": { + "contract": batch_int, + "class": "semantic", + "scope": "invocation", + "initializer": "package.slot_ids", + "recurrence": {"kind": "invariant"}, + }, + "cache_lengths": { + "contract": batch_int, + "class": "semantic", + "scope": "invocation", + "initializer": "package.zero_batch", + "recurrence": {"kind": "invariant"}, + }, } state_specs = [ ( @@ -2492,6 +2596,35 @@ def build_vlm_workflow_metadata( "decoder_step.body_attention_mask", "state.attention_mask.final", ), + ( + "active", + "package.active", + "state.active.body", + "loop.continue", + "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", + ), + ( + "slot_ids", + "package.slot_ids", + "state.slot_ids.body", + "state.slot_ids.body", + "state.slot_ids.final", + ), + ( + "cache_lengths", + "package.zero_batch", + "state.cache_lengths.body", + "cache_lengths.next", + "state.cache_lengths.final", + ), ] if position_input is not None: state["position_ids"] = { @@ -2520,7 +2653,7 @@ def build_vlm_workflow_metadata( "scope": "invocation", "initializer": f"decoder.setup.{present.name}", "recurrence": { - "kind": "growing", + "kind": "bounded", "axis": next( ( axis @@ -2529,9 +2662,9 @@ def build_vlm_workflow_metadata( ), 2, ), - "increment": "package.one", "max": "package.max_context", }, + "service_group": "decoder_cache", } setup_decoder_outputs[present.name] = f"decoder.setup.{present.name}" body_decoder_outputs[present.name] = f"decoder.body.{present.name}" @@ -2624,13 +2757,27 @@ def build_vlm_workflow_metadata( "iteration": "loop.iteration", "max_iterations": "request.max_iterations", }, - {"done": "loop.done"}, + {"done": "loop.done", "continue": "loop.continue"}, {"termination": _effect("termination.0", "termination.1")}, ), - _invoke( - "continue_predicate", - {"done": "loop.done"}, - {"continue": "loop.continue"}, + *( + [ + _invoke( + "cache_length_update", + { + "left": "state.cache_lengths.body", + "right": "package.one", + }, + {"total": "cache_lengths.next"}, + ), + _invoke( + "cache_length_update", + {"left": "package.zero_batch", "right": "package.one"}, + {"total": "accepted_len.next"}, + ), + ] + if cache_pairs + else [] ), { "kind": "emit", @@ -2709,6 +2856,11 @@ def build_vlm_workflow_metadata( "nested_control_flow", "loop_induction_values", "typed_emit", + *( + ["serving_service_contract", "bounded_state_recurrence"] + if cache_pairs + else [] + ), ], }, "inputs": inputs, @@ -2717,6 +2869,49 @@ def build_vlm_workflow_metadata( }, "components": components, "state": state, + **( + { + "serving": { + "active": "active", + "done": "done", + "accepted_len": "accepted_len", + "slot_ids": "slot_ids", + "kv_service": { + "paging": "paged", + "allocation": "runtime", + "compaction": True, + "groups": { + "decoder_cache": { + "sequence_axis": next( + ( + axis + for axis, dimension in enumerate( + _contract(cache_pairs[0][0])["shape"] + ) + if "sequence" in str(dimension) + ), + 2, + ), + "layout": "bnsh", + "logical_lengths": "cache_lengths", + "storage": "paged", + "ports": { + "decoder": { + f"cache_{index}": { + "input": past.name, + "output": present.name, + } + for index, (past, present) in enumerate(cache_pairs) + } + }, + } + }, + }, + } + } + if cache_pairs + else {} + ), "initial_effects": initial_effects, "graph": { "kind": "loop", @@ -2822,21 +3017,18 @@ def build_speculative_workflow_metadata( adaptive_k_max=adaptive_k_max, ), ) - pkg.add_policy_component("continue_predicate", build_boolean_not()) - pkg.add_policy_component("branch_state", build_token_block_identity()) if grammar_guidance: pkg.add_policy_component("grammar_length", build_integer_minimum()) - pkg.add_policy_component("grammar_emit_length", build_batch_minimum()) - pkg.add_policy_component("grammar_rollback_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 = {"dtype": "int64", "rank": 1, "shape": [batch]} + batch_bool = {"dtype": "bool", "rank": 1, "shape": [batch]} control_int = {"dtype": "int64", "rank": 1, "shape": [1]} - control_bool = {"dtype": "bool", "rank": 1, "shape": [1]} inputs: dict[str, Any] = { "request.tokens": { "contract": _contract(proposer_input), @@ -2883,12 +3075,32 @@ def build_speculative_workflow_metadata( "default": int(getattr(config, "max_position_embeddings", 4096)), }, "package.false": { - "contract": control_bool, + "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.slot_ids": { + "contract": batch_int, + "role": {"kind": "opaque"}, + "source": {"kind": "application", "name": "serving.slot_ids"}, + "required": 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( @@ -3092,12 +3304,10 @@ def build_speculative_workflow_metadata( ), ] ) - emit_length = "acceptance.synchronized_length" - cache_rollback_length = "acceptance.rollback_length" + emit_length = "acceptance.length" grammar_post_nodes: list[dict[str, Any]] = [] if grammar_guidance: - emit_length = "grammar.synchronized_length" - cache_rollback_length = "grammar.rollback_length" + emit_length = "grammar.committed_length" grammar_post_nodes.extend( [ _invoke( @@ -3108,19 +3318,6 @@ def build_speculative_workflow_metadata( }, {"minimum": "grammar.committed_length"}, ), - _invoke( - "grammar_rollback_length", - { - "left": "acceptance.rollback_length", - "right": "grammar.valid_length", - }, - {"minimum": "grammar.rollback_length"}, - ), - _invoke( - "grammar_emit_length", - {"values": "grammar.committed_length"}, - {"minimum": "grammar.synchronized_length"}, - ), _invoke( "grammar_commit", { @@ -3180,142 +3377,9 @@ def build_speculative_workflow_metadata( {"adaptive": _effect("adaptive.0", "adaptive.1")}, ) ) - rollback_nodes: list[dict[str, Any]] = [] - accepted_case_nodes = [ - _invoke( - "branch_state", - {"tokens": "acceptance.tokens"}, - {"next_tokens": "branch.accepted"}, - {"state": _effect("branch.state.in", "branch.state.accepted")}, - ) - ] - corrected_case_nodes = [ - _invoke( - "branch_state", - {"tokens": "acceptance.tokens"}, - {"next_tokens": "branch.corrected"}, - {"state": _effect("branch.state.in", "branch.state.corrected")}, - ) - ] - branch_outputs: dict[str, Any] = { - "tokens.next": { - "cases": { - "true": "branch.accepted", - "false": "branch.corrected", - } - } - } - branch_effects: dict[str, Any] = { - "state": { - "incoming": "branch.state.in", - "cases": { - "true": "branch.state.accepted", - "false": "branch.state.corrected", - }, - "produces": "branch.state.out", - } - } - for index, (past, present) in enumerate(cache_pairs): - cache_name = f"cache_{index}" - cache_contract = _contract(past) - sequence_axis = next( - ( - axis - for axis, dimension in enumerate(cache_contract["shape"]) - if "sequence" in str(dimension) - ), - 2, - ) - rollback_name = f"rollback_{cache_name}" - publisher_name = f"publish_{cache_name}" - branch_effect = f"branch:{cache_name}" - pkg.add_policy_component( - rollback_name, - build_speculative_state_rollback( - past.dtype, - cache_contract["shape"], - sequence_axis=sequence_axis, - effect=rollback_name, - ), - ) - pkg.add_policy_component( - publisher_name, - build_effectful_identity( - publisher_name, - past.dtype, - [ - "branch_sequence" if axis == sequence_axis else dimension - for axis, dimension in enumerate(cache_contract["shape"]) - ], - effect=branch_effect, - ), - ) - rollback_nodes.append( - _invoke( - rollback_name, - { - "past_state": f"state.{cache_name}.body", - "tentative_state": f"verifier.{present.name}", - "accepted_len": cache_rollback_length, - }, - {"corrected_state": f"rollback.{cache_name}"}, - { - rollback_name: _effect( - f"rollback.{cache_name}.0", - f"rollback.{cache_name}.1", - ) - }, - ) - ) - accepted_case_nodes.append( - _invoke( - publisher_name, - {"value": f"verifier.{present.name}"}, - {"next_value": f"branch.accepted.{cache_name}"}, - { - branch_effect: _effect( - f"branch.{cache_name}.in", - f"branch.{cache_name}.accepted", - ) - }, - ) - ) - corrected_case_nodes.append( - _invoke( - publisher_name, - {"value": f"rollback.{cache_name}"}, - {"next_value": f"branch.corrected.{cache_name}"}, - { - branch_effect: _effect( - f"branch.{cache_name}.in", - f"branch.{cache_name}.corrected", - ) - }, - ) - ) - branch_outputs[f"{cache_name}.next"] = { - "cases": { - "true": f"branch.accepted.{cache_name}", - "false": f"branch.corrected.{cache_name}", - } - } - branch_effects[branch_effect] = { - "incoming": f"branch.{cache_name}.in", - "cases": { - "true": f"branch.{cache_name}.accepted", - "false": f"branch.{cache_name}.corrected", - }, - "produces": f"branch.{cache_name}.out", - } - branch = { - "kind": "branch", - "predicate": "acceptance.synchronized_done", - "cases": { - "true": {"kind": "sequence", "nodes": accepted_case_nodes}, - "false": {"kind": "sequence", "nodes": corrected_case_nodes}, - }, - "outputs": branch_outputs, - "effects": branch_effects, + 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), @@ -3329,25 +3393,25 @@ def build_speculative_workflow_metadata( "accepted_tokens": "acceptance.tokens", "accepted_len": "acceptance.length", "done": "acceptance.done", + "continue": "acceptance.continue", "next_offset": "rng_offset.body", - "synchronized_len": "acceptance.synchronized_length", - "synchronized_done": "acceptance.synchronized_done", "rollback_len": "acceptance.rollback_length", }, {"verify": _effect("verify.0", "verify.1")}, ), *grammar_post_nodes, - *adaptive_nodes, - *rollback_nodes, - branch, _invoke( - "continue_predicate", - {"done": "package.false"}, - {"continue": "speculative.continue"}, + "cache_length_update", + { + "left": "state.cache_lengths.body", + "right": emit_length, + }, + {"total": "cache_lengths.next"}, ), + *adaptive_nodes, { "kind": "emit", - "value": "tokens.next", + "value": "acceptance.tokens", "valid_length": emit_length, "output": "tokens", "mode": "append", @@ -3381,13 +3445,48 @@ def build_speculative_workflow_metadata( "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"}, + }, + "slot_ids": { + "contract": batch_int, + "class": "semantic", + "scope": "invocation", + "initializer": "request.slot_ids", + "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", - "tokens.next", + "acceptance.tokens", "state.tokens.final", ), ( @@ -3397,6 +3496,41 @@ def build_speculative_workflow_metadata( "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", + ), + ( + "slot_ids", + "request.slot_ids", + "state.slot_ids.body", + "state.slot_ids.body", + "state.slot_ids.final", + ), + ( + "cache_lengths", + "request.cache_lengths", + "state.cache_lengths.body", + "cache_lengths.next", + "state.cache_lengths.final", + ), ] if grammar_guidance: state["grammar"] = { @@ -3448,9 +3582,19 @@ def build_speculative_workflow_metadata( ), ] ) - for index, (past, _present) in enumerate(cache_pairs): + 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": _contract(past), "class": "semantic", @@ -3458,23 +3602,18 @@ def build_speculative_workflow_metadata( "initializer": initializer, "recurrence": { "kind": "bounded", - "axis": next( - ( - axis - for axis, dimension in enumerate(_contract(past)["shape"]) - if "sequence" in str(dimension) - ), - 2, - ), + "axis": kv_sequence_axis, "max": "package.max_context", }, + "service_group": "verifier_cache", } + kv_ports[cell] = {"input": past.name, "output": present.name} state_specs.append( ( cell, initializer, f"state.{cell}.body", - f"{cell}.next", + cache_next_outputs[f"{cell}.next"], f"state.{cell}.final", ) ) @@ -3517,6 +3656,7 @@ def build_speculative_workflow_metadata( "typed_emit", "emit_valid_length", "bounded_state_recurrence", + "serving_service_contract", ], }, "inputs": inputs, @@ -3537,21 +3677,33 @@ def build_speculative_workflow_metadata( name: _component(model, _artifact(name, len(pkg))) for name, model in pkg.items() }, "state": state, + "serving": { + "active": "active", + "done": "done", + "accepted_len": "accepted_len", + "slot_ids": "slot_ids", + "kv_service": { + "paging": "paged", + "allocation": "runtime", + "compaction": True, + "groups": { + "verifier_cache": { + "sequence_axis": kv_sequence_axis, + "layout": "bnsh", + "logical_lengths": "cache_lengths", + "storage": "paged", + "ports": {"verifier": kv_ports}, + } + }, + }, + }, "initial_effects": initial_effects, "graph": { "kind": "loop", - "setup": { - "kind": "sequence", - "nodes": [ - _invoke( - "continue_predicate", - {"done": "package.false"}, - {"continue": "speculative.setup.continue"}, - ) - ], - }, + "setup": {"kind": "sequence", "nodes": []}, "body": {"kind": "sequence", "nodes": body_nodes}, - "condition": "speculative.continue", + "condition": "acceptance.continue", + "active_cell": "active", "max_iterations": "request.max_iterations", "iteration": {"value": "speculative.iteration", "contract": batch_int}, "carried": carried, @@ -3618,7 +3770,6 @@ def build_decoder_workflow_metadata( ), ) pkg.add_policy_component("last_token_logits", build_last_token_logits()) - pkg.add_policy_component("continue_predicate", build_boolean_not()) inputs = list(decoder.graph.inputs) outputs = list(decoder.graph.outputs) @@ -3766,6 +3917,8 @@ def build_decoder_workflow_metadata( batch_dimension = _shape_metadata(_port(token_input))[0] batch_int = {"dtype": "int64", "rank": 1, "shape": [batch_dimension]} + batch_bool = {"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 @@ -3773,7 +3926,7 @@ def build_decoder_workflow_metadata( workflow_inputs.update( { "request.max_iterations": { - "contract": batch_int, + "contract": control_int, "role": { "kind": "runtime", "version": "1.0", @@ -3790,14 +3943,14 @@ def build_decoder_workflow_metadata( "default": eos_token_id, }, "package.one_token": { - "contract": batch_int, + "contract": control_int, "role": {"kind": "opaque"}, "source": {"kind": "literal"}, "required": False, "default": 1, }, "package.max_context": { - "contract": batch_int, + "contract": control_int, "role": {"kind": "opaque"}, "source": {"kind": "literal"}, "required": False, @@ -3891,6 +4044,47 @@ def build_decoder_workflow_metadata( }, } ) + if cache_pairs: + pkg.add_policy_component("cache_length_update", build_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.slot_ids": { + "contract": batch_int, + "role": {"kind": "opaque"}, + "source": {"kind": "literal"}, + "required": False, + "default": 0, + }, + "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: @@ -3933,6 +4127,46 @@ def build_decoder_workflow_metadata( "recurrence": {"kind": "invariant"}, }, } + if cache_pairs: + state.update( + { + "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"}, + }, + "slot_ids": { + "contract": batch_int, + "class": "semantic", + "scope": "invocation", + "initializer": "package.slot_ids", + "recurrence": {"kind": "invariant"}, + }, + "cache_lengths": { + "contract": batch_int, + "class": "semantic", + "scope": "invocation", + "initializer": "package.cache_lengths", + "recurrence": {"kind": "invariant"}, + }, + } + ) initial_effects = { "sample": "sample.0", "termination": "termination.0", @@ -3965,6 +4199,46 @@ def build_decoder_workflow_metadata( }, ] ) + if cache_pairs: + carried.extend( + [ + { + "cell": "active", + "current": "package.active", + "body_input": "state.active.body", + "body_output": "loop.continue", + "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": "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", + }, + { + "cell": "slot_ids", + "current": "package.slot_ids", + "body_input": "state.slot_ids.body", + "body_output": "state.slot_ids.body", + "next": "state.slot_ids.final", + }, + ] + ) if stochastic_sampler: state["rng_offset"] = { "contract": batch_int, @@ -4033,30 +4307,34 @@ def build_decoder_workflow_metadata( "write_effect": _effect(f"{effect_name}.read", f"{effect_name}.1"), } ) + decoder_kv_ports: dict[str, Any] = {} + decoder_kv_axis = 2 for past, present in cache_pairs: cell = f"cache_{len(carried)}" 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, + ) state[cell] = { "contract": _contract(past), "scope": "invocation", "initializer": setup_value, "recurrence": { - "kind": "growing", - "axis": next( - ( - index - for index, dimension in enumerate(_contract(past)["shape"]) - if "sequence" in str(dimension) - ), - 2, - ), - "increment": "package.one_token", + "kind": "bounded", + "axis": decoder_kv_axis, "max": "package.max_context", }, + "service_group": "decoder_cache", } + decoder_kv_ports[cell] = {"input": past.name, "output": present.name} effect_name = f"state:{cell}" initial_effects[effect_name] = f"{effect_name}.0" carried.append( @@ -4152,13 +4430,30 @@ def build_decoder_workflow_metadata( "iteration": "loop.iteration", "max_iterations": "request.max_iterations", }, - {"done": "loop.done"}, + {"done": "loop.done", "continue": "loop.continue"}, {"termination": _effect("termination.0", "termination.1")}, ), - _invoke( - "continue_predicate", - {"done": "loop.done"}, - {"continue": "loop.continue"}, + *( + [ + _invoke( + "cache_length_update", + { + "left": "state.cache_lengths.body", + "right": "package.one_token", + }, + {"total": "cache_lengths.next"}, + ), + _invoke( + "cache_length_update", + { + "left": "package.zero_batch", + "right": "package.one_token", + }, + {"total": "accepted_len.next"}, + ), + ] + if cache_pairs + else [] ), { "kind": "emit", @@ -4207,6 +4502,9 @@ def build_decoder_workflow_metadata( "linear_effects", "nested_control_flow", "typed_emit", + "loop_induction_values", + *(["serving_service_contract"] if cache_pairs else []), + *(["bounded_state_recurrence"] if cache_pairs else []), ], }, "inputs": workflow_inputs, @@ -4219,12 +4517,39 @@ def build_decoder_workflow_metadata( }, "components": {decoder_name: _component(decoder, artifact)}, "state": state, + **( + { + "serving": { + "active": "active", + "done": "done", + "accepted_len": "accepted_len", + "slot_ids": "slot_ids", + "kv_service": { + "paging": "paged", + "allocation": "runtime", + "compaction": True, + "groups": { + "decoder_cache": { + "sequence_axis": decoder_kv_axis, + "layout": "bnsh", + "logical_lengths": "cache_lengths", + "storage": "paged", + "ports": {decoder_name: decoder_kv_ports}, + } + }, + }, + } + } + if cache_pairs + else {} + ), "initial_effects": initial_effects, "graph": { "kind": "loop", "setup": setup, "body": body, "condition": "loop.continue", + **({"active_cell": "active"} if cache_pairs else {}), "max_iterations": "request.max_iterations", "iteration": { "value": "loop.iteration", @@ -4281,7 +4606,6 @@ def build_language_diffusion_pipeline_metadata( ) attach_policy_components(pkg, PolicyCapabilities(masked_update=True)) - pkg.add_policy_component("continue_predicate", build_boolean_not()) token_contract = _contract(token_input) mask_contract = { @@ -4384,6 +4708,7 @@ def update_invoke( "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)}, ) @@ -4406,11 +4731,6 @@ def update_invoke( "update.0", "update.1", ), - _invoke( - "continue_predicate", - {"done": "denoiser.body.done"}, - {"continue": "denoiser.body.continue"}, - ), { "kind": "emit", "value": "denoiser.body.tokens", diff --git a/src/mobius/integrations/onnx_genai/workflow_metadata_test.py b/src/mobius/integrations/onnx_genai/workflow_metadata_test.py index 162202312..a82a2695f 100644 --- a/src/mobius/integrations/onnx_genai/workflow_metadata_test.py +++ b/src/mobius/integrations/onnx_genai/workflow_metadata_test.py @@ -4,9 +4,7 @@ from __future__ import annotations import json -from pathlib import Path -import jsonschema import onnx_ir as ir import pytest @@ -27,7 +25,7 @@ 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" / "branch_state.onnx").is_file() + assert (tmp_path / "policies" / "cache_length_update.onnx").is_file() def test_speculative_writer_saves_guidance_and_adaptive_artifacts(tmp_path): @@ -39,7 +37,6 @@ def test_speculative_writer_saves_guidance_and_adaptive_artifacts(tmp_path): ) assert (tmp_path / "policies" / "grammar_guidance.onnx").is_file() assert (tmp_path / "policies" / "adaptive_k.onnx").is_file() - assert (tmp_path / "policies" / "grammar_emit_length.onnx").is_file() def test_speculative_emit_uses_accepted_prefix_length(): @@ -47,7 +44,7 @@ def test_speculative_emit_uses_accepted_prefix_length(): "workflow" ] emit = next(node for node in workflow["steps"][0]["steps"] if node["kind"] == "emit") - assert emit["valid_length"] == "acceptance.synchronized_length" + assert emit["valid_length"] == "acceptance.length" assert "emit_valid_length" in workflow["manifest"]["capabilities"] assert workflow["outputs"]["tokens"]["contract"]["shape"][-1] == "accepted_sequence" assert workflow["state"]["cache_0"]["recurrence"] == { @@ -55,6 +52,8 @@ def test_speculative_emit_uses_accepted_prefix_length(): "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 @@ -144,23 +143,23 @@ def test_language_diffusion_uses_exclusive_ssa_workflow(): "seed": "seed", "offset": "offset", "next_offset": "next_offset", + "continue": "continue", }, } graph = workflow["steps"][0] assert graph["kind"] == "loop" - assert graph["condition"] == "denoiser.body.continue" + 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", "invoke", "emit", "invoke", ] assert graph["iteration"]["value"] == "loop.iteration" assert graph["steps"][0]["inputs"]["total_steps"] == "package.num_steps" - assert graph["steps"][2]["mode"] == "replace" + assert graph["steps"][1]["mode"] == "replace" def test_language_diffusion_rejects_zero_steps(): @@ -171,19 +170,6 @@ def test_language_diffusion_rejects_zero_steps(): ) -def test_language_diffusion_matches_pr_828_schema(): - schema_path = ( - Path(__file__).parents[4] / "tests" / "schemas" / "onnx_genai_b2157a2.schema.json" - ) - with schema_path.open(encoding="utf-8") as handle: - schema = json.load(handle) - metadata = build_language_diffusion_pipeline_metadata( - _masked_denoiser_package(), - num_inference_steps=8, - ) - jsonschema.validate(instance=metadata, schema=schema) - - def _graph_model( name: str, inputs: list[ir.Value], @@ -287,31 +273,26 @@ def test_speculative_grammar_and_adaptive_k_use_typed_state_contracts(): ) assert proposer["inputs"]["proposal_budget"] == "proposal_k" assert all("initial" not in carry for carry in workflow["steps"][0]["carried"]) - schema_path = ( - Path(__file__).parents[4] / "tests" / "schemas" / "onnx_genai_b2157a2.schema.json" - ) - with schema_path.open(encoding="utf-8") as handle: - jsonschema.validate(instance=metadata, schema=json.load(handle)) -def test_speculative_workflow_uses_branch_phi_and_rng(): +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"] - branch = next(node for node in body if node["kind"] == "branch") - assert branch["kind"] == "branch" - assert branch["outputs"]["tokens.next"]["cases"] == { - "true": "branch.accepted", - "false": "branch.corrected", - } acceptance = body[2] assert acceptance["inputs"]["offset"] == "rng_offset" assert acceptance["outputs"]["next_offset"] == "rng_offset.body" - rollback = next(node for node in body if node.get("component") == "rollback_cache_0") - assert rollback["inputs"]["accepted_len"] == "acceptance.rollback_length" - assert branch["outputs"]["cache_0.next"]["cases"] == { - "true": "branch.accepted.cache_0", - "false": "branch.corrected.cache_0", + 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 not any(node["kind"] == "branch" for node in body) + assert workflow["serving"]["active"] == "active" + assert workflow["serving"]["done"] == "done" + assert workflow["serving"]["kv_service"]["groups"]["verifier_cache"]["ports"]["verifier"][ + "cache_0" + ] == { + "input": "past_key_values.0.key", + "output": "present.0.key", } assert any(item["cell"].startswith("cache_") for item in workflow["steps"][0]["carried"]) diff --git a/tests/generate_onnx_genai_validation_packages.py b/tests/generate_onnx_genai_validation_packages.py new file mode 100644 index 000000000..5fa6cfcab --- /dev/null +++ b/tests/generate_onnx_genai_validation_packages.py @@ -0,0 +1,84 @@ +from __future__ import annotations + +import argparse +from pathlib import Path + +import onnx_ir as ir + +from mobius._model_package import ModelPackage +from mobius.integrations.onnx_genai import write_onnx_genai_config +from mobius.integrations.onnx_genai.auto_export_test import ( + _Cfg, + _decoder_package, + _diffusion_package, + _model, + _TTSCfg, + _value, + _vlm_package, +) +from mobius.integrations.onnx_genai.workflow_metadata import ( + write_speculative_workflow_metadata, +) +from mobius.integrations.onnx_genai.workflow_metadata_test import _speculative_package + + +def _tts_package() -> ModelPackage: + batch = "batch" + return ModelPackage( + { + "talker": _model( + "talker", + [_value("inputs_embeds", ir.DataType.FLOAT, [batch, "sequence", 16])], + [("last_hidden_state", ir.DataType.FLOAT, [batch, 16])], + ), + "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(), + ) + + +def main() -> None: + parser = argparse.ArgumentParser() + parser.add_argument("output", type=Path) + args = parser.parse_args() + packages = { + "decoder": (_decoder_package(), {"config": _Cfg()}), + "vlm": (_vlm_package(), {}), + "diffusion": (_diffusion_package(text=True), {"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) + + speculative = _speculative_package() + directory = args.output / "speculative" + speculative.save(str(directory), progress_bar=False, check_weights=False) + write_speculative_workflow_metadata(speculative, str(directory)) + + +if __name__ == "__main__": + main() diff --git a/tests/schemas/onnx_genai_b2157a2.schema.json b/tests/schemas/onnx_genai_b2157a2.schema.json deleted file mode 100644 index 91f345fe7..000000000 --- a/tests/schemas/onnx_genai_b2157a2.schema.json +++ /dev/null @@ -1,4339 +0,0 @@ -{ - "$defs": { - "AbsentInputKind": { - "description": "Supported absent-input fallback kinds.", - "oneOf": [ - { - "const": "zeros", - "description": "Materialize a zero-initialized tensor.", - "type": "string" - } - ] - }, - "AbsentInputSpec": { - "description": "Explicit tensor fallback for an absent optional graph input.", - "properties": { - "kind": { - "$ref": "#/$defs/AbsentInputKind", - "description": "Fallback materialization kind." - }, - "shape": { - "description": "Runtime-resolved shape of the fallback tensor.", - "items": { - "$ref": "#/$defs/TensorDimension" - }, - "type": "array" - } - }, - "required": [ - "kind", - "shape" - ], - "type": "object" - }, - "AcceptanceMethod": { - "description": "Speculative acceptance-rule vocabulary.", - "oneOf": [ - { - "enum": [ - "rejection_sampling", - "greedy", - "typical" - ], - "title": "Known standard value" - }, - { - "not": { - "enum": [ - "rejection_sampling", - "greedy", - "typical" - ] - }, - "title": "Forward-compatible extension value", - "type": "string" - } - ], - "type": "string" - }, - "AttentionConfig": { - "description": "Build-time attention architecture and dimensions.", - "properties": { - "fallback_behavior": { - "anyOf": [ - { - "$ref": "#/$defs/AttentionType" - }, - { - "type": "null" - } - ], - "description": "Compatible attention behavior for runtimes that do not recognize `type`." - }, - "head_dim": { - "description": "Per-head hidden dimension.", - "format": "uint", - "minimum": 1, - "type": [ - "integer", - "null" - ] - }, - "key_sequence_lengths": { - "anyOf": [ - { - "$ref": "#/$defs/KeySequenceLengthsSpec" - }, - { - "type": "null" - } - ], - "description": "Representation compatibility for the attention key-sequence lengths.\n\nAbsent means the canonical contiguous `int32 [batch_size]` representation\nis required." - }, - "num_attention_heads": { - "description": "Number of query/attention heads.", - "format": "uint", - "minimum": 1, - "type": [ - "integer", - "null" - ] - }, - "num_kv_heads": { - "description": "Number of key/value heads; required by runtimes that need explicit GQA dimensions.", - "format": "uint", - "minimum": 1, - "type": [ - "integer", - "null" - ] - }, - "sink_tokens": { - "description": "Number of leading \"attention sink\" tokens always retained alongside the\nsliding window (StreamingLLM). Only meaningful when `sliding_window` is\nset; `null` or `0` disables sink retention. These first tokens stabilize\nthe attention distribution and are never evicted by the window.", - "format": "uint", - "minimum": 0, - "type": [ - "integer", - "null" - ] - }, - "sliding_window": { - "description": "Sliding-window length in tokens, or null for full-context attention.", - "format": "uint", - "minimum": 1, - "type": [ - "integer", - "null" - ] - }, - "type": { - "$ref": "#/$defs/AttentionType", - "description": "Attention architecture.\n\nCanonical values include `multi_head`, `grouped_query`, and\n`multi_latent`; future values are allowed when paired with a usable\n`fallback_behavior`." - } - }, - "required": [ - "type" - ], - "type": "object" - }, - "AttentionType": { - "description": "Attention architecture vocabulary with an extension branch.", - "oneOf": [ - { - "enum": [ - "multi_head", - "multi_head_attention", - "grouped_query", - "group_query_attention", - "grouped_query_attention", - "gqa", - "multi_latent", - "multi_latent_attention", - "mla" - ], - "title": "Known standard value" - }, - { - "not": { - "enum": [ - "multi_head", - "multi_head_attention", - "grouped_query", - "group_query_attention", - "grouped_query_attention", - "gqa", - "multi_latent", - "multi_latent_attention", - "mla" - ] - }, - "title": "Forward-compatible extension value", - "type": "string" - } - ], - "type": "string" - }, - "ChunkedPrefillConfig": { - "description": "Runtime chunked-prefill preference.", - "properties": { - "chunk_size": { - "description": "Preferred number of prompt tokens processed in each prefill chunk.", - "format": "uint", - "minimum": 1, - "type": [ - "integer", - "null" - ] - } - }, - "type": "object" - }, - "ComponentContract": { - "additionalProperties": false, - "properties": { - "bindings": { - "additionalProperties": { - "type": "string" - }, - "default": {}, - "description": "Semantic role to concrete component port name.", - "type": "object" - }, - "id": { - "description": "Versioned semantic capability identifier. It never selects execution behavior.", - "type": "string" - }, - "parameters": { - "additionalProperties": { - "$ref": "#/$defs/ScalarValue" - }, - "default": {}, - "description": "Contract parameters that are not tensor ports, such as adapter actions.", - "type": "object" - }, - "version": { - "type": "string" - } - }, - "required": [ - "id", - "version" - ], - "type": "object" - }, - "ComponentImplementation": { - "oneOf": [ - { - "additionalProperties": false, - "properties": { - "artifact": { - "type": "string" - }, - "kind": { - "const": "onnx", - "type": "string" - } - }, - "required": [ - "kind", - "artifact" - ], - "type": "object" - }, - { - "additionalProperties": false, - "properties": { - "abi": { - "type": "string" - }, - "artifact": { - "type": [ - "string", - "null" - ] - }, - "custom_ops": { - "additionalProperties": { - "type": "string" - }, - "default": {}, - "type": "object" - }, - "kind": { - "const": "adapter", - "type": "string" - }, - "version": { - "type": "string" - } - }, - "required": [ - "kind", - "abi", - "version" - ], - "type": "object" - }, - { - "additionalProperties": false, - "properties": { - "kind": { - "const": "binding", - "type": "string" - } - }, - "required": [ - "kind" - ], - "type": "object" - } - ] - }, - "ComponentPorts": { - "additionalProperties": false, - "description": "Explicit input/output ports of one executable component.", - "properties": { - "inputs": { - "additionalProperties": { - "$ref": "#/$defs/TensorContract" - }, - "default": {}, - "type": "object" - }, - "outputs": { - "additionalProperties": { - "$ref": "#/$defs/TensorContract" - }, - "default": {}, - "type": "object" - } - }, - "type": "object" - }, - "DType": { - "description": "Scalar dtype vocabulary with common ONNX and runtime aliases.", - "oneOf": [ - { - "enum": [ - "float32", - "fp32", - "float16", - "fp16", - "half", - "bfloat16", - "bf16", - "float8_e4m3fn", - "fp8_e4m3fn", - "float8_e4m3", - "fp8_e4m3", - "float8_e5m2", - "fp8_e5m2", - "int8", - "uint8", - "int4", - "uint4" - ], - "title": "Known standard value" - }, - { - "not": { - "enum": [ - "float32", - "fp32", - "float16", - "fp16", - "half", - "bfloat16", - "bf16", - "float8_e4m3fn", - "fp8_e4m3fn", - "float8_e4m3", - "fp8_e4m3", - "float8_e5m2", - "fp8_e5m2", - "int8", - "uint8", - "int4", - "uint4" - ] - }, - "title": "Forward-compatible extension value", - "type": "string" - } - ], - "type": "string" - }, - "DeviceKind": { - "enum": [ - "cpu", - "cuda", - "direct_ml", - "core_ml", - "web_gpu", - "npu" - ], - "type": "string" - }, - "DraftConfig": { - "description": "Draft-token producer configuration.", - "properties": { - "depth": { - "description": "Self-speculative early-exit depth.", - "format": "uint", - "minimum": 1, - "type": [ - "integer", - "null" - ] - }, - "heads": { - "description": "Named draft-head layout or selection.", - "type": [ - "string", - "null" - ] - }, - "ngram": { - "description": "Runtime-specific n-gram or prompt-lookup configuration." - }, - "producer": { - "$ref": "#/$defs/DraftProducer", - "description": "Producer family: `draft_model`, `self_speculative`, `ngram`, or `extra_heads`." - }, - "session": { - "description": "Named runtime session or pipeline component used as the producer.", - "type": [ - "string", - "null" - ] - } - }, - "required": [ - "producer" - ], - "type": "object" - }, - "DraftProducer": { - "description": "Speculative draft-producer vocabulary.", - "oneOf": [ - { - "enum": [ - "draft_model", - "self_speculative", - "ngram", - "extra_heads" - ], - "title": "Known standard value" - }, - { - "not": { - "enum": [ - "draft_model", - "self_speculative", - "ngram", - "extra_heads" - ] - }, - "title": "Forward-compatible extension value", - "type": "string" - } - ], - "type": "string" - }, - "ForkPrecisionPolicy": { - "description": "KV fork-precision policy vocabulary.", - "oneOf": [ - { - "enum": [ - "inherit", - "highest", - "independent" - ], - "title": "Known standard value" - }, - { - "not": { - "enum": [ - "inherit", - "highest", - "independent" - ] - }, - "title": "Forward-compatible extension value", - "type": "string" - } - ], - "type": "string" - }, - "GenerationDefaults": { - "description": "Author-declared text-generation defaults (sampling and beam search).\n\nMirrors the `search` section of an onnxruntime-genai `genai_config.json`.\nEvery field is optional so only values the author declared are carried over.", - "properties": { - "diversity_penalty": { - "description": "Diversity penalty for diverse beam groups.", - "format": "float", - "type": [ - "number", - "null" - ] - }, - "do_sample": { - "description": "Whether to randomize sampling through `top_k`/`top_p` (else greedy).", - "type": [ - "boolean", - "null" - ] - }, - "early_stopping": { - "description": "Whether beam search stops once enough beams have finished.", - "type": [ - "boolean", - "null" - ] - }, - "length_penalty": { - "description": "Exponential length penalty used with beam search.", - "format": "float", - "type": [ - "number", - "null" - ] - }, - "max_length": { - "description": "Maximum final sequence length.", - "format": "uint", - "minimum": 0, - "type": [ - "integer", - "null" - ] - }, - "min_length": { - "description": "Minimum final sequence length.", - "format": "uint", - "minimum": 0, - "type": [ - "integer", - "null" - ] - }, - "no_repeat_ngram_size": { - "description": "Disallow repeating n-grams of this size (`0` = disabled).", - "format": "uint", - "minimum": 0, - "type": [ - "integer", - "null" - ] - }, - "num_beams": { - "description": "Number of beams for beam search (`1` = no beam search).", - "format": "uint", - "minimum": 0, - "type": [ - "integer", - "null" - ] - }, - "num_return_sequences": { - "description": "Number of sequences returned after search.", - "format": "uint", - "minimum": 0, - "type": [ - "integer", - "null" - ] - }, - "repetition_penalty": { - "description": "Penalty applied to already-generated tokens (`1.0` = no penalty).", - "format": "float", - "type": [ - "number", - "null" - ] - }, - "temperature": { - "description": "Softmax temperature applied before sampling.", - "format": "float", - "type": [ - "number", - "null" - ] - }, - "top_k": { - "description": "Number of highest-probability tokens kept for top-k filtering.", - "format": "uint", - "minimum": 0, - "type": [ - "integer", - "null" - ] - }, - "top_p": { - "description": "Nucleus (top-p) cumulative-probability threshold.", - "format": "float", - "type": [ - "number", - "null" - ] - } - }, - "type": "object" - }, - "HardwareRequirements": { - "description": "Model-side hardware requirements and distribution-matching hints.", - "properties": { - "beneficial_dtypes": { - "description": "Dtypes that improve performance or memory use but are not mandatory.", - "items": { - "$ref": "#/$defs/DType" - }, - "type": [ - "array", - "null" - ] - }, - "kv_cache_memory_per_1k_tokens_mb": { - "description": "Estimated KV-cache memory in MiB per 1,000 cached tokens.", - "format": "float", - "minimum": 0.0, - "type": [ - "number", - "null" - ] - }, - "min_memory_gb": { - "description": "Minimum aggregate accelerator or system memory in GiB.", - "format": "float", - "minimum": 0.0, - "type": [ - "number", - "null" - ] - }, - "min_tp_degree": { - "description": "Minimum useful tensor-parallel degree when tensor parallelism is selected.", - "format": "uint", - "minimum": 1, - "type": [ - "integer", - "null" - ] - }, - "required_dtypes": { - "description": "Dtypes the selected device or execution provider must support.", - "items": { - "$ref": "#/$defs/DType" - }, - "type": [ - "array", - "null" - ] - }, - "supports_tensor_parallel": { - "description": "Whether the model can be partitioned with tensor parallelism.", - "type": [ - "boolean", - "null" - ] - } - }, - "type": "object" - }, - "ImageOutputBinding": { - "description": "One named tensor output produced by an image preprocessing program.\n\nThe output binds a processor-local value to a typed workflow SSA name.\nNeither the name nor the content role is inferred from a model identity.", - "properties": { - "content": { - "$ref": "#/$defs/ImageOutputContent", - "description": "Generic content role this tensor carries (pixels, coordinates, grid,\noriginal size, or validity mask) — never a model-family label." - }, - "contract": { - "anyOf": [ - { - "$ref": "#/$defs/TensorContract" - }, - { - "type": "null" - } - ], - "description": "Full workflow tensor contract. Required when `pipeline.workflow` is present." - }, - "dtype": { - "$ref": "#/$defs/TensorDType", - "description": "Declared output dtype. Always explicit; never inferred from the model." - }, - "name": { - "description": "Workflow SSA value produced by the preprocessing adapter invocation.", - "examples": [ - "image.pixel_values" - ], - "minLength": 1, - "type": "string" - }, - "optional": { - "description": "Whether the runtime may omit this output when a model does not need it.", - "type": [ - "boolean", - "null" - ] - }, - "pad_value": { - "description": "Optional sentinel/pad value for padded entries (e.g. `-1` coordinates).", - "format": "double", - "type": [ - "number", - "null" - ] - }, - "source": { - "description": "Named processor-local value produced by a transform.", - "minLength": 1, - "type": "string" - } - }, - "required": [ - "source", - "name", - "content", - "dtype" - ], - "type": "object" - }, - "ImageOutputContent": { - "description": "Generic image-output content-role vocabulary.", - "oneOf": [ - { - "enum": [ - "pixels", - "patch_coordinates", - "grid_dimensions", - "original_size", - "transformed_size", - "validity_mask" - ], - "title": "Known standard value" - }, - { - "not": { - "enum": [ - "pixels", - "patch_coordinates", - "grid_dimensions", - "original_size", - "transformed_size", - "validity_mask" - ] - }, - "title": "Forward-compatible extension value", - "type": "string" - } - ], - "type": "string" - }, - "ImagePreprocessingProgram": { - "description": "Generic image preprocessing program: an ordered transform pipeline plus the\nnamed workflow SSA tensor outputs it emits.\n\nThe program is expressed entirely as parameterized, architecture-neutral\ndata. Transform operations are generic (decode, resize, rescale, normalize,\ntile, patchify, pad). In workflow metadata, outputs are materialized by a\nmanifest-pinned preprocessing adapter invocation and bind processor-local\nvalues to typed SSA names. A package may name an output `pixel_position_ids`,\n`image_grid_thw`, or anything else without introducing runtime model-family\ndispatch.", - "properties": { - "outputs": { - "description": "Named tensor outputs the program emits, each bound to a workflow SSA value.", - "items": { - "$ref": "#/$defs/ImageOutputBinding" - }, - "minItems": 1, - "type": "array" - }, - "transforms": { - "description": "Ordered list of generic transform operations applied to decoded pixels.", - "items": { - "$ref": "#/$defs/ImageTransform" - }, - "type": "array" - } - }, - "required": [ - "outputs" - ], - "type": "object" - }, - "ImageSizeSpec": { - "anyOf": [ - { - "description": "A single edge length applied to both dimensions.", - "format": "uint32", - "minimum": 0, - "type": "integer" - }, - { - "description": "Explicit width and height.", - "properties": { - "height": { - "description": "Target height in pixels.", - "format": "uint32", - "minimum": 0, - "type": "integer" - }, - "width": { - "description": "Target width in pixels.", - "format": "uint32", - "minimum": 0, - "type": "integer" - } - }, - "required": [ - "width", - "height" - ], - "type": "object" - } - ], - "description": "A square size or an explicit width/height for an image transform." - }, - "ImageTransform": { - "description": "One generic image transform operation.\n\n`op` selects the operation from a generic vocabulary; the remaining fields\nare the parameters that operation reads (only the relevant ones are set).\nEvery parameter is model DATA — concrete sizes, patch sizes, means, and so on\nlive in a model's fixture, never as constants baked into this schema.", - "properties": { - "canvas_pad_value": { - "description": "RGB canvas fill value applied before dynamic tiling.", - "format": "double", - "type": [ - "number", - "null" - ] - }, - "channel_order": { - "description": "Flattened patch feature order (`channels_first` or `channels_last`).", - "minLength": 1, - "type": [ - "string", - "null" - ] - }, - "coordinate_order": { - "description": "Patch-coordinate component order (`yx` or `xy`).", - "minLength": 1, - "type": [ - "string", - "null" - ] - }, - "flatten": { - "description": "Whether `patchify` flattens each patch into a single feature vector.", - "type": [ - "boolean", - "null" - ] - }, - "include_thumbnail": { - "description": "Whether a `tile` operation also emits a global thumbnail tile.", - "type": [ - "boolean", - "null" - ] - }, - "inputs": { - "description": "Named values consumed by this transform.\n\nAbsent means the operation consumes the immediately preceding value.\nExplicit names allow branching programs without tensor-name heuristics.", - "items": { - "minLength": 1, - "type": "string" - }, - "minItems": 1, - "type": [ - "array", - "null" - ] - }, - "interpolation": { - "description": "Interpolation filter for a `resize` operation — generic string data.", - "minLength": 1, - "type": [ - "string", - "null" - ] - }, - "mask_patch_size": { - "description": "Pixel edge represented by one validity-mask cell.", - "format": "uint", - "minimum": 1, - "type": [ - "integer", - "null" - ] - }, - "max_patches": { - "description": "Maximum number of spatial patches for a patch-budget resize.", - "format": "uint", - "minimum": 1, - "type": [ - "integer", - "null" - ] - }, - "max_pixels": { - "description": "Maximum pixel area for an aspect-preserving `pixel_area` resize.", - "format": "uint", - "minimum": 1, - "type": [ - "integer", - "null" - ] - }, - "max_tiles": { - "description": "Maximum number of local tiles for a `tile` operation.", - "format": "uint", - "minimum": 1, - "type": [ - "integer", - "null" - ] - }, - "mean": { - "description": "Per-channel mean for a `normalize` operation (length is model data).", - "items": { - "format": "float", - "type": "number" - }, - "type": [ - "array", - "null" - ] - }, - "merge_size": { - "description": "Spatial patch-group edge controlling packed patch traversal order.", - "format": "uint", - "minimum": 1, - "type": [ - "integer", - "null" - ] - }, - "min_pixels": { - "description": "Minimum pixel area for an aspect-preserving `pixel_area` resize.", - "format": "uint", - "minimum": 1, - "type": [ - "integer", - "null" - ] - }, - "mode": { - "description": "Resize/crop mode (e.g. `pad`, `crop`, `stretch`) — generic string data.", - "minLength": 1, - "type": [ - "string", - "null" - ] - }, - "op": { - "$ref": "#/$defs/ImageTransformOp", - "description": "Generic operation selector (e.g. `resize`, `normalize`, `patchify`)." - }, - "outputs": { - "description": "Named values produced by this transform.\n\nThese names are processor-local data. Final graph bindings select them\nthrough `ImageOutputBinding::source`.", - "items": { - "minLength": 1, - "type": "string" - }, - "minItems": 1, - "type": [ - "array", - "null" - ] - }, - "pad_value": { - "description": "Fill value for a `pad` operation, or sentinel for padded coordinates.", - "format": "double", - "type": [ - "number", - "null" - ] - }, - "patch_size": { - "description": "Edge length of a square patch for a `patchify` operation.", - "format": "uint", - "minimum": 1, - "type": [ - "integer", - "null" - ] - }, - "pooling_kernel_size": { - "description": "Spatial pooling edge used when resolving a patch-budget resize.", - "format": "uint", - "minimum": 1, - "type": [ - "integer", - "null" - ] - }, - "scale": { - "description": "Scalar multiplier for a `rescale` operation.", - "format": "double", - "type": [ - "number", - "null" - ] - }, - "size": { - "anyOf": [ - { - "$ref": "#/$defs/ImageSizeSpec" - }, - { - "type": "null" - } - ], - "description": "Target size for a `resize` operation." - }, - "size_multiple": { - "description": "Required divisibility of both resized dimensions.", - "format": "uint", - "minimum": 1, - "type": [ - "integer", - "null" - ] - }, - "std": { - "description": "Per-channel standard deviation for a `normalize` operation.", - "items": { - "format": "float", - "type": "number" - }, - "type": [ - "array", - "null" - ] - }, - "target_length": { - "description": "Exact first-axis length produced by a `pad` operation.", - "format": "uint", - "minimum": 1, - "type": [ - "integer", - "null" - ] - }, - "temporal_patch_size": { - "description": "Number of identical temporal frames packed into each spatial patch.", - "format": "uint", - "minimum": 1, - "type": [ - "integer", - "null" - ] - }, - "thumbnail_interpolation": { - "description": "Interpolation filter used specifically for a global thumbnail.", - "minLength": 1, - "type": [ - "string", - "null" - ] - }, - "thumbnail_order": { - "anyOf": [ - { - "$ref": "#/$defs/ThumbnailOrder" - }, - { - "type": "null" - } - ], - "description": "Ordering of a global thumbnail relative to local tiles." - }, - "tile_size": { - "description": "Edge length of a square tile for a `tile` operation.", - "format": "uint", - "minimum": 1, - "type": [ - "integer", - "null" - ] - } - }, - "required": [ - "op" - ], - "type": "object" - }, - "ImageTransformOp": { - "description": "Generic image transform-operation vocabulary.", - "oneOf": [ - { - "enum": [ - "decode", - "decode_rgb", - "convert_rgb", - "resize", - "rescale", - "normalize", - "tile", - "flatten", - "patchify", - "pad", - "emit_original_size", - "emit_transformed_size", - "emit_validity_mask", - "emit_patch_coordinates", - "emit_grid_coordinates" - ], - "title": "Known standard value" - }, - { - "not": { - "enum": [ - "decode", - "decode_rgb", - "convert_rgb", - "resize", - "rescale", - "normalize", - "tile", - "flatten", - "patchify", - "pad", - "emit_original_size", - "emit_transformed_size", - "emit_validity_mask", - "emit_patch_coordinates", - "emit_grid_coordinates" - ] - }, - "title": "Forward-compatible extension value", - "type": "string" - } - ], - "type": "string" - }, - "KeySequenceLengthsSpec": { - "description": "Explicit compatibility rules for attention key-sequence-length metadata.", - "properties": { - "scalar_broadcast": { - "anyOf": [ - { - "$ref": "#/$defs/SequenceLengthScalarBroadcast" - }, - { - "type": "null" - } - ], - "description": "Optional scalar compatibility. `unit_batch` authorizes a contiguous\nrank-0 one-element `int32` tensor only when the attention batch is one." - } - }, - "type": "object" - }, - "KvAxisStrides": { - "description": "Symbolic element stride of each of the four logical KV axes.\n\nThe stride of an axis is the product of the runtime dimensions in its factor\nlist; an **empty** list means unit stride (the innermost, contiguous axis).\nThe innermost axis of every layout the converted kernels honor is\n`head_dim`, whose stride is `1` (empty), because the fp16 read vectorizes\n`head_dim` as `half2` and the fused write addresses it as `dst + d`.\n\nThe two historical layouts map onto this as:\n\n| axis | head-major BNSH | seq-major BSNH |\n|----------|------------------------|-----------------------|\n| batch | `kv_heads·seq·head_dim`| `seq·kv_heads·head_dim`|\n| head | `seq·head_dim` | `head_dim` |\n| seq | `head_dim` | `kv_heads·head_dim` |\n| head_dim | `1` | `1` |", - "properties": { - "batch": { - "default": [], - "description": "Factors of the batch-axis stride.", - "items": { - "$ref": "#/$defs/KvStrideDim" - }, - "type": "array" - }, - "head": { - "default": [], - "description": "Factors of the KV-head-axis stride.", - "items": { - "$ref": "#/$defs/KvStrideDim" - }, - "type": "array" - }, - "head_dim": { - "default": [], - "description": "Factors of the head-dim-axis stride. Unit (empty) for every honored\nlayout.", - "items": { - "$ref": "#/$defs/KvStrideDim" - }, - "type": "array" - }, - "seq": { - "default": [], - "description": "Factors of the sequence-axis (per-token) stride.", - "items": { - "$ref": "#/$defs/KvStrideDim" - }, - "type": "array" - } - }, - "type": "object" - }, - "KvCacheLayout": { - "anyOf": [ - { - "$ref": "#/$defs/KvNamedLayout", - "description": "A readable shorthand for a standard layout. Deserializes from the strings\n`\"head_major_bnsh\"` and `\"seq_major_bsnh\"`; expands to explicit strides\nvia [`KvCacheLayout::resolve_strides`]." - }, - { - "$ref": "#/$defs/KvStrideDescriptor", - "description": "A fully explicit stride descriptor for layouts the named forms cannot\nexpress." - } - ], - "description": "Physical memory layout of a backend's KV cache tensors, as a stride\ndescriptor.\n\nThis is a **per-backend capability**, not a cross-backend constant: the two\nbackends own their KV buffers independently and never read each other's KV\nbytes, so they may store the cache differently. The ONNX Runtime backend\nrequires head-major BNSH (`[batch, kv_heads, seq, head_dim]`) because ORT's\nGroupQueryAttention past/present is BNSH on every dispatch path (Flash,\ncuDNN SDPA, memory-efficient, XQA). The native backend additionally supports\nseq-major BSNH (`[batch, seq, kv_heads, head_dim]`), which makes each token's\nlive prefix contiguous across heads — shrinking the VMM granule floor by the\n`kv_heads` factor, removing growth-triggered graph re-capture (the append\nstride is sequence-length independent), and making page-level prefix sharing\n(#777) practical. Absent preserves the historical head-major behavior.\n\nLayout preference is per-EP and per-platform rather than a global constant,\nand a JIT backend compiles a specialized kernel per descriptor, so this is a\ndescriptor rather than a closed enum: a raw stride tuple is unreadable, so\nthe common cases are still nameable (`head_major_bnsh`, `seq_major_bsnh`)\nwhile an explicit [`KvStrideDescriptor`] expresses anything the named forms\ncannot (e.g. a token-major view).\n\nOn-device, the native backend selects the layout by stamping the `kv_layout`\nattribute (`0` = BNSH, `1` = BSNH) on its GroupQueryAttention nodes; the\nCUDA EP honors it on the fused fp16 single-token decode pair. Seq-major is\nonly enabled end-to-end once the prefill (flash) read is also converted, so\nthe two never disagree about how a shared cache is physically laid out." - }, - "KvCacheOperations": { - "description": "Operational guarantees for mutable KV-cache state.", - "properties": { - "checkpoint_serializable": { - "description": "Whether checkpoints can be serialized for suspend/resume or migration.", - "type": [ - "boolean", - "null" - ] - }, - "fork_precision_policy": { - "anyOf": [ - { - "$ref": "#/$defs/ForkPrecisionPolicy" - }, - { - "type": "null" - } - ], - "description": "Precision policy for a copy-on-write fork, such as `inherit` or `highest`." - }, - "rewind_safe": { - "description": "Whether truncating cache state to an earlier token position is correctness-preserving.", - "type": [ - "boolean", - "null" - ] - } - }, - "type": "object" - }, - "KvCacheSpec": { - "description": "KV-cache storage, precision tolerance, and operational guarantees.", - "properties": { - "native_dtype": { - "anyOf": [ - { - "$ref": "#/$defs/DType" - }, - { - "type": "null" - } - ], - "description": "Native KV scalar dtype produced by the model before optional compression." - }, - "operations": { - "anyOf": [ - { - "$ref": "#/$defs/KvCacheOperations" - }, - { - "type": "null" - } - ], - "description": "Cache mutation and persistence operations known to be safe for this model." - }, - "quantization_tolerance": { - "anyOf": [ - { - "$ref": "#/$defs/KvQuantTolerance" - }, - { - "type": "null" - } - ], - "description": "Independent precision tolerance for key and value tensors." - }, - "sensitive_layers": { - "description": "Layer indices that should retain high precision; negative indices count from the end.", - "items": { - "format": "int32", - "type": "integer" - }, - "type": [ - "array", - "null" - ] - } - }, - "type": "object" - }, - "KvComponentTolerance": { - "description": "Quantization tolerance for one KV-cache component.", - "properties": { - "default": { - "anyOf": [ - { - "$ref": "#/$defs/DType" - }, - { - "type": "null" - } - ], - "description": "Default minimum acceptable scalar dtype for this component." - }, - "per_layer": { - "description": "Layer-specific minimum-precision overrides.", - "items": { - "$ref": "#/$defs/LayerPrecisionOverride" - }, - "type": [ - "array", - "null" - ] - }, - "quantization_axis": { - "anyOf": [ - { - "$ref": "#/$defs/QuantizationAxis" - }, - { - "type": "null" - } - ], - "description": "Quantization scaling axis, such as `per_tensor`, `per_channel`, or `per_token`." - } - }, - "type": "object" - }, - "KvNamedLayout": { - "description": "The named, human-readable KV cache layouts.", - "oneOf": [ - { - "const": "head_major_bnsh", - "description": "Head-major BNSH `[batch, kv_heads, seq, head_dim]`. ORT-compatible; the\ndefault for both backends.", - "type": "string" - }, - { - "const": "seq_major_bsnh", - "description": "Seq-major BSNH `[batch, seq, kv_heads, head_dim]`. Native backend only.", - "type": "string" - } - ] - }, - "KvOwnership": { - "description": "Ownership model for a graph's KV cache inputs.", - "oneOf": [ - { - "const": "owned", - "description": "The graph consumes past KV and emits replacement/extended present KV.", - "type": "string" - }, - { - "const": "shared", - "description": "The graph reads references to KV owned and advanced by another decoder.", - "type": "string" - } - ] - }, - "KvPagingMode": { - "enum": [ - "none", - "paged" - ], - "type": "string" - }, - "KvQuantTolerance": { - "description": "Precision tolerance for key and value cache components.", - "properties": { - "key": { - "anyOf": [ - { - "$ref": "#/$defs/KvComponentTolerance" - }, - { - "type": "null" - } - ], - "description": "Key-cache precision tolerance." - }, - "value": { - "anyOf": [ - { - "$ref": "#/$defs/KvComponentTolerance" - }, - { - "type": "null" - } - ], - "description": "Value-cache precision tolerance." - } - }, - "type": "object" - }, - "KvServiceContract": { - "additionalProperties": false, - "properties": { - "allocation": { - "$ref": "#/$defs/SlotAllocationMode" - }, - "compaction": { - "default": false, - "type": "boolean" - }, - "paging": { - "$ref": "#/$defs/KvPagingMode" - } - }, - "required": [ - "paging", - "allocation" - ], - "type": "object" - }, - "KvStrideDescriptor": { - "description": "A fully explicit KV-cache stride descriptor.\n\nThis is the general form the two named layouts expand into, and the shape a\nfuture layout (e.g. token-major) is expressed in without adding an enum\nvariant. The `reservation_*` fields describe a binding that is a **view into\na larger reservation** rather than the owner of its whole buffer:\ntoken-major stores every layer's tokens in one reservation and hands each\n`(layer, side)` a sub-view, so its per-token (seq) stride is taken over the\nreservation's total token count and its data starts at a non-zero offset.\nBoth historical layouts are whole-buffer bindings: `offset == 0` and no\nreservation override.", - "properties": { - "reservation_offset_elements": { - "description": "Element offset of this binding's first element within the reservation it\nviews. `0` for a binding that owns its whole buffer — the only case the\nconverted kernels honor today.", - "format": "uint64", - "minimum": 0, - "type": "integer" - }, - "reservation_seq_slots": { - "description": "Sequence-axis extent, in token slots, of the reservation this binding\nviews when the reservation is larger than the binding's own\n`cache_capacity`. Absent means the binding spans its own capacity (a\nwhole-buffer binding). Present expresses a token-major view whose seq\nstride collapses the per-`(layer, side)` buffer boundary; not honored by\nthe converted path yet.", - "format": "uint64", - "minimum": 0, - "type": [ - "integer", - "null" - ] - }, - "strides": { - "$ref": "#/$defs/KvAxisStrides", - "description": "Symbolic stride of each logical axis." - } - }, - "required": [ - "strides" - ], - "type": "object" - }, - "KvStrideDim": { - "description": "A runtime KV-cache dimension that an axis stride can be a multiple of.\n\nAbsolute element strides are a serving-time property — they depend on the\n`cache_capacity` a runtime picks — so metadata cannot store them as numbers.\nA stride is therefore stored **symbolically**, as the (unordered) set of\nruntime dimensions it multiplies. The concrete element stride of an axis is\nthe product of the sizes of the dimensions in its factor list.", - "oneOf": [ - { - "const": "kv_heads", - "description": "Number of KV heads (`kv_heads` / `N`).", - "type": "string" - }, - { - "const": "seq_capacity", - "description": "Sequence capacity of the growing axis (`cache_capacity` / `S`).", - "type": "string" - }, - { - "const": "head_dim", - "description": "Per-token head width (`head_dim` / `H`).", - "type": "string" - } - ] - }, - "KvUpdateKind": { - "description": "Paired KV-cache update-semantics vocabulary.", - "oneOf": [ - { - "enum": [ - "append", - "shared_buffer" - ], - "title": "Known standard value" - }, - { - "not": { - "enum": [ - "append", - "shared_buffer" - ] - }, - "title": "Forward-compatible extension value", - "type": "string" - } - ], - "type": "string" - }, - "LayerPrecisionOverride": { - "description": "Minimum precision required by a set of model layers.", - "properties": { - "layers": { - "description": "Non-empty layer-index list; negative indices count from the final layer.", - "items": { - "format": "int32", - "type": "integer" - }, - "minItems": 1, - "type": "array" - }, - "min_precision": { - "$ref": "#/$defs/DType", - "description": "Minimum acceptable scalar dtype for the listed layers." - } - }, - "required": [ - "layers", - "min_precision" - ], - "type": "object" - }, - "LoopStatePair": { - "description": "One fixed-shape loop-carried recurrent-state port pair.\n\nGeneric and architecture-neutral: the runtime zero/other-initializes `input`\non the first step, runs the graph, and copies `output` back into `input` for\nthe next step (`replace` update). This models any fixed recurrent tensor\n(convolution state, linear-attention recurrent state, and so on) without\nreferencing a model family. It is intentionally distinct from growing or\nshared-buffer KV cache, which is declared through `kv_inputs`/`kv_outputs`\nand `kv_update`.", - "properties": { - "init": { - "$ref": "#/$defs/StateInitKind", - "description": "How `input` is initialized before the first step (e.g. `zeros`)." - }, - "input": { - "description": "Graph input port that receives the carried state for this step.", - "minLength": 1, - "type": "string" - }, - "output": { - "description": "Graph output port that produces the next-step state.", - "minLength": 1, - "type": "string" - }, - "update": { - "$ref": "#/$defs/StateUpdateKind", - "description": "How `output` becomes the next step's `input` (fixed state uses `replace`)." - } - }, - "required": [ - "input", - "output", - "init", - "update" - ], - "type": "object" - }, - "MixtureOfExpertsSpec": { - "description": "Explicit sparse mixture-of-experts structure and graph representation.", - "properties": { - "activation": { - "description": "Expert FFN activation name, such as `silu`.", - "minLength": 1, - "type": "string" - }, - "expert_intermediate_size": { - "description": "Intermediate width of each routed expert FFN.", - "format": "uint", - "minimum": 1, - "type": "integer" - }, - "experts_per_token": { - "description": "Number of routed experts selected for each token.", - "format": "uint", - "minimum": 1, - "type": "integer" - }, - "representation": { - "$ref": "#/$defs/MoERepresentation", - "description": "Expert graph representation: `dense_fallback`, `moe`, or `qmoe`." - }, - "routed_expert_count": { - "description": "Number of independently routed experts.", - "format": "uint", - "minimum": 1, - "type": "integer" - }, - "router": { - "$ref": "#/$defs/MoERouterSpec", - "description": "Router scoring, selection, normalization, and scaling semantics." - }, - "shared_expert_count": { - "description": "Number of dense shared experts evaluated for every token.", - "format": "uint", - "minimum": 0, - "type": "integer" - }, - "shared_expert_intermediate_size": { - "description": "Total intermediate width of the always-on shared-expert FFN.", - "format": "uint", - "minimum": 0, - "type": "integer" - } - }, - "required": [ - "representation", - "routed_expert_count", - "shared_expert_count", - "experts_per_token", - "expert_intermediate_size", - "shared_expert_intermediate_size", - "activation", - "router" - ], - "type": "object" - }, - "MoEGroupScore": { - "description": "Group-scoring reduction vocabulary.", - "oneOf": [ - { - "enum": [ - "maximum", - "top_2_sum" - ], - "title": "Known standard value" - }, - { - "not": { - "enum": [ - "maximum", - "top_2_sum" - ] - }, - "title": "Forward-compatible extension value", - "type": "string" - } - ], - "type": "string" - }, - "MoERepresentation": { - "description": "Sparse expert graph representation vocabulary.", - "oneOf": [ - { - "enum": [ - "dense_fallback", - "moe", - "qmoe" - ], - "title": "Known standard value" - }, - { - "not": { - "enum": [ - "dense_fallback", - "moe", - "qmoe" - ] - }, - "title": "Forward-compatible extension value", - "type": "string" - } - ], - "type": "string" - }, - "MoERouterScoreFunction": { - "description": "Router score-operation vocabulary.", - "oneOf": [ - { - "enum": [ - "softmax", - "sigmoid" - ], - "title": "Known standard value" - }, - { - "not": { - "enum": [ - "softmax", - "sigmoid" - ] - }, - "title": "Forward-compatible extension value", - "type": "string" - } - ], - "type": "string" - }, - "MoERouterSelectionMethod": { - "description": "Router expert-selection vocabulary.", - "oneOf": [ - { - "enum": [ - "top_k", - "grouped_top_k", - "sparse_mixer" - ], - "title": "Known standard value" - }, - { - "not": { - "enum": [ - "top_k", - "grouped_top_k", - "sparse_mixer" - ] - }, - "title": "Forward-compatible extension value", - "type": "string" - } - ], - "type": "string" - }, - "MoERouterSpec": { - "allOf": [ - { - "if": { - "properties": { - "selection_method": { - "const": "grouped_top_k" - } - }, - "required": [ - "selection_method" - ] - }, - "then": { - "required": [ - "group_count", - "groups_per_token", - "group_score" - ] - } - } - ], - "description": "Explicit router semantics, kept separate from expert FFN execution.", - "properties": { - "group_count": { - "description": "Number of expert groups considered by grouped selection.", - "format": "uint", - "minimum": 1, - "type": [ - "integer", - "null" - ] - }, - "group_score": { - "anyOf": [ - { - "$ref": "#/$defs/MoEGroupScore" - }, - { - "type": "null" - } - ], - "description": "Reduction used to score a group before group TopK." - }, - "groups_per_token": { - "description": "Number of groups retained per token by grouped selection.", - "format": "uint", - "minimum": 1, - "type": [ - "integer", - "null" - ] - }, - "normalize_weights": { - "description": "Whether selected aggregation weights are normalized to sum to one.", - "type": "boolean" - }, - "scaling_factor": { - "description": "Multiplicative scale applied to final aggregation weights.", - "format": "float", - "minimum": 0.0, - "type": "number" - }, - "score_function": { - "$ref": "#/$defs/MoERouterScoreFunction", - "description": "Elementwise score operation applied to router logits." - }, - "selection_method": { - "$ref": "#/$defs/MoERouterSelectionMethod", - "description": "Expert selection operation applied to the scores." - } - }, - "required": [ - "score_function", - "selection_method", - "normalize_weights", - "scaling_factor" - ], - "type": "object" - }, - "ModelCapabilities": { - "description": "Model properties that are baked into the graph or advertised as configurable.", - "properties": { - "attention": { - "anyOf": [ - { - "$ref": "#/$defs/AttentionConfig" - }, - { - "type": "null" - } - ], - "description": "Attention architecture and dimensions." - }, - "io": { - "anyOf": [ - { - "$ref": "#/$defs/ModelIoSpec" - }, - { - "type": "null" - } - ], - "default": null, - "description": "Explicit graph I/O port bindings for the single-decoder LLM path.\n\nThe runtime binds decode-step inputs and outputs from the declared names.\nA port that is not declared is resolved ONLY from an unambiguous io-shape\nsignal; when the shape is ambiguous the runtime fails with an actionable\nerror naming the exact key to declare, and never guesses from a tensor\nname." - }, - "max_sequence_length": { - "description": "Maximum total sequence length, in tokens.", - "format": "uint", - "minimum": 1, - "type": [ - "integer", - "null" - ] - }, - "mixture_of_experts": { - "anyOf": [ - { - "$ref": "#/$defs/MixtureOfExpertsSpec" - }, - { - "type": "null" - } - ], - "description": "Explicit sparse mixture-of-experts graph and routing contract.\n\nThis describes graph structure, never a model family. Runtimes use the\ndeclared representation and dimensions instead of inferring them from\nnode names, initializer shapes, or architecture strings." - }, - "runtime_configurable": { - "anyOf": [ - { - "$ref": "#/$defs/RuntimeConfigurable" - }, - { - "type": "null" - } - ], - "description": "Features that a serving runtime may configure at load time." - }, - "speculative": { - "anyOf": [ - { - "$ref": "#/$defs/SpeculativeModelInfo" - }, - { - "type": "null" - } - ], - "description": "Built-in draft-head or self-speculative model properties." - }, - "vocab_size": { - "description": "Vocabulary size (rows of the token-embedding / logits table).", - "format": "uint", - "minimum": 1, - "type": [ - "integer", - "null" - ] - } - }, - "type": "object" - }, - "ModelIoSpec": { - "description": "Explicit binding of the graph ports the decode step reads and writes.\n\nEvery field is optional so a model package can declare only the ports its\ngraph exposes. A port left unset is resolved ONLY from an unambiguous\ndtype/shape signal; when the shape cannot disambiguate the port, the runtime\nfails with an actionable error naming the key to declare rather than\ninterpreting a tensor name. A declared port is always authoritative.", - "properties": { - "attention_mask_input": { - "description": "Attention-mask input, if the graph takes one.", - "minLength": 1, - "type": [ - "string", - "null" - ] - }, - "audio_features_input": { - "description": "Raw audio-feature prompt input for an encoder-decoder encoder graph\n(e.g. Whisper `audio_features`, a log-mel `[batch, mels, frames]`\ntensor). Declared on the encoder component; a text encoder-decoder uses\n`token_input` instead.", - "minLength": 1, - "type": [ - "string", - "null" - ] - }, - "cross_kv_inputs": { - "description": "Cross-attention past-KV cache inputs for an encoder-decoder decoder, in\nthe SAME order as `cross_kv_outputs`. These are the encoder-derived KV\ntensors, distinct from the self-attention `kv_inputs`.", - "items": { - "minLength": 1, - "type": "string" - }, - "type": [ - "array", - "null" - ] - }, - "cross_kv_outputs": { - "description": "Cross-attention present-KV cache outputs (produced by the encoder for an\nencoder-decoder model), paired positionally with `cross_kv_inputs`.", - "items": { - "minLength": 1, - "type": "string" - }, - "type": [ - "array", - "null" - ] - }, - "encoder_hidden_states_input": { - "description": "Encoder-hidden-states input for an encoder-decoder (cross-attention)\ndecoder graph (e.g. `encoder_hidden_states`).", - "minLength": 1, - "type": [ - "string", - "null" - ] - }, - "hidden_output": { - "description": "Per-token hidden-state output for embedding / VLM hidden extraction, if\nthe graph exposes a distinct hidden output (e.g. `last_hidden_state`).", - "minLength": 1, - "type": [ - "string", - "null" - ] - }, - "inputs_embeds_input": { - "description": "Pre-embedded / routed sequence input (e.g. `inputs_embeds`).\n\nMay be declared alongside `token_input` (see its documentation): a graph\nthat consumes both a raw token input and one or more routed sequence\ninputs is explicitly permitted.", - "minLength": 1, - "type": [ - "string", - "null" - ] - }, - "kv_inputs": { - "description": "Past-KV cache inputs, in the SAME order as `kv_outputs` (positional\npairing). Length must match `kv_outputs`.", - "items": { - "minLength": 1, - "type": "string" - }, - "type": [ - "array", - "null" - ] - }, - "kv_layout": { - "anyOf": [ - { - "$ref": "#/$defs/KvCacheLayout" - }, - { - "type": "null" - } - ], - "description": "Physical layout of this backend's KV cache tensors, as a stride\ndescriptor. Accepts a readable named layout (`head_major_bnsh` or\n`seq_major_bsnh`) or a fully explicit [`KvStrideDescriptor`]. This is a\nper-backend capability — each backend owns its KV buffers and never reads\nthe other's KV bytes — so the ORT backend stays head-major while the\nnative backend may declare seq-major. Absent preserves the historical\nhead-major (BNSH) behavior. See [`KvCacheLayout`]." - }, - "kv_outputs": { - "description": "Present-KV cache outputs, paired positionally with `kv_inputs`.", - "items": { - "minLength": 1, - "type": "string" - }, - "type": [ - "array", - "null" - ] - }, - "kv_ownership": { - "anyOf": [ - { - "$ref": "#/$defs/KvOwnership" - }, - { - "type": "null" - } - ], - "description": "Whether this graph owns past/present KV state or reads target-owned KV.\n\nAbsent preserves the historical `owned` behavior." - }, - "kv_update": { - "anyOf": [ - { - "$ref": "#/$defs/KvUpdateKind" - }, - { - "type": "null" - } - ], - "description": "How the paired `kv_inputs`/`kv_outputs` cache tensors evolve each step.\n\nThis declares GROWING/append versus fixed shared-buffer cache semantics\nexplicitly, and is deliberately kept separate from `state_pairs` (which\ndescribes fixed recurrent tensors that are wholly REPLACED). The KV pair\nlists are the authoritative sparse layer ports: the runtime binds exactly\nthe ports named in `kv_inputs`/`kv_outputs` and never expands them from a\ntotal layer count. Absent means the historical growing-cache default." - }, - "logits_output": { - "description": "Logits output.", - "minLength": 1, - "type": [ - "string", - "null" - ] - }, - "optional_inputs": { - "additionalProperties": { - "$ref": "#/$defs/OptionalInputSpec" - }, - "description": "Optional graph inputs and their explicit absent-value contracts, keyed by\nthe real ONNX input port name.", - "type": "object" - }, - "position_ids_input": { - "description": "Position-ids input, if the graph takes one.", - "minLength": 1, - "type": [ - "string", - "null" - ] - }, - "sequence_source": { - "anyOf": [ - { - "$ref": "#/$defs/SequenceInputKind" - }, - { - "type": "null" - } - ], - "description": "Which declared sequence port drives autoregressive execution.\n\nAbsent preserves the historical `token_ids` behavior. Declaring\n`inputs_embeds` requires `inputs_embeds_input`; declaring `token_ids`\nrequires `token_input`." - }, - "state_pairs": { - "description": "Fixed-shape loop-carried recurrent state ports, distinct from KV cache.\n\nEach pair binds an input port to its matching output port and declares\nhow the input is initialized and how the output feeds the next step\n(`replace` semantics for fixed recurrent tensors). These are neither KV\ncache nor fixed conditioning; the sparse set of state ports comes from\nthis declared list, never expanded from a layer count.", - "items": { - "$ref": "#/$defs/LoopStatePair" - }, - "minItems": 1, - "type": [ - "array", - "null" - ] - }, - "static_cache": { - "anyOf": [ - { - "$ref": "#/$defs/StaticCacheIoSpec" - }, - { - "type": "null" - } - ], - "description": "Explicit port binding for a fixed-buffer TensorScatter static KV cache.\n\nA static-cache decoder scatters each step's K/V into pre-allocated,\nfixed-length buffers via an integer write-index vector and a non-pad\nsequence-length vector, rather than growing/appending a cache. These\ncontrol ports are integer vectors and are therefore SHAPE-indistinguish-\nable from one another, so shape cannot disambiguate them: the ABI must be\ndeclared explicitly. When present, this spec is authoritative and the\nruntime binds exactly these ports. When absent, a graph that exposes the\nscatter ABI is REJECTED with an actionable error naming this key rather\nthan having its integer control ports guessed by name." - }, - "token_input": { - "description": "Token-id input (e.g. `input_ids`).\n\nA graph MAY declare this together with `inputs_embeds_input`: some fused\ndecoders consume a raw token stream AND a routed pre-embedded sequence in\nthe same forward pass. The two are not mutually exclusive; declaring both\nis a valid, explicit contract.", - "minLength": 1, - "type": [ - "string", - "null" - ] - } - }, - "type": "object" - }, - "MtpHiddenLayout": { - "description": "Layout of the target state consumed by an MTP sidecar.", - "oneOf": [ - { - "const": "BSH", - "description": "`[batch, sequence, hidden]` legacy layout.", - "type": "string" - }, - { - "const": "BSHC", - "description": "`[batch, sequence, hc_mult, hidden]` Mobius Hyper-Connection layout.", - "type": "string" - } - ] - }, - "MtpKvMode": { - "description": "Lifetime declared for an MTP sidecar's private KV state.", - "oneOf": [ - { - "const": "proposal_local", - "description": "Reset sidecar KV at every target verification iteration.", - "type": "string" - }, - { - "const": "accepted_prefix", - "description": "Retain only KV corresponding to the accepted draft prefix.", - "type": "string" - } - ] - }, - "MtpTargetInitializer": { - "description": "Exact target-model initializer reference used by an MTP sidecar.", - "properties": { - "name": { - "description": "Exact initializer name in the target ONNX graph.", - "type": "string" - }, - "source": { - "$ref": "#/$defs/MtpWeightSource", - "description": "Initializer ownership source. The Phase-1 contract requires\n`target_initializer`." - } - }, - "required": [ - "source", - "name" - ], - "type": "object" - }, - "MtpWeightSource": { - "description": "Ownership source for an MTP shared weight.", - "oneOf": [ - { - "const": "target_initializer", - "description": "Borrow the named initializer from the target model package.", - "type": "string" - } - ] - }, - "OptionalInputSpec": { - "description": "Presence and absent-value contract for one optional graph input.", - "properties": { - "absent": { - "$ref": "#/$defs/AbsentInputSpec", - "description": "Tensor value supplied when the presence key is absent." - }, - "presence": { - "description": "Opaque, non-empty request presence key; not a port or model name.", - "minLength": 1, - "type": "string" - } - }, - "required": [ - "presence", - "absent" - ], - "type": "object" - }, - "OutputStage": { - "enum": [ - "pre_adapter", - "post_adapter" - ], - "type": "string" - }, - "PerformanceHints": { - "description": "Publisher-provided speculative decoding performance guidance.", - "properties": { - "expected_acceptance_rate": { - "description": "Expected fraction of proposed tokens accepted, from 0.0 through 1.0.", - "format": "float", - "maximum": 1.0, - "minimum": 0.0, - "type": [ - "number", - "null" - ] - }, - "optimal_k": { - "description": "Recommended number of draft tokens per verification step.", - "format": "uint", - "minimum": 1, - "type": [ - "integer", - "null" - ] - } - }, - "type": "object" - }, - "PipelineSpec": { - "additionalProperties": false, - "description": "Executable package described by the universal typed workflow IR.", - "properties": { - "workflow": { - "$ref": "#/$defs/WorkflowSpec", - "description": "Required component-centric SSA workflow." - } - }, - "required": [ - "workflow" - ], - "type": "object" - }, - "Precision": { - "description": "Weight precision and quantization-recipe vocabulary.", - "oneOf": [ - { - "enum": [ - "float32", - "fp32", - "float16", - "fp16", - "bfloat16", - "bf16", - "float8_e4m3fn", - "float8_e5m2", - "int8", - "int4", - "int4_group128" - ], - "title": "Known standard value" - }, - { - "not": { - "enum": [ - "float32", - "fp32", - "float16", - "fp16", - "bfloat16", - "bf16", - "float8_e4m3fn", - "float8_e5m2", - "int8", - "int4", - "int4_group128" - ] - }, - "title": "Forward-compatible extension value", - "type": "string" - } - ], - "type": "string" - }, - "PreprocessingSpec": { - "description": "Declared, architecture-neutral input preprocessing programs.", - "properties": { - "image": { - "anyOf": [ - { - "$ref": "#/$defs/ImagePreprocessingProgram" - }, - { - "type": "null" - } - ], - "description": "Typed image preprocessing transform program and its named tensor outputs." - } - }, - "type": "object" - }, - "ProposalTopology": { - "description": "Speculative proposal-topology vocabulary.", - "oneOf": [ - { - "enum": [ - "linear", - "tree" - ], - "title": "Known standard value" - }, - { - "not": { - "enum": [ - "linear", - "tree" - ] - }, - "title": "Forward-compatible extension value", - "type": "string" - } - ], - "type": "string" - }, - "ProposalType": { - "description": "Speculator proposal architecture.\n\nKnown spellings are enumerated in the generated schema while unknown\nstrings remain valid to preserve forward compatibility.", - "oneOf": [ - { - "enum": [ - "eagle", - "eagle3", - "eagle-3", - "peagle", - "p-eagle", - "mtp", - "dflash", - "d-flash", - "shared_kv", - "shared-kv" - ], - "title": "Known standard value" - }, - { - "not": { - "enum": [ - "eagle", - "eagle3", - "eagle-3", - "peagle", - "p-eagle", - "mtp", - "dflash", - "d-flash", - "shared_kv", - "shared-kv" - ] - }, - "title": "Forward-compatible extension value", - "type": "string" - } - ], - "type": "string" - }, - "QuantizationAxis": { - "description": "Quantization scaling-axis vocabulary.", - "oneOf": [ - { - "enum": [ - "per_tensor", - "per_channel", - "per_token", - "per_head" - ], - "title": "Known standard value" - }, - { - "not": { - "enum": [ - "per_tensor", - "per_channel", - "per_token", - "per_head" - ] - }, - "title": "Forward-compatible extension value", - "type": "string" - } - ], - "type": "string" - }, - "QuantizationIntent": { - "description": "Runtime-independent model-weight quantization intent.", - "properties": { - "default": { - "anyOf": [ - { - "$ref": "#/$defs/Precision" - }, - { - "type": "null" - } - ], - "description": "Default precision or quantization recipe for model weights." - }, - "overrides": { - "description": "Layer- or component-specific precision overrides.", - "items": { - "$ref": "#/$defs/QuantizationOverride" - }, - "type": [ - "array", - "null" - ] - } - }, - "type": "object" - }, - "QuantizationOverride": { - "description": "Precision override for selected layers or a named graph component.", - "properties": { - "component": { - "description": "Logical component path, for example `attention.qk` or `lm_head`.", - "minLength": 1, - "type": [ - "string", - "null" - ] - }, - "layers": { - "description": "Layer indices to which the override applies; negative indices count from the end.", - "items": { - "format": "int32", - "type": "integer" - }, - "type": [ - "array", - "null" - ] - }, - "precision": { - "$ref": "#/$defs/Precision", - "description": "Required precision or quantization recipe." - } - }, - "required": [ - "precision" - ], - "type": "object" - }, - "RuntimeConfigurable": { - "description": "Features whose concrete settings may be selected by the runtime.", - "properties": { - "chunked_prefill": { - "anyOf": [ - { - "$ref": "#/$defs/ChunkedPrefillConfig" - }, - { - "type": "null" - } - ], - "description": "Chunked-prefill support and preferred chunk size." - }, - "continuous_batching": { - "description": "Whether continuous batching may be enabled.", - "type": [ - "boolean", - "null" - ] - }, - "kv_cache": { - "anyOf": [ - { - "$ref": "#/$defs/RuntimeKvConfig" - }, - { - "type": "null" - } - ], - "description": "Supported runtime-selectable KV-cache dtypes." - }, - "prefix_cache": { - "description": "Whether prefix caching may be enabled.", - "type": [ - "boolean", - "null" - ] - } - }, - "type": "object" - }, - "RuntimeInputRole": { - "enum": [ - "prompt_text", - "prompt_tokens", - "media", - "max_iterations", - "max_output_tokens", - "seed", - "sampling_temperature", - "sampling_top_k", - "sampling_top_p", - "sampling_min_p", - "constraint", - "session_id" - ], - "type": "string" - }, - "RuntimeKvConfig": { - "description": "Runtime-selectable KV-cache representations.", - "properties": { - "dtype": { - "description": "Non-empty list of supported KV-cache scalar dtypes, in preference order.", - "items": { - "$ref": "#/$defs/DType" - }, - "minItems": 1, - "type": "array" - } - }, - "required": [ - "dtype" - ], - "type": "object" - }, - "ScalarValue": { - "anyOf": [ - { - "type": "boolean" - }, - { - "format": "int64", - "type": "integer" - }, - { - "format": "double", - "type": "number" - }, - { - "type": "string" - } - ] - }, - "SemanticInputRole": { - "oneOf": [ - { - "additionalProperties": false, - "properties": { - "kind": { - "const": "runtime", - "type": "string" - }, - "role": { - "$ref": "#/$defs/RuntimeInputRole" - }, - "version": { - "type": "string" - } - }, - "required": [ - "kind", - "version", - "role" - ], - "type": "object" - }, - { - "additionalProperties": false, - "properties": { - "kind": { - "const": "opaque", - "type": "string" - } - }, - "required": [ - "kind" - ], - "type": "object" - } - ] - }, - "SequenceInputKind": { - "description": "Primary autoregressive sequence source for a decoder or proposer graph.", - "oneOf": [ - { - "const": "token_ids", - "description": "Integer token ids supplied through `token_input`.", - "type": "string" - }, - { - "const": "inputs_embeds", - "description": "Precomputed floating-point embeddings supplied through\n`inputs_embeds_input`.", - "type": "string" - } - ] - }, - "SequenceLengthScalarBroadcast": { - "description": "Permitted scalar compatibility for attention key-sequence lengths.", - "oneOf": [ - { - "const": "unit_batch", - "description": "Interpret one rank-0 value as the canonical one-element vector only for\nan attention batch of exactly one.", - "type": "string" - } - ] - }, - "ServingServiceContract": { - "additionalProperties": false, - "properties": { - "accepted_len": { - "type": [ - "string", - "null" - ] - }, - "active": { - "type": "string" - }, - "done": { - "type": "string" - }, - "kv_service": { - "$ref": "#/$defs/KvServiceContract" - }, - "slot_ids": { - "type": "string" - } - }, - "required": [ - "active", - "done", - "slot_ids", - "kv_service" - ], - "type": "object" - }, - "SessionLeaseContract": { - "additionalProperties": false, - "properties": { - "optimistic_metadata_version": { - "default": false, - "type": "boolean" - }, - "policy": { - "$ref": "#/$defs/SessionMutationPolicy", - "default": "exclusive" - }, - "ttl_seconds": { - "format": "uint64", - "minimum": 0, - "type": [ - "integer", - "null" - ] - } - }, - "type": "object" - }, - "SessionMutationPolicy": { - "enum": [ - "exclusive", - "copy_on_write" - ], - "type": "string" - }, - "ShapeRecurrence": { - "oneOf": [ - { - "additionalProperties": false, - "properties": { - "kind": { - "const": "invariant", - "type": "string" - } - }, - "required": [ - "kind" - ], - "type": "object" - }, - { - "additionalProperties": false, - "properties": { - "axis": { - "format": "uint", - "minimum": 0, - "type": "integer" - }, - "increment": { - "type": "string" - }, - "kind": { - "const": "growing", - "type": "string" - }, - "max": { - "type": "string" - } - }, - "required": [ - "kind", - "axis", - "increment", - "max" - ], - "type": "object" - }, - { - "additionalProperties": false, - "description": "The selected axis may grow or shrink between iterations, but never exceed `max`.", - "properties": { - "axis": { - "format": "uint", - "minimum": 0, - "type": "integer" - }, - "kind": { - "const": "bounded", - "type": "string" - }, - "max": { - "type": "string" - } - }, - "required": [ - "kind", - "axis", - "max" - ], - "type": "object" - } - ] - }, - "SharedKvGroup": { - "description": "One shared-KV binding group for a shared-KV proposer.\n\nA `shared_kv` proposer graph exposes `shared_kv..key` and\n`shared_kv..value` inputs bound to slices of the target model's paged\nKV cache. `target_layers` lists the target KV layer indices feeding this\nslice.", - "properties": { - "key_input": { - "description": "Proposer input receiving this group's shared key cache.", - "type": [ - "string", - "null" - ] - }, - "name": { - "description": "Assistant input prefix, e.g. `sliding_attention` or `full_attention`.", - "type": "string" - }, - "target_key_input": { - "description": "Target decoder past-KV input whose current key cache is referenced.", - "type": [ - "string", - "null" - ] - }, - "target_layers": { - "default": [], - "description": "Target KV layer indices whose cache feeds this shared-KV slice.", - "items": { - "format": "uint", - "minimum": 0, - "type": "integer" - }, - "type": "array" - }, - "target_value_input": { - "description": "Target decoder past-KV input whose current value cache is referenced.", - "type": [ - "string", - "null" - ] - }, - "value_input": { - "description": "Proposer input receiving this group's shared value cache.", - "type": [ - "string", - "null" - ] - } - }, - "required": [ - "name" - ], - "type": "object" - }, - "SlotAllocationMode": { - "enum": [ - "static", - "runtime" - ], - "type": "string" - }, - "SpecialTokens": { - "description": "Special / control token ids declared by a model author.\n\nEvery field is optional; `eos_token_id` is normalized to a list because\nonnxruntime-genai accepts either a scalar or an array for it.", - "properties": { - "bos_token_id": { - "description": "Beginning-of-stream token id.", - "format": "int64", - "type": [ - "integer", - "null" - ] - }, - "decoder_start_token_id": { - "description": "Token an encoder-decoder model starts decoding with, when not `bos`.", - "format": "int64", - "type": [ - "integer", - "null" - ] - }, - "eos_token_id": { - "description": "End-of-stream token ids (one or more).", - "items": { - "format": "int64", - "type": "integer" - }, - "type": [ - "array", - "null" - ] - }, - "image_token_id": { - "description": "Image placeholder token id (VLMs).", - "format": "int64", - "type": [ - "integer", - "null" - ] - }, - "pad_token_id": { - "description": "Padding token id.", - "format": "int64", - "type": [ - "integer", - "null" - ] - }, - "sep_token_id": { - "description": "Separator token id.", - "format": "int64", - "type": [ - "integer", - "null" - ] - }, - "video_token_id": { - "description": "Video placeholder token id (VLMs).", - "format": "int64", - "type": [ - "integer", - "null" - ] - }, - "vision_start_token_id": { - "description": "Vision-segment start token id (VLMs).", - "format": "int64", - "type": [ - "integer", - "null" - ] - } - }, - "type": "object" - }, - "SpeculativeModelInfo": { - "description": "Build-time support for self-contained speculative decoding.", - "properties": { - "has_draft_heads": { - "description": "Whether the exported graph contains Medusa/EAGLE/MTP-style draft heads.", - "type": [ - "boolean", - "null" - ] - }, - "self_speculative_depth": { - "description": "Early-exit layer depth usable for self-speculation.", - "format": "uint", - "minimum": 1, - "type": [ - "integer", - "null" - ] - } - }, - "type": "object" - }, - "SpeculatorConfig": { - "allOf": [ - { - "not": { - "required": [ - "num_speculative_tokens", - "tokens_per_step" - ] - } - } - ], - "description": "Configuration published with a standalone speculative proposer model.", - "oneOf": [ - { - "not": { - "required": [ - "method" - ] - }, - "required": [ - "proposal_type" - ] - }, - { - "not": { - "required": [ - "proposal_type" - ] - }, - "required": [ - "method" - ] - } - ], - "properties": { - "backbone_hidden_size": { - "default": null, - "description": "Target backbone hidden size `H` shared with the proposer.\n\nFor `shared_kv`, `inputs_embeds` is `[B, q, 2*H]` and\n`projected_state` is `[B, q, H]`.", - "format": "uint", - "minimum": 1, - "type": [ - "integer", - "null" - ] - }, - "embedding": { - "anyOf": [ - { - "$ref": "#/$defs/MtpTargetInitializer" - }, - { - "type": "null" - } - ], - "description": "Target embedding initializer shared with the MTP sidecar." - }, - "hc_mult": { - "default": null, - "description": "Number of Hyper-Connection lanes `C`.", - "format": "uint", - "minimum": 1, - "type": [ - "integer", - "null" - ] - }, - "input_embedding": { - "default": null, - "description": "Relative path (from the model directory) to the target model's raw\ninput-token embedding table, as a little-endian f32 matrix in\n`[vocab_size, backbone_hidden_size]` order.\n\nThe `shared_kv` proposer builds each step's `inputs_embeds` as\n`concat(target_input_embedding(last_token), hidden)`, so it must be able\nto look up the target's input embedding of the last drafted/accepted\ntoken. Required for the `shared_kv` proposer.", - "type": [ - "string", - "null" - ] - }, - "io": { - "anyOf": [ - { - "$ref": "#/$defs/ModelIoSpec" - }, - { - "type": "null" - } - ], - "default": null, - "description": "Explicit proposer graph execution contract.\n\nThis uses the same architecture-neutral I/O vocabulary as a target\ndecoder. `sequence_source` selects token ids versus embeddings,\n`kv_ownership` selects private past/present state versus references to\ntarget-owned cache, and the output fields assign semantic roles." - }, - "kv_mode": { - "anyOf": [ - { - "$ref": "#/$defs/MtpKvMode" - }, - { - "type": "null" - } - ], - "description": "Lifetime of the sidecar's KV state.\n\nDefaults to `proposal_local`." - }, - "lm_head": { - "anyOf": [ - { - "$ref": "#/$defs/MtpTargetInitializer" - }, - { - "type": "null" - } - ], - "description": "Target LM-head initializer shared with the MTP sidecar." - }, - "logits_output": { - "default": null, - "description": "Name of the proposer's draft-distribution output.\n\nDefaults to `logits` for `shared_kv`.", - "type": [ - "string", - "null" - ] - }, - "method": { - "allOf": [ - { - "$ref": "#/$defs/ProposalType", - "description": "Proposal architecture used by the speculator.\n\nThe deprecated `method` alias is accepted on input." - } - ], - "deprecated": true, - "description": "Deprecated alias for `proposal_type`." - }, - "model": { - "default": null, - "description": "Relative path (from the model directory) to the proposer ONNX model.\n\nUsed by the `shared_kv` proposer to locate the\nproposer graph. Optional for forward compatibility with proposer\nfamilies that do not ship a standalone model file.", - "type": [ - "string", - "null" - ] - }, - "mtp_hidden_output": { - "default": null, - "description": "Sidecar output projected through the shared target LM head.\n\nDefaults to `mtp_hidden`.", - "type": [ - "string", - "null" - ] - }, - "mtp_state_output": { - "default": null, - "description": "Sidecar recurrent Hyper-Connection state output.\n\nDefaults to `mtp_state`.", - "type": [ - "string", - "null" - ] - }, - "num_speculative_tokens": { - "default": 4, - "description": "Maximum number of tokens proposed per verifier step; defaults to 4.\n\nThe deprecated `tokens_per_step` alias is accepted on input.", - "format": "uint", - "minimum": 1, - "type": "integer" - }, - "projected_state_output": { - "default": null, - "description": "Name of the proposer output threaded forward between steps.\n\nDefaults to `projected_state` for `shared_kv`.", - "type": [ - "string", - "null" - ] - }, - "proposal_type": { - "$ref": "#/$defs/ProposalType", - "description": "Proposal architecture used by the speculator.\n\nThe deprecated `method` alias is accepted on input." - }, - "shared_kv": { - "description": "Shared-KV binding groups consumed by the proposer.\n\nEach group names an assistant input prefix\n(`shared_kv..{key,value}`) and the target KV layer indices whose\ncache feeds that slice. Empty for proposers that own their KV cache.", - "items": { - "$ref": "#/$defs/SharedKvGroup" - }, - "type": "array" - }, - "target_hidden_layout": { - "anyOf": [ - { - "$ref": "#/$defs/MtpHiddenLayout" - }, - { - "type": "null" - } - ], - "description": "Layout of `target_hidden_output`.\n\nMobius MTP sidecars use `BSHC`: batch, sequence, Hyper-Connection lane,\nhidden." - }, - "target_hidden_output": { - "default": null, - "description": "Target decoder output carrying the recurrent MTP seed.\n\nDefaults to `hidden_states` for `mtp`.", - "type": [ - "string", - "null" - ] - }, - "target_hidden_size": { - "default": null, - "description": "Target hidden width `H`.", - "format": "uint", - "minimum": 1, - "type": [ - "integer", - "null" - ] - }, - "tokens_per_step": { - "allOf": [ - { - "default": 4, - "description": "Maximum number of tokens proposed per verifier step; defaults to 4.\n\nThe deprecated `tokens_per_step` alias is accepted on input.", - "format": "uint", - "minimum": 1, - "type": "integer" - } - ], - "deprecated": true, - "description": "Deprecated alias for `num_speculative_tokens`." - }, - "verifier": { - "anyOf": [ - { - "$ref": "#/$defs/SpeculatorVerifier" - }, - { - "type": "null" - } - ], - "description": "Identity of the verifier model against which this proposer was trained." - }, - "vocab_size": { - "default": null, - "description": "Vocabulary size of the proposer's own `logits` output.", - "format": "uint", - "minimum": 1, - "type": [ - "integer", - "null" - ] - } - }, - "required": [], - "type": "object" - }, - "SpeculatorVerifier": { - "description": "Verifier identity embedded in a speculator package.", - "properties": { - "architectures": { - "default": [], - "description": "Verifier architecture class names, in preference order.", - "items": { - "type": "string" - }, - "type": "array" - }, - "name_or_path": { - "description": "HuggingFace-style verifier repository name or local model path.", - "type": [ - "string", - "null" - ] - } - }, - "type": "object" - }, - "StateInitKind": { - "description": "Loop-carried state initialization vocabulary.", - "oneOf": [ - { - "enum": [ - "zeros" - ], - "title": "Known standard value" - }, - { - "not": { - "enum": [ - "zeros" - ] - }, - "title": "Forward-compatible extension value", - "type": "string" - } - ], - "type": "string" - }, - "StateUpdateKind": { - "description": "Loop-carried state update-semantics vocabulary.", - "oneOf": [ - { - "enum": [ - "replace" - ], - "title": "Known standard value" - }, - { - "not": { - "enum": [ - "replace" - ] - }, - "title": "Forward-compatible extension value", - "type": "string" - } - ], - "type": "string" - }, - "StaticCacheIoSpec": { - "description": "Explicit port ABI for a fixed-buffer TensorScatter static KV cache.\n\nDescribes GRAPH STRUCTURE, never a model family. The four per-layer cache\nlists pair positionally per layer and must all have the same length: index\n`i` in each list is layer `i`'s key/value input and updated key/value output.", - "properties": { - "key_cache_inputs": { - "description": "Per-layer static key-cache buffer inputs, positional per layer. Length\nmust equal `value_cache_inputs`, `key_cache_outputs`, and\n`value_cache_outputs`.", - "items": { - "minLength": 1, - "type": "string" - }, - "minItems": 1, - "type": "array" - }, - "key_cache_outputs": { - "description": "Per-layer updated key-cache outputs, paired positionally with\n`key_cache_inputs`.", - "items": { - "minLength": 1, - "type": "string" - }, - "minItems": 1, - "type": "array" - }, - "kv_sequence_length_input": { - "description": "Input port carrying the non-pad KV sequence length (`int` vector).\nShape-indistinguishable from `write_indices_input`, so it too must be\nnamed explicitly.", - "minLength": 1, - "type": "string" - }, - "value_cache_inputs": { - "description": "Per-layer static value-cache buffer inputs, paired positionally with\n`key_cache_inputs`.", - "items": { - "minLength": 1, - "type": "string" - }, - "minItems": 1, - "type": "array" - }, - "value_cache_outputs": { - "description": "Per-layer updated value-cache outputs, paired positionally with\n`value_cache_inputs`.", - "items": { - "minLength": 1, - "type": "string" - }, - "minItems": 1, - "type": "array" - }, - "write_indices_input": { - "description": "Input port carrying the per-token scatter write positions\n(`int` vector). Shape-indistinguishable from other integer control\ninputs, so it must be named explicitly.", - "minLength": 1, - "type": "string" - } - }, - "required": [ - "write_indices_input", - "kv_sequence_length_input", - "key_cache_inputs", - "value_cache_inputs", - "key_cache_outputs", - "value_cache_outputs" - ], - "type": "object" - }, - "StrategyKind": { - "description": "Generic inference-strategy vocabulary.", - "oneOf": [ - { - "enum": [ - "speculative" - ], - "title": "Known standard value" - }, - { - "not": { - "enum": [ - "speculative" - ] - }, - "title": "Forward-compatible extension value", - "type": "string" - } - ], - "type": "string" - }, - "StrategySpec": { - "description": "Generic inference strategy declaration.", - "properties": { - "acceptance": { - "anyOf": [ - { - "$ref": "#/$defs/AcceptanceMethod" - }, - { - "type": "null" - } - ], - "description": "Draft-token acceptance rule." - }, - "draft": { - "anyOf": [ - { - "$ref": "#/$defs/DraftConfig" - }, - { - "type": "null" - } - ], - "description": "Draft-token producer configuration for speculative decoding." - }, - "kind": { - "$ref": "#/$defs/StrategyKind", - "description": "Strategy vocabulary entry, such as `speculative`." - }, - "performance_hints": { - "anyOf": [ - { - "$ref": "#/$defs/PerformanceHints" - }, - { - "type": "null" - } - ], - "description": "Model-publisher performance guidance." - }, - "tokens_per_step": { - "description": "Number of draft tokens attempted per verification step.", - "format": "uint", - "minimum": 1, - "type": [ - "integer", - "null" - ] - }, - "topology": { - "anyOf": [ - { - "$ref": "#/$defs/ProposalTopology" - }, - { - "type": "null" - } - ], - "description": "Proposal topology, such as `linear` or `tree`." - }, - "verify": { - "anyOf": [ - { - "$ref": "#/$defs/VerifyConfig" - }, - { - "type": "null" - } - ], - "description": "Verification configuration for speculative decoding." - } - }, - "required": [ - "kind" - ], - "type": "object" - }, - "StructuredOutputFormat": { - "description": "Structured-output constraint-format vocabulary.", - "oneOf": [ - { - "enum": [ - "json_schema", - "regex", - "context_free_grammar", - "choice" - ], - "title": "Known standard value" - }, - { - "not": { - "enum": [ - "json_schema", - "regex", - "context_free_grammar", - "choice" - ] - }, - "title": "Forward-compatible extension value", - "type": "string" - } - ], - "type": "string" - }, - "StructuredOutputSpec": { - "description": "Structured-output capabilities and model formatting conventions.", - "properties": { - "stop_sequences": { - "description": "Literal token sequences that terminate a structured response.", - "items": { - "type": "string" - }, - "type": [ - "array", - "null" - ] - }, - "supported_formats": { - "description": "Supported constraint formats, such as JSON Schema, regular expressions, or CFGs.", - "items": { - "$ref": "#/$defs/StructuredOutputFormat" - }, - "type": [ - "array", - "null" - ] - }, - "training_format": { - "description": "Format in which the model was trained to emit tool calls or structured values.", - "type": [ - "string", - "null" - ] - } - }, - "type": "object" - }, - "TensorContract": { - "additionalProperties": false, - "description": "Typed tensor contract used at package and component boundaries.", - "properties": { - "dtype": { - "$ref": "#/$defs/TensorDType" - }, - "optional": { - "default": false, - "type": "boolean" - }, - "rank": { - "format": "uint", - "minimum": 0, - "type": "integer" - }, - "shape": { - "items": { - "$ref": "#/$defs/TensorDimension" - }, - "type": [ - "array", - "null" - ] - } - }, - "required": [ - "dtype", - "rank" - ], - "type": "object" - }, - "TensorDType": { - "description": "Tensor-boundary dtype vocabulary, including non-numeric pipeline values.", - "oneOf": [ - { - "enum": [ - "float32", - "fp32", - "float16", - "fp16", - "bfloat16", - "bf16", - "float8_e4m3fn", - "float8_e5m2", - "int64", - "int32", - "int8", - "uint8", - "bool", - "string" - ], - "title": "Known standard value" - }, - { - "not": { - "enum": [ - "float32", - "fp32", - "float16", - "fp16", - "bfloat16", - "bf16", - "float8_e4m3fn", - "float8_e5m2", - "int64", - "int32", - "int8", - "uint8", - "bool", - "string" - ] - }, - "title": "Forward-compatible extension value", - "type": "string" - } - ], - "type": "string" - }, - "TensorDimension": { - "anyOf": [ - { - "description": "A fixed, non-negative dimension.", - "format": "int64", - "minimum": 0, - "type": "integer" - }, - { - "description": "A runtime shape symbol.", - "minLength": 1, - "type": "string" - } - ], - "description": "One fixed or runtime-resolved tensor-shape dimension." - }, - "ThumbnailOrder": { - "description": "Optional-thumbnail ordering vocabulary.", - "oneOf": [ - { - "enum": [ - "none", - "prepend", - "append" - ], - "title": "Known standard value" - }, - { - "not": { - "enum": [ - "none", - "prepend", - "append" - ] - }, - "title": "Forward-compatible extension value", - "type": "string" - } - ], - "type": "string" - }, - "VerificationMethod": { - "description": "Speculative verification-method vocabulary.", - "oneOf": [ - { - "enum": [ - "single_forward" - ], - "title": "Known standard value" - }, - { - "not": { - "enum": [ - "single_forward" - ] - }, - "title": "Forward-compatible extension value", - "type": "string" - } - ], - "type": "string" - }, - "VerifyConfig": { - "description": "Draft-token verification configuration.", - "properties": { - "method": { - "anyOf": [ - { - "$ref": "#/$defs/VerificationMethod" - }, - { - "type": "null" - } - ], - "description": "Verification method, such as `single_forward`." - }, - "session": { - "description": "Named verifier session or pipeline component.", - "type": [ - "string", - "null" - ] - } - }, - "type": "object" - }, - "WorkflowBatchingContract": { - "oneOf": [ - { - "additionalProperties": false, - "properties": { - "kind": { - "const": "none", - "type": "string" - } - }, - "required": [ - "kind" - ], - "type": "object" - }, - { - "additionalProperties": false, - "properties": { - "axis": { - "format": "uint", - "minimum": 0, - "type": "integer" - }, - "kind": { - "const": "stack", - "type": "string" - } - }, - "required": [ - "kind", - "axis" - ], - "type": "object" - }, - { - "additionalProperties": false, - "properties": { - "kind": { - "const": "ragged", - "type": "string" - }, - "offsets": { - "type": "string" - } - }, - "required": [ - "kind", - "offsets" - ], - "type": "object" - } - ] - }, - "WorkflowBranchOutput": { - "additionalProperties": false, - "properties": { - "cases": { - "additionalProperties": { - "type": "string" - }, - "type": "object" - }, - "default": { - "type": [ - "string", - "null" - ] - } - }, - "required": [ - "cases" - ], - "type": "object" - }, - "WorkflowCarry": { - "additionalProperties": false, - "properties": { - "cell": { - "type": "string" - }, - "initial": { - "type": [ - "string", - "null" - ] - }, - "next": { - "type": "string" - } - }, - "required": [ - "cell", - "next" - ], - "type": "object" - }, - "WorkflowComponent": { - "additionalProperties": false, - "properties": { - "application_overridable": { - "default": false, - "description": "Allow an application to select another declared component with the same\nversioned contract ABI for this invocation.", - "type": "boolean" - }, - "contract": { - "anyOf": [ - { - "$ref": "#/$defs/ComponentContract" - }, - { - "type": "null" - } - ] - }, - "effects": { - "default": [], - "items": { - "type": "string" - }, - "type": "array" - }, - "implementation": { - "$ref": "#/$defs/ComponentImplementation" - }, - "ports": { - "$ref": "#/$defs/ComponentPorts", - "default": { - "inputs": {}, - "outputs": {} - } - }, - "resources": { - "anyOf": [ - { - "$ref": "#/$defs/WorkflowResourceContract" - }, - { - "type": "null" - } - ] - } - }, - "required": [ - "implementation" - ], - "type": "object" - }, - "WorkflowEmitMode": { - "enum": [ - "replace", - "append", - "event" - ], - "type": "string" - }, - "WorkflowInput": { - "additionalProperties": false, - "properties": { - "contract": { - "$ref": "#/$defs/TensorContract" - }, - "default": { - "anyOf": [ - { - "$ref": "#/$defs/ScalarValue" - }, - { - "type": "null" - } - ] - }, - "required": { - "default": false, - "type": "boolean" - }, - "role": { - "$ref": "#/$defs/SemanticInputRole" - }, - "source": { - "$ref": "#/$defs/WorkflowInputSource" - } - }, - "required": [ - "contract", - "role", - "source" - ], - "type": "object" - }, - "WorkflowInputSource": { - "oneOf": [ - { - "additionalProperties": false, - "properties": { - "field": { - "$ref": "#/$defs/RuntimeInputRole" - }, - "kind": { - "const": "request", - "type": "string" - } - }, - "required": [ - "kind", - "field" - ], - "type": "object" - }, - { - "additionalProperties": false, - "properties": { - "kind": { - "const": "application", - "type": "string" - }, - "name": { - "type": "string" - } - }, - "required": [ - "kind", - "name" - ], - "type": "object" - }, - { - "additionalProperties": false, - "properties": { - "kind": { - "const": "literal", - "type": "string" - } - }, - "required": [ - "kind" - ], - "type": "object" - }, - { - "additionalProperties": false, - "properties": { - "kind": { - "const": "artifact", - "type": "string" - }, - "path": { - "type": "string" - } - }, - "required": [ - "kind", - "path" - ], - "type": "object" - } - ] - }, - "WorkflowLoopIteration": { - "additionalProperties": false, - "properties": { - "contract": { - "$ref": "#/$defs/TensorContract", - "description": "`int64` scalar or rank-one broadcast contract." - }, - "value": { - "description": "SSA value containing the current zero-based iteration.", - "type": "string" - } - }, - "required": [ - "value", - "contract" - ], - "type": "object" - }, - "WorkflowManifest": { - "additionalProperties": false, - "properties": { - "adapter_abis": { - "additionalProperties": { - "type": "string" - }, - "default": {}, - "type": "object" - }, - "capabilities": { - "default": [], - "items": { - "type": "string" - }, - "type": "array", - "uniqueItems": true - }, - "custom_op_versions": { - "additionalProperties": { - "type": "string" - }, - "default": {}, - "type": "object" - }, - "ir_version": { - "type": "string" - }, - "onnx_opsets": { - "additionalProperties": { - "format": "uint32", - "minimum": 0, - "type": "integer" - }, - "default": {}, - "type": "object" - } - }, - "required": [ - "ir_version" - ], - "type": "object" - }, - "WorkflowMemoryClass": { - "enum": [ - "default", - "device", - "host", - "pinned" - ], - "type": "string" - }, - "WorkflowOutput": { - "additionalProperties": false, - "properties": { - "contract": { - "$ref": "#/$defs/TensorContract" - }, - "role": { - "$ref": "#/$defs/WorkflowOutputRole" - }, - "stage": { - "$ref": "#/$defs/OutputStage" - } - }, - "required": [ - "contract", - "role", - "stage" - ], - "type": "object" - }, - "WorkflowOutputRole": { - "enum": [ - "tokens", - "text", - "image", - "audio", - "tensor", - "event" - ], - "type": "string" - }, - "WorkflowResourceContract": { - "additionalProperties": false, - "properties": { - "allowed_devices": { - "default": [], - "items": { - "$ref": "#/$defs/DeviceKind" - }, - "type": "array" - }, - "batching": { - "$ref": "#/$defs/WorkflowBatchingContract" - }, - "concurrency": { - "format": "uint", - "minimum": 0, - "type": [ - "integer", - "null" - ] - }, - "memory_class": { - "$ref": "#/$defs/WorkflowMemoryClass" - }, - "preferred_device": { - "anyOf": [ - { - "$ref": "#/$defs/DeviceKind" - }, - { - "type": "null" - } - ] - } - }, - "required": [ - "memory_class", - "batching" - ], - "type": "object" - }, - "WorkflowSpec": { - "additionalProperties": false, - "description": "Sound, component-centric workflow IR. Tensor math lives in invoked components.", - "properties": { - "components": { - "additionalProperties": { - "$ref": "#/$defs/WorkflowComponent" - }, - "type": "object" - }, - "inputs": { - "additionalProperties": { - "$ref": "#/$defs/WorkflowInput" - }, - "default": {}, - "type": "object" - }, - "manifest": { - "$ref": "#/$defs/WorkflowManifest" - }, - "outputs": { - "additionalProperties": { - "$ref": "#/$defs/WorkflowOutput" - }, - "default": {}, - "type": "object" - }, - "serving": { - "anyOf": [ - { - "$ref": "#/$defs/ServingServiceContract" - }, - { - "type": "null" - } - ] - }, - "state": { - "additionalProperties": { - "$ref": "#/$defs/WorkflowStateCell" - }, - "default": {}, - "type": "object" - }, - "steps": { - "items": { - "$ref": "#/$defs/WorkflowStep" - }, - "type": "array" - } - }, - "required": [ - "manifest", - "components", - "steps" - ], - "type": "object" - }, - "WorkflowStateCell": { - "additionalProperties": false, - "properties": { - "class": { - "$ref": "#/$defs/WorkflowStateClass", - "default": "semantic" - }, - "contract": { - "$ref": "#/$defs/TensorContract" - }, - "initializer": { - "type": "string" - }, - "recurrence": { - "$ref": "#/$defs/ShapeRecurrence" - }, - "scope": { - "$ref": "#/$defs/WorkflowStateScope" - }, - "session": { - "anyOf": [ - { - "$ref": "#/$defs/SessionLeaseContract" - }, - { - "type": "null" - } - ] - } - }, - "required": [ - "contract", - "scope", - "initializer", - "recurrence" - ], - "type": "object" - }, - "WorkflowStateClass": { - "enum": [ - "semantic", - "advisory" - ], - "type": "string" - }, - "WorkflowStateScope": { - "enum": [ - "invocation", - "session" - ], - "type": "string" - }, - "WorkflowStep": { - "oneOf": [ - { - "additionalProperties": false, - "properties": { - "kind": { - "const": "sequence", - "type": "string" - }, - "steps": { - "items": { - "$ref": "#/$defs/WorkflowStep" - }, - "type": "array" - } - }, - "required": [ - "kind", - "steps" - ], - "type": "object" - }, - { - "additionalProperties": false, - "properties": { - "component": { - "type": "string" - }, - "inputs": { - "additionalProperties": { - "type": "string" - }, - "default": {}, - "type": "object" - }, - "kind": { - "const": "invoke", - "type": "string" - }, - "outputs": { - "additionalProperties": { - "type": "string" - }, - "default": {}, - "type": "object" - } - }, - "required": [ - "kind", - "component" - ], - "type": "object" - }, - { - "additionalProperties": false, - "properties": { - "carried": { - "default": [], - "items": { - "$ref": "#/$defs/WorkflowCarry" - }, - "type": "array" - }, - "condition": { - "type": "string" - }, - "iteration": { - "anyOf": [ - { - "$ref": "#/$defs/WorkflowLoopIteration" - }, - { - "type": "null" - } - ] - }, - "kind": { - "const": "loop", - "type": "string" - }, - "max_iterations": { - "type": "string" - }, - "setup": { - "default": [], - "items": { - "$ref": "#/$defs/WorkflowStep" - }, - "type": "array" - }, - "steps": { - "items": { - "$ref": "#/$defs/WorkflowStep" - }, - "type": "array" - } - }, - "required": [ - "kind", - "steps", - "condition", - "max_iterations" - ], - "type": "object" - }, - { - "additionalProperties": false, - "properties": { - "cases": { - "additionalProperties": { - "$ref": "#/$defs/WorkflowStep" - }, - "type": "object" - }, - "default": { - "anyOf": [ - { - "$ref": "#/$defs/WorkflowStep" - }, - { - "type": "null" - } - ] - }, - "kind": { - "const": "branch", - "type": "string" - }, - "outputs": { - "additionalProperties": { - "$ref": "#/$defs/WorkflowBranchOutput" - }, - "default": {}, - "type": "object" - }, - "predicate": { - "type": "string" - } - }, - "required": [ - "kind", - "predicate", - "cases" - ], - "type": "object" - }, - { - "additionalProperties": false, - "properties": { - "kind": { - "const": "emit", - "type": "string" - }, - "mode": { - "$ref": "#/$defs/WorkflowEmitMode" - }, - "output": { - "type": "string" - }, - "valid_length": { - "type": [ - "string", - "null" - ] - }, - "value": { - "type": "string" - } - }, - "required": [ - "kind", - "value", - "output", - "mode" - ], - "type": "object" - } - ] - } - }, - "$id": "https://github.com/onnx/onnx/issues/8184", - "$schema": "https://json-schema.org/draft/2020-12/schema", - "allOf": [ - { - "not": { - "required": [ - "speculative", - "speculator_config" - ] - } - }, - { - "not": { - "properties": { - "model": { - "required": [ - "io" - ] - } - }, - "required": [ - "pipeline", - "model" - ] - } - } - ], - "description": "Portable, runtime-agnostic inference metadata for ONNX generative models. All top-level sections are optional, and unknown fields are permitted for forward-compatible schema evolution.", - "properties": { - "generation": { - "anyOf": [ - { - "$ref": "#/$defs/GenerationDefaults" - }, - { - "type": "null" - } - ], - "description": "Author-declared text-generation / search defaults.\n\nPopulated from an onnxruntime-genai `genai_config.json` `search` block.\nEvery field is optional; readers treat an absent value as \"use the\nruntime default\"." - }, - "hardware_requirements": { - "anyOf": [ - { - "$ref": "#/$defs/HardwareRequirements" - }, - { - "type": "null" - } - ], - "description": "Minimum and beneficial hardware capabilities used for distribution matching." - }, - "kv_cache": { - "anyOf": [ - { - "$ref": "#/$defs/KvCacheSpec" - }, - { - "type": "null" - } - ], - "description": "KV-cache storage, quantization tolerance, and operational semantics." - }, - "model": { - "anyOf": [ - { - "$ref": "#/$defs/ModelCapabilities" - }, - { - "type": "null" - } - ], - "description": "Build-time model properties and runtime-configurable capabilities." - }, - "pipeline": { - "anyOf": [ - { - "$ref": "#/$defs/PipelineSpec" - }, - { - "type": "null" - } - ], - "description": "Declarative multi-model pipeline and its dataflow graph." - }, - "preprocessing": { - "anyOf": [ - { - "$ref": "#/$defs/PreprocessingSpec" - }, - { - "type": "null" - } - ], - "description": "Declared, architecture-neutral input preprocessing programs.\n\nCarries the typed multimodal preprocessing contract (currently the image\ntransform program and its named tensor outputs). Every operation and\noutput is generic, parameterized data — never a model family, vendor\nstring, or baked-in shape. Absent means the model declares no native\npreprocessing program and a runtime must obtain it elsewhere or fail." - }, - "quantization": { - "anyOf": [ - { - "$ref": "#/$defs/QuantizationIntent" - }, - { - "type": "null" - } - ], - "description": "Model weight quantization intent, independent of the packed representation." - }, - "required_capabilities": { - "default": [], - "description": "Capability identifiers that a runtime MUST support or refuse to load the model.", - "examples": [ - [ - "kv_cache", - "grouped_query_attention" - ] - ], - "items": { - "minLength": 1, - "type": "string" - }, - "type": "array" - }, - "schema_version": { - "description": "Schema version of this inference-metadata document, e.g. `\"v1\"`.\n\nAbsent means the initial `\"v1\"` contract (readers default to `v1`).\nBump this only for breaking schema changes; additive fields keep the\nsame major version and rely on the forward-compatible \"ignore unknown\nfields\" rule.", - "type": [ - "string", - "null" - ] - }, - "speculative": { - "anyOf": [ - { - "$ref": "#/$defs/SpeculatorConfig" - }, - { - "type": "null" - } - ], - "description": "Standalone speculative proposer declaration.\n\nThis is the preferred native source for speculator discovery;\nHuggingFace `config.json` is a compatibility fallback. The deprecated\n`speculator_config` alias is accepted on input." - }, - "speculator_config": { - "allOf": [ - { - "anyOf": [ - { - "$ref": "#/$defs/SpeculatorConfig" - }, - { - "type": "null" - } - ], - "description": "Standalone speculative proposer declaration.\n\nThis is the preferred native source for speculator discovery;\nHuggingFace `config.json` is a compatibility fallback. The deprecated\n`speculator_config` alias is accepted on input." - } - ], - "deprecated": true, - "description": "Deprecated alias for `speculative`." - }, - "strategy": { - "anyOf": [ - { - "$ref": "#/$defs/StrategySpec" - }, - { - "type": "null" - } - ], - "description": "Generic inference strategy, including speculative decoding." - }, - "structured_output": { - "anyOf": [ - { - "$ref": "#/$defs/StructuredOutputSpec" - }, - { - "type": "null" - } - ], - "description": "Structured-output formats and model training conventions." - }, - "tokens": { - "anyOf": [ - { - "$ref": "#/$defs/SpecialTokens" - }, - { - "type": "null" - } - ], - "description": "Special / control token ids declared by the model author.\n\nPopulated from the model-level token id fields of a `genai_config.json`." - } - }, - "title": "ONNX Inference Metadata", - "type": "object" -} From 7c6a20153910236cb69cc43183824b03c57b0514 Mon Sep 17 00:00:00 2001 From: justinchuby Date: Thu, 13 Aug 2026 02:20:47 +0000 Subject: [PATCH 030/151] Add KV service contracts to nested TTS Publish talker and code-predictor cache groups with explicit aliases and logical lengths, and validate the real tiny Qwen3-TTS package in cross-repository CI instead of a representative arithmetic fixture. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: justinchuby --- .../codec_workflow_metadata_test.py | 7 +- .../onnx_genai/workflow_metadata.py | 238 +++++++++++++++++- ...generate_onnx_genai_validation_packages.py | 46 +--- 3 files changed, 252 insertions(+), 39 deletions(-) diff --git a/src/mobius/integrations/onnx_genai/codec_workflow_metadata_test.py b/src/mobius/integrations/onnx_genai/codec_workflow_metadata_test.py index d47b4aa99..2b37ceb00 100644 --- a/src/mobius/integrations/onnx_genai/codec_workflow_metadata_test.py +++ b/src/mobius/integrations/onnx_genai/codec_workflow_metadata_test.py @@ -174,7 +174,12 @@ def test_real_qwen3_tts_workflow_carries_trained_transitions_and_kv_state(): "code_predictor_step_embedder", "talker_text_step", }.issubset(workflow["components"]) - assert workflow["state"]["talker_cache_0"]["recurrence"]["kind"] == "growing" + 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"]["kv_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"] diff --git a/src/mobius/integrations/onnx_genai/workflow_metadata.py b/src/mobius/integrations/onnx_genai/workflow_metadata.py index f930e9867..18850307d 100644 --- a/src/mobius/integrations/onnx_genai/workflow_metadata.py +++ b/src/mobius/integrations/onnx_genai/workflow_metadata.py @@ -545,6 +545,8 @@ def _build_real_tts_workflow_metadata(pkg: Any, config: Any) -> dict[str, Any]: "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( @@ -666,6 +668,34 @@ def _build_real_tts_workflow_metadata(pkg: Any, config: Any) -> dict[str, Any]: "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, + }, + "package.slot_ids": { + "contract": batch_int, + "role": {"kind": "opaque"}, + "source": {"kind": "literal"}, + "required": False, + "default": 0, + }, } for iteration in range(num_groups - 2): inputs[f"package.setup_predictor_iteration_{iteration}"] = { @@ -930,6 +960,20 @@ def frame_generation_nodes(prefix: str, hidden: str, logits: str) -> list[dict[s "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", { @@ -1027,6 +1071,16 @@ def frame_generation_nodes(prefix: str, hidden: str, logits: str) -> list[dict[s ), }, ] + 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( { @@ -1101,6 +1155,25 @@ def frame_generation_nodes(prefix: str, hidden: str, logits: str) -> list[dict[s "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"}, @@ -1240,33 +1313,120 @@ def frame_generation_nodes(prefix: str, hidden: str, logits: str) -> list[dict[s "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"}, + }, + "slot_ids": { + "contract": batch_int, + "class": "semantic", + "scope": "invocation", + "initializer": "package.slot_ids", + "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": _contract(past), + "class": "semantic", "scope": "invocation", "initializer": f"talker.setup.{present.name}", "recurrence": { - "kind": "growing", + "kind": "bounded", "axis": 2, - "increment": "package.one_control", "max": "package.talker_context_limit", }, + "service_group": "talker_cache", } for index, (past, present) in enumerate(predictor_caches): state[f"predictor_cache_{index}"] = { "contract": _contract(past), + "class": "semantic", "scope": "invocation", "initializer": f"frame.predictor.{present.name}", "recurrence": { - "kind": "growing", + "kind": "bounded", "axis": 2, - "increment": "package.one_control", "max": "package.predictor_context_limit", }, + "service_group": "predictor_cache", } 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": "slot_ids", + "current": "package.slot_ids", + "body_input": "state.slot_ids.body", + "body_output": "state.slot_ids.body", + "next": "state.slot_ids.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", @@ -1363,6 +1523,11 @@ def frame_generation_nodes(prefix: str, hidden: str, logits: str) -> list[dict[s "nested_control_flow", "loop_induction_values", "typed_emit", + *( + ["serving_service_contract", "bounded_state_recurrence"] + if talker_caches or predictor_caches + else [] + ), ], }, "inputs": inputs, @@ -1377,6 +1542,71 @@ def frame_generation_nodes(prefix: str, hidden: str, logits: str) -> list[dict[s name: _component(model, _artifact(name, len(pkg))) for name, model in pkg.items() }, "state": state, + **( + { + "serving": { + "active": "active", + "done": "done", + "accepted_len": "accepted_len", + "slot_ids": "slot_ids", + "kv_service": { + "paging": "paged", + "allocation": "runtime", + "compaction": True, + "groups": { + **( + { + "talker_cache": { + "sequence_axis": 2, + "layout": "bnsh", + "logical_lengths": "talker_cache_lengths", + "storage": "paged", + "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": { + "sequence_axis": 2, + "layout": "bnsh", + "logical_lengths": "predictor_cache_lengths", + "storage": "paged", + "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", diff --git a/tests/generate_onnx_genai_validation_packages.py b/tests/generate_onnx_genai_validation_packages.py index 5fa6cfcab..497720a8f 100644 --- a/tests/generate_onnx_genai_validation_packages.py +++ b/tests/generate_onnx_genai_validation_packages.py @@ -12,7 +12,6 @@ _decoder_package, _diffusion_package, _model, - _TTSCfg, _value, _vlm_package, ) @@ -20,43 +19,22 @@ write_speculative_workflow_metadata, ) from mobius.integrations.onnx_genai.workflow_metadata_test import _speculative_package +from mobius.models.qwen3_tts import Qwen3TTSForConditionalGeneration +from mobius.models.qwen3_tts_test import _TINY_CONFIG +from mobius.tasks import TTSTask def _tts_package() -> ModelPackage: - batch = "batch" - return ModelPackage( - { - "talker": _model( - "talker", - [_value("inputs_embeds", ir.DataType.FLOAT, [batch, "sequence", 16])], - [("last_hidden_state", ir.DataType.FLOAT, [batch, 16])], - ), - "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(), + 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"])], + ) + return package def main() -> None: From 5169a15ab0f127666953c05fe4ebf71db51471cf Mon Sep 17 00:00:00 2001 From: justinchuby Date: Thu, 13 Aug 2026 02:55:46 +0000 Subject: [PATCH 031/151] Check in ONNX GenAI workflow fixtures Add reproducible executable decoder, VLM, diffusion, masked diffusion, Qwen3-TTS, speculative, and codec packages targeting onnx-genai c9bddd6e. Validate checked artifacts unconditionally in CI and fix decoder token append and masked-loop control contracts exposed by runtime conformance. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: justinchuby --- .github/workflows/main.yml | 11 +- .gitignore | 2 + .../onnx_genai/auto_export_test.py | 6 +- .../onnx_genai/workflow_metadata.py | 36 +- tests/fixtures/onnx_genai_workflows/README.md | 8 + .../codec/decoder/model.onnx | Bin 0 -> 251 bytes .../codec/decoder/model.onnx.data | 0 .../codec/encoder/model.onnx | Bin 0 -> 251 bytes .../codec/encoder/model.onnx.data | 0 .../codec/inference_metadata.yaml | 64 + .../decoder/inference_metadata.yaml | 433 +++++ .../onnx_genai_workflows/decoder/model.onnx | Bin 0 -> 2385 bytes .../decoder/model.onnx.data | 0 .../decoder/policies/cache_length_update.onnx | Bin 0 -> 386 bytes .../policies/decoder_state_initializer.onnx | Bin 0 -> 6244 bytes .../decoder/policies/decoder_step_update.onnx | Bin 0 -> 2037 bytes .../decoder/policies/last_token_logits.onnx | Bin 0 -> 631 bytes .../decoder/policies/termination.onnx | Bin 0 -> 3730 bytes .../decoder/policies/token_sampler.onnx | Bin 0 -> 521 bytes .../decoder/policies/token_state_update.onnx | Bin 0 -> 1240 bytes .../diffusion/denoiser/model.onnx | Bin 0 -> 268 bytes .../diffusion/denoiser/model.onnx.data | 0 .../diffusion/inference_metadata.yaml | 251 +++ .../policies/continue_predicate.onnx | Bin 0 -> 910 bytes .../policies/diffusion_schedule.onnx | Bin 0 -> 497 bytes .../policies/diffusion_timesteps.onnx | Bin 0 -> 493 bytes .../diffusion/policies/euler_model_input.onnx | Bin 0 -> 1926 bytes .../diffusion/policies/schedule_lookup.onnx | Bin 0 -> 607 bytes .../diffusion/policies/solver_step.onnx | Bin 0 -> 3060 bytes .../diffusion/text_encoder/model.onnx | Bin 0 -> 213 bytes .../diffusion/text_encoder/model.onnx.data | 0 .../diffusion/vae_decoder/model.onnx | Bin 0 -> 174 bytes .../diffusion/vae_decoder/model.onnx.data | 0 .../masked/inference_metadata.yaml | 237 +++ .../onnx_genai_workflows/masked/model.onnx | Bin 0 -> 1008 bytes .../masked/model.onnx.data | 0 .../masked/policies/masked_update.onnx | Bin 0 -> 13182 bytes .../speculative/inference_metadata.yaml | 333 ++++ .../policies/cache_length_update.onnx | Bin 0 -> 386 bytes .../policies/speculative_acceptance.onnx | Bin 0 -> 8736 bytes .../speculative/proposer/model.onnx | Bin 0 -> 135 bytes .../speculative/proposer/model.onnx.data | 0 .../speculative/verifier/model.onnx | Bin 0 -> 233 bytes .../speculative/verifier/model.onnx.data | 0 .../tts/code_predictor/model.onnx | Bin 0 -> 103645 bytes .../tts/code_predictor/model.onnx.data | Bin 0 -> 65536 bytes .../tts/code_predictor_indices/model.onnx | Bin 0 -> 942 bytes .../code_predictor_indices/model.onnx.data | 0 .../tts/code_predictor_prefill/model.onnx | Bin 0 -> 366 bytes .../code_predictor_prefill/model.onnx.data | 0 .../code_predictor_step_embedder/model.onnx | Bin 0 -> 817 bytes .../model.onnx.data | 0 .../onnx_genai_workflows/tts/codec/model.onnx | Bin 0 -> 146 bytes .../tts/codec/model.onnx.data | 0 .../tts/embedding/model.onnx | Bin 0 -> 4639 bytes .../tts/embedding/model.onnx.data | 0 .../tts/inference_metadata.yaml | 1455 +++++++++++++++++ .../tts/policies/cache_length_update.onnx | Bin 0 -> 386 bytes .../tts/policies/code_frame_update.onnx | Bin 0 -> 1580 bytes .../tts/policies/code_history_append.onnx | Bin 0 -> 853 bytes .../tts/policies/codec_layout.onnx | Bin 0 -> 428 bytes .../tts/policies/continue_predicate.onnx | Bin 0 -> 910 bytes .../tts/policies/last_token_logits.onnx | Bin 0 -> 631 bytes .../tts/policies/predictor_body_sampler.onnx | Bin 0 -> 536 bytes .../policies/predictor_prefill_sampler.onnx | Bin 0 -> 539 bytes .../policies/predictor_state_initializer.onnx | Bin 0 -> 14690 bytes .../tts/policies/predictor_step_update.onnx | Bin 0 -> 2039 bytes .../tts/policies/setup_predictor_sampler.onnx | Bin 0 -> 537 bytes .../tts/policies/setup_talker_sampler.onnx | Bin 0 -> 534 bytes .../tts/policies/talker_sampler.onnx | Bin 0 -> 528 bytes .../policies/talker_state_initializer.onnx | Bin 0 -> 7457 bytes .../tts/policies/talker_step_update.onnx | Bin 0 -> 2044 bytes .../tts/policies/token_to_slot.onnx | Bin 0 -> 438 bytes .../tts/policies/tts_state_initializer.onnx | Bin 0 -> 2194 bytes .../tts/talker/model.onnx | Bin 0 -> 30846 bytes .../tts/talker/model.onnx.data | Bin 0 -> 2048 bytes .../tts/talker_prefill_embedder/model.onnx | Bin 0 -> 21644 bytes .../talker_prefill_embedder/model.onnx.data | 0 .../tts/talker_step_embedder/model.onnx | Bin 0 -> 7296 bytes .../tts/talker_step_embedder/model.onnx.data | 0 .../tts/talker_text_step/model.onnx | Bin 0 -> 1759 bytes .../tts/talker_text_step/model.onnx.data | 0 .../vlm/decoder/model.onnx | Bin 0 -> 656 bytes .../vlm/decoder/model.onnx.data | 0 .../vlm/embedding/model.onnx | Bin 0 -> 216 bytes .../vlm/embedding/model.onnx.data | 0 .../vlm/inference_metadata.yaml | 608 +++++++ .../vlm/policies/cache_length_update.onnx | Bin 0 -> 386 bytes .../policies/decoder_state_initializer.onnx | Bin 0 -> 7283 bytes .../vlm/policies/decoder_step_update.onnx | Bin 0 -> 2037 bytes .../vlm/policies/last_token_logits.onnx | Bin 0 -> 631 bytes .../vlm/policies/termination.onnx | Bin 0 -> 3730 bytes .../vlm/policies/token_sampler.onnx | Bin 0 -> 521 bytes .../vlm/policies/token_state_update.onnx | Bin 0 -> 1240 bytes .../vlm/vision_encoder/model.onnx | Bin 0 -> 208 bytes .../vlm/vision_encoder/model.onnx.data | 0 ...generate_onnx_genai_validation_packages.py | 148 +- 97 files changed, 3575 insertions(+), 17 deletions(-) create mode 100644 tests/fixtures/onnx_genai_workflows/README.md create mode 100644 tests/fixtures/onnx_genai_workflows/codec/decoder/model.onnx create mode 100644 tests/fixtures/onnx_genai_workflows/codec/decoder/model.onnx.data create mode 100644 tests/fixtures/onnx_genai_workflows/codec/encoder/model.onnx create mode 100644 tests/fixtures/onnx_genai_workflows/codec/encoder/model.onnx.data create mode 100644 tests/fixtures/onnx_genai_workflows/codec/inference_metadata.yaml create mode 100644 tests/fixtures/onnx_genai_workflows/decoder/inference_metadata.yaml create mode 100644 tests/fixtures/onnx_genai_workflows/decoder/model.onnx create mode 100644 tests/fixtures/onnx_genai_workflows/decoder/model.onnx.data create mode 100644 tests/fixtures/onnx_genai_workflows/decoder/policies/cache_length_update.onnx create mode 100644 tests/fixtures/onnx_genai_workflows/decoder/policies/decoder_state_initializer.onnx create mode 100644 tests/fixtures/onnx_genai_workflows/decoder/policies/decoder_step_update.onnx create mode 100644 tests/fixtures/onnx_genai_workflows/decoder/policies/last_token_logits.onnx create mode 100644 tests/fixtures/onnx_genai_workflows/decoder/policies/termination.onnx create mode 100644 tests/fixtures/onnx_genai_workflows/decoder/policies/token_sampler.onnx create mode 100644 tests/fixtures/onnx_genai_workflows/decoder/policies/token_state_update.onnx create mode 100644 tests/fixtures/onnx_genai_workflows/diffusion/denoiser/model.onnx create mode 100644 tests/fixtures/onnx_genai_workflows/diffusion/denoiser/model.onnx.data create mode 100644 tests/fixtures/onnx_genai_workflows/diffusion/inference_metadata.yaml create mode 100644 tests/fixtures/onnx_genai_workflows/diffusion/policies/continue_predicate.onnx create mode 100644 tests/fixtures/onnx_genai_workflows/diffusion/policies/diffusion_schedule.onnx create mode 100644 tests/fixtures/onnx_genai_workflows/diffusion/policies/diffusion_timesteps.onnx create mode 100644 tests/fixtures/onnx_genai_workflows/diffusion/policies/euler_model_input.onnx create mode 100644 tests/fixtures/onnx_genai_workflows/diffusion/policies/schedule_lookup.onnx create mode 100644 tests/fixtures/onnx_genai_workflows/diffusion/policies/solver_step.onnx create mode 100644 tests/fixtures/onnx_genai_workflows/diffusion/text_encoder/model.onnx create mode 100644 tests/fixtures/onnx_genai_workflows/diffusion/text_encoder/model.onnx.data create mode 100644 tests/fixtures/onnx_genai_workflows/diffusion/vae_decoder/model.onnx create mode 100644 tests/fixtures/onnx_genai_workflows/diffusion/vae_decoder/model.onnx.data create mode 100644 tests/fixtures/onnx_genai_workflows/masked/inference_metadata.yaml create mode 100644 tests/fixtures/onnx_genai_workflows/masked/model.onnx create mode 100644 tests/fixtures/onnx_genai_workflows/masked/model.onnx.data create mode 100644 tests/fixtures/onnx_genai_workflows/masked/policies/masked_update.onnx create mode 100644 tests/fixtures/onnx_genai_workflows/speculative/inference_metadata.yaml create mode 100644 tests/fixtures/onnx_genai_workflows/speculative/policies/cache_length_update.onnx create mode 100644 tests/fixtures/onnx_genai_workflows/speculative/policies/speculative_acceptance.onnx create mode 100644 tests/fixtures/onnx_genai_workflows/speculative/proposer/model.onnx create mode 100644 tests/fixtures/onnx_genai_workflows/speculative/proposer/model.onnx.data create mode 100644 tests/fixtures/onnx_genai_workflows/speculative/verifier/model.onnx create mode 100644 tests/fixtures/onnx_genai_workflows/speculative/verifier/model.onnx.data create mode 100644 tests/fixtures/onnx_genai_workflows/tts/code_predictor/model.onnx create mode 100644 tests/fixtures/onnx_genai_workflows/tts/code_predictor/model.onnx.data create mode 100644 tests/fixtures/onnx_genai_workflows/tts/code_predictor_indices/model.onnx create mode 100644 tests/fixtures/onnx_genai_workflows/tts/code_predictor_indices/model.onnx.data create mode 100644 tests/fixtures/onnx_genai_workflows/tts/code_predictor_prefill/model.onnx create mode 100644 tests/fixtures/onnx_genai_workflows/tts/code_predictor_prefill/model.onnx.data create mode 100644 tests/fixtures/onnx_genai_workflows/tts/code_predictor_step_embedder/model.onnx create mode 100644 tests/fixtures/onnx_genai_workflows/tts/code_predictor_step_embedder/model.onnx.data create mode 100644 tests/fixtures/onnx_genai_workflows/tts/codec/model.onnx create mode 100644 tests/fixtures/onnx_genai_workflows/tts/codec/model.onnx.data create mode 100644 tests/fixtures/onnx_genai_workflows/tts/embedding/model.onnx create mode 100644 tests/fixtures/onnx_genai_workflows/tts/embedding/model.onnx.data create mode 100644 tests/fixtures/onnx_genai_workflows/tts/inference_metadata.yaml create mode 100644 tests/fixtures/onnx_genai_workflows/tts/policies/cache_length_update.onnx create mode 100644 tests/fixtures/onnx_genai_workflows/tts/policies/code_frame_update.onnx create mode 100644 tests/fixtures/onnx_genai_workflows/tts/policies/code_history_append.onnx create mode 100644 tests/fixtures/onnx_genai_workflows/tts/policies/codec_layout.onnx create mode 100644 tests/fixtures/onnx_genai_workflows/tts/policies/continue_predicate.onnx create mode 100644 tests/fixtures/onnx_genai_workflows/tts/policies/last_token_logits.onnx create mode 100644 tests/fixtures/onnx_genai_workflows/tts/policies/predictor_body_sampler.onnx create mode 100644 tests/fixtures/onnx_genai_workflows/tts/policies/predictor_prefill_sampler.onnx create mode 100644 tests/fixtures/onnx_genai_workflows/tts/policies/predictor_state_initializer.onnx create mode 100644 tests/fixtures/onnx_genai_workflows/tts/policies/predictor_step_update.onnx create mode 100644 tests/fixtures/onnx_genai_workflows/tts/policies/setup_predictor_sampler.onnx create mode 100644 tests/fixtures/onnx_genai_workflows/tts/policies/setup_talker_sampler.onnx create mode 100644 tests/fixtures/onnx_genai_workflows/tts/policies/talker_sampler.onnx create mode 100644 tests/fixtures/onnx_genai_workflows/tts/policies/talker_state_initializer.onnx create mode 100644 tests/fixtures/onnx_genai_workflows/tts/policies/talker_step_update.onnx create mode 100644 tests/fixtures/onnx_genai_workflows/tts/policies/token_to_slot.onnx create mode 100644 tests/fixtures/onnx_genai_workflows/tts/policies/tts_state_initializer.onnx create mode 100644 tests/fixtures/onnx_genai_workflows/tts/talker/model.onnx create mode 100644 tests/fixtures/onnx_genai_workflows/tts/talker/model.onnx.data create mode 100644 tests/fixtures/onnx_genai_workflows/tts/talker_prefill_embedder/model.onnx create mode 100644 tests/fixtures/onnx_genai_workflows/tts/talker_prefill_embedder/model.onnx.data create mode 100644 tests/fixtures/onnx_genai_workflows/tts/talker_step_embedder/model.onnx create mode 100644 tests/fixtures/onnx_genai_workflows/tts/talker_step_embedder/model.onnx.data create mode 100644 tests/fixtures/onnx_genai_workflows/tts/talker_text_step/model.onnx create mode 100644 tests/fixtures/onnx_genai_workflows/tts/talker_text_step/model.onnx.data create mode 100644 tests/fixtures/onnx_genai_workflows/vlm/decoder/model.onnx create mode 100644 tests/fixtures/onnx_genai_workflows/vlm/decoder/model.onnx.data create mode 100644 tests/fixtures/onnx_genai_workflows/vlm/embedding/model.onnx create mode 100644 tests/fixtures/onnx_genai_workflows/vlm/embedding/model.onnx.data create mode 100644 tests/fixtures/onnx_genai_workflows/vlm/inference_metadata.yaml create mode 100644 tests/fixtures/onnx_genai_workflows/vlm/policies/cache_length_update.onnx create mode 100644 tests/fixtures/onnx_genai_workflows/vlm/policies/decoder_state_initializer.onnx create mode 100644 tests/fixtures/onnx_genai_workflows/vlm/policies/decoder_step_update.onnx create mode 100644 tests/fixtures/onnx_genai_workflows/vlm/policies/last_token_logits.onnx create mode 100644 tests/fixtures/onnx_genai_workflows/vlm/policies/termination.onnx create mode 100644 tests/fixtures/onnx_genai_workflows/vlm/policies/token_sampler.onnx create mode 100644 tests/fixtures/onnx_genai_workflows/vlm/policies/token_state_update.onnx create mode 100644 tests/fixtures/onnx_genai_workflows/vlm/vision_encoder/model.onnx create mode 100644 tests/fixtures/onnx_genai_workflows/vlm/vision_encoder/model.onnx.data diff --git a/.github/workflows/main.yml b/.github/workflows/main.yml index 63f399053..124505e59 100644 --- a/.github/workflows/main.yml +++ b/.github/workflows/main.yml @@ -23,7 +23,7 @@ jobs: - uses: actions/checkout@v7 with: repository: justinchuby/onnx-genai - ref: 8215649100a0a27be15b04045fddde777c8248fc + ref: c9bddd6e path: validation/onnx-genai - uses: actions/setup-python@v7 with: @@ -35,10 +35,15 @@ jobs: pip install onnxruntime pip install -e '.[testing]' - name: Generate representative packages - run: PYTHONPATH=src python tests/generate_onnx_genai_validation_packages.py validation/packages + run: | + PYTHONPATH=src python tests/generate_onnx_genai_validation_packages.py \ + validation/generated + diff --recursive --brief \ + tests/fixtures/onnx_genai_workflows validation/generated - name: Validate package semantics run: | - for package in validation/packages/*; do + for package in tests/fixtures/onnx_genai_workflows/*; 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" diff --git a/.gitignore b/.gitignore index efcf9d8cd..38ca53b2b 100644 --- a/.gitignore +++ b/.gitignore @@ -212,6 +212,8 @@ __marimo__/ *.onnx *.onnx.data *.gguf +!tests/fixtures/onnx_genai_workflows/**/*.onnx +!tests/fixtures/onnx_genai_workflows/**/*.onnx.data # Common test dirs output/** diff --git a/src/mobius/integrations/onnx_genai/auto_export_test.py b/src/mobius/integrations/onnx_genai/auto_export_test.py index 9560af08a..1b3c369fa 100644 --- a/src/mobius/integrations/onnx_genai/auto_export_test.py +++ b/src/mobius/integrations/onnx_genai/auto_export_test.py @@ -178,7 +178,11 @@ def test_dispatch_decoder(tmp_path): ] 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"] == "sample.body" + 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]}, diff --git a/src/mobius/integrations/onnx_genai/workflow_metadata.py b/src/mobius/integrations/onnx_genai/workflow_metadata.py index 18850307d..6f312b786 100644 --- a/src/mobius/integrations/onnx_genai/workflow_metadata.py +++ b/src/mobius/integrations/onnx_genai/workflow_metadata.py @@ -3009,20 +3009,20 @@ def build_vlm_workflow_metadata( if cache_pairs else [] ), + _invoke( + "token_state_update", + {"current": "state.token.body", "update": "sample.body"}, + {"next": "token.body"}, + {"state": _effect("state.0", "state.1")}, + ), { "kind": "emit", - "value": "sample.body", + "value": "token.body", "output": "tokens", "mode": "append", "effect_name": "emit", "effect": _effect("emit.0", "emit.1"), }, - _invoke( - "token_state_update", - {"current": "state.token.body", "update": "sample.body"}, - {"next": "token.body"}, - {"state": _effect("state.0", "state.1")}, - ), _invoke( "embedding", embedding_body_inputs, @@ -3095,7 +3095,15 @@ def build_vlm_workflow_metadata( }, "inputs": inputs, "outputs": { - "tokens": {"contract": batch_int, "role": "tokens", "stage": "pre_adapter"} + "tokens": { + "contract": { + "dtype": "int64", + "rank": 2, + "shape": [batch, "generated_sequence"], + }, + "role": "tokens", + "stage": "pre_adapter", + } }, "components": components, "state": state, @@ -4687,7 +4695,7 @@ def build_decoder_workflow_metadata( ), { "kind": "emit", - "value": "sample.body", + "value": "token.body", "output": "tokens", "mode": "append", "effect_name": "emit", @@ -4740,7 +4748,11 @@ def build_decoder_workflow_metadata( "inputs": workflow_inputs, "outputs": { "tokens": { - "contract": batch_int, + "contract": { + "dtype": "int64", + "rank": 2, + "shape": [batch_dimension, "generated_sequence"], + }, "role": "tokens", "stage": "pre_adapter", } @@ -4845,6 +4857,7 @@ def build_language_diffusion_pipeline_metadata( } batch_dimension = token_contract["shape"][0] batch_int = {"dtype": "int64", "rank": 1, "shape": [batch_dimension]} + control_int = {"dtype": "int64", "rank": 1, "shape": [1]} inputs = { "request.input_ids": { "contract": token_contract, @@ -4881,7 +4894,7 @@ def build_language_diffusion_pipeline_metadata( "default": 0, }, "request.max_iterations": { - "contract": batch_int, + "contract": control_int, "role": { "kind": "runtime", "version": "1.0", @@ -5021,6 +5034,7 @@ def update_invoke( "linear_effects", "nested_control_flow", "typed_emit", + "loop_induction_values", ], }, "inputs": inputs, diff --git a/tests/fixtures/onnx_genai_workflows/README.md b/tests/fixtures/onnx_genai_workflows/README.md new file mode 100644 index 000000000..f8eab0434 --- /dev/null +++ b/tests/fixtures/onnx_genai_workflows/README.md @@ -0,0 +1,8 @@ +# ONNX GenAI workflow conformance fixtures + +Generated by `tests/generate_onnx_genai_validation_packages.py` for semantic +validation and runtime conformance against `justinchuby/onnx-genai@c9bddd6e`. + +The decoder, VLM, diffusion, masked diffusion, real tiny Qwen3-TTS, +speculative, and codec packages contain graph-only synthetic models and policy +components. They contain no downloaded model weights. diff --git a/tests/fixtures/onnx_genai_workflows/codec/decoder/model.onnx b/tests/fixtures/onnx_genai_workflows/codec/decoder/model.onnx new file mode 100644 index 0000000000000000000000000000000000000000..f6224514cd03055262813f841be4615a590875d7 GIT binary patch literal 251 zcma)$yAFad6o!FFsK*dYC&NGtVU;jDy2&GGOjAk~BD7757iM4BM-T=aoLv9whZftn z4KEv9CQEKGdeA#ROWlUSzT~2-yPk-_Az6$jA@^Y}Xv>YFjAJWLAX3*$Dus9&rfb!u ztf9t`s^*$%R`u99p7J$>KdEtp7$(*5J0!4I@)1HGHWN((F|u(49(D`bu?ieyLx6pH XD{Dy%ZB@ffk-?W5|JB&dz*??8ka9_W literal 0 HcmV?d00001 diff --git a/tests/fixtures/onnx_genai_workflows/codec/decoder/model.onnx.data b/tests/fixtures/onnx_genai_workflows/codec/decoder/model.onnx.data new file mode 100644 index 000000000..e69de29bb diff --git a/tests/fixtures/onnx_genai_workflows/codec/encoder/model.onnx b/tests/fixtures/onnx_genai_workflows/codec/encoder/model.onnx new file mode 100644 index 0000000000000000000000000000000000000000..de3b4f103a00d673d42a1a09237bb147d7c1a9a2 GIT binary patch literal 251 zcma)$y$ZrG6ori_#w(&^k`6^mSH;=U%|3!6Ax*Eg*d`&j{>{FykDx^aadJF|k8?gg zSUwos8SslbOgBec>=c2Nr(Ud literal 0 HcmV?d00001 diff --git a/tests/fixtures/onnx_genai_workflows/codec/encoder/model.onnx.data b/tests/fixtures/onnx_genai_workflows/codec/encoder/model.onnx.data new file mode 100644 index 000000000..e69de29bb 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..032e47071 --- /dev/null +++ b/tests/fixtures/onnx_genai_workflows/codec/inference_metadata.yaml @@ -0,0 +1,64 @@ +schema_version: v1 +pipeline: + workflow: + manifest: + ir_version: '1.0' + onnx_opsets: + ai.onnx: 24 + capabilities: + - workflow_ssa + - linear_effects + - typed_emit + inputs: + request.waveform: + contract: + dtype: float32 + rank: 3 + shape: + - batch + - 1 + - audio_samples + 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 + 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..7f620e841 --- /dev/null +++ b/tests/fixtures/onnx_genai_workflows/decoder/inference_metadata.yaml @@ -0,0 +1,433 @@ +schema_version: '1.0' +pipeline: + workflow: + manifest: + ir_version: '1.0' + onnx_opsets: + ai.onnx: 24 + capabilities: + - workflow_ssa + - linear_effects + - nested_control_flow + - typed_emit + - loop_induction_values + - serving_service_contract + - bounded_state_recurrence + inputs: + request.input_ids: + contract: + dtype: int64 + rank: 2 + shape: + - batch + - sequence + role: + kind: runtime + version: '1.0' + role: prompt_tokens + source: + kind: request + required: true + request.max_iterations: + contract: &id001 + 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: *id001 + role: + kind: opaque + source: + kind: literal + required: false + default: 1 + package.max_context: + contract: *id001 + role: + kind: opaque + source: + kind: literal + required: false + default: 8192 + package.active: + contract: &id002 + dtype: bool + rank: 1 + shape: + - batch + role: + kind: opaque + source: + kind: literal + required: false + default: true + package.not_done: + contract: *id002 + role: + kind: opaque + source: + kind: literal + required: false + default: false + package.slot_ids: + contract: &id003 + dtype: int64 + rank: 1 + shape: + - batch + role: + kind: opaque + source: + kind: literal + required: false + default: 0 + package.cache_lengths: + contract: *id003 + role: + kind: opaque + source: + kind: literal + required: false + default: 0 + package.zero_batch: + contract: *id003 + role: + kind: opaque + source: + kind: literal + required: false + default: 0 + outputs: + tokens: + contract: + dtype: int64 + rank: 2 + shape: + - batch + - generated_sequence + role: tokens + stage: pre_adapter + components: + model: + implementation: + kind: onnx + artifact: model.onnx + token_sampler: + implementation: + kind: onnx + artifact: policies/token_sampler.onnx + contract: + id: onnx-genai.token-sampler + version: '1' + bindings: + logits: logits + token: token + parameters: + mode: greedy + application_overridable: true + termination: + implementation: + kind: onnx + artifact: policies/termination.onnx + contract: + id: onnx-genai.termination-predicate + version: '1' + bindings: + tokens: token_ids + eos_ids: eos_ids + iteration: iteration + max_iterations: max_iterations + done: done + continue: continue + token_state_update: + implementation: + kind: onnx + artifact: policies/token_state_update.onnx + contract: + id: onnx-genai.state-update + version: '1' + bindings: + current: current + update: update + next: next + last_token_logits: + implementation: + kind: onnx + artifact: policies/last_token_logits.onnx + decoder_state_initializer: + implementation: + kind: onnx + artifact: policies/decoder_state_initializer.onnx + decoder_step_update: + implementation: + kind: onnx + artifact: policies/decoder_step_update.onnx + cache_length_update: + implementation: + kind: onnx + artifact: policies/cache_length_update.onnx + state: + token: + contract: + dtype: int64 + rank: 2 + shape: + - batch + - 1 + scope: invocation + initializer: initializer.token_slot + recurrence: + kind: invariant + logits: + contract: + dtype: float32 + rank: 2 + shape: + - batch + - 128 + scope: invocation + initializer: decoder.setup.last_logits + recurrence: + kind: invariant + active: + contract: *id002 + class: semantic + scope: invocation + initializer: package.active + recurrence: + kind: invariant + done: + contract: *id002 + class: semantic + scope: invocation + initializer: package.not_done + recurrence: + kind: invariant + accepted_len: + contract: *id003 + class: semantic + scope: invocation + initializer: package.zero_batch + recurrence: + kind: invariant + slot_ids: + contract: *id003 + class: semantic + scope: invocation + initializer: package.slot_ids + recurrence: + kind: invariant + cache_lengths: + contract: *id003 + class: semantic + scope: invocation + initializer: package.cache_lengths + recurrence: + kind: invariant + attention_mask: + contract: + dtype: int64 + rank: 2 + shape: + - batch + - context + scope: invocation + initializer: initializer.body_attention_mask + recurrence: + kind: growing + axis: 1 + increment: package.one_token + max: package.max_context + position_ids: + contract: + dtype: int64 + rank: 2 + shape: + - batch + - 1 + scope: invocation + initializer: initializer.body_position_ids + recurrence: + kind: invariant + cache_9: + contract: + dtype: float32 + rank: 4 + shape: + - batch + - 2 + - past_sequence + - 8 + scope: invocation + initializer: decoder.setup.present.0.key + recurrence: + kind: bounded + axis: 2 + max: package.max_context + service_group: decoder_cache + serving: + active: active + done: done + accepted_len: accepted_len + slot_ids: slot_ids + kv_service: + paging: paged + allocation: runtime + compaction: true + groups: + decoder_cache: + sequence_axis: 2 + layout: bnsh + logical_lengths: cache_lengths + storage: paged + ports: + model: + cache_9: + input: past_key_values.0.key + output: present.0.key + steps: + - kind: loop + setup: + - kind: invoke + component: decoder_state_initializer + inputs: + prompt_tokens: request.input_ids + outputs: + attention_mask: initializer.attention_mask + body_attention_mask: initializer.body_attention_mask + token_slot: initializer.token_slot + 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: last_token_logits + inputs: + logits: decoder.setup.logits + outputs: + last_logits: decoder.setup.last_logits + steps: + - kind: invoke + component: token_sampler + inputs: + logits: logits + outputs: + token: sample.body + - kind: invoke + component: token_state_update + inputs: + current: token + update: sample.body + outputs: + next: token.body + - kind: invoke + component: termination + inputs: + token_ids: sample.body + eos_ids: package.eos_ids + iteration: loop.iteration + max_iterations: request.max_iterations + outputs: + done: loop.done + continue: loop.continue + - kind: invoke + component: cache_length_update + inputs: + left: cache_lengths + right: package.one_token + outputs: + total: cache_lengths.next + - kind: invoke + component: cache_length_update + inputs: + left: package.zero_batch + right: package.one_token + outputs: + total: accepted_len.next + - kind: emit + value: token.body + output: tokens + mode: append + - kind: invoke + component: model + inputs: + input_ids: token.body + past_key_values.0.key: cache_9 + 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: active + next: loop.continue + - cell: done + next: loop.done + - cell: cache_lengths + next: cache_lengths.next + - cell: accepted_len + next: accepted_len.next + - cell: slot_ids + next: slot_ids + - cell: attention_mask + next: decoder_step.body_attention_mask + - cell: position_ids + next: decoder_step.body_position_ids + - cell: cache_9 + next: decoder.body.present.0.key + iteration: + value: loop.iteration + contract: + dtype: int64 + rank: 1 + shape: + - 1 diff --git a/tests/fixtures/onnx_genai_workflows/decoder/model.onnx b/tests/fixtures/onnx_genai_workflows/decoder/model.onnx new file mode 100644 index 0000000000000000000000000000000000000000..13e5e0d904296e7a1ca2710d4d67f2b0f98bd520 GIT binary patch literal 2385 zcmc(gO>fgM7{}|ZrQWlyD@Ge7lg6kR8ktnJTibPr124B;IB*+6mY13}U6+K~o3`7; zSAZjjegN)>FTm%*Nnc`@b+a0MPU3sFM86hhz#2-k>5h)jwkb8sR z)RV&An9B9;leUcqm6@)S=gtP?Vl^GWsDwJ%jQWjGHZ9!mw|85^=(fhB+Z=!;?sv;- zi!G-BIn^9PnW~x#@R|l*a}Lik@FLi>)cb1Mq&07omS&UthD{V>YA%_hI-I9ETu6me zr23_Kw_$jmd`8p7CoM zA1$IZIbLHI^^>$-`%Cr6Ybm*lC9_8vCK3xkAz0u^oRWg3)KzA zxO>Sx@8;a$n4NBHs<-~}GlX~OmzWVmPUi)&h%AY61}mkwu?qS9%}IMF6Brq>L0<`p ztd}g{s(Q_1Ys}6|rz?tODZJ-}MkkysY8F1fBy*S@QTq<=rLEBK2{zV-Fh|FQQtOLO zg`L4_~e(FQ-eEJ2Sf_QEK literal 0 HcmV?d00001 diff --git a/tests/fixtures/onnx_genai_workflows/decoder/policies/decoder_state_initializer.onnx b/tests/fixtures/onnx_genai_workflows/decoder/policies/decoder_state_initializer.onnx new file mode 100644 index 0000000000000000000000000000000000000000..f509433bc9dcc4a6173749fa0a6b615d36e1e27d GIT binary patch literal 6244 zcmd5=PjBNy6pusNxUWf?PStMfMN4W`WtEkp*h!NV32B!VmtK%qabbll$8lS?sU2ds zY8P?offE|_Gosn0`@`X2aF8z=G zZNqmGR3}sKViFjEcW%2rS-CP^othKdXe*WI#q~xwQkTNfQ59BP^TPHgre%}D(H=aR zoS(Ej*S+?wsWS;$*4XrY>Z+R2jrz_xYF%XRdqRoIb z0p7FS(YACa6`K`U@&j`kY!`6!01$N!qAoxzau6xtydeH1L25s;ab>*pTqMQ~K0A(V z*-^H{Poa4`>YDJqsld{eIli=;HCd{ck7TJ3J$F(mFZjPoFb^}SnL%K?f#bQxh3TIw zt1nGIFhby6wL}Na6(|PYb~#6Gf0RfiCDb@j62>AV2dW%X*jxt)xicI$@Smf_%Xv|T zag$hF0}dd85UyngMo+0CHu@Hl)>D_GVRIEq=C$MJ8aIDS$OiSq5Xn5R>@V=j{{O-& z-_G^QK9O-#)DH)&KMwBAAODa*hzm#%k&U4@x-o>X;e6mQ;{$mnrk^BIWr9a*EA=>7 zBA3+K4-!7VapA*5B&&(nYMN4G14z>{?B!W`{*-{lLLjWwhClWK*YtOX z&McoJ@Z1vmTfyv|jM+Pk+3T3fXvLvfGuMnHdSf>OtINRRlR46rTQdKJ@kSjuZ|u_I z^XltM+kRsk8Wx}UJu>%!TAqE#tI&w1>}5FXF`RLx%rmk4j(sDWSUR)@#@uKJ{A7=o z#LrFl1S#w*)#xoEu&*wK!@TGdU2c$}Z@|!Ze=djp1H%bV*MQWnj4xb2_PRr5Jsm}~ z9jYtIr@W~984DecBW-!wG)aGYJu%&pp|_RV>@@_Zx7FonoENem!n~aPo>V5D?}Tlt zGxC*6+NJ8+y|o43vpIyG=0X%Yo=_M#`n|ys{9eyg0G+J>yx-^9H!m*ydWR|edY3Q! zdiP!nKOTru3o<=`-X-f{hsQY_@BIgPd%tMfOhwSyiolx|kHooEQ4v>POtpIdK~k-V zD^mjW#U((_k*|N4jG=KAK=1LxK0h3=fpRDWO7?t_XNQbu%(O$+56R|)p3Khe8zbx} z{Z_k$g9XiUr=4wWYJL0WCQI^jhu{{NVj5ta-QDi4Hl1@nMgw9Ru&cDQpWeZR^ zQ^cE^adTGZ7T@x@jod+p}zeeUbKy4 literal 0 HcmV?d00001 diff --git a/tests/fixtures/onnx_genai_workflows/decoder/policies/decoder_step_update.onnx b/tests/fixtures/onnx_genai_workflows/decoder/policies/decoder_step_update.onnx new file mode 100644 index 0000000000000000000000000000000000000000..b2a4c7288de43c04c5ae83a2119d81bd95d87b9b GIT binary patch literal 2037 zcmcIlL2uJA6pp)g&EDE=bNE1fDntYBIciqR=~*3Mjd7=fZcwIPfSiytkQGDCx<%OFbj*ZLZHA5R z0qpq~=M9hZsjvgrmkoPFg&+f$1vIb+S2*+GW48*{ovg415w_eLL33nl&>H)Jp)DbTnbau_5_X$hYnVBQ2$lc&v}oQtFrm(!%*`$wXQYF^m>2b2gPE>9eFT zj+cZH`YSEKUy_b$-*r@b!BIO(=CPx+>C*!B9VJ}UhJx=2SNA2?5xND$jCSr`awq*u z{-)q90!mzQ2bc5}FD6)U>Nt_@U)I*tI3?n6*z)^W2PyddILz$9NAOs@%#vw}>}qb*s;+&E8X3 z=4Ryv!}|*z6$6;6Y!=kcLo1<-X5wEZ(c|W;GQ<6q-Q#w#0miIwEUQ-H-)>^lHeQFU z0%Q!Mc)X?)$FI>7*pD98#yR8ai3_FM@V$|1Uy+cIfZ7sgK7GDMr*QD!OyWBD!7M8( h(x8*6J924ob<(^>kD>a{8m!-Ew$!iZ>ns>MmA^m2fFb|@ literal 0 HcmV?d00001 diff --git a/tests/fixtures/onnx_genai_workflows/decoder/policies/last_token_logits.onnx b/tests/fixtures/onnx_genai_workflows/decoder/policies/last_token_logits.onnx new file mode 100644 index 0000000000000000000000000000000000000000..f674183330680e7f8e3468568778f80f8327ea0b GIT binary patch literal 631 zcmcJM%T59@6ox4ohv5KGyHLPr5Q#)$CJ=Y7yev$31dVArr9fq9$7zQ^j2m9V#Q1Q& zhYk#=kQi5X)ysGOd_Bj>VSA`N5$W@H3tkY;kI1!>+E8gov%LAWm#V`_v+67_8YRe% zs2_16q$z$Z7kWeKb^)@I4!I5}<2co>!9g(Sxk}2hW}yg-%Y3Re=?fmxko6}xeRtok zKzWUtWf09&!1auPR|K}NdctT}vBbGJ&XDA(T)3g8&qFd7ZdL8%*vNrJ$3hqXQ4nv> zJ&YP~xv;rCTPouds5=NBK?ZMl)UZAPb~JyCxMUpi7)FXwFY;+P@lN6AkMKV1#g)6? zl`Hha&RREYu$&A|O;H^VlhL|8F5^}kp`;rqU$6;bN}7-|gVRsXQ*;TnKS?4wxU`^> g5SK<{;fvV#qBTXwQ28~3ce|XK`4q1#NNp8A0X>A#xc~qF literal 0 HcmV?d00001 diff --git a/tests/fixtures/onnx_genai_workflows/decoder/policies/termination.onnx b/tests/fixtures/onnx_genai_workflows/decoder/policies/termination.onnx new file mode 100644 index 0000000000000000000000000000000000000000..3693811ae428bef7c7f17652199285b2ca557a9d GIT binary patch literal 3730 zcmcgv%Wm676!oY|NM2c%JZ>FfY6K2pxPZ|JqGZ`|R%wEwfPlh4i*8yJ1Dc$~lqFIj zDIsamjW-2~0tK4KqUa*sW|2S1*YpF@nc+hmDn?g~HsZZ^9%s&-duB#&;Oo=i%$rAd z-|fKP2(K^bgTRks+mC6FwDNC1a2f5HdgeGF01JCEXVmlK=-#{Rhh8CV)C9w~r!1P; z4#R7sJ-9i$c+?Gi|4HP8-Yo7q6FZ9NxyM2~bk3jR+7CaDc7c>yi7{y81Tz+i791WJ zad5$W>bVgpNO5yPzw@KVbH;vR)FSPrA(30Ak$b9w_e^E4UuLggU~f}lZ>>Vo{z9O! zAX=fUuTnNxp=|I0D0_iGO`K{Du(6VhD3N| z8oB471^NYJGuN9&9lqTv1N*011jpqfI4%@Ho0FVLNTr~74X9ix;&F@{3;I>a?3ji0 z4%x^)C5d-TEw!skZ;7#!GGixKF~+5~vNZnA3&7)i*s&M*3R@?l9KLW}YV}A%{7agw zp7{|ss^aL4EP-TJ*qGW+XbvvITY2_Q!$c!2i)^P^%UlUc%gnNy`l|x`2ZG|Jg7;y> zbvN)CfzvSMuy58*!>SPdia^dMBvq=IH6ZItIkN^IOwaz2;+kT>8<~iG&9iD9lFdZ3 zgq7s%Fq1LvQe>>k{r?;ZbiQBvi7z2+IWlXAb$&j?o4toi(i<8z3Ak6q-7uiMq%EwZOFO906GuWp{Ap8M>%sv zL`x$?5yqbEmBsedn~a6GV3T)j=IfMWt0VN45%$#)2KW=$;v=*~#{whc&x%(S0v!(g zc;!aYXlbUv@6=}mJZ7mn#bG9!jJKFsPoa*#y1twll_Vmm-?OL{XP}Q5_Mml zZ2$e@S?~Gj8R!7+>7mmX9jdB^KE4oDgBDj@2kh+hUpB{+aSupXOr&QEhXge`} zi+(o4V8YD1yJn#Z=qzNe=kSY#`8f}Wr%NQ8WRDo(PA})epM7#~zR87uY4C`8`iXGK z!`q01MMUI3w=+=?zWkf}*~5p-iG{a-+Wh4bjbQJ8SIwUk*c6Yt+ngEO>lTK%oiF^J Vb&2+2_sYQ%1oC}%zYc5d&VQs*nfL$z literal 0 HcmV?d00001 diff --git a/tests/fixtures/onnx_genai_workflows/decoder/policies/token_sampler.onnx b/tests/fixtures/onnx_genai_workflows/decoder/policies/token_sampler.onnx new file mode 100644 index 0000000000000000000000000000000000000000..9e16b22bf4e14fc26641479c047034a83731b6d3 GIT binary patch literal 521 zcmaiwPfNov7{=GRKi#Wy5(LXcsbwe(SM<7*?l|xZC{ofi>%!WkB$=#;mmU2S9{h}6 zwQ08<1Tm-N$@9MN^Zx98)JjFjmTLU92cH(S5)m`4P*aNq<^Arxj34kS;3A^Ijn}H4 z!yecRO0$Ti%Bg)=_2Drz@QNwtdwZ-xopzzaaY|K&2}QN(0M4^T>SOxmN+L3s2w>!ofE62 literal 0 HcmV?d00001 diff --git a/tests/fixtures/onnx_genai_workflows/decoder/policies/token_state_update.onnx b/tests/fixtures/onnx_genai_workflows/decoder/policies/token_state_update.onnx new file mode 100644 index 0000000000000000000000000000000000000000..4f3026a099f6c88007fd54d27c96dd2b224c70b7 GIT binary patch literal 1240 zcmcIkL2uJA6t26{BoC?;BcUxwQ?V)zh(&G3X{QZtl{j)4LYC`zYr4!1u~TgzEim-uv`DKfCvF`$|osOdtR9;5)((=Jd6a+AwM8z~8NG zsX{>q#9aZ;yWq`PoCz99v)8yq&2a}j$*zP>7#FxP?!!THc@ZcnuQgAj!~{HMTGMGH zQkL@R4Q{@BKkk9Qb}P?8HCKty`3?Arz~Nb%3TdFm4m)%DCX4CNhkPp-4#|$Sw?*Y= zU7yjqKBGJO_yCnY`sDX&p?%%Y$OJ}k?e3>WAb S^!}NnELr>)uUgRP?)?GSSBa1S literal 0 HcmV?d00001 diff --git a/tests/fixtures/onnx_genai_workflows/diffusion/denoiser/model.onnx b/tests/fixtures/onnx_genai_workflows/diffusion/denoiser/model.onnx new file mode 100644 index 0000000000000000000000000000000000000000..4c43d9a1434fc502a2423e4464ee5b8c0b3b3717 GIT binary patch literal 268 zcmZvXJqyAx5QbH1jR#TM`jI=>dXn6v&z)U+PKS7Fh)kt(gX&F+4r-IP)p8J9))-}eG@OXa& V2t#029^hv}pZeeF9Km3oegG0sMuz|Z literal 0 HcmV?d00001 diff --git a/tests/fixtures/onnx_genai_workflows/diffusion/denoiser/model.onnx.data b/tests/fixtures/onnx_genai_workflows/diffusion/denoiser/model.onnx.data new file mode 100644 index 000000000..e69de29bb 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..5f29b60b7 --- /dev/null +++ b/tests/fixtures/onnx_genai_workflows/diffusion/inference_metadata.yaml @@ -0,0 +1,251 @@ +schema_version: v1 +pipeline: + workflow: + manifest: + ir_version: '1.0' + onnx_opsets: + ai.onnx: 24 + capabilities: + - workflow_ssa + - linear_effects + - nested_control_flow + - loop_induction_values + - typed_emit + inputs: + request.latent: + contract: + dtype: float32 + rank: 4 + shape: + - batch + - 4 + - height + - width + role: + kind: opaque + source: + kind: application + name: latent + 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: 30 + package.false: + contract: + dtype: bool + rank: 1 + shape: + - batch + role: + kind: opaque + source: + kind: literal + required: false + default: false + request.input_ids: + contract: + dtype: int64 + rank: 2 + shape: + - batch + - prompt_sequence + role: + kind: runtime + version: '1.0' + role: prompt_tokens + source: + kind: request + required: 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 + - image_height + - image_width + role: image + stage: pre_adapter + components: + text_encoder: + implementation: + kind: onnx + artifact: text_encoder/model.onnx + denoiser: + implementation: + kind: onnx + artifact: denoiser/model.onnx + vae_decoder: + implementation: + kind: onnx + artifact: vae_decoder/model.onnx + euler_model_input: + implementation: + kind: onnx + artifact: policies/euler_model_input.onnx + solver_step: + implementation: + kind: onnx + artifact: policies/solver_step.onnx + 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 + diffusion_schedule: + implementation: + kind: onnx + artifact: policies/diffusion_schedule.onnx + diffusion_timesteps: + implementation: + kind: onnx + artifact: policies/diffusion_timesteps.onnx + schedule_lookup: + implementation: + kind: onnx + artifact: policies/schedule_lookup.onnx + state: + latent: + contract: + dtype: float32 + rank: 4 + shape: + - batch + - 4 + - height + - width + scope: invocation + initializer: request.latent + 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: 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: euler_model_input + inputs: + sample: latent + 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 + derivative: denoiser.estimate + step: loop.iteration + schedule: diffusion.schedule + outputs: + next_state: latent.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: loop_0_active + next: loop.continue + iteration: + value: loop.iteration + contract: + dtype: int64 + rank: 1 + shape: + - batch + - kind: invoke + component: vae_decoder + inputs: + latent: latent + outputs: + image: vae.image + - kind: emit + value: vae.image + output: image + mode: replace diff --git a/tests/fixtures/onnx_genai_workflows/diffusion/policies/continue_predicate.onnx b/tests/fixtures/onnx_genai_workflows/diffusion/policies/continue_predicate.onnx new file mode 100644 index 0000000000000000000000000000000000000000..3177d8905c97dc120048da84e39b00bcae00c3f8 GIT binary patch literal 910 zcmcIiJx{|h5Ur!oCReIf9-ssXL4`o-5Vf!|@gbH>m zU9gO*mp+iysemhKpX-P+jtiY#*oq$Rd?n>jvsgsNXCc*^^aPJ-%z7hiUEg%J!TW{E zdJxT2#C3*%cLXy9xH2G@JV+S7phMF3YCl1)s@-sBnPwH72hO8F^mQFAQQ#}WtLciA z)P_n!4!r8@o{KtYINyn84cq|@6HbIQ^}-TGy&+iK6)b4P4pY(ynQ(9Q{4t>+$)z4P zR`O^`e~HUor9w_+`kJ{~pZco`701Ocl<-QI8rJKs!AunNk3AdU?Y-kD5lI{>*;xnPGI+6BoP7L kn4%_4e3}eJC}=!7ZjaF(Z2#GV-`;X&a?79EP^i_v08>02xBvhE literal 0 HcmV?d00001 diff --git a/tests/fixtures/onnx_genai_workflows/diffusion/policies/diffusion_schedule.onnx b/tests/fixtures/onnx_genai_workflows/diffusion/policies/diffusion_schedule.onnx new file mode 100644 index 0000000000000000000000000000000000000000..74c32b348f8e298f13299c0755c5726b615162b7 GIT binary patch literal 497 zcmdvta;&1&l0Syj+}l ziMgr81&PV2LJVFiTyh23>3aEjc@@RUMVSR9ddWG7#l`U%nW;sIMadbJLQK)IUJ6`N zXsSSF#1|*$7o--0R0s*BWTvH+7H8(?0YxDplqALl_qPxi7Y7HU5DOO*hkO&GI+r3O zQ1#MN^MLl00JZ5Asj=am#CCYJ~?Ro6B#I&!HHR2iR{A|we_1u;Y~v9uyH hCo{3A(!sEaQH@Ihrx~efX{pI2U@thaa4`r-001ORvu6MR literal 0 HcmV?d00001 diff --git a/tests/fixtures/onnx_genai_workflows/diffusion/policies/diffusion_timesteps.onnx b/tests/fixtures/onnx_genai_workflows/diffusion/policies/diffusion_timesteps.onnx new file mode 100644 index 0000000000000000000000000000000000000000..96f6134f2fc6bc5066af61d0c0d2cd3eb16d8568 GIT binary patch literal 493 zcmdSF zJ9D1ybOxFU6!!y)=>YkhAT>@vaUkY&VgP~#j4WQfT%37{xv9kkiOH!#3|=Z+as}Dx zdii;I6~)O#nFS?!$vKI|#qk-LsYQuJ$r+VGOwqAk3S3fXsz7GM7boWzq!xoz2nnTR zrlpk@XXfVtMIoV+B*q2zvk(^-2M41N3l|fIToa=@mm(xU_0m)GfcBICwdob)=VT^V z#wX|Jl@ui=mk2Rc*ETUaa;Xqh8K0RVBnedoF+?x1v?4PnGqI@B!LW%@jY|Qi8L4S$ QsmUc^FF3JqF$hQi0BOmnd;kCd literal 0 HcmV?d00001 diff --git a/tests/fixtures/onnx_genai_workflows/diffusion/policies/euler_model_input.onnx b/tests/fixtures/onnx_genai_workflows/diffusion/policies/euler_model_input.onnx new file mode 100644 index 0000000000000000000000000000000000000000..cd1ecb2db40cadbc9aa5d62530db9fa3a9eefe0d GIT binary patch literal 1926 zcmcIlU2oGc6s_CPrni)BMq^th(1^x|rm5OzHz6T9hJ-YUC&bGTvb@x5JW?l%ov?u> zp*?}b`yP1Nf8-C~Bx~1nAjCr+qI+%M>vME;a&0znweKJBNIv^@2QCnp(s5Zg@)#IW zu>eC%UQ^|=kgQnEt4#P^Mpo=@9#bR+ZdAq~ixnCKd1Tq0Byvu_9n zYCw)+BGOK`OV%w&LP~SpuIRidDrX6k6cU>;d3zixvSBsSn4)jkMzSahUi)PTn-JKg zQjtxok;at3rfnpPr8<1V1^rcKT; zZQP8>J0XuF#y&H$Wldf$QuVg%`f#Txc0UkQCG7`#7UkLCMqf`Lk8lx0%7Rxs&K$jM zS9D$w9X!SKR1C=o^#1mMAa233iSSdX;adk(IW7P;9hVBhJP9VQx{UW+l^lQKUFGi2 zT#eTWXqq)#ANK7(x}Lyz4rcVFcVP-o45l*inCL0@;5m$H@W1d!_j408A%R9n99hFL zf^VGN%Nb=2SHaNDrK~6W;cxWJkNk{(KFmCL6V1w@wK0IH^)jGrR*sg(Z7~mXn#%N0 zOkITr-#K~$53`r9)nh`xK8pLI75E-^PKo0S6;emx@|Ux7v;&L(tR%dP=dP-(6@6s9t?&9p-x#tmB^z^xDD zLpcMeNMc;M>&^LcdhWSi3FoJ3AQE%=?!qeq&#;gWlZb>#EJ%kV8U>oCdY{PzaQ)&guUI7=OfcFN}vT$$&UmM{}L|fwH@sHJa|DkqA(gGirXX&cJS6ifdDxbyJQ#$MP?$oGNp>N zo3`jBMbWL70_|bgS`7PJ_MhxODM>{!?|Y=x#PvaZ6b|{`KTGf# zfyMoR_J$F~jS+bgcnKw&_H1UbxFmuyth}SO2r%-ZAtfObyel!zl!fX&f-zH20ElNKg~KDVWrJ!o z$StQT?XtjpEK<~&q^M(1v>^1uewXcP~esH$majDw+9pO2%CT`jbc@4~x+1)fGf4x{0yOnOW#LD=m%%_|kD5AvQ zjL46SKO0i|k`mXo7ssAv&~=^0MORr;{2+6xN7HcRiEu<=0+YJhgJDG6ExW1O4416m z(kTnkI|Oy^4U&i|>ayq<7(Eq4K6--M&5K&;cFHW4c#+CKKuUNVjZ`Dc_u+1i4Y%Q* zh45XN!P7^c@B;vB`GLn6jW}4i9?9&5#w8uC!ivNtDao@%vZol-EbbYJe#NvpLw_UqHv7G{PbibWNfE}Zoy1Y-zaW0h zMC$nhe?hODG>)jVx8@l6?A(*C5|!@PqElUzi*Rt>J1!2{WSZ!?QViNZ|0(Of+?G3x zlTD2np4$ERnEFDm(8_XNpP??S|K}k{*u%@}2fZN=>V@s>92-iPUH1&#fwh0#A`7X% L&_C2*YGL^+<@Bd( literal 0 HcmV?d00001 diff --git a/tests/fixtures/onnx_genai_workflows/diffusion/text_encoder/model.onnx b/tests/fixtures/onnx_genai_workflows/diffusion/text_encoder/model.onnx new file mode 100644 index 0000000000000000000000000000000000000000..95fbfa02c117e7b02c14d664bab3bcc9da3045d6 GIT binary patch literal 213 zcmd`>Qg9=vHxQ9dg2W4AmHr0Uts%OzQ7 z8;X)WaUHUNlT7Cj8sYG#qi0*Ac_UC+cQn-=0HvU1FnOGrq5q#lLQZfzmg^~1_Do_I@Lf9ie> za*3KgU{pmNCu2v9oV0M?*W6YKBNZymd(@~Qm)fmd^|&;`yw0UO_3)yQ5)i38F+H>Z z_FN@2k8Y1vzr#N4;R4x#0?Bom3Hks8YyLduGT>bwe$*xh(7=eef40FEo}z9OHt$OQ l+k3FPVmxbjY$+Vb;_SfPwBd3N4eD!F_+7Pi;XtNZ{|0)8H8B7H literal 0 HcmV?d00001 diff --git a/tests/fixtures/onnx_genai_workflows/masked/model.onnx.data b/tests/fixtures/onnx_genai_workflows/masked/model.onnx.data new file mode 100644 index 000000000..e69de29bb diff --git a/tests/fixtures/onnx_genai_workflows/masked/policies/masked_update.onnx b/tests/fixtures/onnx_genai_workflows/masked/policies/masked_update.onnx new file mode 100644 index 0000000000000000000000000000000000000000..4e0bc858a5bf66d50ba9b15fb31731206717335d GIT binary patch literal 13182 zcmd5@&2QYs6&Gd6yE~FBiLPy#jva+gnntY!kwbEKb4r{B0gNPv7)1dW1wrjn(n_o4 zX2})o21W}X`dt)AKY;Wjaleke_1Ou}@IrI2?AM<-} zM)ei*@@cZ)8>g@T^MHI}kjtl0dJ=cd###9AC>qA$p0&zv21yrM_A+bj5OOIVbhj3t z8;d)uWI4^EQMP43>%kgn45HIG9Y&qlTzK#@xidUD+)W09$7yHO8)mzmew3!+Q7;}v zqt4M2bMaf>ey~lf(ywrgFzqD6IK^L(j}3G6EPN{&z*K`QbgXr8&m?o~dhuXqgEY>f z{x}YMgKTR-qfoz4g>uV6xsp(ugb1P1H$bYIE~tKO!1ovGBQk2PxdI&@j^p?{aY(Hz z#VtdS+O@o`OR$m-y)s;$1ec;iord+?P#|)M43X01p;vqm4G!bbw^sQLL(aFCSxXmg zB{K`kxCIh!^E0#GF*D2YP=E@|wX^UWgETj^rnNC?FuQDdVap=j(gqj>jtDZhJ&|HhfBXnGALo_%Lau~Oi zF#bjH2N6nJJpa@{fH*8f71g{VJDNkSP2pF$66qXV?PCMQn#|lJC9#BBJz1@uR<-`B zq*lM8R$p?!O_?(Mqo&aVfn*ih^i*y7%ncK;{s7z%uH5sj?Ch!*!WdCB2_2DQPCTZLa zv*aWmq>5TyM~$_vTZy%f86Mwb^7;PS%5 zPqFN}u$~Au<-&Rb4PAk#6trC+|0oewzbCZxO{DUnKg9@SBl?s zudG(OC#m%Ud==7o2^_^m>7E+8sZ$Bv)R{x*E@N0pH0bO^<6p%ptuxU$4@0+vVKvpD z^GG#5Gt>o-(o);OaQx0Fjxvx;*9QEF2fJQdTr#Jdr}th~02W ze1%8r%sBo!;i}Xpq=Plo(CLQEEk205<4*i=d>Ycg+MKkwkb&J0k9O8c{UnZu-QH=c zH4FKX2!tXBp-K)y)f{AA*fr0=u2sdZbrIP8Qq7rYyAr2omYsCce<=2>WJwnFLo7a} zyat8zS*mCVEdfqy9(iVs`IS~HZV=!uAl$F?)xPqNU5 zNRC&Ak`Iv_8oE%D)CF?z+06z61Gm&Ab@{GVb-@ZZnN|=epUULZrx!s!eL9;!`Cf%U z`QHBlf$}Tc2EM#)psdI`Pa2e0WMADo@O^Wsn+(9^-u`e5EQ?!wU@x<#E}O^AM4WYE zWeQFA=SS23FiCOzCN3&KJ}i9w}*!PEu9kdVYnP`!lJzPRo)PK#rHP zPWuo|A!E#XgRyE|0R(id)>YcT3o2U%fxKlfzjF_;fd>n6emY9QDqzKjq7oaxaiHR! z`D5T}uLWc3%+EW>HQ|8qb>yo7@O4oM4uG#~tOS?rJ*brT1G&6Ew@CU!JOP0RL*$+G zz+>4T9DN9Yy^A9d0kC(CM<9MIwt@qm(Xjc=6~-F48fA59E!>+|`N==d|VPj&@i68A&xTAUQ@2S|( ze2u(XH_W?Jhb5~}6HhwP{(vL5K z?NmE%7JgUh+JB8mocP>t-v5eB-4LBwU$75?N>);x6a-}}Gk!rs`ojI|M1*sbkUDfO zk;S@kyrpn6?5X8FH>7aKUMV`RD~E_2^DI|wbSM>|b6u5s?TEjF$J>Sg2R8;m$^>Af~vdBgz(^S}4r<6Fw7z>d3YhK`CB z>~z0KZUY!MCtyh7h9MVB-M(e-8RxxMlEYKv_^Q0Jx;}GgV@9CiSfO#i4&xN8RNc8v zfidrpHvZ7J5d0VAU9M?Esw>k|wY7;>diRJ$#YLaU>2U%QJ? zNjo8E-VNEj7r-I6g&k7A>HzqQyC}6(jO0+3+yFvnZqBrMZje2GsWe=x>b9z0Xekop z<>K<<0OmWXN>)16A)$g=L`(J-8+O$qN@_1^DEJYruaZQapEJPHQ zCp6!N8X|aP;w6_^-FYQsO7}!#KK-9G=1;*wg%{G{0^S#dT|i=3(aH-GTj05m;x45@ zWiyb$1==q=8CyjP;t$ljkVNP*n4)NPDxjmvTS!O)zasuKGwKMLQLC6Q71G1xcHGhO8T+=`nH|3JkIazuWj3`REIrkArHIwd8nw5y`{5J z*o{L?hU8<#LwtgpJlLMp0`aNTD*B2cSQl$RPmyloJK`Rmj|PYE=p!2$7@GyJ@SHPc ze3IW4qfa^{u4F%j?)cy!?qv8%4#d-E#(nbge}ppZbxjNZm*pP79rboOn)mo1Z#d74 VFO%(=gYZe&zq8z5A`2T^{|1aL7~%i` literal 0 HcmV?d00001 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..97f600015 --- /dev/null +++ b/tests/fixtures/onnx_genai_workflows/speculative/inference_metadata.yaml @@ -0,0 +1,333 @@ +schema_version: v1 +pipeline: + workflow: + manifest: + ir_version: '1.0' + onnx_opsets: + ai.onnx: 24 + capabilities: + - workflow_ssa + - linear_effects + - nested_control_flow + - loop_induction_values + - typed_emit + - emit_valid_length + - bounded_state_recurrence + - serving_service_contract + inputs: + request.tokens: + contract: + dtype: int64 + rank: 2 + shape: + - batch + - 4 + role: + kind: runtime + version: '1.0' + role: prompt_tokens + source: + kind: request + required: true + request.seed: + contract: &id001 + dtype: int64 + rank: 1 + shape: + - batch + role: + kind: runtime + version: '1.0' + role: seed + source: + kind: request + required: false + default: 0 + request.max_iterations: + contract: &id002 + dtype: int64 + rank: 1 + shape: + - 1 + role: + kind: runtime + version: '1.0' + role: max_output_tokens + source: + kind: request + required: true + package.zero: + contract: *id001 + role: + kind: opaque + source: + kind: literal + required: false + default: 0 + package.one: + contract: *id001 + role: + kind: opaque + source: + kind: literal + required: false + default: 1 + package.max_context: + contract: *id002 + role: + kind: opaque + source: + kind: literal + required: false + default: 4096 + package.false: + contract: &id003 + dtype: bool + rank: 1 + shape: + - batch + role: + kind: opaque + source: + kind: literal + required: false + default: false + package.active: + contract: *id003 + role: + kind: opaque + source: + kind: literal + required: false + default: true + request.slot_ids: + contract: *id001 + role: + kind: opaque + source: + kind: application + name: serving.slot_ids + required: true + request.cache_lengths: + contract: *id001 + role: + kind: opaque + source: + kind: application + name: serving.cache_lengths + required: false + default: 0 + request.verifier.past_key_values.0.key: + contract: + dtype: float32 + rank: 4 + shape: + - batch + - 2 + - past_sequence + - 8 + 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 + role: tokens + stage: pre_adapter + components: + proposer: + implementation: + kind: onnx + artifact: proposer/model.onnx + verifier: + implementation: + kind: onnx + artifact: verifier/model.onnx + speculative_acceptance: + implementation: + kind: onnx + artifact: policies/speculative_acceptance.onnx + 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 + cache_length_update: + implementation: + kind: onnx + artifact: policies/cache_length_update.onnx + state: + tokens_state: + contract: + dtype: int64 + rank: 2 + shape: + - batch + - 4 + class: semantic + scope: invocation + initializer: request.tokens + recurrence: + kind: invariant + rng_offset: + contract: *id001 + class: semantic + scope: invocation + initializer: package.zero + recurrence: + kind: invariant + active: + contract: *id003 + class: semantic + scope: invocation + initializer: package.active + recurrence: + kind: invariant + done: + contract: *id003 + class: semantic + scope: invocation + initializer: package.false + recurrence: + kind: invariant + accepted_len: + contract: *id001 + class: semantic + scope: invocation + initializer: package.zero + recurrence: + kind: invariant + slot_ids: + contract: *id001 + class: semantic + scope: invocation + initializer: request.slot_ids + recurrence: + kind: invariant + cache_lengths: + contract: *id001 + class: semantic + scope: invocation + initializer: request.cache_lengths + recurrence: + kind: invariant + cache_0: + contract: + dtype: float32 + rank: 4 + shape: + - batch + - 2 + - past_sequence + - 8 + 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 + serving: + active: active + done: done + accepted_len: accepted_len + slot_ids: slot_ids + kv_service: + paging: paged + allocation: runtime + compaction: true + groups: + verifier_cache: + sequence_axis: 2 + layout: bnsh + logical_lengths: cache_lengths + storage: paged + ports: + verifier: + cache_0: + input: past_key_values.0.key + output: present.0.key + steps: + - kind: loop + setup: [] + steps: + - kind: invoke + component: proposer + inputs: + tokens: tokens_state + outputs: + proposed_tokens: proposal.tokens + proposal_scores: proposal.scores + - 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: cache_length_update + inputs: + left: cache_lengths + right: acceptance.length + outputs: + total: cache_lengths.next + - kind: emit + value: acceptance.tokens + output: tokens + mode: append + valid_length: acceptance.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: acceptance.length + - cell: slot_ids + next: slot_ids + - cell: cache_lengths + next: cache_lengths.next + - cell: cache_0 + next: verifier.present.0.key + iteration: + value: speculative.iteration + contract: *id001 diff --git a/tests/fixtures/onnx_genai_workflows/speculative/policies/cache_length_update.onnx b/tests/fixtures/onnx_genai_workflows/speculative/policies/cache_length_update.onnx new file mode 100644 index 0000000000000000000000000000000000000000..72ad12c9f4e344e870cd12220f36d4f2593cc2b3 GIT binary patch literal 386 zcmaiv!A`C9_8vCK3xkAz0u^oRWg3)KzA zxO>Sx@8;a$n4NBHs<-~}GlX~OmzWVmPUi)&h%AY61}mkwu?qS9%}IMF6Brq>L0<`p ztd}g{s(Q_1Ys}6|rz?tODZJ-}MkkysY8F1fBy*S@QTq<=rLEBK2{zV-Fh|FQQtOLO zg`L4_~e(FQ-eEJ2Sf_QEK literal 0 HcmV?d00001 diff --git a/tests/fixtures/onnx_genai_workflows/speculative/policies/speculative_acceptance.onnx b/tests/fixtures/onnx_genai_workflows/speculative/policies/speculative_acceptance.onnx new file mode 100644 index 0000000000000000000000000000000000000000..912d023e99436d4787610497e7782c1fee02fd73 GIT binary patch literal 8736 zcmd5?%WvF77>_rb?CxZ@NlewIE=p6YAhBxF^&|V}sSh|1QIJ|8q#$HD-bu2_#@^cA zD53Py6R#eq3I~u-L2y7^xNznV;g2cbj9)W$n=Xe5?8P(V@yzGn_xm36d8>2s?18s6 z7zfuMzm5KoP$RVc9UNMLj|Rcg^W-PN=tcw8T>AhABiq69)W#LGJlfw`^IZ2RaQwk2TyutY5Lmkd z?AyMxdn}he{$yhXDY;o>4ITVNq0)+b7+Z>XImh;Uf7Gb8J8 zM%EE*WM4{|m3300cUhu$PmSomWmM`SABfw`^zM8e&&jxQqK^2NPaV zHucn6I#O_@qLlenCaXb|M;(LKQtFw9&{`aFHG$N6D#;DTMQEC`khp{hBTai5`TNT25{5uTN-$Ja;5ar40RehqKE_*ScteF(!E`K zge}cf8c{0}#Z+f#SClY+Q5Ti$EuYOe^}%*Lx~`m0sFjRGvaPFWT#9cr6*NPmQS(ey zTBvMPr{do;sVwRZ7`yCdcG>3avVCfH`78Aj5v8~!>xIY?i}+|4`=IxZ(ui6Swh#tr zR}^YL(J&BMM7zL8B6%VgBV(7Bu`5`{UrscX9&aeU*if1XO=GveR08MQtfLu7q5N=3 zq|bXSAue1I@M6>H$>qN1V$Rv=Ey2$27|~6MYoi&orp#q2v?er%j%+d7HMK=uXKhi} zIa?H&Mt_}GOcJ$qo%;QL-_nWShbuKYo%sE}FuxBxA<>5J%5kJ_I}RR&aQ-66&fJ}W z8_Q#=)6fy+0#cqp-W3kqF=lzPE|e#^>_u5vY#C6qiCPpu1H3iuih?i4Iirxs8*CzfdYpSkIhSC< z!DVRjnM3PCH;7qh=*nW&i(u2$S~?^Ox*vISHH=}5oM52 zxQG*2qeoiN{_>oQk3uU_Bcm=8%Jt&1Y580*!p?YThl4|GMV4>74sI%>@NkQ%C2mv` zU%z=H|E1;{TCGa*Tc{$>Y}ui+ix%WFec#@OPdONZ8?1-wkQ#W7y#-tO$D3y1s4QPb zm#P){;t6YPo{v;0&!YxRt)){pFC_FY+}eVja{g?BowhFV0jIbuh(b=J^hS*lWosF+ zdMrK&E3-A(;)IK+nj9nwtM4sw)m?NYW8ho5@F79MSWAAMP+!ispzeeGU&fwsGTmvylrQax6S15PO>dsKCX*S zviy0@4QW~`CskuwRh<;*s7^M?#z3I7s$$9+fl7PoOVWX4&J>!8I1?AUCf+S8fW>qr zE7A8%5|K#|W(9|12E%+Xn505eG?>>3%c;artjPO~J^E-=Nhsdrb);upy2vRHPL8ZJ zl*!roC{=83wO)GT9MMQ}K9H+u7KmR0g!ohe|!lGJ!5m8|6hb4<nkc7^A52;*oF`P>eUsM z`~Lx%$Pd7b+$VlACqD$UauJ8j$`7%#pj^dGA$_<{0t^|aSLi2UfsCbX828;B7 pDou&Ur%nr|E-@#*I61#4wOB};OO%6AND#$5pqUCzEL;o%5&*cn9zOs8 literal 0 HcmV?d00001 diff --git a/tests/fixtures/onnx_genai_workflows/speculative/proposer/model.onnx.data b/tests/fixtures/onnx_genai_workflows/speculative/proposer/model.onnx.data new file mode 100644 index 000000000..e69de29bb diff --git a/tests/fixtures/onnx_genai_workflows/speculative/verifier/model.onnx b/tests/fixtures/onnx_genai_workflows/speculative/verifier/model.onnx new file mode 100644 index 0000000000000000000000000000000000000000..7fc3949d227e6ad5bed743ad25a73484c231a177 GIT binary patch literal 233 zcmY+-!3u&v5XSKo!(~Fu38|xW5UG$3oq`Va7+yqeDlyz`cNKcso~121*m2!sTo$^zfz6!k2%#vH&Kr8>eR1zdrQ6gT zm_W=dTn@DjK6+oR=+RMIk$J(&1Uy~HE9qQST#~{$gk~BHr;d?Pc>qrsK&}2oHn1`> V9!eW#82m+|b+r40P7aKwi#KhYI5q$P literal 0 HcmV?d00001 diff --git a/tests/fixtures/onnx_genai_workflows/speculative/verifier/model.onnx.data b/tests/fixtures/onnx_genai_workflows/speculative/verifier/model.onnx.data new file mode 100644 index 000000000..e69de29bb diff --git a/tests/fixtures/onnx_genai_workflows/tts/code_predictor/model.onnx b/tests/fixtures/onnx_genai_workflows/tts/code_predictor/model.onnx new file mode 100644 index 0000000000000000000000000000000000000000..9bbf0d6b7c2917052bb237ca76a56607e5203cf1 GIT binary patch literal 103645 zcmeHQ|BoccRlogmcY9@@y`I?ddYm}xwc~U4?AhtBGjoaUyEFF18=t?O6DN0BMsu@$ zw==gpv!0o?effnTktl!=hyW5ufcX+g1QI_X@k=CxkdTlN5(p6g1N;y`;uj=S-Bn#( zRo$<;dZu?XJ$ox9zM1Lwy1L%yy?XWfeY%(KE1lZw-|QR=%O|T9qf))_7e6|M{$n0J z)$VV5&HaJ5-D!{dgUX;kY7Oo-y}g^2c7NDxx7s@%x_`ev?2J17UbC}3EIs`}6BfGC z8tr(4X0?1yJZ2+Oyg`rY`Dd=KogyDM-$ZA6tvzqJ-)eiM%cAs!)we(Nde*z|UV9k} zy%mo3D*mv`X}+-fGv2M%!S3jt{=SD5KCyr2cBS9zeKc$jI{Twad$%PZ7a(`9HN1oFYxjG@QPXI4T)XsG4971Y>|Hz9YntWHi16E>KLm(1 z*G`d-n>8M>>!OGXv6~lDk@_J;2t zc;2_Xrd9r&c%2*3Z^7(XYiHOu#hResm`7(jz5RpH(2sr|4scPwv{nwS$|*!pz?)V+}9`;G{W1bV``?-A>zU*5cp` zz7Hm`hQa55D%LdqZX(3U)pB!UDkFA1En?RT8?pap9;w!QM>3{%RjkO2FuZsn;czoxeWt2V}nM?x>CqUTF7aXSzbfxGl$+fGMYV($5q9@76=dFreHS4BjyRKDt%zEAR z&Rcc=kc6Gedw7%WH4Q7aQz1W)iL=&D1|N$xb-$LBJcQQ^I$b;`=bX+KoUV8|`0v9` zXX{66CpB6-g^kv~z9?Q(4n6Vl@pc9L}<@NSbI|Z~LzWv>{GLB^0O(z1>?)Tpj9F-pAF? zem8_c4TeEX3*wpB)NU}b4RWT9NQ{`iwswYnwSE>Y?t6p1i*pxG zd^~r0bHH0RFi}*i9mle4Fc(%Yg{3h%t_!QLcY0oHuz5Sx$ggBIGCP-VC~_KX@M}tg zJ81@I*WxQBbBm4s2a9MjS4QJPpk!yj&ehiF>cMW)sF!Qnz@ah(HTAVqH3}Hqv^5X?GR(8kj&Ye<4CqI108wNKSO{;9A zXNKX?K~88|Yp3ZKn{9sM@HJ6UUK@wat<;^u7suZzY;C@#+#j6R-ydutUHi$4@C5d= zv)j+UXFpn{5p(w~^v{bZqxa}-{rL*ei#qQ4J}ye*qLEyyX_`*?lIUg06#!ovr~#SI zBsHKR@9?+8q2%T8xKiMcKdJ}qqy~Uj#}5E_`xX_4@~sv4SAkdl%&4+;P<4kugw_?c zP(?-JsB(2sMWi8AO;lMt5>%1#7adjPDln@4ZUJc-9S$Yjs^ezN8fw{Bo!Y%s2UvA< zEeX5t1eUC4w`BdEB>_wRQ`Y>UhP2w8wT?1Mu1;$4PSAS&ir;&<-Ra$KnvHTpT*s4A z3R#Nm1R84(N!>k7WmU%pl~qjv)zUpsF{%O249K)o{oj*l%xE#5 zoK@AUY{*q8b$6759@+Psq3(Ehn;*1x4?O?QJN|$kjei~dUNUt;B)gB;w2Wr8bQbeg zhPc0<=q^1NeGGJ#t@Q92X$%6Vg-fgKEB^-kuzzdRZtXvE0{{6p;6H`#zw<#e3Zwz% zajeA2oETgm#4X?L)~@%+e=q*!3%(?dza5`Yb1-b}l0NWzH5%^zobK)Ve)i%!6^V-P zmZ?*GCt`B(-7$CZ!MxRJYVk~eS+QfD(+N}R@^O1r{Ydzgo?Xm z9c{&(3}Hrq@e`Sjqqj`SrYo(&>&m6=?WSdyB?@GeWc$C9Rl&wBsN36{Z}2rXsoXb4 zjcxO-)bd#+Q~7Lb^GnK-IcAlKE27XY5hS#!(}kpdm8hi(cS4Q)mMTiP3)(P(VkxbJ zy$V|b6pdBHET{aMcny)gAT2;;%$lg)y3b>wEe>f6ENno7kIG_)+I;tQC2%e%frAUR z(CBf5|F;G7oH#uZWcGIVE4MM_LRixkH`>JGxdq*YLo}AzVCRUf=B{Ej@9-x6rg#nG zm>7@b!s^x6--?FAVs$Gu61Qq?7Xb_-siWsT8-F)3BqGGe24RoA?9kVE#MEw)l~nmfTZUK zh->7*zOX_nc1-@vhet>2qg^dKBJRb&MGi9W4GPm&VOC44A=-{c4bgUvM-A}@i)l8n9f+K3 zWG3gvC1Mpr8zko%GSL+%6xtKD&`p{j83iZj8veF63pv-wPR>mb;UYCi&NWiWx%Od_ zb7jjivSP8xxq-piwfNT8BIjzkGJNAN1p@GmzXZWB@`L2u1mI8^vOCI;oO}E*0(0PsgpY(IX3|~d|+_$Bj+AJjOYah+NYK!Fi?>Q6%I(w#X ziku63zTQ3h1iR~1+^Ha~;<{jCWdd6Gio2U1IrsPh;}=pJ?QoM9*e1(EN46&1YiCq^ zKCTP$b&Uq>^5D!#6P{S$ATz<$SYx*)!;QGqy{;j0ipZl6LlR(f?ozODgM zUV&nnJ5dRCQogQ)1#s|DrEE(1xb%`CQetC7Q9VTBFGSvnh z%u%k5wd*c#^Riuj<5Tn*&wis0ba#9DrQ-C4$gim%H@Ntmw-SWGAQ z)*H2I2Af zx_l63vyZo)D!y*yel-POmmI}(a(vyIg0Cw}g@G32>k@}#rxEzNwfy+Hvuk3@EyX%=IaIqXV;?5*VS@m_+DTN z1mJst34-Bz4f1spfFtaV^5g3sKa9Yg0lcLscO$JcAVGu+3?N^Zw=X$UtTb|%d|fT0 zlX0;Gy#TblcwivcfHQR!UpL3VKvBa(Krb**LF(jnkguBn96m7A^W*CtKaA)F2HLuo zB{0Ounstz`%iEWT$$_DMn0#F=qmywJJ1~%IIJ&@)00F(gKn26V0P=McfFlHk{P?=Z z4vYsPfSI5ND%e;}&05xBipM7}|=J+-W zzBR8Qv`8J~>%#WSk*_;r_EYh76K5af>+(ly@^uqLVkMwD$k*k~t>o+0Cu$a(lv$|* zm_qV(rC6ZL*X4~K`831==Z~iBJ(?bZ?^R6@iK7bS>++}~Cj?p$>l0NLj|5d@{81NT z5%{`sRFSK|`J{^F$r~Nw>(({-y1XuiP>`?7113D0WDPm;b#;PBfwBf&OHz+da9=Ke zQ1EpVEeY~mW2UsphYBt?9BI{CWC9aKle*HwY4 z0rGWILDiT}zV3{HO2yYz8MdzA>#F1$3cfBcCs&~|^L5Q+zOI?Z*X0Khv1=OwUpLSl z@pY+Ah_A~dJ<=Ei&PI{=x-*7+z2ZA%-(@Mj$9S6!kgv-_gNVt+_r_uJb#*LSK)!B# zvqa_=ZzWZ?w~-THS2FBW_1c-(KUI9)aJnGAt~owmS0NhPWj8>+t_f0Jfnr%> zq7v++d|e3(;NYW5*_85igLuCWyjvNA3T}rtGxBxi8M5y?>piN-DE)fZNxDAl6ORN_j zNA`9(@^vArO(PG9`Kp?)3mOCF>&nJ}DP!bFpNrc7`MShpjIZ0siLYBwQ~e|OeqAz` z5?|MX-k&NS9j%W>wOV$Q-=Rtj0T6bvoZ|J4c{35%`dBE&F7r5nvHLy68+J3UjpJCm zRMmqBUFVFd0oTubK4FLteqF>u?9Gq&px_OLm1HRVK=GDOK2)O2HYe-Ku=L4}-u8SuL-*ip_F%yzY zkw*0}nY)n5ipJdKJ&;W$-h}>S0qOY)8qcUD@UaRqc38x@x4EjQ z$f{Xm*sb&=?Dzg7{*0% zrb|}AIlNXfhxdOLpq5Ju`*Wp_Ei^)D}rIL7!>S3~YW!p05Yq43pfyvpm z2rOPq@t@G0(by;`f$Mz&e%RVwhmyDi@QBNIeoWruhZDW~2diULHFZZ{-w+7yJr#vUxM0Lj_0T#J{ zsBf_d4Bt4g$W`F{^6wXrmf0bOFOCJ=^5bUBAZlYQKRJ?obwWvjG6-E`DtvenW7fwo zW{y1H88ha;WX>h>CHmHM(u{W`%C2F6PYL9U$FdmR8%!tF_t@hqU#{NOl9M9PT58{WQ-|jR`qHgu?kgY!mpJ~__ZkE7wlo=uREgW8m7Sb z4U|WWU#hfhr^7u`Ndl-@B+BoM!C$XFPYHZk>hm$`r(v2p^?4!}tIwNqEO&G-fwNlx?SK!SOL>{fg zYqO?P-h@_tNDWj28eZ7Y1!uM2t z{tA0TZhW;rOa`!QWX5n}<6$@9r}~Y(fsOIt{kAq7sm(j&m z`Wn=DZ%JHnahcDL1w8Ah_>~0&w$;N{DSZK6A$$OY^j^cPmnEJ~bbk&XxDg&mcYHY# zfg!_9BN3Rn>=yd-vFw5L?E&-9h`9qEn%RJOXvFU#9-5gG8MvV4g2}*SHg(Cs;?Z31 zy9JVgi6ug^XCy*g$vR9jFl4NddU})#<@m4fWnL$cfqyZ_Wm2)1_M>D7j>vX~RfXBW zB1Xa`(gNAQB0z)%(#nqwJj+=469p4dh?X!N7&4kx7`8w@SZ;saAVSl}Cd?aX}OxJ10FXn}lSTbB=PP1Kh+DIeG_I3L)~$_KWy z^MMmY_=Onc1KUE=@BX0QjjzRr$p@Bg%Q&>f<^u;NXV;?52i9|E_~v063gDZE36g9L z0QtZP;1M@X`SF2|A5P$g0^U}X43gGRkRZZ`2apfUTbNQhBRN_(IZQsVmf6X`Sk7L6 zcl*Y}1GxsAysP-Yd4>lHoXDDn)zAwMR5Ad zgad{R@_~5^Q%YwfN5X@Bn0#O@vy*`pJ3NqUIKuFdgaWz$@FPlfo3DkUGe*R2I+!UWA6SYAx_n^X?3`mmM05UdN(@Zn11E?i01M;;^S~l!gqj%biRz9= z0xWU?#o`0UfkmzY=a(u9DQ|X=4{U4lfq7*SN zft9#&Kt6CfuAJ%Q1J4+)RD58iX>A1`SSjcz_`tkitU{HU4{Rs%f$cOtFh7cjylXfD zA2?7R@qwvQh!4!eJyJ;ms8b|9@QlG>Dn2lqFo+LqkIx5Ih{m_y4#)?# zK}Ic9N^>Tv(N4+-mf!%6KB}Wl>HC0#&_5O**v`lYmXD7XZo>ijz_vv4jD+ZFf0%q= z*~qxX|HA6k*WZfx$?PWlRP%uY8{@%e;R9#%Yq;}XX?D2tUKV33J;(=^p&~Bx`SF2g z9Th4*Fl-f=4@_i2ny=x2d|+b0@IdO{Ir4!a!%ZWziMgzr4-DG_<^#v<0Ta>4p#dM* z0r|khY!M&W$%zkKP;ADDAO zSHTCCK_D14ZCAp^!f2%~$Ok5F$s|NRu$vzrc-C#M;sZzZEszf^+ncSGc50o`1^K{y zozcyS4=mYn>gtStq~K*!JGc%ii%!Z+3Cr^7Epe4b@9r+=+VYUg6)} zB)#)p`R*ovNKuvVE9i;+JGU$SUhkt}d(hb*Roc6);jp>W@dm9yd*^QH3pdWeZ-VkG zG_-7bZfmoS%91kw^-|4YyT9)ZOXVBq#M$wM{n6kE`DgR!3Gv6uaMWtw@wS`0d(9oM zwT;kl*y>0KK(&I6cTxpGVyg}1&Yh_8%2K|Zb&s|(sAz!TDN8g=%!hdCQji1hS zaVVvq%IyXkFjcY3!=%hDIa@~SAg@!Gc68+Ar^o5byE)wQMYt3u66=>hS=);mkzo?E)O zJa_pi@rTMFXrc2bsnYHbo9$M6$Ll_hmUjE?)~M6(l^)#l-<=nP+RAoo)au@cPW5ly z8hWGBT($ckI<@QdZjW|KCoQw)*7xVu_cTs5?DTX_g|*XD5$#E*du;ji<@?(G-f+}3 znjP2PT$(#U{?^M&bNEL@f81^wir>u5rTI|qQ9Ms)iBz>xZQioX%|-vW_4D`#9EJRI zJ#SU)s#!NJ+jXtFW7g}gciyU#gFpG?=RR3CmllPyJDaQY1nfVb(7$ifZ`RkB76ci0 zbA=sD{}IYux6#>7Z~tI4BnTgtonDG=z-GxL+||quh;gP*h*-A zY3cgI!isbXEiRodow)u*^gwGg^0D9VHTPP>JEbq6bNH91B4sP3Pw%(zg=3Arveome z^k3IE(Ea$%Rwt0{Y!6GHL!ZT>k42)=IIh2n9tl+4@$NQ1Xzd<&!%Effj4C~k&f_aM z7fIp~qfeC{3Qrte$^r)2^)I7G<9Y&xSv4h$$jE6#25s)jcx?$IGIJV{ncIlSD=0Ev zS%THV7h+`8lrTr62`;0agb{@;dD%4Gd>)6xV6D}0##1++L9r#NR6@(x2kyttx`6+8 z%6HnEFQK#h18<1^lR8QTJ%OjFf~+dIF%bA?RltoQ#C=v3+!zSpvnt@m5F$OJitd-R)2;+W zckQ%|Q<5KaU8Iw~lpaK9v3Hz7{vv_>=b~DUNUE$a*MhT%Myxc9~ z6U_;~n8GiQA~B-~l+Be+bkC#ojDj8fC4{rIZLf))nY`GzvfH}r4gB>p$OW3g$6%vR z>#hFi=7P~JbuU9)zAE_OC~hj3{ZOQ*gHns>5I7x6`Z{R0HKc>Lfy2%8V)sS8tK=u? zK6D08(#ZfYdh(R|L`Nsc^CNS)kuB)W9qrJQuG>$9Fr1gWuRzUZ)zs!P9ttlZ>1p8+ zko+<{dz^jv7*EcID}fDPKvKCuR3P1cjftn6pA~ZD^V4})!Qlt+1Y;skFmmt&UdOne z!0Sly1Q74M4w3pT;8jarfp;*f@aG>F)kZ$3HjsVYF@`UsOIIzK{YQwE~?FZP;E|%YNB9lR1@(DP;E_wYW^*W^m#PWs_Czo z>WqWeF)pfk9pj>!*D(&NdHHcr&C8FA>h#xLX=NbfQPUG-AmWqY!o8#zi%+V_a18I>tpcuVWlk^YY`MnwK9J)r!`$1hayuB6r3oKs8Y? zcC41~%u8;;^6F=mH z>pqP>9oKhov&nBd0d2b-d`qCl&Uyl!=vI^$TQ5C;?#E@0C4`T~OX^(rwS*TVUKx0; zkvDjA@XGrKxDMGUJ&o3u&Xk@Agy4yJLrtWvW}WEQ5+?d(|K5dR1gg&H+w3LuXl#xz z_3*7EQ}fy`C1CMAe9zx&8djiUE_He?f|p}!YIsdO!49q@KZ6S(Ys9qp|6yBgqZeWm z0@k47)HGMegO|FS@TuhSl+i5{PMsuWe z73N0TkVm@I)f_3+Mep95RjrX~&naE&qmi1%*dwL7!jTf6@d2-z)1~fpRMD|m_&g`m zV)zC1&Mz&K>RwjewvO(gHaFMpq8D|>O*5gGbC{7IreBBVxoFlDKbM=~y1tgt9yKjQ zE_bu))$_@v?$07u#}4mEPY*K5dXqb_ZEk{wnGFcR0yXqTyk(KJAQ2@UONT*Sudj+k+M!dAW-Lz~yMSJK~SP?|YjlX5t za`aAjMKvj6ZJnE4m(4EA#J_9SddRDDEWCvP$F1ONsJz+Dz(gJ_5 zK5-2{DP9MpEBGdd!U243;Wj9mLAx5C!S@jOAmVCJIt$g%`Lhkp@o8PHt+(8*V5BP& z6#r_YS_F5!7Q0^TkEEpW^)$ppioz?_^r2K|v^Bj$kv_(6+{9Kbd-C*=zbQdeK5kO8 zrH=`^!g7$@mN*^7))lHc` zs`R_Mo~lS6dHqrP=;o3>^15T`qnj;#OwdK6fIE5mXao;7ay6GeMf#|?IUQTNT+Ibi zkqpQ;v7>d%&6YmOio&VYb3A4GNK}XEqpRm~%Jh+F^wURI&%u=GBheP-ebEfovZarr z8ajWq9MZ?=YHNDSoicqS+M=wFyWVW+qpT>rVoe`Xq>scIYntAnNFU=J5-}gFO`bmT zcmHGQV=Y_yn4l}1S3SQ`rjNY#So&DYmOduvq9d*AE=r0zDX%@2KGw6Pj|sZMk?Q%A zGJWK=hv`Z^r}QyVcTVXeuNz7q>p7*5{JEg?v7S@<$e)YZLvIyRrjLnw^?Xv9J_>rJ z^s%nDx+&8~mHvjVrz+A%UVoH6HgZWHdEK$}v5_r(OwdK6Kx6Xs(G1eZhUU_zNFNnl zV6mmEp}9aR(nq;lEJ_9%+0sW@Q8=}Fj;Bl?iRv(YZ0Nb1GJPZ({q(V+=U^+{R}uL5 zAfhdvJ{r}8{PFAX!)Zi4bOMd4ZUT7~eIBacM%Qap^;W!B|ZI2DHLVJkqFY`iKYLri(Wu#yrueYG#pNfgjGyciY9XNTZsNMsC3- zN>Ik;R?mT!ar;dj0u--4mPZ=ZghWzvtm18K!_ZxyG{+jtBn{mC5jK3&n=DIR&9RC% zvJKk(k$xxd27CvCpVu&VF>p^v*ou*U3plg#EZO881il}RA32mw8b+3MQuAa`I%ybL z@=3*AA^yrDo9L};%iYY1^}KVb+l0Gegb_;_rD5o;c*|p-K)d(k2E3&5R}kfthN+)a zY6gT@QfZh8S>-qN07yQ#Sx_RtkLV(Z`4HgMC4=E%V(x5*9L zQDSMB37Ms0e<5oM*D+Ji1y{On!_Q+Sy2I4cF!h{o1vi0$xfW6B=azlKPlSMYsTiKM!WwK&D?^Gt2f?_$j#2qL)E~QK^l_${dOi?74 z{0X?^GI(*C-sly{C9gc5Tsqm3%LHX%@WKr$nkm|oP)f)i-KnQZkz6Y1Q1R6ax1?x} z984}{-V~Qy;wBU+_oKKlt5|u%6so7bgOfHGKgya&p zmBeOIkz9&;!Z~zvNG`bx#@!-1-l?oy5_O5mCGH^!FIgwKlr^!7o;vhF3jIHicXmsB40iTX-?qkd37f204#ytJ)+RdI{3 zFdm;dQ4}>Rgd*0>kx{QVY^S-QX;Ki@e2K(~Y4Nz#egj-26WrTc|JpyYy%!VGa&UX|D5CSj^j|7 zYa)i@FqC8PlIl;sDO4JDp1MoDpbFN+Z^xlv4F1V+Xg$P2{%B;7gG;;12Z=wGBSmwK z9$#EUJGbYe;T$vMzUYcJCPVS2MF2`CMd0wvI2flTVNFvF9Byh6@?M9loAfAcZh$7r zfYYlCNIc539#V*|oobB}y5pa{nxfoq{W$Dlg-k5FyMrlvMfC#YN0Q|c}Ck@`&i z(`%uH-AplXQA7D?j$UpZJzVyzTU6$jD5~g7_80YbKNfHMHAKbGO?hQI%BjKRLlc(~2`qFSW0k&k|I9V{OxpBI0k9Jw*u*w|Dp z+PBIVJ1bg1^Qi}%OkI&OHV~cMBN5*u9_L;s;X#lFu|;*Ts;GyJQ4g1v23$@wVEQ`l zh2sY7x=F4l2J~c1IK?sW*Xhr@@lpR?|QCq0fR5taD`n3-J$(UgL*#&;z z`r%S;GbQ2EaM>_eBO6q(mWkW8C?8U<+8y|`TC84E8g3U`W7-@CI8GRjX^|nY>=}*s zfeCniC>b^1Xi(KdhfI4tw%*rcb1wr{jW^)R4g+kNBTD5MFrPW17IVaB#shQah&GG| z<2eQaEUJSz*Z7ho1OvLS-YOI*5#ndA@giY4tb9u&gU*?5{ z2JAE%@aBpEr5_q_E!Th!?+qBneNd8XAcJG@gk!Lp{PoEmPOYSlP}ix))NASkRS*jb z^5pOHLP5U#E4kq+LZsF8N?3g;2vW8_SlGrv_Nr%)rzZB6OQxMuY#zT$Z&Y%c_@NZW z_n|G|bhkIES8~Vq%pjC#7ljSG;*sZ&jOVL07~Ncl*-iCG&(x!mwE<`JjCGs2woe(b zgt@U8b3#?-gaKRwY1{|TIR>ue?@j)DR4=lOr*=}AR1Q^;6Mml?3+DgtS_A%>F4E!Q z0AM(VUVodr+*Fa41nYLFh50fdyEs)ZCIP7lCGwwY*SZ=

3LENCIjbl3|;nLF+C$ zZ0@MX<19UjfpLAb0SmVpV8*&&Q?>!Sxdz5@4J_muc*HeOn*8D9cP9UHsvmhLQ@g3( z)&>968n{qm7`6{dl5K}+RCD-uHkr3oG&@p5-gnl?N3lNg&4yC4(Uls?h)_Ebu;8w6 zZEA{#^MIaXhoF$TKaQRX$7<6!*!4(4%};9VIH1Mt937IE>ET|Rb2!3)2Ftk*4ja(; z8gtzv?)z5;T;Li==Nho*9(cnsC`$G;su_QVcxnxmLEWUDQ2$##{IhW&(lHn#?0zV7 zf0mb(+IgaNleQSUZ;sONWr(~uK0!V^Ybhth%vL;nr>O$(>=u^WE27u8_L%GMh-pDy zII?L3R#uC_tXB!}s73x}T1*+GLvtTJQr_wD!N~xpX$G{|!`kP90k`imAHHDS%UqDm zej$N-Ac*`Mxdyc4FGf}mY94ivx=cNwa;f*!uYTe0ii6jSl@m9o%n)N9N6Txug|PQ@ zI5b;2t9DfDDQ~Vd$n{zW+4{pZrO9hsqntfnXfOQ`WA-(}isQXd-q{U11_r_BUL>H8 zN6lu*=rC1-^g23xs;`IsgdRmY7|>K}fERlNfA$CuZn5rVkC4H7pf2~o9PWWDTmuf| z7v#T9btC6k>c7z=JXLp--8=iCv2&~}(|@+DCWAQ!heXKK%JPtxRz@H8lPF$J&YI9j zNwRJ$@~++%jk4;Y#C8Ggb7usn`=QIYaD?T>Vz@;TdOuPlW~&yDuIbQkp&s9>GX?~6 zZ7(&T`62G@ECZ@>4~*d+Sj#=Il6#;d`5$r(eB~HyCw~*N$55%%G3o~OYY+U&g@05R z|Cv2t*(_d`50Am}U6tk6JKvQGZ~U=qPc!UxTdQ=q6(;*$i<6bwG?NYAO;%h+ty8U8 zykB(pu8ggXJL0G109<+Lh18;<=`CE)*&}YkEb8?xGv&kjCL__Pfw?XF!*+Tcy`slT6>FMNjB#6;*G{uP$mTj=9tdR~ zn7}=7lY5{D`Gd(nko-@nzGR+6?V`?6zxuGjaSu{*u%m1W!nOrp%iSrh9F6D|7r_B*6yFpns6dNW6U}IVYY8{9} zwGBx)XRSfPBQ3Uk)nVFpJ!&^$tc&4%ryAgRjP?Ev#t)wCTVf@ zkPh~f^sp<>d1XC3f1v?-BlG$d_W7)ZPO%oUWDFR@7*Le_N4W+baST?GzaCi&>fsgC zKUoi(y}2M(JQ^XxmW5#I;KEX8UrvrJ<^|K%ZQyufp7NpR2cI%0hEyDSrKD zs**#~MBM#~cv8OuUTPdsEzt{w_Kd)=Ix+b0E&N+Eax88$+~5`%DC}T>76k8pw@B$>*V5Sb>9jr7WBdS z4jw2!gC5b&C=@N3fFYxk;c`NQ>0Nc0-boMjH9h>n7|vSoFk^s>M11-tlg6tP48_xeo zY76yG`2sSilw9E#ixvT$UCIA~YoHL>cTr8r9Y?L9GN|9y!hcsC@Wbkc z=zi7*aaHuvk$>m$#2#`~;|0nwpLpD?RUSnxZ!4pfIN5Dugw&|($*Y0EO68^>RTds6 z#r`GL;j36<&4PiLHo*tqtAwFhp;)LVBw~+Fjgp~Sl-sPs8Lb{k#n{(*u+EvwUf`es zxXhaVK6?YkK!tTc8S-Cc3|K?{OI!o<__LUhIfzo=n^ynJDx+rx`kLs`WIOxc_e+t)ingP-0 z>60=Bj%5to#TYo8G0>m<9^_w5{zUQ@CZ`)UhdMz06TQIbC5=()NMU*V?J=dp{=&F@ ztB|pBrn4MwXo1C-+|ce&nsL>{Au?f{PNGg1c_q-O-0692fBgjuMdKI6@yM2FH5nu*O=CtJn4D)rE0Bg>$}*zAt0o!)&f;#=vEa zfg2bDUo!?&B7ZdbfA5p;CjW2y;XMBS@9@bx&z`G1sgNwEyc{Vs(~H?z96KU%DkaGQ zBOZw63u93DW16a5E65eO2D$X7qdYe1uHyeV$yj^!C{a4)ix}0pDast_i9%~#vG7hH zj?RrlM!$HRFO>|dBn`?}(4lg5J<<;A(WMRjy=3P7^_=Gu26Vo~JkMT8#~7H)7#L6f z!;Arf{2e&|Sro`Qn%YX8q3%%6se(G_xBTDv`#<9e4Gm9_%UchX8!h^vWzTQQo5R+a z@J)yTc?A@TphKAhNkz`wCK#ilE=j)97tOEzWX3Qsl zK4U-x`Ad_(D%lTFt;v;4t*1^GnFX>a2q6tHJ1O z;)f13!{M?i7JWY@V!?4Wk{4<5^fb>0X6P}hGPwh|&lhvOD+bKG%0BNQ>tE&oALf8K z^4B8&Mec#u9E0`bZ%EciY87>q`n!G6n(yn0HLg{$Yo;UISPPV@&NI04%=uq-D-H{z z(Y31yDs0G9YMhLbqjXX7jBFqq=;D=XOYW*JJ0BA+txT}%XlImvV>GWY`7MKSLe1P@vB?D5J18%YxoWUHB%NTf=F>nF- zPm_NJ`74t-fLcsRsvrkEqzZDtd+M+HA=4~JpmZf2Y%FJq6ZbsjX7+xYFFsH*@79pt zI=EuMj!qahe}Yo&e2`qUC{g-Zx005V7AWgSsa3^R?hwN#RKV!g?a}sDe>5xS1#RRA z)X0s-fh!3ZT}q9;^|c6A>oCqk4+C>Rp#jWwQyJgV$bOz;4w%Xuu%9`=oBV;~AO2s= z0UM?d!QMs=FzdHlIb6Sp%uA{T`z!fI)4V=1vrk8~cX5U~y09`=y2{qGwK8juom8Dj zQ)18dvum?%nRuFB5}t-usK320>b_wPC>sKsqS3f&n}AcP$vBa%K~pCkDs>c_a&$$Me11>NJlp=ov@@qN&o}B+TloRK9Dz%5YK;5HWQh!Gd z*gMWOef!Y+s`1-~$@8Vn(jTtz7R|p77Y9Fu$zkKlV6*)QL@gL^H{+DO?CN2VvQZxy z?siETZBxy7HDaDwIK|%YHvJ9efcMM+HsoK=9B_{@U@-Z+lRby(LDq594l0w%p$c-q z@4bLOlLMaEx=MU7mqA4hvP#ldVfWra>9M>dG(+^VMD^`r)U_mx`C%q*H|Zc(zBWj^ zXlGfZ*(=3l-4=$srnS%GV;L zq#oEuFQ5h2dOXj+*avLM;9TEijmJKqB>7h|2KONU3i3B2e^at&Qr2WOP@Ab!R5tZ@ zoez2q3zE&e4Ki}PN)~nXz`XV*^0u8lEEKJ*-n_VI&b{h-V}S^*Z7NImg^;d%CUCnt$is-oxegylB)M7%XEJ?K39P!?LcNEuvqSkagnKa$#VQ zTsNVzY`xN23GC2NRG6G48ZWPdS(YkHe>@oFEBWDEwQ$tk5R0oH6R{>kjav(~7{wg8 zfH`mobKqBhuKUHT>lJ!jSLt&w2i9Z`^kNPy=)3-#Ine4|cNFcFo8EQvBeA`X7nHer zS#T!%ZZ>`4HA|F1SL2{(oxkMG4aGG!PPTm?Av3AIl#iF*sm`UJ5CiE8k!>#SBPJBP`S^ovDyGuhE)s*-$bl`7=mK9N|dGETngfFh*^ z!0w6{!t;10(me)!$|u6#N{#f6T9lfk!`47OYS9nwGnl>1bo$(T>33h?`R6^Zb=Cl7 z$v-dOfQjUfBmYkFk0O6b@_SPMgdZF>c7hnj{We0SK_M;d6nB-h`qmS5zb#Wtb6d(8 zNe;NWrw8oXIV)Y8`N|JJlI7?69ppySHA+J7#wumf1~DeJENt7h#l6w}P{hI${jZ1M z>BML}oJ}w6PBK~*)j}Vv!(AsmUOmynq%S@FiHvo-SofaeT<7qf3v-|wYXEcdgZ$me zKb!pF4)EDTxSjNf;q4~Yk*Ydz#Oj0f#mN- z{yWq^RRi2TeOBprva1a1QbwNtVvzM0W{C!7Wt1|NK8Xf7S~<`6kQh`m5#^?p6aC#f z%a`mCW@QYL7SEq4-R+kdZAvAG`a3^}i<6q5?Xn*Dcz+n?4G+YZE|KVC5|4!slQ7R) zLqAK4>gDw4d{B>4R$S8w%y;Y93(^Zry2+f+8o-kL`#-SuA^%45w<5nK*$cctTdwzi z!V7#_rvj$j)yww*QPN|Gm)uvXD+YCXt~~nG4%L!5*I9!^AyckdF;W~VXCc?sH^{_h z?lSe%D`RVoi&)y?xww>KhKqN*;&Az)C~6UaGjS1EH9ZdQQ7 zFqe1ba?DukN0D_E*8%(B4c8g(*awHO2CyN2Ve-e5zdZR1?#dme+K^vOZJb#xC9!yhVO` z3duF;1%}Y)UdEVqh`vsi0bN)Ftz-?bk2Sy+=0HLIXN&NOa1*Z5-3xiGZS z#3Ji@B5Lw}R72j6vN@td$H{tRm1q3(WiP*w^}dm{{1pR=vj)1t8lWBdUC3X8{2Anb zLjHpLQT50cPOYGhP}ix))NAS&|5yI{Kjn|;m95M;>xi!7)u>jSI@=>r_S&oeOcZ)Oef zgE=ts3t7p3p8T`PUzKcu)Dr6VHBf;M@R|BE{JTo)Wol(DCYm+C{=4o-+foTXEE45+ z5r@OmP2>-oe8tHj5SMORKvjRE@;N?SX2i!yzu=~_vGXLQ_LI%3d&>`q+v%0z8PE}r zhYo;SX>Y_-55)qv7&z2TgjFXs>Zr7^nx=zyh#oicd4JlO{?81q`Ms>sFLG~lyj!vc zYDoT6^3Ni_n*4jnpF;l9hCxcA9>eY`VXuC3`|02#+@&A9mW1rgXGj; z7nFN`QfWS@p&XGv7^RA+;8rYI*|s7;CM`;mZ|Ap>H>WREKFk}Y@*S~Nr23Xe$-s74 zI;cPFcX(oAyAe2+5sl&p6QJZLBeSX&j4P}G z;9Ro?vM2vg@((8eQu4=={|Dtt=GoMK>Jn8@1LaZ$KER)?fqFbyC=Oh9SDuH2;>k2G zIJZbpjt!_T_snm~I||myfHz%bLkAU%sRNLZT1WXd(nDrtX=Iwajf|YSU2$;QYWJwd zD$(?FDb#aojj!DuuyW9FH0&3G7go_YG&BJ}_atNWGY#$z)uBo+-if%Y$Jn0y`NxrS z2lxM3*28xg^H~EHA^&sM0JX{Ag8V7u_a%ECHHe(ksXy{T3i!WJ1u^i)|4~8g``_YU z(Nc-9h;)-P-M))mjztlAyPn+CXQomy{I)1vsIF>Ol!pv7n!!EA3oEz98h@G(lGiNs z(qn5k8P)Ts5)#`iUA=LxcrdaUu6DITP0K!r%yozG4n~#PQP94N$D6ohEKxMLZ>vKN z`^jooc^;r*%{Pj@-d67M(`3))eqaq4$r^AH`EPM87A1cO`Txl~5tqKJ5H&kd?i*b~ zc@U(wGmQ0>8=Z!zzN%Y_YQ<7S;qQ@h)s12}7ZiqY$I_}53wp^xlMV9bq`oqF`&FgL zwHC(ajV6kf?|+Kmdd*S3Y;PDxx#3Y{5ITQ|gw5o5%xsg41G6-6Hr1hjLp_S}?Zl4` zyd$fn&$*E~oi$*k+x&OdfY+I8o05MC`EN1@IFi2$=RTY2&iNcmZKuvscc{PW166p^ zL>cqqlj0-V%RUu9E4}pw*Fg8ch;Vt#@HD>b13GoRuLj z2TF$%Pn3DH<{8Z@B#WD8K8lLDP4Mn+56rFTiaMG=_&7)6Tf=xfc$tK60UC5HqJvij zJ$9$-@zctHMTzXYSOX@p228!l`|zv*LH-QZB<|$jM*eo>Z%Ot`lpT5h?iz4zwNUxS zf1WVc>ZC|@m)^S~Wwr7xaK}1b>Grl!fqx>qn5T#$OSNcT*GpWov6KO04RYGKVe-J} z&q{-;*NiC}y+rhcTv23)Iqq%gh9dKaa%~6j?hD^b?1{spwMp>ptii;GT2%R}L$mG7 z0gZY78^b-jn)d;Y@torZ=aeuUbtgu{G1y3*q;69MKG0k0 zBlT36+|nBKJI)&Li4H!i^_X0bYd4Hu7=5shhZ)nZ(Gz10ScLrftbwvv1G$p_3u8bT z@*kiq$(TT`qmENIslV!jwdkH^w>&%=OC{iQJQj@ppCYB_ zr+euof|;jlo))IOCu*LdksXQ!$m~burNg!sN}(i6u{G_cNUL5C-K+#o&2mPM(SG>m z6%O~CvG`d%2}^FOd8bSZpUb=tFpnO93GaUevA$cvwJ+KCWpO@P1A4Oto<{zey75W&VyT!l&s)@p4}+Jj=0# zP07J{sPRR&U14Y&5{nv{iHO>$M&k)ud{cC&J(2eT%5hKou&-Und~amkcZD(T0q2=D zpc(lc$p4c$@DTYQl7A)n>yj;u`c(s8~oRb z&=U?)FMD$QCqC&aU2K;rIjiC@=5ht3G`XS77#1h3Dn!ct&FjnFjYE|MYu>6xjyWNQ zJ5@)KN!EynABYywKA2lH43$2{;NGZ2ge0qBLm$wZKHw7ifV+wsVCunnp3A#^2gr7r z^T`^x3i;Dm0~RCyMe?sA{{`~@)(2cdW=UmHf7J)vTkXASbJuYBg8R?B>nG7Qzc;29 zwZQ0_JnN~&yw$s*bZK&5dF?a;HhB#Zk(s8P=@ThG+QiC^ZyU>^J;y4Kq)}DQ;A-|FQ6FC3n$m>Hbq>R*6>LK-tD(C_JObs|Dpr^dJF+{Fv zIS^}t+%Y;Q5lQp*3#avga=v*WV$Jgv$Gw&0*BlSD-P{g`B4#T7?jf??l?2&8!AkbM zI!9S>V4|wk_gx~#t0KM}YY#{M>!_m_f)|g#lhQGm{v-jXtMM$bi55!?Iy~{xW9M7; z08adwr!s!0F^-+*o@Ncak2UZW)_^hOk0E~m`PY+QL;j*<^q}Ta2dT@{1FE0~{MJc~RF7rq{Y-k#3msbe6A- zyOt~q?`<#RkgA-uYoWSPXQPO?UKW*H+M>$heyHEc6R+-v;Qowg#4boc+k45VUtEjB zAv)|Fs7Jz6^7rL>pTrt(7i$~#0K0P7_mba-{1)WzO8%bYpF{qC$_IVic#JZs>O)cZ z;U;BZVQ=(Z>WtJGiav=KQ~-oH*Chs6Xlf^7-e# zYJh*Y2gn~`iNJwzINHNS?yXr2J}XKBD{D*hIrEjSZSRSeN}gTVba#2+oH_jRJ#n+- zDC5PcgQU?!FTZ$olgi8^%Gw;0^d{!>gwN(;_~B)RlrDWx?yEcOLW1#SK@?m+#v?v4 z8C?%)kY}gEjP`nrXAjUo@Z4iG&-b=+AD?Eglg*ltJwP1!r;+~-YoKD}A3^>>WPe6E zkZ}^Vn>t6`rCw0)s6WHM&rgeb!@Ti+xIrdwSY+2_WLa_K^bk4c(#`ZomJuSl>|jx` zZkW`SD~C1PLvVY0ce@teePmrngS>sjUe=Ads7$(8-uR*X6w$j>VN|@)oOgbEV{?`p z%B&AUt=ds=UmcGDU6b*33C|ds=#bi2k1IS2zTAnqTE}}1o7n%JGQjUPdDsJdWep7S zZzBH#<{)SC_aOUSswX+eQ9GziDu?=`9^l`d|CE&{MbXSLB7Afg`7;~w2O&y#Epldqkjx?F6>Y_CqRyUVUP>69A$RhyUO0G3C1%7bP#wws>poNpnYU5Dy$|*K`Tuuo;4j}Q%4$Bvq(%I)Uv z(!{QbTz*f1j(vTHm-Z76~sLCv))pn-&Rl7Z=jdEOub}n z_u}$yP+8@TYcEl<%R_M>ssSD*bb)sr7bK4JN87Ry(DjW&C+{Q-F2g(F?1APy(Lv81 zXlXt6K4J9V*#muL548Up`+xR8#mN7OHSn(<=qqF3Kh*<;yWJD5Uv*NZSJYyou9zIQ zJycFwbXU>1C*xLBcRA{qAL^EjmFpf#n{f)g=9Y&=cBjg6?Le)xz3V5BwJj%Sc4({& z-P%UXUUgfT9IJ=zs|7mUcSi0WKdc!Yj*4$%aj9_Zr&4t0t(q+Y049H>5k5x0I9TKQmTHXZ(HWD|_6C!-KK`ve3_`#yYcgc-QQq z7}L3!Ed4xLzUmb$r$4SDFFz5={MU6w>W3?$TctXxTU=kEq6&mquDQXS6(Hki(}i z7Wvr9$&I!s`?klRYuT#sTz65KIV4tAEfghFZ<)$lPotGGDUVgtA~VGF*(T7Gvc|yr z1L4%k2MO7ssFodrY0-%|7NO?9%l-3ZM%_w4*jkJ^t~=|UfBcYzGN*( z{tv8yPmq5J`OlJn7Wu1?HIVwF9`Mif0O@?|-M`{#)r?CKQZnx>s`_0-yY@lzwT)5! znnn)g9j=vA8ps=?9w>JzkH9KtGi<1srtG$il*t2P<Wd{fbCrv_hz!@-^Y8<7n#%9 z1NxD_JNXZge;RtP5U;HwLiwoI&L+2Lju;-Q$uU1#qbmzhWYAY&v#Mr0~zb4vA*9!5AXtge)fP_ z?17%L2QrXfP5u$&-%NfT`HPc(I5m$tNL{8LP=B@u+?>@)l$~Xetv=V3qwg2PCpY!~ z|0Xx3k*TO&)k$t~bOyQ~RqEK9OUvtnkZU5i-(nTF<^j?_Gf5^^Yb#qXT&6snHBeRL z&=zr|ae2NSY=^!b`y+ChCu&)Yz+z)G-@+uI$NOYV2jXp{@to%ub09(WPZ7rphJl-Dj^Qu(Zwz0P%# z+PsZQ&B|Zx_8(p&lKPav^loiX>Yf8s2Zy8O%n+#Dqu~>ufWiEplr_I6Rmq$2-$4)0 zhk6|AO`ZvS@41s6M<&nzb9fJoJzz`nmmvRdzb6$$_I%2j%rmIH)J5tZ^^(e?et94T zIk2DyD9C%i^hf#6hhE*g4U^OBEUSw;SsxTbpd)Hs?}6nJN0nvSpA^TAjA!4k8e5d< zCp);b$JSXx;JWI)@#EKF((;a;h} zJpN#Ya@X{>SlP6`%Bc5{O*WgM@f0tVIT2&*Q*V&`)Ko8>c6E~{oQ^7wwl+_{wrZ|0 z(-%W7tZ=@0AIyH~j?G@d*f}c-DxQO#;5pc6o`YSm)!{PF!OF1*zNq5QKZ>>gR`#>( zf#b6O=Nv48{Qr$}ux|_8r8o2K7qf8GdS#Hh(uY+W-q?#5D;?#_DdUZ4_r{5{-K~XR z(Gl{%!ivc49t{63vu#XdUpcS2K_0!;Tc!^5TLs_>Xz7>_%)z=Ko0VOrQ7 zwT=a0yk!(Zw()IrAKryprGZN`*8diIs513XI@8b6bM0^)_q0r>y_*x>hu8W(ErKMV^KTabM*dP z9)pa0573gne1*BOzGn~Ii~OhA1A3BwC;8ixzZKapQFi1VMQx!@Q`yur z>Miw=`sIQBJNyr?^g@xN-Q~>EZRGK+6+6B+g3XNXh-x_u4Ic*} z{#^tf-HStoqr4B^R|EAMes8IW9v#x?fi*KADvq&j4gL2Fo@?D;93=m}*PMIuw;W#zwozp*i1`xu#Eyc22IMV}}>w)Lk<) zXwVJjorj`Y_W*PmAA!A#;~?fH;l3r`#AR!F&rye;n^^;!^X^U*ebCjcagMT{xz0Nd z^Z?Z4?@9hrtgJpB3c4e`o%q#~VbU%XOvIHXrF3>xQXa zgJeTLzH^9DNwXji?CoVD>)G2Q>!4OH`C3fupUR%&$pW#=#8f6a=w+d4p0dmIV)9FP zDP{S7d*N&PNK}#y(0_9mSnxeS3%&<9Q#%5?_#WVXU=lW0(7=!H0YaZK2l73De|`E~ z;q<}i0kk>7wS0}e413`J)_Z{B)mzGsS`+cy^{QgWy6bZJCi3yt1etK5q^vq47T?l3 z$@@P&Fv%@JZdBD&ZD&pL!EBGnZ)hT`ztzfd(Z2Fb%`!4lRZn?St+O!wc}L7!P#-7S z1BUKH5R&hQYFol_vs@fLbxJ}lzKhoJU9|eP4!xG>kz13oD}MLFb4)B^J!ZO6;@Chr%zTFOnr9RrT%Bc3t8g^i5hgv4`5TAv+~(P?OXB$DGO|!R zjjUQHPzFq{D4&0AqYU!15NTJgi?Hyzn7dzvRz00@-q8>BMK~;v#qvE^5}IFB!+oU| zzUOr~I$Mu`s`R!4S^Fl-*-ASf%fdHJggVD!QBqxm8VW;So^IL!Q+;jF6R5lCbVB&->riP%N}s zJT$x#!`c**Pbw$N0~14}>Emj$)w*8F&o$LW_=-znwW>CxjUB9e494(OUp&hU!+Lcr z47>}OvqOzS zhj}U9_x0j@&S!i}=g)ANx!?i4AJ$|q^SIVI|4&&1@8JC3YLG3sFT13u;)5aYio^An}NJb?1Q7vLSgTT3$WUa-5$vg)RVh{LHkG+HW26_h10rs*7y1;iZ z^nhxSzbg3`lmDOmKGdCQdiiK(o?>xou(Yr1BZpn=jnu_H=-WY!HvKjV`-6URK!zVY zSH4z$POK_xB)DU3*ADz%%~a*W;b3{*EK&BbXel%8=P4#P#;BfD+9j5it%$di+QYJ@ zBM!Frf|xV{sXwC8?rs9US5#xZxfZi^ybJHiyWns1xHEv9tSi#eILGH1F5W)HP<@~SZ{3mn%i*TOZskzi&^?){dHjv?+mIyY8k2RB!6_-sI==> zL{+%sTG2SR4BWihK>gkUZB7oyfrTNs;mvyh$qCqdHW?@L`Mp>lzKQO~`j;M1S$oF5 z3H<)oPWrl;?0v|8h#p8g@|PihOY(On{}l2Ell|}XfOIXBkQ@+&l(n8piI6;z{AQmL zU^g5+D-D8=Zj~~9St)tNr9Rq?tEfCr?kye0cSh*FL0Gh>v@&mno6K<4%GA%cvb@=D z8_>z!k{fi2Qo;dy)Sw<;1l=joL$9p#H1}wC<%1 z{61xhDdVq-a-UprscJ`=P4D~AMl);a>Hw*5|^>=z%{#tqTlJ_Xw;-7mLBheBxeuwPYuRl8HKcp2`D@v z870qX@UgoNrq=vM;&naBcVX;HVcg%w{eFgdKAUy;Gp={?uO)vjd!P#Bk0QSd*5|FeUQ?y zq$2X=KGns2KJwt(I&iq=i(WHp8BJ3L$k8+Pa`D_AGGqM-#iPuu^tEZT#FNOP7&XQc zPd3@()qHmxvfw?C@F-Z{j7Ob-Wc1&yfn^822W+K>xWv2AwybYPvJT$Dy-p9TA^F|O z??C=o@}DLDXV$=eJ)MmmZQ^d;3g)u#+IkrvfjU_kS@O48F&x51TI+fp4;yn@^#|K0aM?xqivql1EFm;Bb48dn0< zLwsf0=VeOAV^*GcyEw^;v@e+&7`k$)2B{~YJP7y18g53F+?bJgVUYn6e!o~jlW z$+Mfdxt_dHzOFnt|FmK;L?@e7NfoJ14#swKuZuR_)pAMfRWWD`zwz}xTXm3UN7cUZ zt<(f3x!Uoe(xuuYW5_$52%h~(+?;HIbNzZ^(@eS66>AyT-)@(UfiVU0eO)B`Ul2*^6w*mXY#k@+P^}nxXueau&vY? zDw}#n6?kC3UH{+t``?}e{J#Hvb;tu>$37EpRu+*l3*RZ$h*u6~xX9ME8_K)%cXs#6 z7jK96P&)T3fhw)_a=^D8qWRh+WCb)3v&MChU*8&Jg)Yu=@TOd)SI#D5%biie_fVc_ zmfaXFn)X12Im6)HED#&&N8<92IJ~)(KnYp z-xDV}EHh5V4Q+^|H|v$xDfMyVnL%C~GgYk2(ZD)tpt#Y$r5w(05~yyc03MUb|j%8G{}9b#jEeU2eDI+3KsPL zV|mvlmAUU2*ZvLm^7O#=zveiQ-<3h#0Kjc)xs3m$i+bmj+NcNOR-gkw!?MtQ3^7hyhW{~q!oWyP`e&_!CND*=$q2a>-t`3sRhj{N1wZ$kFNR2wq?@&Hfp z*SDyGdq8if0uT5z_3wCq_Fe4d)iXLFCe@WSrufK}v)qtp86?jI(C>6t$7L%Wq{(_msw&#iVm$DaC51y;xuOk*K3Iz}ih+ zuvfU?{!)J|qz5>g9$+Rtz-9CR-RJ=Z&;#5|574(h-#>*j=dbu55AgQ2|M38;kpG|l zCdk`>7#aGtglt*2T>6UZCFJtY%uD?`$PFvhvQmYr((Pv?8jort){q8Uos8P!BhQpBE!V4RDM$L)ip^Gc#oZ#N7^njp#0^1e4eqfc;ZSPF zp-k5#6#1xz%YH4s+~Hm5rTjKxE%ve_cy_srb9RWm0QtiovA3lMc7`6t`e3Zd8bLT zw#i~FXDt>&{@LW8Oa4~m&msRO?ty~0AdSfpL#0y3s2kMpJU^;7&#`V7Xu%Jg>p+4C)Cu2*=rah2alBYze0dyqeh{58mbf&8z?zn=UJ$QDVh zq6$3F>(rm|K%?ht5b(5xGWJvi%(SkEZ=NL@y$=-)MirCQn)X*Uc(38NfFoqWa5b9D z_K;f}SSfbt8oXZpUX)Dvsw|tBBprK?ko{{^mpTUr#*YDee1L0P<{>c+H8ZVod@yT zWA(*yljlKzz^@3qWl&QB-%WuEtv-=F-?=>b`g z{|Wj3l?Qr=bAOF`ME$28c<7$0at`D0hwKdFiK9L!{O+uh5Zf1LUqsl}TJ~1-5A?vc z!Up+(y~z80U1W(LE0m{v(`H;;5xX<4DGQs$NqgH!d01UvX84CGW(!`c9L-LM`Bv4@ zG{zc!0Ru56&<8WVg(C574DRX^5u2z+ULfQD2EGT@=;2zFKf617zq!oi2bkx`fBrtl zg8bXbUzq&-HV0O6{a@hvpTqTEjjTb`5=v58R001l53nHr{X6^r3eWYLd-cZn<5n`$ zt|X?7QAv+lvAExLiWqEffd&t{qURFUxGPw%+P0QY^NuM)J4d42;aYg(m!V8v94#*< zN6Rv^&1AhBS|uR%it1$fBVwv|H5}~F8RKU-p}`Gr1TGJS|Lz!kc1c7e_fk<8zJ*)H z_u#R73-g2D#2ChSKASb|e&+m3tog{lgC1Z4`ETd5&nN#(@@J5LD)}pt&!1XE{S^;X zt}P^|9p0gOTU9G%=TfTG?y+(z=ejEY8DLXW`$N;MA?t6oy!49S_3l_-`k#BEwC^2) zN)yfDxqGLw`h0}!G$U3XU)V(6SvpSn)?>eFbY8mHFrx|z=Xb={{R1$yuQ$TGhT=p* z4ED54MC*|T{!!#V zNd7;%2X}?z@R0g7|DX8le`o$r#yt~x3!2KFNB!jX&u^5$)AjOWN+jM?vBUI|Be1+i zz6fd)ErVZ2AdhEW8@CjbWwpLIlWPfgWx2BCc&Ieh$IEfknoDEGG^J66WvX?T_loz% zO8h2x2OPfVh|B}@fba90P@Q8?t3)DpH&tV6JJ$abble|$bo;;_a1i;Yv+vo%9B_es zF!{f|;Mynubn=fSe>C~G|9|AYXG~U6w>1i)A~ryz*b6Eh5rt>(RcT^HKu}St6cv>! zy^4ZV6>Nwo#R90)##lc|qhu)I$6R|5J$naFeeP|GW5`;4b#|nq;XX;=a6<&XtB5cuJcO z+49RhK<8G2I{o1eRvPUqjX&+n4ddU*g`@o?a(^Dre{I2?GjrrZ$|6Z-8!g=rHI)jV zZjujdUO|2KA7-(y`tq9a3H*-o4F1r-oA+P27~X*k18oB|;68DD%E$y>laR=Jc)@&d z(13S1!v0?k-}6zxcoptAK7s)J=RgBo1o)Q%ejmVJ2KZwE|9{tjYXENx#QzEYmz#(3 z>P@1gn(n~r?<7kGN}4d|Q;|C=#B(jkKOsM~*t`Lb(y7%hywc{hy!C~y^f+uL{}RRc z6ER%wtPmhsK8%sPt<0ro-;3lj9Y=a;z_ZmL^x^5;4$uHP|V8rFT9-0GDzbGC(Z8160WZunX&1AX^nhBa^Z+#ugz zY%blkp2BN#?RkabGVzsH5~VOdf9?(JB3<076n1%L+Vn8Yu>NaAy*!1(DCHXAQe~z2M z7su+#w>o-Ccecb!R=Y$gwqw71d9Ao2lBlrjBEKs z@sL`Wjr=~Io6Jk(d6Sa((|h?#(CEyQt%W{Cf;0sr6e|M%bEzvB<_5~YE_x1UpEp zJMB%#`{&zORJ1am;W(BTubs*x?|SgYKa2S7U!nZ41jfZr3a|4;Qnj^1m{&gL5Oqg8v^ zSA$P%L>C{PUqhtYUBh^L!w^38Xb(yI+zI(x``_%-gYepY2WLyagAMrppYD8KaCx0u zt*c~tAW`b5wvrUD%H{i-C!9Q8zL^zo=*B-QnDX#_Q}`iUn3!L%h{v4>L0*g3i#iE2G9rahXDRrfd4JT z|AYoU{a_Fu*0Qxu>+KeK%0D+gauE2bo0dqG`VCC2qK!SjWh{wa0rKIa&#_f~OK5S1 zkJOwogkRM3<-fyC>*9txNwT;k$#t}y^n2%d`Sx2YYqbK_Gm8T-OXZ+3-+5s&KeyMN zw@p~YH%5i>yALCI*g2mCVt`~Q>%ZW!#u-yNGM4W!Oe3DjLTQo=Y-NtP~qD6oQYzU*P@G)cX) zqE6LsHCvi(%HoCuNK=4OZrZ24w)fCA%)UX9 z*LoUpi)WMgp?)w|>FPqRHzAZSsEXvl(_=Z>8PAU!Ci0{DzCiVNNGwgx}RDLjULPt;?x&!@2$zM)cV->$s<;8Gs?u^$~d)n1b2B}=1f zr%Kz#+>*~9Jhtx2^$d2P`X77O!w}}!+ViNGdAvSwAs^Qs!t*mC`JJ(`yf81G7Yt40 zH-^JJlt%CYEkNf_06%9BepCNug{D$~{)4-NJ1HQ9hpmYUr6n5>w*SefxXHPmw!;W5)f70}lleP${ zCd`o)zn(;EhaO=;GomERQSVufei;8%RX|PpSxd_Ip%3(hgJgW|zTA9IKwZYlmF)e) z4%Tr~O)w(~ID=Wc%f!^6*k4Ubp5zB8T#dF_&@HT=D=pXe-Jk%WO{y5M9cLSaV z&^_A#C*a=z`2T!MhC7sS7Y6X#0Db|mw?bF}zLgLMApWC)|DWLh;Or-<=Zs)&I-zpu zqZT!<=vp^5KttL#aFn#+)E)UfizI2F?S7`WFr!ZP<^jvP9Vgv&xW+E_iR6AqU(>4< zSaSTEECqKsNoGr*%Nw-T)E)ClV6zl|vUWE;jtwXBT&uZ!Md1SQ6GQmhv5~yvTMU1_ zB%Z5xPvno)lK5;nXdq(&dkm~=Ip`VZVcqXS-SZOeKLGygfZreR9|ioDfPWldZ-JnI z`M*4Xi}3gR5dY-?{6_=(-&y|%MR{G;?3sMD$7R-a-bZ=s#z*p{3$o;2hq*{jUj|D* zvXiA|)XwT+MEO6zuDtq3qV(iQ6`Sf6!*5p(X8x*{Qjc%RQot`4Y3Pr4@(<3t>pmV1 zV?Ty`W!E}Kb8lwHUu<{d=HnOep`$~1AEiisxHX2~bc^TZzY_Qr@Bp-`lX#gC)Lc=J z`^o?tcmN{cj|TkefWI8@PX_!u0KXC7HvsJaQy#$X@+t6kUmVQi_2wr|xk=?kcGC4p z=`<@aTsj>%Q96AcxNV6xZymWwzQ#m{N9;+Keuw6>-a1elSlcn*NoJD8>0~MX(QGLs z?5jK^?QGre&IRn9#V1x=GK#-*vgJBm-T1(Z0ldefV15(k36DJ)!?#R?H^Sa0@VyGq z19v2eTk6A}UIu5tPT+B8p$@nM`WN88`Vn+~z;6ur^8o)az&`@;Ux2WLHBN-s3vmJB z9>o9G_5WNE#pT;G2V z7kvWw;w{14wjzRuY=b^<^LW1MQ38+tnaHE}Lx0^Um~9sZ-}4U8aO#1xwgRpfQ0oJJ zXTYxp_>%#@2H+nE*v~?Y2fXnRdm#SH1Nn~!@E;HCe|P_9e+*(Hwl0&eHlD!q!TY>B zHc<*bV!~6e#7Yl_#7M8Ex=Nv^rt#nbic;iG%%8(tm_c=xO!YLpX(eW|mE%T9{qH78 zrx$rh8U-*fx%qe9hh%4_?)HjlKN`WOgjsXV`q_N#H-A3SHJHzii{KH-G5o;Lcz(4d zft!Cy^7eXOg)$;I{+(Qvv@N!2k6pd>epY8SwW5>=J|-;EaW+ zfcQW6KotIe{b$$WIq*KKz|rM_yyN~q)}DKf?XTFwbfU6YtF8}2x)y8%?Z-TSy`_FL zR$5k&Woy#O)VO&hLe-2^Kk7p+?S6(m>`vgE z#0~hZa|k|i*9t2I4#F>Dej~Tr7tq(T?I^h4N~FhKkS#VuT})Mwp7wpQ!{D$OzHY6! z&VQlk5{*P%TP^XK@n@k*uU^<`yG^(enJCElbYXA)DB&tIN#v-dLVKSZmp(?`F#|wMfg#_=%R!bcb2KqY9R>2d-F5UsGmrykQYX z7ofh_6vfAtc=CQ0oIAf5&gVFE1Fm+LeV=rg1=Qs-<%Nq`K|EzG3Io{tH{a>n(GB!T zUJ1PrnnG9YaiSq6`t)$7GIfr;OYZhQNIV{{Cf*hPWSgx(^b`A&mChe=!rfE&<)T8| z7!!vl-=2csm5s!K{gm)X-h!<09;9HJiAHOCqTdGQD14nNTIc@~@YewTO@KdKtkSa; zciIjSZ$x$i{!4&=C*V&7{Ii51w#LHy=AMFc%|pOn1^Dv-|6|?{URj;0$VIbknp z7?(xP)OeBNgz=iH7UTZf5P#tsmdKm9KB`g&e4cPSBOo2Ll= zqLWb9q$gy_lm(afcVzzC4#@ISSId+%{bjQ*$z-Po_m|x|`oZ#v$0^G~r9#W_o3WNx z>#U{pspWM?R&?e4+dTO$sK=A0fj?8|$3vuvympu#@3`5MA2)f%URu}^&!b^o%;=jteQ4$OXJot839>nU19{{aLQ3yg zk%&cuh*QjOtaAGTZYtZ3-TJM>YTO0Oup#a+QNhEt@1q|Ehf(OdwJ6trA(|PD&`4V? z^wjt>tbaYM|2A0vL=pDC2>V}z{V&4)hxM<7_0NL!_ZDFP3$XtM*#EHpwXpu1Vf`az zu>WPS|7EcMWw8HY{mU%X5+bFG+4j=zDYxqsW-0K(NA7&`bI_%mL8Hy|=X)>ObCnVL zT)(y#H?4loic(IpS?-&e(!(frs&ouYqa&2GAGSq7c<%biY{HjXMh z?MweUydcU2)kHsUBO$6`WQ>gs`JFYGbHNoy#v0);)1G)) z)dSQs`UtvHl#9k_2BNnu81?bcMkVcE#MK+liaP^~#Ql!R;-)$W(W_~+xMxc@@tMnQ zK|gJuu5x-G65o09Na-$zOJO!)vokWxDl0%d(xz%56J!4BSdd| zF3B4dNXqVD^2T8S|3`_vzS@iCF2lvy=8B@f*bMu>9QJ=Y?Ekrf|F*G0 zUTPmfN%I-(|KqU#H^BZ6mOb&Xk`*cql7-*=Wod7w&-=|AAh|9}kgC=H$vp>IGmB?t z{M)rqe(W`z!KvXqvULt0JzV5zmeBh)O@Z$-xy~$%_Oca^vzT*>4{JMQ!Hym5$6Tu3 zQD~>816~)Mm($D};Np0y8`O~p5&)UO@k1&EhrE!UT(bPcR z*KRrwiZtOj($#og=X(|yafW3NDq=IoEN50uj!Y{?kDW49W^dH*((A1U=!mnc>BDS) zdSrq?=SB9XbEbSGmbXrkJqrp+SX3P8dus}rTsD$u^iv{@tOfVOd+;lhOx#=B6Tdew z$BWmg;w=A{$T7MGjj-K>o*0LtBYL)IhwTv564?pg|0VeTcf$9d3g7=M5%RwX`Co+m z58rJjZsHW_o5VGrhEk z`I!)Qzg?RpPWVRW)}Nz4mle~IS<7jrr6bKy)T3YXyVH`7t>nV>{lw}?4hd58Blmw> z5*(pHQp?`sfE5zn=CTDpua3cELMG#*sUvWXeitk|e+{Mf+Krf)fhuy`k@hqxtD_jzZ_@EBO}z|4fv1D@NzxlVlg|amco3lII)@g_1TxvJ(zMu8=*=9Fp3l;xYkIVka(0G{$=piHrzQ|@$7#gvmM$5j+?5=;a05HW?!}&_ zS=g(}3;&B6k6Vh=@Nmu7=vB)})Y~H;6}K-#IUDTIj=-U4pW{Esf0rTu?SlN52KjG} zsGnvcF1yi7T<7r^^50R&f9oLsEfH3zGhve903rYGcgTO|A^(*?{!4-U=OpuhHFwqN z!Be}~@OL!|(hmoy<(pKcHnpxy*ssSGaS&GoUXZmag5NKm$47T3JnHoj?zg84e{OLL z`iA$jhr3p@7PkQQ%}ZpFdo`JCRXZJi@eEzvXB+L9kVp-*r&F!OQ8X^MD`k6bk{4qu ziRYOta(Rb0St3jzACWrwqxJ^>E2_l}|MKyMUlDjp;6yx7a~SSwu7Ia*y@Kv0m7%Fd ztB?Y9LyFd>sAWlSlPO#On*)&nXM@uHBW5nsrdNLuZYwLMcGj>Y*adu8NT-r-G(jx#T3P2$i&rW=Q{{ zG~_$7KQp5kH(p|t$dyXtd9S_;c>9XU{K-WF-VOEPW>a4=xAbDx9fXPe;YJH7Be|Hr7j^e8i40`k^ zF4N}ky8LF=TQ0HK^fFeYmd^Tbp2vbp&6w?`zRX4a1=Y>0rbQb!(&hSL^q02{{ZT%c z#w`CsUM4gV(?2`NxG||@$sbp;V!IJJ>DZIZeDeTr+9en{*>V}%Q=7VJc6&T?8d*%Zetg&9%c(n z@)#bvgq5o?wl8)7i`oB!=F~UPMe3!r;&uvM=HyH#9MPw%zA4iYx^3js=7Yr6X$_g< z6+mom31s${{^W=1C;W5SX}s}SAs%iOkKaa5#a{ZlxH?V=+q`T+C(rCbriGbEr@<4Q zJ8q7SbyY(r2ET&y?_@>oUxZ)K<{&i$up4Rr7`TUl%6jaDwt&e9YdKz2fHi`x9=*nU`Zqn!3 zmDG1&Hnp(wq5b3&Xz?C(I_AO~ve}}JOp4n~CY*{Szg6tX+fBnr(-s9XZ{k&KtXz)G z3s&K=AKY;BRa5-tL2sgUyi25A*DgIm7D85Z96(8uX z5G&=*Vo)~&(RD%(u`abuIQsmcFmm-8VX+h-wCjpO**HyMV8ka`%GuMhg`>C1)EBLn zUryP{dL_-2W;l8A$|29$xe6}13=WrW*`1NEPu<8Ghg$LZR#7|{G#c$q!MyDFOdj}T zJa?V1#bu*@GRtmF?4IQ=_Nya}W!`XO>&?e9ji5fP+Uyza?0SOE_t-!uj18d~jaD?p zdl3B+{F~@rzd-QL?L<+1C8?k;Bpw-(0uvQdt#u!tHaLu3*R93j{tK~fG{OUIwQ!s9 zXLLZX9&NYXhOR{>qB=eu?a3d7fd3=l|8V|S!ug*C=fAfI{9gqAF9QFE^S>6(|IKjz zM+(6I1>pYz@P7gLKb-$%ve84Mq+@p$^6RIPrDMYtc@nbc7i@mYr)*4?%8UC+TKb=< zD(INmIZOZ7918ab5xjA*JKvE>c(DIaeppM1PycX>%^z`qZ639T{nlH+7AqsxWr7wP z;rE&5C)Cr;+qTj9Ly};|(sXLne>An}*^RCVzD0&!sU#1qv&oDwAENwV0`WhkPBdS= z!Lz5=;mGXG_~_k8Ts71lPb(jWZTU?#P#C$qe0?NrI13}S41^}@9zxibHrcoz z2W3aBtfbI;!F<;XJ87eI5Wh2oH@k`qf z>~%~@ykma@_|IP8KbgRPyab;f5NOze{i321bp~Z4f>k62_4K1 zM*|~lk$K7xqSaK^>Y}cMG=BzRn?UZ|o&($9Z+jFag`jkB3-ipOS z@oyqzOwkd1Lchs=H=dI<9t-3@);7sIJ+85^4<3@usM-9WUk3|6CrB4}MM!OnugDu5 zv)LgV%D+qo?`j@+OnxC;qs|51EStawyK3_VO~07+l1pr{PZ?YBX%*Y;H;>(nGGq12 z`m&5W&uKr8YHA&}k=}kEN-NE5=*^VD^smhy;yJB}@Ebcwm+q-#b)zde8e>HInD!(+ zs~+H%s3W+xC>Qt848(0M7H&!*7`>R$x)FYv#8 z!2hy=|IHVn{uQDA6`}qG{wD|iR{;DkN`U%TfcjT}`d5JZSEe?0HXq#{0@|(*uL2(@ z`nw9Rx^2aW{8W@;Wy#WkZmLpaFJ&g^>ho6U;r)-^93LFPZzRs+eorv(**Jv%fY~uu zm2R;H&;2YKX5i!x4`67V$oyAnGCsAP?k_k)BVKH!FTW*FhiB8MhU+N0G^8v27IBk& z>sm=p6=xBprQRf}*93A*OP%E3dW}Da)Z%-0@^MUU1nwzJ#7o}~#VY+2@aaWYkkhO( z)S|Kq^;LF5nVP2P(&XMqOXrCwGe0IyyuMB}FI_6`KET8iy9bD!m`Cg zL5eWwp_5?mtS{*8P!`T^xhvD`re15GwLsdgm&nVqyrh}oT5RUWAfEhvg?#CvM9I{} z9Qg2rTEFvZ{7Ql!KRq72i48IQN3{<>xz>hn)ELDxhW6mLX%E@k8OPXf)eTI#6~fe@ zhekhtFjK7lL+x%g(P8d8>EWbQx;=juH5g+|?dp2c=%Rh`=Lqq z-U0uW1OF`m{u?dI&Q22Be-9Vasue|BjcdSvD}eu|1OJ^TjGAR8?9k~eR46@{y(o2) z<{yU{fLqP^l2wXQN=24Dq#1Pgsu-f<7$UVQI7zPqhsdjcd|=n}=kTj=Zzbq@Jg*8` zz2B_i6GBA#n?5)G&8 zi@l1KMcoZ|g=@M8g=e2u3!Cfx1*N%yFlcbXF>0>s@uaB>0^yUuiqD8ELuHt8)>(a{<#TH*IR!{)JrI zbe7y+T12d8CzFzT2ePAiG#OOfjZAdEjm7MJ*!NC0&i9{>x0YI9$+I60e*PA@>&Vfi zu>~mUX%y=D)gFy=8IC~zL7@Kt|1SsrpAP(gt_b>%2>Oo*`Va8`ray z{}Dj{kqPO_^1ne*{9T2DRKNERn_{p)x;wyuS8eIR-(-%H#(jyB`la5L&;7ZaWw;=2 zU>CvLGNBL0G?X7W?aF_BvEVl+>F|pY9c)?4W%ew)oK?L@XSv_pK{p@AToP1SK-x>X z>u3!fk+F%^9tfwwg|>8Z*$^5&@GqIuaEZis-$kzcNhLF9&n9-&#>C647x8v|glC_x z!UOl^;q#4)vEq1&|3vBF^6GC$rR#ZgA+H!EZd#5km?Lsg)I%W~x})$Ht>W4X`^C09 zIpSkQKe6b8rKr7FLwvgNz0evV2}P5)2>mN!gwsBg1?|`og3+)pg0jq(m+bT4+Z>ik z&4)dB?!#p1&y2o&of_jE_qs|;I+LU&f56%|WeWGvl&d`#qt@QN#1bSoXGbLt8oNgaMlhBxkK^8 z)&J13kC#!ymt9CanTEVG=Ag&}CMe^3FO;(WG1Natq5fG1_0LlAHfQ2C#{pu!)(@zE z8le7x9O#e&^^cRV?t`9Cryu3RJ7Cv-kRF!rP9rB2`tVdfImpsA;0z~ zNjlFBC6nK$Nkg#~&(m1IpREEzX}p*w#g-e+lkRV+Ae9aG~$ z>}wNeqaF`rqrd*5X$2SPxt#6vsLD#(lethmYeVW_sX~$deR5DdOoo-LC41r*l9Vh& zp3KxDYwbVdwWJ=ea@~eU<|X2bv!>(7-J@{tLtSw|+f8(CRwerWFbic}_ePsNC!i8H zb+p#%jX2?4t(Y)$v)E}EDXQg76m>0ziPxtoh!4+Q5$2YZ36t$t2}z6Gguec!f?7sz z!LTZeMeOYx_Z`&&($-N79^giJDeNOJ0{(%TiPBa?Y?3w#n+y!yHeYQ>vZ$v zzlSDq)im(T75#bSIC~zrS`VtKUVO~0XUwp^hCQy!XX^JO*}<#!>|ow-c431e3+)N| z_?imp_9UJ5QFW&wYs_f<)V_3++jGKOPLMP08;En?P%>zlH5vJA5Gm&!xH9@8rusW@ zMchhk_R<9(I%9|p3RSR1!+mt-_+eDlH5VNlya?6w!stYe7AkuG8S1}!sQ

{g)(y z{wad~DT4kfg8m8hUnSIk*--!a2%vuopnnRlpBhMHV!Whmv$nSTJ4J3ef46-6vYYa< zZqSEt@;x;jyiocucdk@@a;7{e6z=t$n90{e@7jh73A{dF5igQW<6VM{`E8huH^=uK zn=$Y-yZvV?n;VtHJ~unCPgY~tb<^%_XIv}&wRS(fSd&BVKlG(r+bpS(w+4Nh@t&Mt zDiI&sE#yl{45@OTOnf3o5YNF~$oa$9@aK*SEH}=;8)Dq?`!+LdpVk+jI{X}+8c>bG z3^t<23qsMoYHM_D@L;rWONV&p)J1WK-wv_Qo|R(icvo?(f{}Rml8X33>w(b!=V2jT z$Q4NFBB8q@7QTJb5_Wz}kfP=|OFov`eD`z>67r&w9yc|Yd}BlSUx(f7^&dxRbtv3# z3D_+@*j!Ns&vpVHK2I&^sTH)7xYJo&w$m^2nFC)4a4$+qr# zr|5WkA2Hz~PZ-@pCE_#pT@e(qfu?6)$6odXQoQ&QijX)1{ zyP!(>HK>1gL;afp^{=~l^tqWha&=#EvGg43-)g9TH$wdzDx{pX78Z^kEOdWxMQ*%r zpghRdWptRToNNN3N@Q8S2K6{3_bnQfx^wH*yylbBoEM`06kG&)K%rxjd z9}vpR%3S%72fcd+;SF0%tm%USQs>1=Pt4! zo6Ey#ailH1ls<$yyzV5m#g|C#i=CwZ(^T@ncs9A4Z%o28dy(JzkFdv%D*QG#4|^0Y z#sPgPPMWKO2WNdlC-0v_Z^MgGVd8Q$T+0zDJsyJ=#C1nWwXLGOe7{&I&k?`h^c8WF zrMS*PLwvvRy>Qi063!WJ5q?F-2%X|&;qKfKf`efn-q<*Z8*B;T3eg^t{qyNClPpOZ z<2r<2tsBn=D5y%4KPO7}w{@5HwvS<}I*j>>nIZh~Z|FZ;u#6AvGnZHTAnxTem{&O| z^4qUln9a6IHl$w;`+e4zh1knj{!a~d{QL)cEAJG2Y*R@8IK|P3GgD~ytt08t-b&O^ zY#}3rJ;cj6lZ0q`5={eh@;gtJ^!I;>)uU^$g6$@J$~YXC>e=EgwnOm6$WBzvFQJP3 zo#@c^RAha278-18j6nZKp#MYtUj_Al9@PJfMbQ66(Emly|Dpar2lan3)c?z&{+}U$ z{x67tBRx2Y>&)!Io8T^YOg+3+_c~T8tvbz8C(q&OURUHWhgA};4wb&_&!v0Zhw}J< zK(28Uda7qe@qgHpf4ac9U&?T<^R+7aeThztiHp z^E75)30qP|Y6?DPG7_IW*aaUey^aEEccb$S8OSf(16k}Eha&V;(WnzIM6Y+%qCv|> zadAwTIBAoO7=C52`0d0WA+x+mP#UsRXyd8Ew^6eM`57bON0>72RcOc4b{wfY{ntzS z_#lp#4x1|tuzyV_`-XDM$WVDiZGzPG=>(}!`*`ig+kLrfkuUGmfPSn8F)-(8K2OxJ z<(he;xZxU@`(gKpefe;V4cM@ObvY2qv^LwYccX_e^^#8dBli-$*}jub{Rx+h+-Fn! zv&QrU=|$fxc|mC89?5#Kn50jmL?v5?*xmhxV~3u{#^uGhzI-{>ndgYFYwO|l zW!-U=(jBDxct2`;mxI!V`=J@iGSrc%fxaDm5Bi4$`o|W~KVrno@sq_TYetB*W?jUE zE!RN**bVwe2IwE|!nhx1!Wq;iH~q4PKGHIlQk}a<|GGKx_~pYSryC>Xi7kn|k9T<8 z3YQ4!-|ERy-p#jlZhbnLzmhxGJdwyvK@WKJ!JijFPVBN)pYIvko3HhK!Op1Gvb)Wj z*}CFrww_I99w$ezlWAR8+uG~&TjOr3@F)ZN3_Ykr>o~gCRh5!%FUfJ^8e-gI6X}07 zjL7fVkhRJ~h`VMdHu>Cyt>^B=(x#%)~hB(u4jQF=}chTAj1*^6y=C zrf?=VJrmE*Ux8V0p>PJ;P2=)7V?Jk>DmRaM$GRAuW|cj+F_qOxEcB@Z`{Fo;{YKr{ zlr^pNR`GtisUe4^e)Xj`?=0!2r5ZG~=skfpOS05?3)yl!hI|g1Of*wR5JA5Sd3f#` zp51FV9wTPpwK?v%=Q}fOwXQGT-|!r5HLOP6WgF3jgiy4+$r|-E8jQ9a>Hz)cBIrLm zK>t|@`j4v!{+9^;mk9nB=s!n5|H%dYXORH@m+(eDN=kW_EV*3rmu9JVutb&HG{W~e z{^veM3aVblBl9-0d8b{a8D$Hlv3)kn-&7rEP9alx-1u02Vq6kmcpT~=Lr31`V#eD? zXmIu9&n(C4Jd3j~WqVJqVB1%@u-#7$+1ZaO4Cg$c(NRZe&jY#iuvH)(r9fy@sy00% z`%3m-Jxg|PDI!BklF68V4rGhM7*c<(8(CLy8!xHbhg%hMaMP~&Snrkv*3RgMD-GVE zd5&^4*rNcwHH=0DCX-NJ%5Zdijv^|FY!)Z}DHkie(nW9Md7^pqSTSW^AMxMpXToyD z6GFkH4Z^eE!9vCcE1?*;n&TQvxWD&U-V=N9eeT8b13mV$M-8sh*1aD5Y0GPNi!!NT zcewQE`dRsP-}P+UD$Y~;MDb`dxSM=Bm~ZCK++A9TZB8NlK2Zzr&i;(FBI?n&lx?WbokVnK?R2Em zItrEC>WbR6Z-M?*3Hnzy=wCkK_S^|#eTuqxZ^awXziL7M+6?+vr0~0OqA(5Kr%jk{ zOFDAw;B155L3?9q)YOqY&~2IY{Wani?t@?^pRTkhEKXYe?4z8-dayU3d&y>o^AhOO z-?b);yR4nf*IkmqjG6&_NyuOJd~P$V-nE;llxDKK4PMN<#{~9%jyfw@^@e^LSw}My zH`C(QNV;Foo(`xQMpch0&^6&#$pDjbvgX7pk~v~7Dg9+iZY%U5->shFjU$d@w0=Fd zZVkd-dpJJ3Y#=Us{}Z)rZbbHFrD(J53UsmE8I_nCpv+l4kabO)*rIkwe0*Syxa3EG z*w0xMyUo@VYxaK<yhV_=pKB}=%Wd{6~QtH1<{KGfKukV33q;#+8u@6_eW zc^W;fgm%bM=&Kn{bp2C3YA{2Y&ept3Ky4(-rK`ycTYs|avy8mP{mGnHA8>8jDZF)5 zAwJ$S4sT4Fg2TK=;^x*a*x|`_6!vL1s&3Cf9Y;OTfwpnT!$lQ6`TIij9bF?1Qrsj) zp9vGcT(c1~{tgx+x_1hHo;L}z9d`u%s}3Vt zxuqvPyx<|(*Q1J9YvvL4n}K9Z2qEEjv`LTJuee!u4m-Rr!oSp)<3o#P;7PN_;HxU# zaZlw|l&-lSU7VbQ`s?^2OLI$PcfB7PSNab0PdVtH1)zUMi+<}Si5IUA7Yzy&#X%37 zLI11({WBf(&v}B`M^`DcP^$Z*c#yre{Vo47T3@=h(~n0(U;A6Pm35a_!mNtQA9eR( z_Qj+>*3zegyX#UNx^m4kp8SVR5_g;#$8VJT@mrTB@)Bb`zWi}do@D)uCEC}po3rzo z{e(zH?CqK9Na&SAiYzy#nO@YXptDw{)8@!|^ie-EYWT7bRX*{Iv^t+4v6&kPUJycd zy|p60q6U!@Yk%XXuP@-vqucSOp)2vFnJ#$9G(&tDso(~q`=}5dMs3AwkxAS_RF{E} ziL(}3Vf$J9B-V?rGq;H=S0{=sGpCEa%SVYi2fB)=^`>CqQYk#YpCwpc^%hE%@`!Wf zWXWmuCi%_3FblqUgFG^SEO)o;l#hD6R_+J=FIyX{$cszC65Hz}WvY&lmzBL|@9gIA zlg|_Qq-SuCD;M5L)S1FnJwOYs?ZclLzGm$Ka#ksAVd)p+n10Jt7TBW8cFMXk>Tr_^ z|Mt?S6bf{a=;ku2{&j2s!KK)x4V!TXcS z@OYzD_!V)(->gjWr6s*_s?!s6$om+2_kA6*zqb@ke8`Y?^8j@H`47;48$kap0sS`x z^j~KY`X5B-e-NSn0rcO4p#QD`{Wm~pzA=}3!2PxK*N#%J-^o&Dydr!2`U>-X5G8MH z{~>>7mdNkq>BD_CU#UIYQ@ZH2RBn9Z8XGv#iC-NR&(B&U@{ulqeC~|te08!3@61)> z{(U~Mx|B0)ovYL8r^`{Cp0@YaFpBB&i zNOYf_BG1AKNoGbIDSbAD^f)|{WDQaxDl=NJ*Q7o8nk*A5>3Kp)XO1s!RmE13FVXqb z8r12q33;<{bjiXNDL4*6x*45f(u_;unjJgE$(5<%k(OEFT1R8?+ufdG?8S#d|2b8{ ze3v|7kg!3&E7wI0{ zb^dEyD9@uu*~9ab`E!L>9%c_SEI%ya^)|2i2_JD3~Izr`Na_Pzgfi!6@p-Gpt>Epd$iLJ>wqWz+X-0HoY44prN ztZ^Jee7kiggA`hEU$y->$SwyDQ}@NC#+LZP<$m~Y!8>$sw;c7|Re)kjqLFLPB=q|7 zaD+B0qI0dypnq3@{+$l`_dM~Mu9^7kQy+12{WH+NPk{ct0rc+>;gv=womFqn+)qxD zGzKe43(Q^ktH3n5gGT{NJnJY;G4bL1w?1SVQ>~?A5nf-U**TR>1wZi?zd_kc!LYQ4?Cq*{7@<|;r@;`f))Hc6{`Rw-UQnyFF zxL5QNDdfC8kIYf!^L86c`1!gc1XpMe@{+WLx4&p$~S%cp5Kmr<7tJ8B&;lse7$ zM`B?DMf8;N3~^gHVQ z^*kEuQi5igryxJk3Dq{~A?W`>(EkJa{{hhdSA+iVFGBy12>m}I^#6eVe+u;fLeT$X zg^@77L%l~FS2ei6G@^bn?W7o9Cpt=tc3JTn5Shvb>hbY&jrsMHn$p_tiPEh;U8JT5 zGPW~yEccxr$_?OtV)vWjT*KatOL-#iC?3Qwmn-l~7q2tz{d*a5$YxVl&S!6|EZMG6 z8ti_^d-^m%qV9fMsPC&7y0>^TEjToSN{(IV+8)=*E}PvXeL)6Ud&`{+)gMP*oaswu zYQ4Zi7guAu#T#+R)ljV4*9O=748{|B{6T{UG$E@!JJ9^6E79wOvmToGg_Ybm4K^Ci3v&U-CC_7tGtDpEM@m8`Y#^cu3h&p8N;ytX_-Y zOMANWX;Fl`&K%06?n-=A%Pm&j^8ovxzJ^7pEMNs6MD~1`7TaX=nNIhur}Hzm(a{}= z^xLuN^mXSbI`d~&TH<_*Ts>AvG{vsLeR5Iuid@ zvd3rghT+&v3fPQYMb3)lC}_hf6#Bvqt+`-|+V1p5j}@PS|8WfbkM-bx1c|K?Tr8SA zQ0!mv6a0@x@IOky|5zdDZZP62ZCA)gG{T!qaK|Pn{R2}kPnLFAC-5F|aOUH~Z0C)6 z{MWLxa-+6nDJyG)baqArEy*3sC+vXt!Oz3pv|Ukr_;F8O>cDwS=Wu>)emCxM`z|vM zILzX1=CZ&wi&<9O~I$mjKsI}m9V{E3-Wc_gJLIU zqMyVQ#Y{FwN_$n2e$Gn~<=2RJJvNChv%|%KQ*6c6UPHw3rJce--%G-cBRhqF>Qtd* zeF9&8)LYtj(Sqmf_2wB(%jwlkOC_6fc%yQppH$Stg88pq!prvL$V4*34IP zFzbBLLGz58=*a6ksF_+S^?mM2BQ_b)rBiy+6HgzI5yeM{>8V`uby^_Vtw2b?a&3}k z{1qFYJ&OnA7U5H?lX2Nc2Q2?G8e5#~hTSu6qh$y7p;h0rQS$ovXnunQg8nxYto083 zA369x1>pZgi_rfjLjRix{cj@lzk&Z#4*pNNaOBuU_P#1YY8tD`rl%h)J6fXBxm7_O3FmZUP0{XE-+X?(v<69(;}n?`2Rd<{?f%+#F5Y*D4BB43h4={3aG zc@w#B7fvd%Em3nFLU!eL;>=l>u*>e9`1iq7JhN>UJ~_)6w?FKO)2=^68$GK~ky{>G zW3?EKKSxmfP#yHg@EiDF=fM9e2LEe0_+K-`xg}%7$@bmFq(!aZf9(hVD@XVWbzQ9e z9u||mSfUj_JFCTtWzS;G#Tn+AFp7hVACoPvq{#|BCKYT>F>g8ul zrZJ1p^hw}n=Opm9=!N{l$f+>r!iab5>dPbR-Y~}#605mg$o6bZU{+J6GbN?bEXbrA z%e;GwcCp$=kL}H-+YIMZ)p!f~)21Jl%iof$jyiH_Kml<+9Ysv;*pp{;IGHy>kqASZ z@yDa(xW|}utbc7Tu3j}3JB9bbTYfx6WmU(~UX=|Ei=W*#|*^QGd;xM#rK3uMTdkbBi9OEsokVK!vne5?HK7}n!9v%N3t}( zPR5K9YuJEcu5!Vmi?m=O%qshgSnL%aDWkWy6m&0A{t#bb1Kb?>R(R91mq#KWzbBAS zJ35`;duGDR->LCedp|HRvrF>3G|2xTZJ4lzuCorH z6B+~>`L;hjq5O$dgq$WWPYOvrjVJq~rjiqSy2LL|iM)K#g0s);!J37c`0sg7eEhgM z-rrRX9~}G&9q)A#9k1Dhw!IHW-Ro^p)2SgSoBsv>?-KZbyTJcT1OIQf2>qWT^nZ%b z{|WwI75IO7LZ{P1rX>Y&^IIdN(C^7o=K(*-H{~z*zZa=>vmLoHN>6HzUdCSIV?OEWFJ{qQ;Kcc{@W;iAT28RNgm8;Bwou(iPNAJWZxlY;@zQ7!i;+m+n6?7 z-*yn2rLDoc4+r3?0U};vpo#A;_=H-kPow6+ThZPv@u>CGR5aL67xmergi^=f0RM9@ z_@9~Je|m}if0~Qwf|^J|UxEL568z6if~}%1FCBeNzG=M%pYcLTO8k*5ACMmmvxy(b zRjeMd^W%M_79$tVmnw37cS~t~M5JV#eqFxuS~}}uPk2Xr1mCd)W=if3;Tttw`7JvO zey~iNN4o$1zc@P6xEQ~O4_DHjR-u$sR7#STxzCyuLMqynv=ODGXx~JI79=TaB#~$l zLd_gI{SYBU2xX`2WP8s4dEcA*Of&bKb6wY$*!Nx|S1(qPyHnPZqi_95V6QcCyFP?? zN4~%c?Wgd*8zopbHx|d%IN=*TBXRz--{{-btLPc5MC;~dBTIL0WYcSbl5&*LN5>~H z=v5;`1Djyx!3cPC27;~4!(d$8yi@0~93+p8!tNd*u6ZM@EwLi;Qh5SltJhQkZa=Im%&D%?& znda>J6dy~oqJ3yv zwF7N`ID#5E{3WL2Zj#^W)x>wzdXltx0r}lEiCnPNAZOj)V41HiIN7fZ-`W|E6TZ7( zP&f*oDv`n}jyKRE*#k&Ynu`v-^FxDgSfR|vDyUrXIsEnI1ayrlhTglQVLjLHO=s&u z*@horVDT03Dq}BrWV{;GNIXH*U~|CvXTbSqe*as~?|(P&``<+Z=br`6KMS0H=J&s6 z`Tg%sQO4Xvw!&r+n`077?H+t5bGHtm$9h-L<~Kg;;yzspI}RZ+LeP20m{R7*Zd z6@?P}JN9JoPA&4)`UCDBeg@O0+i_}RGLD_%iuc$VVT;o;*x=YLe%4ZrzI3cZWAo;t zdnYHN90xTt>(&ccDt#LI^pwDHTVi4DVJDc=Jra(%`x_{;tKjQ|N-)|z8x&)2u;s1= zNK94&{^OqrN%tEC|M#1O!^IIow}1q@+F^oH!WR+!cR{p=?G}9kn034;leFuo(Alm5 zw9kOII^6z^IP7StGaR47L8%v|FM*V-2 zLyoPad;^>o`Xd@7GN8};WT^e&yM(-{B^KWc$>~Rni2PN7s1MR1eOjOJZvGDcEn*ke z{4X6Z+&Bx%JT}4GAIstiW_M9|$zf!sosagJ2BU)&wy0o}1{&D@22LCzhGFx{;J@|p z@bDuS*b_GjmZnQVJ)Ijs`e_wdn34<9TKs_cfECywQ32mNp7Hzd6a43KRti&i za*<@u^r2gn+a+^$+!POxt#NH!~R?8sZ8k4>p3MH_1TD=v!dO=4xc1zfp? zwD0RsQ_FC=hx2fxJdwKE1ySkOj#OA}L`_c1(M`sWNQcKUBC~Wexi==7L^n~Ac~h6z zb^gR>*R*20)C#=r*J?b$bS^$+H4ev`D&pztk5JgyqiEBnjc9x5Vl*=Wplwb%$jkB* zTw;12rcK!eYvR-4K{^ZWDK&wID`X+Q&;fL(90qT1<^lK9K|pneE#Uk=m|*&b-~YGp z`~Nb2{~s@K{$Jqyzrgu_f%E_T{=Z73Tr-#N@&?fC8!60orYv=+bD}BdFNtp}E?`a@ zQt0NDeyl=xircG~)2%O}CCyc7Y__!>Q_wxvBE4$}ZCn#d)zx{Y>17g4a}1#)ji=Bb zb#{C-T%@pR8*BgG~-x!_9;C<7T-WtgyldtKPH3 zoimlOXLdJw^5Gad-&TaQrIw=aBeuc8Omtfw>a(L$WDp@L_!#_({liWH)w%tH{J```{OS1d1nsGm%i%u$yt z?o49!mnl0idW-nDxg2%gyMXFO@s7WC3hjkq)XU3-K0G;=4oOj>JsY2sl7>^{;=`?^ zDRLzlRpCMuca0*K+N8+h|8C%k1qX1C##($&G7rzsoPgobVU3UmOHk>t>Y`;q59RZa_<~Czr`G$y{rHi#B~9i{S6?_bR#f284d#9 z3Bc_3P!N&yk@FAdIRCJd^ABml`U5kCLHZ`bLn~s}W9yDhk}Pt}py~rI;=lV>usb!btS9ulSZmNPayQ4H zcJlX(i(^x1i0L9)ls1jtxMoIGd{k)r*%#zQPz(8?xsAM0P9itAPAA`Gj7Tk$A?x4Y z!ll}Wu$=ySyjv*{A27AS6dCF=PzN%ayFl2|WG3e@l7H#-kw1G9YW?-PvQ-LSDA*9UY?{R~=rc{} zqtX3*859$K3pN*2!^pks;BMRbFk!_+81ioj+?4nN_?|okNOcK#WE%_a204M^ijjcE z{T4PYyDI4R?Gqm9W((f0y@aEiErfGUN`h6#V^L&Lqo`4TDt*6gh@|UA25t5UX0d5; ztTa2BeLTB|>`=%iUY|aS2S3$dm-u$?rzw%-_pbo<`f4bvwp=S&+;xF0do_z*w@;z- zPp8v&>mq4{#EquCv7*O1HRyz?{p3mT1ya{oPA0@>5!*OV;yuZn94JyC7Cv1#@JIu; zAscbi`*6I(5a4EM9qe=EBXa0GhnfcOLQz`jXxzt{=#-xcvRW>SlvZ@W%hHFT_s%@H zb5RhyGuRfs)zE-@+g^i(5iKC=PARzDxDx!eb^#||8h~Z0QebXK8|PoDIRBE%`4?Yd zvYM4}dWwqhPVE`zUrvZTlUC5qi3U{q{yaLOUY4cZ&0vCsA5+vbiJRWc)RY zIG*7R^01R+MoWDQRQPG-Aa)|O0#i+2$pXu-OGf_ACr&FcEuI%oSNmqt9G@8aoj*_P z?X;!ioQKiUx}W6K+-v0itSYjsJC{Vd`ja&w)+9A_2q_cy;x#U(uyRldPP(}q{~qpy zTb7N)_9nkkferVf&g?_IU)P{5wO;6Qm<5`kuY@|w9>aIRjZnLA6D;|$1b%3NFwb=u zynUN{QE3;zOrPB#%Q_PzR=R=FO=G~kyuo0na8DSBIU)p~&ll7uhY0=IcEaANn!?yK zZ$<6{Vo`xj8e3}_jZZtfvB@|5sZFQ1L^fm*YI>eVy^V&lseH@sc~ul^uydi>6E*0v zWD7QDPZAqy+apOH8A~EwJJ9?kE2#M)-mPj|P7m>Y_(iSubc4zWT6pa@IrXiLjLSGk z`i`t4T_BLOT(lt>bJWSTC$I3*>&^I0L@EBMw-Otlb;c^i26$}QKlE$hI$HX7KiYjE z2L+e;qN^?AQMR-STCx8r{Mm6F{#&~l7Mn-GS?&Z5oj(FTANL)UjK2&vX6*qde6v7# zya#apZ3+%9kOx7Q4>_R!NJ5ZPT=co`DaXe>x-b%)caw*0W-ve=^yM zKHfpP>_alc8^~QDgML(_Ov=xOPP0g2M|CFBr=}(};LJ#d9;UJxWI)26x{_pT8>-?S zOC5Nxy3jt3R>t_!bB!V@b4!;7zW7Iml(m!XwTH-$lst0NBA7Iq+mRqAP0oe9!_-@X z*~)Er`Qk+E95D^gYcj-2xzae|{7qDf4kEwewa9X{Ke{p08a?P$MXE=8pzO?(@b&5~ zaFc%wJoMBNmM_zTU$X|l-p8%r_TdUJZrB=ddFouSWb!!h%|;Ozjd~!uYDmyuy)aiIelX`@Zf8~0(UM`v(;<_`SmXbKM2n1L%F8sWcm z5bnsojn1STLK>m#(a)uU==L%jWc8mKIw0XZ@R`%_cKTMhAwLc(tegVPv-F`{=O1un z<29gjzY>(+%?5qLd_aD#CE)%e;Qk}$zm9SKtBCVokplN01@1oz+<)Z!*G16=e%9f= z*p@vq&5=CU9!wn$E}$#l1<+4HuGCX zM)xb8?l6`PuZ-d!wY*n%i+=}J37{%Th<=}IK&=%9(d`{~$pw`uwiz(LvP z!h`-3LeJ7-;c8j5!1MFMsu{Y1LCp`*q%eLDwq+3Qn|>aJ^qwIHk{E(r{62qhDz!Hp zMrSSOUGy&rH1$#{Q!2YbqCd=~nQr$aUwJd*`>c^nddz-2bYK*X(2Arx`QK4aK7lU0 zI-fgYgf3P$q?0$v(o3@mK8n?jjykQjY(tT!nK( z+_ANZDeh^M!v{xpBCYSWDC=JV+DOAtr=bJ7zgY`Cxc45K9b)j{itTXy`y_b2Y&s-W zqhawB8F=wyJJ@Yi4c2@9Jub$S_T2neRI zw^8zZ4L=k4J%(97GRGf_N6?!Ai|N>EzSY6q>HVt~P?fQc^g`=sTCF2TtyQ|nrfFXeY2FjJwZr@lrGu3?gzfH;tH-S--}~stj3!Z=iu1s=J-yu0#@sNh@f2q618nW z{^sH6VX}b6Ck{m`d-|d6sB>`hvYjy3APpK>%!I8cjiFxOAb6|d4tO@N2K?5}1B>$) zfN$QDfOWPy82{##Fr%Sa80%On9K5hXa9i#y#LO}fwi*5tRh3>BrN?Bj^9g)!z>9Y{ zH~X`iG)QyXL+RV_8mu;H858}^py?3{SyB%sU+zUx{S>~BJtBjJOr5~89%#Z;IW2lV zeIXt0!`tI9iC(e@riFYnYVay!x_E~?y;b>`=oB9(A6FKWfA^P>f)SHRO`jh5T>T4A z{`DXBJhTsM^sm7s4PLlD$^xG=Rl+NeJVqfa8d2ZQO=yf(1bTc6qO5>nNcH1q=(yzq zOpDkJv)wb{DzO`!b9D@?+&&mSo_i0B$vFaI@8p9`^FzRh9d>~Gzo5VCE$9Cvod4U# z`M(5#`@aJBe+BOU3f%t{`8tkfpTC*Vg>OXv+nZoxmR;tb_f|AHcCw_RXD+Mk`rJK=*-%C-q1wnV|2z21jH#kf2G7Mk12Rb@t!AI^M@X;evXgO6LzL9+Z zBI@fvZ*d_wIesC)-45WYr8a21@j+;gJtHj9-XY}nCkxRLu0qgkBSHD9j4=G#Em3aJ z9y0OeJ@JcPLvsIeF#GqkzQuj-Rq}giuVlIV1U7ZVVrm@;Si?sse7dq_uSo&YN56sb}uC3eSF zlJAdP$d>D)h~_sb@>208UdayN$#5+$l=jCa$0y(&J5(_sJ?KbZ6PgdUpr#kg(DQ6Z zG|og1IVuf6-;7r15mW)${nZdYn**2ho5N!t72vnyUEuJ;2H-z;Bk2Ac4noENpewBd z63%_({Np*!KknrGW4iFWcc$RuZX!f3krh0{J47un5^0L>40fXSqr@+HF#Emzp+shP zA&#)NWP1e}_UDH?^^Y6HZmAke26d#y4tC-&;`I>{5MdF=T9U6O298!EPwW+_6x zr0G^B{kn2B(wZE@=4p8_jT`zBeES1=u)&v}{gy@p@1)Yq*l>FK%XI21jH6R_sZxWl zugH0kgv{8!olGuFA$ac$(sRg|bn3|RSWyR#x>JMa<>ujN&mf%tXA)kfrh$L7y+&V` zwxCP*N|Ed7mFR`N3%dQz0G-m1Ld%x4K@;yPC|1mc-{kyY`cNx)aXRm#41WeD*_;4V zZfyqEd!xX?Is%%iM}Xf4zH|QbGUq?{aQ-uk^Pe69_umEXzYE-d7mahXmGmvuW}6-) zv*PkEq`4)Y^-DBa^!|9}bgf0ac3=qY-+xB3MB}{V*f!o_?K_2Mc}B3^I=(D7#6|Ku zw3`_Hm`f9M(`ed;G}`@YF;%XbLAN}ypz}wl(Y4uc`5Dj|GB@GqwTR*t0wj?eTRk}lAzS8ZAiT`5xv|v4UKI#M1#wu zQSzgkaEQl22rJe?uWkNN&deIReNlyatvw(y@FX}{v<1wIivd@^ID)|$dVp;o5O#cO z72dW~2!T^9j-WxF zvuO6}@$^WsI#tyCK=$XHBe%_W6GfX_wO{cFxbwijaCy$kUR z7ZDb%(Z=;beMt1`46=&ff$kTkpjWS4QOa>68@Syx+Dndx!J-6Gc`CH`v&}NYyo@K%RseuJkZQ^ z0Tp9M0b@rgpxoFdj4!GZ)IhHApO>F7-pfjGPE`?NBA$tE{87iP0i5IPRijo?A@q3I z5LQ%}$#k{nvtDf_>hL)h`-*KjuVKl~osXs7L(0j6%sFhIbOakAyGyd|N+Zz8E&z!U;-7 zqv=ogppyMr=*?#j6d!7a>>A`zRl@^#9M(fS*$wd8vxU$)P6T5rwBgLcK9Hky7UXT* z0e+lF0rM1R0CK|!ys#MrQdis-Vip||{<^OhUb_VfBYkWHnGx02fp39yF{YzQbo+)&S9a07SMgadx>&5SloRx8;iVg|%$XI?4yn7kFc6JV}em9ApiqxT7Y=00Ht*a!^X+Mb?oJ+nv_v2*E z1m4L~B_@-5Fx5PXhc4cN6C7i3++9a}DpC(CW(}a@-K}V#t^y_Mu0gM7%|(}|jzg;* z6j8(2N3h)fC_J`vBQ%R&3{PbP_{~iR`jJoIirsneZ}u)QdR;n5pFIl{?lS??56gna z+Z~*LKFs;&JkCD{3EIB4LdaYVK~MBX^h>Q1TZ|83M{cT7x0|0O!{rU>p6?$e^S1g) zc16X~7f)VG?wov0hQLs^$j*~GwGE<{S8Umnukq~9xDJWirW`WA0#bTvB^`+8dtDtd zG;X#R-IrlUJ0A|CG-o_SHoykf}Uje#w0x6j^FD9>#^U~+B9?RGWNIB ziEirFqCr{a?7^sH=HdKOGF-frTrhH=%A;4%1&x_BQaYBl-|?oSRUN42%n`I(;J)PC z8)WL$gT(LsIx-h*tk4K&6?94c8H}?$0UNGwhCSP&U|ThT^7}`?<$JyZkAlnK<@G&4 zD9HjB?t1|9S!RIuPXO|6VBY{)xc*Cj#%Eh<2ZgW4FI_OFRO`u)-Ia z>_-*<4B*&`x4}V@&%5lq-XPYKYE9QxC9>q1HdM{Pn40a>XB9_NS)9yI$$0%)M3QYo zOP9scj>mkn%P5XM_VT5AWg@!pq%Mu;J0Kro+sWe%hscBAJYp&xOl0Nl$kef#B>c`> zyv|yJ4ScuZ3Wr485!;0w-I>zaHEGFO!FK7#Z9eXeqjZ;{%19KsXrGa8IA){!xcfmmoDLx>QUig z>PErTak21g9uP7tbc9QCpF}>!sgiB7^<-#wGIQYVPmiW$RIW{mp77M9>KEeZyi7l~ z!OWEoO-^U>D-G%WrJ1_v_G$`uRxn{zO#op@zB!UcjkFr(x~i5?C&Yh0ofY;IQ8#;d+@r z;NG*V;O4AK5E7mZl)}8hr7sqstw;$Jdp+U&dn4!Hi#Y!tA$;~n!lNt01gX7WM3~=I z7zahtZHJez1x|7-{mUdu_}TZJ!!9&_cmP#%jArYP2h!q+jJGOyx)vAG)Mpc?#lZqybU$kpDx)4sY9VYXUApev6>DtOS_hQoLWc@ix-mC zZV|C@8A>j7_u~}jbGUBrPW(AH4d2e2i9aiw;AwWU7!2z`QKxFq)6hKBqZou1wM;_4 z@2aC0F|XkS{T66hS_=Potc1}|o#CDX25?K`KcG9L4Y;UOfiG=2VEjQ}Ftl|%n4qiz zt{!_T*giWh{8zMDm?}gGZvu#b!bb=K{4QE@EtSe{85CU?Mnu$&-+Mb44z4kyT6yb<69An-V9}-uJ`eUU&gemc`4n>cM9H3 ziKho@1Lz@ROdGx!@MnQR^!25C#PDt%$vL%w#Of_3k=FsqxuZkc!$0HT*%$D;sk?EJ zRVMDLb;B1gj=@W}4aT0{_mJMYBk1hCeDo|d1U2okL(>8^QGd@{IAXX2R!`Ul{huem z@1Lhav$=-Q(?%LD5N?9gR}X-74r_tu`*|R2*95@(KY;gtIRD?o`Ts4P|6eBX{*S=> zKLYRnh&~0JBd;n>6XhAB$rS0&WT0-f_&eWSe2|b%NgnSyJ;M z6{7CuFr=8nIgsaC=-0dtaFEs+Xm({gJQ<%1v*cW1{om2>Zk`O>yYLpU%C83cd)9$? zoB5z2b|U!nX9$4tFNA~>r-YLSN`!EmSm8mSlQ6Jnq)-#{TV&d*$r>8(O5XeuQQgjY zRQUWsa@}tQT_#SavK=SLOTLHkxI2Tcd=bZTJGPO#O2Jg}^q3?gFO#{P;%8t}t#A_G z0$m=%zxUK+(jcQmdTHeXI%R|-trw4`-<0L(%byR)FW*LTYF7~%=vd1AQ9^n@jUc<@ zeqg!KE7-eeFYY3%@w0E9c#&X^>q8W9M*BmQIHm!uKDz-u)(%JWBLws~Vkqaq`r$j( zbFeCSCmgDf231DPgrY;naQ^c_uy^Afux~~UXi&%lWr+)b)YM79D^VT9b-d#KLo@du zO1b~ALTF#$EX*SYLXYA<(co*DY^GEkEAh>s=G=7~Q|roV`)1I|dP`|`j28Kx8OI#z z`1a2C#q7iuUt-y?gm#seO8%s0Fc*Gzbkb%L+L*3Q#m0;1(@4HWDwRb4To0lJ8z$3n zqm1c77kPR_<1zU&>^SjND<&;A%Lv}=NIquik>DQ#_~O$4aEARpyl3+o{6yCaPy9R% zw_jAmMQV?c$(N%j)P56s7`+4?af8UbUk9l>e})|j7vPwmyWooL8L-5{4W9KL1E0ac z(7gFB7_M{#xWwjzmX=^JTh|U8@7Dl_?A{7Zd1B#!xJ<~gNDx|srV1w?juO;vNeNba zrJ3KYFcxG!pIR*sp|u4`?6QY4z4>`6ReQceJY-xjYoe)i;Nv_tdh%ynr?8wB?pYy` zyO+kkueN2wVyj!e8>-RSER^0H&bNwgC(%+v&cR4cp}R6n=;?F?I@PS3gxotp_O04N zOgAnkrWsD8S6!bdWcsSUY z-jPh#M}^W4(au!5YYhE*Op!L?XJpLQlVsb+62hY`)oClBj19s=q0^&r}50~j=I5qNt}1QMTX1Kn?Z!p{X~g{`-D2$d$O!pV>s zLcET#Fge6MgVK9lQ@^9`->_)n_%N^1`? z3qFaYO17Z(q!?5&;E0sg=%MlZ24K<8R#<+n0$wy*18u|RLO;K8aHWeP{EQxf_~}Q1 zbj?OEwqP+3N&)DcuLI`zeB%Dgd7*H@F7Cgi3-zJ1ggcEU!lF~M!i2O@)P7JBZQVG7 zZJqc@Vl;IytN-{=!WI|cwbLxwWn&qp_uif6EgHqDCL2o@-c9GfXRO%DeI;`$na(yD zz7y+9{ULxm^U?ea@Ebn^bnjnC5t&A}Y%-%7hAQ;?!WZOBLGr znoL?NL(t}1_{o3OxbW#ZTzhCfc7HGt&+t*h_vBxqTUMvh6MhDGy*(CvdhCR}74?zB z@INs1%T;I)Pzk^AGeB{?H+(wC5)SY)K>H<6z-#$qphT+(%xI4Q>y{z_?hXTaC%y=9 z(nTTr^={!QKLb1%Hd~nFIabJ$lM@$*x2ZeeEdErjb>W9Xt8fyvqe3Giy0>ubIu9$LUB)ia(O|(2vR(q|@BmRH|+k zP7m#wPEWoxr|;&f($<7mq_Ry+Qk}Mwxqp($FN+z3%`_&rTL+P*9d~eVObsS>dAO`? z0oL3y2`@UW&d-8)C;U${5@(d6)^RJ*l=IGL&qf2}SojZy^<9TMKktXe%{g%TI$s!3 zHy&nvQ-&$UPr-uA$ANj;W)Q6t1@2M;cDj!M=XAac^Yt%t|A(Iex@B?y$3x)#H-Y!x z1gZEWc1+%z6x+G64ZMBVnRk+R2cnqb2V032(5It!Kai|pM+szKIGSI|d*1Scg*TiV&=Q4DD zO#<3-fH$It8lth%(#Y-24X9gi0D8)aZxvSj+bhK8uNF!d%@IE4n+vD+CDDbMMj3?F^+X*^XFO66(#Lj_h`hMu(sbvhQt& zWhT{dmD@Tv|40B#mYWC*4-A3QZ+gMBsi(jkw-Rt{&vLNh)npLlG!od{`^EjI|G59O zkNZz+xc}rOc$ixV_ueZCQ4W~&1g5dT%_m8-k~}k8yM)G%OlH%R0+{4TCJXSh;|2!Wo1Z}CGMT5G> z)2*-6>C?& jJ4nXzd%8J@d}guHVn_rIBvkRo|kXKIMcD1E2sgaY z#wt?%C@JDB@_n=eowiFwcFSj=Oe147Az=_Q=)Mhg#D`$r*7fiO3xv`~ZJ?%cT z3bYJv1}~0p1^0X7z}c25;7qeVh;;ZXNFH7jLWk}bq!n_6wTpd(x=WVA7oaQ*OyT&u)@B(T3V32bu6QS$q7Bq=#|R&4QjD0?|SiGJLkK;|Cu zXQw56J25ykJ({lzc*lqVvdh&MtIpLpsXw znuY#5Xo4E+Ws$bu5)|?tS`2J6oWGGYg{}~D{=uw zcB26A9|GP#C(X=~kwo$h(&Hf?$*!_U`nYx^E8=(D zdm2NS0+gYF^cddxuNnnqnXzYsR?s?~ZAAV3T&7^Qgw5&MCJCN=f=pI%rB{X~(*c}8 zweByaYv#C9FMe+Dz(tGti$9SO@t1fLeJ_cBvYJSrol7*dEQpzo5^?+QG2VKx5vSkV zgg<0OU}lT(liFc;RN5DGS?(gzxV#(HzRW}&qS;7$!dN7FI~aZJxCb>q9Dyq41#qAz z1ZpbU!%+>IyovA*?9b*t{Mv0`)bd2IV9_+Nrq&RANR&o5gP0MKROK)so55hltl17kW%HiAL|{`%y2VXzVTzntgF1ZH&{VwU@t; zU*oTk4-YCxmvS~4T;omT_gaz!KV{xb?Zz=~$MG$4!xVj6CYR*$Fr0)J6!!ls>_FLS4tLFaOI_|&C7q0D|D7YG_31n^{t5wLPi#OU3_K>q6 zuXdAfyvdxP;m@We4QFPzr_xogTqTzjGs%s~t}NMUA>G=2i)egxW@#^CSY6yn$-T;5 zL_Nrf!oLaB^>YS&^*)*|=<%d!s*|YS?xFPL&2Pl&MJri0y^?5%a>zaYJ>sdA6*;k8 zh2(WT!%m-0;3FBuxS%x}n;oQBK2aBs|M3H9S6)H?)$c_&XRSu@(sNKXG)I$z70|h> z58;qe4RA^G1~^1192x}+&?0mwoO7-p6b?EE?EQ9vMFXkem+DL~b+0iP+cgL{W#19# zcW}3VQg-tBfWg)wWRP#2yJ+OU;N|jF>>h0cqYoZAyIG$ zrrFzTB)6_Nkt}}a<=*E{=W@4i>3PJqiB_6g^yk!D;-K9!4_x50w(v#Tn))wsZJ_dJ& zOvbto_3$M7U+7faf2eBAKIE3X23_o(i$0$nhc0bbL^YQl!GflvuuOInEE~B59#V#| z{9oZ+vb}EmdG7YaJBF*WN;&pL!Xv+u5h|2fm?2s^K5aK~ouFBH! z+>?DUE`jm5wPam+Hu)3?>9V~m>G%}BlkqZ!jtuvrO?|S+mgpQdr(gWgKbl!5($Y zr{5R2N_JnnfSycBry*%NOxjGH{#J`-st25@^+gT3Pt}5{wr+C3X$6rxTNzQ4lj zx|^|Iaw&djyAs=VIAgfa01HR|A!+qC6sT5(TyEu{uv%Y~+cqAxsi~l}mZxyttK;y? z*3GbHauhrrN}zb@2sqODJ9sTO_pjXbY+eanzW<*Iznp#xyqfOQhbWu&E1D>6)24aEx{mX?>PSe`MP*@6D!k zsBAn_z2!g;*&9&0Wi$i)XIFyH5+#18)30nx6{X^60ly>fyAVtB?)lK?IwE??QI}qS z`IpF_xkAA%*up zzJWy92hiESx#;NTd8plM0-7^J6)kLf4&}0%pmBFGwD(#DorXBVDHn8MaLi8-EZqun z2lj%LwX1<`&>V0r*Bor!sQ@lnb_x2?4FVn5AUMZ_3kMnm;r_m%LZELPX%uYO-!U7- zdky$^{1eIS(b#layu6P1p6?}9nJKicX9k;b7SYho6jt!XoVsd{qwOb$vGk@i_V2J1 z`|<&iEiTsde#vs`{63TR9OVBqN`5qNu0S2a^{8RI6u%z1MMA@Ch`nY$d14eo)_t-g zhP|3(q~d$rW6p4p;dbnvmW0<%pN@T(j>h`KWbm{j?MUjklrD0eEgCWB=f>Cvy;dvwWCWn&C8*u5R{EF2h;HONx|b~X5g?J zMp)l|5PqF{8@-D@giiXdNA3ZEXkmyA3Ser;{>V!xbM!RKjN1xbljGo=MN^=Aygqz+ z@egoabq&;BtpvNTWrLgYK44?AC20Gt1mvZ=1?4rzgtmktVNz_Qu;?8UG{|s46q3r? zUJxp=jHP#VLYecu?~>^6mi+8$Iz3kAN~L1`=n1xrWqSqD>p%E+#8_o&{fg4A&oXR} z3Ex6qJcM0}R3sBunbGeXqUjm#i8uwUq(?*j>5GeyMjz3qdGRvz^TQ6Z7S$5%C57bH z=Y^z0MIb#(hLRjzzLmA&9QMAv6W5vfqg{Z*|1OD0{M3Xk$BZCjrk>i{BM$E%-a$_4Hb%%6_q48&Ybl3&ls=FKe zzs$g~@oqS|bPU#AIT%}xxQEJpkD$8xeDuyZ1bvIPLt(a>NN(R-_~)4zK9b!A19vCD z(3Yt%V5}i@7$OZP{<{J0?mPg_XsiX%_4B}-s0lzWP!+g#JQoJanuG>y+%}@0;U_^jNY^4OWdaq^T~Ew49%zw)Dr-{6q6;uN9$d`}q0F3RybY zv6GC}Xdp)7jpXh5CFIChM3%lAMvho~#jxcf2FuFvs+?6=eS$mA`ZX30FP6jCUfxH^ z*J{zYM+FEdgdxW-_UP3zE!4N`JxonyPm82a+4z31s zY}bM2>;Q1A*BbEtGvNJa?*E?R{%;BQfB9bCg~Ba=ORGn%QWFT+QXQy9axeX5quHJn;G;Gd%RDJYIU@0rH$&kM=8WK_r3LVZX)Q|Bd;i`dKKQeyv6_ZdxWAsb|KDt`5d~42IL$r^ER@NhURnNThL37tkde z9O+6+BdWVhj+X81B4#~}BtT;`8B-iZG#n{87NbkHqMz7ya4S|FT7eC!SL3M_bMVIV z=2)X$0W0QoA@dUrXq5d%)NwN$-I4;R@A*)~ia)}H!{^|V?wxR;BMlBnX2SOqOyEd& zS-25*fX#PnKzwo@@Eacl67Np}{_oX6P4R2)KeurIxs?0QD+T(`S=hr2gruam63?q? zO#4w1OTV2-%OCO0Pj^q&mFi5lH7uq_TR!3$P0Ly8GS2rWgtPRVS>)#42zpk3i{$y- z4CZNQ#aQlJsMe)L7jc$iXlN!4)J>x4-}pAbX}+5}-I%6D%2Q*{$3)BRIN|kT5}mS) zY`x}4oEr4VC+A;yU-f@DX30MM=JXo;D8LJEF}1)d{fhX&+{Z}9qLF*Bn~>4*B}i{A zL@N!3p#vK}!!OPk;8t_a!M@6Xt0UcD+`2JPvT!hz{c{(LrAL7L&U{cM9Rluq+JOXp zO>iy!t#I#xSV(_YCdecv2)lPq6~rorf?sb2+xt#~Ey;>x{cfSOhHv~iCnd1mSJdfJ zGJ|fYSRuA(3t`F^_!g*M00Wn#NO^7yy)Z6GGXHlP`|)EEQ+{01vTnT^UH&1I?i2aZjk-~1`7xxuu?Q{S7Ky@CG5YQ|9Bo+t6~;cg2rtH!!`$>$ zP)Wxf{^%YHgHz;SQp$fxhiGrjI%IzG`bFG5l>r@10-H&$nhXP42fpdFaNf9X#>Y z{FQY2nrRY)$*GJVDKRUjN#Zr8%5+PAFx}FgL2rg9(}P-}bg!~A{hmIC8iXj)k4K)8 zM*WjSe?tlJoE}G}^-dvHwgx28`yY<@bRC~I!=Bk9^ebxxmf>A@L6HDZVGF)2luf6rC;R%Ik?_;Qd08v0RDu^Nf2(8cbP z@<@cuQdTPQ$JLcDN&cEP$;?c9mX#PuTW=HQdgOe|znUa^J5)z<$0(U?D+q!_vkeI4P5{$AHbL76{c z)R0y;lh{8bHhc1t$U6{sWIXk$z@6{k=Il$*mPS zE3MuPyL3xA6Z9iFN&a|0Z>D%ZSll%~Q1R}*A7dB3oeey{{or4| zl_`8a-w+SK<@_4IZ4g~Q0VOWK6)eiV1LS)@@y*G-SfX~nx@bhd1i!((;x+<5u*FrzHPp{*xB+*Aucst4shvh2D(HsF^&P8D`Oe*H#1+O{A*-m(Ec z3_H<1&9$gK0Ro6UB0^g|o3}7M^J@-07mnpTUH-T{f)tfJRz`3ZZy&P)?&eR+)o4q8L;=is~vmM7s$))%x8!cHy!*P9~5#;8JKhUN<{B<=qxKPbTv@ XweZ6^oaLQ4;Cg;J8WK%8FF_?av-TAc literal 0 HcmV?d00001 diff --git a/tests/fixtures/onnx_genai_workflows/tts/code_predictor_indices/model.onnx b/tests/fixtures/onnx_genai_workflows/tts/code_predictor_indices/model.onnx new file mode 100644 index 0000000000000000000000000000000000000000..22966b7456015175624d153fc5249fdb8b5cde0e GIT binary patch literal 942 zcmcJN!A=4(5QbSqSOy7LmEazXA)Y*JmerFt@g$F+G0oBzsIu**Ed-7{fF67TAJtuD z2{hrjx5?k>}Yi(;GDZ@d}&4b%E z)ECig&M+6I>-;$7VH4^iikVI#!cb{=16Rq%#Fs)WHBmeUM(jE|Mv3SHNws6(;~P86e@ z7^yN2PMCf<+qo`Og^F84rO_qW#M`HH*g3jwdI7tZ^Z31E-Row)m4|u%paP|KcW??# PB4a=1L`f~jrtSO!P1z~A literal 0 HcmV?d00001 diff --git a/tests/fixtures/onnx_genai_workflows/tts/code_predictor_indices/model.onnx.data b/tests/fixtures/onnx_genai_workflows/tts/code_predictor_indices/model.onnx.data new file mode 100644 index 000000000..e69de29bb diff --git a/tests/fixtures/onnx_genai_workflows/tts/code_predictor_prefill/model.onnx b/tests/fixtures/onnx_genai_workflows/tts/code_predictor_prefill/model.onnx new file mode 100644 index 0000000000000000000000000000000000000000..3d18439dc78a95a69f11baa66cf48ccf19548f70 GIT binary patch literal 366 zcmZ{f!Ab)$6h)iR*721<8&JzEM8t)-7{;y3x|>f>BxK&CozR&iBpIlG;o&z#Uz zi!L6#3x|6-*@Dk|Tk9d@v$A+D%9kGs-vn5c-ClXo>ZVZ!*v{L*iBhP(Rt;y`xFLpP z3OReV24xh+HcFKD^J9Ekz*HV|Sd%Y8D+ZTF_9{4uichKsxOMwoVU0Ni^xDM&I~l^b zr@ZvoescQ$QQbj)Mx_oB&^i?o!LKoGIJeB*8tVi<1`>#%#IyY08EZ3VAFtmNqYxFkQZCsLH&H85l literal 0 HcmV?d00001 diff --git a/tests/fixtures/onnx_genai_workflows/tts/code_predictor_prefill/model.onnx.data b/tests/fixtures/onnx_genai_workflows/tts/code_predictor_prefill/model.onnx.data new file mode 100644 index 000000000..e69de29bb diff --git a/tests/fixtures/onnx_genai_workflows/tts/code_predictor_step_embedder/model.onnx b/tests/fixtures/onnx_genai_workflows/tts/code_predictor_step_embedder/model.onnx new file mode 100644 index 0000000000000000000000000000000000000000..15ba8c92fd116d0815d56290eb91bb6e6ff32d32 GIT binary patch literal 817 zcmcgq!A`JGGUakbi|T`i?M?CIwxEtY>wfH zT*Z3AauWLPDidakgeT6j>@}ep&$;SaA9iKnLPNwUQ&~(Iwg#tgnB9*9Da2gSoM$?q z53y2Y!dV{Ybn=WVH@AZ$@PBj_97Iu>F;#THn++u;P}B0BiITWzQ(W|~;@->RA^<9d z(j-bqh+QtSnN~}SmG4~(^*Cefg^`7c7blJN&mUQb4~iv|EN2O)TINJ)maSIK@}7@u zuih_79w6IVzTw&L1Hl+PWaAc8(3T0u1R8d~_t6Fy)o}yrxHgJ4oecMXACv=C(_@X! zVbdGD2fK#3vmQ3oc0oVwZa4~n%zU*#rw8Suz}wI=1jrx@S<~YZ8sQ!27}PD(qka`E TuiNi{OJy3QoaR!=JMGzDEL!Ql literal 0 HcmV?d00001 diff --git a/tests/fixtures/onnx_genai_workflows/tts/code_predictor_step_embedder/model.onnx.data b/tests/fixtures/onnx_genai_workflows/tts/code_predictor_step_embedder/model.onnx.data new file mode 100644 index 000000000..e69de29bb diff --git a/tests/fixtures/onnx_genai_workflows/tts/codec/model.onnx b/tests/fixtures/onnx_genai_workflows/tts/codec/model.onnx new file mode 100644 index 0000000000000000000000000000000000000000..9647d0279b493baa6c31b0ec20cdd3878842c13f GIT binary patch literal 146 zcmdC=@C{>H({$LU09&EXP@hDX~-A6#4{J z--2i9GxQ1i1XZP~Qq_z45Z3DuY)1kKaTjl9W@o?I-G7G}Kz((iqMcf7Z$T&ug_U2E z@I4I^xL(24s-_QK5Z&<5ZXz~;FZ;SC8#jasfF0eBsKt61eZ<^>v z70U&#&;OL(LPl5W*lH*$MlXo-N^a*2>!rQD-B*|SZ99%9EkRdu zc{&6=4Q$reQ|nz9seVbKnXi$s(agist>oTmI#x{Sbu9I%DC&L`2l-|sawi5!lOPXg zVWV1=OYAHO)wT$(kyn8d*X{hizxY#-3>bPubKTXB>V{U8=h?A<{oZJz6I zeoF#c9g6z?+$o0 z7PxNb_gy9tr$u5yEfNJHC-P#n7KyD|w0NspB;GkK5*;o2lZLB%C{B4okmHm`6q<^% z6srcv8&saVN&ZKy@&x3fl}8jja!?k41KX~>cLP*9zx zxk@z$&DJ5~kGs9YAZ22)!DB8@jk`@qs7qKTUBVVEHLGOe>!hYN>MgYiFC!%>M@y2^ zCbX=SFQW`u$bNz9v!;M!`q4P_quz>Qt2KCtM!YR(;g6?SSFwD8^?ea2ywfo0xJ(2t zqX_x)(LR+Wp!YJGg&8J=xWKF3M382f$})rsh@JrNcuoeMZRc}ck4v#FCd$IF6Ahq; z`FoM^oE8%{P{T0jcrvb!=U0e6&_(}112B$PPm-C4Ja#x!^Xi~yD#f1|^ZC9_8vCK3xkAz0u^oRWg3)KzA zxO>Sx@8;a$n4NBHs<-~}GlX~OmzWVmPUi)&h%AY61}mkwu?qS9%}IMF6Brq>L0<`p ztd}g{s(Q_1Ys}6|rz?tODZJ-}MkkysY8F1fBy*S@QTq<=rLEBK2{zV-Fh|FQQtOLO zg`L4_~e(FQ-eEJ2Sf_QEK literal 0 HcmV?d00001 diff --git a/tests/fixtures/onnx_genai_workflows/tts/policies/code_frame_update.onnx b/tests/fixtures/onnx_genai_workflows/tts/policies/code_frame_update.onnx new file mode 100644 index 0000000000000000000000000000000000000000..76ba875d751abeb777abedb6463134ac1114ac56 GIT binary patch literal 1580 zcmc&!O^*^W7^Z}e;RV8OO;(&lvpXT0?13y_9z3m@y>Q)l!I-9$vXq6+I4xw@jUG7r z12oaZALIYhGJGt>HJ;F(`hNE5)2FYK$CQQ>Ps0Z7K%Vu*@!Ad27-q)WP=$_zJ1p|26rcG z(Ss=F2V5;O;1hzhkUbv39RLhdauEeo_DIbwu68k+nwMFc`-1Zo;lh->_NCHPYEpMg zn|%zR?m276sS7ZrqX{QMYIXi|z0h}*-50!X$eg+11}vz!I3;hSdOzX(hLeVSaA(BO z8(wbX*%!VY`kK4+HUBStpN-gJwOh6)+?}v-e!OXW)U_Yj9wlGK@ zX%gqVN0*Fh&BKckA9JHb62JCyue9So%|nVtLY0j`lY0^XK3t$*0wa^4M>X$w*2VAx zICQeO2w9x&Qq2aCbW)x2dJ2`GXrj4Iv{PD*2b_VFi6puwuyy%+yA2;&vjiMdmomqo z;nDW*!470hBngS3Yw4kmVXNk`&3ywGOf+lYB;uBKY11mL*rl!e(4skd1}DpTT)pDb zOiwL*Sq=P=U^62IsYA*%PTqc)qnA+mYbFsr{AiW6Y*M9@NQ?vxXXmv!dJ4yX?7{ne O&U8GOU!@>%Q2qt$*!zP3 literal 0 HcmV?d00001 diff --git a/tests/fixtures/onnx_genai_workflows/tts/policies/code_history_append.onnx b/tests/fixtures/onnx_genai_workflows/tts/policies/code_history_append.onnx new file mode 100644 index 0000000000000000000000000000000000000000..41f2f9fe6724529e403e852d55a18346e830dbbc GIT binary patch literal 853 zcmcJN%}xR_6on~47_Sha#)#t|5TizuiR020>W;k|jcH0LP{eke76KYKyn-8-`YJw| zodF?`7+37td%vEXdwQKUoPJVu(bJdTX?RCC+b1_lYD1+V=kEISELDS(^NRCZI4uD5 zX}8CTkVaS0C(6`vuqx>j*8yc5Cu)bVA3U}^CFMY~Pz1(fU8*%{3m(#twTC!)|4=J| zyFf)6L^Bm|9TD(>V5J$_BFve?xjwm-`nkvX3n#vtpIu_Jer0u9S(ddg1evzbMukI| zw}Z2i4~(5PZn`;}8>5EA&MqqH@%3~KQgk46(Z-itl1IG3qOY46Av6fMY}}nped9gY zaWZ%dGI*s<4Qqpg({bDyl1}PwcM?9{fP&?6GtPGFKFrOD%WzspgpyvMx`GV}Q__T#8Jv85 t8=)&W{Ffx6fww2rSj3~ff#?bv4y*nM9YN_&4c=^WX5wAFPC+7H`~mk72`2ym literal 0 HcmV?d00001 diff --git a/tests/fixtures/onnx_genai_workflows/tts/policies/codec_layout.onnx b/tests/fixtures/onnx_genai_workflows/tts/policies/codec_layout.onnx new file mode 100644 index 0000000000000000000000000000000000000000..31ce6b3a9b69ef63df9a7b26d8b599edb29c19b0 GIT binary patch literal 428 zcmaiwJx{|h5QdYGHt8Wm<$;ePP(f8e3=ynMfiW`^2w9HnG!}_%*^bl*vE)B6u<=t# z2v7-xSnnRa_uLz9qG2ghUYV=c4R~=O$homn*JwcHf|<$qRuf^AGz^P&fGPHwpj`;iASIUO5TaqGI~VXmKb=GW5BkNm z+cn#`r*Pc&5iOZ$EiIQKQnKK*##D;-HnqrmdT!innEgv)zCin*R38$NYQ+mqbRAzd V?i{9nX0R+{)P59j>m zU9gO*mp+iysemhKpX-P+jtiY#*oq$Rd?n>jvsgsNXCc*^^aPJ-%z7hiUEg%J!TW{E zdJxT2#C3*%cLXy9xH2G@JV+S7phMF3YCl1)s@-sBnPwH72hO8F^mQFAQQ#}WtLciA z)P_n!4!r8@o{KtYINyn84cq|@6HbIQ^}-TGy&+iK6)b4P4pY(ynQ(9Q{4t>+$)z4P zR`O^`e~HUor9w_+`kJ{~pZco`701Ocl<-QI8rJKs!AunNk3AdU?Y-kD5lI{>*;xnPGI+6BoP7L kn4%_4e3}eJC}=!7ZjaF(Z2#GV-`;X&a?79EP^i_v08>02xBvhE literal 0 HcmV?d00001 diff --git a/tests/fixtures/onnx_genai_workflows/tts/policies/last_token_logits.onnx b/tests/fixtures/onnx_genai_workflows/tts/policies/last_token_logits.onnx new file mode 100644 index 0000000000000000000000000000000000000000..f674183330680e7f8e3468568778f80f8327ea0b GIT binary patch literal 631 zcmcJM%T59@6ox4ohv5KGyHLPr5Q#)$CJ=Y7yev$31dVArr9fq9$7zQ^j2m9V#Q1Q& zhYk#=kQi5X)ysGOd_Bj>VSA`N5$W@H3tkY;kI1!>+E8gov%LAWm#V`_v+67_8YRe% zs2_16q$z$Z7kWeKb^)@I4!I5}<2co>!9g(Sxk}2hW}yg-%Y3Re=?fmxko6}xeRtok zKzWUtWf09&!1auPR|K}NdctT}vBbGJ&XDA(T)3g8&qFd7ZdL8%*vNrJ$3hqXQ4nv> zJ&YP~xv;rCTPouds5=NBK?ZMl)UZAPb~JyCxMUpi7)FXwFY;+P@lN6AkMKV1#g)6? zl`Hha&RREYu$&A|O;H^VlhL|8F5^}kp`;rqU$6;bN}7-|gVRsXQ*;TnKS?4wxU`^> g5SK<{;fvV#qBTXwQ28~3ce|XK`4q1#NNp8A0X>A#xc~qF literal 0 HcmV?d00001 diff --git a/tests/fixtures/onnx_genai_workflows/tts/policies/predictor_body_sampler.onnx b/tests/fixtures/onnx_genai_workflows/tts/policies/predictor_body_sampler.onnx new file mode 100644 index 0000000000000000000000000000000000000000..e2685bfe5abf6f9a185f76a54debcdd4d02e72e5 GIT binary patch literal 536 zcmaiw-%7(U6vmga&F<)&1i>;%!WmB$=#;Hv=ESJ1=}s z?Vqg|f|#r1%XiNCozL3Goml!}r6=z#_(IT$WGIZrZ6lXF8T9T|_(0brkpU++Zl`*V z`d}?NPXiHayYY#dqhn~%wb1tW?x0+q^`M*3nCp}>jvKQfoTtmslS#7HOo`NZETUSI zh2V-Rw#abv>3QbB;FqiDK{S&o*M$QghKlnbBbvr(#MR?fO??KZ7Q!d6huc0iYyk%t zUL>RbDxxa$ZK&45edt+xxLf-+=mWt;6@(Y^BRViaWtn^#xKh`STV zE%VtSxYv#w=Mfbj6w&0 jfCoji(Pq5+wnh^e n{2Pym0QU>JqkIJ_yz;NnhSKfL8eM?%=a4E!33$?g+M)de#-y^; literal 0 HcmV?d00001 diff --git a/tests/fixtures/onnx_genai_workflows/tts/policies/predictor_state_initializer.onnx b/tests/fixtures/onnx_genai_workflows/tts/policies/predictor_state_initializer.onnx new file mode 100644 index 0000000000000000000000000000000000000000..5f30409921a77135af5e7155745f39eabe49d79d GIT binary patch literal 14690 zcmd6uTW{P%6oBLHChJTtO(u|}t|;kJ2&5vcJ(o>DLZuZ56T8FqHC1m*+p)DOE}(}yz$cfhUg<43!lAh-ydy>>U+ZJu~x=pMEGp4V$0 zUBc2=Umt7%OS&qS;fBpl*AL?>a&@!?9c@X+lAvR(aZx({YJk*l>7fn5x%;3KG`+}m ztfj$G;!meK*FW7^0A=r@9Tw`tzYM@>0@o6)*wh8lXNUbsM%FshNBAqz6IHKE|L<5} z_S`!;^LLg_W2JIn8Z-UhOQT$re}6GRmdnUp$)X{S!#zG8_pU4+|7450m=|EA5nO9@T9>ZaEXKgT z5(7o9oX-uo#3heivsQ-5($5~d#-+~+TU&?eQ|N;Eg3BHofM&cjx^~O0#i{=An={{k z;=8tOT^}ANYmRMKD}z6ZN^99)j^Zz8oWF{k#y<=oP9qTfBsMzn0`$2TJYp_H)>8kd z??6RcjR9e01{+*A|e zz7E;^vDGc!4%&Ox^5C#Pe|yTZTjZviZK%>21iV^sH6cW|6Sk9mRl5~Rk;VW~Y^42Q z#2SqP9WDyWaN*|ErWJ0EEr*+9j}14oz-^wC29Cl5hv$Kq1}wk|+kUi>G^hGVM#rq~3uPiNe2yWZH)tYBD8RH71=^w4L;SV7YVCDUwtAlEd)cfht;fp4T;H{p8n z>W=kuPyAr(?$+R_@g5LLNZ{+>HP-b;+I1N&4o8c-Ha-Mdb&@dFw!&+>4quVAWy9Ky zw}Dzk1J}SR>%EkYS~88sCm{RfJ$x73WuDpQuQtINTpgO2xF%K>Tr-~l zWCqL}TX8MwuoT86Tz4Z@awM2?En@`NN0-Ec*SWfCe!vZ(*1@j9IZ$TC#~p;K*wOz| z2WQl<9;-L3l!YdkV`55die^aa>~UhH5%;+!c|??F zBg(Upy{JzO5l3XN%F`xg#5@~Oo{j8nW-3-1G0#SvXCuzDk-hL4HzIrGk~YaU;yfF1 zo{j9Ocq&%3kv*4Ao8%Emo{c2WM)u$^6)TR&9^Iu)%7}S3k~|yP)0I@LG-95OG|xtw zXCu1-9XBGo1x%Y{8)=@6G|xu1OHaj$HnM$k+EjjKY=hhVuNZe9`2qW$FfLWz?RHMu z%}cl038J3Yj9}^8@17YSfjj?urQ2@78-uI*LUz5gi}p#|>s{WnpBZm~t+6uv!$ZFr OC7&wZFN2wt&Hn&Kn7?xX literal 0 HcmV?d00001 diff --git a/tests/fixtures/onnx_genai_workflows/tts/policies/predictor_step_update.onnx b/tests/fixtures/onnx_genai_workflows/tts/policies/predictor_step_update.onnx new file mode 100644 index 0000000000000000000000000000000000000000..5c2e73d7a71c4f373a8eaac8f1919cc9f7b06eb7 GIT binary patch literal 2039 zcmcIlL2uJA6pp)g&E8tJc`!5Qdd_4(U52gh zE^G&9rww0-v2;TosD?YFQj$JrAr0OBCC+{L*sX$dD=Vr&r0WMvMj7xG;nD?p;|sl& zP^9Us-HZ*gX3b6l^$mavIvg>=g_4IFV=*gK*9fLeK?)!S1J|XB>^WtPEgBi9d$r|w zT3>}c9dnr}alesfBO=Zq!JAVO-<`xayYFlz)FmNgOOXCkAA1V#&Q3N}_G?X$En z4(5ar`70~HU($|h-*!}c#!>59=82<>>9Yd$9i?2Ma^1RsgFKbJBoDp&OUJx>myUGuVR4gE)z@v&CKcDl70~Nc7 z??Mq5dQ`c6Siwct7fOrs_)+?FcoQ~sZh4xUvr+E=7AYBT!$X~0ndX-9>ZES6TbIp)-g)4pw9!octE9TzbX9J!v#@{MEVjU!6i$TIO8whPZQ91m zkk^2$K^#xkbdvZLdICG~)7m&?LO*hmcpHH~4!jlKz zQ}d_oK@j(}`|UR~-~8J9q?L$(FZKAn0bdrhVi9qpNz;f0rv2`{iXQ07Peq9SjoYf8 z!ydF3h%)4fwrihOeRvEFy5idY-X1Ggr(NizG(jyXLsFX#;9M>uPo(KeGsUIxSWLC{ z=NuJPY@U(&)AQ7U{x4V2!`Dnm)P;i_MhY>^d`%M>qk6pBr9Okxwna{0k2C{n*c=WB zya+~vWlU8T*ifxSdeCj}kdlLpj|>>pLlvJn6P literal 0 HcmV?d00001 diff --git a/tests/fixtures/onnx_genai_workflows/tts/policies/setup_talker_sampler.onnx b/tests/fixtures/onnx_genai_workflows/tts/policies/setup_talker_sampler.onnx new file mode 100644 index 0000000000000000000000000000000000000000..f609a5b8a6a49d4be6dc363d14cf83440ca7953c GIT binary patch literal 534 zcmaiwPfNov7{=GJ&Fl4PcQ`+ z{j=>Lh&d%sp7+UYNsF|_E4E9-l^j6$9DpqtVJRYn<-#%u`ZSrob=O;?IZo@tlGR4H%4 zQBuhkIcYvU&ukd{a+MjLVj@FTI>=!t5rf=QG|6I=k5{|aXK-p7;DEr3 zVANm6ROY?~^|weLdgdPK?tE+G!+?t#gd1X7^wC_TZYE;Ra*v5rOUg8nuZ|QkIuqNe z_-rdU*S3=sh{}t)T&dJ3OKr0&v}$ZOU%GHUN2W_jpnkt?ej|!r- zKj%ylMdul=KRr(!=>KvRJ$y}N%5>qthmm3|%zRDaR50~;wM%^lrxwB|(8kSx7&?ao z3@?(BU@3^o0voEea1Xjx8+UfTHF`&IQGxIxmgIR1mx-53!D;4GnHWW=!Q$0bQZRSy zxFw%$1ozr;;~Y`(QH3j%3T2^fZn;*L?dl5`o6Q(Cr9{MnsZCG!ZH*={_&1Z!L)&u=3&6wah+Gj-E6b!nR&1Zsz6E5a&@$N9NPsO)Nml*>xH;sQdW$;54^O(!8U zt=dIgdPYcy1!<9x%7F_94%`vqKf|8MWNc5eh{I&JQYCwAKYxDjeec;X>Ka@*ANHIJ z@4hE384OHv-Xnbvu3nnYPOK3zYg!>}xWhi}>B~W{0>Bb+`(@=;$yN$r z+4HTjU(V28cO9&{);aM;mQ7%$dl$SrIygeFJ$pDJ9-Hw>j7Y$UB#c=BBSV}P#$OZ=>n&~~2rkV>L)W%^)6mu@tw5iK zz7lpTYanM`IbNzd{7V5G61WwhqNa2GCr6=?v9?3;2%iF+Sl^JoS2VDEX$>w&rKl>U zd{w1oLDytyt56k%!jsJrJ z3$a4YhBnn$h@_#bkqRrD)DBE{rsMkFA$7@oS}MC0Tv7nYd1X`d%I53y%H}MuY|Zh? z)*IoKFJ^jW3#!x`m>;%ze{3(zAAhGoUE~FrB7$4JVgK3`heaRgBzz#v$oWEn>l}Gd zO)EuY3E6{c9DPbOwnp?xHDLb4u?GQQv(RXwzFA{bfA-Y{APu0)D3_!J@HO7lkgrx@c56bH7NU|PfA;+%oSIA`ER z$2q~xxF}*^xEb80MnlgD-XKoj6Jv={Pdk=;uP{W@PzDPE7%rd~jfGu6We02~h|!!& zj5Hlk#??yTYR$z(b7^W%H5pesfvY_iS1v_d|8RLqVrWw&Y)kC6dH!}J{wCiSY2H(4 z*vssUT${l|!mHM*w4?1tkB+Dw5+o5^OTIIAqUr~r)eohcW3`VID<6_? z;6#$a^OMX1xr7r*2G3723*-{UGM6xZ!?=WzKpt#}UCeO_*F`R2nqI;;z(T5sfAd_z zb-0`4Qb4MCY-+r2Ho>xQ4bI5eq;*c8)!=B_wg%2~GOlQ>midr-=y%mgSu<17)vPdv zx}|6=s(pUmNfgo#!27BK_dy;m_blH&0T7m=TDwQ+^T@Rc$f|1ZBe`4yKL7`69=;bZ zEyJBa*L0mZAnx9Wdy`fV10mznMf-H=J5=yeyu1ht(T;3+?_)5l3z9l23 zTch56ki7W7U9e48Z;Cr-Ri*bB%nrKnKB!O)-HUP3*p1=J+ zL8oQzwOWG!s1$aNZlMSxw01uo$)4_VcU5jc;RpB?fvo9UofbSeS^KkV+0m2?$ee1d zd=jN3v%xJF1vj~o?CDIwt?3jdc!w(m@*D07K?l{~3dm8lumwUPXZl@QU`z*g___d8 z!3q_Vr8~bZ?ogk?Vv4`RC27PTcj;7@-|?^nic}_G77PBFZ1E0EPnM^E=z;0Rf$1hM zY++LgQHL93+X~ly2tkR;EKE0Ls6_#WXn$#tcLDLd^cfLZf|6Di%3Cu+w9Pr%*(oM}3JcxN&z6jHS=uAG5mjjs>UjdBm{&6Popj)?6CzcfC~flR6V EFC?YyR{#J2 literal 0 HcmV?d00001 diff --git a/tests/fixtures/onnx_genai_workflows/tts/policies/talker_step_update.onnx b/tests/fixtures/onnx_genai_workflows/tts/policies/talker_step_update.onnx new file mode 100644 index 0000000000000000000000000000000000000000..5812ac850e80197463e5e98a44ed90a6cb4bd74a GIT binary patch literal 2044 zcmcIlO>fgM7>>Jk&E9oi9&9B-Q>i8(`GTKaXuMik6pt*iTi^z8xe5^3ErHN`2H+D*+XY5p)N@ooQKt0jwkt=0|mouZPiBW#g2`v z_`5;*tSJ74;6BWAijBypQoZ8?Mw(h3No=;#)asZdvcma$&O}BJ35+J*6l|f=Xzos44y6dRcjHA}I%o9f$(`NptmtsSLLwCl$SClrVyF9aCUKx_7x5z~ z;zE}yw+BnO==wrwWuES(w}!Xjv98^KoV_ru&)KMZ1ouf9@4!=CU!K;k;N?m0%8N>) zyASj9Jh(`Id*))r0Jd(O*HBR3NeQHp`d3a}b<=S=9AE>dtk@Q2*JA`vMebrsHi{&I~@VE4b3gnRhO kq$+Mwr=u|+avEM8H?Glhs65z%ot-gPnMd?Z9<259Kjcw>BLDyZ literal 0 HcmV?d00001 diff --git a/tests/fixtures/onnx_genai_workflows/tts/policies/token_to_slot.onnx b/tests/fixtures/onnx_genai_workflows/tts/policies/token_to_slot.onnx new file mode 100644 index 0000000000000000000000000000000000000000..eb8e4b6599f5d6095617f9fe0d294eb1d098d216 GIT binary patch literal 438 zcmaiwK}*9h7>3un+4PfQjh9(LrVIrSWxIOP9eekpNXW8mjq8%yq_9rB^zN_lhicoE zp@NuW-iMdxd(i>jU#Y3c_07i~yj##SYRM%4Q&Jl;o{O`5qznA(Y=X}G7R5MOS z7w+yS^1zjJ#dS&<$DL>hXX$bgDk;~RWg;~pOQ_Z)7CfUFiwkT&K1Blveo=J?(M+XW z*97(&8DbQv6|Mth`Of5HNf(^F^;z{hY@1>Pkfwhl0?k#NlR@@LYl8d bB4}3Jj!Wwj27h+&`J6M;Ofu<0$M1gu>uQY1 literal 0 HcmV?d00001 diff --git a/tests/fixtures/onnx_genai_workflows/tts/policies/tts_state_initializer.onnx b/tests/fixtures/onnx_genai_workflows/tts/policies/tts_state_initializer.onnx new file mode 100644 index 0000000000000000000000000000000000000000..517b0f069f2496c38f31b74aecf3eaffec9cc871 GIT binary patch literal 2194 zcmd5-O>YuG7-qImwlAPC4VJEnX-U&2#1JV6>Z$6DiwAESvl+HS849!O&Okupg}-3D z^CQ&1r?U%%*^R^#ZTB=Y@5erQp7+Dd;N79yVI%SCPaft7a-Q!FJ!wn#nsR~n#`gOj z@u>aG$}SGvy;QeyD!f$yb|}9rp*b=tSx5;#M5E z`9wHA^JLW-5FzXyqdxJS-V_@jKedm*ig^V&Y~i>b6@dqQM>spS-@07=ic7m@JzSkb zWVNz+^H3=OjLBd`ZN{azNX}7uTc{LqYo;Iq5JL#pA<{m#atd2;3Zb5tcNf=629jjL z#8!#>5hbo3C$6p$w;v%caYxkEHR>8s>KeC7-7lp02x=}A5ofx7P?U~p1Yte8skV1r ztv{ZZS7o_;uUL=@i?Q;+MEP{vOf0`SN^YwvW4+RhOWxEZkM@e{=R4??rY3kRDtPNw z1^=lF-og+LLt6}7xk2w2_d@UAI_w6{1{JCXtdOoh-2nXMeE^P&QVLrQbQz3lCK(y9 zFVz3=G`@`y;(aJ6-`v{QFq55Q(6s!Q3+kLf zgg1NuMU|(Z=ShId)j!ZiJ_azAS5g^J--<5)%~py`l6903i&C?16=n4GsD@_QiAAXf zT1BBqlq%wt#h9TpI9_c1>J{Z`0ZZocs^<=vGgWhqOP@Fr8(+W8&}%sPucyr}@!`^I ik)%pS6E7&-6&6WZqQ-yn#)?&@c4Q{CY})3}?s7#cO>Eh9oHBOeSPlZ%6}BsKEs>^3 znxyPFm*mt#Pwl1Wpg^zfp#^$rt_9jtilV2Y2--u@77gGQXq)s<_s!1i?r?W@b}2cI zek>q?cjwLQoA3S1o0&Jm1aUbRX#E+H!Vr)##XdK6Q)~<=mniDQdNAP@djnJyoutiI%Zu_Buw@ zl+KYl^D}SVH(Sb$8;dur=D%T9p-%TI{IAKhH$U^uR^RM4P2-+fyL!FdF)KIFvCf^9 zY`fKZw^!}fJN<06Y4m#fYTfJ_-RkNCNxn53-1zL7nOSd}vr82Xr97=-mi1m0KG~C| zZ_N_6ac#G^vrFiYaU_BWp=!I;>+6M@o@f4DVE!!!VBD-%O})Uy*ebFxE&#^(X-8xH zi@0=*r-m0Xv_`A1%c{~vpyYJYevvtZ&^@EMZtC?`zjr?IVQhRy!}n<%g}@h} z!bAh5uLfXWHTt@edY+3sF0*h~(qn)-JZAQyQRCfuZ+hgzSbRt0{Y#Gz9kBK1d0@WX z>fK#8&G$@QO{M*Ky2TcxYI?$Xv!j9hQyhiCg@nqXAq zUNx$#CSRLvph|qdOT|$R)8MkL!8;mZJaOlRHD=_}U z5N=nIgpd|OAuZUD?jFbc+W>rRJsh0oawj3S&nR%`aCRwCT6uT)UI=?Rp279KAvlWl zSP9i*n6VN0?P&MQ2ibMMn*v!41zFt{AP?hzRk7X{Lm`Ff8qwRQ{@3=u#ev9ItkN29 z4Q!eouU*oCeks{hr*dOR1Z$Ib+97lIts|Av6foyJk4AfZIbb{0Dt=1+CRrw z)z^&P9sU%~_8f|EMqTDl@ZNz13fQvXPVh(>0l%EyV?FLd?C(dwevXg5jjjlLInSG1 zGuX?pDYx^n|L=zNUKN`~UI>0641CzF+xg)0olX?RFNqXiE(o!g!?2fk1NNWB&;g_0 zH(S>Jz2)ZjR>=mmUe{Dfa#rgvuCFbwuj#Uqnxaee^o`}BWxQ=?QB-p5<~`k<5^dgB z&WG;L9~17+FI7$o_u`SW7r%r)k0DN%n5b8grkpN2n`XMKrs%ksb6=8<4{WY&>AwFv zhW6K6o%MbX=VX|Wt8vKbD#^l-i>bp@xJ+7xT%>s!8H5sMt2P~n4cW{-v>;WHxfoXzEHnb(C6!r^GT2C7(Vp_Z# z5U3P_6hbL%8-@6Lj4OGiAWb@wzi2@Eog$~6vBMnvOr$KQM{#Lp0>%67>4_~r@PU|p z{XzmAXNG|VNO(~cXBfbb1Op|VfQPNE-rtKO*X5I*p#3DQ&1duaZAC`M$>UiitLEjB ztf-o%l!|hxq?xCc5+-3{1Tr61RA&@mCdi6v&l9$^EBp&P3fTSeX?>BW^~C|L-xlGX zSL@kz?Ev?t1-dtF+qk#CB|H&6IzgORwrM;tyS;2aYfS)j>#^1Du4V6=^_A7WRCL^5 zAm#O|i*G;;qgj8?ux5+OT&hHUq+KcJX+GzxM4ZT6deU10Nw(u=r&n*bTT>s+eW&;&^V;XpLO%-P({&(NgtD;Nzguk$mMrOK_)329k4A%&E0_Cq)Y zoHh-4@UormPz@k)c0*wV%&<}fio zvtTb;P=4lTuGCwm(XFh63h&ae!gIcb1wlRof_z;l$g@F$boziBBUpkk|B*l&GeI=0 zXr(g|bm6Mezq$@lZ#k6?fbT>%U`pzkpRVRb229Ak z2CU}4ECVL|Y)=Cwynk^9%n_~FfN2G4lxjg_z=Sr(8n6mUX-~p{$)~ru%*6(LFS-F! zQpSK~kpUBOp8?BXmH`ufwx=)&?_Zn&b3`jPU|PX8U`1rWgm%z?VI$^A7%=(tHkbK_ zqiCpTNVwgY7naIWRup$#x6aC5!rL;ndTT{j)l?xk^72w!v6rf*4|)nRlMr~rrroM( z)Ouvw<~elOJ&sC0R;#aBhwL&E+6yNf?bEa5;N2}Yc7NmX#_pA+&|v+gjR)&Xl?7pl z{z7<&zJ$V5#5>*!mEiHJNO5=-DTb&Bo?mC4s)%!pJCch0d=DDZwuhuIioKL~#NaP+ zYiO#2SF-Pz4|MA=GhPjX2Xw^$+A5baB8TP7iFH-e^V0rK*X%(S0cle@=ssF)rA`Ms zGq)w;#$dV4d4(5+dhOf&s?nJmfj_(`^bFc}=brAmHA`o~bKPa$Sv%ARVTZbEG)-Ld z9Igr8u#V3ot@n&3*|y&0&P^P^wxb8Il#~fz=R^T4A@>Eab6-{fOZeHI!3^OY0@yz% z&=3|YUN_>%QyjI@3U<_55=E_qb}(v%^x7v8wUSS7b3y+q?p@^<2C-hzNYj45zEG>_ zYB@zu5jhLyIsBe@rWEsMN-LCg|%iUy8;;^17jq|4hQMn0kX(h znFHFK)24THEJBurQ{B74Vz8pSv2oD>2lF#muUvE4aD1f^YO&{rvRHs0J~XhS@|sX` znIOpl)3_N2Igr>V3A9ODjab)G zBlPk5XjB;D^Krxz4x*V^}618y%M(01t<$H!TVzypgz>vMggOs2;khB<6JN$0n@ zK-*y%+}}qDX2ZkQVb@INg2HI0jVvpvpfK8{wjErVA4WTbgJHCrmGo48GqVmzAiv+h z!B_u~LY9Mq0&J&#){#cgnv#|mtSM>vow=s;Fc4EMkLbCoEW2!U+>wCw*0(aG6;gEl zA%_V97JszDW)^>5W@j|{2s0YYiMyindg%NGGF!!WEW0VPd0}iOv9THE0@V1saTE?B zlh992H)^tz(9axq)ZpN+^|)P+JN`8$9d&*EV(pl*c+$x6E4(uZ4q2Iu^UZfBb$+r>_NcdzKP3*2}^b8nla zh2Bt5gS}{n-n-^{j`f%{UtW9+fw1NTdX-S03DZ!q&vcT}6`{+VVqGTfD|pSHM$#lP zgwj~Ey;AS@oa2!;ln<4 zO21$=%K1nd`Y>@J@YbJ7d?c0$q#>7d0L@4ym5fb}UtoT{ppyI_J+|X3D^!xS=zU+A zTuI_g5B%pVQ!2?ADe{-8l_aoX{bjxaoUqqeFx8HLzy^;3Rz`89%xNeHKNMIQee~%4 zM>kc^4y^P7I-ZP4N71;nXW8giSJ44!U(bX~f?M#L3?Nn%PbP1^1H4#dtE|R`-1M=B0q7AL$Ir0KX^Pn(<@|JtG~m zKh82fdh>E{lF(2MEoVW;Eq5Wkh|<8oF_!@k>nM?YIJtZV`G+m(2%3VbkQroSjd=%r zmtR4iG^DfWG*HeCQXZD1K4e=qb&q^Tm(L?_-jXz23yrMsB2MWb+7J9pK;J&V9LM?e zj%oYpPXu4jB_WhdPDrCxyIeYAtOIn@%(3shRWfe15Nog{g}@RyutW(gQ3pz<5NoK? z7#anw9%&r$oY%DqVo-28%E?{a3MN(~pMS*Q#jLI;8|CQd>JFfNP%VkI5C12wN>fGY ziA%RC7FkGDn&P*SsxSCugQ^6*IxEL2*=jd8cLlU8TdNShnpS^-~k*CV){oNUZ+ zXWeW+P6ow~uEjnioSzE#=@aq;M~rGG8bPG=S5olziM7QWXTpAPg=sj}$S`YeZZdZk zrghdYtT42RMkVwYF~C254H!8Ng27HHcr zhpkZXu^%9JoBe|8Uorn7{3yfnqYhi4ICKZJ0Re<`fQWYp36y*t2+$vL;X%PKVj;ck zm#~O@@Dw}(SJ;@~ zsAmmAo(rx|ZzXaFu(r?gdqQN0eE>1jSmT25N)sGtPC%fEse`-AhYjw9!2#e9ur2s6 zT(O#%h>rXmgB!WK#0r+f`%q>?=G^dM$JyfUmwap&od{iAbAc-*MUzgV6A&7uEmw+t zbujkd+ThXmNRzI{f;iKe-jm?0^1$T>j#1iC-6C%n%Q@3+U$sCNJe2Zd;fn@~E?=ta zj0j(&jS#-PN0XxPH5!H&jYJu~TqhQhFYkGkXnegr%-n3QX0*?|!n0zWEafDQvU}0)P)@*` zY+MW44|VI?`=N3ajwTzIBc#kl5W!^QA`>4&c%O z`)KaO)gR=+mBDOro|!@MHej3;E^K6Y#qYu>8kwBFa0oB|tYz2g)ovS~|C@~cAN}c} At^fc4 literal 0 HcmV?d00001 diff --git a/tests/fixtures/onnx_genai_workflows/tts/talker/model.onnx.data b/tests/fixtures/onnx_genai_workflows/tts/talker/model.onnx.data new file mode 100644 index 0000000000000000000000000000000000000000..d3d841223b6dcd164d80d6b7710ba7835fed5123 GIT binary patch literal 2048 zcmXZbiC@k69>?*zX+cv-H)~NxNxQC;bUyDWS&~Xxw2Jms=d`MGe!`RoUCW3JjVrW~ znIWQa&hMMdSWhHnz0+h#g+#a1v@q^{-1p=0`Ul>R$5TmZE0n(G;xhqtqa5rewelT( z985P$*!CA3Z2ji4qB|TcocGhTA30oZOo#3+4)QN9$P4H2X-togw&u`^IM(nYs*qLeYOOa^U_(*FS%IwqL5$d zm5Ut{=E3^+9IzP@lt$*j_(Kzwb!@??gC=v{umve&9E)$s28+>P|5%!h&sI<9XjK-n zm;_}Lvaqeh1b$JOa6E93KUbH5fTRp|MLh#GTfwru(lNE)fO_|(VUJ}zj1^*3?e)V{ zWic-0ua~vjh|sxsJqr&^MR;HgyRv07-e{=sy@yg@)#nQ3v&p!~rD0H(gv4)m(xdB% zNRyefirWcrYsh8EH{;QMTfy7a$DwP67Dmp+;>rMryFbJrFGmd*wnrn`{dfNSrYO{m zm$GJIB(}B-*fiZuaO?b@_P0ksMI#d~RU6PeCWK95IJ7!V`RI9JIH;GvdfG!!?zxUB zz6*wh>qTcx1mQ$t5NbODQ5_nM=KKI$((4x{82e*K)G}6Zc^&-ki<#X@U%WQi!KbzO zU@+bi0YTolyeSv2p02_6(wk)J?}_12ZDy+QKt~tHdYxBcZ{aw$t8xXBf}WGI_HwkY zks@OAGT0j#qekHhC&y#F@|>lR{gBB5Gh8q`)0thpMyP&hL++YFyc$fvC$GhbQ1`{t zJ&O>;#$;b^I$-p0FiZJtkD#z9W^c0q7j<9A!-MVcs>K~c+iam{D28UaHBt*Ws{GXw zs`}RKaho}2D&?}uelyHU`;|9Qn+vmNlMpo71U(Nq_|FvJ(~-Zy)W`^3aUFcC(JZv> zk+M`n1Dp<<%`~RWz_|PqR5o!MB-UA2KdO(4cTTYB*2A?1N8Y|!7g-MzS@tCzw99;$ z#`juqoFNc?%+SR9`jG$GNDhjGnwbVHG@FFpBdS>8vVvLMQNhP35vwRx#*jRNKY8^t z9hBQZV;jz%9YsJ(=X{}S~vYR=?}Vd zOA0gbQ|dfyi2vMerQJ>C{4$gK78=avvWnNz6QBGAAE76-;N_~~77++LEyL*$! z=zK9L{eNZg0SI^~aK3aS7SwXm`C6V1&8i!4z9(P8p4|}$b#l{Lf9+~vq4r+>eo4B} z-PR0tv{!i0T>|gZCx!Es9}wwY77n^iWEtrSVaOMbi5|2IABT)*)fNLn!R$xW#E&D< zVF@B%X^~sDDZJFC(WKY|y!RXdnHXj;0a;Oviwi68b)dNHY2*~Sgl^WxVJO9ebZq>w zJ;RUwt~yC}C^MX1mj<#AqBt6T5W{?8#pGC^&SwTlNZr8|I>)z@R%;qm_wOY03%f`n z{U>$5oX6J1en+Qzb6Mh=Vp40n%a_@fkf-@KXwxdCMZ+8hyUJ)VRUOUzX=-qK#(VDj zi5xp5Y-Q*Ldc)6QRVH=xFL^Po>Xwtol?-TAHjtbV9!1=sktAdOTYUwYmBzE---_PpPS8JsNA;=*Y8Zn0vmYuE}G<+V)QJuXbZ=li!i?yfoH6@PWSg z<@0Wt10*Q2Kv>_Obo@XrbYnhKP3>(u5wDC!O>Gu`Oa=YD9BaL+im)B3ta(HoW8Qy| zugL^-T1v6nLldsb0-R0K!W^9vKDJN?Qr}FLTA_w{zf=B}P%3DFg#~ zl+tL19v3S%c*z`vcDc;>m?eC6U+155)<}M*1G7+D>}lul*4_?Nid0}YVF6?tU-Gq2 z>>){!vhq_72yiiCFX9&=Km0gtuv`q;>rCi&J0s-$V!S_2Uq4&Gcl)~FdQ2i4QeKLR zjcb|C8CTpJ(-E3@EyFge5Lmr;$KbysVX$)rMzo&@-_BfxDZ7`m!fFp3YY?%)g`V&i z4Vih^qRqKPk;0CI-@WtF)x-9$XI;bi*rcm{Vb*VD@uZ^CRPWtm! zAcm@?s22pmCU+M4GJ|pM%qc$qdI*%#vzS;r3@WBh?B})NNOTa=!hIVs?3#qp8xfeO z?}L#)H{zayz1-Y16671owylmr{liUcBrzIZ>Nn(j3u3UPe+3ql#Nrzb5ni8;Lui1A zj@HCu;yzpERht0Eh#WS2J`q2f@%*8xB=nxpgTawx1fS}F(Tz2*Dd iQxVu9Wv>lIXt`j(DnE#@r@xZs%fwjuISZ1!H2ej;sFaKV literal 0 HcmV?d00001 diff --git a/tests/fixtures/onnx_genai_workflows/tts/talker_prefill_embedder/model.onnx b/tests/fixtures/onnx_genai_workflows/tts/talker_prefill_embedder/model.onnx new file mode 100644 index 0000000000000000000000000000000000000000..27a7b8a0d36a6aa8eb99138a75ca0e4817234d04 GIT binary patch literal 21644 zcmeHP?Qh%09kxU}rmsn=UYj*wfy`BtxOH2Zcck9Dbx_h4>y!jb8~4i!kDw(crYcLC zsI=Ms8h`A6uulW_Z7YTW0|pG}fORO=eHglS-KT9WhVA$w>PYL3cgh^VoC5?fMc(oF zJa@l)?ztbSJ_gSm4ENmop0=V_P(|PV_=5uYKmjjL?EbsX*cy$UyKcX4IR|@Ax9f~6 z6X(GsI&ug1D(^e)z5NMLgP+~52hn?$VA1c|lYMdrT{9U4wvzYSiV6&dUB}YPXTrnw zE?67b2aY$gI}W_fsA_xt&G(&wv9oh$CwQ*g;bXlXJyq&X+w0*8@__Bz;L_;bdzIl} z@WAVg-O;4d>D!)X?Yqv{9(VSC3@h(!a4X+oDcj&< z1uTd^&>8s89r1gv?LkA1>1*AtWoR|VZRXPK7YR5v$_n}Wlwk0iB3J}=_#B(7JMO)M zq1&}g?JA3#=_LUdW_dL_Zof*0R~EAw;gwX9H5ope6+(C0O>W?Ph8;L}!Cw}G6o^0y zzsbhk4|`Ru%);D19pGIpE9CEdU@fOC@IuN0S7JF})nt?fPAiMrqgNLAWGD+vl*LC1 z$m-p4IwC=woQ}YiND+gHx0s?}wfj@33I)&)5@we7%fi(nz(;bk^m zQ&SGrG}i9U*7TZC#?*4=+k`^9n9=EPILHx?<3W@<4 zT@jghQqx=LCShGhS5Ph<){j+JoFI<>E;$2Efv4>f>Tg8<4a@BL?@r(CI2K0Q71j(G zUJ>XHjELLk7s!Wo@;|_w#u-auj>Ry?ro z*Vv@I>q`mm`m<8&ucg-KrPimV)}KrNZ}v4dss1QvU*oE*XtlVv2*Y^CmXya}f_xd}HU zeU#Z9dFhcI5bsqkvH;-BG20lG2yQJf?Otpx9$tB z2qS~P?(+|48BBgYQ$8XCKNLqfU))X_K`9c7IVtFgP)xRuQ=vFNDPx}qbX}3iVe$Zf z{lUl{kibx7&wbG)VGpmC&xFVM=+yk_E=W2`Bt1mgkJNNWFEM1NyPD8)!>HwIxxt5A zL}Ab{S1Bq)6o%Le&&byb?<0YAic#an?!e|Z>X`zgAtErMGar1Kv9YqLsfE(&b?QBs zwz?GkJ8}-S)dK4rZmUJsIisyAtaD~tU1puDw$&BZxmsKO`aj{nAL!IaUoL~Px(Y9W zRk*TePdfYHN%%z26+DN$!ZRTGCw*tI_X5vvS{nPfhg`T@(~mb?-GW zx1Yh6z$U5W>a3Ed;f1(Erfcb4ll-&|EP?Z+dH|NyuMXGzhV))r=#wkw!C6v|>emp| z@*SDXttcZL0F~6SEZ>Sa!SkV75#^83t++_7?sS5ZO8SA{zCw>agP3P9^Nf)=b4_*= z89olq5vAkVY%UckF5*fi~{RrpR z5{%WGa;evev1%M4|JK5<7|m^5d%?F^<~``P;M*AU%GP+^N%)S&-1bX?OAjnW+?dp) zD+z8qvQTqw<0QeIru5FOtC6k8F-n#6`=Cn3K&Fz%c;S-9+>&C071Q(Vis)b^nO#0_ zv%$L3`yOw0VLH;Omahb8eMSqYQ~uQ^Ua0 zK!%z6_R17q347&UJ

@gk9(u#OlTzWoBvrp4 literal 0 HcmV?d00001 diff --git a/tests/fixtures/onnx_genai_workflows/tts/talker_prefill_embedder/model.onnx.data b/tests/fixtures/onnx_genai_workflows/tts/talker_prefill_embedder/model.onnx.data new file mode 100644 index 000000000..e69de29bb diff --git a/tests/fixtures/onnx_genai_workflows/tts/talker_step_embedder/model.onnx b/tests/fixtures/onnx_genai_workflows/tts/talker_step_embedder/model.onnx new file mode 100644 index 0000000000000000000000000000000000000000..58e017b452c5aa41993a77a5b7e8bb5feb777368 GIT binary patch literal 7296 zcmd^@U2oD*7{}MbMt3%p9W|0hw}~3eVpvZLEzFoGi80~MY*$;Brj)Y*8!ggKrys>D zKY;Pl#4G&*#tRc;{1ko!%iBV24^PKkk%WZu^t}Df|D6AMjBETg9xC^m{%hn=+aLEyH$pKvP*5WrF zJq*>N;~M4(Z46fQq$V%1wv`Axl>pCp%kL8pty;ymJ#I3#$K97_l)j~UQM?;AJ=^2> zV1q2?QlsbiZi+v71hRn8%FqfgMNnjE&J_6tR3w48mhdfxOrJIyx+-pjO0*vv$WzTF zyx;MaCfia#!D$f$rv(J}-&zW2(d7s{%K@JAZYkwUSxO~}rBvp?QhvdfQh~NU4YOQn zl%Wubfpap0A;Al{<_jJ5EMOQ4GD<@MkMz9?})>4hNy{MqQ#64$u2Ov}ta_PA{yllkjmeoCTo#B%M!qMXGz4&_A*W=X{su_!?RhQaFuLc(HD{W)ty3a{s<> znr@;S0*BWoU*Wvhg!%*IpFzxev#mwDdd62#+G}bVC4@hjeiQBLP>-gQ3f1di9I}R`gxyS5agoG$fZ0Zmr8)kBPPr7F-T&MuMAif yyQBs-%aDu;77W*n65+@)1Dq+U0pO@iIMpOt$YiT)C~exELZ@Z+Y{x!wGu&T}N1o;Y literal 0 HcmV?d00001 diff --git a/tests/fixtures/onnx_genai_workflows/tts/talker_step_embedder/model.onnx.data b/tests/fixtures/onnx_genai_workflows/tts/talker_step_embedder/model.onnx.data new file mode 100644 index 000000000..e69de29bb diff --git a/tests/fixtures/onnx_genai_workflows/tts/talker_text_step/model.onnx b/tests/fixtures/onnx_genai_workflows/tts/talker_text_step/model.onnx new file mode 100644 index 0000000000000000000000000000000000000000..85ea4a6b1901c83e68e99ed2807de57a204ef019 GIT binary patch literal 1759 zcmcIk&2Q5%6mR;mrmqc4y^WO!A&3B>9IB@4io^+(15zbUh|3VN+|*mWwN8oMs_+MJ z;)FPGgA?5O3vtpUxI^QXT=KsC^Uv=+vx)W=;e@S(Q|op8ZtvM|3w{}JC?mq=j9*|$ zHxkpugiZx&tnup^SyJ3{T3dmKQ>D4JMDqYDlutY5A4b`0K}|>!$&R74u>}oI7E~-r zKv8M@5biCnF1jJ-8xcfoDZ9a(2!UseMkESmA5i)2yYYQ+u0_Q)un58>6>$W7H=x0! z%8@MO&&ZKh0j`IDXR^8*1I&NAP!#IQAPZ(#}uVd+IiFOF-d;*V61G!K~wF!~I3=U8`U$BBI zE_gtpY?^0xwS9*6!7}Tpk*7@_fkq=bh9gzAox4>5BQA1zw7UuSpsLE0E<4(T_z1vM z)@q!z<55qgd5*R1wc|L{X9HzNsFSs=mYZ*18|5rqVNZuT4^2$lrp}@sXh#jv0kqYv cE&bN!s0t;!Girksgp2Ni1yLx%b7>oY0Fs6WN&o-= literal 0 HcmV?d00001 diff --git a/tests/fixtures/onnx_genai_workflows/tts/talker_text_step/model.onnx.data b/tests/fixtures/onnx_genai_workflows/tts/talker_text_step/model.onnx.data new file mode 100644 index 000000000..e69de29bb diff --git a/tests/fixtures/onnx_genai_workflows/vlm/decoder/model.onnx b/tests/fixtures/onnx_genai_workflows/vlm/decoder/model.onnx new file mode 100644 index 0000000000000000000000000000000000000000..59df2e6aa626e5aac02f9b52bdb3a967c5b9f925 GIT binary patch literal 656 zcmb7>Nl(Kt5QS-=#2QFjjfkK@Kvcw1B1Jt{MPfPe7jlWRSYS09waMj9kuxCRVknz(?6XQEB~MB)o&~eWDHL(%4eQ twYacLK04tPn!Cg?+(_h7A!XJ9_Czj%UnnO9I+5}%n;EW`_D7RRUNCZ(nnO9`dsW&$NZ(l9Y44$qX-ypqh4 zN+C|Ls+7#U^e7EvL!`MRIoO3nx!8qRlM+jkGq|{fIEqsXOH=cbQ=_!F_%d@7(^KQq zQWHx`i&Be)B)G&l7=?ro%9%MDn7EiY6p{?MP~9UW&m{{~Dv3~vWHLyH6AKrEfG7Zz CdOHsQ literal 0 HcmV?d00001 diff --git a/tests/fixtures/onnx_genai_workflows/vlm/embedding/model.onnx.data b/tests/fixtures/onnx_genai_workflows/vlm/embedding/model.onnx.data new file mode 100644 index 000000000..e69de29bb 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..4fc8779c0 --- /dev/null +++ b/tests/fixtures/onnx_genai_workflows/vlm/inference_metadata.yaml @@ -0,0 +1,608 @@ +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: &id004 + dtype: float32 + rank: 2 + shape: + - patches + - 1176 + - name: image.grid_thw + content: grid_dimensions + dtype: int64 + source: image.output_grid_dimensions + contract: &id005 + dtype: int64 + rank: 2 + shape: + - images + - 3 +pipeline: + workflow: + manifest: + ir_version: '1.0' + onnx_opsets: + ai.onnx: 24 + adapter_abis: + onnx-genai.image-preprocess: '1' + 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 + - sequence + 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: &id001 + 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: + - 1 + role: + kind: opaque + source: + kind: literal + required: false + default: 2 + package.max_context: + contract: *id001 + role: + kind: opaque + source: + kind: literal + required: false + default: 8192 + package.one: + contract: *id001 + role: + kind: opaque + source: + kind: literal + required: false + default: 1 + package.active: + contract: &id002 + dtype: bool + rank: 1 + shape: + - batch + role: + kind: opaque + source: + kind: literal + required: false + default: true + package.false: + contract: *id002 + role: + kind: opaque + source: + kind: literal + required: false + default: false + package.zero_batch: + contract: &id003 + dtype: int64 + rank: 1 + shape: + - batch + role: + kind: opaque + source: + kind: literal + required: false + default: 0 + package.slot_ids: + contract: *id003 + role: + kind: opaque + source: + kind: literal + 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 + role: tokens + stage: pre_adapter + components: + vision_encoder: + implementation: + kind: onnx + artifact: vision_encoder/model.onnx + embedding: + implementation: + kind: onnx + artifact: embedding/model.onnx + decoder: + implementation: + kind: onnx + artifact: decoder/model.onnx + 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: *id004 + grid_thw: *id005 + token_sampler: + implementation: + kind: onnx + artifact: policies/token_sampler.onnx + contract: + id: onnx-genai.token-sampler + version: '1' + bindings: + logits: logits + token: token + parameters: + mode: greedy + application_overridable: true + termination: + implementation: + kind: onnx + artifact: policies/termination.onnx + contract: + id: onnx-genai.termination-predicate + version: '1' + bindings: + tokens: token_ids + eos_ids: eos_ids + iteration: iteration + max_iterations: max_iterations + done: done + continue: continue + token_state_update: + implementation: + kind: onnx + artifact: policies/token_state_update.onnx + contract: + id: onnx-genai.state-update + version: '1' + bindings: + current: current + update: update + next: next + last_token_logits: + implementation: + kind: onnx + artifact: policies/last_token_logits.onnx + decoder_state_initializer: + implementation: + kind: onnx + artifact: policies/decoder_state_initializer.onnx + decoder_step_update: + implementation: + kind: onnx + artifact: policies/decoder_step_update.onnx + cache_length_update: + implementation: + kind: onnx + artifact: policies/cache_length_update.onnx + state: + token: + contract: + dtype: int64 + rank: 2 + shape: + - batch + - 1 + scope: invocation + initializer: initializer.token_slot + recurrence: + kind: invariant + logits: + contract: + dtype: float32 + rank: 2 + shape: + - batch + - 128 + scope: invocation + initializer: decoder.setup.last_logits + recurrence: + kind: invariant + attention_mask: + contract: + dtype: int64 + rank: 2 + shape: + - batch + - context + scope: invocation + initializer: initializer.body_attention_mask + recurrence: + kind: growing + axis: 1 + increment: package.one + max: package.max_context + active: + contract: *id002 + class: semantic + scope: invocation + initializer: package.active + recurrence: + kind: invariant + done: + contract: *id002 + class: semantic + scope: invocation + initializer: package.false + recurrence: + kind: invariant + accepted_len: + contract: *id003 + class: semantic + scope: invocation + initializer: package.zero_batch + recurrence: + kind: invariant + slot_ids: + contract: *id003 + class: semantic + scope: invocation + initializer: package.slot_ids + recurrence: + kind: invariant + cache_lengths: + contract: *id003 + class: semantic + scope: invocation + initializer: package.zero_batch + recurrence: + kind: invariant + position_ids: + contract: + dtype: int64 + rank: 2 + shape: + - batch + - 1 + scope: invocation + initializer: initializer.body_position_ids + recurrence: + kind: invariant + cache_0: + contract: + dtype: float32 + rank: 4 + shape: + - batch + - 2 + - past_sequence + - 8 + scope: invocation + initializer: decoder.setup.present.0.key + recurrence: + kind: bounded + axis: 2 + max: package.max_context + service_group: decoder_cache + cache_1: + contract: + dtype: float32 + rank: 4 + shape: + - batch + - 2 + - past_sequence + - 8 + scope: invocation + initializer: decoder.setup.present.0.value + recurrence: + kind: bounded + axis: 2 + max: package.max_context + 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 + slot_ids: slot_ids + kv_service: + paging: paged + allocation: runtime + compaction: true + groups: + decoder_cache: + sequence_axis: 2 + layout: bnsh + logical_lengths: cache_lengths + storage: paged + ports: + decoder: + cache_0: + input: past_key_values.0.key + output: present.0.key + cache_1: + input: past_key_values.0.value + output: present.0.value + 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 + outputs: + attention_mask: initializer.attention_mask + body_attention_mask: initializer.body_attention_mask + token_slot: initializer.token_slot + 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: 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 + outputs: + token: sample.body + - kind: invoke + component: termination + inputs: + token_ids: sample.body + eos_ids: package.eos_ids + iteration: loop.iteration + max_iterations: request.max_iterations + outputs: + done: loop.done + continue: loop.continue + - kind: invoke + component: cache_length_update + inputs: + left: cache_lengths + right: package.one + outputs: + total: cache_lengths.next + - kind: invoke + component: cache_length_update + inputs: + left: package.zero_batch + right: package.one + outputs: + total: accepted_len.next + - kind: invoke + component: token_state_update + inputs: + current: token + update: sample.body + outputs: + next: token.body + - kind: emit + value: token.body + output: tokens + mode: append + - 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: active + next: loop.continue + - cell: done + next: loop.done + - cell: accepted_len + next: accepted_len.next + - cell: slot_ids + next: slot_ids + - 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 + iteration: + value: loop.iteration + contract: *id003 diff --git a/tests/fixtures/onnx_genai_workflows/vlm/policies/cache_length_update.onnx b/tests/fixtures/onnx_genai_workflows/vlm/policies/cache_length_update.onnx new file mode 100644 index 0000000000000000000000000000000000000000..72ad12c9f4e344e870cd12220f36d4f2593cc2b3 GIT binary patch literal 386 zcmaiv!A`C9_8vCK3xkAz0u^oRWg3)KzA zxO>Sx@8;a$n4NBHs<-~}GlX~OmzWVmPUi)&h%AY61}mkwu?qS9%}IMF6Brq>L0<`p ztd}g{s(Q_1Ys}6|rz?tODZJ-}MkkysY8F1fBy*S@QTq<=rLEBK2{zV-Fh|FQQtOLO zg`L4_~e(FQ-eEJ2Sf_QEK literal 0 HcmV?d00001 diff --git a/tests/fixtures/onnx_genai_workflows/vlm/policies/decoder_state_initializer.onnx b/tests/fixtures/onnx_genai_workflows/vlm/policies/decoder_state_initializer.onnx new file mode 100644 index 0000000000000000000000000000000000000000..a2d78bb4b680638f98b2ec428b1443e7f3b2dc25 GIT binary patch literal 7283 zcmd5>O>g5w7_LLwxNnnwOw~5^qE$93Wl5zd9_On_EA3Ko=>>@umqo~O9Jh6w+97tU zb`h5zIB`KNLU5{txbOowfe`sJa-Aj!UrY7?+keG)#*u9Ph#7 z;l*j)b(|Z|9N9y^ZVn93)6Z;cWQ@$STax?yi{lC?!dHO~-80>x@m3&*mw5dH3@d{*#(mS7rYvTFFg<@mPi=$D3f zp{yJkp09_%o7H?Us4YXzcguwgx&2-u#e`73rA!!$khH3DOkr&UAmmQB9p8J77B6H) z8OF`T;wo?e350M>!`ItN39$uFA!+UEQZ%fsK;F2qy-efgPYKxwJTXKv%_}=Iyt4DZ z@X9w+y|P1O+*II)UDhAFbMwdFBoN{PCWy#p&+XspLfCLVaG3IeEECfY5-Bmkqc)WF zI9MW=R2vTxKEH6`!$Tx16R%ZOWt9yeRZX#%W##!@0u~E_uvY8d!1a}#n8!H9hY_yQ zw};L+pCj{`m~8e{f#Ol3jE39@FU8ADxL*M*d4*Lg&6P~UvS-sXjbG+E=4wa2$ z6w!8AU7ozkin<@M(D68umZuRU>5p%QhSS$+Ls=a^hTv$Ux)hDGLiSylmnX}U;?VW% zuuZl5o>EM@RI1LcE%=trA;D=bL_zU{!oZ<(gCqDoO;rHORsi1bv+SE^7k;WSg`YP0 z!cUuXE&O;ON-W6q0JKRq!VZsfINtkPS$n@|+EhiLY(?Nri$~(js;G#oBc@u~d5}~q z;!2eOI=ckO8S?cvlQC4T0?;-;?C`@b8z_fDprp?iS$4>H#!Nb7dPug0!O85xy4Ayu z(yKS>IGE8aYmIcXydB?u0N#E`$N8XH&ahDyd_z-3pVicvK0C))(r5}7orx>UPFdtD zRaUgwW#wUFN#jFHC*gn`JDkDtM`yBo3(dni*$WVXnr027x~dru+V z7?4faK;cp%oRg*A6JV|prhoDRrsq-e6!v9_R9S^NSR*Bl0?5g7?+9kO&wp2PzQZ>B z05WbSkJ@wWgi4YLiJhR_#FOba%(3hOESxLiHR@zR>AcKcWyVqgl&Ok4leGiI+sp-B zEE>Fj#FbUGF%d#Do$oS*O=UM3?4CdRk|oF^H7SKnPT@9xLY@c-+S-$eZT5snWLUyA za%tFIZk5O4nT-^2%LP=5n?m7mTgl0Zr6aWV_|(Gbr&B9@xg{_PKf|?L%dz T4(sH^GX3yxy^r#+uwMQbCi2d} literal 0 HcmV?d00001 diff --git a/tests/fixtures/onnx_genai_workflows/vlm/policies/decoder_step_update.onnx b/tests/fixtures/onnx_genai_workflows/vlm/policies/decoder_step_update.onnx new file mode 100644 index 0000000000000000000000000000000000000000..b2a4c7288de43c04c5ae83a2119d81bd95d87b9b GIT binary patch literal 2037 zcmcIlL2uJA6pp)g&EDE=bNE1fDntYBIciqR=~*3Mjd7=fZcwIPfSiytkQGDCx<%OFbj*ZLZHA5R z0qpq~=M9hZsjvgrmkoPFg&+f$1vIb+S2*+GW48*{ovg415w_eLL33nl&>H)Jp)DbTnbau_5_X$hYnVBQ2$lc&v}oQtFrm(!%*`$wXQYF^m>2b2gPE>9eFT zj+cZH`YSEKUy_b$-*r@b!BIO(=CPx+>C*!B9VJ}UhJx=2SNA2?5xND$jCSr`awq*u z{-)q90!mzQ2bc5}FD6)U>Nt_@U)I*tI3?n6*z)^W2PyddILz$9NAOs@%#vw}>}qb*s;+&E8X3 z=4Ryv!}|*z6$6;6Y!=kcLo1<-X5wEZ(c|W;GQ<6q-Q#w#0miIwEUQ-H-)>^lHeQFU z0%Q!Mc)X?)$FI>7*pD98#yR8ai3_FM@V$|1Uy+cIfZ7sgK7GDMr*QD!OyWBD!7M8( h(x8*6J924ob<(^>kD>a{8m!-Ew$!iZ>ns>MmA^m2fFb|@ literal 0 HcmV?d00001 diff --git a/tests/fixtures/onnx_genai_workflows/vlm/policies/last_token_logits.onnx b/tests/fixtures/onnx_genai_workflows/vlm/policies/last_token_logits.onnx new file mode 100644 index 0000000000000000000000000000000000000000..f674183330680e7f8e3468568778f80f8327ea0b GIT binary patch literal 631 zcmcJM%T59@6ox4ohv5KGyHLPr5Q#)$CJ=Y7yev$31dVArr9fq9$7zQ^j2m9V#Q1Q& zhYk#=kQi5X)ysGOd_Bj>VSA`N5$W@H3tkY;kI1!>+E8gov%LAWm#V`_v+67_8YRe% zs2_16q$z$Z7kWeKb^)@I4!I5}<2co>!9g(Sxk}2hW}yg-%Y3Re=?fmxko6}xeRtok zKzWUtWf09&!1auPR|K}NdctT}vBbGJ&XDA(T)3g8&qFd7ZdL8%*vNrJ$3hqXQ4nv> zJ&YP~xv;rCTPouds5=NBK?ZMl)UZAPb~JyCxMUpi7)FXwFY;+P@lN6AkMKV1#g)6? zl`Hha&RREYu$&A|O;H^VlhL|8F5^}kp`;rqU$6;bN}7-|gVRsXQ*;TnKS?4wxU`^> g5SK<{;fvV#qBTXwQ28~3ce|XK`4q1#NNp8A0X>A#xc~qF literal 0 HcmV?d00001 diff --git a/tests/fixtures/onnx_genai_workflows/vlm/policies/termination.onnx b/tests/fixtures/onnx_genai_workflows/vlm/policies/termination.onnx new file mode 100644 index 0000000000000000000000000000000000000000..3693811ae428bef7c7f17652199285b2ca557a9d GIT binary patch literal 3730 zcmcgv%Wm676!oY|NM2c%JZ>FfY6K2pxPZ|JqGZ`|R%wEwfPlh4i*8yJ1Dc$~lqFIj zDIsamjW-2~0tK4KqUa*sW|2S1*YpF@nc+hmDn?g~HsZZ^9%s&-duB#&;Oo=i%$rAd z-|fKP2(K^bgTRks+mC6FwDNC1a2f5HdgeGF01JCEXVmlK=-#{Rhh8CV)C9w~r!1P; z4#R7sJ-9i$c+?Gi|4HP8-Yo7q6FZ9NxyM2~bk3jR+7CaDc7c>yi7{y81Tz+i791WJ zad5$W>bVgpNO5yPzw@KVbH;vR)FSPrA(30Ak$b9w_e^E4UuLggU~f}lZ>>Vo{z9O! zAX=fUuTnNxp=|I0D0_iGO`K{Du(6VhD3N| z8oB471^NYJGuN9&9lqTv1N*011jpqfI4%@Ho0FVLNTr~74X9ix;&F@{3;I>a?3ji0 z4%x^)C5d-TEw!skZ;7#!GGixKF~+5~vNZnA3&7)i*s&M*3R@?l9KLW}YV}A%{7agw zp7{|ss^aL4EP-TJ*qGW+XbvvITY2_Q!$c!2i)^P^%UlUc%gnNy`l|x`2ZG|Jg7;y> zbvN)CfzvSMuy58*!>SPdia^dMBvq=IH6ZItIkN^IOwaz2;+kT>8<~iG&9iD9lFdZ3 zgq7s%Fq1LvQe>>k{r?;ZbiQBvi7z2+IWlXAb$&j?o4toi(i<8z3Ak6q-7uiMq%EwZOFO906GuWp{Ap8M>%sv zL`x$?5yqbEmBsedn~a6GV3T)j=IfMWt0VN45%$#)2KW=$;v=*~#{whc&x%(S0v!(g zc;!aYXlbUv@6=}mJZ7mn#bG9!jJKFsPoa*#y1twll_Vmm-?OL{XP}Q5_Mml zZ2$e@S?~Gj8R!7+>7mmX9jdB^KE4oDgBDj@2kh+hUpB{+aSupXOr&QEhXge`} zi+(o4V8YD1yJn#Z=qzNe=kSY#`8f}Wr%NQ8WRDo(PA})epM7#~zR87uY4C`8`iXGK z!`q01MMUI3w=+=?zWkf}*~5p-iG{a-+Wh4bjbQJ8SIwUk*c6Yt+ngEO>lTK%oiF^J Vb&2+2_sYQ%1oC}%zYc5d&VQs*nfL$z literal 0 HcmV?d00001 diff --git a/tests/fixtures/onnx_genai_workflows/vlm/policies/token_sampler.onnx b/tests/fixtures/onnx_genai_workflows/vlm/policies/token_sampler.onnx new file mode 100644 index 0000000000000000000000000000000000000000..9e16b22bf4e14fc26641479c047034a83731b6d3 GIT binary patch literal 521 zcmaiwPfNov7{=GRKi#Wy5(LXcsbwe(SM<7*?l|xZC{ofi>%!WkB$=#;mmU2S9{h}6 zwQ08<1Tm-N$@9MN^Zx98)JjFjmTLU92cH(S5)m`4P*aNq<^Arxj34kS;3A^Ijn}H4 z!yecRO0$Ti%Bg)=_2Drz@QNwtdwZ-xopzzaaY|K&2}QN(0M4^T>SOxmN+L3s2w>!ofE62 literal 0 HcmV?d00001 diff --git a/tests/fixtures/onnx_genai_workflows/vlm/policies/token_state_update.onnx b/tests/fixtures/onnx_genai_workflows/vlm/policies/token_state_update.onnx new file mode 100644 index 0000000000000000000000000000000000000000..4f3026a099f6c88007fd54d27c96dd2b224c70b7 GIT binary patch literal 1240 zcmcIkL2uJA6t26{BoC?;BcUxwQ?V)zh(&G3X{QZtl{j)4LYC`zYr4!1u~TgzEim-uv`DKfCvF`$|osOdtR9;5)((=Jd6a+AwM8z~8NG zsX{>q#9aZ;yWq`PoCz99v)8yq&2a}j$*zP>7#FxP?!!THc@ZcnuQgAj!~{HMTGMGH zQkL@R4Q{@BKkk9Qb}P?8HCKty`3?Arz~Nb%3TdFm4m)%DCX4CNhkPp-4#|$Sw?*Y= zU7yjqKBGJO_yCnY`sDX&p?%%Y$OJ}k?e3>WAb S^!}NnELr>)uUgRP?)?GSSBa1S literal 0 HcmV?d00001 diff --git a/tests/fixtures/onnx_genai_workflows/vlm/vision_encoder/model.onnx b/tests/fixtures/onnx_genai_workflows/vlm/vision_encoder/model.onnx new file mode 100644 index 0000000000000000000000000000000000000000..aedac17fe11658323b62d6957d238d0e0bc7c033 GIT binary patch literal 208 zcmZ9`Jr06E5CvcXm0@BC8Ic$(VqxJCOl<7D0?me>5jNp3?1I{P3@_xpL_2L?UcUDL z4=)3cxM<~#OEJV{&x%}G#W{-$$GY!WZI)a~Eq=o-oHMSq)OT_Rsbo_NuG1#tF53NZ z0s3H&hYs0otkZ%OTJUtD5FzhmCUkKn9|>%75orG1QD}md?D6jktuYuaLYrpE#~P)F MId;R-6er&54O$>L{Qv*} literal 0 HcmV?d00001 diff --git a/tests/fixtures/onnx_genai_workflows/vlm/vision_encoder/model.onnx.data b/tests/fixtures/onnx_genai_workflows/vlm/vision_encoder/model.onnx.data new file mode 100644 index 000000000..e69de29bb diff --git a/tests/generate_onnx_genai_validation_packages.py b/tests/generate_onnx_genai_validation_packages.py index 497720a8f..467102571 100644 --- a/tests/generate_onnx_genai_validation_packages.py +++ b/tests/generate_onnx_genai_validation_packages.py @@ -4,18 +4,20 @@ from pathlib import Path import onnx_ir as ir +from onnxscript import GraphBuilder from mobius._model_package import ModelPackage from mobius.integrations.onnx_genai import write_onnx_genai_config from mobius.integrations.onnx_genai.auto_export_test import ( _Cfg, - _decoder_package, _diffusion_package, _model, _value, _vlm_package, ) from mobius.integrations.onnx_genai.workflow_metadata import ( + write_audio_codec_workflow_metadata, + write_language_diffusion_workflow_metadata, write_speculative_workflow_metadata, ) from mobius.integrations.onnx_genai.workflow_metadata_test import _speculative_package @@ -37,12 +39,127 @@ def _tts_package() -> ModelPackage: 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_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 _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 main() -> None: parser = argparse.ArgumentParser() parser.add_argument("output", type=Path) args = parser.parse_args() + decoder = _executable_decoder_package() packages = { - "decoder": (_decoder_package(), {"config": _Cfg()}), + "decoder": (decoder, {"config": decoder.config}), "vlm": (_vlm_package(), {}), "diffusion": (_diffusion_package(text=True), {"guidance_scale": 1.0}), "tts": (_tts_package(), {}), @@ -57,6 +174,33 @@ def main() -> None: speculative.save(str(directory), progress_bar=False, check_weights=False) write_speculative_workflow_metadata(speculative, str(directory)) + 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, + ) + + 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)) + + (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 against `justinchuby/onnx-genai@c9bddd6e`. + +The decoder, VLM, diffusion, masked diffusion, real tiny Qwen3-TTS, +speculative, and codec packages contain graph-only synthetic models and policy +components. They contain no downloaded model weights. +""", + encoding="utf-8", + ) + if __name__ == "__main__": main() From 1b66dcc2c0a1fbc6607a7c3070e48b326df2554c Mon Sep 17 00:00:00 2001 From: justinchuby Date: Thu, 13 Aug 2026 03:15:53 +0000 Subject: [PATCH 032/151] Bound speculative grammar forced-token emits Use the grammar adapter's per-row forced length when appending its token so unconstrained rows do not receive an extra verifier token. Add a producer regression assertion for the ragged emit contract. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: justinchuby --- src/mobius/integrations/onnx_genai/workflow_metadata.py | 1 + .../integrations/onnx_genai/workflow_metadata_test.py | 6 ++++++ 2 files changed, 7 insertions(+) diff --git a/src/mobius/integrations/onnx_genai/workflow_metadata.py b/src/mobius/integrations/onnx_genai/workflow_metadata.py index 6f312b786..b84097152 100644 --- a/src/mobius/integrations/onnx_genai/workflow_metadata.py +++ b/src/mobius/integrations/onnx_genai/workflow_metadata.py @@ -3662,6 +3662,7 @@ def build_speculative_workflow_metadata( { "kind": "emit", "value": "grammar.token", + "valid_length": "grammar.forced_length", "output": "tokens", "mode": "append", "effect_name": "emit", diff --git a/src/mobius/integrations/onnx_genai/workflow_metadata_test.py b/src/mobius/integrations/onnx_genai/workflow_metadata_test.py index a82a2695f..e2cb739e4 100644 --- a/src/mobius/integrations/onnx_genai/workflow_metadata_test.py +++ b/src/mobius/integrations/onnx_genai/workflow_metadata_test.py @@ -272,6 +272,12 @@ def test_speculative_grammar_and_adaptive_k_use_typed_state_contracts(): 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"]) From f20885b5ec4fda7b614584ce12ece7dd93150f98 Mon Sep 17 00:00:00 2001 From: justinchuby Date: Thu, 13 Aug 2026 03:16:02 +0000 Subject: [PATCH 033/151] Execute every ONNX GenAI workflow fixture Replace structural VLM, Euler diffusion, speculative, and codec artifacts with executable graphs, materialize deterministic weights for real tiny Qwen3-TTS producer graphs, and exercise grammar, adaptive-K, ragged emit, and KV contracts. Add an authoritative seven-package Rust conformance runner and make Mobius CI execute it against onnx-genai 5f151fd. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: justinchuby --- .github/workflows/main.yml | 9 +- tests/fixtures/onnx_genai_workflows/README.md | 7 +- .../diffusion/denoiser/model.onnx | Bin 268 -> 1721 bytes .../diffusion/inference_metadata.yaml | 4 +- .../diffusion/text_encoder/model.onnx | Bin 213 -> 1287 bytes .../diffusion/vae_decoder/model.onnx | Bin 174 -> 825 bytes .../speculative/inference_metadata.yaml | 451 +++++++++++++++++- .../speculative/policies/adaptive_k.onnx | Bin 0 -> 34702 bytes .../policies/grammar_guidance.onnx | Bin 0 -> 1926 bytes .../speculative/policies/grammar_length.onnx | Bin 0 -> 394 bytes .../policies/grammar_sampler_logits.onnx | Bin 0 -> 631 bytes .../policies/proposal_metrics.onnx | Bin 0 -> 1026 bytes .../speculative/proposer/model.onnx | Bin 135 -> 1019 bytes .../speculative/verifier/model.onnx | Bin 233 -> 4264 bytes .../tts/code_predictor/model.onnx | Bin 103645 -> 110332 bytes .../tts/code_predictor/model.onnx.data | Bin 65536 -> 2529408 bytes .../onnx_genai_workflows/tts/codec/model.onnx | Bin 146 -> 991 bytes .../tts/embedding/model.onnx | Bin 4639 -> 5618 bytes .../tts/embedding/model.onnx.data | Bin 0 -> 4984640 bytes .../tts/inference_metadata.yaml | 2 +- .../tts/talker/model.onnx | Bin 30846 -> 32472 bytes .../tts/talker/model.onnx.data | Bin 2048 -> 72704 bytes .../tts/talker_prefill_embedder/model.onnx | Bin 21644 -> 22707 bytes .../talker_prefill_embedder/model.onnx.data | Bin 0 -> 4984640 bytes .../tts/talker_step_embedder/model.onnx | Bin 7296 -> 7521 bytes .../tts/talker_step_embedder/model.onnx.data | Bin 0 -> 69696 bytes .../vlm/decoder/model.onnx | Bin 656 -> 3198 bytes .../vlm/embedding/model.onnx | Bin 216 -> 1649 bytes .../vlm/inference_metadata.yaml | 4 +- .../vlm/vision_encoder/model.onnx | Bin 208 -> 652 bytes ...generate_onnx_genai_validation_packages.py | 349 +++++++++++++- tests/onnx_genai_workflow_conformance.rs | 145 ++++++ 32 files changed, 943 insertions(+), 28 deletions(-) create mode 100644 tests/fixtures/onnx_genai_workflows/speculative/policies/adaptive_k.onnx create mode 100644 tests/fixtures/onnx_genai_workflows/speculative/policies/grammar_guidance.onnx create mode 100644 tests/fixtures/onnx_genai_workflows/speculative/policies/grammar_length.onnx create mode 100644 tests/fixtures/onnx_genai_workflows/speculative/policies/grammar_sampler_logits.onnx create mode 100644 tests/fixtures/onnx_genai_workflows/speculative/policies/proposal_metrics.onnx create mode 100644 tests/onnx_genai_workflow_conformance.rs diff --git a/.github/workflows/main.yml b/.github/workflows/main.yml index 124505e59..af9524bcd 100644 --- a/.github/workflows/main.yml +++ b/.github/workflows/main.yml @@ -23,7 +23,7 @@ jobs: - uses: actions/checkout@v7 with: repository: justinchuby/onnx-genai - ref: c9bddd6e + ref: 5f151fd0465c25593e17e824d9680a04adb07eea path: validation/onnx-genai - uses: actions/setup-python@v7 with: @@ -48,6 +48,13 @@ jobs: --manifest-path validation/onnx-genai/Cargo.toml \ -p onnx-genai-metadata --bin validate_metadata -- "$package" done + - 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 + MOBIUS_WORKFLOW_CONFORMANCE_DIR="$PWD/tests/fixtures/onnx_genai_workflows" \ + cargo test --manifest-path validation/onnx-genai/Cargo.toml \ + -p onnx-genai-engine --test mobius_workflow_conformance -- --nocapture lint: name: Lint diff --git a/tests/fixtures/onnx_genai_workflows/README.md b/tests/fixtures/onnx_genai_workflows/README.md index f8eab0434..9e396e658 100644 --- a/tests/fixtures/onnx_genai_workflows/README.md +++ b/tests/fixtures/onnx_genai_workflows/README.md @@ -3,6 +3,7 @@ Generated by `tests/generate_onnx_genai_validation_packages.py` for semantic validation and runtime conformance against `justinchuby/onnx-genai@c9bddd6e`. -The decoder, VLM, diffusion, masked diffusion, real tiny Qwen3-TTS, -speculative, and codec packages contain graph-only synthetic models and policy -components. They contain no downloaded model weights. +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. No downloaded +model weights are included. diff --git a/tests/fixtures/onnx_genai_workflows/diffusion/denoiser/model.onnx b/tests/fixtures/onnx_genai_workflows/diffusion/denoiser/model.onnx index 4c43d9a1434fc502a2423e4464ee5b8c0b3b3717..42970ba7bd4b16331995083b31d41589651eacd2 100644 GIT binary patch literal 1721 zcmcIkOOMkq5U%@5I(=KJwgiEyXoW-%t9IKi5*J{(RpJ7-B4l|JYm?F@A$H6112}Qy zfVlO_h2MysZj&@+kBJp|#_`O2`R1Y5U;b^vZv<6N*TIV6+LnB9Xu!zGXzT^gx20xQ zB5epzVxHYcJ5(Gtpuz{q2I9h6QO=3u zvXF+B^9>h2ewq)!$lOvs2)F!zadiWJA>7!Kx4tJsd4i0Mb}UE8#-@Ha81?|RbhTl` z^#p&7)Y>5}FJwCe30rUk5<>*pqJmr+O^HuQ5!_3&x>tvFDAPNa=StpB1iEl#QUnVm zH;Tpx6RvN`d&Z+gPmI<sXUXhO-8t- zAle&?y<^lPQ=^;GBA`&PWh4JB8y#o;a{B>q$ioho)ZaXxFKlns-oP{}k`{#VDS9 z&){4~cmNu%+)!aTP}IwyGC&H8Cb1r4b7J9Q5D*0b&Y&2P diff --git a/tests/fixtures/onnx_genai_workflows/diffusion/inference_metadata.yaml b/tests/fixtures/onnx_genai_workflows/diffusion/inference_metadata.yaml index 5f29b60b7..01a6e47d5 100644 --- a/tests/fixtures/onnx_genai_workflows/diffusion/inference_metadata.yaml +++ b/tests/fixtures/onnx_genai_workflows/diffusion/inference_metadata.yaml @@ -87,8 +87,8 @@ pipeline: shape: - batch - 3 - - image_height - - image_width + - height + - width role: image stage: pre_adapter components: diff --git a/tests/fixtures/onnx_genai_workflows/diffusion/text_encoder/model.onnx b/tests/fixtures/onnx_genai_workflows/diffusion/text_encoder/model.onnx index 95fbfa02c117e7b02c14d664bab3bcc9da3045d6..7cad175b663f65067d567706847007491169e9ee 100644 GIT binary patch literal 1287 zcmcJOL2uJA6vyk(EzWDJb(Rp7HK8b)IP@^qZU~7B(-4ds{@w5Hxpa1XwaOch60Dd3{5~XLBC9xsFoPC&bE!i>cRs)sA zE{b-_Xwri~@oQ-`7m~Clhwxl~{c@BkwJ;)2v^^46+!!`ZWX^Lj{YJb`pC=*EtzE?# z%!o`&Q|y3Wh;--dU8aoX%CZsNTaC;6M-jfC814hi`PEFaL|JndE|K39)zV3`(gQXC z36!{kTQ;UW$5xC=)N!<1g~JZC`64lc))IM5A^xsc4?q+OvQf;ANpH?Bl)0Hn`CYOT z+Fu7H=1CN+W}1TgiC~k0V5b^)gr~Gq{!5h8s9glZ06d!wP=jjyL9illE-DOHR**cE zN;seA&bV~PU($L=fh2|=*3i~SBt-^-nB3&yNmfevIDMo$&yjIZ(_FC2V@nyYb zBM7li4!|cnm)wde?2|p6XV=;?BVGNKke6@a;TMt@aEK8JA4>)ghv{I$TL9Qy*>)=^ ueOBucvR?2~y4^9gBzIg*{dLxp^!2u73PS&%geV=aOZ%U%Lg#I0^#`|Pw}^89 delta 109 zcmZqYy2=>K!DV%dk;{XNGc&KCv?M+=rC3NbH7_|oCABC%BQqr>H7~xnB(Ws5SV}52 lH?t%jtDF*tCs1KYW=W+GPf2P8T-D^A%tdTYEL;o%q5vIrC5He2 diff --git a/tests/fixtures/onnx_genai_workflows/diffusion/vae_decoder/model.onnx b/tests/fixtures/onnx_genai_workflows/diffusion/vae_decoder/model.onnx index dd4d47b6f1a50e45b52cd31dc3fb40aad0de5b82..97b376218d632bd1aba0daf917238f70b75a2dea 100644 GIT binary patch literal 825 zcmchVzfQw25XR#+f6TQA@Psm;s;U?mN)k3E1Y1@%5VBmyN(`wJ#U_;J;00J%c}AL| zA`-Q1W%GC6-RXDwj0cyWF1#T!tm#dmv|-B7z(4qTt0ECJ7~^%}>8GU8F zBQM^EJ;mlimy8Qy#S=IwpVDce)I#%0mS)OxrZvr^s943bSJJzGh)3Xex|$tS^P&{G z>F830U8-;^Rsa8}qApdml?v^U>H~qBGb5CNR+23hwp2tMIcKTxhqt-pg07Oj7&|Mq zqfE)5W`ZVyFMm~a7p$4dDNHaTBk+hbV}@tY!!}^zW#DYi4d@da4vxYHnspJeaM-;h6%I$t9WZE5r@qmuIGwWH_;K IF$jnP0MeEqDgXcg diff --git a/tests/fixtures/onnx_genai_workflows/speculative/inference_metadata.yaml b/tests/fixtures/onnx_genai_workflows/speculative/inference_metadata.yaml index 97f600015..cfff66939 100644 --- a/tests/fixtures/onnx_genai_workflows/speculative/inference_metadata.yaml +++ b/tests/fixtures/onnx_genai_workflows/speculative/inference_metadata.yaml @@ -14,6 +14,11 @@ pipeline: - 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: @@ -117,6 +122,74 @@ pipeline: name: serving.cache_lengths required: false default: 0 + request.grammar_state: + contract: *id001 + 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: *id001 + role: + kind: opaque + source: + kind: application + name: adaptive.current_k + required: false + default: 1 + request.adaptive_estimates: + contract: &id004 + dtype: float32 + rank: 2 + shape: + - batch + - 24 + 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 request.verifier.past_key_values.0.key: contract: dtype: float32 @@ -151,6 +224,228 @@ pipeline: 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 + tokens: + dtype: int64 + rank: 2 + shape: + - batch + - proposal + valid_length: + dtype: int64 + rank: 1 + shape: + - batch + transition_table: + dtype: int64 + rank: 2 + shape: + - grammar_states + - vocabulary + outputs: + next_state: + dtype: int64 + rank: 1 + shape: + - batch + consumed_length: + dtype: int64 + rank: 1 + shape: + - batch + logits_mask: + dtype: bool + rank: 2 + shape: + - batch + - vocabulary + forced_tokens: + dtype: int64 + rank: 2 + shape: + - batch + - 1 + forced_length: + dtype: int64 + rank: 1 + shape: + - 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: clone + effects: + - grammar + grammar_lookahead: + implementation: + kind: adapter + abi: onnx-genai.grammar-guidance + version: '1' + ports: + inputs: + state: + dtype: int64 + rank: 1 + shape: + - batch + tokens: + dtype: int64 + rank: 2 + shape: + - batch + - proposal + valid_length: + dtype: int64 + rank: 1 + shape: + - batch + transition_table: + dtype: int64 + rank: 2 + shape: + - grammar_states + - vocabulary + outputs: + next_state: + dtype: int64 + rank: 1 + shape: + - batch + consumed_length: + dtype: int64 + rank: 1 + shape: + - batch + logits_mask: + dtype: bool + rank: 2 + shape: + - batch + - vocabulary + forced_tokens: + dtype: int64 + rank: 2 + shape: + - batch + - 1 + forced_length: + dtype: int64 + rank: 1 + shape: + - 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: lookahead + effects: + - grammar + grammar_commit: + implementation: + kind: adapter + abi: onnx-genai.grammar-guidance + version: '1' + ports: + inputs: + state: + dtype: int64 + rank: 1 + shape: + - batch + tokens: + dtype: int64 + rank: 2 + shape: + - batch + - proposal + valid_length: + dtype: int64 + rank: 1 + shape: + - batch + transition_table: + dtype: int64 + rank: 2 + shape: + - grammar_states + - vocabulary + outputs: + next_state: + dtype: int64 + rank: 1 + shape: + - batch + consumed_length: + dtype: int64 + rank: 1 + shape: + - batch + logits_mask: + dtype: bool + rank: 2 + shape: + - batch + - vocabulary + forced_tokens: + dtype: int64 + rank: 2 + shape: + - batch + - 1 + forced_length: + dtype: int64 + rank: 1 + shape: + - 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: commit + effects: + - grammar speculative_acceptance: implementation: kind: onnx @@ -168,6 +463,40 @@ pipeline: seed: seed offset: offset next_offset: next_offset + grammar_guidance: + implementation: + kind: onnx + artifact: policies/grammar_guidance.onnx + adaptive_k: + implementation: + kind: onnx + artifact: policies/adaptive_k.onnx + 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 + grammar_sampler_logits: + implementation: + kind: onnx + artifact: policies/grammar_sampler_logits.onnx + proposal_metrics: + implementation: + kind: onnx + artifact: policies/proposal_metrics.onnx cache_length_update: implementation: kind: onnx @@ -227,6 +556,27 @@ pipeline: initializer: request.cache_lengths recurrence: kind: invariant + grammar: + contract: *id001 + class: semantic + scope: invocation + initializer: request.grammar_state + recurrence: + kind: invariant + proposal_k: + contract: *id001 + class: advisory + scope: invocation + initializer: request.adaptive_k + recurrence: + kind: invariant + adaptive_estimates: + contract: *id004 + class: advisory + scope: invocation + initializer: request.adaptive_estimates + recurrence: + kind: invariant cache_0: contract: dtype: float32 @@ -272,9 +622,44 @@ pipeline: 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: @@ -297,18 +682,72 @@ pipeline: 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: acceptance.length + 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: acceptance.length + 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: @@ -321,11 +760,17 @@ pipeline: - cell: done next: acceptance.done - cell: accepted_len - next: acceptance.length + next: grammar.committed_length - cell: slot_ids next: slot_ids - 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: diff --git a/tests/fixtures/onnx_genai_workflows/speculative/policies/adaptive_k.onnx b/tests/fixtures/onnx_genai_workflows/speculative/policies/adaptive_k.onnx new file mode 100644 index 0000000000000000000000000000000000000000..58de18d70f6992b9663deaece49f999960fe8032 GIT binary patch literal 34702 zcmd5_TaP5kRi5hUneNWH>`Zy>ZOVYhn2&LR=*TFT5D>16ckBPe{D;0$)Vr<(#-w^-E;CU%DbA zD)Y+|=Mv|fI8nQ|hBu!s9v@#UKmWhq+zP%m7~XzXeQ_~guBP)_1cGnor62U-hBFWwmh32&W6_xUJicp?A^B?E#~v*%fs{I zv(=-+)9G?qy>mP}pPnDS^Nr!wSHE`fAQ)+*#67Cz;o@wz6eDz0p^hqaqPo!^sz^r_ zIZ+KfqWZ@{us2(-j-O6fvt__0leI$DiiUg7s<%#$4`)?8x~)FQ36IA&WUHsCO>`nA zP9pX=wW*FOb)pJ9POT=KnT{%RqPiM}OBjo#td(^PmotXT`(?O7Ct~3wVvkc>>ZnR5 zsuw&m+;;}UJH#4D_ljZ)&spAU9WM#nN4}(U@h2iK9(Q;wvnd@{) zMg}F_)}qfDjtyL~lk;Bek@FG*m)kJ-;HFW8oPS}pc>J{)ZMS{i?w zTi?jA84<|0AxMqgYPu>G4?=C74?A zK}~IfsYT1zd^%-;@_*!z1#UAU-!_BELfXkf#$+MuO%|NAqx6a$w9C03@$F$S+T|VX z3P!u)>S;&ml{jd}?M39|H{lba?k!^(*t` z*Dq$XUqj3!9Nq6)DsU2x_u5OoRElOIV#7qlKME!yHcUk9G!a-p^<}wY!>Nd=Q^73P zmPc5>JqUJ>&ZkdS)zf9Lx0;^6Jrf^?yU(g$I6A80baYGnS6r2jucP6MpMDEanx~3| z*lWu-rf*bnHoBueDGb^8hHUnS;g-}3!Ryc+`&#E`i ztGE~i_@~emRkDlSL2#DDeha(o>PUr+y@@Oj1B5)tJiwz7&N zkuCV8^BH>1s|4AC=Cj0=OvXF)1z*@c90d0c7f+uauYj#uE#95Ym+K%Rb-*Ep-XVUE zFXIzKM<-NAe>}LQmhn$c7t__pui@{*!N3>!{)EUoBU{2}HrB}&dneoP4n!!gnXDuh zib-+4RE+bZDoLOi@n7*;0>wD@rWn672tM)T`1BM^c6PovTP&xi)#HmJXlki2C7KqO zP}nUlp|Jho(x?=+Xwnu6R2dEZu%6 z{na|<7HF%<;m5%G)73k(^Dmvwo?`L2tder{3GI?%{BnGkUw?Qj*q%N=Uhex`OHugw zCE%wMJN$HVCGh`^#CpQ;GWbIW{_yJYrK zT36$}pY!W#utiRqSZJA)Vxjof#p5cAN4La(MYwo;9SvV(gC?Vy+$<97Rs}-9dHn)2 z3136}S8(8Kpy3P7AKADpDFrA6MJ}cxz8Q@SWqu<=NUspK<_iu?C9sJYRhA0o|5dD8 zXH|x%k@}>FhNzKj`Xbxk4>+o%5QrNrc4q4aAh<-qBDmxgEF5nx4W$KOI%%b$pwduW zku(&oG$8uaNdux!eM>_sT}vqqQez0-%O`Lx@n3RVK7niLFHX8v8lem721ZPu?;z~4LP?B`BmD6JhL-fCk@;-05;37qpIEB7yX6j7VVjHhE-yo5=Mx(bc@o zM{5Iv5-hr1PYmrd42l4Dv*s4|CPid8-s=~jfjX?Ys0c!x{nfpSR#nnR6!$Rh2MzS zzcQ64f~;!)oh8vx6!u7UQCvx+nj+6-T+xl(k)kQCCPkwOOKC*X$~BQR&%$C|M-;qX zqm=Hikd`teQd5RrrN$b4Ddb{8L21ONj7UJajJyVv#TJZKW>~O_Z^0<_M1AQ$QI9Qq z0-M|_&G(_l(nS&rz*3~@n}uL0QuY28g5TE`f^kS_NB66gDC=wX7f5|vrrvio=rU>4 zVwAKNLo;B#Yz7gQq2%VP%oy2noospUWHVf%a*a#$pC5hMr%QzBacxj_vbFKDlE%wQ zzwxpTPN8f*!UAO}oSiSvPyT?5|f1oKWc&;rZS(mNQAt_F1gs7KCNbgCV zUO!h*MeMu^Cn%R_Hyxtob=M(kJZt~GOt{5MgsW03?l;d>el47&qnEU6HG`*EHmV(p z^yDiJj@v;hAO*5*n2_U~av(|CEESOM`P%uFtA?tyb zugUwtK8){=ToWuJ{qqd@eGVK?IN^dSG^eEN7MXvTm zM-*D?bVRXrcZi5a$N&E2=RR#UtWjsH86Cq=$fk&)har@|_)yag${#IX^!kSu5zKXm zz^39t1edrbYaqsbsW@yMF}cmA>O#5n$L{zr#L~anC69*BZStsIpq5>htx0tpVvkg3 zxI$z@(DF5-x+V!TQv%j5@?Z5S0slgRLb(jrwpuO>Gb`KWF1E}2Vms^>Y}O!QYGt=v zc;%AakSXuj&1+Q#;%m)A5D?3mzD_)$A1sQD}Xaq=V4*oDcfT zmXV?m^tOVb{kVh3l_~^P#J=&>lR_*6b=HGQcW`>D8=?m zH7w2>s>KsdwMd-fRr?O09O|)~h7uRlWBabu{I*F!j;UQZp)Q*&RZx@dbE!gDO(&XU zHL(piL<#oob~7?sUiV+BA{4o3i>2MFA&|+XT&9_a*tdB=C8V3npiAm_9ICbNlCAGy z45Q$6@4N1Hp*mYP3NbYbTHpb>l~+A(4xnyVA=jauEr<)7+l2~mySE>UY;rFFO>E1g z6KZq3;?BcB`8N?;)PDOE+cgoyL>}EF7fM0cRtfY`ce)x-3HN|hjREVC({i+oJF>POA=;M6YNCdCN=S&OokAU zc@q}Dfe(&JpCEXoKneXNE{5QRZKo?jP)3m~pq8LMHwK{&D;!HoE zN)8jgGZ?rJW72BkGg1eni0L0mOMfdhr~-o6*r3@p@!RUECjLOI>H5zEOO(D z`lD8xL-bvt*TsSXk%>b@$cp`)o$bi9kP^zd{gUyoct$lQz#`f}(K|t@r6w}psR``( zq)%sO8i25G)rT~ODB`}?G?b;vlJQP`#g~#Cs4lZFI@sNE^~DFf*Ij+d z!S1%JFFn}(zxwYhKThZQh2YNa&Tuc-8Qyq2T^+vjl~0mWwTw=DwLSZEX!L%@Jblz0 z^fRubBjwON4N7wouQ~Wh6uEK2zs<4ozIA9bhL$`vzsytd^Bzt*pzwqAeUkYZA)mre3;nHL7Mo!$ z^MB(+`#vX@*2!m7>*R!@PfjoOWKJ;c=J~pJ=ER~;E^SZVJ|RQ!qaYTt3{K>Psoh12 za64S|1}85u!bj2YX0Q#}xu$o-ZHKyNQ)&uvHttD*S*gZ=WfJ<@n^>8R-K5xj6Dt+z zO|Y{>5Ya#a6~s@}w-4i#gm-MGBuWn*9ZA~(y|1=##HpcdPkus3Lv#GOp{Seg`aE!< zjHw=nT4)dKHu~ga#cpKePL0k_(H+dlWNsJ`ngck2klGV;7=YFItz+|>iXLK_l&tGi z#A2`8eKTA4&D`wNW<60Hwx`R_L(35@I4jO7jkd7}{$GBqRrw80toM!V_l2ph*TRTJ zi#epG`b2wa1c;$C zHT={O%0Wr;e4J-uh4?wE>uROyjT8G|D}7yqt+aLxw$lBEwno5NB8(v7P=dliZAme8 z)Q6oMHl0>D7}+}3F%x_aagc)nxwO)J04oK#RyF8$$e*DWvQ+dn^h(PrrJ|{^+TeFl z{A}DsdyA}n3;u1z$D6*b@v#szlx16)RPwRRPoq3r4l!;nmxnUhIdP~~GD30ex_dat zDDNRFs3wYz#yzBn@Pst#gdOxGkQBojw<9AQR3#%EYlVHWkVF7QIf5wUY^Uv~EyZhg z@@n^hKh!X=e7=}GlnL5r&q@;1!-=UjURW}!EF5wWSUyiqLFk~Wjltw8;b2%G|FN#Z6m5y_J`B8;g#TZZ5rWVUJVNIy@*Cn95mIx-1(_|PO!6Y6TFSe zHF4+w9m57%)*hH4<=SoFd^OBY=#V;yXq`1|HQMmf@Uu3JSjVc;l%@{SwyH5*N7;Ip zhJ;~Aa3`JEydi)ZBd9#2m(T-nv^L(ScDvR-D|2W9il=tMjQc#$(kar_q_K6T-D*e7 zxC?vII!&OQwJ>w&SazRFP&so@S+4Hd42+e`by%>;<;j`Q{O9cxeji0kQH_80 z!MOdcN-?l&pr`ra5P66CM?}Oo+U$nt=o;c)+V85#Tk6*^&i=btVzi?N-)XWOVpRJX z*V_b#=;#{gY13Y!r)!``gc> zJ{o>@F`qww1yeRXe$;&bm1ag>QNIoLYWV)(r-BF0`zZv=;nr8b_J!@>+P(e%2VDnB Apa1{> literal 0 HcmV?d00001 diff --git a/tests/fixtures/onnx_genai_workflows/speculative/policies/grammar_guidance.onnx b/tests/fixtures/onnx_genai_workflows/speculative/policies/grammar_guidance.onnx new file mode 100644 index 0000000000000000000000000000000000000000..cb4ad9693ecdab7f45483069d859c941cafc93d6 GIT binary patch literal 1926 zcmcIlUuzRV5YL7*xtla?Tx|`fEo~?*1RIhzt>BYZL`qCYR(wt50`Nv%}oXZ+`om9W9TuuHUi;{K=m+_>S=EkUaN2E~qC+ z%~+q=p5JDqR@Nd%Wdrg<>I@j^I6f5{qtTbqHA2k-=pJ<$52(p-s<{W-LEozSo_E2` z&<;e^bSUSfYqOAsX7>`O-@R|{fUy)c&Ox~82aJy+;2VM%GRQiEz@3)$nI|cDza)1t6f-J?-)A0&1xex+ z<%`Y7<7DRpZ(1RBT^bT=V7KK3bLs{T3oC|(in+4|GMieY|MP=r#PJq9&=9^4IlR)M z!j!_0Bj%g5HE?KndHyK&FsNw#_UOGOYVcjSqvdd6o;|Ih^PO4LM2mPG3NlKcMYSHn zgvT5I@0DT-0np^n8IYNKn=x6}a27JMxzyA@f5{?$W$|VCHxs?x*gfCW*i!k1DxZ_( zw`2NN?dl8d_9SgFfvL9jCEBjg0qjn;T-9Qp+&seey=ve)ws}e9$cT`d0;fNHzCtfx x@4u16ZsV=-N;{OY40m{R-{F&aHV^r^8ISyX0a!%Rs7F`U}-(V^IJA literal 0 HcmV?d00001 diff --git a/tests/fixtures/onnx_genai_workflows/speculative/policies/grammar_length.onnx b/tests/fixtures/onnx_genai_workflows/speculative/policies/grammar_length.onnx new file mode 100644 index 0000000000000000000000000000000000000000..36bafd06a4dc6b8b1db3c42a428a6ab6ab0b2ac2 GIT binary patch literal 394 zcmaiv!AiqG5QdxBns%gA_R^YzR4IB3diA6p1HM2-md$pO3?{o_cLNQghrF~`+oru# z#M{7p^UwDW&)8|H6RPdQw+~;?FEDj5FqGwv1tlpe>oSBFQe5d3^1JJk%^Z(mq(q6f z5)xT2UchDbo<&-zwUve{7s*0c%X2~_jLaM6zrMwD2#=_32DegIXtxBLkaC!z;rqwS zX(z-cFdhbM1OXc)!pVF&J2)P1k`NBszph&E7OuLsQHE;c=18lk(gn$eORby{(lP(z sv-O@}@h{0~#V+<#7ZQnjO$7~0MAQ&I{*Lx literal 0 HcmV?d00001 diff --git a/tests/fixtures/onnx_genai_workflows/speculative/policies/grammar_sampler_logits.onnx b/tests/fixtures/onnx_genai_workflows/speculative/policies/grammar_sampler_logits.onnx new file mode 100644 index 0000000000000000000000000000000000000000..f674183330680e7f8e3468568778f80f8327ea0b GIT binary patch literal 631 zcmcJM%T59@6ox4ohv5KGyHLPr5Q#)$CJ=Y7yev$31dVArr9fq9$7zQ^j2m9V#Q1Q& zhYk#=kQi5X)ysGOd_Bj>VSA`N5$W@H3tkY;kI1!>+E8gov%LAWm#V`_v+67_8YRe% zs2_16q$z$Z7kWeKb^)@I4!I5}<2co>!9g(Sxk}2hW}yg-%Y3Re=?fmxko6}xeRtok zKzWUtWf09&!1auPR|K}NdctT}vBbGJ&XDA(T)3g8&qFd7ZdL8%*vNrJ$3hqXQ4nv> zJ&YP~xv;rCTPouds5=NBK?ZMl)UZAPb~JyCxMUpi7)FXwFY;+P@lN6AkMKV1#g)6? zl`Hha&RREYu$&A|O;H^VlhL|8F5^}kp`;rqU$6;bN}7-|gVRsXQ*;TnKS?4wxU`^> g5SK<{;fvV#qBTXwQ28~3ce|XK`4q1#NNp8A0X>A#xc~qF literal 0 HcmV?d00001 diff --git a/tests/fixtures/onnx_genai_workflows/speculative/policies/proposal_metrics.onnx b/tests/fixtures/onnx_genai_workflows/speculative/policies/proposal_metrics.onnx new file mode 100644 index 0000000000000000000000000000000000000000..0084a107f34b0d3c63c7a853108d6ea7a925af8d GIT binary patch literal 1026 zcmcJO!EO^V5Qei$)9g+YwX#%51FA|?QPIPy3pY-H)Jv}1ijd`aCyB|%c5QE!2=$T| z-~~7!UXWM7*@duaRXF7~zvlPMH)A{6BfB#-6oq~Bvjsl_*w;pBWig{pO;K7B-_Q>e zrjcIs+EpM`=2gF0s&jxxD6?+xJqXX+(6o*j*KK$;=s+yl46SAyNn>yd$9g(Wl$3MJ zjnFRPIkT2d1R7@eWI@7@p9Uw;yW?8bpq49*cIEgtWV!*F-hm7^AXmuSg8Ul5dU5IB z@Pfp6!}5YT%zB;6xn?q>UUX_qyLK)t2zc!Mgbjoe|t#CHLRrQ#guRT;zh8#oKflynvGjYv5>v+?7A+{SGt^ Gy1xMaltV@U literal 0 HcmV?d00001 diff --git a/tests/fixtures/onnx_genai_workflows/speculative/proposer/model.onnx b/tests/fixtures/onnx_genai_workflows/speculative/proposer/model.onnx index 77cc6632345686ea3719755becbded28b5aacab7..9a0ae971283481e5c99debc21e1aa798967a540e 100644 GIT binary patch literal 1019 zcmcJN&rZTX5XM<7bQvJP8i6#?kj8lEfrxtYWb`B#58g0lvy>GIxNW)%!IKZ*(HGz) zd@mot{;`V?uQX{gGy8Ss*KY@xufI)rhobUlTu4${(#Hu4I6b#2X(0Sjl-Y$O>sO&9 z*ejO-b2-WNPvIn(jk~@OE9r(_pt|myNl7Oj4_W9=-bnuGx!(b6M-}&=()9x_V+DL7 zQd!bFUr5D-qCKmo<79r%##zB>04&*j!Ko*dyy>iwxi6|Uf{h=r1xO&lb(x|URyAUa zqXhNBE~cqd27|3U+0Ly|eqV@RC~_k5$ zG@+2V>HtP{)Hia`vqr|gxId6yxnE{JqUjD=wfhkl%2RKtp;tD}3VWgy{_8b#GM&CH(t;z5NJICE mXGU*l+WNDI1OP`9aBD!W-u?l5C^T;X delta 24 fcmey(-p1VwvjE3_}sLm!|I(D!L~mlVyVtRo!?5O;Qmvmd{iof&GIU;O$h{DEN0 zo4KCr(=iF$bL#kbV@{r(*&Zb=y&Mi4cg!TCAS4qYPo42DS|G(NL%|R1S+JX9(y#&> zj{S=Io;{*CH~b9ldgrJ6uH#(#quIm@_D2`C?~}6$o!PU|*&D39csXo>p3>qrh(B^Y z>T?bF5#jQjeC0Z9D<>e9z7>xJGRx4G6J`xI=Jv&FN+wRWpxz_4(~>BF5l2U?F=^pV zHvQdeO24P1tyh;rfwXOuBKd~6>ZG~qtmBHiGb>OF1d;g8}C znClt^A!P+I6=S_LZ@mw}+YchfB040Vdbl(v565G2paZ`Njt-1G8{ZH^{WL@UbqwV* zjI4JEmPAPW^f^D@OdN8k1HTD=4vjn;XT{IY0*cr(kq#_enUilFKS~G#y>=xCPy=Hl zp2-T@d&E&Jb0o1^dMy?N2&=Wud7otHyq^%ruo4-SXVJ%(p6!f@W$D#u7=pEoVkl?D z*KdetsCawHozgbL1T*z?8*h(#!o1og46JM15tawkyX?G<~yX)e#nPHs%o;p6g6=7yt3>Seo2wtUm-hLXzBoQyXPM%N9vKbK&7l`{!{Atq}qTeMwFS*j_-kd{BZ;7?9uJyqbk#E9zOt?WBfUoOmi=(slX22 zmNseLZ+7`c)7?bhv@U&9#M82RC!Vi! z`%zO>L- z@2*sM+2Z}{D+0bQoTjPmyYcpSWL!PrBqZFYOjC>fga`A}k2vtFBy{CNb-K6aBv+x# Q!b)S5h~Fa>a<$$60X$wF`2YX_ delta 31 mcmZ3X_>z&IgWKvJqYy_~YEfodW@^#osk~R%oLIOR1S9~PaR}@H diff --git a/tests/fixtures/onnx_genai_workflows/tts/code_predictor/model.onnx b/tests/fixtures/onnx_genai_workflows/tts/code_predictor/model.onnx index 9bbf0d6b7c2917052bb237ca76a56607e5203cf1..8b27a601fd3723e0b01e5ed1445b2289dc493d94 100644 GIT binary patch delta 3531 zcmbW3PfQ&}9LJqf_HAF`@t`7)KZQqWt9iV_{M*@Glag2;9x4UYq^2dN)e5CZp$$;c znD98*#EZ6>#EFTa^ytkV)I$?XJQxxi#cLB1F5ZlZcau81%Pcc}``W^u_xESM^Z9*$ zv-zF%-=@`OA8fL&&VF!ZVXmzf_Mtxb>hGIY`tN<|9)7RCP0W_Sf#s7woCKLtdd9A? zE!)n5Q{v7Lcv&o#KxZ+Nwq_?L=Emon_B#l}srs}vJw9`JezK{?aoIIcn24h~#p)1f zNljn9FggDIhhg{;J)HThGJIMLmB8r+?LUTHEE)&Fi{i%;c(=xfRNO6rW^rr~d}Eyr z`{Ur(7T(p~gg+j?a(Qz8bpC~Be1%W#ayANWZ0+I0%_;t1)Cybd%f=_Pi0@O+n`(&9 zSj!xT1gHGOF54;|hH(w@uWe;&UF1QL9R$q>WS#iG9$r4261UESs8FL;2tg{RTHezCWfXD>y$yX2qS&D6+iGGq<{#lM4~ zn3gKoQitt2yw6SXnQKvkqShQ`4#YG`g$tnsE?;iXiT+X0b>cacU$o1)Fs#V`e7%D| z7)T3vc3TqI7`L55(4ihR#P>Y{sXX-!m7*Y19JQiL4==$;fEV*qdJot^)L}^}cEA(=Oo2U;*CJLw0*4aOa|m=8)GI6Wbfj81yp-W*9v%!& zZxBUXLTN~p0bvkP$Ma*=Ky)EO(YU4NW4OXTGvd7w5EXhw8&4FlQ30kFD+?Zq!7S*J za+E3~btc1R#cyv%{CDb5IzEHaKN3b$L3#UPj*r}E6{};w+9jCFOC4hL9gwYLicsG* zuSGGXLAktC9ZE_)$c$l92{b5!`Wvdg0i;72VTnVb0eNC!&p`66l{#@}0VMk&V~7!B zqLQK9TG@dzt2qgV*z-d7eu7C^m2Z5lF=<2*htbeZbtJ=1l{I$iL}DmyHrZN+i<7E;6r1iKB>o?b{rGd~YvAisjG1@hU#aC}pm3DT*i!Np%0Hm{k^M zxTH(5L|Pvt6|Mn^?sxF5U#hvJM-lc6uTm*uDEo`gu^RNH+`{wgt>OIX3)%It3ytYf z8k?@T_9(;44{OERaKf_1Og0=X+b1y0uh;RVzcdVWMydpUXi_TG^iCVQSon(rKGT(R~x3Jy@f-eNzXqUtg9!@q@lK zOI8@fAd37s&|BfE7!ubr-&rbhhe6SO4s;0k1*k$HMuZ_haSkF1MdLoOdkh8K4=OqI zF>%eIQTiGjij9{*M%=PN6%q-Nl11CrBr&L3C%);jiDX8|RFJ+(an4nzuR z2!H_yKo9{!Ba^W&8kbFH0ST8-MFAd{?MML)moriU6PJ-i0apYvH#RtzF--v!mPP@Y Tmt;r*G?#!$0V}sWU;)5RVSFIT diff --git a/tests/fixtures/onnx_genai_workflows/tts/code_predictor/model.onnx.data b/tests/fixtures/onnx_genai_workflows/tts/code_predictor/model.onnx.data index f6d131792b339d592bd38ae7d51bf16aa4ebbd5f..cdb0d2e4321f773d3e97c3fd7835ccebe02abe60 100644 GIT binary patch delta 41078 zcmeIwu?>JQ5CcF3F#uv>?~pDf8)PIlU;@GxFN!N%TedFQGe^79`jMv~=iJ`aTUJ3= z#U33sZv5kDL)D~AOo(n>*+9s0jY|Wq6UAuuc{J3)L$Eo99s}s03N*f=U;XMn=yk5f`-Qwmyn7?4DKCBm_16^-8jQ^)N*O>;#sq%@6R9Dce-c3}rQWE&h} z51Hm;FtG)gtc}v*gAb2j&Ra@RRf-nUTesqNmjK{=vz0GIwOR1XEUe7@W#_&@UAli} Ifw9?s2Oq90T>t<8 delta 108 zcmcc5K8ewggUhO&kxPqewOELwJh3b_Ex#yNiZ?Ykvm_qIRpRhWNzE(CEUBFQ ypUJ>biwmw$NRCT}gHcF=i(QB{DX}CugNuoSk&9D^y*M$qASbohiG_Gaq~RhMiv0XK26U6 delta 41 xcmeyQJzqtfgIkC#H$N$}v{;JOK+jOm!0MAS*A=G8zCwx{Ez)>57Ya7A008=}3>E+Y diff --git a/tests/fixtures/onnx_genai_workflows/tts/embedding/model.onnx.data b/tests/fixtures/onnx_genai_workflows/tts/embedding/model.onnx.data index e69de29bb2d1d6434b8b29ae775ad8c2e48c5391..885cdd9a062995e9b753010d104a5fb72e2baa58 100644 GIT binary patch literal 4984640 zcmeFtfdBvi0Dz$VsTV1P3IhfV7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd z0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwA zz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEj zFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r z3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@ z0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VK zfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5 zV8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM z7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b* z1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd z0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwA zz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEj zFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r z3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@ z0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VK zfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5 zV8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM z7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b* z1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd z0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwA zz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEj zFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r z3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@ z0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VK zfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5 zV8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM z7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b* z1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd z0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwA zz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEj zFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r z3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@ z0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VK zfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5 zV8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM z7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b* z1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd z0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwA zz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEj zFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r z3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@ z0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VK zfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5 zV8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM z7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b* z1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd z0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwA zz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEj zFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r z3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@ z0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VK zfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5 zV8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM z7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b* z1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd z0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwA zz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEj zFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r z3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@ z0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VK zfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5 zV8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM z7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b* z1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd z0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwA zz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEj zFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r z3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@ z0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VK zfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5 zV8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM z7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b* z1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd z0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwA zz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEj zFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r z3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@ z0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VK zfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5 zV8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM z7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b* z1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd z0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwA zz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEj zFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r z3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@ z0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VK zfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5 zV8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM z7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b* z1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd z0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwA zz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEj zFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r z3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@ z0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VK zfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5 zV8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM z7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b* z1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd z0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwA zz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEj zFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r z3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@ z0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VK zfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5 zV8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM z7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b* z1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd z0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwA zz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEj zFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r z3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@ z0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VK zfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5 zV8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM z7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b* z1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd z0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwA zz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEj zFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r z3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@ z0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VK zfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5 zV8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM z7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b* z1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd z0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwA zz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEj zFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r z3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@ z0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VK zfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5 zV8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM z7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b* z1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd z0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwA zz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEj zFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r z3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@ z0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VK zfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5 zV8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM z7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b* z1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd z0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwA zz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEj zFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r z3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@ z0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VK zfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5 zV8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM z7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b* z1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd z0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwA zz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEj zFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r z3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@ z0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VK zfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5 zV8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM z7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b* z1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd z0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwA zz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEj zFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r z3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@ z0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VK zfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5 zV8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM z7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b* z1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd z0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwA zz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEj zFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r z3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@ z0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VK zfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5 zV8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM z7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b* z1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd z0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwA zz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEj zFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r z3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@ z0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VK zfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5 zV8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM z7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b* z1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd z0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwA zz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEj zFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r z3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@ z0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VK zfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5 zV8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM z7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b* z1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd z0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwA zz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEj zFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r z3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@ z0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VK zfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5 zV8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM z7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b* z1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd z0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwA zz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEj zFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r z3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@ z0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VK zfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5 zV8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM z7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b* z1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd z0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwA zz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEj zFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r z3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@ z0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VK zfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5 zV8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM z7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b* z1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd z0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwA zz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEj zFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r z3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@ z0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VK zfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5 zV8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM z7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b* z1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd z0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwA zz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEj zFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r z3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@ z0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VK zfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5 zV8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM z7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b* z1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd z0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwA zz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEj zFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r z3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@ z0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VK zfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5 zV8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM z7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b* z1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd z0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwA zz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEj zFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r z3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@ z0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VK zfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5 zV8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM z7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b* z1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd z0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwA zz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEj zFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r z3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@ z0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VK zfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5 zV8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM z7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b* z1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd z0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwA zz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEj zFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r z3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@ z0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VK zfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5 zV8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM z7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b* z1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd z0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwA zz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEj zFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r z3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@ z0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VK zfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5 zV8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM z7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b* z1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd z0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwA zz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEj zFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r z3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@ z0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VK zfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5 zV8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM z7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b* z1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd z0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwA zz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEj zFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r z3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@ z0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VK zfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5 zV8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM z7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b* z1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd z0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwA zz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEj zFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r z3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@ z0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VK zfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5 zV8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM z7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b* z1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd z0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwA zz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEj zFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r z3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@ z0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VK zfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5 zV8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM z7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b* z1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd z0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwA zz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEj zFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r z3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@ z0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VK zfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5 zV8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM z7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b* z1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd z0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwA zz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEj zFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r z3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@ z0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VK zfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5 zV8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM z7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b* z1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd z0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwA zz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEj zFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r z3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@ z0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VK zfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5 zV8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM z7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b* z1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd z0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwA zz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEj zFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjxE(qc00000802p~iJgUz z0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwA zz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEj zFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r z3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@ z0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VK zfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5 zV8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM z7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b* z1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd z0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwA zz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEj zFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r z3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@ z0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VK zfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5 zV8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM z7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b* z1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd z0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwA zz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEj zFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r z3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@ z0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VK zfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5 zV8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM z7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b* z1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd z0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwA zz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEj zFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r z3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@ z0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VK zfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5 zV8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM z7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b* z1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd z0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwA zz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEj zFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r z3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@ z0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VK zfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5 zV8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM z7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b* z1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd z0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwA zz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEj zFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r z3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@ z0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VK zfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5 zV8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM z7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b* z1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd z0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwA zz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEj zFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r z3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@ z0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VK zfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5 zV8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM z7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b* z1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd k0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|wrK028zT0RR91 literal 0 HcmV?d00001 diff --git a/tests/fixtures/onnx_genai_workflows/tts/inference_metadata.yaml b/tests/fixtures/onnx_genai_workflows/tts/inference_metadata.yaml index 2a8fd5ab3..cb0219254 100644 --- a/tests/fixtures/onnx_genai_workflows/tts/inference_metadata.yaml +++ b/tests/fixtures/onnx_genai_workflows/tts/inference_metadata.yaml @@ -205,7 +205,7 @@ pipeline: shape: - batch - 1 - - samples + - frames role: audio stage: post_adapter components: diff --git a/tests/fixtures/onnx_genai_workflows/tts/talker/model.onnx b/tests/fixtures/onnx_genai_workflows/tts/talker/model.onnx index c5504386668a145f33a20d6a4dbda4b85037c572..9c99d083000771135e7a1c2a3887624a2e93f168 100644 GIT binary patch delta 1287 zcmezOf$_#)MiCBfA-3H7q|DM{DOLkLLp=klRev{%{7INxlqfTKgCOVR>4{p*rp9KI z;}a8Ej7-f8HXliJXXMcGWq{28okUQqcTj z_R0Iw*+zOQ;sqxJt#EJ>(w`hxqDg#~Oy>{)Yg3u5;4Q}n$;rSxFQ>-M#g?CzR-9TQ z#AuMk%f*(HnwMUZA;fHIXjH(c1=BXU&|ev(Tob4QsvM*OQ$rRg;{$U($R3f&zF}Ml zyY#{^jJE>2NvI?-Cp)!BFE76+SC8mICXV9+#7mRs#fbskrYUZ0YGIPa5BCDAnWdqT MK>_1t?&M-o0E*Rf5dZ)H delta 60 zcmV-C0K@;-{Q>^)0Tc)e5(aH=Vrg}A8U-*eF)lDV(D1Po{9plKlM!J;lh%2(h86Ds diff --git a/tests/fixtures/onnx_genai_workflows/tts/talker/model.onnx.data b/tests/fixtures/onnx_genai_workflows/tts/talker/model.onnx.data index d3d841223b6dcd164d80d6b7710ba7835fed5123..2869428927c03c8e598c47fb258bebcb1babe593 100644 GIT binary patch literal 72704 zcmeIzi(k%nAII^!lNv2)-8MtirIt#yZYgwq-%ZIOsg@2Z#gxvabE)gPoR&RsJIryI zTQ=mdlL*uG`&u@PkX8)$S`I05$hu2S#oYJfasL&2zyE~im(#Yj^bZ`rMM{c^PgOm3LG4x9vPF%`^?qmPvDDTx1sFA`|7Q#f$KpsgcsQ zB@<)YePQ@T29D`+(N&d>v=2WL*UzRQx5`oOzmN*Q>Jpi8J_St|Y83a=$!HyBj-GFl zaH>;>%bzAorXUA_<*{=Xr=xuEe_@__R7LBvDj>tDjzq+U|-;DS@Ur;oP3XYn;{B&(xOn=9Erno7ox5-0>`a7 z)Ty@NSdlPAt~@at;a78{`}9!Uw_Bm))`y@g#TgM%!8oy?1ov*v!t%27!ah6@&wB?+ z`t-xgcw$PLzM|4C5g#^5;fJ3^yo|RUy*!6LWYCQ>lN?H^Y2LBPt&K z9Zcb5R;v!~F~&5XY0~MU5q?j|mixCFqFbM@eD>W7v03km?7-(DYiluX)OCuBlTV9h zA-{{%$~WZg3y(y-)XCH8Rw4Hp$S)4w73Xswh#v<3DlY!0g+tCw(X!1N|GC^CR@d%R zrr2E-e%}^Kb$*S=jGZ96zdI-F4!tL~n4S?czfQ;a8H#Ys2u4@wNudp$s_#~giky3M zWmx-H;z`H?>0xt7CiTxTugSwsoodyn>!McaC$hI`5P#2H_!ZeB zsQ*wgIBBC2JX$5}tn;Nxj1?6=K5}KKr${auDkjFiC(c(UqdPM|Sh$8`d48DqyYb+v zEd_Dnep#e^nw>0qe@c`gNjbu6rHN7yp%EsYzOdN6T$nfH!g#|+!tv-Tq0#goI)!SJ6Jp!{5W6>)z%ucLsHnUk_M{l%lGz}cvdah^Z93U-*%+}ajAdPq3HpM6 z6`^(m(BiDc%m6d^8mjO`x;aK$>{OE0SU?+EAhY&cqJ5%To>Z(55a%S`ts8>+(o|f! z{T2ou3We_R+qkmyM%9GQ;fOJcmTleE*xHyNH+7Ffpw&bD+a0gkAK{Ojhoccvn1fA? z_E2>$6`7YD@YrXp>^kms*lW$N~;?-LSSvhld{Scyp5xtOtxk z)dzQ#${QZgWNPKU&pi?0V%?wM^}qVr%=4}qopj`F&oA;I$2{J4wo`R z`QJ7kXSRqxevCx-VJ%LpqTpIG672=i_~y{(O6l1c7~~booI$ZLvY#lwem@Rro@z0E z{X9JPO-Jv!`4~7P1U-LzfGeIJdPnrP3Q&Lo6rcbFC_n)UP=Epypa2CZ zKmiI+fC3bt00k&O0SZun0u-PC1t>rP3Q&Lo6rcbFC_n)UP=Epypa2CZKmiI+fC3bt z00k&O0SZun0u-PC1t>rP3Q&Lo6rcbFC_n)UP=Epypa2CZKmiI+fC3bt00k&O0SZun z0u-PC1t>rP3Q&Lo6rcbFC_n)UP=Epypa2CZKmiI+fC3bt00k&O0SZun0u-PC1t>rP z3Q&Lo6rcbFC_n)UP=Epypa2CZKmiI+fC3bt00k&O0SZun0u-PC1t>rP3Q&Lo6rcbF zC_n)UP=Epypa2CZKmiI+fC3bt00k&O0SZun0u-PC1t>rP3Q&Lo6rcbFC_n)UP=Epy zpa2CZKmiI+fC3bt00k&O0SZun0u-PC1t>rP3Q&Lo6rcbFC_n)UP=Epypa2CZ@K+J| E6PZJli~s-t delta 8 PcmZqJ!O|eGh>ZgP4dMcQ diff --git a/tests/fixtures/onnx_genai_workflows/tts/talker_prefill_embedder/model.onnx b/tests/fixtures/onnx_genai_workflows/tts/talker_prefill_embedder/model.onnx index 27a7b8a0d36a6aa8eb99138a75ca0e4817234d04..c3dd371ebaf2533c28bf1d122fe404ecfb6f6bfc 100644 GIT binary patch delta 887 zcmeBK$+&qVqZ9|X5L<43Qf6tf6sv)rp`L+N>qbVd70g_^lNm+SC!ZH)-{_(5%3aOz zYXK(*hXA7!Oq|U>OPY%#CqFr{Br`uxh(9+!B{fGcKQFIBFD0=gF-w4pEk7--IJHEG z&CuA;z}zTHkc%xRH7~s+Lx|nP!qnK*%%Xr%YcCTA5CF|Jn#{_s2G4k#V{!|vDixiERKgj4$DGio%-_B%=C#H2?qr diff --git a/tests/fixtures/onnx_genai_workflows/tts/talker_prefill_embedder/model.onnx.data b/tests/fixtures/onnx_genai_workflows/tts/talker_prefill_embedder/model.onnx.data index e69de29bb2d1d6434b8b29ae775ad8c2e48c5391..885cdd9a062995e9b753010d104a5fb72e2baa58 100644 GIT binary patch literal 4984640 zcmeFtfdBvi0Dz$VsTV1P3IhfV7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd z0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwA zz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEj zFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r z3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@ z0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VK zfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5 zV8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM z7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b* z1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd z0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwA zz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEj zFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r z3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@ z0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VK zfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5 zV8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM z7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b* z1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd z0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwA zz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEj zFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r z3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@ z0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VK zfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5 zV8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM z7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b* z1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd z0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwA zz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEj zFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r z3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@ z0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VK zfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5 zV8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM z7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b* z1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd z0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwA zz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEj zFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r z3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@ z0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VK zfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5 zV8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM z7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b* z1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd z0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwA zz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEj zFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r z3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@ z0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VK zfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5 zV8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM z7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b* z1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd z0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwA zz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEj zFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r z3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@ z0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VK zfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5 zV8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM z7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b* z1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd z0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwA zz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEj zFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r z3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@ z0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VK zfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5 zV8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM z7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b* z1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd z0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwA zz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEj zFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r z3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@ z0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VK zfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5 zV8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM z7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b* z1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd z0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwA zz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEj zFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r z3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@ z0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VK zfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5 zV8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM z7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b* z1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd z0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwA zz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEj zFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r z3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@ z0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VK zfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5 zV8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM z7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b* z1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd z0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwA zz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEj zFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r z3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@ z0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VK zfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5 zV8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM z7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b* z1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd z0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwA zz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEj zFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r z3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@ z0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VK zfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5 zV8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM z7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b* z1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd z0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwA zz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEj zFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r z3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@ z0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VK zfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5 zV8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM z7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b* z1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd z0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwA zz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEj zFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r z3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@ z0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VK zfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5 zV8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM z7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b* z1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd z0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwA zz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEj zFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r z3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@ z0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VK zfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5 zV8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM z7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b* z1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd z0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwA zz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEj zFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r z3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@ z0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VK zfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5 zV8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM z7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b* z1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd z0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwA zz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEj zFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r z3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@ z0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VK zfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5 zV8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM z7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b* z1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd z0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwA zz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEj zFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r z3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@ z0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VK zfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5 zV8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM z7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b* z1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd z0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwA zz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEj zFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r z3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@ z0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VK zfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5 zV8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM z7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b* z1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd z0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwA zz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEj zFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r z3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@ z0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VK zfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5 zV8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM z7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b* z1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd z0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwA zz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEj zFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r z3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@ z0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VK zfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5 zV8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM z7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b* z1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd z0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwA zz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEj zFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r z3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@ z0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VK zfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5 zV8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM z7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b* z1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd z0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwA zz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEj zFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r z3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@ z0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VK zfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5 zV8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM z7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b* z1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd z0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwA zz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEj zFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r z3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@ z0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VK zfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5 zV8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM z7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b* z1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd z0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwA zz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEj zFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r z3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@ z0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VK zfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5 zV8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM z7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b* z1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd z0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwA zz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEj zFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r z3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@ z0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VK zfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5 zV8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM z7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b* z1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd z0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwA zz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEj zFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r z3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@ z0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VK zfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5 zV8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM z7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b* z1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd z0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwA zz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEj zFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r z3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@ z0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VK zfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5 zV8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM z7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b* z1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd z0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwA zz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEj zFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r z3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@ z0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VK zfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5 zV8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM z7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b* z1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd z0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwA zz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEj zFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r z3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@ z0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VK zfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5 zV8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM z7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b* z1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd z0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwA zz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEj zFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r z3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@ z0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VK zfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5 zV8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM z7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b* z1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd z0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwA zz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEj zFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r z3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@ z0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VK zfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5 zV8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM z7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b* z1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd z0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwA zz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEj zFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r z3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@ z0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VK zfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5 zV8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM z7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b* z1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd z0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwA zz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEj zFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r z3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@ z0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VK zfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5 zV8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM z7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b* z1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd z0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwA zz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEj zFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r z3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@ z0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VK zfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5 zV8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM z7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b* z1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd z0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwA zz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEj zFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r z3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@ z0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VK zfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5 zV8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM z7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b* z1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd z0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwA zz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEj zFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r z3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@ z0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VK zfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5 zV8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM z7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b* z1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd z0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwA zz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEj zFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r z3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@ z0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VK zfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5 zV8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM z7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b* z1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd z0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwA zz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEj zFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjxE(qc00000802p~iJgUz z0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwA zz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEj zFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r z3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@ z0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VK zfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5 zV8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM z7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b* z1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd z0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwA zz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEj zFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r z3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@ z0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VK zfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5 zV8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM z7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b* z1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd z0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwA zz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEj zFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r z3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@ z0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VK zfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5 zV8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM z7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b* z1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd z0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwA zz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEj zFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r z3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@ z0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VK zfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5 zV8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM z7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b* z1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd z0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwA zz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEj zFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r z3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@ z0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VK zfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5 zV8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM z7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b* z1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd z0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwA zz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEj zFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r z3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@ z0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VK zfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5 zV8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM z7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b* z1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd z0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwA zz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEj zFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r z3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@ z0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VK zfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5 zV8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM z7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b* z1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd k0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|wrK028zT0RR91 literal 0 HcmV?d00001 diff --git a/tests/fixtures/onnx_genai_workflows/tts/talker_step_embedder/model.onnx b/tests/fixtures/onnx_genai_workflows/tts/talker_step_embedder/model.onnx index 58e017b452c5aa41993a77a5b7e8bb5feb777368..19ccf3900c5ecda91b68bd5aaf811986f7d77acd 100644 GIT binary patch delta 201 zcmZp$d}t-Y!7ap=o1c_fS}etCpl7IOV70|+@&YF5jrwcExN32BPTyOu_QA;Pl!J^KP5FsFF!A@LN6t;Br%JZi!DDbtvIzrh}qQKEQ_Cu zEhjZEy(B}3)y&e+$e@5xtB`}4gAHt^4%kc^E^%x8W^%*LWHbO-05uV0K>_3DTVl&t E0qaIL*#H0l delta 38 ucmaE8)nF;Y!7ap=o1c_fS}etCpl7IOV0GDI@&YF5jrwcEHt!K%$_fD6%nTL) diff --git a/tests/fixtures/onnx_genai_workflows/tts/talker_step_embedder/model.onnx.data b/tests/fixtures/onnx_genai_workflows/tts/talker_step_embedder/model.onnx.data index e69de29bb2d1d6434b8b29ae775ad8c2e48c5391..3faa8a536d481e35e98281cfb05b9b76304cbd76 100644 GIT binary patch literal 69696 zcmeIufdBvi0K=g9Qy<|1g-~I@fB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM z7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b* z1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd z0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwA zz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEj zFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r P3>YwAz<>b*2L8YR5KjOB literal 0 HcmV?d00001 diff --git a/tests/fixtures/onnx_genai_workflows/vlm/decoder/model.onnx b/tests/fixtures/onnx_genai_workflows/vlm/decoder/model.onnx index 59df2e6aa626e5aac02f9b52bdb3a967c5b9f925..cb5913f3ec6f24cf22c7f07e914bcd69006ae50c 100644 GIT binary patch literal 3198 zcmchZTTc@~6vw-zrOTl=#sFoF*i=zj64C`MmuMmoUz!*beXB8(?d}NMmhIYE8hGO~ zsEP57cYQFv`MLZ8o}H!a1&aExwrM)IGxN)t|2cEy>78E{IG3Pc+0L=cDLZO2lar|- zeSN4qjMkOB->~~8_G($tQvoKJZC0dnDXr!q%UyloR#Mn&7GTQOj~I7!gOOD8J}f%j zgVnxmpKxPfIqs^_(>bSy78~dTV0hIs@_1-a~mhJLqXNZ+c3hS!`!EFXj4I;1v z1iGQSbX_SRr`HnL>$OS0t4>2kKe2eC_xvL+MF7L|DAi0_C&eLs#mr-ayS5HOpP(B%{bNOSSmR|s-KY<2?s`FT*PdKB4#TZFp0sP{@Hh5KEC`K!2sI|N(`fW=e{`RXf_ zATr6B%)okQpeh&!xxLU6)Rfa+OyH+gj~7jM;f^ek1;~+XTX&5^03vht{+QVYgS1?2 zJ%ys~x_B#A-=;@8?~)Z*mUCn&Fi}WF_-bTVU8#-4dbQa>ut^Kk0_V=5#Rk(QmNEH_GU?n&>gbNe$nbh8a zsBFnHJVcXwfr-ajCPlaJ>b(eT`w7Gq^bgjZN~FvHkkPRjbb{2G4<3LzR&Cc#Xpspe zURi~V@SIR#5NtQm(c&cp`GcIKPV}J04G{QPje*z-&15=D*Cstuggjn@3@QJM%dVtB Y0%cs1_gpsmAzzS>Vq5HGAXTpX1_aE$QUCw| delta 208 zcmew-F@d$7gUhOwl}nF{H#4uGw4^vbH8&|WrC5k9CqF&2q*#hOH8-;)9>P)L@Jva~ zE6FUWqRrNvXBgR>ShyGjL;1Jrk}6@x0qc<3 zlzouub|S=A!8){_DsmL)70+F*bZg*BJ`^hEz95C}0X&XB_m4(W&Xhk63Qx7i}+6{myznKUY z%0#_v%~3fos<}?hDj%>0NWkLyJYi?FrrC^SQO~^6LTy)|$Y-Hy7ILiKPXs&RkoN_< z7Cf1Zb(if-**g(TeDPLr$|RIg!7fImk(xWW2o7%bI^d8?hR=-^)Y3Dot`tL7}rNBhn+jNSunQG$ZGFN;v9H~#>TaqP;?OOA`?f1y*B)& zhIObJIVYP)MWQt9z)|ur?#9NftV(r~b=6E-ixW+$Ox5H`ly7c_Ap`~0x(8YnCyu}= z;FSq~i5Ic4nKT*pgXZR(L-gDHCvkEBkh7yUnH{i!HwZsibz5Zb3%()HxG1D02@g{H-<$fR9rh zVK#Xb4e)tYaCgbvs@!#;@EIYX<;wkUh3EEg8Ws)7`>;(PHVq0E0dV)uy1H`iL8*24 E2@X5KasU7T literal 208 zcmZ9`Jr06E5CvcXm0@BC8Ic$(VqxJCOl<7D0?me>5jNp3?1I{P3@_xpL_2L?UcUDL z4=)3cxM<~#OEJV{&x%}G#W{-$$GY!WZI)a~Eq=o-oHMSq)OT_Rsbo_NuG1#tF53NZ z0s3H&hYs0otkZ%OTJUtD5FzhmCUkKn9|>%75orG1QD}md?D6jktuYuaLYrpE#~P)F MId;R-6er&54O$>L{Qv*} diff --git a/tests/generate_onnx_genai_validation_packages.py b/tests/generate_onnx_genai_validation_packages.py index 467102571..cdf3ca12e 100644 --- a/tests/generate_onnx_genai_validation_packages.py +++ b/tests/generate_onnx_genai_validation_packages.py @@ -3,6 +3,7 @@ import argparse from pathlib import Path +import numpy as np import onnx_ir as ir from onnxscript import GraphBuilder @@ -10,32 +11,50 @@ from mobius.integrations.onnx_genai import write_onnx_genai_config from mobius.integrations.onnx_genai.auto_export_test import ( _Cfg, - _diffusion_package, - _model, - _value, - _vlm_package, + _VlmCfg, ) from mobius.integrations.onnx_genai.workflow_metadata import ( write_audio_codec_workflow_metadata, write_language_diffusion_workflow_metadata, write_speculative_workflow_metadata, ) -from mobius.integrations.onnx_genai.workflow_metadata_test import _speculative_package from mobius.models.qwen3_tts import Qwen3TTSForConditionalGeneration from mobius.models.qwen3_tts_test import _TINY_CONFIG from mobius.tasks import TTSTask +def _zero_unresolved_initializers(package: ModelPackage) -> None: + """Give real tiny producer graphs deterministic synthetic test weights.""" + for model in package.values(): + for value in model.graph.initializers.values(): + if value.const_value is not None: + 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, ) - package["codec"] = _model( - "codec", - [_value("codes", ir.DataType.INT64, ["batch", 4, "frames"])], - [("waveform", ir.DataType.FLOAT, ["batch", 1, "samples"])], + _zero_unresolved_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 @@ -89,6 +108,295 @@ def _executable_decoder_package() -> ModelPackage: 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 _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"]) @@ -160,8 +468,11 @@ def main() -> None: decoder = _executable_decoder_package() packages = { "decoder": (decoder, {"config": decoder.config}), - "vlm": (_vlm_package(), {}), - "diffusion": (_diffusion_package(text=True), {"guidance_scale": 1.0}), + "vlm": (_executable_vlm_package(), {}), + "diffusion": ( + _executable_diffusion_package(), + {"guidance_scale": 1.0}, + ), "tts": (_tts_package(), {}), } for name, (package, options) in packages.items(): @@ -169,10 +480,15 @@ def main() -> None: package.save(str(directory), progress_bar=False, check_weights=False) write_onnx_genai_config(package, str(directory), **options) - speculative = _speculative_package() + 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)) + write_speculative_workflow_metadata( + speculative, + str(directory), + grammar_guidance=True, + adaptive_k_max=4, + ) masked = _executable_masked_package() directory = args.output / "masked" @@ -194,9 +510,10 @@ def main() -> None: Generated by `tests/generate_onnx_genai_validation_packages.py` for semantic validation and runtime conformance against `justinchuby/onnx-genai@c9bddd6e`. -The decoder, VLM, diffusion, masked diffusion, real tiny Qwen3-TTS, -speculative, and codec packages contain graph-only synthetic models and policy -components. They contain no downloaded model weights. +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. No downloaded +model weights are included. """, encoding="utf-8", ) diff --git a/tests/onnx_genai_workflow_conformance.rs b/tests/onnx_genai_workflow_conformance.rs new file mode 100644 index 000000000..9b32f3e1f --- /dev/null +++ b/tests/onnx_genai_workflow_conformance.rs @@ -0,0 +1,145 @@ +//! 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. + +use onnx_genai_engine::{ + Engine, EngineConfig, GenerateOptions, GeneratePrompt, GenerateRequest, PipelineGenerateRequest, +}; +use onnx_genai_ort::{DataType, Value}; +use std::path::PathBuf; + +fn root(name: &str) -> anyhow::Result { + let root = std::env::var_os("MOBIUS_WORKFLOW_CONFORMANCE_DIR") + .ok_or_else(|| anyhow::anyhow!("MOBIUS_WORKFLOW_CONFORMANCE_DIR must be set"))?; + Ok(PathBuf::from(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 +} + +#[test] +fn mobius_decoder_workflow_executes() -> anyhow::Result<()> { + let mut engine = Engine::from_pipeline_dir(&root("decoder")?, EngineConfig::default())?; + let output = engine.run_pipeline(PipelineGenerateRequest::new(GenerateRequest { + prompt: GeneratePrompt::TokenIds(vec![4, 5]), + options: options(3), + }))?; + assert_eq!(output["tokens"].to_vec_i64()?.len(), 3); + 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(request)?; + assert_eq!(output["tokens"].shape(), [1, 2]); + 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("latent", Value::from_slice_f32(&[1.0; 64], &[1, 4, 4, 4])?); + let output = engine.run_pipeline(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(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(request)?; + assert_eq!(output["waveform"].to_vec_f32()?, [0.25, -0.5]); + Ok(()) +} + +#[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(PipelineGenerateRequest::new(GenerateRequest { + prompt: GeneratePrompt::TokenIds(vec![1, 2]), + options: options(1), + }))?; + assert_eq!(output["waveform"].shape()[..2], [1, 1]); + assert!(!output["waveform"].to_vec_f32()?.is_empty()); + 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("serving.slot_ids", Value::from_slice_i64(&[0], &[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(request)?; + assert_eq!(output["tokens.row.0"].to_vec_i64()?, [1, 31]); + Ok(()) +} From 7420d7657e61c75d41156bebe9e808096a49e422 Mon Sep 17 00:00:00 2001 From: justinchuby Date: Thu, 13 Aug 2026 03:28:27 +0000 Subject: [PATCH 034/151] Stabilize TTS fixtures across serializer versions Materialize floating TTS test initializers as deterministic zeros while preserving integer shape constants. This keeps the real producer graphs executable and makes regenerated external data byte-identical with the ONNX IR versions used locally and in CI. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: justinchuby --- .../tts/code_predictor/model.onnx | Bin 110332 -> 110332 bytes .../tts/code_predictor/model.onnx.data | Bin 2529408 -> 2529408 bytes .../tts/talker/model.onnx | Bin 32472 -> 32472 bytes .../tts/talker/model.onnx.data | Bin 72704 -> 72704 bytes ...generate_onnx_genai_validation_packages.py | 8 ++++---- 5 files changed, 4 insertions(+), 4 deletions(-) diff --git a/tests/fixtures/onnx_genai_workflows/tts/code_predictor/model.onnx b/tests/fixtures/onnx_genai_workflows/tts/code_predictor/model.onnx index 8b27a601fd3723e0b01e5ed1445b2289dc493d94..b4039d6c580b82239edc97441ecd8482370f079d 100644 GIT binary patch delta 37 scmex!m+jA8wuUW?Q{1>2fIusZgF}GPX)@z6@yYViLfcQcF`8rm0QH3nZU6uP delta 37 tcmex!m+jA8wuUW?Q{1@!|F8eA6~@6K!00rY@tF8zd1;~TC)^lKG5|-m4(b2^ diff --git a/tests/fixtures/onnx_genai_workflows/tts/code_predictor/model.onnx.data b/tests/fixtures/onnx_genai_workflows/tts/code_predictor/model.onnx.data index cdb0d2e4321f773d3e97c3fd7835ccebe02abe60..7d72fe4f2011031b6b0cc022f53e07a512cab7ae 100644 GIT binary patch delta 1627 ncmZqJIE%3Xh&MDzY!oOO#iL;`ng&MGz-Ss6O#>9Bfesb`FTD?6 literal 2529408 zcmeFabyU_{{{C%uE4HGbVu0On?@Pr%QA9Q2u65Q~^Zk8i&hgAYk2C99qvyc#_-y`sQ1b&ae?-BSt0>4M#_XzwRfnPlWBT%5g z|9{nGKq7>>;hMsN2hREUp9`l%!&f6$4yQ+sXf#M>Mr$y%B4P1 zU#PFtck0I<^xv45))lWPuHhEOV>2d*!e#|fs8ds<*Xaf8sjg_87>L!MBXE3b9PYJA zg3W9V`q$IpO?^GWPU>OX-hfY910HWMAnk+!wQm`){fPltuQ?7M3|RWffSX@9Rve42 z9EU8*mV9HVZPXd+Hua2pOXX3Ysr<3{M*VR-{;9EO6seQe#~{7eq$%Y$hjjL>zEBLR zQ4Dw8n`7oZY*wi$AscSS3cA$Z-)A0-pRacD*?j8hY_ zy0HfKH?#uCVSK_EOUw#R`kXDi8{L; zi8p=gqj1@-cx7;g$w@yH?GTO@fw9PqPQ-+&8jL)vg?ypIjWv3lF*U$DoXjf>n0?rQ zIvj%<9E0W@gUTF(Cm#*)=NJ^?7?dZ!k!nuHL~0%N$1!-#-~TtpU~fhoT<3mPE!1_E z4}WkSEE^!76@9E6zCO#?&{Qqjw#XAZ%3DD5u{#`1U64E`03F;S5Z5RUXI~}aexL?1 zg>|qhuSaL29?s1TxRhYPw6)v|#|+qcgItdd=)st9l4Ib(b+Dad(3WG+g8UaL8?uh1 zHdCjlEb1xsa~=GXF~Rtw3w*!!#l@T^O8m!RvVM?8)+^UZCT!iTyid7ev;X5Nv1)Zm zxL#<9sk80jFn$=OMg+sMM-<)##N*k)BvgN`L1lLxGVJu&a!-#N^O=^*)HX5Q@**7!+EK$6yut z>ykZ;T0tGAu2GMuSJZndKNjTY$v@|X{CxR$a>Eq_ORH-YuF+nK*?ol>>Z&+jzl;j$)<{Eg%eXx&fAe&>bl>9Zx6+$he4pCPr#-x9A4MZ26VSV}1 zP?=oR0?ECG!RF>%7}rM`uwaV;W~>V~W*M-HYhWzbzyhv; zhg<_C$sb03C-Og|`jU4NwTt>=UGPt>f%C^!emdSn($|cqgGSt z)D7w}^}prAzZwT39D*?1=DRZIM_E~+jR#sbYK<{_XDjtz1j`HK;^osbmU4XbEXCbt zsw&|2E@8Q?JbHd@i#dJ{m>TGb!yAWVMU`mGd>Ie-8su-H#pIDXH1*aa`K=!B9Sv}t zYCyBytbNWKaO)29;d9o#%mqp87vi}G0?EIDYd}l>B4l-^=28c!OVoWThk8f->=%Aj z9K2SfjJPp*x)}W^N?y$=fIX+epxN9}wY^FYd1H-1uG8AfmhZ1Bjb2$B<*ac+d-1y% zy|)RLAM1&-POjKKAQ0YnBLICIsy9hOyD1u^)zaZ(T|M;2^(fTNfW}$_JlP}ou}8Rn zlXWk9gml&ewYdjoa}QkR8n7q7Apb3@8#%{N|BW8uiMpG#?dXe!PBF4nzggCrbmkZw z6v0y}%7dO-8Fk24qG%a8b9@6Ou~Tc2d*zmBkXZ-Cwh3sTIUy*`7hT4NAv8Az!z>cf z>!BLaTeNt1RfoO{^!QelF(8O*dx-%}4{~p38c>CMU^Mr@8t#D=+ym{&|A1@Y3&&s^ z`5Tcvno6OLQrD@Ud*Ej-{8x4HU)ckeP2yzPuxKpXSy6tu{Y@$N+7CN-H^DB~HA=gi zp|a1_SXrS}6IuW5B*l5;TGi@B`-H7mMQmx<9zQJmb%u z(js?~4mp8(Jo%u<`9WO&)42xr7;uGq!1FHW{)GW1xdtrZ9+<*C5KsPHqlJ7mv=hQCXH5CJTJ854Q)N7_m0k`0)Aw*}bh^zJJh79zJ|T zNv+u-t-JFaF~+9|PFPyuTAki_dEX6rE*@#lb2j2?<6G%&plUw$bCrWj$7? zSksJTjN8Jzc8dK$7S{puKnU}|cra>ZCE)=Sj7m?Ms# zGDnE)24&?Cte@bI4XNR%u|F17)+gd@Ck^5sYO(o?4pX=3QL_{(x_&KPcM^G4)#{I29&NtAxKgXO0%6S;PeoigcT74fpbMd4PoCakO2z`f}pw4LRH z#fL-DH6{k@FC}2_7ByZ@)Z)xR9qcCRVN;g#%6fR-0t56$=Jm_$^H~d>WG!UL7%-4A zpfLH5a1A`<7_20J9kS-v!^^3EvK}^jeO@eoI9!G<4aSr~1*Oiej2uzK6Q(U&!QuE^ z<$aIg(r-w-Oz^Rk1!m7weEZE%C55Dl*n8#ixNbYV&^Vw9*&{4qVfK1JUZ1@ zV`?)k4v*3ymwB*Qt{!~`u(o9$bl7b`1IBKFr(7z2XIzmfbp^5^Hl zVbnjJ2fubGBdcq6DCfq%7gal-J#Zyuh-~K37^kPW!>{UcW7I$=Ip>H@)-K&u#*ULp zulP|1w3hQ(Cl^JFdsbLJzc#D=F4tl7s>fr~* zaMpr{7z0$S1=F$&ILH_H8v0T-GGuiAzPXXj2DS_^Ux`>kpDkFFsSErmrxLQ({G_&hJHzS0fx4*&Q#> z4aE<~0PN}*f$Z{esBkwCb{-nE`KCqf(t6C@ug5Y==D>IZGS_gQA7kHpgSnr+K!ewe z`-}mO4ELn< z^;D5=Z6b@k*UCqMKJrYdQgUR8I*L0wic{sXMb)KsF}@>Eq~~CmeDa0WmN3jL8;i(} zi3obF#z%EBT*u4HO{z zPO34vW2x0tI`zj|_*c~d->t3-+cVyXt)!O@{410(R@vC?29k znN2#J*6NX1gngYm>zq041r8X1ORVYdu{U50R9FX;BL5Y}fYs!`$TcvRKZ^;O1F0p{ zK`N7aNWG+f^8Z~=AW`a&ymg}ZJYbkKuJyp$#&?vpt!u~=BZtB&q9abf9IMPo4wUIi zf*dfVg><)Gpe)LaQ$1<1T@>&thiz5bBH~0pZ294VE`h_b{c#jpUx-J$Vrq1(qeXa< z4)@*ksQOBe0}ibFCv#n=8W44kJ}G117{|F56>Sct8dpvjEaS)OBx-e$mjjH- z?H)(>)t$dUG<;qZ4{e&GaAt3$dAP%5V-VUMjl{A_@pwET3Ej_VaJ0J)t2^m&<(eKn zyD-itbI!NY_hk%xki|947`T)%a6Mz-E5?8dNfR^%CCd|$p4MM|0|x* zkg#~UtmP27!J;>s_xP&3KGX@5-|OYcV=A#`ksgP?j1{I&8_0?;^fI)DyS!vsP@dnM zXFR726vZ;0h_WTk(0WQ&q+fBytAhRrb`Qr!Lo8m6N#J?r&YU!VIr zf@~`dsCI<&e2sCRb>M(kjQQlxV+;r3WQ;NN#`b z^F>_miUBjOu+MwI`jiqs{58mbfqUQ;$6y`#>ytHtT1g$Dezh-J{aqcg+NCmf z&TxP$Yk?9~c?NfmIsfx6#eP8)x^^`|x%C-JwG+{Dq%KmPmi1&kU7S*7@g3DAr=!BT zg$Z^Z>4>tg9WlGQH`<>JLB_dg+zU!TonSSV4%1@LDjiG{^f*y~In~tw=h^Im_p=_p zXh1S^zzz0-)0qQu7y}P62F@q{De_Mze?>C;Q;R4`<>!D0RDKS4NBvztWUA$Gl&qjb zXUm!5_+1aViM`*Z3-^_bJJsaZb}s0@y#t2M8?RJ37bq7lOpv~vT1d-@^ObcY)v6*Z zwu@oo%VAW@wrG8`ADWc$gf?P0s^vss|K)g$Dxt=nx>^LObr|cehk-euK!4`CDU5Ha zWIsnS2TWlO*vA~;Mg9Qt5Bo3Xfc4V`V^0HnnDyPI9I9JL<|fvF-Q_%^X>Mr|wcOs9%u-_KbB&+cxB$YTVYL@?1%?vWWH>1EH!E ziS|q5@J^M4GYd7C-9U%E4fL=+tw%%$*8V!q^F{;8on-HKi~a_4z&qxE&g5Um9B`L0 zU=aCj$(~JhC+k>hJC#9YQ~5dI&tAY^$pMe8T_oO{%fP}0St;?0uz6>%bYE5+n!$Qm ztm-x~@@gVRe>W4i8nu%vUKylKl#?vf;pEZbFOc&#$z8)oct>ogS(S| zIr*EAzcJY}s7_=xP@AZeR2KEC&Idh*2FfO01{pC zEItJ!;y^hK?@DM9^;CylDS8CfWj+gI99wQc>%(NZ%2@x9`HnfT5c#v11BQ|R8}~pN zvLB$V$eTc|rH)bgeZaqy1Ap+p_rK-AXC(?N)<%`QHO*RS$v<;XuVHdqZWO8y2$InY z_Zk!FVRfpQC8AyxleOa1azQ|#Tsyv^Y`LP7640)`C^sonG+b5-Gc8q^_Gl2wR`A8y zDq*O-J_c7lBw%&A8aEecF_JlOK6Bt;=D;t0T=$Du*DLh6uF&UT4y?``=*b+I-*^2t zbD-5*TNLh@lh$?9L$R%vCzLsQnSUnxb{2i%)r*yZS7M=Ooxk|?b;TtnRo-VhhxtiSd9D^83R{wE-r8n%;C>ch4c8c5Biw{e^nop z%37^cVjHBM_eX=#g>c+1M&3LW3B|dbj888phbg{j)UO%Z*IK1ibPAOp=ogEuCbGTr z6ea2AN>%WIy&|DPMVxrq4uwkehs|YAgyr%~#5Nke%O=3jN{zJkT9lZm!;mn6gQQ$`q~4vzb;iwbDGQPiT1d%yE|;!I4ND5_{jI)ljNtm?c@g2)k=J?hAL&^ zdNDesG^|^<#@$hUQOLpr{jLS$$%H68m_;w_b`n|?)0YOI@JuQ9pCw->HFO)0)Z_6Ws2=MiInLX=Z#re~`>MZVB`1u2{WSF@}|NmYqX&GHegzqtR)~xEnJ& zZ*IO&xYRF(XrJZ?e$^WrBizw_cMv9@jznm+c$A!y1P9)is@_wFQM@nZcU_NF>;reR z5479L-u5(m`7G9S>;qS`2KaejsuKC5$v>3r|KxqC?uWI~we3)2P{0b+``Yeu)q|Ec z86WzK9ljxAa=~z^N-2T+1wzp3TdGZ5c{_P6)F4wA_m-&%mz1ToY8xAFoG9*%D}XD8 zra0BF7y4{*h0Byc%qS6wd-LLOsZ$ba&DY?vnGUnf_0XNt<6V260q7Y2Sp(#ppdWsV zah)~5bLPOdtN~J(1GBj%2avxT`EOJIR1I+F)ETAkiLNrVODTEovq9EfkSXe!l~PJo z_$caSYvo*@gJNLy1eBRtM)Y&-C||Hgn3+CMT0DEI*xD^Mb}kVw>h5?iE=+8M)=RtN z!@Z%HJ1hX7yF{S3NgNhDNW@$(4gD-Fs+Q5C;{iQNSaD6qGvBRcFGw#i@dk4~YXD2~ z?|aYMhx{AJ--7&>WY6~kt-0R+2`}(Tt#X)rM=#&`M@sj>o^o%At{B+iney;s8&pZ+ zTxSjx1x&eS#Rze*jD=iV*B}#`xXF~WKNP3#FO4lV&SFWkXX0YA87|!EibG|Gpse$!yMj~%Qj=JA4%4gTnFrf*I#42V;>yM8lW@z3z9#M z{AJ0Xe^>4h)r$OTYCUzF`c-}K0seU-JDJM8TVmz7r-kH!x3T!#yR#fN!UOfU#>t-L z-rAJnykw8qD;oBzBJIv=<@kd>^5vvbQuU^;(q&J3(OiF9*oD-^H51@S^TF8g)E6aF z!tkgUSKeN?xl=r2kGl%8qkF`&ocVd@(7htWP30+7)NJ8sr(~hkpjPYgp%CZM zjdjo*&h>uQgY*H8+~XM-`8Tl!_|6=d@tLgTKS%yqGgYqdZOs2=hO5dQy zvZ2#NrRL*Js=Lefi(6?G;o;vNkB0PzYe_FeR}H~@*J#++PJmShHEOA}u$roaSFj$} z^LT&SiT=-YuK7Kz(JydsbG)0g2C7f~6!Om`znc8J$)8OAlH~KE=2L0ZW$IU)iI2Er zF8v0S0|q3bBje6zoAzSg$$@f8kTc3WJE1fkSYHm$8-x;tRB$blq-ru0Fao^!5Q1KE*(2>Az*e+l_x$p4*kA@eM1A9az+uYq!? zd>`Pi)KFMW@==rTW1+DWt(E}xW(pSwUwgrrxK{++7e%E?XhCuFx2lGjOSKSI5;F8KXxZ! z)l&`b4AG%dPu_{RqsN#Y{Q1X{b36C{8P>zM8S_~K79#&M)&Moh-;DgpGfme21g2Tn4`8x#A;q-|G}LRXs^n>LsrR=oQmg6cFy+0wmW9O;UO5rOFVJ_4O5 z#bHLPBDC&uN6Wce-Ovc^A~u z1i;%V0$=OL;r@$6eD&9$eIXq@%jvNzO^+W|1}scq-^Cg*ku_k-4c>=m4G8k5vnFvP z|5oz1A%AnSU!-iv`>SifIaNaBYrna|T&t5J#Z7wcjF45zHpA^sX-c;@4f6dH*~L6r z6k4K1)7qZmVrNU~KgJ-Zo*gRpkNTw4t9;d%yunjMjn5H-Q(-)hP`r|mKDy+)phd7&D$n`1w6i5N7sjx1DLFKb(Q$|qq(rNfreN^1R{!l&5- zp$V>s+^8<7P~90lNBiMW$#6{S9g78?iO4CbLEmGn0Uzt&y-JTsb+~pz>4nh;`*4Uc z?J7Mn)_{e`pT`;~lQobF`9CuTlp_Cr%94!n)LQBob%Xl5K3Fr`RGVdCQE1dWLuq@$ zQ<^!pegJwDz|D;C5&UHz0WI zl89GFirYI2%lFnUD)(y|)E*u#lcLoy2D{5_^XAHe9U7c}`A)R#^+owKDpB@x94_xR ztSSeOwpZqDuOz%1Ul1?$*1*$jYjiF?2oE$q=(aNyt%GAwEh7Pu8`NkxUW>1a4mBt6 zK0q1nX>azm3z+YXtotrA#@**Uvj#LHzdiYXFb5ta{{!-`Ab)MLg;GCj;H&)gBkJ$e zz?;~sZoX|)ez^9;(HE1_ZkB!~vg&)lUg~8Jj{k(mU8S@2QYCw3EJj}{hvY`rmFYuc zrB%5Ixvyzm*{flQGJo|O)rir@#W2UJC^WGX!s7;@S(G>C6b?m&57D?gG6BI!YILR# z*oi*iV)}qP3mahS&Uv20yM6n~c8T-J8n_bqQ&|HRA^!#PuO$C@^8e8XTuf$3Wm13F z2i#NTooZ9pF!`MO&%EnL(KW9ZrWCfosOmiHslmL}tG;w@bWeHZI2@gG>mxiPRXN=| zLVoBRBip}eC=2%(qdb&GRn^l*(Icx8HZ<#i$gKXT>*IwFjv>e!5e<)a2^iN)jf1_k zs4<5=SePEiKI<{anRWgQp84)&tYZy)fi-Zz3%(U1e_9^d$v>L>M)Hs6{FfoGH?@E= zQdg)4)JrPA2ly*B;AsCI^1_B-xw82HtPXU;sO$tJ&fO;**9FLV<^hN?&r=-sRFGe? z-O+ke8ypOuq4>E4%Q~0iWxseU+2_h^W&ZvNsutgNifqsF_!y*i(-kZ`cDk@@JmH_?^l)c8+_RHSk{6z?WGAMw369{Ql%$ zM}7_Y3zN~EnnN9+E>ZWX{2KVL@;hIMkvruk#rnWP5%iup$c1-54}P~vE8azZENq7- z&sQl=eXZmLU0;}9wZ%f+Q02)?9~pZ!NfzAGR>mSlIb+jIb-mUG5q_;SDmb@BrA2*F zw}S^>-V4UP=~0N8ACJ~|lTf#)7Keg$*fBtl_$TD=!}UIqHQr9vHtYd*X0z`lzc=|U z$lsOxJ;*`l#V(Wn$$AqTqv#%7B7i=(EHLDbp7zo2!+UpIg?$!;00EsiS&I zH=j;u>o^dtAC*zId~lT(wpyv(W-VtLb}HuQ$Jm%RUM`yVDuLqJEpX~}AGEDC3==GZ z@z^8^W$fZ{Xj2mM?y&|Kq(c{5Jt}4E;nbb`o;^V3HhSA<*z4S8jsJ{(Cu`uFtN|;K zzajYzcW1e$dJ%XpFAl7zDa~ik zQ@XakE1E00Hlfqphf@XhnUjp8GX7p4r9MiagK?Ac8!GY%_jvQ5$&na>s8n~LDO zrxlXB^hTL4Zm>M@QzKs~{8k5N3|+roW(ioH%2 zYeMz_vE-jh{@bj9ijaRe`3I8yDP>Q_iPSFYEOm!^PQ9i63jbbTE$R&Q!nT6rMETmGQdg!7R&NW&t!=h8&AfWc+71SJ>#&`y9eqKW zc%iKEec8#PSBZisf4wR1{Px18Ojndz7l;})BjL6x4*k0(;mcy4F*MR4rJ){|c@})B z19P>G_Z&8||2=7d?=AAM2l&Do806nb{`<^9PUP=S_B&J$a*n08QyEk?^{;w>e|P?q zSDX-qGe(QBQC;MZ=se}Z$7#yy5s5NZrxs5}2G~s1928S8B*@{{ABgsT(a1i(UF8%H zIp(=RF35C{S+(yddn(5pS3eyoy6^rhT>3PI#n~RHzQqOa?gyaPf(TUVPySL#ut?Hi zTR9yxRrUCMNRQA~jCDz@d)6@q9A}?%lewS#Bgp@Xu_%H3Kkq}eC;QL)P|0N8O8xOZ z)Su`7->rc^e=RSoco&iT&H5>~nz~67n?`ckDZNx(`7XwmAE>O{S^|b>y-b|4MO^&O z_pT-l#kA{fz23b-j$sH@RMDtvOT<;brC%+DqmG`vA#AcV3+ZQZR?H2dOtq_qJrNgPg zk>c>ICep>pAgznL%9)41Dj$m8H&%P?Dl{3-MTM&72<+Ysp^ih)c#uCX%?QUT)=ozj zB;ryF4f1Yl5tFCGxy^bEXuxM-K9HTqD)gG0 z9}-y|D#|qjw9@*HuRPkijGWP~p)zDkD=};3En#xB4z{fl=y1;oIlFzadQ=$7zlp)c zhKcC$K#fzIwMe|GL(v6#oU6)M9>lf3gnQs1<9{aK#gl*JOZGVAU&b1^9q0c(`M+=t zY~%bl;(SL_Db)Ye1AXP+`}6+)!N2ye`rs=LdfdBWg6g}hP}XC8n!R~*IceTgW4ScO z-v>Uj`}J7dFYPZ2{CHxlHA{!LO&*BR9gE14&ywWJoOIzbv{{ zsDsc7&sMplY%NTW>)=0?K45v?0rF$bw~#rIV{3PX zd!9YuSM~t?$R9}lisU~-{^#UhLw-|ogi|Z1UvUo5uD%&UM(SWvC0!Mn@nbjX}_m1X-b-9`if~%Ahwv%G-8HaC}(^tu;ByfoJj3WMHUFi>)DD)I*iWrX|Fw zP8q^$aCL;e=!^}%1L0iW2k9Z9aP1re-IfIC=Bm+Hqs7GCIxI=%U5FB#YfpOr^T^Bb zE_s<=FMGgv@_%3t(471)SOe`M|1I(_CVvfb52uz6ld%(^emntbedc!%s zn9R&^#l9DZRin*IBDI?nnx8hvVUrn)yshP=2Ah?=+oI96bY*y~y`an(93v|ih?FTe zP36reQA(-gN2;k2>0;U}6KG0w!hpI1;Ml<%@mV3Lk`;}qQ3*I2uIB%j7JFCoE<~ap zrwejzyOMu4{jUA=y)JTI*#qXhU@b`g_pE`BlYcS!&yas6`74n%fcjTG;9uzh()iZ9 zU-?t2=@-MLWZqd=`J0Gx>5Zmq8lvn~jU2)|Tq`Emlh;SyS8i7vj+IVkSYJL>*<~3a zllsTVB{dq!dGp37zMIoj>Ul=->Qp6EFztYMm-|EG?gh~=1mpD4INpZy-&2js>;c!$ z)?prdz`5)J+d4Du&0x*Hm-nDAFsHKz^d-M7`S+851^JW7pGN+1kT~yov#=5Dj?|0J!JWrpWJzyq#peO8s4CGgne>nL!kzYstqU0Y&&7}@dm#F*H zU+n=mWwsEdXBuRSPu1n9dqwckRsH|I$xUuxDymj?l$#x#0NW!cPUKk)(R!=+Cz0<7@E%r zhRQ7p-f{65#P3OU;`gK~crpIl>*4W0kE6ZFGoJ4~chKX=;Q4&ysI<<$oK;5NYP`T7k4kdGDwGRsQp)SNC@Pq_kQ~Yok`?d&Lmo zfEw4jV_EnSWogz&#i2dp*|#gkW~KVdcFt|FW#(YGtbAwu@MWm9Jff8)TvT#ct3As6 z!wqdR_AC+i0*WKkwgs}r^p$jrSbSa2?!c&rAMuZ+O@918X79e;v-hn)B~Y_SaN@GEJdUsq@qy`Tx@%&{cj* zE4i{A9=1%t@@=ltWW!hCSE4Z7D%6$7?oU_lnBEdA8rM}B_3pCKCNngg?1@sxqm8}m z43r-m>!st)Zt}R}5#`~QrfFAK&Jkw%BFKRi&Nc0gSufnM$ukH$W=2BAbFkw)2OGt6 zu=CbBT;e%c8TP;zRQ&lzvi9G?ewICOY}Wrd2MZ_vf8!kN>wGuq#eDnOEDSYX8lJaTM<;MZkVXA{{+&TVRtM{f3#Y5mVB zpT~Yn^Glv0+@}@9qxQ{k|70(i7IZ_6qk$M_8HwPnd>h@HccE5l;M|1uzl9#E3_X;N z^t1F_yBitj*aJ_w#aizPYk%_FbN+X*2YA97Xb9)OC+Gey)su5Rj`|fnu&al^e09kZ zE|Uvk(MB)nw=hMSZe39c%ZQK*M%NL66H`^O^H+&V!(*k-%n!n_IRdsHG^$GvyUAt8 z4YKrwzS7$ArZO{Qps}dy7~$UTn>gRN3C#R^;%z0~n`=Z5AUy&~NF0jRNy5-E8Z50! z51Ja<~OQsOYS%_milxt%5bXZd)>0BQH&krLxI%ecTnEt0zAiG3y(nEACkN}IW0 z;J5$;xktdpA`XtPdB-3~gAIjs7*Jl1fkwUuXinchfootbeZFJ#wQg|Vvj^@;{!{D$ zJ;=X<{B6nKg6tP58}g2%HdCjlEb1xshWbGL^uYcd{s)(PqReEETPmbSyHt8$O$>;P zWo%nb|2>`OTGtr|$$$41=brp+$v>C;X5=>~`)R5J8Fkb~>LhiG`m6j^>c-&QOuc-4 zK`YIRzZd&n8|3Y~!{z5Ue)8~?&KTVNmeMV*1stV8wm;+}O6<~Moprp}zpk-d!h1qv zqFm*W>~G2()5k^|M;Fn)-E(pBjv4CJ>jv|VLr}%mA6>?UW6z>kh&hS4XUR8lSz6w6 z)Zxb_*1+bxyAw$tbQNoyBdlkx@s0yM05$o0kiP`^)#R^0{_12uLbWEZn);tSfKU8; z`T6hPng6J922tQrZE3aDTYAK}VoKLQS>KoM93oZHEYKZ$dYZ^Oc6P`-pp}ci6cPKT zu;+L@Uo15-l?nEGSzxM%>@uy0{2W$7S+>tk_?SKv6=Xg1+tdXXd=Jo!?*UHN498Bs z2e=oIh)v}*@a21ekf+Rnd=KDPmp)e*eK2|etqyZ7Uu7@D9{9iY9-wH|=JJEqL_Bl3 zqS&zRx>UB2e6%HA#-A@PD^HKX*VGR3-Vb+7bd8rARMl16Sd+Xr+b#0yo5-qfv~p~e zk33zyl#Ec-QC?T+C`^Cc7IWv<#qqX)!FDi$^L$Zda~N)viN(hbiKxMM(K^11R^QU0 z=VCo_sxx*4)6-naeNGSHekOZC@^5(gf8IrJ`^b35`G3wD;0M>hPR@U0&TTBUno6f` zP>-qC)O+fGn*aPd=+7~*b^)y%I>!D^K|D$aqLnlS)h$ZR<0Ew{U?=|&%U%$2Kriv)XUdIXjpB`*{4E_ z9!@yt;ETE<433#}Fp3@^pdRt4naib-<+gkWjJyTXpjh743LS>D{HDxK4n^Na?Q8B~kjIb|R9eww9#)>WjQT3xY z;@m^w*fIvM)+V6COf|d`wK%{y8aYZ2^Af!8>&f|?$M}}UpWza7!F_r^tjS*Fa;`7lP}I zA(%16Ae(Vtc1c#r8?Nh=h{62k-}Z`duY6f)wjoB|-Vq_=9ZhAM0};xO8qZXIwT_9H z^Csxlx)VY=58z#5Zyb3V0`rH_xEP;+HqmN~=b4M;S{;@r=@DGWfZlHGYv(Wq>}TG) z$UT0Kdzc=;YVv<(yg5bwrR2{b|19!XCiAcIKt3!igb{_@4&ml?}Wzej+4V4)SOq)5?!h4hEc1 zwLUFH^I4U#aa{+@Z{>*63%zhDE(B%^qS2&B0xk_uBi2ES_ys!bj?|M&kFi5o`^;qg z+RMFsfjN)-6UgsQ{?p{&ME;57KSchCh?j4} zRw>PR7GQTi0Y2WA_`>h|rRNlf%cWRZqFyv~CKctD=ie3cy?%IIuqomWtyTgKhssrt zV&!eeCNi_lWaaScHL4!b`^4k{6*09|dssE<5Bn#cSj+bl{rX1ZU8Mw6ZKXzoPFl>L z#B<<4_JAMs*fWT4pr`X3U=Mqs^Lz(G52yzDE0cc_`Txo9L*1ULmk(#;Di$XPNxRzK za_EI#NLl2KKJCu3~b1wCZt%onmR}@^~|`Ei9`$;6PhXh>63I@;wS|?!@C;c{S#lYcWg5yYL>o z3;tS<+x^MOx*|1|b9|05o*u|r=Hip|0C*=0;hcXT&i@+De-h`v5a-#AnnV3v4`_o& zJsH+vu`r|xW?m4=k)S9F4a{;o+uujDN|y zz-6@vjL^YjFz>-V<5{3R-?&Yp*S(7#|5@(&Y}UQxH;~_&{0`*rNB#xmk0Sp!>Q{L{ zMaHfa33>jo4Y!rodOk=~hYwc8aKHCH8-<)|7mObdTgqioolthiNTuqo4su_=9(c3A zH-^0Fte9@~ls9gvWrVha{N81Q(xz`ARl)XaM8lX;aP@2j^*ejCIx!6U7X;(F7w-Wi z#beKzB%H|O_hP;ICb}={UwS~L?HK#U^ZQ>r=<8;%_aXm5dLV7cUyA(A$={LulgS@M z_Fw4%>6#}Z$v+awYdn-*Jjv2b@uW$toUneL*MDW9xmS+iZrt7j_*z~pgui7jSeNPod zgGSA<A1m{(|aD=1J5p>Kyg+9z?ze_@8)SM-F#XqQ)LkiXHIA zJ=R_8PDDsMv0L1FbUw?xkj?tx26J9!C_$$kCQ~ywMIX=egma1@D1`MWWM)_4Y>-50tli!W}_T-Nt{~7XsVh!v|{(fYCMD-@;1ZpRBmdd7{ zQ*WvNgg-GeRPOf)m$?~Qbae7YgxVlqhNs&!OfN4?9S6$+9p9wA7#1n2cOM{>yb!6_ zm%;IQ!Kf6`!Dd!sZ+Y0>AfNBBlY{zRP%0lUV?1RsS*)m75Yw`oV(ZjiSbW_TU)Bfm zJU9|9Q}|6K-UD%7qJecIeiw=NKyps&alZrKxaxTaXe0OfN!B*E=xIM;oO{i>i~K+D zfea@9uf7LzCvAWn6&NJD?2b@EmhhdEw0otjgZH^tBGzSq^j@M72;*H zSed>jPvq$$kYpLBT2j|mu03jyceeDE{jT3osv8`Q6K9PPl|8=+rw&co}$@%1iCb%19_Rlf4A)Obg^%He_1rOHHO@LQb-n)OlKm}G&oeS2Wz z5Eq=@5P&@~5zx1fLyymi*cGn9{o*>5uf*>mNIj;u;CWvn>z#F6+w{Pm-=OFLxs(6e zd&Ye7?dgKJQc53 z6q3;k-YT6CryNLkmMv@6mv`vz?CO~(-VE!mbnIIUm0Id$|F7Fc(>006^sgsoj%_Eu zyfMgfU7X~gjX6rs?2X3eJ0gY8!CcWKt09^-?v8S^hr+8#0M^%yz@_i8czq=iC5LFB z{m6TUCG?oTm-j%M(Zh^m%wNNMo$0&_MgDN|mnQ!z^7kPB3i3B4e`B(rr8<$zKy9M_ zjt3a->WAK(m$8;%68l}`sb(ESpB(mlj~(UEj93{vq&^Z~uTx$n*Twax26=V#6tN;( zgHDkH#Pxp7)9bVxW*r+PK0CW zu~>B3o{0L;Am@n|FTe2~#11{mS9IxQfmqyT~)qHQo%f zuXe?(;zRJ<%pWzQ!*P2`EZR;^#P>!Tbhxg??p*o{8}%@naqc6@vXXT!J-}7hnCr+N zK>m*8FF^iS@|PjM3E2-(t;qb-13b=O-=y;I0llX3J>XB&zvBVgb+MCIPV0o2SX)+` z>@8Q!bVaUZpgikOztc@6OUCj&`w$b^f1n+deOj4#p@^`a%Coum^M!3iQ<>X`-^!Zi zA&oPNNT-Aniq#G~v99(*QA??ZH5*|81nF;dE2yaym$2M`=HsMjZnlxDJWaMsddAejtxlUC>Io#V? zY_hr|?i4b`03A>-b}&+^agQAigHkgVrMf1f&<8b~_i6F@Ht#|&;kOZMu$LXqv&*HN zvxDpf$RGBQy)8Yk)AYdh(F5tp`G3h8s1WBqm1@GD;eYY~|8MjEga5lf#=wupc*QJ0 zEtAXEK;`TL#uD`DcX?$gCe3}NjX0t_{-{A{@ji06E)W&YMatX?tJ1tivF1{ai?`Ry z$Aq+Z90rd&F?e1x5zVsH zXu6&rz-4|LZk`?!O?Z|YNMC0O?=(r)Hkpj&ti^)KKa2cx$lrqe+2sGoJ&^wvq#-$? zsTArcb)EXN2l|D7?a%T5&-`iu88;*z+x7u2*=rafRPV zBY!3GyOTeX{ME>Rp8PM#zmEL%$QD7Zr1Cw`Yt&!yK%?eq;QyqVGUjAB%(O0uYnmw< zybBTaMi!A(8uwGxd#B;IfWu|{Ff|&@a+h1`St&MY8oXNdP83i6qAZ<|C>?qYm;I_$ zl{$NS<#KiU7`5Sw=j$O7UuW{HHMAXqQ2CjHQ&OF zE5rEh&3bnM>v|*o(96t+^Z*a!F!qzbFZq9b{%798obGl|B=oJ0PhQ2aC5wLZ_WCy8 zSC_)M_2ZSS-0>opZx}CGCdsxp_}z+0{pF`FiHcQ1B2=bDp}zQ3X`B=SA1pFGfmocpWPL+U^Ez(aOdma`d$-)E&8 zk00?y!MA6W_?SL8^E}+T#?m*UUw}Kd7BtBF>_y(~?IMePU#>jin>OR3^4OJrRawv^ zR@zxd$V2M7GTkp!F`NHNpuXa1H3W)YX}nFMB|P=0Wk?`YJ7y&_Y>H#*Xq9xY5rcbOCyPOL7N~c>D|#$ujk}!n zs&z~GB=@K?q+R`0a|uC+7r2a4!{h=3BU>d=DPOw=m!NO^l(8=d)PT?qkls$eNG*+vx$u zlmAv8`+V}xAb𝔡xuq`TVGb)Zg(y<(dL=>Y?qbHa+e<%Zo4hUGMgFrQg}dO50w+s4&4C9=moZtImbX4%1`gu?3Cf?ImNCuif{l zM&+i7_0ub%U|xHC+1DRa`gkF%YY2|VM`L&M1hlkO<3)FVV|*6x0fy@F;}ie6a)asZ z%^>?;)<75Nh0p^XLVg?a?a!KQF+%5yW0l>4+%Yd#_3pD*}Rwf`Akil`6f8f;00Qge?e+b||1o&6N`gaF>9$>j(HDLe2{}k*$ zvB_7k|6Tlz(2KpZI#KEfyDx8{^Q2+=3#6@wZTOWQz;mlXp8jwbD~|M*CY1*Xe(h^B#6Di#fHj(n5=Ew&%t)gBPM_BZ$zPvhgBEPFVlRwn= zaS=c|HahuR)LFBQU^!9yq`y zfPXpQ_X7MSfIk}W|91|!8t^uQ{h#1}xpf$?$`K{iRC`u;H&N18(u6smirhgVmTQ6k z3HqVM<_~m`POo$3<<@89EiZJW$04)$mk7q6h@o;P1wYB+VU*-)X(m1UULcogJJ8Fk z4=}FQhnpKu;C4p#{CutluN<|MS5=4bD@S1_)R$;pqXzyz1bQK%4zToX0$(`;{9h)- zK9#VRSAh!#{5F6;9q?NL{#k(E5%A{&{{Peg{!ekhnN1~f6P>bCKQ~%(y^e*_@b1&) zmanXslMUR%(6_9!`D={~_}vd_R{Ya~&GOAgW>TxwR9>BB$IBF_%NuPMO1GV2q@91v zC6#|A@^@>rskTcQb2RV8S3DojdpXSDiCsN->KT9jJtl+?&xqiyccS5qi&)<12OOY7 z0>AVkfzO=^_g^~f-6N2{UxBm!5aO9PaQ6fLL4aQb{2cIS0sdux{}>SKi`i+xRzH052}IL$nCM*cz!(3o}9oRKThDfb`bxh z!ulSDSP$@jd;s+?fL|5xe+CX{IN%=#_)`GCFJS))HVZIq06PKJ1oq!K;Qu@R|NagB zJN_UKQ5ppC_EU>Fd}fYnZOPt&d~o(hdB}AK{jm=SFXNSdwth* zspdqVf zEa7*41@j~3Fc-Bbnj5yk8yPdmdy| ze^UO|?l=4RAhhPdp*hm;Kz&~E(~Zv$EUk5|agi(z#!DSlmXhK%x%@!W#8YQVx3a>` z-S}q(6CQeCDnD!k6Z7+y@R*aq{7cUW{vtV=cexnL_gKesTLG}&N#KV$Xnw0<4J%+> z8v*luh=pE5{sZ`n0skA|0Qvy_AizHx@V^E7pWwh}917sWo440$z1=2H`sd0=4F)|m zXPH#4SI^YS+Sv2kMv~~^Cm%89JX_PZh!&=KNlj@(`6W$n{yWs9HfDsQB#TLqT*lZ+ zzjt4d@4UUbM$2y#Gd~EkR1O*O-500u^ZVU++r%Y&OGGeleHhLgmPGTz$7A^^i+G-5 z3Vjg3fiEL)rdPtasfC|z+^V0q)c*~8@Nl6q%Zt*XyDwmj2>#SHbA0#w!cKH#*EOm;st z%VE0oc2J_UGIp}mJ>aU`q|c|C-ow^2yLv@lv%rv>Kby=C_k+1g*A{c#iNSnPWjGI< z5zW!=SboAFo*&bL*azxBhl}vP6QB;P9M1SfhY?bK!Wq++G`l*y532YqRhPc-&~Pbn?*>udIRLEE_TN|>U|`~&?x#nUYE_e zkiuSg2TRd?yYf$mmh<2b{phG^c9JYRQ5sz{P1-T`wtV4`akW=(q_Kll|Jb`81~A9g zjz>hz=XL&z`S?#kJU1fW&^gO}(Z-e!3f%UhC^|u7<_rYucZwlBUu*+Z%z}|p;2Kxc_8|?p$16%P7;+;i< zq^p1?jME86jDWZf9%tLin%B zJZjv}N>aWLb)YZoC8O*2;62htOWZ7YG(u8!@06X zG{2P)%f0)-+X&jgf7B)LU^B@3V}J+T3wY{*_iO{4fPXXK|MM*odMKe62Jl-0egUwz zfLQ{*)nEs~{=9sN}Is8qO0y`Wf)8)_Q^;+v|k9)mQ~lug#b< zi;r=?!n)4?AaB|7NWOeghWzVrXUXx)5a~x|qSS;wvDzq6{?Df?ulf-$Jvmy*rg=p1 zJLN-|ud0RA<6ENS_sdxt_T!!Wz0=;>4@W}SkD*`L_0BQeliBhYJ6*ZighhPVm>}Lq zDV!f^iQ>0hV|nSXIDQp00IjM7USbG2R|NRJ62Jx;fC%^_0lzxnF9rNl0RJw)ZwUDH z0sH@y2C%nuD!knn1M_%2`N`9+QfYy$bYpTV&F~ME&iGG~&fI|5wrBuv8I>bnZ>-J3 z_9aTcgL7GLZO9F*Y?=3DQ_1{Hq7?gRjuaH~RUVXbuJ(85BKFSwBP%Q#&EGlNaP6+H ze9$F7-s4dqzXkJz$DNAe+a|#qVeN5zzXH_29ZldCda$Qgz#Xs~;<$5=2iyhz3-Dk2 z06ah7Hv;_GfPXmP9|`y`f?2>C$Aj$$y9jmN?E9PK=O- z8eL*BI};@}dtWKpa;cP%I)TqizasBdV!-X&6D2dN6>O0Rxm(&YMrRpF54{s5k3;jM z!qyIX=fQ`yn( zd1M9D*NukRb|G+{cL9e}2Qh04;CcbMKHzr({Az$d5%6mO{y~8K9M}ZF8w<7%?7uXS z|7ZaJ(ZK$9_kZSx05)>_3i&#ti98py&(?AAQs7Zzo_sZ0dN4FfdNs{O3O+rZ2M$z} z!gpi-9Ol9duC-vQXW&gMF`ca*KU(VFnjoE7;x1|A!MxmTKNIo^h zifh)*;TyjB@<}d%d`?Ul4@->V2ZzP-Yt3=o>|;D1us4BQjfArv1oa*};oP20pH6_+sCvcGFXd_7ehMcY$WZ+7W{fkeP+t3MNl41_K2kF-jnJ2!y~ABk|KS0s)?+& zC?hM@r;=$g^GTShDXDtYhg{zK47=N&#B1X>r4dO=M#iDa05_N2}#AimIg-YEzVYkf=;bwTe zAm=lL{kfxsBRjhaHs@~1hT4?N+M64q`&^htIRy%?NC*X(noLB@LYNV+n03U4K?y$=!h zhwF%EnJ?L4BM`m#{$#b&2OQUW8oykUj~k+5@RU1K@w<{y*uS3=9>trH72bywOw!Sq z0SnM?eKQogQ59|UeF^xh0e=qQ4;3qQZN%L+L&clnoq+!`;NK1SlL7y1VW^Ff(B9Nj zaH@U?_$vW_HsD_>J9CK0y7bkSg{XZ4{O19GA>dyL_-9(oSU=XHE=i3o*`g2qB?0_D zoT<;qFz%o}pJxVRh+l{D<>R{WV&7Y=``rUGZ;N9qQA)E8W$2lO)FNC-vhq$k}QSQW!UZ>kk(V!XJIb3 z|Gffl%D2TXLBnvH>p#?a?h2ZJeGiH$OhJ3y=AwX1W0Y{W7pnAqEbcEpCe|<5C`LSA zCVFZ!@#?sN;>xGrgO*?*X13zdh!#-&)7@zYBpkOF1tB2oPC>M$L48_VCg1`Oy94Gj#Me5MCwibGHvI-}#K}v^+_+#%?B$9D+#kT}u+SWH51z`i)iYT*QqfJF#oO)mV)?V;MHU z9mXnn#DM$ghyD>1ym14{@?DH(MItoHMhiVP`V8w|2kXBB)<0f^{V&4)7h(U4u>WEG z%VGU9VEsJ>*#83Te*yMCtbYxx|5jN4a2f1>8SH-BUW|p0k3aQ*5@|R;KhYf}Jj&%!ZSZ ztY%mj=JWbG-C(_!Zd{Q@r~h)JQ>Txo3QzmezxFSPa$XhD%icnWY6ux?ZB2e>3?bPw z{$MZ1M*QIRF6^zGjQ?D8!Qs({IK-qUUQzh~^^81Bnutefd8Z zh$UDQPc!x6(l0CiC0U1GfOkP0${#TIykl%g%qAB3IFJR6uwvHV2D8p19dtvNMk*ZI zMVr5_rdCxh^iilGm5h4QgkulL{ph1acV`yK9_&v_?qc%GYXEuh>kFQje-2v&72q1z zMBIMP9y?thgCh#N;XQ75(6Y<}DDG}1I^w$!?JG7%^$Yr;@aJz~|I10n?F!;SrI2yssEEN7;MGNpPBM+*Ms@-*Kh|Xhw|{2xqQqBk*8Qd?b~z( ze!%z!GdJANRz1#OPR(Ad?XWpJeyAUFu6#$KoSF`NokvIOMA1TpDb#cPNb2)Sk$ziv zop{ugky*b}$yhx%lAUi#b~*JW^`6hM(w&pI`pag#P%RjnhFRew&jw>=*?~GDFQG+x zyO2C)HL`o@jB3vsAk%ymG^qZ*c;Unm@p#uP@#K&t;;CL(tg6-$3)?>n%5`BVN zvvhP|T2Z>}w1F~vquxqyv>c=(&#j{mGkxjNi2|J;-k;8$`hi&7K27#5$|oTaF{JP9 zsbosYD5B9%i8Qcg+z;==uZ+`i?*R*NyS^D-x=|Hp_`XCAk=1CVO%8fu6pD`O+Mr!F zLs4^hC!GJwaQ=70`A>%PKU)O4nbi=X(7&Z|35f37H`qcT>~bPETX zrl?E5=60t=A6m%88x_Rz$yySi=0ooPwjel6gCv)9rKd*|yV}qvPW62}& zTD>kA!Sb$(Fv{?O%&W&y*}h3xvLz`?WLjQW_QST9WVRpX zQfSX%^_m`1ps-Lf`;sIte$>E@eYNLHSH4sXDJU?Lk=e+>Tb82G=9;Qy8htJIk= z*%yO#--f#33ifc%I@av!$G&-pEPTHvldbtgM_f8f z*Y(*!J0`|c{Q)zmR{Urhlhu{7eYePqaph#e*$i@JmnT^!Oe7zWI{Bma2LCIl!S(-g z@#bG)c&h&-JV*rJ$rF=-Q!Zb2ij@UR0}N+M{~|QxyD~pB!zfo?WEjttieq`NzKi&$Ra5wr zOZvPU>cdT^y<$dtYgwGmHn!?gG|THUjp-Tau$>i3>|D@I>OW&Y-JzX9mlb)?+kq3P z_jWb9M)Ng^zkP~0dFB$GpDRenHd|5?IE+x&e|X=yD|p=XJ=nJ}1#fhlix*}ZSPI=^5e3VXf`d1y0qW!ylN`1CvY{|n&%i@^UUf&X_DA^s5|{t+Sm0snsx z{Qo-e|GommKLW%*0>nSy|4)Pe&jEz=oshvl1lNSa+{^BI%=GS(uK=6?^9E-Ez1 zSNuq4YaSDx2ffS77eYODLJ+rj=*)}ynDY-S2Jm-Xelx3Wmsxab2`f-bW&O9#X92~g z%qFKVb5?&rb<(S7!Imv_rCtdAtw!vhZDde_WzS@SBwbu$u8#6n5?$;#mdg<+?<4>VrMH_+<>jr@Enq zX?Mhkg9pU4@0nuKriG$+y}4)@*G~-6dMDh0?)n`Lc|x65q~JbuvJe+GLeRHY6pDPC zWE@Der=-k*;_|)+c=xH%pT1D4q3uGs)2*N z7r}S*^xy`ioVPoThKDU$s?!QX7|FUJa(PMdh+grIzt57}% z&iUUy#&oT1fE2xTmUOJfQ2w(1HRFzRctUv`mk)^JDc=|IoOM&V;#C9wbyQ!T()OBp z{g$&Nl+Rvmie)HrI$Q5Lng#6Y%Az`M(dU`v)O%1SHMjJl{p1sA;XZXb_Tn3|)x4HW zj@e2ko(?C!RqV*yoZ+N#n*y0X=^8dtF2!bfYw)=DuDI!%34ZgSH&*TW6q$cJj+O*% zLg57g=u9j}<7N&*eGdN=|1EA1-zF4`4|GqPo7^TGdwxh5wQjwz zRPqx(>4-wfcuiqY*hg8?xihlGW46oGmu!+>N!rbNCCrj$I(qQ(q0ia*GA=m}36*Z! zo|SJ(-ohG&S@MOJ5j+q$ngKb1yyV0z?*C*0cUh>#Wut#Gi*AkVp2Z&at0RS_-*jc0 z%*HW|fIh6s^cn5!dXg@5-%KZt3!-TamNdz8F#Qtvo9NuQNbv5RL{WV;DWlFL78#H{ zV--@RbswM6KZ0F0Zor|wi?K~4!h>wIaGTL*bWpbr?X=l}u7}5?T0R5q%N>m%{znl1 z!~I_l_kRZ5|DGbm|02ZyBEuV*h*WZL-?HmyxCQt zFL&t@BZVkX>B*(gS`^lUYaH?9j@uLXqzln}^f7N9D7WPU^~dt+Un+d>(I>3)=SgO5 zlfym*uV80x+cMRk!+ zo1`WzNqGhF&ncNd_^O7_=ULmv#qzyTo$~%y;<)NTQ`T^0p|sV)Lvq~^EMLo-*rj_8 z{9un*-T?7(QM5n57CVF6mK*bsDm6aW`90fNbC&tcEnrLMu4I#o99Z8fUFPPg%)Tya zrExIoO|emV6+GMxG%FyYfjwOAJvTGmU(HK8pBEQX(bhz!x?CVphv8v!z}oY}Lm#Y_HFJb}Pb^)vf5u((XQ|{oJdlRmc{4r#+aKn_AOb zNkiyg>px_{^hUyO?jl{fCzEvzF63C0A?ah%lk}{7fSV(Z;+ldi+(XkJw>4wj&3ypg z`soWw+k6fc`xl`74vA<>tvy=MI0ixfiy;4n`0oJ3f0+>fEfgXD6(Ro>A^(N=PY&^4 z9>jkU0_48}qLX={BohP{AT$22Sn>KZeGO%y*3@ z?9We>RQ2wjg6vN@p5tRK_;HuJ5cN_{UP{ltQRibC=&GZl7ztz9R)ilJwbPuvT$x&t4y<- zdW~JiB59{?JTJ-ckYwhT zs=WBA4c2^%#%P{4tOvJAdC1<*JkEZrZf4T$Af^U2G8A{#YYwTicUH7CaAFnYGBuuHqIP^R=;_M+HTT6h9x0B$$q%ho7LNo5)GpeEqm zE2D_ELy**>;3&QJA1bf<@t)ntoy)I5-%7xZSY8>hhV zJH2HayOt8ejPs|lUwd@ey?$L;z4|RGUfWOgR%Ot=(>$rx&k0nympXk~|C;bcHN>zv zmpnKUM&6oEB8snukvn}9h@tmYJjk&Gmv>u(mngVmnVJb+W!oE@t3N^IM#s^$D;rVQ zZOhTay$q@D8Hl2beu%Tz){Aeh6p3g{l6bzwQ8bvLC-y2-7Iika3fFZG3C}*R6Smg* z3QF??Vfgy~!n)-jWPh3$OP{SiL2buxy0-l7sg=V=OW8qTJh45SrEQxnsSjT)rH;*! z-Me`J{o5CZCS=|DLQa%LsNL8pdYVepP56?3oO2WF{{a0#THF-X3`M@ zHnu^98O?Y=YpsvcnJHP+qjL$pdlA!0*8$Yx(-(3*=N!4SyntBENhC#e_GDMn7&5r9 z8=2&G2aA~pu=m|eoa?&~Z!b2-(t>_C@cCQhrY%R8$K|1frxB>gKv z51XpLNNOEu&nvg};BV5$OXI&pNd1yq<@0{7WNFTb>)VF$wsfe&F$v}e&$#ekU(ETf z$=duo1eo?t949Kgnd)oH@j{%7}QF_9C7xkMNuem3YvBY&?pa!tv7Me-7)Ng|pH9zz(V>?{cBR*T+#~~w_LHcY8ARc= z2YFUJffS!nBT@FR@e0Kn+_W+me>Doji_Y0%A!`_Zxb7c1{^1I$|FQ>(r&5q-+FTTV z&={qC?}d^!J%;?}800@2A^%w}-r-E#;V@9F)A|AVPd(&6-~;WGApda`Hon&t3Y?XN zzM#i$ZCTD{o{p5-KwDEgvs}9MEsn*w`|$^HyX4pZBuE#yfn@yq45=^F;@KLD__H;j z$(BX(eKwx(PKPD`o;8YF)OY7^#QQA8r;-K6Y-DOYfPHP`Z1m$nY|PhRG$rpMJ->D* zJ*Kjn_GHdf*UEs}Tc}W^cb^;*kC5R-8_2%c#Uv>MktegX$OgO5cmt`!Yg~5VQQ7hM z((D;HeD7%7`*2t6*LDk?pIwf=Kg>WGH$2hS1rt$`t2)|X`9_?0zDA52wpHvj2p83| zCy6@d!^IoZ6~u?uN`xtPYlMU)u0mg56G1Jlw_s42!NT_Ul)lX!CS5z{!RP6k zN{ix8)*i`<|tBgNL7s=Zp%g==?Pz<&=*;HoL0nJfD8@bPxsf1NI5 zRlWGw+s~LmT{U}Lo6FSihqFW1?AW2~5$xh-MHbu>`0@2+)b&X!?W5{OgVvkUx@moB zj_Y&6n@^InpEeVxzQJVh3M(?|**uRrjr#lO z?1>|&vTGJPK4b~1?uF6GYAsaI{u%PWI>`TaK>n8?0{Z~ zc?rOO3c!B~ub=8mWKyi8Y`vkT`#VK$v2d?^!irn+l5S9kaH^e}3|TC_pEpmcIyFlk z5Da}iCui|ZP`kGIVjQpYTfz%u(|MObBYp>Fu2?Na*U(?_18(*vtei2fGzcu_FAS7n8+4;g|EZ0iv3 zp1vdw_1Pu%*|%Cup5P*mQ!o^dTvicZXgv`6|2!h33RwaPULtgNz{0nWTEgxRaZrtA#LYt1L`bO-!Um(9X7m|j&m1MfD1KH7Cm%Q5Aojh)D!DZJf@T~`H@lO>W ztpCjddk1OYz=C!pkCl+g%x$Rm;VATX@f7qXVI+E>(*>2wuS5R37xLdU$ba3$W6w>+ zQS17OOQq+K|5ic%y9M&!U?J(8m9Tis5TW~nt8$|QgXDp(;r#4_4*4E<1Ety8PfF=O zlZVAB^Eorkr0XXmqz~42x`M_XaQsTmwx0v%N6K$S2rGvE} zxWW!DFJ-;2q_T=eH|Cx{o=IM+?3(LKT9jQ))tBYaTPs6pVYm&woH~@+zwRV8g_lXz zi`}IE(`53%Xbx%3H6kIJy~uC9N7#K=C4QTgjotH?V!yr=C(P5vLo&XhQ}@rKx1oh7 zKYk?|q2+*-9*;$fV!ES*nif%BS|R4k*NWe7d5gHwLfmMtA+|4W7p^%-!g+&j!mr3E zp;Me9w9XqT*cjD~Lb-4fSV>R`5Z6=J9eb#67%*@JdHTe&=;Fv))n8hW1;_exLJZL3T2h z`%{CRxbU9d&OS{aTj$e1jxjXs>{Qx)`zU&>w-PlFo5@IFAMr3sCqbGEh^D?7`JJsw z`uo1b>XFr0!6pZvHVVbXx;A*5%}{(Pyc1RN%cv}OH#)pC8Cjj1jfU75A>jWJ@c)qi zS3>@u4f+345%_-*_%<JA&sFsmceftyR!uR;Xw9{DY+cvBtkH+^cr{d$rqwuLiUGVYZ z8_2I_FS<~lhI~TZk@=qSC`?xsjXwE8^mtb#>Njr@mqvw%lXI-a(5pklZzul<>7|W= z($L*P8&4L#jh-#Y&l(CpLX>&0d|RHf>uBwnzaG+u2Qj>O_&jN#-D^629wKi>K8w(@ZCN_n+pBl+drF3C!8*b2V z4ST7=qco^9bf@+$9f>8)H$@oBgrUJGp7i#%cwg zX;&)W-T#hUfWDlq3qDiVSBs>6>*h(0?auP@zpdd`>s#pU!U~#Gzm_I{^`_SEENISh z4VqlgPN2+^EO**Qww;I~p97{4&E%0p(Cb1Tp1+Rg^xBKZifMSmS~uMDohi25*cVsS zKS$dQs!(^?7IZN#7_DryLOl(Kplydcfd9D!{Le1ne^vwk<069oC4&AXg8l{k&r#ri zvVi|tB7puSypfNVlAa|>&X;|q+3FoEUgZuA^L~#1xs8@95m85e10$zo|- zpRMvYl_!{E&{Q5XA)22YpTHNNfIP^+fwwuE@=qf*xO(Dew$}0ji?J$Z`%kT6JJ&d~ zy-y9;xeqD~uYEuxBaYIZ2earAOMg0AfzXKL0rafwE2+44j_lo5K!z42lCl5n$u@IjpyU z{=LWYp4goqa4VD_>`}oU)w@XB_q+3_&9B*Q%A~x#q0*xp=j1oMH?bXSI8W{q!6Qwf zH~CB;-|>DH|1fYOr<=9-%ebFRbZlg%;d_|rxHW9J-aHl^F^*~N?!zk5pV3=0PExm= z&GgjTAo^LuirQ}2)B2S*~Bo4N#$-FFQGBn$D z*NB_D4ThO~I?|Gm7-`+J4{{Ri&fWm;C7Tt>i=a+_&-xJVykQRCcv%KBY6kLUL4Vov zc}=Wp&t9fdoX%S7J(y>YiL8CDIxAZ9hJGAXOVi`G(!!Q-TA^!42UZTJs>c-Q`p|1+ zpm8Z#e{u~;A32W{|1u$W6#9_wmQV4PktZ5 zYHW0gKEJIL_`hatJ{%gxJ$ujR3zlL2+I1*5G3>$@Ew}|URSz%&^>yrmpD%OF6PRr; zP1bw$M|%F}8M@JTJM|bHN9W+_^j@P5EmQ<9tnMbUkJwKdjWS5p5f7phF@bm$sFBkC zukqgIQ~0f0E{^-O0>^B&#oPRc;av{@P*&{~-4tU4`Y!_g7x=$p z!2fLo{%@HG`mYH3uL$}t@P8M8|0@FiFA4ZRN1-ENlzjJYs4JcbbC-^H;U*9R5a`eT zH#%Ah%c^1LVqN)H^+tK~heT;v)o|&BK{{FefEFF^jb-E@jj8Dbto~vrh@%sq@tfG;(?o?T{tWS2G>yrl-18f2J~> zquEM;Z6wOY>&OcmU$W=3jJ(GE$=p}(aZSo;ynS>&KG8D}rj{<){>cp# z@^LS!`jmz`j=7_QZR3%_l4*^rdk}Vm?K7>4H3Uww-(d>4iUq;cM5-=Hwtqc zb_+|FBnwW9W(#`7OQl{@b$Mv`T;5ZhN%jORkvhxb`O?px(!Z_4Svn2i*YBsuk8X;W zatp^xT_d{G^qZm1OHTUm-xy-Rl`zvZ&x=1^W)1z#I^0~L2cPHkkVWboXM^8uVk+B% znB7S$wm5hQn-=qj=I&~w6U~4NT#!r`d(Wnp9fq`gTTgmq(L-{eMgOqngC8~l|5*n7XDaZY^99onE>e2FRQpHq z5PNO&TmEB=o^*Y;508Yp_P4I9Yp<+^Srz3!YVX7Bi^+eiq>qR8)+XC`<(g*~@E_I* z++kJ>zggG^wzKA=I?Zi35 zR^!XFobk}<2KWk6!S#msQ9e3?+6p%yk?GRV5ix-<` z%@BK+juy2Kb`?>}Ey3KmTzGyzL$J8!DHJPZ6Q}SglHWf_D^lSUX!q&p|5)43{d z$iHnhnAkboX6dve=YU3gVgJHqLd!1$R5AG%Df*$$QwTWkUukx=XbO9pr6fK`jojqy5zB3 zZglfH8#Kw0UmG3E&zZ;bQO^E+-pm<%U7|7X%u?gNecrR$q_b?x*aEh2+)Ae8>A?CY z>$32P$}E3mD~*4Ckh)!6M?JRtQUz;)YOL%}3+I0zI?qm%XCe6{JuQY5KbuN=92rG2 z1}hPjna$W^@;-cBmX4Kl7eGp9hA(Yb#g^eO(S_t{)M=lCJXt8ZY;J=T9EKvDv`#T$ z=4EmHuHE94@?`O7^K5a0gOT{HwWk<;>7menZl$o$Ia?SkEETr3E#kMHYe~alwvleA zr_>{+8;^gJPrsJ`B%Pi5(%@C0{B1@)tC;I7-Gg4|zef4;YGvkSd?5Al@w)fiuy0!5rm4C^i zs}K6qgn5J}TpmCl@Bd0{jL(w+FAB)*-Ydzlg)_-|hq1)FTX!;8p#}FV( zZ(M9-fiGU^hyUiiL-+Q|QQtjzD5@wDxvZUxUSAo3&=y5>zNHEH?=s-OQ-S}UFJ9L% z6`y_VBW|sG2K@I);J-Hm{~aW}(&(hK>&%$jsmYSY5G84mnKOUopCY$+&tviD9Hgnn zUc6%aL#8p!N;)1EB2Aty$#-<-vWBac{ANr9@2&;?iawy9M$Y25iYIWrv08jg?GM(c z;u2f&em9Fgox%cExU$1fO_;H9ALeuMDNQ?kf~pMMOta1g(htFw^ojOhI;8U#@rt-e zrp(_-7T;P$^6t+f8twW-{Lq7#9Jz<(Er)SWrwtf&EW%Nv5U&5LiL2$G(8sH1(eOXp zQC63D)b?aLx-nxk@?O*x_4c_X)_>hEUd_%BH@SF<8b2q9H@d5fCu&{`U*^{c_ZxBr zh2k(_PPwfV-sr+7cSOnm>|0jT^cv=~+o?<49`)iLk;|l@3wAtwtty|t*GM|x7$wb~ z^F)q4!`Qxs0^h!A1$TtI)A&#@U$}HO)Q?;6buYB}kc>afUw)N2Ih3)*scGy}hdayP zXvX^KtFgxIujq{2QnsRFebxyOTR(`{%SWeT)N(kKxh<8}aPt%kWHXhR=^1h$EhUN8P_(K;xW?&}_3Lky#;QxI^sQ(e6{zruRAK?E_1OJ~7{C~7C3g&mH_lV)D`WKl-#4k1=A&S?E z4$_i6mb@BBrjkLre8M~#?lN}}*@_8K`~*^R44c<`jN@LoyURIU~iE@i7O=3z~VQi*;S9p630ev3}pECLPOzN|y7)KhU#!J&Z5!>Bgr=5biQ-7?-*$@zKqQtNaYj`*{;gz(&t?zD(K!m(lwCtkilr!E^BNTV z!WFH*XoA}A_C}8tpMw5z9Q2P(pnn92En!?Nm@-K0U-lF9j|R{`ib4NaCFpE6PRO{)1!}zlxfAbR`Sc>5ZU=-9Z9_AOUC#JBs-@+35)rF<2z2{ z_|^G%&#oBU^lK_UeRdSSt*3pm1cDINVJ3sBS)Go-X%73r;gDWcqJvDH0C zbeh%ULd2 zm%_xeOFYbp94FnMWLrHtr60Gm_2oNff@b0v#h*?0;_C-k z^C`tT{PRg=?sew@=wHWJTka+{s$UT6GsTK|YYt(ZFFI(pQ6n97V;415OQzn>U1(U2 zAzePTCq4P}0U23%l$e~(B44Nblf4Rr_^lj3QjET0qjTr*z^np%dR-ze`CyObU&dhb zQ{Avz+8wmw-~qJednQWUv=A+0{U;CfpGXnv-$bZ?6QTZ1g!(to ze@a3BNfnMBzr@-r!=%P>s%%D*5`TRDsC-2vlPdZJ@~ibW($(HACtrp~^65%rtZS{yA+%waY@7R*Sv#&?<&<1b*u-TUHfpekOyC)yk{Diq_R56u#Ic(eJ z73}gCTed@WICF-+>SgNJ=)LbHbZOoin*7I=9(rOzzkcdXV*DPj!g0@?aP4S)yxG174nEL^R(C&)HkPhO%5VJ8G&>PF z+iRj@MIXhmS7$^o>+Rx%JF#NHr)lES3>~rRqLR2_(M`cof4^|2JzcocTUCl!{z9&o z@k{QxKAD^tK2d5O_6>yW>AbCdkkmlUj5dm$CES z-T3#RWIU^FHa<1m2!DFm6Q|sGh_)=KL<`qR1} zmCN6fjE-7zd0-xKIuk*R@7j@PbOf0{QjrM5n(&8XrMSn~RIGP>976L)LyDVpZdanTuvW$lcPWWdmIu_;z^HvzL23pRmuLPd_$;-+N}vOW&#SSNq>H5aXEXmI5{; zVcxw%7V@i=O_83j>eLT zh-u`ct`6~uQ6ev2G~>*(`>MM%_myG z4AZ)0e5UBYExk>-$~X;vVasP$>3o5CIu*0y4XfCjd}sFQs{yO)RAH?y+aM&a&&H49vW>Mi9T3OMqjc=px3^N=w(KexV%r97_&T8++#go z9P?zH_~TR`aYo=XVM_m#Ld~4b!b_tBZaDKaTMm=Eo1S?}C+-ZFCU&*vWAx1Ai$*MG zFR$83N#%aLQ@F~e+t^97or0tx-;T+D9V}$MKUnjhMv;6l%&++e^Bo&2ocNAAW?Vf$ zlmChN&W`W6$XM^)tf+emGq#<}=CvEMoqoO9(KS!#$MWM;Znud#1qDz)PfoqaAXgc$W8Hc@T3uI&&uO<9lk9`VDK14X<{UlX@3`iNSp z&Y-3t+tL1Qv8d(rG&ICV2ld&fgpw!R1pRYA=%4AJe|m`hf0~J@f|^KzUxEI43iQt$ z!A4Ps7mv9v&)KBGXTDI9;(sK{2j&LCY~lxU70XBL!UQj=+0dEu<%(R-%|hA~7A_g3 z-jJ`pp31t|k^hUM^Ny?ed;EAM?X*KFsYpta#(kfe6hacIl(dLaQrgoS1vGl-l-juA2qJ=u1^q#XFt*sbJ z5q!(?ibf9gIhErpQzr<|ZE%H09l6bG)L=v|x zBELH)kxRB3g%{C5`&@@g`a#dk9I2v(S1(hnkfWN++g04{oaQ@y%Sj+W$v)Ou3w(bWQSb7b-PTvn6o2&&@5)TkD*c@>F z8F2oY-~ZO~``^v{{&$JM`DcOi&jRP4`Tg&Ce*e2$ls+$kt+H9d=9)%RyNBP&yrLoW zM9(VP_|}VEshLJ!*rc(;2S?G5Io9<2R}EIaHI1#_FU#J0TM^}9Ry6lq6m|N}+dXEh z=;Bj8G+YkQr79!ncR6XgI^qtwWqXv2Y|0@iqF`cw*PaaCtwp|Cf5ctG&SCm&Cr+tP z!qHP)@jg3aY;jfw8=kn$&sr+cm-dZlO!h)_|MWz(-a!q`y8RLsOP_^a-Gy-6_Gnmr z)CsQd8Ucsj`wf)Yb?|jUIT+=h2@3Fhu>GC|NJvrwzT=+?i4W=p-w#`bqXl6?mw*Jj z>Y;*C{1*}ZcS*F5?G=3jn6X&BmTe1A;)4V-vFn1zK8~h4C(V;8ESv@9wBe5 ziN*Ika`y2OB7a>V>VtGiuhu8Lm%qb*3)_P=|4YM*x6HyaPfYR7C$f0L*n23n@F*Iq zor4aT1){@cwkUV21{&D;7ET-@h9L_|;J;0A@aSU~*c~$x7N1z=$_oWhs5rL7ev1k%Wm`iPhRL!av$bci1pe^X)5k z47`H(d@99meXH^8NA5T<%M8m8lE-Lw2kP;!LH={|(8$XnC^^>wZF{4IDndTO8+zwp zRO?Q7FCz&a8R7~j4>pFQw#q=u$lGAZwn}j5;6@M%7lNL|i9j+)4HTrm6#k0O2>%@` z6egf(LA=CC@H{j^c$4&7)EIkRr1j;1=x)|>YJd~yAvTuXGYF)WRdG~BCZ2p3C9-wA zxsW7^WQUfnCWg{WXxGVd$vTA$wxida`CV&4+7EQ8nPn*5$9XtXnm}Fc0;u#GM=Go} zrlx1*=vI@*q8aZ8Wbb^OHV*EM6ilrp^W*IGQmY#u&iH4aCc zDdOqsk5R~&<7n&FEof)(QZzFjpdC)S$kXx@TxNC=rcT)dtK!n&VLA)$D>j8k%VZ(G z)DHBf90l)gWdrxK0YG(zE#Uk=m|*sn-~Tu9`~MPt{~srC{$Jqyzrgu_f%E_T{=Y(` zTs4pH^7_%to5{>}rYv=+aiYl=uZZs`E@Doblj*itA66zj!)=u-Y0=9_Nn=GSn{92! z6!b1MN$(j#>(>QSb#>lpdX-319fRlylPUB^sVNQHr9h4Ry2vN#26FAdc2aw7CHZsI zi9EA6AgdeyVADe_xN*=y+$gsmE3ERus`o8%$4q7Hk=ccwemsFLw&o*kspV)yAV#IT zh9T*%U!kn+6_|ax6rMY|8dj`yhkwn@pul+aCN1c$e6qV;1dMK{-UG@bntND?oHvTO2X=(*)1 z(pmpNa(eb8X2P3gzdI4DDUcUGGDxD!=Bi7Ub|kXetCSrYwOxGDT#h>LUqtmHc*kEm znfAaC>gnl1ADtdUha@Y}?kz7!Vci*W`B4#R2#+Nr%Up=!o{{8As}x!K-%UI*_Ym&Z z*nsa#7T|>$6EIz=ioJilK)wwP=$p!Rw0X}8^ug5;9sHt)j){N5QpaXkG`UboCs4w(-Oy z$r8tOsy^T%{(Ep0yIbYTx`V%qwFdnn_tyK;4*s5Tc}xloGFw9PQ>W3JEn}&QmkMn= z|B{>vXd*u}caXQriR4z%bn;!sm{cJp zgc;vUVJ@F+NPvnTGkkJYvR*Zm#hhO%ar7BYe>A?4WWC)d`CnI5Q^f&GYe%w-J9F6| zx$%-q?r%tnjTc>cFqJ+tNTt_hmeMo3Tcyj0X1WdW&XG7bEXYCESaML%9Z1QBB22 zwDyn`y1U>etQmR;uCB;}`vxt5!=72ed+$}?YNHomrhEe!NGkwePDX<2%@joW>w&W7 zADsWV#`%x^oc~zM`410)`!53bUj**IaQ>rKj+?(&tq*W39%eM?vDj=KVna*+( z0L~imfww;kB{QqLC9A#Csa2^8>9SnHOzi#G+Ry%y^^W(6>tPQnw~n(Ayk#`Fbs0TT zH2NKk9p7f>f;a{gSNN4(NV%R!{jNC6rA|`cU%h5G>{G>drS{8!G zCpzH3$6DCe^CP;TehzK8x)XhkOG3W#uIS{yQRqRA42lZ94I6VSVfg-yaIft`7{6*F z4Ei?&ZcTUzyicD2q_Pk^wv7h&0-QiW*$6;mehXVyTo?3u4+xL-GKKkXJcZ-iEQAYA zN`h7U6H$16y{O(`D*doyh@|soI&Jg{WYMWHtT;1?^`GBIb}3{M&rkj0!Ot|<6~5j3 zX-YWx{mYNNxgN|aEjLJ(c3vVYUeBU8?33xjvuX7G#&8-Yaighkt>}pk4LV_JA9)&h ziPY4Wk_mBZh;57qnLo*#9LiT97G9m$|5zQiAzN_6hfuuC2;fF(UF>zOA31bfKn;WU zpa`urG_HRpI^$!CtX9e*rB&_ls`ODfe|I+Ay(9qM9c&BVX=uRxt#82MuqLqPZZWu8 z9}E6kyMR-#48aOjDKIanmGdtZoPWvU{EN3RSBE3JSCEEWWHTVZ}pXt&6_=#BrlAuS1pyC9(SA= ze3?f7oJ*ve19-=OQUpDB-kln}wdNf+ExIV;Gr6L0m6W{NPgI=O5uGKT#J$9VsI65Z zR}-J&RJ#*+f@?nZ62fumHH5Xt48yb0S5&$GGU}dEilTj1qaP1vBUU#Ct@D&a2QwbP zyYbb~u{am*C<=y?O6*~rsTMpw`U6<0$H3)(J3y^XA~2Jg4jx*L0^(=VK-aTP$gVjo zjO^VYOw#ca+H$Ogi8^%MmTyF(rkNblD<=t>>q(loH&OgEo_NKm5T{+w@h0&pZ0uQp-)BbRTYMrr zY=$0ouKj`Z!>=LDwfm8U*;;h@k_Wn=WR5yS3W(%Bf?e-w;ka#^Ve!=^aO*e$N;&Aj z3#ok|_4;|BXt5iV)TIFTJ2SwaXcNHwFTnjT&i_<#{$~^Ce-;Vc{}Qs-!Y<$yw*cxxs8JkN=k?gmCRj0Ov` z2ZPHXZXN zvO!Bw!;4fp-&luDrLZPrnCxLb?_Pd*?toK+prB5 zm`A`_?gZ*A91dTM`wj}nUj;tE~*MQQvIl%e1892O19t2oE72+T=ZYj!---qQ$zZE`c?apL7s&{%Blm=K+OI~Ll#dOaW|7E_ z>rSN4%uK1@xe*LKN?|j|fP_7BB}vvcRK-1-I`Cd~o_!20kMgD$>P1xMwjT9=`Hu`K zX(KzUkB}e9+2oc*AgMREBLPmDoC|r6>3j)hu{-d}r3u(MY#LtBV1$*jq;c5ATc{Wv zMm_}_kmXumbW_I~J?v3Ms>i#b?99{f&D!m7t8WxM^2`yIuF!{H*9?IDPnyA(W>mqxO6uAfjnE~ zZA?D6TGD9MD0-0J1C7q(ThdMo=;0DT)3%SG8R^or#^Ek8d2*C&oR&*mMud>1c@AWR zjW)Sx(u<`YpTjvByYS=V$v9AB1}=YOjQ`R>xIO0%I+uI|X#{UVKbQNXJ1cCE)qiT} zkc9KV=gz`AX+>~zP7G9todV6*7(ls>Kj7Gw7NGl}9F*S61ieGOKu(q=;Qk}v{v+qV zPH_G!pYvbg0{0&U?mr6Lf8_kvWzlAS)-ivnEqiRXUh+bFFm*V*h{nG6qn`p?so|zL z_I1=!YIHM$m21e*V^uEnYsOE>wcHFQOEsB|-d8-`VGJEs9>G6qd9Us^{|>C;M^zFL z{XWl-S}P2qJKOJ(ODZ*FuhwQFsDzS^K7pJXrAs!se8Sa>FJiasJ(&DW!*OS3;dlQ` zasN+QOg-s~Vc8eL!@g5O_woYadP$_f^Yg;$8G3?Y)eq66 z5PlD~eGu)Peh~%roFj)48G@bsK7ViuwKp0{XRYL2^e^!=q?v9+lg?yi@^hevcEt?$)n&A(i5x>6f(27#;N|`{u-B>*Y?!|h(6fHvw)#Zy ztY!#U{=J9uZ)Z6FR>=9cm4f8+WFc|p2;ouBFV4UHCu$0QAsPM%(ed^R>A@9WCC`?{ z&~;(FOAG4B@6Vij{+muSx5u#N=_RDjFOb5XddZ7*{7mHcXlDJ`9QPLtr?>o;(lM2M ztAo4K2UjnmDq|e!rRGtzQdf>zt8|jdOY4c}k$h78b~$->ijZ_EJ+gn}4}5dgHC$7= zA4ktvi?=Gy#nIEv@!d!Ttk&}gLAyF6YTb-{%|p?nBms?2&_S`?eb9E~1-Na+ZkT15 z3XLsh!sgQ^P``H&yj^w|JYP@+ersogr8$ehxA~KRb*4HP|Ms;oqpndH<5(;lzO+hk zTj?xB%`y~r82uAf6yFe~MWwTg@qBNATP>tTt%{6a7x7X<>_5 zVmBsV?uS!@WWJ9*Je>tioxs*S)PyN=TJ&PtVmi!=x5r^3yY{Dmj~`VV^?Ie<0#*5SfBPh1;efiIXT;n-tO zP|&J+)Vq5t8m$$Ep4^6Ljo(nD+W#3kZodRm!}h{V_YAmN>;~sv9}UZQ4u((W-3O!B z9|O^MbHLVxL16eUJHY*4(AW8n^M4Y~|Lx%XU%bHmUxE9-0{4Ff?*EFs9Y?Xx-%RP^ zx1#^;O|a1`uJX@&E1DENSyI+LkCpfG&CV?zY+PVI{x)+J)m`H$sXLd#JjaNbc>PfE z;(03cY*G+SdX-My5brv-1W`}%6lywYG+kw=NGsqo@?G3ORA?c2awVGdq)Z{lpBRvl z*Z<rNJwqGoPNIX$x1o=}m!mEVg6`4b=+K{U zaF*g#7`k{LbaY$;AG^n|$b+ka?BBE6Cinf<$#0z= z$x8PLZ0hi()H)onx_&#c)WIaG(y1WPPUB}oMUJfS?=Epb>R@WiyAfXg>D0l3vvMIJ z)MK>^9r|Jn{jXGs9=`R0JbQJ9R45b^yA!eG`(qcf{l-Y5`Av$vQoMy@*&#d`Zoqlc zzS#8S1iWjPDh8w*9qVmC3&D2O@Nxxuk?DxWnd&1)r2*)j-VEmiltK1jErieK!exEt z@I=1?{C2Vv9DP&=eCKTeU4KJC&}abkq;*03g?`RIUf}%WZq7fZ3BP-03SRD}LijRS z!6UR?)buieCVS6dr>grUK1qYw@12h%GJErIn6)L_FUYVzKisKr%t&@y)kHF=J&oSF z7A+%|r+c6QZpRa~@%DqBu^L<^5Y`g(F2Q*T(t9@(}w z9lhtoUO(9=c`CJlojkZllIdzg#dgvxS;&zz+|Hn1W7i_B$x&>9)*PmB(?Ejnd?XJy zd(-pZQmOyl6q*qoO3!|oPQ8V3bjlu8YWVdvxhRs589R59$$7~H@1H@skC>1SU0EJ0 zYR3_GtMG!XY#iwkfOGy#!YkA?@Q>Cv=sHekI!0vxU(prLX&_nbB@6McY*uw0{7oVH!j5=|C)FpfF3G>JD14557o z&q~sAbfU|J@#zY z#Pf^aqoGG6D5YWtQZG+HueMA>W7>?+;1X$+^!OGWGUqUaWgDR94qqrY)*8BfQHA== z-5|mLG&r5V9W0260@uGdg2CzffbAR*c719V-Zhm8MyBh8D!+LGZ}tjVj*7w`+sC3S zOCwmwB|^SV8_jlk`qKLCY3zZ^99GWHx$CqB5iqxv{CXHCF`e^6lE{1QTlLLJ%7$Pz zuh^fBtcZ|IInze0$sFFoNujUo`KI^$FdE=7i)OAJPmdL-Q$@{>pUUW~Zm$>oG|GL6H?Cd}s zb0-Px#;_W4-*+XcWdDf=8R@Xl>_pl(GnRaNEUl~bd>P~{@Jfd_w4B-TaH{J&3&cBVL1DY8wplr-YVB#nRl~6Z?P@sH7wbM zi_vuckW%t6V=g-&9ma;s?vd=cR!_7-rqeUe6KTv`?!%}>&_{*tG*DtqyVhz^V)mJ= zY`sDzHS8zuBi0dTQ%_PAVnL$3mB_@Q6NeFH{-_a6Z7KXCr{Fz0_aaQ@d<;Qa@I_a6k_e-N!5zmWRBkfE0j#Lzd7 z2h+no783QB%bD*wYo`0oj>3XZ;@xe8sJA$>^*zhzU5$%mu!k%2=#FGt2GvStZYn2W ziZdlU6KLAybZWUUlD@O@pwQcfHh$Npud2V2sr}c8dfEXpC_j@pXL^wYkMSh%l`_db z`wS0JIE8D^Z^N}E5!mxM!JW0k@u}0_QAycVH2v8=RCsU=di!|}iVGf#?CRuEMcqSq z64pXH+0C%!`C@1tBZ5(7+HhuGFIca19%OIX1%8}L1`8Bt0CLk9ytEkvQdZp&qLv&H z{QbEk$&1J!Z7SRL0dx&x<oR@K=iVpGg*cg5fymtj{IX{`EU9WVvckII7jJMEl*KFay$Y|cLi2-@DZLeM-7L0|M% z^h>P+TZ|84$8M=nw_Bej!{m+VzVH2#1w}rRJrU9L<F1}QxgO9$fkURQe*jhXF952V}Cjz>c&IXyrwPP#$#Rq5r(cS|OE zb8tS=v#>-r#7d}f*HgG@;t6PDmk+b#!=ZCAf*C)D!o|H`fJ)G1uxdytIB2;V>?@cJ z9wm$cm80dryGaj(TA^Baza&>sm=-Lk_}B~Q`ZNX4+wVm~bEYws>#JCOfRFf-t}i=r zJ%m=o>r3`-NyOvr_`Ob`KKpI0O*1yEV1J99=+-VR8nDKkJsg?D<~YBS3==OWmy8{# z@~BmGQGEstmyV`wcjwblst#0h=5X32a9?uXO)~ZRVdC>)BUx7FPr8OrBGpCeM055V zoI0inx6~EmzAdr%fT0Us`NR;nD@frF3tQ1bw+e)0v(W3m-Y8Pd3JrHwL08nD!x+m` zud5c(giZ$I( zk-(B>+E6t^6FPRc0V_M6!eV58O2!+^B9crSTD&5fwm;#UUB)r=iKjQ!FA>qjr}bzY z-vRj;-A0~lK0+P_W)m~%Kq4z|N2ZR^B%ycT;f>Z3Z0NlMmpLTh2B&FwdWjMK9U+Ys zcHKg%BMu{jqz&kmw=X*Pc>-GUpDJow&zs_Ar{RRq?QpVX6kO2c2)E49hp8R|prD}{ zEX*qdH~y>zuMFmaM5A#aVwfWE`_d_VQavs_O4%ZqIW84`EdWBgg|2W#?vu#NBt^1A zwwCC0B{2uy{+!dWg37f@(Ni9pRQ*y6U6A3!Hjj0sI!S3vKGujnSUHham~al_Xc~hP z1~V;JOQPI7o@$jvQGd?Q-V&{%d96M)XC$CkOh(W&eQ6rF+wctX+sSA528G zbJS4hz)LvQ_$;jcTL?=f(eQbz6CC<`1l%O^2i$*t9o(8#4uV26fl|nPaOH~yXw6pw z1)fhi|6b4e_k7O3hY6p3k?{E1P(f<{7ZK)n6(#}UbjQ(UY>|^3OZzg35`Ol5_oxev z8|FvV93$DLlm2x2y9}lys8aV!gl3q?vKo~PcJ{C;Yn`q_R`!jhb&Zk#zq|5b3U5PA z_NB`XK%a?Rko83=l44z3( zx_^+o=UWj=-s-Sm*9Z90FB4kPxSSU8or3pM;^?7jKYGLj)4DH){8?ZSeRJhLF}hbn z)}Pr-qV<=O@Ed@vzpG2yLO?cEgu1kH*V)48|Vw?<4(<$I$uv zIp}$C5Ng!qQ35TaQH9@temg|`o4&V-#<@_y9y+a%rm`S@F(0Xx!MjN9KCQnRu! zI=zE$w2I>CfVMxKxsOn5Un3gRCQAppI>^G}I&#foD|wa}MjoF*$W>C$}q#5|H$EG1s$l_s|J5!LHhK!wjAB{zIl(G}t}D%*aFyyANpPrA}+ z?8_LI)xLw=QwpS#XD1}-*%{2`3_k;#YK0T|7U;?-{=KIvg9aEU&?~Wv=#=4(v{pQd zep8mCuYNuvzr5?onLYVrpnW;_M+xcqG@R^>`GMtvui^Rm`*9~(i=Thx1uAFTjex-B3p%6{-xM2}MUt;KCP!V9%Dj z;J}P3P^XX$N)i?Usi~8IXM#G2X@AZAheqx{6m$P!mC&}xSy(^}g>J=vqQNa0Y^GEU zEA&pM=G=7~UG2)MduPze`pao%lot7(5yKp6`1a2CrR>ypZ(>=ujCPh5Oa7##GZ%h$ zblPST+LESC#U@MWvv9seDwRn8+z6n#nk&{8Y~qPy9R%w_R4m`D#y)>6hau*nTT|6uAr?bA!md zPZy~>e}?S}m*D81d*G^_>9Ekk4W6Gr8a{`Ep?TvyFihzfaEZ8;OG zsoINO;vwS#Sp!X>1N{rwsL7vkjlxQrw{Mk1?tUu!zSfovjjn9^Zlp$MvtW8_7~d+o zlSqq=I0qv&h3-i=rDxL==+v=YBv zmzlVOc;R11EpeZ@GCt_mg>GFqft;%H(Ss}D=(Ww0Kcf|kS65&|kcC@=eemT_E3CFz1vlh8N2&Rz&?ut< zba!MVnxjh*`#l_)Sp9&9|6YZo#_fm2PuIXBb39<%(Xp`iy*#YZcm$+3)q+T;&0x^9 zCE(o!5lDEU4fMYC3O^T}7mDue63R_egwsJYgg9LjVS2=|WV-z6XR($fftg!7v1zqc;+dml z>C|KXbO=qS-@TJ*b6p7iecy$8IhawccS^L^qnk|kJWF73G5It#j;x$Nm8`QeBI}){ zN&UuK`1a_-_yRTCMe z|G}MRFB!=yCYwkW-%I1aXSCSLJ(f8YPG?(<-ir;S{t&>O`AB{S_>G?dy7w)nh)ko~ zw~nP6Mk@6C;+N!_Mibd^yM)XsOCYU|(}|UM6q&R^hM;Y?@zeh*ao)3yxcbOK?EY{f zp5dj2@5{eJx2?{ir~C}?Mq4!c^u!6xS2REl!~VdOFV~@=Upf53&j7`7^Wn2WmT-Wd z0opHn3f{<{0EJrlU`AUQ*ti@4aBnEcKJ`U_lP(LHZ}tk$_!;2gklDf{k1@g;IXR)O zs$6nl{Q|m+-+wynTEJYgr6hlwY?*1z97$oW2it3PT(T!YMBg~dvuR|tWY=!KHScl> z$jn;F+^c3Y=W)7{!h(LX3HneO!!(*zokG>ehSDSZrqk1}%<21ism8zX?)tiR^^@d{SWN#y0czU4wHWncW`B z>Sf+b-XDr!iXUwyTEKvg+WSzljvbT8e$JrQ>1p`aspafhlqX{b_L3)md&ug5`LrV@ zm6lCUr7{baQu5E0N{|ID%NjyMHoPSdvl$tFd>0w;P9;BeXA%2{reulAU^4pPJ=~*x z48Px)gVR3+;yZ43c(tJ>{<7{JGHVf|-JeU)#dYy$`yt+l)-gh3q@|JD+nZ1?_Ym}u z-2iXzTmbuhCqQ|3RVbFc0JpOmfVjH=Upwg)E?q|Je^#To=vtd8^e2{awN0*0q#7*n=qZZI8ZkP=gT|b9c#6) z^o0-T)=q|!m+eHXGZFRVOh9;*# z+SD^(u3I5Ev2P_<^?EW0a2f$@?*HQc(|_E5I>7y>b=-gQ6y}&)2=_lI3K0&Nbo-|= z|81v9qmn!uyI~oP8WYuRnxakcCG&^U$8Dgdlsa_Qc@3HdH-eWZi@<}P7;wI63OLtj0Ky&q z3X(@HLa@$3L0VzGuwki}P;crx#ZFFPyY+lhhM65Egq{C|5E4VCA+_HAjj z=0Z5#IcPTRx0*n24b`Nqu8(9lULu`;OUaYtYe@4w43rl=7;6}%{!1{6%NH#11zsAJ@wR{&~Y&R0{{vqJ~L+*c7aQ`cd`(Hi+?;i@h ze<&#HdeDqDGLmq2i?~nveK1C+0xYBDnNpt|GQ>_Qf>AJb@)RUhZJao~bzT!`0c-$4< zMBh*1o~|X*=jRa(EekSMSBbd&_XHOmuE%Njx8jd$!Z5Q%_-XY}JTmnQx+-@WXTbNYvu+aKFn8$l(rVex2p;_9qhzTg+}!IFGvUZFQCeL zmee?PpG46vp1r*`nf05SQNQOw;sFh1f@=h3?G{1r&N)RU*34qJ>>`-iSe3OH$o#G~a=$E^;y zHLC`Km^_&IJp}GN;Q$XRYQrJLAHi_vbD*quCvZBM1da->pvTM@C>@mnqjuir{#zyY z-!^jpZK2SzccS2GtR|3o{;XOdgD%}-L)atEg1p{KzVRk=yoN8EmN<-!y)%`re(fr` zqL@K$PIhHUPK#+#*KMNl)tRNfjAAu0rzQ8x_Ym~}Ckp?@Q`gVw^!0~Gx~SWOrm9Y& zK6`cO>0953)yrnGVtP5z5UnTo`S*xtR#xQHP8E{f`5ZfaI)#s=7vS9HNIdp1#qtSy zc>IqaNW1(R`mc6Bx;1Mpij$s;Dxo==9H@XUTz>?IjI4vp8aKlsN}l1)3y;MyU(u?uD>< z<&N~~xi^x$r$Myt!vpb;^C!rWW8;}<{Y{C2LmwS8;HFsnQWC$ZoVdT)q^;ijrfN9#}(AsY#tTn+0v`mbZPv9AEa*2b+YB;L87}j zi^QcbAf-(c$eT`8vaP2Zs}!Hcj<>gCuMbhUBWN<#d!&yi+5bXkV*W!FqYogrq;=?W z&ph<`{5W)Fry{Dl@)+he9ET;cTVctFW$=hHgtcdNVbObBu4R{u} zjxNIb8ZmgI{uFG%jSTzPKWNe97PMw(Ich$biM0NkkMML$G^1S!6`X$xmFAp)SKRX9 zs0rcFrU5|*tzocZ*jJFS_A>A{F9qICt3mOp*+6mo7_bP*flQwVf|Yl*P?MG`%#REf z-o)DrrHWcYGQX?s^0a2b|0T2RSIRitd=7iuzL0)j7b^2Q^ zlBphYrq-7==mAv=rrwmuN;|tH316c~r^JCyS+I&)@%`qz9V=utYj?)*fFTx+|3lL1 zt;k=k0=e8?k3y=wQC91C)T*X}Qk$N^jjvC_FGbs6)#M0xHkd&1^5JlV^LOxW@>S5E zzYn}#yav3?o&%Pvjs?jP^5CG;L+(G;aQ`un`;Uu->DnTp$xB;MRfwS`{Js9l(XmWp zTc1SgT{4@xD21+@xeG^WCz9sp8T3b{4V%BsjOxh7G1c1+^oYG7rQ1g_z<)Lte3mHj zJDomdTdF7(Lv#5ZdGDoYntk7kzR(rXGmd)n#;dh(Ov}BosqTLGMie;x@ z5LqYCfz5(*RH$&MP7od(&=LH-V@SPV!~TxmBHnMvzvG`uVvon9(SnsV#QS0osmMsC zHQh7VjPr;FcO$J(Q(2q_TfUrP!B`h-`PUrVk2NQs)mDwEHmspHcFm z+4BVI5UNj&;-vWX$ZZlFT1D(NbI4QUAhPk39WmT;|cx}W8 z&H?=bchCI?^g<2*l_TrG+DV?kO3DK8{w?7BTkd}z=lh-5{aZnKt2J|4 zsfiUWRp@EH0e$ zl+Nz&l4qg7giMUEq}TMKs26{JI=X%p_of!mXL*3eWsjgoJkaRH8(Pod&Q-gx%)VqiUS<8gDDR*^LzeZ3s)ZOH_;`NE{=aAC=NBxsOff+#43wZ0@&Vi`^E=>{|B``;yz-!1vs)pUBI z#Fa`#`_NNt1TX`No@6EIoKzo*|hk@hq=vr19Qz8bg-SkOzofPfubWoFZr*--9Up_?4*q z=SsWB_enNg;#{bj4%?%^+s126Y5FJr+{azez2_$IY@C#AHn}+=zJNW0h7(SNW0sZ&J!{DZ=&~J%{0njN0sC02gg+Qu;dQ$?43)!o_9$0^Y_hh`T8tLw+gGq z7||4$a9YaGP@DSVXwH#^w8x6j4ZZw)WtA+Q?ASp@Y19#8@fPy#;xck<3?j?l4<*Me zzGB#P8G{w2c=h_#Sbc&!Uh``V9#$ZSTV6dtNiEfA+~ZsX6he^W7kl)2g%;}F^8u!$ zFsM9xCwx_#2;VK84!hDv!3QQXutL%X!UtD^xwacYW2PTC(PIsG{~7T9GxvYbaR0ZE z`@eiI@6u#}_n!s7Pd@tF=JDyr|b0tbPY4kCFEkCGvL*{SDV8MC3 zlcJx%4r|tuBzHfWqIpqryFG(FdOM1#sIS1+=8T|98<$a?OMJ7wmiJ(;`P2IWl$H+} zMP;G~(^>rw$^GQxWWlSgq`WGetf)yU z9)jXsF#2LLjM16_L+=~IWML50 zPrn2HMjip(UYmg80)L8qNagcjI=!VzGHzN1 z8=*gzxrMnHZr;( zf@nBWaw1BP6rrEkc5pLR)hWYy${^$j`tZO$MXitR$l9}+s1XDP|T^4S^?O@xzDiD{H4SdE2fP@E=fbR!& zP*w1T`_E0>e=g?!bF4uBISc!kp^%vPPU3Mrm1#dtWNCLYXz63V`RVS#I#ZnKj=H7v zcvC-~(Xf)Gtl)fqd?-s>Ka1SjA4bm`Y?r*4m(DzltQgCB2h}>Y=n~FS=mcj_f4xMS z_Kj}?oaMWz(@khGrmCdMh-E9@Yw``<0#P_q5T}Fi4U1AA3gX4Nen-&4t8cdJLJ?Glf{$8j^(h|8Ur+ z8+fc*1x|jl9?RbJ#yfk*V`(cD{H*U8s?|P)HaBiV;Q~y9z%y z?t?8^Yv8u?b6~nK7D{iGhwoD!0(bRVV8I_grW{!eb~%f{u{3RNl7C%*6oXfwufrYD-y3=;Amb;D7}5+!4KDk?cJ8JnssM}vaP~n)Ut*ad zfl+7}VQn(*y)$z$WFR6F5fQ>7iV#9*5h5loY!fZY0*OeWlp+G51S>QhDq0x0XwfE= zCTTVnk=UwInxE0)d4IyQd^zVtMSZ^ zlR-E=QW91t7W4i0H~IeiK)(Op9ZXz49&G!yH#oM^ncp$%j9X?8#|?`=rpnKiP7Q8O zWR|mMruwV8Q=f~rW^#YaGVMjJaY03OrsP0A3;6X)yf`yG*)r6T8E>o0tj)ZhoUd6O zKYf?qPcH7sZ@+ZKo#p4_Z;k2r!N|UNwrOMhWPL6*`h7fAlf9FwFX>42ET>bSp6|~8 z!wOQHq7TWLvB$}R)z^}Hj-E>1dvPeaba6|v^hPexy7Ds7yS+croW7Ju-LFmb7Vb)< zyZ%HAkKRWO^}|u~-K){i&Zg-4!v3gkVng)t_gvWVWi+g>> ziRql*SDI-(N}~0S(IQk*E-|6*-21xEpYXh1&y!cLlvl5G?i&>@4+;u-aPpdnhA|E{ zliT^O0S?An#BBQu4%R;NSn(YWW{&%5`cE9{o717ai-YW|6Y@hjd>Yr`BP}_!BZf77 zk|4ayXDf6oD%DMnd%4h%oEP4zFAaOlV__&y z1y{Y-2h#+pxLmM7+HRGC?j;*osDCm-{iE2Gty}O$O_A?Elmv?bX9&(F;u4pJVQB*5 zzuQTVuE!%yYRW2a$HAp3mnGheMdxifZ_^lqo|zgLJs*uLLmck@7=`>C6}r_A76sG7l!tCL>^q5&9l-@SP>Zrz79M*uVfi zF(mF(UmNn8Rua6X0H-$aF{f^og#~MXF#3~dxcJ}m3 zx|!BXzfJyw?%a~VB=sqEAJ)ge@3zzKmI{8k(S36HC5wqN~o1q~$eULhOdFH-8OVCFOMGrjlP%xn$LP?~_3)w)$sb>%AbCzsJc zRWc?Nm(uQjE`ba#l#s&zDUbDo-$S9}kPo6ZL7 z)`*HU_wx5k(?zbcA3O7(&eI{uN+Hx&?$NxG=bHa z4T*$v9#IRgL@9^G2!EwPE?l-TJXEIB1Qw@Sx~K$6D#z#qnPaJ$W$|lEh*Rmf%a>c1lDp}*=+G#m#h*D0_mtCc zvMO5n)70efjCbGnGud~E*{a}+^oF0ys*M`xpRy8K-76!vD;dzJY9bjUJPNx(qX~xm zce-*iDT`&Fe`_J@{`E|EOFNY*?st^wJf)Vl4QQ_Ipd-&BVe0;pqn;__qPos*@~v@U z>k{9Q;ruk#IrM?P`V{aknL{KjHbcn3pLF~{F0`Y*P;K39IuR>Cv$`gWJ*J4kevY-@ zRYu4TW!5^XigC|B$lGWly3Hk6dCG%#1Il#edbg2X$MC0A-=z(K?scuogy zp=Pw8RTpgqacF)%4U^7zgX^1td-+eKi-%?*NHLK0kLaVMBZ3u=%z?YkTiJ}kzaF6H ziqyAr;gyw&;tpd7hw>?@*#vz~7Hs&kDT-`znc*>WcFW7 z4Wbx)z*hF<^-k5hwn`Iwe4`n&W zO}J-g%gIcQ!$ID`Y}=X$G(OzSM&l#lp?X8Mw=fD@ z2UlWYX*9l7OTp{YF$nfcp`*32n6%HDdDO+hJ}ifQz7UU}OnClKbprZN=s@pCA_6aP zDBO_*AHT13KY9zYvflCouF3H4lCanMDQLT>$ErT0V9%gN6)ljaV%3){i1X8sFDMrX GH2wpu=%Dcc diff --git a/tests/generate_onnx_genai_validation_packages.py b/tests/generate_onnx_genai_validation_packages.py index cdf3ca12e..f900d1c39 100644 --- a/tests/generate_onnx_genai_validation_packages.py +++ b/tests/generate_onnx_genai_validation_packages.py @@ -23,11 +23,11 @@ from mobius.tasks import TTSTask -def _zero_unresolved_initializers(package: ModelPackage) -> None: - """Give real tiny producer graphs deterministic synthetic test weights.""" +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: + 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())) @@ -38,7 +38,7 @@ def _tts_package() -> ModelPackage: Qwen3TTSForConditionalGeneration(_TINY_CONFIG), _TINY_CONFIG, ) - _zero_unresolved_initializers(package) + _materialize_deterministic_initializers(package) graph, builder = _graph("codec") codes = builder.input("codes", ir.DataType.INT64, ["batch", 4, "frames"]) waveform = builder.op.Cast( From b7f9936d1108779b2b09e0762c20cc4bcfdd3548 Mon Sep 17 00:00:00 2001 From: justinchuby Date: Thu, 13 Aug 2026 04:00:24 +0000 Subject: [PATCH 035/151] Fix component symbolic shape isolation Namespace artifact interface dimensions while removing non-contractual intermediate aliases before serialization. This prevents linked execution islands from equating unrelated changing TTS extents and preserves dimension denotations and in-memory models. Regenerate deterministic fixtures and validate against the latest ONNX GenAI contract. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: justinchuby --- .github/workflows/main.yml | 2 +- src/mobius/_model_package.py | 94 +++++++++++--- src/mobius/_model_package_test.py | 116 +++++++++++++++++- .../codec/decoder/model.onnx | Bin 251 -> 323 bytes .../codec/encoder/model.onnx | Bin 251 -> 323 bytes .../onnx_genai_workflows/decoder/model.onnx | Bin 2385 -> 2577 bytes .../decoder/policies/cache_length_update.onnx | Bin 386 -> 467 bytes .../policies/decoder_state_initializer.onnx | Bin 6244 -> 6641 bytes .../decoder/policies/decoder_step_update.onnx | Bin 2037 -> 2199 bytes .../decoder/policies/last_token_logits.onnx | Bin 631 -> 756 bytes .../decoder/policies/termination.onnx | Bin 3730 -> 3742 bytes .../decoder/policies/token_sampler.onnx | Bin 521 -> 584 bytes .../decoder/policies/token_state_update.onnx | Bin 1240 -> 1304 bytes .../diffusion/denoiser/model.onnx | Bin 1721 -> 1885 bytes .../policies/continue_predicate.onnx | Bin 910 -> 929 bytes .../diffusion/policies/euler_model_input.onnx | Bin 1926 -> 2142 bytes .../diffusion/policies/schedule_lookup.onnx | Bin 607 -> 669 bytes .../diffusion/policies/solver_step.onnx | Bin 3060 -> 3248 bytes .../diffusion/text_encoder/model.onnx | Bin 1287 -> 1331 bytes .../diffusion/vae_decoder/model.onnx | Bin 825 -> 957 bytes .../onnx_genai_workflows/masked/model.onnx | Bin 1008 -> 1104 bytes .../masked/policies/masked_update.onnx | Bin 13182 -> 13264 bytes .../speculative/policies/adaptive_k.onnx | Bin 34702 -> 33930 bytes .../policies/cache_length_update.onnx | Bin 386 -> 467 bytes .../policies/grammar_guidance.onnx | Bin 1926 -> 2054 bytes .../speculative/policies/grammar_length.onnx | Bin 394 -> 460 bytes .../policies/grammar_sampler_logits.onnx | Bin 631 -> 784 bytes .../policies/proposal_metrics.onnx | Bin 1026 -> 1146 bytes .../policies/speculative_acceptance.onnx | Bin 8736 -> 8835 bytes .../speculative/proposer/model.onnx | Bin 1019 -> 1095 bytes .../speculative/verifier/model.onnx | Bin 4264 -> 4371 bytes .../tts/code_predictor/model.onnx | Bin 110332 -> 108810 bytes .../tts/code_predictor_prefill/model.onnx | Bin 366 -> 465 bytes .../code_predictor_step_embedder/model.onnx | Bin 817 -> 888 bytes .../onnx_genai_workflows/tts/codec/model.onnx | Bin 991 -> 1040 bytes .../tts/embedding/model.onnx | Bin 5618 -> 5622 bytes .../tts/policies/cache_length_update.onnx | Bin 386 -> 467 bytes .../tts/policies/code_frame_update.onnx | Bin 1580 -> 1648 bytes .../tts/policies/code_history_append.onnx | Bin 853 -> 981 bytes .../tts/policies/codec_layout.onnx | Bin 428 -> 508 bytes .../tts/policies/continue_predicate.onnx | Bin 910 -> 929 bytes .../tts/policies/last_token_logits.onnx | Bin 631 -> 756 bytes .../tts/policies/predictor_body_sampler.onnx | Bin 536 -> 626 bytes .../policies/predictor_prefill_sampler.onnx | Bin 539 -> 638 bytes .../policies/predictor_state_initializer.onnx | Bin 14690 -> 15717 bytes .../tts/policies/predictor_step_update.onnx | Bin 2039 -> 2213 bytes .../tts/policies/setup_predictor_sampler.onnx | Bin 537 -> 630 bytes .../tts/policies/setup_talker_sampler.onnx | Bin 534 -> 618 bytes .../tts/policies/talker_sampler.onnx | Bin 528 -> 594 bytes .../policies/talker_state_initializer.onnx | Bin 7457 -> 7875 bytes .../tts/policies/talker_step_update.onnx | Bin 2044 -> 2200 bytes .../tts/policies/token_to_slot.onnx | Bin 438 -> 480 bytes .../tts/policies/tts_state_initializer.onnx | Bin 2194 -> 2339 bytes .../tts/talker/model.onnx | Bin 32472 -> 32118 bytes .../tts/talker_prefill_embedder/model.onnx | Bin 22707 -> 22703 bytes .../tts/talker_step_embedder/model.onnx | Bin 7521 -> 7537 bytes .../tts/talker_text_step/model.onnx | Bin 1759 -> 1860 bytes .../vlm/decoder/model.onnx | Bin 3198 -> 3469 bytes .../vlm/embedding/model.onnx | Bin 1649 -> 1715 bytes .../vlm/policies/cache_length_update.onnx | Bin 386 -> 467 bytes .../policies/decoder_state_initializer.onnx | Bin 7283 -> 7747 bytes .../vlm/policies/decoder_step_update.onnx | Bin 2037 -> 2199 bytes .../vlm/policies/last_token_logits.onnx | Bin 631 -> 756 bytes .../vlm/policies/termination.onnx | Bin 3730 -> 3742 bytes .../vlm/policies/token_sampler.onnx | Bin 521 -> 584 bytes .../vlm/policies/token_state_update.onnx | Bin 1240 -> 1304 bytes 66 files changed, 192 insertions(+), 20 deletions(-) diff --git a/.github/workflows/main.yml b/.github/workflows/main.yml index af9524bcd..0fd3ec7f7 100644 --- a/.github/workflows/main.yml +++ b/.github/workflows/main.yml @@ -23,7 +23,7 @@ jobs: - uses: actions/checkout@v7 with: repository: justinchuby/onnx-genai - ref: 5f151fd0465c25593e17e824d9680a04adb07eea + ref: 923530877b440145f5fc1b848d50cac73e02e2f6 path: validation/onnx-genai - uses: actions/setup-python@v7 with: diff --git a/src/mobius/_model_package.py b/src/mobius/_model_package.py index 3d94c8eef..8fa6df5e1 100644 --- a/src/mobius/_model_package.py +++ b/src/mobius/_model_package.py @@ -25,7 +25,8 @@ import os 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 @@ -167,22 +168,23 @@ 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, - ) - 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(model, path, **save_kwargs) + 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) @@ -209,7 +211,11 @@ def save_policy_components( if check_weights: _check_weights(name, component.model) relative_path = f"policies/{name}.onnx" - ir.save(component.model, os.path.join(directory, relative_path)) + with _namespaced_symbolic_dimensions( + component.model, + f"policy.{name}", + ) as saved_model: + ir.save(saved_model, os.path.join(directory, relative_path)) artifacts[name] = relative_path return artifacts @@ -336,6 +342,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 169b856c1..d9e076f03 100644 --- a/src/mobius/_model_package_test.py +++ b/src/mobius/_model_package_test.py @@ -15,7 +15,11 @@ 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 @@ -375,6 +379,116 @@ def test_policy_components_roundtrip(self, tmp_path): 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_namespaces_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]) == "policy.sample.batch" + assert str(saved.graph.inputs[0].shape[1]) == "policy.sample.vocabulary" + assert str(sampler.model.graph.inputs[0].shape[0]) == "batch" + class TestModelPackageApplyWeights: def test_single_component(self): diff --git a/tests/fixtures/onnx_genai_workflows/codec/decoder/model.onnx b/tests/fixtures/onnx_genai_workflows/codec/decoder/model.onnx index f6224514cd03055262813f841be4615a590875d7..b6d63790dccb59f94949fe266f3ceb0f5ef732f1 100644 GIT binary patch delta 189 zcmey(c$kTmgWGBc(?r$@^?_Wh$@wX%#X=rjZXAq4PF#{g;>r2B1^IcYc_n%&sUXQB zy`;pF4vN*s(ra$M{}tVxL_$r)Ts9E@E2LcEEk kDVh24#fiBEIjO}-23#EFiDjv2`9-2PAlO z$@wX%MS4kzCCM3FOdO0{ibC?3(ut)hnfdX>iMa(isl`cwT&y5n#RQGiC)!LW7A^(> F2>^5^I!gcm delta 116 zcmX@i^qY~DgWKu@<3!d8X$D*z<%wmfY57IDLMmKJ9E?J8TiMa(isl`dUT&&6YDXGQ8DR*MwVi1r30J44@P5=M^ diff --git a/tests/fixtures/onnx_genai_workflows/decoder/model.onnx b/tests/fixtures/onnx_genai_workflows/decoder/model.onnx index 13e5e0d904296e7a1ca2710d4d67f2b0f98bd520..217185c90abcf532637fb044057f560bb325d85f 100644 GIT binary patch delta 507 zcmca8G*N_=gWIZ2WFxBxd%Xu2XJ%eOX-Rx$O0keNmn8?gkSUkAkZ5v#Zb5!tYF>$6 zZhlH?j$Tq?Npc34l#m3vL~&|iX=+|_YE&>6Ut&oKP)BBdUVLt1akh{nmp#xpYuv^u z2-VA>n^%%wl9&?@H`SYqry##L6Kp8OPE6zyEl4abiO){0j4w;fDNQY|*E7%q@`OCN z+&CD8oNzmpiGzttUPu<*wP535j^$$F;7Ib~V#~=-&nzhxa^SK9nrVgGOkilHqZgmM8PzLNZs#GH7zCT%XBg8br4 zupVqS)T3D8$R%2kSX>gHomv@RmY7qTTC8WF2jmE8aH(-H3MnCMV&Y)p;uqottA<&| z#l*ppq{YRSlb@bhQY@svB?r_ah0ufK87^jy2F4^CE?%{QqSRubpP=?=b7=zgs3P=$ O>=6(lEi7%#0QH2 DWYjl; delta 97 zcmcc2+{Dbw!7ap=o1c_fT5R>6X(DfryEqq1PHI|-5I+|$2fGj_7rPK^QesJRMwA2> ZYf)x;2Bu6BP^KimBrykB&tyZ!MF3MT6yyK^ diff --git a/tests/fixtures/onnx_genai_workflows/decoder/policies/decoder_state_initializer.onnx b/tests/fixtures/onnx_genai_workflows/decoder/policies/decoder_state_initializer.onnx index f509433bc9dcc4a6173749fa0a6b615d36e1e27d..3a042bd0b7d887ff7f856e266bb7a5bcf4b5fce7 100644 GIT binary patch delta 795 zcmb8tzfQw25C(8HNMH_%#8nj=>V%3ah4haY7#Wb585mHp?8GQ;lQ`6uN@dBwtMCGd z_hI4%cmv!BieO1QTC(ov@BDoPpTj2G!|l{aUD%8F;cN4yRTj|9XC}>wlbNc_Vg!=y z;5~HlBsYl;7k;Ed6R8=uBuH_cY0@ImPip2%L17F73~FzF8&*9QDhrh?;6MaES6)qbL#B|Y7b9q)sj%C&;RIsVa0OhydLHv2K)-%5OYrAYz+-H31iN#76Tcm|ueQDb D+!+^c delta 318 zcmexp{KSBlgIkC#H$N$}wAkvR!A9OC!u7^nyah%1xdkQhCHdK@dBs9XTnZfQLNZ+J zLaa%NCCM3FfgJzPQ@aMiYx`9;Mv)wooE#ws9`hzbd# zIYvQS!7xc<@&r*4OJOcSplV)(Y9O1+fSqWI#H)PnfZf|SIPRK29clH?376(J=&YLfHwN>VFIqU^YM z3i69H!B%9Z6btEb=>Toe0NGHFM?1s{CJx4=Y%bwEph59yE)C%d0-EDTluI>))Ua9Q g2lWorsS4T(hDpv`f?(Tlxfz=|P-8YvWJ_TL01X6sCIA2c delta 199 zcmbO(_?4fRgIkC#H$N$}wAktq|3=_Q@3 z>_V(bi6zMyT%1De$@zIDsTCzr%3M4J`Nf%F#hEF^Lc&~vK()LGwM-n0Nv2%Fc|gVS nXy(ds$pF<#Ak^{-aVKj-3|7!qFig_m5(MkRWy5A;_7p|{s@F9= diff --git a/tests/fixtures/onnx_genai_workflows/decoder/policies/last_token_logits.onnx b/tests/fixtures/onnx_genai_workflows/decoder/policies/last_token_logits.onnx index f674183330680e7f8e3468568778f80f8327ea0b..4142d9050c9fdc97b12ded4b6b105e73487316e5 100644 GIT binary patch delta 241 zcmey)@`aU`gIkC#H$N$}wAkt)>qg#S#`-!gww(O*%#vcE60RZ+Mxi_|1tGbD{G80> zO1+%K;*$81{Or`cc!(Okq{Ncs3@&9MMVu;%QwvK|^O94!RE3mrDk{rQPE0DzNi3=i dON!&-2HOL5g+G@s&=sCUxB{OAlXo!g2LSQqRMh|g delta 115 zcmeyu`kjTBgIkC#H$N$}wAkti%SPT{#t36Bww(O*%#vauO)hl~Mj>S`b|KcJ#FFF; zE-oRC;?%;@)V$GSvi#)4q|%(kqRJ#4F7BMf;*xl%PFXH#piXgwPN3?^rcC<* DBibI7 diff --git a/tests/fixtures/onnx_genai_workflows/decoder/policies/termination.onnx b/tests/fixtures/onnx_genai_workflows/decoder/policies/termination.onnx index 3693811ae428bef7c7f17652199285b2ca557a9d..72763e63ade7f01e72cff99022dea83883b2f75b 100644 GIT binary patch delta 557 zcma)&yG{Z@6o!EicJ}OIG2#WndLdR8;0C_%yL06Xb@ z06XJL`54ZQ28bSOGoF@@clL|1ITBCRTGE5dJ zIE116v#ZYjn73SYZ_^W^d!6oCrMr?%$9!rl)a;PE{QQg3q#7fE8hGiVVR?pJ Yj6MUp?i}r}WWjYlw>{_H-tCwC2a<=600000 delta 521 zcmZvZ%}T>S6osi&letMOo%+|>Rt1SWF;+#~iHLiDHbS{fn}H_U32i1p-1H4xn7t3+ zPTc!c?sXDtiNx)kbMCqG-IvyL%hCd;lnm}Z5#|3(BNGo9J8YQ67 z`5u$~!EF;X9)y`zg^JP)AXF345Y@9%i6{A(NiPt^>=rBz*#piT#qqsWsFvM8@GMMz ziWAqko@|2^iH2d#_4qb+t z5-X{g`-8s&B*ybm#P0EcHegLBbP2k@0P;>7bk)W}iHDg{Q^x;%n`!eUccVVPQlO3s zWRk8~-)Mso{jlAxAD5A*q7=oXq4( zy^{Rw)V%oO#N2|M)FQp4#FFF;E=3`Etg>bK$%#p&If+G;NxEFDV4XrLTuL15LUQ;_ W{N%)>(wxMi$|MOc){^|})I1@6E?y3HAx@CG&5n%hi~s>46#D=G diff --git a/tests/fixtures/onnx_genai_workflows/decoder/policies/token_state_update.onnx b/tests/fixtures/onnx_genai_workflows/decoder/policies/token_state_update.onnx index 4f3026a099f6c88007fd54d27c96dd2b224c70b7..9a8aafaf3172c300217724e5f79d073252f30bbb 100644 GIT binary patch delta 254 zcmcb?IfILrgIkC#H$N$}wAiYZYa{PZrg{r5_TN* W8sxqv#x6!K#mVbfBqx7kegpsp|3_i~ delta 189 zcmbQib%T?agIkC#H$N$}wAkt>=SJS2On!1)?8&7?MX7luLc&~v9PC28TFwk~9}fUTOu3zR8Q2m85}2mc{#)=ENIf7&-X{ qv!b>N7jId7XkKw)X=-X!YP>OqDv-rZj9rXeij&{5NKQ6mc?1CV<|@tr diff --git a/tests/fixtures/onnx_genai_workflows/diffusion/denoiser/model.onnx b/tests/fixtures/onnx_genai_workflows/diffusion/denoiser/model.onnx index 42970ba7bd4b16331995083b31d41589651eacd2..b2f0012c61e80b875e79bade7be93ad228a581a4 100644 GIT binary patch delta 397 zcmdnVdzX)ugWKv7-$vFyEcKaOY{iMW1v#lg@m#SSj6#uIQbH2R`MCx8d8v6NdMT-S z`I*J3MS4kzCCM3FOdKp+(n6A0_V(bi6zMyTudA+TpU7d z8L64+86_a$^30Tyj3{X?j*`sW)Z&uV0wI1bUZ5UMxSlBc$#YpHA~1{+6yh%^%Fit* zi7!qqEKSWzP6e8-kYvKem6xAcoEl$Hl$t_>119IN^(kp`@s-5~rKXf7r~0NQ=Ea)` RiF1hpT_Ff}#bi--4*&@cHrfCH diff --git a/tests/fixtures/onnx_genai_workflows/diffusion/policies/continue_predicate.onnx b/tests/fixtures/onnx_genai_workflows/diffusion/policies/continue_predicate.onnx index 3177d8905c97dc120048da84e39b00bcae00c3f8..3f2c3d0e62b67b86c520aa2590b90d5985642f53 100644 GIT binary patch delta 123 zcmeBUU&zkO!7ap=o1c_fT5Q$FzL7VHDZ!A7B_%&ERY;Rdor6IB{Mm(Bvmgdu_QSoNsNmFCMU$j#R1gD!o|eFm?gr+Q5Nrfs TAjHAN#=$Pc#Kkaq57Q+84z3^m delta 79 zcmZ3;-p9_%!7ap=o1c_fT5R>3Z6j|Ildw1!OG*P61ax&6f b9A)v&iNz)H24FdMh+I-)Npi;I^Gufj$pI1- diff --git a/tests/fixtures/onnx_genai_workflows/diffusion/policies/euler_model_input.onnx b/tests/fixtures/onnx_genai_workflows/diffusion/policies/euler_model_input.onnx index cd1ecb2db40cadbc9aa5d62530db9fa3a9eefe0d..4dfc8aa475265467b9bacd3f9c003cd9079f4924 100644 GIT binary patch delta 617 zcmZqUzbC-U!7ap=o1c_fT5NS%U?Z;~YyEsiF1F&t+=84`p(%`96B#)eg?bpd6oljo z@^dniEA>)Kb5e`qbMsSDbK*1e3Q9}#k`hajGq{w66mhCZ&PdG5OU)_fQWTQMsVE~g zGd-iE9iy^z6SsR delta 420 zcmca7(8kZp!7ap=o1c_fT5R>1ZzHcEYrO>*TXAA;K~AcW9+wUWqmTv{yAW$qVo7oa z7ncx6az&)0EwTPmtK+)rN_lmT#{NKq{gMf!7il0B_||26SG%>1kDNg>#s5torV>tjk9YVqY delta 154 zcmbQsdY^@tgIkC#H$N$}wAktl%SK)+#(GsQj^gBu)RfYkR3S+&aSlcyVJ<--ewbK% zPHJ9yNk)`77fW$TYJm_x7cU395GNPA5NlFmNpeP#G#5unW^O86k`YNVOO}hPEZ#k_ WBqOyb-T%{6(onZdX*1Zz%^b~UCg;=_(8To|nNfGh@^h0=e|82&MMMbOF}96!u+C&NFZaASQRFWa?d9 zq+CnQS$>(}X~=zKkzNMY0LA3i6p~z}^l_figo?reS9@?tj1vri)Gb9QpmnbU)tEeR q@f1-yrNl6B4X~QK3TsN&7Qfe+W%nbW7^h0xQVygX>G%BuJ?#(Ji{O+1 delta 595 zcmb7?ze)o^5XL#tWHTWKSLBF6MMSXBNJLExHez9wD%b>J*vk#sJ1%R??yX{TU~iQd zQG5g6#CPxkyhIWLSEZO{e$0H|{Q0Q-x+PZRxxu3s#Jct4ewOxF(uABAgJDFx2FJ|d zN08@ho-jQLz*$cUr6|(CcurE#mu7;W1K-H*HGof}zzP%C<*Qv-)rN)%FGGptd0{Tw zJp`tElKeRJQ~@{jv}Wqi?B86*N!8spDlanH|N`d{r|72^-0kvz4+IA@xb{5 D@3@cv diff --git a/tests/fixtures/onnx_genai_workflows/diffusion/text_encoder/model.onnx b/tests/fixtures/onnx_genai_workflows/diffusion/text_encoder/model.onnx index 7cad175b663f65067d567706847007491169e9ee..b6a0068b2cd1e442121e496ffaf2a8db6967af62 100644 GIT binary patch delta 253 zcmZqY+RVku!ELpIYa{Dc=7d--&dj`m(vtYhlwzR(EVINFberm zVD4lImRcuyF5a^E(7a-x*{M~j@y0@YTs%Pca&R#KArl8R$zpORV>pOKl8lA0G^T#{IlS}df-rNY4|q=2G%atTXq zkvSJ{S$t?-G0=e2s?>O6uu66`m1x#*F>x?v8F6ux#XBb!m&BW3H64o{Cl)RS0SN#> C={&*! diff --git a/tests/fixtures/onnx_genai_workflows/diffusion/vae_decoder/model.onnx b/tests/fixtures/onnx_genai_workflows/diffusion/vae_decoder/model.onnx index 97b376218d632bd1aba0daf917238f70b75a2dea..773c8b4b2b4432b5f206ae035966db2fd0edc2f7 100644 GIT binary patch delta 245 zcmdnVwwIligWGBo`$pCZrurf-ww%P0)Vva*46ZZ|Mxi7wSt04<{M>^4JfNUnSz>B@ zN@{X`N@|f_QesJR1{V_t3zwXb40a_MshQ~+CAf_!&rB)FNGjxF&CE?qPX$_;M5LX} LUO`gmg3jh$aOa}k} delta 146 zcmcb>@qwL{gWKvU`$kql=6Ve-&dj`m(vtYhlwu)iE=dk{AyF=NA=aeClH?37E+LNM z)WXu#yyVm*EiSg4{PfI{Vj%@CISxi4DXhx4m^m64lk~Xw3ySg!@{3bb;!ES?j1Imgr>ZrjipE$SMp9>KiuU|`pw?_FwE zY2_VEv~8BdDd5f_{6fJ_idAh53POJgchdZd0>QD|?U% zb>ISJ{#~FY7lr$}FZeiNk2yx#3W72c^cVf!qy$oQ|6NZp57`FhYcZ^tP`HxsLJ~mPxr9nQI(6 z`}r)DoV*Jo9(iRLWL-!@3N0OiL?>IAElfh%BRg)Bh^tYj!K=Y>QH-h`Qpl@$8289q z)4`WSKTsmd!WaPVOF;sL(#giXU>~1U1Z-eu7;ic+J%6c6Y7d4kQ7NPoqW-w7PR{^%{XpIe?yqf%skVu%aPFB4?v2^5n6esk21%o6=-fUv+$t@* zEg}zc0lz#(L>Qp5s}to?tai{dB^Imr zl0Dp_NCURWQ4FJSDoTGf$cmyy)y|KVyx$jDWi2L5;4%;38|pYjS?3mKf&e1(&|=dd zN@!zGeb-_{6?@J{Nmgav;1cd$)=-TD@6}LMyX>dBUaD&`O21cITiy3riVM@bsibr_ Pbsu>trp?(`@jUqlT8!Wf delta 1863 zcma)6%TC)s6vZtK_Jxo#h*BO2DHH($0)8dXT}7p;Qt6_os>Ehw#}gnVHuX#NA^9Wx zB^5u=UAycrbj?3h9gkBdjvd7Mo_prrbI+OIXMbBoei9YBmVV(mFK%1ki+h{FE=)H) z+a{)qUF(vV4qAc*ege(G1bU=vZu1-f@ecXq5wl57Yf!qftt-nR#=Y<|6gt+~h3lX# z*x)&|Hbfi5yfB5px0;&nbu|0)={gj;nsZ4EHjFT8n0Xj?TyljNlq0R*;Sgn>r+)R_c1z?hEmvU+@xl<*569fj0mbh~TU8bWl=%`oe zsD@<{HjWVQc=L6;zy*si;p3Nvfe$jHwx7V1kB=;qj!N%(hzc-H1sZEF0q^5uuY<+? zY!fR+grb;NEDkV=5(;m3rjFKO+Q%PEhaQo9Ay|~rEX>fCipeS90nfF|(X)(MK{HUI zfXM_9?(8lH3~+4Gr|TI2jI}Eu(8o2$#bSfixd8%NpNG_N2W(%C?#p;zUHDi$%l2D8JYsE5_kT$$mx5x8)xw)$q9 zjn>muQkjXk5jp7VE5-5+$`gL8sg)G~AAdY2Ho;O>iSSpgy1J3xMHcB}vox@!WS8*a z-t%A>CTC#Ln)D`X(E_h+LbGCHSz*Qy28}UBLDl52vYNF{gjkEldQS|mNlv^j4&+>e z(7g2wA~)_#X8T#1PtFNX+D&pQJpuKLr(c1X549uDoE!;bSe$z~>G-F&`dwXA{Dc5QaN z)3BW8vx9Atv6@YL;MlFpT<7j>*XoX~g!`WE-yjRk{>e$(;RejaxbT^_A^N<7^@)Z$hdyw2mBh0>M&=_vLj?%CUyr`6cg{coIfV`-CT9 zl)p{3GO7G#kYg_{$z(cnwQ-ls_UyCERi$B)EAnzun4~Bjf4ld3WsPKAw$JNrvvPip z5dOYdGUeave~BRzL4WlRR<98N)>4(A?`61@rYg`G32p&ff*vMo7JS*VSKnKIZU7o+?3e{ zx#F_Vd!y5_ZJ+ZhLsm_Y1!o$&?wa3a;d7j1R z4|_+fY@+iel6KkSAuHF=p>)01cUZZOPOp+Vmwhs{xon55a=VxaQHA)gM%TQkVOgOa z;LvNaN+tFbm;nQI7xPl5m{lZy1)Ei%Fj=7xS+gzw<@Qp(qL)M*7Z!D;iPPl;O@3QG zEiQZ7;j-u)sYwH9LQx{mv6%s)lwBoXgUoO((J@d4Z<6vyZI>D51DUn>AV66K(I0`h z`_}1_|CyVEKK=M~%*SOFG$&Ig@n75=uO!#p{WH)D#ZpArlhY{8 z;53IQ`oBNlFrWt38LR0uBqEesLmXr(p&x5{Ipy(O4RP|HoNs^(AMprvyOf+6kJJ$% zC9_ZM(TLF^6iGczewB9ljHE@>lF|%8KsUx}hZc3R(oBqZS_HPr?3Vw**Q*CY5$X&@ zfr7?Dh_?PVG<^NVZwl`g>tRZ6l0}yZ;*WnNJsmfQH?P(bTGqv&)LV{KgP7lxuTaT+`5qMm26E-gR5olEkq!UVx4D#*(1(0` z3L#A>6{K<(A{tL{p&BDJe~ckP2WH>^?tACbprAr`zo4EPiJS)boR`?8`SULc^#$Qdfht= zyoPRMpLzqg-FJ=FHS}7&dvMSUn3V+G2VTe5U0j<$i{X3UFH&KeQ9?I*?!AB+LGyqm=4W$gE}hEG4r@Xa>81S(XsY8K$0c8lYu>+t&Iy~*Mfhju z#UAxH}jP^2g!b1f--xt{>g4%U)_4gq1qk zk|2e=36g4wdbV6YAO%QeTH-5w_AnRzd$djg)ovb=LQ(jUQ7lv{5lT;3K%`LR@lIs! zC`eNEQXTKbNQ_*y1H5SUFeMDgXC42j?|F|FShFqc%=qt6HYo#gqtPHvNoZM?YwHqk z_Dz4=-6l?1llrOoX6GJpD!Tb1njVnbJ>pb#n+Utz4R(oB)2%D$?0|gO^VlHtNHISS zfHgq)kAIOSy*L`E1+{k?q*xqf1Y+bNm%cR661jbs`8BYTXpTpUAcPr7xk4%%;nTNs z7m5}QaO?b%Wh^$rXTjp4VSl7r8zIgwQ1djg-83L#no7V`LJCUSB^V`TL8V%f6N8?> zO8LZ=fVe~qLEQVz0Vx4*38R_-@TMi8!VCn4&mU&5mGjvtG|4vXD=_{a9QEBcDZ`}5 z-QatKmJ37l3LluA{yr(|J5+Q|0B-e}GYdxoz(=lZSvF$z?cMi!9x01$L)Dic+F21s zUN_RVi8Hlu6=!N-i7&>}Oev{Yv?irm(VqjshYTs0r8ms9V&;@WCY;pDI8rO)Nv*&N z*R&&Nn&1>j_pxH3Px4o_@V6Pd#|ukF=FYUhXygox=S@={eV5gV3V%Ib z-h!sP$E0duK>+v^RudeE!%A9BRg0*^;>uNZN*I0fG@iO)9yMe2k@$EXH6sJ5)|&iE z)Id4%KzAPU)!fQVz`278xqUf)M50?0Y_vCYq46=I=Agkf(m3j8zLyHoDN}t|9>S$?ik=IxvpuIyjFiIly^j&~KHj5WL7$076`cS2X-S0wDoqRz<)T?jv> z$V-w})F)9*XY8h8H@liJ5B57nFwME>CRa)Mk7m}zgAX~E^r}QEd?TEvY6K-qc7TsV zj1qx(I!YEU_|lQ_^OP!V@Kf*gb9xmB0X} zBZIL^eY_uP*4FIAt`IbY!79QTzWOx~vfmJd5LH ntJ|r&?5*dQ{GQ&_=noh1`wyE>`AN9AhK}@i^|!CS%T2rl*tG?T diff --git a/tests/fixtures/onnx_genai_workflows/speculative/policies/cache_length_update.onnx b/tests/fixtures/onnx_genai_workflows/speculative/policies/cache_length_update.onnx index 72ad12c9f4e344e870cd12220f36d4f2593cc2b3..45764e632822b23fc759eb608a1f8079786bf246 100644 GIT binary patch delta 178 zcmZo-zRb+a!7ap=o1c_fT5NTIc_MF5y%85nPHI|-kQSE)2fL68my(b|L4HnVa;08! zVsb`md`@a!dPzonX+cV2Nvd8_Vo7pFlra};QD%AuNt%*?no9CZ5_3>(lEi7%#0QH2 DWYjl; delta 97 zcmcc2+{Dbw!7ap=o1c_fT5R>6X(DfryEqq1PHI|-5I+|$2fGj_7rPK^QesJRMwA2> ZYf)x;2Bu6BP^KimBrykB&tyZ!MF3MT6yyK^ diff --git a/tests/fixtures/onnx_genai_workflows/speculative/policies/grammar_guidance.onnx b/tests/fixtures/onnx_genai_workflows/speculative/policies/grammar_guidance.onnx index cb4ad9693ecdab7f45483069d859c941cafc93d6..e79efb9b0951cf011dd6e1582fe7c860ff045b82 100644 GIT binary patch delta 510 zcmbV|!AiqG5QfvzrkiO)WeygEB8o*s=s~N9co7dhNN-A^D8nXQ)3iypZZ^Ss>8&r& zXYc`hBo97;FW^$~BvSD@|9sy+%%Z;7a9I`YSbAZWo`2L|-Bx+t2Ro8Ms8j60H7j5j zwy`;tQRvUwf#Pw@6$M#1*M{uU~B#PMOGM`ebe_8i6qb?+=rIIch=%w2g<9)zZ5Rt%Qy1| zI63u4LJ2xT0A`+C0t>A17C=syD(LYiI>wrbDrlR;95~*DT25CoNi|P2bfKEw=j1w~<$swO*ZzEhj%cv!qx^mP?w0QAnJNU5GU)u_QTz zi${p7EI&ChsWd0Cs4_~2iyNXaJ~y#A8?2fWw`vtG-n9IpAr~-@y5s|fjv4|kX_bMj*Gi2-n}R_u_U!9-W2Q(P9%3^0X5X(I0=PjN1moYb@uA$~4i4t60qg#S#`=CnF1DQf^vsfCp%zB2Mn(=sp;|6gA?1Sn zoXq4(z4W5Q+}y;X_~OLef}GT%c!*NHq{Ncs3@!~JHT=qoQwvK|^O94!w1hOUtE`8a vSeBohm{gjRSX7ym!Nr}ESX=^?5sKsr2RbU4SV!Ts3e%=6F|Nr+88-m{TAO4u delta 115 zcmbQh_ML^7gIkC#H$N$}wAkti%SPT{#t36Bww(O*%#vauO)hl~Mj>S`b|KcJ#FFF; zE-oRC;?%;@)V$GSvi#)4q|%(kqRJ#4F7BMf;*xl%PFXH#piXgwPN3?^flT`W D_oW_A diff --git a/tests/fixtures/onnx_genai_workflows/speculative/policies/proposal_metrics.onnx b/tests/fixtures/onnx_genai_workflows/speculative/policies/proposal_metrics.onnx index 0084a107f34b0d3c63c7a853108d6ea7a925af8d..a8ad0be82b3c1f4cc5281fb4a4f14ec3a89ed4ca 100644 GIT binary patch delta 296 zcmZqT_{G7?!7ap=o1c_fT5R>4Vk z)VWlFW-G$YuE%9~k_i`QYFT1VX(CV~hBlnK9J$2OGIMf(cEMZ_pH!NXo>~Gnf)mMR Qh!MQ`#Yo0O~Vhi2wiq delta 175 zcmeyx(Zs>a!7ap=o1c_fT5R>6eIsulQ@tJ+e?d`xL4I*+N_`d>1Rh-f~AymE|$v~;LD z2C^cS71_+CYAtu^ip(c4>C$6?`#nYDG~A{3nd6wmI{X{x4bUy-inbf{!EgTgapUD; zIhq3qal_x+Cdc?80?_4$k%knc4rU=2;;o~AIM`4Jx!yqmvLW7K0#}E|GoXiflW>7+ zLpn2%mviQF`DY1SM;S;!5>IE?8OZuM8lq{lDHC!7TNv7U<%uCPa2DDcPVdC$l;=GdF-KRdq#t>4xv`~D0y<3%P z&_Ft5^cMC&X=U2heTyZf~NbuX{vNP5V%6IZ zNu47mQRv)Zf;6D6q{vgqkb4@H`Z3%mtA;}dA%LB^%+?v?33yzmnrl(#W4{B%k;6w^ zFasq&Wws!Xp{is^X#q!n9aaQmhEIx62(j4&+uwjR9}EPOGe=q|MRlEx!5hJONGSt zGHxD1K7?~yNNP*$kTk4>6TM?bv3A|zdv}Q41w{V!VwuxS+X>S|-WT%{+#X?lbtawTV%>s~OC#b*`zNZJ-*r`VWn6<#i*<9sAl|B6sk)QffxN4v(3` zP+cZ{@T*cf4A=B(LDZQ?V3x<*7P)A(rXJ6D!_DYw&oHpno;eJ*s+W6DVAaD9wg^8D zcEPaSArD|JT#!QY4qgWGBk$3|9e=6VAzwvznp)VyLLH7*qnb|D2WDItmE{M>^4ywtoBy@I0r zg8bstBE6)ai2 delta 157 zcmX@k@td8MgWKu@`$kr7=6YE!wvznp)VyLLVJ<-qb|GFab|KcJ#FFF;E+!6^C?zib zf};F_{Nlu%_@vU5^wbg|elA|1I!?H{BsG|t)RcHUW+iFC)B(*ZPR=h%Efx~z;u7Ux N6cPj(Hu))Y5&-2%CU5`% diff --git a/tests/fixtures/onnx_genai_workflows/speculative/verifier/model.onnx b/tests/fixtures/onnx_genai_workflows/speculative/verifier/model.onnx index ec6b476f1f1c650f957a14d194e26e3adc90f235..d994492e184284f711782abf7a235bdcd1a045e8 100644 GIT binary patch delta 326 zcmZ3XI9Z96gWIY@X(Q`IUK>j;{(_?Xg8bstl=zbT?9{wsAvG=)4t600E-4|261}q2qRh0+)FQp4#FFF;E+!6^sPxJEc_r%oxO_Mmg**tTX5vy3Qov$bL1J-9 zd~s@FX=+|_D$rDpBx^3-lEk9)RG@Hjeo<<%kQSE)&;k`47R18c1JtdMl+49j0MrC> zv4LK8YNb#RR{&7A4-VZh%}iWsLMm9Cw>gBbhmlcb@8U00#mV_asl`I#T%z_sy@JSkS%69vl5Dwn3xFz8^Gcw$=x}KP iRjMJYWa1JQ65M=@uZNLIn``n1c2PkLmrpJhSO5T!eK~*t diff --git a/tests/fixtures/onnx_genai_workflows/tts/code_predictor/model.onnx b/tests/fixtures/onnx_genai_workflows/tts/code_predictor/model.onnx index b4039d6c580b82239edc97441ecd8482370f079d..7bc2e7571afecf81e958f73cf849c061f5e32c46 100644 GIT binary patch delta 6492 zcmchaU2s#^5yx2<Cl-b^s&=YwKJLafe!BUrJd=Jo^!8+ zb+4{m9_$P+V)xuVXZQSnd)EK`li~dD4K|@rGDM>NgQIc#hQ@kNedE?Y-Z7ASdFVL? zHPdZ#ZP0t*Wnvl}j*TYbfl#DBG!U2elYIh6ox~+o2BVQ!bT~Aes1HU5LV?)v(7<3Y z5j|et-=-o(Ah+!ytY%sBE?d*OKo<1@{ zOsogSc|&ZhXD~>;m1TeVhwZf zu5}LSHnD~|_|`gybfZ|q9QXS=g5^lzT2gi*pc%b7h6e&3ucQvVncfhc@@DXGvrR0hU6dCt8$Ij< zRi%hgUu&dNHN>UB-b)e5P2v_IsF=^pg8yA7pI`3RQE&Kw~=0}wU_CC+`+fCZqc4c z40DG^YZk&U(CLVsWo5*PNyth=%<$i3ev7s!;a!CO;+N%efNfJyq^fIV3VPsKT%3X} z#hS|rdueyVW%cBz3Zz`GuU+oT=cVTgt~SVlxQ=tg}+HQ(f77h z(x0~zDcegubyk|(X`)NdtvCCg&Rl@qZ-JYNI}A8=8U%FyklpU?(_n=i*n1kh#hRRL z;fz=YV`ik;Gtdb;`9b}<2|s4d%!{!q$O`_;Q9BmBuCL%tIQcr%$ik(KK_mZPK8?Eb z&Hk3mowDS7SrA%1zAcE zWX;;;Gn_gHh5GX3qL55`>D_MA=r={ie(tQHt-+N9t{lVk1{7x2fYYPs z{NWL-sv^l4%ZhF;X;Mi*`INKJV%3IZnJXG+<)<$lH>++uU3VkB9JAoP^B^(7izib> z`0a5~+noibd{}TCnlX49rZL z==Ia<&CRL;a0d1$g6FXfxDWxsv8OMwXLKmAC97|w7*JwmG2q7D3sA3sVCTO@T7RxU zzY;l+xb%)4_g>UT!imWqFyrDyc${-lF*^eeRo{_WNM2P~cy%V-r2_?VP^Q&v z3Ilm>eFi-9%_S`=38WeD%;#jlL1(^CsQIQI1IepKT#Q0KCeJgeZho6dHAO<&L%SA; zZdu0 zwT^KHVQ=0}Lc4Eyvbv!y+(Jhl*lE{iCUrZQeh0epp5gt%m)+}q#7&v7?HvPb7;0?F zdJDP;_56B!k;aZQbF44Be#6HYGnQXwox0-&n^m|a>K4q6svd{!*Psp;F2hD;)#4w_ zbn$mZ>Qgpcfqk+YyRIrPq(b;v{`%Q*<_A!kYmCWQV$ad+zfHE&j@ujY{sGuP zyK0J^9#4~g=T}x%)0w-D9J7u59?K`PWea}USt_Y2B6odd7-YGaaN#P5W{*nxau?^B z%3h18dGQ78{1BY1;6@a$fvD;=tlazHIsE55Yx!&eZTovZU24fM^LUjJdwDJw!u zHdFh({A{%!qRD^1AZqSj!1fz#h!x+1l-!fop#%@S#im$hb67P8wS`;BW_ssem6iJj zk42(`eCE!og`(XWvsYaLCU-XvYM>MUk%yovMf~Myqm+Y!TD7(q0jW6m9}tXPV9lWX zM^K_bp#;NDX0T*1El)zr)!9(u(}M#$@<5l4&FD(7#&goZUWF6yD|}R7%*eLp4w(|4 zst%ZZUm>6ZyXL^Dx={v)8ynZNx#^;tMP<)H6=y{bSj<@wO8j1}{?fS2&NB`DSp=e$ zaB93!h)J(1Sb-YsHZX>lj;68k=f;LFt;XwiSQ(s$yXyUUpy2axNiKxgf(o-=tsMT` zGP_F`n@Z&#BJ&z<*ag(#o@(~W-koPdlT1)c*6ZNI&L2Z{miZ7a`9U9+$MCBLcJK1M z)b>Fdmgp8RaEH#w)QTT{08eS}lR<*hTcO6Ro)h__;6^#xg(l^1g+e^@Az0I_q+PnO w@Q8Y}TnjL~Wl~PuSSE#A>BR2*5|7TTC>O?UZO-PT6+JhhPW56=+=Pk$0fy%vr2qf` delta 8088 zcmb7|Uu;v?8Nm0YjeR*W#&(h>cACV-u^s1+|C5AHX-KC*P)C+9sV$_P^Uv$z5XZqz zwM{_q(m`n{r5xa68zf;ZrD)QYhpH7)3!&;m)rYCw<0RBD4XLZDF4bQ4&~-cKT-ynV z?QKJ~Qf4zuOb(AtXB?%WKsXS3>@T+sCQ+PQZF?j=l*z_3iK+O3M9SIh%s=f95>t9IGbUum*fRB`mWkR| zMcFcJ?sOtGlHj8IchadymEG1#nwdGi)p;d~LdtDcTp5C@Xb-6p>y0OV8Gm8uz;q%L z2nEE$#)7bunp0ZZp$ESk21oUVCNNDmu4!p{kev}ZJ25>9)fy)B;FqJ?cH&JJXc@{r zc{!n>>XG{4R$@3X`TSURooxBz&tSX1i)>CGOk`NPu#m9$*aD=5MY8eUNLJWv*!4VA zH+PdQtJBp8ZYZ!;e^{UXaNhIM`Xl=E7a}XIKdMiERImLpefnb(`m5d>ABPePJV}@S z?)5zo8f&YWlqfrre6FW4EH`;4O)E7)9fV2lA&=5;CTc=aQToPvnE_;7G~siHpa%O6 z!8@>`K2C`MzWUC|JW50E!lD+iw-MUCf=VHm-bjfO;2@A1J(l#c$NYxARk;&0S zoK?x_*hqHrV4yoS8V^Toa^`ju`>F!`0*63emw(q1v`$`8tLGGu_>ZP zVk1kq5+%7G!!Lsa(?_AIx!|bvY-j$na4A?LRT+pC9lU{@qRJ%+y)|fOR-2IIpr%EB z#k}O8uN<9cz#53O=~(uu;zw5c+l$GwuDR?JEy=Yq9a1t%@~1R6A$++RyME331erKm z?vvWm91Te>x1V|QL)DunDG*OI9M=|A6OcuK2HK6Z>q~+yvrqv8JdfMG59+t>L#LA5 zMfA39A}nE)d>)0t)~F0v)^ae|stbX*;{>gwB&z;47$) zc3~-CshDeP+NNDYJUlAjI^+tJ1wmyR%Pb7N0$Zf^(^jmR1uMSvN`Ye{3(<51%>K5h zT=wsxu;}!+N!HZ_b>$?P=A=u9EjW3JmG0G(%s@dn2OG3E$?-s7uK!nc^uQrKn9I$& z-EG<-;+Wf0fZEbw0}dVp1#irOvAI1gX&{4YZdMv^Qzv07&L3p{5PVCYD)5av#h|`n zkR!5t5iIH<-0K${G<8O@y>sI7UNF;zQ#L&DT9GYrc#(SxBZ72uQb^*kUS zI=2eQ^mSc;gm>%k$T=u$lxrPDA84n2H;H?@cEs@3Tj0XbG?;lKSYPrCq5U6(<@$>1 zbGkR;xDXnVN#Lv}G~XBIk1y(1mdW|XlDtQXNFy;$?>xm5!$ddVEtg$u zt&p{e-DyeX0t;XNo5GPM{O|9+Bu3VcO$phpOVbRByIt)9@ah7j=Lf`eTX>f&4Y4Twl3Umj~AeK^C=!(hVC z-+&6qMFh`^Qr0HPYYo03F`~6efvgp8zR6NR0L;su1-+WW+ayoFJs7+K4RY;W3Rhwv z$t@>Nwu1pbzXJXO!NkD-f}jgwEmFSbtybX+JpKDc0LZx5_`(Yx8cbD0iBR1;s5{u diff --git a/tests/fixtures/onnx_genai_workflows/tts/code_predictor_prefill/model.onnx b/tests/fixtures/onnx_genai_workflows/tts/code_predictor_prefill/model.onnx index 3d18439dc78a95a69f11baa66cf48ccf19548f70..c6809bd50f5181889c04c5e26595dd5746b95f0b 100644 GIT binary patch delta 190 zcmaFIbdgzvgIkC#H$N$}v{;JOK+jOmz-kHeM3EzEo?N^oi861~a)7-bmsCKod)s`+#AX66-?mK4XQ<|d`46bspL*>ErlSzy($fyvE8 zii@i(-aWA-Bef{rP>6?%3#gini-C)YgTsl1i$OrbNsx;dXq8@WW^z$}aei8f03!fk C3M?J~ delta 121 zcmeytwvkPkgIkC#H$N$}v{;JOK+jOmz^Y}Vun&`n1Q%;bes*e}5I+|$2fGj_7rPMa z#RNGe$` zDX}CugNuoSg-b$63|%s!*5v$@)M6o7E@=*SA#pBtA=aeClH?37CJq)Z4k5O* zqQuxsBcmt7Z1c9sF8tO{y-zW@f(T5!JAEmLKzuNCwB|0ar1HU z0L|gxV&Gz!>>#Q(`H!$1s}+|E=(5-p2I zL`HCCdTxGZO1y~>x$T389z8*p)yq*j!~XQmVj zscT#6 zj$EQ;@nF}dWag#o1!tz`=4YnFn;?Y(9;bq>wB-^+Ho~_wMpIQ>Acx`4f?Fqwy& LON~fPH$`Oux*eXC diff --git a/tests/fixtures/onnx_genai_workflows/tts/policies/cache_length_update.onnx b/tests/fixtures/onnx_genai_workflows/tts/policies/cache_length_update.onnx index 72ad12c9f4e344e870cd12220f36d4f2593cc2b3..45764e632822b23fc759eb608a1f8079786bf246 100644 GIT binary patch delta 178 zcmZo-zRb+a!7ap=o1c_fT5NTIc_MF5y%85nPHI|-kQSE)2fL68my(b|L4HnVa;08! zVsb`md`@a!dPzonX+cV2Nvd8_Vo7pFlra};QD%AuNt%*?no9CZ5_3>(lEi7%#0QH2 DWYjl; delta 97 zcmcc2+{Dbw!7ap=o1c_fT5R>6X(DfryEqq1PHI|-5I+|$2fGj_7rPK^QesJRMwA2> ZYf)x;2Bu6BP^KimBrykB&tyZ!MF3MT6yyK^ diff --git a/tests/fixtures/onnx_genai_workflows/tts/policies/code_frame_update.onnx b/tests/fixtures/onnx_genai_workflows/tts/policies/code_frame_update.onnx index 76ba875d751abeb777abedb6463134ac1114ac56..102314d5d32107012502c49a271bf2dd75ed6280 100644 GIT binary patch delta 237 zcmZ3(^MQwlgIkC#H$N$}wAkwIMxGRwdTTE3w4%h^)cEB5l+NCJaedTwHmn6(tZmh;qQ>8!VHgWVv|D;zRR_3rkZ|t5V}lg?PBQfCjR0 KF--1eT?hbe=|uwo delta 191 zcmeysvxbL=~KGCHX~_Lcv^t9PC2ATuMR;1^GFd z$(4G^`6;RKQ0e%@f`Zh%6uqRxlH?37Wg$g8YSM}lb5o1Cm^fIX%(+;>JRvWaFSVisra&m1D->uzAesR%X;kN_3#k%xo`SZ5A=l)5CIu-)F5a^E d(7fWp($v(d)ObT70WLnERa{&QlQ%Fe1pqcFWJ&-4 delta 192 zcmcc0ewB@vgIkC#H$N$}wAktp+eY5kjP)8^>=~KGCHX~_Lb6=a9PC2kTx$&iaDFSVj19%h&#mpo9BG(r)N z5EsNS1#JaGkbzlJTs&oy|1n7@YH;zE#fRn<7nY``R;9)p0^J}A)GP?sJlTMGDFB}F BE&c!i diff --git a/tests/fixtures/onnx_genai_workflows/tts/policies/codec_layout.onnx b/tests/fixtures/onnx_genai_workflows/tts/policies/codec_layout.onnx index 31ce6b3a9b69ef63df9a7b26d8b599edb29c19b0..dfca6fee18ae6ad030cbe3c796112a7eafca77fc 100644 GIT binary patch delta 181 zcmZ3({D+x`gIkC#H$N$}wAkw9M4nCc0bJ}EnZ+gfMU_IXT+STqLiSwJLXrjfIho0o zddc}IsmbvIB{Mm(Bvmgdu_QSoNsNmFCMU$j#R1gD!o|eFm?gr+Q5Nrfs TAjHAN#=$Pc#Kkaq57Q+84z3^m delta 79 zcmZ3;-p9_%!7ap=o1c_fT5R>3Z6j|Ildw1!OG*P61ax&6f b9A)v&iNz)H24FdMh+I-)Npi;I^Gufj$pI1- diff --git a/tests/fixtures/onnx_genai_workflows/tts/policies/last_token_logits.onnx b/tests/fixtures/onnx_genai_workflows/tts/policies/last_token_logits.onnx index f674183330680e7f8e3468568778f80f8327ea0b..4142d9050c9fdc97b12ded4b6b105e73487316e5 100644 GIT binary patch delta 241 zcmey)@`aU`gIkC#H$N$}wAkt)>qg#S#`-!gww(O*%#vcE60RZ+Mxi_|1tGbD{G80> zO1+%K;*$81{Or`cc!(Okq{Ncs3@&9MMVu;%QwvK|^O94!RE3mrDk{rQPE0DzNi3=i dON!&-2HOL5g+G@s&=sCUxB{OAlXo!g2LSQqRMh|g delta 115 zcmeyu`kjTBgIkC#H$N$}wAkti%SPT{#t36Bww(O*%#vauO)hl~Mj>S`b|KcJ#FFF; zE-oRC;?%;@)V$GSvi#)4q|%(kqRJ#4F7BMf;*xl%PFXH#piXgwPN3?^rcC<* DBibI7 diff --git a/tests/fixtures/onnx_genai_workflows/tts/policies/predictor_body_sampler.onnx b/tests/fixtures/onnx_genai_workflows/tts/policies/predictor_body_sampler.onnx index e2685bfe5abf6f9a185f76a54debcdd4d02e72e5..5f78f3304d884cd9404c5d2beb36397e77ca81e9 100644 GIT binary patch delta 192 zcmbQi@`;6qgIkC#H$N$}wAkwYM4p}X$y{tX`RSP@#X^x>;T()Y!Cb0B$_4p3naP!U z1x2YTnaL&jMe#}bDV6cXiMa(isYQB8i6zMyTv|dJ_?4CACnqMA<|GzXCYf=umgHxr i<_YO?X>+g(sUz%A#%oCv<4i^_#mVm&B`2#fegXhhEJ1bv delta 87 zcmeywGJ}POgIkC#H$N$}wAiY3BF|1+buPA?{PfI{Vj)>BX%0pqaV~Zt)}+LePx# diff --git a/tests/fixtures/onnx_genai_workflows/tts/policies/predictor_prefill_sampler.onnx b/tests/fixtures/onnx_genai_workflows/tts/policies/predictor_prefill_sampler.onnx index d67d86e339afd9306caf255dc6666a3de6190f98..b5ab9658c9b24c00658df8c476b003a4f910a928 100644 GIT binary patch delta 143 zcmbQu@{fgwgIkC#H$N$}wAkwPM4p{8nOtl+`RSP@#X|91u^fy-kz5)=Y6bZ@naP!U zlY1E@go^Z%5=)XZxO9cIk%Tph^vd#+6O&4F5{oL6EV)=q^0QO(gbcX!IM{`>(9Gb4 Un$g5KhmlKh@<&F=$!d&W0TN;>4gdfE delta 87 zcmeyzGMj~mgIkC#H$N$}wAiX+BF|1+buPA?{PfI{Vj)>BX%0pqaV~Zt)}+LeAsY!9v7RaM<09I%IFznFvX<2$mrd=1GFR zU?bK(f?#WJAvjq`p;+t|S-UCD@crL^&YS-}SW%Wx+I9`vV~4MU@nm~aK2w7LW<8%= zQp>``Hb|Qx8+H_cu3#DMdagyyzA6XW)a0&@8P_?%)S;Z}799{@HFR!X!+lhbFgpV= z((Mt)Brv4}ig62IR_B~JoVpISb=E;Hw3Q5^u_7FyR+I=j<5=>sMV!)JiHQ$26ZGJ+ z=Q8Oe`BXMQpegY<7QTwK2>6-#6mQ%CxM8^MzF17Jo`cYiWdxa<)zp$zbmkb8c||DX@y_}V8NU2tV8q>&e^RSO3knk-Y6rEdks71Fe0gzX1kiZ)gAj delta 785 zcmaD_^{9xKgIkC#H$N$}wAku=(Zr3i8?}P<>MgnW3W`$GGIMg`Q*)D2Q;LPuxl}nA zg%r8ig;Qu9hO^Yh|!6N|HjG`Z9{ z*oBl4S|x^w0c5p?GP?J-Q!nA~U}A|}iw2sD-#VJs5|W0KqC^9B+Ix?I{o zkE$a}ac~I>3Bp_q_9@6}E@qAf#w3r)afVWcBx-Wwnp`L*W++6WHjl}5MpC3Y&}}lm mF&Vy^XH14o`X*%9w9kYLoBU15u<4$ulmxCI-yCMf&kg`Y5XKDv diff --git a/tests/fixtures/onnx_genai_workflows/tts/policies/predictor_step_update.onnx b/tests/fixtures/onnx_genai_workflows/tts/policies/predictor_step_update.onnx index 5c2e73d7a71c4f373a8eaac8f1919cc9f7b06eb7..306634d708b9cb702261aafc596e47aabdfa9891 100644 GIT binary patch delta 374 zcmey)zf_QygIkC#H$N$}wAgBr;6~m%to3PJe2FC`sd*)t`FZiViN)DMAzVQm>_UEA zDnd#H`8k=%m3jq5sVSMsCHY11#U-f)@udYRi6yCeNr@%N8C+^YDtJ{T=jWBAR+L0J zaPbu67iWU4$xJC0GT_n!+M-oYq%}+&j7j-i!g)ZW;?dk2!4(EHDTqjew1hN}-5Z7B cEvSPPv=t1KT)70nHez!*LHjmOW=mlN0BsL|U;qFB delta 199 zcmZ1~_?@4ZgIkC#H$N$}wAkti|3=_Q@3 z>_V(bi6zMyT%1De$@zIDsTCzr%3M4J`Nf%F#hEF^Lc&~vK()LGwM-n0Nv2%Fc|gVS nXy(ds$pF<#Ak^{-aVKj-3|7!qFig_m5(MkRWy5Ai_7p|{x=%G< diff --git a/tests/fixtures/onnx_genai_workflows/tts/policies/setup_predictor_sampler.onnx b/tests/fixtures/onnx_genai_workflows/tts/policies/setup_predictor_sampler.onnx index a8bbba82c4edcee282d0dba1e7a6fabbd0cdd57d..6a5ca34f02a182e41dfb043ee54e9a92c401c67f 100644 GIT binary patch delta 195 zcmbQq@{NUugIkC#H$N$}wAkwLM4p}Xsa$M1`RSP@#X`|sksORdpzngrC9BX%0pqaV~Zt)}+Lee(0PzwO^8f$< diff --git a/tests/fixtures/onnx_genai_workflows/tts/policies/setup_talker_sampler.onnx b/tests/fixtures/onnx_genai_workflows/tts/policies/setup_talker_sampler.onnx index f609a5b8a6a49d4be6dc363d14cf83440ca7953c..df3db0d1e71e9b526405b67134c204d485b37f15 100644 GIT binary patch delta 172 zcmbQn@`{CrgIkC#H$N$}wAkwUM4p}X@my><`RSP@#X_N6!5oZ2{#?pJiUs*OnaP!U z#i=Ew1@R?`IoYX2@x_U`1v#ljdP#{T$r)T4LTdOFmE|WVCY9zS7F8yhaIu!;XQ$=~ WX>)0EunVc8nlFTD!Dc2#5k>${;5ion delta 87 zcmaFGGL40YgIkC#H$N$}wAiY7BF|1+buPA?{PfI{Vj)>BX%0pqaV~Zt)}+LeY-0P6r0)c^nh diff --git a/tests/fixtures/onnx_genai_workflows/tts/policies/talker_sampler.onnx b/tests/fixtures/onnx_genai_workflows/tts/policies/talker_sampler.onnx index 83f3e3cdb794a651f74c3f9b7191e0fe327971f0..9ea6d838f9d735672f69e84200f381839be31248 100644 GIT binary patch delta 169 zcmbQha*2hPgIkC#H$N$}wAgAt%S7HC^}$?hIr-_CCB;IXT<#o7#{%ux#u=C delta 88 zcmcb_GJ%DcgIkC#H$N$}wAkt|^F-bqw(4ALIr-_CCB;IrT+$qjLgHNPLaa%NCCM3F pJVIP$`N@e%r8$X3l}Qp@tR?x`sd+;DT)Z6YLYyFVn_U?B7y$_V6&(No diff --git a/tests/fixtures/onnx_genai_workflows/tts/policies/talker_state_initializer.onnx b/tests/fixtures/onnx_genai_workflows/tts/policies/talker_state_initializer.onnx index d16e72d0909679244091d68a91161a4a10207b2e..75a589d22319597d99196ab56d2af2e5ff653506 100644 GIT binary patch delta 536 zcmZ2zb=a1dgIkC#H$N$}wAgBc?Zl0;8@2SrRVumo3W`$GGIMg`Q*)D2Q;LPsxl%b8 zg%Y{cg;Wdjb25`F^(OxplMvKPN-Rmv;4%_2KoYgmgJ~{KEi6sVOHSot;^0WC=Hg2% zDM`&M$;{7-&rK}O7RuyG=U^8~#$l+1kQtJpc6tSg#U=4T6XSDI^Axldpsb`~E}nw? z;!LoqnLw8&b0q=|jRjiCjKg3(Asr-x1@vI91R0!E$2EyfL>TN}cAWk-6EZc4gIkC#H$N$}wAiZ8YT`!Ojaqu*^_E_W;2tr9|F1&PHa@j!j?IjMOH+6qusk}(%gL4I+j0a$Y;&{}0KMW9|;pvlY#{d__^ zFe^d&lkB)Av5AO+oyd;tL@^| diff --git a/tests/fixtures/onnx_genai_workflows/tts/policies/talker_step_update.onnx b/tests/fixtures/onnx_genai_workflows/tts/policies/talker_step_update.onnx index 5812ac850e80197463e5e98a44ed90a6cb4bd74a..4aa8d4d74f9d3d5eeab65e1e99f653491abf3fa9 100644 GIT binary patch delta 347 zcmeyvKSPk0gIkC#H$N$}wAiXua3k+c*7^i4zQmG})Vz|+{Ji+w#NuorKQ12*b|DWg zMIrfu{G80>O1+ZAob1%1_~Me(g80&cl*E!$y`;pF_Q@3 z>_V(bi6zMyT%1De$@zIDsTCzr>Rdbp`Nf%F#hEF^LgHMaK(&HgOdQMz^-LU$Nv2%F qc|Z;EXlBcC$pF<$Ak-T33UNaXSI|~4Ow!>J1na|R%Vrby6h;67o;9KX diff --git a/tests/fixtures/onnx_genai_workflows/tts/policies/token_to_slot.onnx b/tests/fixtures/onnx_genai_workflows/tts/policies/token_to_slot.onnx index eb8e4b6599f5d6095617f9fe0d294eb1d098d216..2ac337e271fb50fb203de6db46cf55db8ee7532b 100644 GIT binary patch delta 111 zcmdnS{D7H<~Z1rheyah%1xdkQhCHdK@dBsAZT)`adLjGJT zLP`bsIho0odL<>r@x>*HC8_b5d6^}di8+~7sYQB8i6zMyTcug;;JXk8YT{wBs(rHuoL5pbMi~z1`w;CG0BCC2V{4A cMrLtIeo>{636~KEyRwiTL7#5cWlv=U0Ns^nxc~qF delta 173 zcmZ21G)a(`gIkC#H$N$}wAkvuz((E)Z1p-^yah%1xdkQhCHdK@dBsA~T#_8@LZV#k zLaa%NCCM3FTtXbhsfDGfdC93sN?hD&MTxno@yYoqsl`IVT!KKgya=^S94twSTwGwI q;)`?gORy=Q8&$LC6P`E>mfd& zEU)4w;;;E43u^@-yDV^NA&ZR&Exy2xgrZQdn9FjRgq$KXcnVo4z6#t@MeZQ&HDrh= z-sg)8*p|=9qDB*oDJ?Xm5t*uiRlJLmF0Q~~tZ|Uc;`?;w$#Y0X#NZ$a)TWC`-USgi z<`c5QrGzbRBato$awr_qe*7QVv&+xnx^!+uRAh1w$7#lqb245uEjFQJmcNUgX*KMs zm`n;;PF50%KyKjc6dusScPp;0Zh#(^U&8kMx*)5sA{?R8OV@eI1YruV)tE}jtk$f%|JE#l) z8GuHefTP>}HZ29Cr7(bPbJ8@7yI^$iDlBCz@I7MH3xbe~THx!Zza&)uzwgrthhCOR z;eHG3$E>8X69gyAICRj*#6oHq#mxMlB^-#Co- z;sfUY-^=HSEhNUeSteoS}6^ft>3!pTz~*k@f%otKrMX*TGr$U#H1 zQV5RcueoZpT51=BXUvd)*C78$G*%!jAJ4~Bq0hhn#I_3tdW*|K|5drnKv&c z65q2%JQ`zSI2%bk_%C=e`X9i>g9(YhHn1rM+NGDKd7t^qxy0ue zcs9%Dp^2tbW6;6LGM|@)LY~cW(h@TWKAIxn8myaf6gf#|308r*{H~+7R215a1xZk` zL&6ayMJStSV;z5jHZ_TRiLMY6fFUB!Tc2y94RD^Z9drGyvpF?|-^HG4Lzr|BTC;`6 zf-Er}=p$Ucb*>Ydz)0^~(*dw$6{CwnCd21hN#4HuIwW-jm8l$3h#pb2)f4{=%VS}XUjK@Hxj_CvP9MrzY?4%<>pH{HWF>%_Tnwv zS?vr*Vyb59r0Y{VgPo=Kh8y^G(CJf-Vj@R5N1jvXRUSt*+#rpZ1Z&-&a@1%UuUZNiNXs9{?{YT_4F%slZuD zKE$6Raku7g0WE9OhI!;FcHE8vei)|k$r`w{{w4(G!9tY6PixK^-1c%EJX%O1ycr*r zT^Nc}IKEy-gZkx7!sOEwt?Ur*;g9vU31JCO-E?c!d<(Nvu@F@ zT?{VU_>1Z%)My`$gxm9;_j%44>Eu->KI&u+r7g+*p5Go3qhe%tpc8np0S4Cti)~(7dS-q}IPyR7R-^OBY9;iE0?Vs3C5T?!Vk;M>WW?Fdq+dy_)UoJ{o19O^X5{*99u*k8ww zQxHt3X(0+5IJEgOTW~l{8Qv7WHBq;S?}ZhQDf<(NSS(U=Yjj_HS|SM~rI0I@o?S62 sUW|z`_mnry>li!_W2fvTG0dieLS|MlKAn6sCAkw3bAU=IjG@cn7nF94od5s; delta 714 zcmZ3#k#X}zMoA8CA-3H7q|DM{DOLkLLp=kl){Tr@E10<^H@fO=beHz3H{jwZNv$Y} z&rB&6QsGkKU>B0(Vi#gfN-Rmv;1U)R1WOjD7M7;wC8x&cq~;~ra`6=urKV-()0EFbb(6^oa?Hz?2|s{&Lc_Oovwr^=pyeV;Z0X<|gSs<8a@^^nBu*VGf^d=_+ zv+J4@ryuS)y~+HlY`PXC=+d8j-c)(>zJL-ICTp(Ag;t8fM&?4IT*AOu;D^V;=ETqg uj1YMxAu|l=$pPVxjFyuf%oU|AFeJGcz^;QBE^1EOCa`*OQ!MHYxMU^=vWib`5MdED z!y>Og`2vp;C){?RX=;P+XTFQv@+fQ3Quk N3Y>6XPIi!01OPUvR9FB2 delta 499 zcmexp_0UR$gIkC#H$N$}v{;JOK+jOmz-o)t%BrxhjUrp71dr=%7O z33CZ@unX~Wu?w*#C6**-a4~VPM5%Idm84db#HZ#arKSjpbBS^=3JD@BVFW7UNYdcq z&CDw(Eh&a5EJjv2`Gdq~CNr*yh6)nc>|o5Y0?NtD8Y1g~IRI>W7D%(aq!D)ArXb15 zv!%S4z~;!A;x@;eOJ=g7jI@jyZY5@uFYqXdV+jhzEQm$(q`ereCx4Vtlrg|^4ywtoBy^_S7?9`(8lGKWl_~Me(0==ZflH?37Js}-DYD$U{GjlTY(&LL$ z3rkb;l2f^uI5?urxi~XRQi~ExGV}9=0A delta 162 zcmX@Ycb`|BgIkC#H$N$}v{;JOK+jOmz-kjO*K)?m{;c8~^%Phc9Y8D*4K6hfMj<6G zb|KcJ#FFF;E@2_T$s1V3%w@PZGfPs75=%1k^Mv@hcsbaGIN_?2RJpiHQY%X0Q*)D2 uQ-s91M1h6~f}|%Kv)z?b@6ht z4z_kafp26h&PEbL7Gj~@48wf?Ki`?pxySq@Np3ylpI!>D#r#Zg32Ddmd?^fUnP>~T z54$8z%P>!K7H@gnMOW%JTHMCAZZc`zK!Gkqj|ATO=voK~qSY8=n3U?K!(F4rM1xkL zLZ<2QC^i`8mgs7f>oFmXK3rSV`pa4sQXUshFTzp9cc82iHb)bN1ER$P^{4?6VYZjn zz&oSEnm!7BT~`}=1}DnN;W%`HB%nx_e(A5|kXjvbgP86LWP^Xjo98gq z_V&2ip^+dm=Drb_Aop?e3esL1MfeZCbyVO~agWAvLmkUhi|UYEvgTl(p%1fq4cVW7 mo_P9~i0%hGFKUqK7#C>!7FLkC#sEOnfZC~xrxQuLfTxK9PC1>2rW`V5|i(6 zO4e(0@f742XM$B`0?n4@k_0LgMJPlvJIawuv>>s#BtAQ}GQKP^r!=)#&p;2z5z^pN z1G-2Fp^J%wiHl!|7pxk|X&h0`T;j+^Oy0xAQ%|(+BrPtsoc#36l42nRE;*o$QV1K7 z+{wkv(ZHBw#l>4tlv)f71*ju6xzvGLlo47$jua5$vn$CjNz92yc4m?d7e7KX*sDUg jv`-f0nW|^V#a9*|l$uhSoa&pJm=|w?HZEcF3goSAtAr6uv1DaAt0T#g*3T_tCCM3FazZlLB#ToEOH=cbQ=@FT_%d@7(^KQqQWHx`i&Be)w7E1n z7==`^85@r`6 wykejOQ>#+r4TSi(c!2)o;9>wmCJx3dDK3t(c<038l6XU~93w<-@-^1k0JhFpMgRZ+ delta 241 zcmdnY`;mv0gWKvE&qmg3EcF^(oSAtAr6uv1DaAt4T#_8@LZV#kLaa%NCCM3FTtXbh zsfDGfdC93!np}LDxryni@oA}vC8b5F#X{m-q8yAuf(YeI94uT+912MWT)bdIi{n#s zlTv{O%X7&BRZ1dM!VI3gk5wrOsJtvbG_M$FQEF9cya8A_J2vG&cK|J5%+luKD2sPa REG~&R#BDZE@ni|M*#ILiJahm6 diff --git a/tests/fixtures/onnx_genai_workflows/vlm/policies/cache_length_update.onnx b/tests/fixtures/onnx_genai_workflows/vlm/policies/cache_length_update.onnx index 72ad12c9f4e344e870cd12220f36d4f2593cc2b3..45764e632822b23fc759eb608a1f8079786bf246 100644 GIT binary patch delta 178 zcmZo-zRb+a!7ap=o1c_fT5NTIc_MF5y%85nPHI|-kQSE)2fL68my(b|L4HnVa;08! zVsb`md`@a!dPzonX+cV2Nvd8_Vo7pFlra};QD%AuNt%*?no9CZ5_3>(lEi7%#0QH2 DWYjl; delta 97 zcmcc2+{Dbw!7ap=o1c_fT5R>6X(DfryEqq1PHI|-5I+|$2fGj_7rPK^QesJRMwA2> ZYf)x;2Bu6BP^KimBrykB&tyZ!MF3MT6yyK^ diff --git a/tests/fixtures/onnx_genai_workflows/vlm/policies/decoder_state_initializer.onnx b/tests/fixtures/onnx_genai_workflows/vlm/policies/decoder_state_initializer.onnx index a2d78bb4b680638f98b2ec428b1443e7f3b2dc25..4565c7336226903f9bd4b082c6dec88bd6d001b5 100644 GIT binary patch delta 844 zcmcK2PfElv6b5iwLBrF!2+mA#H7ijDN*xCA&#c^e0zuJ`w9#hTCbh{dbkl`+f}TK6 zpm+pN;0-*3UcexZh?`o$T^5i0^1Xa-+fQQ;E#moHF)`J<_u~iewOdWWD@v6Yh8UG{ zsc{SpE#OmF$7_YkM06R%JW??)i8j=5BBU^aX5xyM0iz~LV2H!U$xlP`iz$~8XB8|{ zW4Nr0P?F?Sr#Oey!sjjEJ-pKlQ*FqM?GO61BgZF&Fz diff --git a/tests/fixtures/onnx_genai_workflows/vlm/policies/decoder_step_update.onnx b/tests/fixtures/onnx_genai_workflows/vlm/policies/decoder_step_update.onnx index b2a4c7288de43c04c5ae83a2119d81bd95d87b9b..390f62573f15a272042ba9b3688c9fcdea316696 100644 GIT binary patch delta 362 zcmey$KV6WQgIkC#H$N$}wAiXea3k+6*7_tazQmG})Vz|+{Ji+w#NuqB04_fcb|Eh= zB_V}^{G80>O1+fSqWI#H)PnfZf|SIPRK29clH?376(J=&YLfHwN>VFIqU^YM z3i69H!B%9Z6btEb=>Toe0NGHFM?1s{CJx4=Y%bwEph59yE)C%d0-EDTluI>))Ua9Q g2lWorsS4T(hDpv`f?(Tlxfz=|P-8YvWJ_TL01X6sCIA2c delta 199 zcmbO(_?4fRgIkC#H$N$}wAktq|3=_Q@3 z>_V(bi6zMyT%1De$@zIDsTCzr%3M4J`Nf%F#hEF^Lc&~vK()LGwM-n0Nv2%Fc|gVS nXy(ds$pF<#Ak^{-aVKj-3|7!qFig_m5(MkRWy5A;_7p|{s@F9= diff --git a/tests/fixtures/onnx_genai_workflows/vlm/policies/last_token_logits.onnx b/tests/fixtures/onnx_genai_workflows/vlm/policies/last_token_logits.onnx index f674183330680e7f8e3468568778f80f8327ea0b..4142d9050c9fdc97b12ded4b6b105e73487316e5 100644 GIT binary patch delta 241 zcmey)@`aU`gIkC#H$N$}wAkt)>qg#S#`-!gww(O*%#vcE60RZ+Mxi_|1tGbD{G80> zO1+%K;*$81{Or`cc!(Okq{Ncs3@&9MMVu;%QwvK|^O94!RE3mrDk{rQPE0DzNi3=i dON!&-2HOL5g+G@s&=sCUxB{OAlXo!g2LSQqRMh|g delta 115 zcmeyu`kjTBgIkC#H$N$}wAkti%SPT{#t36Bww(O*%#vauO)hl~Mj>S`b|KcJ#FFF; zE-oRC;?%;@)V$GSvi#)4q|%(kqRJ#4F7BMf;*xl%PFXH#piXgwPN3?^rcC<* DBibI7 diff --git a/tests/fixtures/onnx_genai_workflows/vlm/policies/termination.onnx b/tests/fixtures/onnx_genai_workflows/vlm/policies/termination.onnx index 3693811ae428bef7c7f17652199285b2ca557a9d..72763e63ade7f01e72cff99022dea83883b2f75b 100644 GIT binary patch delta 557 zcma)&yG{Z@6o!EicJ}OIG2#WndLdR8;0C_%yL06Xb@ z06XJL`54ZQ28bSOGoF@@clL|1ITBCRTGE5dJ zIE116v#ZYjn73SYZ_^W^d!6oCrMr?%$9!rl)a;PE{QQg3q#7fE8hGiVVR?pJ Yj6MUp?i}r}WWjYlw>{_H-tCwC2a<=600000 delta 521 zcmZvZ%}T>S6osi&letMOo%+|>Rt1SWF;+#~iHLiDHbS{fn}H_U32i1p-1H4xn7t3+ zPTc!c?sXDtiNx)kbMCqG-IvyL%hCd;lnm}Z5#|3(BNGo9J8YQ67 z`5u$~!EF;X9)y`zg^JP)AXF345Y@9%i6{A(NiPt^>=rBz*#piT#qqsWsFvM8@GMMz ziWAqko@|2^iH2d#_4qb+t z5-X{g`-8s&B*ybm#P0EcHegLBbP2k@0P;>7bk)W}iHDg{Q^x;%n`!eUccVVPQlO3s zWRk8~-)Mso{jlAxAD5A*q7=oXq4( zy^{Rw)V%oO#N2|M)FQp4#FFF;E=3`Etg>bK$%#p&If+G;NxEFDV4XrLTuL15LUQ;_ W{N%)>(wxMi$|MOc){^|})I1@6E?y3HAx@CG&5n%hi~s>46#D=G diff --git a/tests/fixtures/onnx_genai_workflows/vlm/policies/token_state_update.onnx b/tests/fixtures/onnx_genai_workflows/vlm/policies/token_state_update.onnx index 4f3026a099f6c88007fd54d27c96dd2b224c70b7..9a8aafaf3172c300217724e5f79d073252f30bbb 100644 GIT binary patch delta 254 zcmcb?IfILrgIkC#H$N$}wAiYZYa{PZrg{r5_TN* W8sxqv#x6!K#mVbfBqx7kegpsp|3_i~ delta 189 zcmbQib%T?agIkC#H$N$}wAkt>=SJS2On!1)?8&7?MX7luLc&~v9PC28TFwk~9}fUTOu3zR8Q2m85}2mc{#)=ENIf7&-X{ qv!b>N7jId7XkKw)X=-X!YP>OqDv-rZj9rXeij&{5NKQ6mc?1CV<|@tr From 532b83dd44777a267773e4882594587fb0466ddc Mon Sep 17 00:00:00 2001 From: justinchuby Date: Thu, 13 Aug 2026 16:31:43 +0000 Subject: [PATCH 036/151] Support reduced-precision decoder logits policies Build the last-token selector with the decoder logits dtype and cast its sampled policy output to the float32 sampler ABI. This lets BF16 and FP16 decoder packages execute without regenerating neural model artifacts. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: justinchuby --- src/mobius/generation/_policy_components.py | 10 +++++--- .../generation/_policy_components_test.py | 9 +++++++ .../onnx_genai/workflow_metadata.py | 25 +++++++++---------- 3 files changed, 28 insertions(+), 16 deletions(-) diff --git a/src/mobius/generation/_policy_components.py b/src/mobius/generation/_policy_components.py index c1fcf4866..7eb363a99 100644 --- a/src/mobius/generation/_policy_components.py +++ b/src/mobius/generation/_policy_components.py @@ -159,15 +159,19 @@ def build_greedy_sampler(*, effect: str = "sample") -> PolicyComponent: ) -def build_last_token_logits() -> PolicyComponent: - """Build ``[B,T,V] -> [B,V]`` selection for decoder sampling.""" +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=ir.DataType.FLOAT, + 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, {}) diff --git a/src/mobius/generation/_policy_components_test.py b/src/mobius/generation/_policy_components_test.py index d5245953a..fb04025e1 100644 --- a/src/mobius/generation/_policy_components_test.py +++ b/src/mobius/generation/_policy_components_test.py @@ -162,6 +162,15 @@ def test_last_token_logits_and_continue_predicate_runtime(tmp_path): (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, diff --git a/src/mobius/integrations/onnx_genai/workflow_metadata.py b/src/mobius/integrations/onnx_genai/workflow_metadata.py index b84097152..e3da5d910 100644 --- a/src/mobius/integrations/onnx_genai/workflow_metadata.py +++ b/src/mobius/integrations/onnx_genai/workflow_metadata.py @@ -2572,7 +2572,7 @@ def build_vlm_workflow_metadata( token_state_update=True, ), ) - pkg.add_policy_component("last_token_logits", build_last_token_logits()) + pkg.add_policy_component("last_token_logits", build_last_token_logits(logits_output.dtype)) pkg.add_policy_component( "decoder_state_initializer", build_decoder_state_initializer( @@ -2736,7 +2736,7 @@ def build_vlm_workflow_metadata( body_decoder_outputs = {logits_output.name: "decoder.body.logits"} logits_contract = _contract(logits_output) last_logits_contract = { - "dtype": logits_contract["dtype"], + "dtype": "float32", "rank": 2, "shape": [logits_contract["shape"][0], logits_contract["shape"][-1]], } @@ -4000,16 +4000,6 @@ def build_decoder_workflow_metadata( if len(pkg) != 1: raise ValueError("decoder workflow requires exactly one neural component") decoder_name, decoder = next(iter(pkg.items())) - 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()) - inputs = list(decoder.graph.inputs) outputs = list(decoder.graph.outputs) token_input = next( @@ -4052,6 +4042,15 @@ def build_decoder_workflow_metadata( 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]] = [] @@ -4344,7 +4343,7 @@ def build_decoder_workflow_metadata( body_decoder_outputs = {logits_output.name: "decoder.body.logits"} logits_contract = _contract(logits_output) last_logits_contract = { - "dtype": logits_contract["dtype"], + "dtype": "float32", "rank": 2, "shape": [logits_contract["shape"][0], logits_contract["shape"][-1]], } From 37bba25f427bebce6a30e20e76d23878518538c6 Mon Sep 17 00:00:00 2001 From: justinchuby Date: Thu, 13 Aug 2026 16:33:24 +0000 Subject: [PATCH 037/151] Add paired real Muse H200 benchmark harness Record one immutable 68-token workload and provide synchronized native ORT GenAI and metadata-workflow runners. Results include package hashes, token identities, TTFT, and steady decode throughput for direct parity and performance comparison. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: justinchuby --- benchmarks/README.md | 29 +++++ benchmarks/muse_workflow_h200.json | 44 +++++++ scripts/benchmark_muse_native.py | 188 +++++++++++++++++++++++++++++ scripts/benchmark_muse_workflow.py | 108 +++++++++++++++++ 4 files changed, 369 insertions(+) create mode 100644 benchmarks/README.md create mode 100644 benchmarks/muse_workflow_h200.json create mode 100644 scripts/benchmark_muse_native.py create mode 100644 scripts/benchmark_muse_workflow.py diff --git a/benchmarks/README.md b/benchmarks/README.md new file mode 100644 index 000000000..44fc50b52 --- /dev/null +++ b/benchmarks/README.md @@ -0,0 +1,29 @@ +# 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 same rendered 68-token prompt, greedy parameters, token budget, +warmups, and steady-decode window. + +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` 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_workflow_h200.json b/benchmarks/muse_workflow_h200.json new file mode 100644 index 000000000..f7bab0255 --- /dev/null +++ b/benchmarks/muse_workflow_h200.json @@ -0,0 +1,44 @@ +{ + "scenario": "muse-glimmer-30b-int4-workflow-vs-native", + "package": { + "repository": "justinchuby/Muse-Glimmer-30B-ONNX-INT4-CUDA", + "weights_revision": "bf36a94a4519e14e3c48ad005c6ff1972ab44ccb", + "schema_head": "a341c463a0090298238102506796bc73d86b84fa", + "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": "ede24ecc6a254ef33354bc0eb7bc90d1daa540bd", + "execution_provider": "CUDAExecutionProvider", + "cuda_graph": true, + "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_tokens": 68, + "image": null, + "max_new_tokens": 128, + "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 + } +} diff --git a/scripts/benchmark_muse_native.py b/scripts/benchmark_muse_native.py new file mode 100644 index 000000000..102203e90 --- /dev/null +++ b/scripts/benchmark_muse_native.py @@ -0,0 +1,188 @@ +#!/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.setdefault("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 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"] + ) + prompt_tokens = len(tokenizer.encode(prompt)) + if prompt_tokens != int(workload["prompt_tokens"]): + raise RuntimeError( + f"prompt token count mismatch: {prompt_tokens} != {workload['prompt_tokens']}" + ) + + 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..9903f0633 --- /dev/null +++ b/scripts/benchmark_muse_workflow.py @@ -0,0 +1,108 @@ +#!/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 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 main() -> int: + parser = argparse.ArgumentParser() + parser.add_argument("--model", required=True, type=Path) + parser.add_argument("--runner", 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: dict[str, Any] = json.loads(args.config.read_text()) + workload = config["workload"] + sampling = config["sampling"] + 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", + workload["rendered_prompt"], + ] + if workload["image"]: + command.extend(["--image", workload["image"]]) + completed = subprocess.run(command, check=True, text=True, capture_output=True) + 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"), + }, + "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()) From b4ab3bcc98be7c24264c5efc39fb8fffd10adcfb Mon Sep 17 00:00:00 2001 From: justinchuby Date: Thu, 13 Aug 2026 19:04:37 +0000 Subject: [PATCH 038/151] Guard metadata against tiny config leakage Construct a 52-layer BF16 VLM artifact while deliberately supplying the tiny test config. Verify workflow generation admits all 106 decoder inputs, 105 outputs, 104 KV aliases, real head dimension and vocabulary, omits position_ids, and writes the required policy artifact. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: justinchuby --- .../onnx_genai/workflow_metadata_test.py | 128 ++++++++++++++++++ 1 file changed, 128 insertions(+) diff --git a/src/mobius/integrations/onnx_genai/workflow_metadata_test.py b/src/mobius/integrations/onnx_genai/workflow_metadata_test.py index e2cb739e4..ffe36422c 100644 --- a/src/mobius/integrations/onnx_genai/workflow_metadata_test.py +++ b/src/mobius/integrations/onnx_genai/workflow_metadata_test.py @@ -7,6 +7,7 @@ import onnx_ir as ir import pytest +import yaml from mobius._model_package import ModelPackage from mobius.integrations.onnx_genai.inference_metadata_test import ( @@ -19,6 +20,7 @@ build_speculative_workflow_metadata, build_vlm_workflow_metadata, write_speculative_workflow_metadata, + write_vlm_workflow_metadata, ) @@ -101,6 +103,132 @@ def test_vlm_preprocessing_is_explicit_typed_ssa(tmp_path): assert all(output["source"] in declared for output in image["outputs"]) +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 / "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 = _model("decoder", 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. + path = write_vlm_workflow_metadata( + package, + str(tmp_path / "package"), + _VlmConfig(), + source=str(source), + ) + with open(path, encoding="utf-8") as handle: + workflow = yaml.safe_load(handle)["pipeline"]["workflow"] + decoder_invokes = [] + + def collect_decoder_invokes(node): + if isinstance(node, dict): + if node.get("kind") == "invoke" and node.get("component") == "decoder": + decoder_invokes.append(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"] == 4096 + assert workflow["state"]["cache_103"]["contract"] == { + "dtype": "bfloat16", + "rank": 4, + "shape": ["batch", 2, "past_sequence", 128], + } + assert workflow["state"]["logits"]["contract"] == { + "dtype": "float32", + "rank": 2, + "shape": ["batch", 202048], + } + kv_ports = workflow["serving"]["kv_service"]["groups"]["decoder_cache"]["ports"]["decoder"] + assert len(kv_ports) == 104 + assert kv_ports["cache_103"] == { + "input": "past_key_values.51.value", + "output": "present.51.value", + } + assert (tmp_path / "package" / "policies" / "last_token_logits.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]) From 632ab91f54943c4e7f9d2731edb5444dd08267e2 Mon Sep 17 00:00:00 2001 From: justinchuby Date: Thu, 13 Aug 2026 19:09:09 +0000 Subject: [PATCH 039/151] Derive KV storage mode from model ABI Use shared-buffer past/present aliasing for conventional decoder graphs and reserve paged storage for artifacts that expose page- or block-table inputs. Cover the real 52-layer Muse interface so metadata cannot incorrectly advertise paged KV. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: justinchuby --- .../onnx_genai/workflow_metadata.py | 51 ++++++++++++++----- .../onnx_genai/workflow_metadata_test.py | 7 ++- 2 files changed, 44 insertions(+), 14 deletions(-) diff --git a/src/mobius/integrations/onnx_genai/workflow_metadata.py b/src/mobius/integrations/onnx_genai/workflow_metadata.py index e3da5d910..5672e2f35 100644 --- a/src/mobius/integrations/onnx_genai/workflow_metadata.py +++ b/src/mobius/integrations/onnx_genai/workflow_metadata.py @@ -458,6 +458,21 @@ def _model_cache_pairs(model: ir.Model) -> list[tuple[ir.Value, ir.Value]]: return pairs +def _kv_storage_contract(model: ir.Model) -> dict[str, Any]: + """Derive physical KV storage from the admitted model interface.""" + 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") + ) + return { + "paging": "paged" if paged else "none", + "compaction": paged, + "storage": "paged" if paged else "shared_buffer", + } + + 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"] @@ -520,6 +535,8 @@ def _build_real_tts_workflow_metadata(pkg: Any, config: Any) -> dict[str, Any]: 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( @@ -1550,9 +1567,14 @@ def frame_generation_nodes(prefix: str, hidden: str, logits: str) -> list[dict[s "accepted_len": "accepted_len", "slot_ids": "slot_ids", "kv_service": { - "paging": "paged", + "paging": ( + "paged" + if talker_kv["paging"] == "paged" + or predictor_kv["paging"] == "paged" + else "none" + ), "allocation": "runtime", - "compaction": True, + "compaction": talker_kv["compaction"] or predictor_kv["compaction"], "groups": { **( { @@ -1560,7 +1582,7 @@ def frame_generation_nodes(prefix: str, hidden: str, logits: str) -> list[dict[s "sequence_axis": 2, "layout": "bnsh", "logical_lengths": "talker_cache_lengths", - "storage": "paged", + "storage": talker_kv["storage"], "ports": { "talker": { f"talker_cache_{index}": { @@ -1583,7 +1605,7 @@ def frame_generation_nodes(prefix: str, hidden: str, logits: str) -> list[dict[s "sequence_axis": 2, "layout": "bnsh", "logical_lengths": "predictor_cache_lengths", - "storage": "paged", + "storage": predictor_kv["storage"], "ports": { "code_predictor": { f"predictor_cache_{index}": { @@ -2531,6 +2553,7 @@ def build_vlm_workflow_metadata( present.shape = value.shape cache_pairs.append((value, present)) cache_names = {value.name for value, _ in cache_pairs} + decoder_kv = _kv_storage_contract(decoder) rank2_integer = [ value for value in decoder.graph.inputs @@ -3115,9 +3138,9 @@ def build_vlm_workflow_metadata( "accepted_len": "accepted_len", "slot_ids": "slot_ids", "kv_service": { - "paging": "paged", + "paging": decoder_kv["paging"], "allocation": "runtime", - "compaction": True, + "compaction": decoder_kv["compaction"], "groups": { "decoder_cache": { "sequence_axis": next( @@ -3132,7 +3155,7 @@ def build_vlm_workflow_metadata( ), "layout": "bnsh", "logical_lengths": "cache_lengths", - "storage": "paged", + "storage": decoder_kv["storage"], "ports": { "decoder": { f"cache_{index}": { @@ -3199,6 +3222,7 @@ def build_speculative_workflow_metadata( raise ValueError("speculative workflow requires proposer and verifier") proposer = pkg["proposer"] verifier = pkg["verifier"] + verifier_kv = _kv_storage_contract(verifier) proposer_input = next( ( value @@ -3922,15 +3946,15 @@ def build_speculative_workflow_metadata( "accepted_len": "accepted_len", "slot_ids": "slot_ids", "kv_service": { - "paging": "paged", + "paging": verifier_kv["paging"], "allocation": "runtime", - "compaction": True, + "compaction": verifier_kv["compaction"], "groups": { "verifier_cache": { "sequence_axis": kv_sequence_axis, "layout": "bnsh", "logical_lengths": "cache_lengths", - "storage": "paged", + "storage": verifier_kv["storage"], "ports": {"verifier": kv_ports}, } }, @@ -4002,6 +4026,7 @@ def build_decoder_workflow_metadata( decoder_name, decoder = next(iter(pkg.items())) inputs = list(decoder.graph.inputs) outputs = list(decoder.graph.outputs) + decoder_kv_contract = _kv_storage_contract(decoder) token_input = next( ( value @@ -4767,15 +4792,15 @@ def build_decoder_workflow_metadata( "accepted_len": "accepted_len", "slot_ids": "slot_ids", "kv_service": { - "paging": "paged", + "paging": decoder_kv_contract["paging"], "allocation": "runtime", - "compaction": True, + "compaction": decoder_kv_contract["compaction"], "groups": { "decoder_cache": { "sequence_axis": decoder_kv_axis, "layout": "bnsh", "logical_lengths": "cache_lengths", - "storage": "paged", + "storage": decoder_kv_contract["storage"], "ports": {decoder_name: decoder_kv_ports}, } }, diff --git a/src/mobius/integrations/onnx_genai/workflow_metadata_test.py b/src/mobius/integrations/onnx_genai/workflow_metadata_test.py index ffe36422c..8f35aa1fd 100644 --- a/src/mobius/integrations/onnx_genai/workflow_metadata_test.py +++ b/src/mobius/integrations/onnx_genai/workflow_metadata_test.py @@ -220,7 +220,12 @@ def collect_decoder_invokes(node): "rank": 2, "shape": ["batch", 202048], } - kv_ports = workflow["serving"]["kv_service"]["groups"]["decoder_cache"]["ports"]["decoder"] + kv_service = workflow["serving"]["kv_service"] + assert kv_service["paging"] == "none" + assert kv_service["compaction"] is False + decoder_cache = kv_service["groups"]["decoder_cache"] + assert decoder_cache["storage"] == "shared_buffer" + kv_ports = decoder_cache["ports"]["decoder"] assert len(kv_ports) == 104 assert kv_ports["cache_103"] == { "input": "past_key_values.51.value", From 7f017bc685a54d09c8a03bcf923dfa1a0c37c2dd Mon Sep 17 00:00:00 2001 From: justinchuby Date: Thu, 13 Aug 2026 19:24:22 +0000 Subject: [PATCH 040/151] Emit capture-stable decoder workflow state Derive fixed request capacity from prompt plus output budget for shared-buffer KV workflows. Emit invariant masks/caches with logical lengths, parse packaged ORT Extensions processors, suppress YAML aliases, and pin benchmark length semantics. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: justinchuby --- benchmarks/README.md | 4 +- benchmarks/muse_workflow_h200.json | 2 + scripts/benchmark_muse_native.py | 6 + scripts/benchmark_muse_workflow.py | 4 + src/mobius/generation/_policy_components.py | 130 +++++++++++++--- .../generation/_policy_components_test.py | 56 +++++++ .../onnx_genai/inference_metadata.py | 27 ++++ .../onnx_genai/inference_metadata_test.py | 56 +++++++ .../onnx_genai/workflow_metadata.py | 141 +++++++++++++----- .../onnx_genai/workflow_metadata_test.py | 17 +++ 10 files changed, 377 insertions(+), 66 deletions(-) diff --git a/benchmarks/README.md b/benchmarks/README.md index 44fc50b52..d14701064 100644 --- a/benchmarks/README.md +++ b/benchmarks/README.md @@ -3,7 +3,9 @@ `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 same rendered 68-token prompt, greedy parameters, token budget, -warmups, and steady-decode window. +warmups, and steady-decode window. `request_max_length` is prompt tokens plus +new tokens (68 + 128 = 196); `model_max_context` is the independent 131072-token +artifact admission ceiling. Run the native ORT GenAI path: diff --git a/benchmarks/muse_workflow_h200.json b/benchmarks/muse_workflow_h200.json index f7bab0255..f0adf243c 100644 --- a/benchmarks/muse_workflow_h200.json +++ b/benchmarks/muse_workflow_h200.json @@ -28,6 +28,8 @@ "prompt_tokens": 68, "image": null, "max_new_tokens": 128, + "request_max_length": 196, + "model_max_context": 131072, "stop_on_eos": false, "warmups": 1, "runs": 3, diff --git a/scripts/benchmark_muse_native.py b/scripts/benchmark_muse_native.py index 102203e90..fc77ab79a 100644 --- a/scripts/benchmark_muse_native.py +++ b/scripts/benchmark_muse_native.py @@ -127,6 +127,12 @@ def main() -> int: 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( diff --git a/scripts/benchmark_muse_workflow.py b/scripts/benchmark_muse_workflow.py index 9903f0633..fde40f648 100644 --- a/scripts/benchmark_muse_workflow.py +++ b/scripts/benchmark_muse_workflow.py @@ -38,6 +38,10 @@ def main() -> int: config: dict[str, Any] = json.loads(args.config.read_text()) workload = config["workload"] sampling = config["sampling"] + 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 sampling != { "algorithm": "greedy", "do_sample": False, diff --git a/src/mobius/generation/_policy_components.py b/src/mobius/generation/_policy_components.py index 7eb363a99..967d4bf53 100644 --- a/src/mobius/generation/_policy_components.py +++ b/src/mobius/generation/_policy_components.py @@ -539,8 +539,9 @@ def build_decoder_state_initializer( attention_mask_input: str, position_ids_input: str | None, cache_inputs: list[str], + fixed_capacity: bool = False, ) -> PolicyComponent: - """Build prompt-derived mask, position, token-slot, and empty-cache tensors.""" + """Build prompt-derived decoder state, optionally with capture-stable storage.""" graph, builder = _make_graph("decoder_state_initializer") op = builder.op decoder_inputs = {value.name: value for value in decoder.graph.inputs} @@ -556,25 +557,63 @@ def build_decoder_state_initializer( 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])) + 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, op.Constant(value_ints=[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_value = decoder_inputs[attention_mask_input] - attention = op.Cast( - op.ConstantOfShape(prompt_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]), + if fixed_capacity: + attention = op.Cast(op.Less(offsets, sequence_length), to=attention_value.dtype) + attention.shape = ir.Shape(["batch", "capacity"]) + body_attention = op.Cast( + op.Less( + offsets, + op.Add(sequence_length, op.Constant(value_int=1)), ), to=attention_value.dtype, - ), - axis=1, - ) - body_attention.shape = ir.Shape(["batch", "prompt_sequence + 1"]) + ) + body_attention.shape = ir.Shape(["batch", "capacity"]) + else: + attention = op.Cast( + op.ConstantOfShape(prompt_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", "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), @@ -608,6 +647,13 @@ def build_decoder_state_initializer( if position_ids_input is not None: builder.add_output(body_position, "body_position_ids") builder.add_output(token_slot, "token_slot") + if fixed_capacity: + cache_lengths = op.Expand( + op.Unsqueeze(sequence_length, op.Constant(value_ints=[0])), + batch_shape, + ) + cache_lengths.shape = ir.Shape(["batch"]) + builder.add_output(cache_lengths, "cache_lengths") for name in cache_inputs: value = decoder_inputs[name] @@ -620,7 +666,11 @@ def build_decoder_state_initializer( if axis == 0: shape_parts.append(batch_shape) elif "sequence" in dimension_text: - shape_parts.append(op.Constant(value_ints=[0])) + 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: @@ -634,7 +684,16 @@ def build_decoder_state_initializer( cache_shape, value=ir.tensor([zero], dtype=value.dtype), ) - empty.shape = value.shape + 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, {}) @@ -644,8 +703,9 @@ def build_decoder_step_update( *, attention_dtype: ir.DataType, position_dtype: ir.DataType | None, + fixed_capacity: bool = False, ) -> PolicyComponent: - """Build one-token attention-mask append and position increment.""" + """Build one-token attention-mask and position update.""" graph, builder = _make_graph("decoder_step_update") op = builder.op attention = builder.input( @@ -653,11 +713,33 @@ def build_decoder_step_update( dtype=attention_dtype, shape=["batch", "context"], ) - 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"]) + 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: position = builder.input( diff --git a/src/mobius/generation/_policy_components_test.py b/src/mobius/generation/_policy_components_test.py index fb04025e1..32c4a1c18 100644 --- a/src/mobius/generation/_policy_components_test.py +++ b/src/mobius/generation/_policy_components_test.py @@ -243,6 +243,62 @@ def test_decoder_state_initializer_and_step_update_runtime(tmp_path): 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) + 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": np.array([[3, 4, 5]], np.int64), + "max_iterations": np.array([2], np.int64), + }, + ) + attention, body_attention, token, cache_lengths, cache = outputs + np.testing.assert_array_equal(attention, [[1, 1, 1, 0, 0]]) + np.testing.assert_array_equal(body_attention, [[1, 1, 1, 1, 0]]) + np.testing.assert_array_equal(token, [[0]]) + np.testing.assert_array_equal(cache_lengths, [3]) + assert cache.shape == (1, 2, 5, 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([4], np.int64), + }, + ) + np.testing.assert_array_equal(next_attention, [[1, 1, 1, 1, 1]]) + + def test_decoder_policy_chain_generates_multiple_tokens_from_prompt_only(tmp_path): graph, builder = _make_graph("decoder_stub") op = builder.op diff --git a/src/mobius/integrations/onnx_genai/inference_metadata.py b/src/mobius/integrations/onnx_genai/inference_metadata.py index 15b15d727..432cdb200 100644 --- a/src/mobius/integrations/onnx_genai/inference_metadata.py +++ b/src/mobius/integrations/onnx_genai/inference_metadata.py @@ -587,6 +587,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): diff --git a/src/mobius/integrations/onnx_genai/inference_metadata_test.py b/src/mobius/integrations/onnx_genai/inference_metadata_test.py index 4b750b2df..e763a69a3 100644 --- a/src/mobius/integrations/onnx_genai/inference_metadata_test.py +++ b/src/mobius/integrations/onnx_genai/inference_metadata_test.py @@ -28,6 +28,7 @@ _input_source_map, _match_max_token_grid, _port, + _processor_values, add_explicit_package_io, add_policy_components_to_workflow, build_diffusion_pipeline_metadata, @@ -41,6 +42,61 @@ ) +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( [ diff --git a/src/mobius/integrations/onnx_genai/workflow_metadata.py b/src/mobius/integrations/onnx_genai/workflow_metadata.py index 5672e2f35..c2f87a782 100644 --- a/src/mobius/integrations/onnx_genai/workflow_metadata.py +++ b/src/mobius/integrations/onnx_genai/workflow_metadata.py @@ -45,6 +45,15 @@ ) +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 _contract(value: ir.Value) -> dict[str, Any]: port = _port(value) dtype = {"fp16": "float16", "bf16": "bfloat16", "fp32": "float32"}.get( @@ -434,7 +443,7 @@ def write_audio_codec_workflow_metadata(pkg: Any, output_dir: str) -> str: metadata = build_audio_codec_workflow_metadata(pkg) path = os.path.join(output_dir, "inference_metadata.yaml") with open(path, "w", encoding="utf-8") as handle: - yaml.safe_dump(metadata, handle, sort_keys=False) + _dump_yaml(metadata, handle) return path @@ -2125,7 +2134,7 @@ def write_tts_workflow_metadata(pkg: Any, output_dir: str, config: Any) -> str: pkg.save_policy_components(output_dir) path = os.path.join(output_dir, "inference_metadata.yaml") with open(path, "w", encoding="utf-8") as handle: - yaml.safe_dump(metadata, handle, sort_keys=False) + _dump_yaml(metadata, handle) return path @@ -2471,7 +2480,7 @@ def write_diffusion_workflow_metadata( pkg.save_policy_components(output_dir) path = os.path.join(output_dir, "inference_metadata.yaml") with open(path, "w", encoding="utf-8") as handle: - yaml.safe_dump(metadata, handle, sort_keys=False) + _dump_yaml(metadata, handle) return path @@ -2566,6 +2575,7 @@ def build_vlm_workflow_metadata( 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" legacy = build_native_vlm_package_metadata(pkg, config=config, source=source) preprocessing = legacy.get("preprocessing") @@ -2605,6 +2615,7 @@ def build_vlm_workflow_metadata( 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, ), ) pkg.add_policy_component( @@ -2612,6 +2623,7 @@ def build_vlm_workflow_metadata( 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, ), ) if cache_pairs: @@ -2784,12 +2796,16 @@ def build_vlm_workflow_metadata( }, "scope": "invocation", "initializer": "initializer.body_attention_mask", - "recurrence": { - "kind": "growing", - "axis": 1, - "increment": "package.one", - "max": "package.max_context", - }, + "recurrence": ( + {"kind": "invariant"} + if fixed_capacity + else { + "kind": "growing", + "axis": 1, + "increment": "package.one", + "max": "package.max_context", + } + ), }, "active": { "contract": batch_bool, @@ -2823,7 +2839,9 @@ def build_vlm_workflow_metadata( "contract": batch_int, "class": "semantic", "scope": "invocation", - "initializer": "package.zero_batch", + "initializer": ( + "initializer.cache_lengths" if fixed_capacity else "package.zero_batch" + ), "recurrence": {"kind": "invariant"}, }, } @@ -2873,7 +2891,7 @@ def build_vlm_workflow_metadata( ), ( "cache_lengths", - "package.zero_batch", + "initializer.cache_lengths" if fixed_capacity else "package.zero_batch", "state.cache_lengths.body", "cache_lengths.next", "state.cache_lengths.final", @@ -2905,18 +2923,22 @@ def build_vlm_workflow_metadata( "contract": _contract(past), "scope": "invocation", "initializer": f"decoder.setup.{present.name}", - "recurrence": { - "kind": "bounded", - "axis": next( - ( - axis - for axis, dimension in enumerate(_contract(past)["shape"]) - if "sequence" in str(dimension) + "recurrence": ( + {"kind": "invariant"} + if fixed_capacity + else { + "kind": "bounded", + "axis": next( + ( + axis + for axis, dimension in enumerate(_contract(past)["shape"]) + if "sequence" in str(dimension) + ), + 2, ), - 2, - ), - "max": "package.max_context", - }, + "max": "package.max_context", + } + ), "service_group": "decoder_cache", } setup_decoder_outputs[present.name] = f"decoder.setup.{present.name}" @@ -2964,11 +2986,19 @@ def build_vlm_workflow_metadata( *audio_setup_nodes, _invoke( "decoder_state_initializer", - {"prompt_tokens": "request.prompt_tokens"}, + { + "prompt_tokens": "request.prompt_tokens", + **({"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", + **( + {"cache_lengths": "initializer.cache_lengths"} + if fixed_capacity + else {} + ), **( { position_input.name: f"initializer.{position_input.name}", @@ -3061,6 +3091,7 @@ def build_vlm_workflow_metadata( "decoder_step_update", { "attention_mask": "state.attention_mask.body", + **({"logical_length": "cache_lengths.next"} if fixed_capacity else {}), **( {"position_ids": "state.position_ids.body"} if position_input is not None @@ -3205,7 +3236,7 @@ def write_vlm_workflow_metadata( pkg.save_policy_components(output_dir) path = os.path.join(output_dir, "inference_metadata.yaml") with open(path, "w", encoding="utf-8") as handle: - yaml.safe_dump(metadata, handle, sort_keys=False) + _dump_yaml(metadata, handle) return path @@ -4010,7 +4041,7 @@ def write_speculative_workflow_metadata( pkg.save_policy_components(output_dir) path = os.path.join(output_dir, "inference_metadata.yaml") with open(path, "w", encoding="utf-8") as handle: - yaml.safe_dump(metadata, handle, sort_keys=False) + _dump_yaml(metadata, handle) return path @@ -4131,6 +4162,7 @@ def build_decoder_workflow_metadata( ] if unsupported: raise ValueError(f"decoder workflow has unsupported non-request inputs: {unsupported}") + fixed_capacity = bool(cache_pairs) and decoder_kv_contract["storage"] == "shared_buffer" pkg.add_policy_component( "decoder_state_initializer", build_decoder_state_initializer( @@ -4139,6 +4171,7 @@ def build_decoder_workflow_metadata( 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, ), ) pkg.add_policy_component( @@ -4146,6 +4179,7 @@ def build_decoder_workflow_metadata( 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, ), ) needs_token_cast = token_input.dtype != ir.DataType.INT64 @@ -4425,7 +4459,11 @@ def build_decoder_workflow_metadata( "contract": batch_int, "class": "semantic", "scope": "invocation", - "initializer": "package.cache_lengths", + "initializer": ( + "initializer.cache_lengths" + if fixed_capacity + else "package.cache_lengths" + ), "recurrence": {"kind": "invariant"}, }, } @@ -4481,7 +4519,11 @@ def build_decoder_workflow_metadata( }, { "cell": "cache_lengths", - "current": "package.cache_lengths", + "current": ( + "initializer.cache_lengths" + if fixed_capacity + else "package.cache_lengths" + ), "body_input": "state.cache_lengths.body", "body_output": "cache_lengths.next", "next": "state.cache_lengths.final", @@ -4531,12 +4573,16 @@ def build_decoder_workflow_metadata( }, "initializer.body_attention_mask", "decoder_step.body_attention_mask", - { - "kind": "growing", - "axis": 1, - "increment": "package.one_token", - "max": "package.max_context", - }, + ( + {"kind": "invariant"} + if fixed_capacity + else { + "kind": "growing", + "axis": 1, + "increment": "package.one_token", + "max": "package.max_context", + } + ), ), } if position_input is not None: @@ -4590,11 +4636,15 @@ def build_decoder_workflow_metadata( "contract": _contract(past), "scope": "invocation", "initializer": setup_value, - "recurrence": { - "kind": "bounded", - "axis": decoder_kv_axis, - "max": "package.max_context", - }, + "recurrence": ( + {"kind": "invariant"} + if fixed_capacity + else { + "kind": "bounded", + "axis": decoder_kv_axis, + "max": "package.max_context", + } + ), "service_group": "decoder_cache", } decoder_kv_ports[cell] = {"input": past.name, "output": present.name} @@ -4617,11 +4667,19 @@ def build_decoder_workflow_metadata( "nodes": [ _invoke( "decoder_state_initializer", - {"prompt_tokens": f"request.{token_input.name}"}, + { + "prompt_tokens": f"request.{token_input.name}", + **({"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", + **( + {"cache_lengths": "initializer.cache_lengths"} + if fixed_capacity + else {} + ), **( { position_input.name: f"initializer.{position_input.name}", @@ -4736,6 +4794,7 @@ def build_decoder_workflow_metadata( "decoder_step_update", { "attention_mask": "state.attention_mask.body", + **({"logical_length": "cache_lengths.next"} if fixed_capacity else {}), **( {"position_ids": "state.position_ids.body"} if position_input is not None @@ -5107,7 +5166,7 @@ def write_decoder_workflow_metadata( pkg.save_policy_components(output_dir) path = os.path.join(output_dir, "inference_metadata.yaml") with open(path, "w", encoding="utf-8") as handle: - yaml.safe_dump(metadata, handle, sort_keys=False) + _dump_yaml(metadata, handle) return path @@ -5126,5 +5185,5 @@ def write_language_diffusion_workflow_metadata( pkg.save_policy_components(output_dir) path = os.path.join(output_dir, "inference_metadata.yaml") with open(path, "w", encoding="utf-8") as handle: - yaml.safe_dump(metadata, handle, sort_keys=False) + _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 index 8f35aa1fd..3e9c1c5ba 100644 --- a/src/mobius/integrations/onnx_genai/workflow_metadata_test.py +++ b/src/mobius/integrations/onnx_genai/workflow_metadata_test.py @@ -4,6 +4,7 @@ from __future__ import annotations import json +from pathlib import Path import onnx_ir as ir import pytest @@ -182,14 +183,20 @@ def test_vlm_writer_derives_real_decoder_contract_from_artifact(tmp_path): _VlmConfig(), 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): @@ -220,6 +227,16 @@ def collect_decoder_invokes(node): "rank": 2, "shape": ["batch", 202048], } + assert workflow["state"]["attention_mask"]["recurrence"] == {"kind": "invariant"} + assert workflow["state"]["cache_103"]["recurrence"] == {"kind": "invariant"} + assert workflow["state"]["cache_lengths"]["initializer"] == "initializer.cache_lengths" + assert policy_invokes["decoder_state_initializer"]["inputs"] == { + "prompt_tokens": "request.prompt_tokens", + "max_iterations": "request.max_iterations", + } + assert policy_invokes["decoder_step_update"]["inputs"]["logical_length"] == ( + "cache_lengths.next" + ) kv_service = workflow["serving"]["kv_service"] assert kv_service["paging"] == "none" assert kv_service["compaction"] is False From dc3f3623eb5a7b7a4eacf3da0844a23eb24f3792 Mon Sep 17 00:00:00 2001 From: justinchuby Date: Thu, 13 Aug 2026 19:29:31 +0000 Subject: [PATCH 041/151] Add executable text-only VLM branch Represent media presence explicitly, skip preprocessing and vision for text-only requests, and feed a typed empty feature matrix to the embedding model. Document the required empty media tensor binding imposed by the frozen workflow contract. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: justinchuby --- benchmarks/README.md | 2 + benchmarks/muse_workflow_h200.json | 1 + src/mobius/generation/__init__.py | 2 + src/mobius/generation/_policy_components.py | 12 +++ .../generation/_policy_components_test.py | 11 +++ .../onnx_genai/workflow_metadata.py | 87 +++++++++++++++++-- .../onnx_genai/workflow_metadata_test.py | 7 ++ 7 files changed, 116 insertions(+), 6 deletions(-) diff --git a/benchmarks/README.md b/benchmarks/README.md index d14701064..f74fcee26 100644 --- a/benchmarks/README.md +++ b/benchmarks/README.md @@ -6,6 +6,8 @@ paths use the same rendered 68-token prompt, greedy parameters, token budget, warmups, and steady-decode window. `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, bind `request.image` to an empty `uint8[0]` tensor +and `request.has_media` to `false`; the false branch supplies empty image features. Run the native ORT GenAI path: diff --git a/benchmarks/muse_workflow_h200.json b/benchmarks/muse_workflow_h200.json index f0adf243c..0629336d5 100644 --- a/benchmarks/muse_workflow_h200.json +++ b/benchmarks/muse_workflow_h200.json @@ -27,6 +27,7 @@ "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_tokens": 68, "image": null, + "workflow_media_binding": "empty_uint8_with_has_media_false", "max_new_tokens": 128, "request_max_length": 196, "model_max_context": 131072, diff --git a/src/mobius/generation/__init__.py b/src/mobius/generation/__init__.py index 57d4f0726..8d94f27b9 100644 --- a/src/mobius/generation/__init__.py +++ b/src/mobius/generation/__init__.py @@ -18,6 +18,7 @@ build_decoder_state_initializer, build_decoder_step_update, build_effectful_identity, + build_empty_features, build_eos_termination, build_euler_model_input, build_euler_solver_step, @@ -57,6 +58,7 @@ "build_decoder_state_initializer", "build_decoder_step_update", "build_effectful_identity", + "build_empty_features", "build_eos_termination", "build_euler_model_input", "build_euler_solver_step", diff --git a/src/mobius/generation/_policy_components.py b/src/mobius/generation/_policy_components.py index 967d4bf53..b69213b1d 100644 --- a/src/mobius/generation/_policy_components.py +++ b/src/mobius/generation/_policy_components.py @@ -531,6 +531,18 @@ def build_model_token_cast(dtype: ir.DataType) -> PolicyComponent: 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 build_decoder_state_initializer( decoder: ir.Model, *, diff --git a/src/mobius/generation/_policy_components_test.py b/src/mobius/generation/_policy_components_test.py index 32c4a1c18..ae6699c17 100644 --- a/src/mobius/generation/_policy_components_test.py +++ b/src/mobius/generation/_policy_components_test.py @@ -17,6 +17,7 @@ build_code_frame_update, build_decoder_state_initializer, build_decoder_step_update, + build_empty_features, build_eos_termination, build_euler_model_input, build_euler_solver_step, @@ -299,6 +300,16 @@ def test_decoder_fixed_capacity_state_matches_native_capture_layout(tmp_path): np.testing.assert_array_equal(next_attention, [[1, 1, 1, 1, 1]]) +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 diff --git a/src/mobius/integrations/onnx_genai/workflow_metadata.py b/src/mobius/integrations/onnx_genai/workflow_metadata.py index c2f87a782..0deb15122 100644 --- a/src/mobius/integrations/onnx_genai/workflow_metadata.py +++ b/src/mobius/integrations/onnx_genai/workflow_metadata.py @@ -21,6 +21,7 @@ build_codec_layout_transpose, build_decoder_state_initializer, build_decoder_step_update, + build_empty_features, build_euler_model_input, build_euler_solver_step, build_greedy_sampler, @@ -2596,6 +2597,15 @@ def build_vlm_workflow_metadata( 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, @@ -2606,6 +2616,14 @@ def build_vlm_workflow_metadata( ), ) 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( @@ -2647,6 +2665,8 @@ def build_vlm_workflow_metadata( "contract": {"dtype": "uint8", "rank": 1, "shape": ["encoded_bytes"]}, "role": {"kind": "runtime", "version": "1.0", "role": "media"}, "source": {"kind": "request", "field": "media"}, + # The frozen contract forbids optional tensors without literal defaults. + # Text-only callers bind an empty uint8 tensor and leave has_media=false. "required": True, }, "request.max_iterations": { @@ -2709,6 +2729,14 @@ def build_vlm_workflow_metadata( "default": 0, }, } + if text_only_vision is not None: + inputs["request.has_media"] = { + "contract": {"dtype": "bool", "rank": 1, "shape": [1]}, + "role": {"kind": "opaque"}, + "source": {"kind": "application", "name": "has_media"}, + "required": False, + "default": False, + } vision_invoke_inputs = { name: preprocessing_values[name] for name in vision_inputs @@ -2974,15 +3002,62 @@ def build_vlm_workflow_metadata( } ) + if text_only_vision is not None: + feature_name = text_only_vision.name + vision_setup: dict[str, Any] = { + "kind": "branch", + "predicate": "request.has_media", + "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 = { + "kind": "sequence", + "nodes": [ + _invoke( + "image_preprocess", + {"encoded": "request.image"}, + dict(preprocessing_values), + ), + _invoke("vision_encoder", vision_invoke_inputs, vision_outputs), + ], + } + setup = { "kind": "sequence", "nodes": [ - _invoke( - "image_preprocess", - {"encoded": "request.image"}, - dict(preprocessing_values), - ), - _invoke("vision_encoder", vision_invoke_inputs, vision_outputs), + vision_setup, *audio_setup_nodes, _invoke( "decoder_state_initializer", diff --git a/src/mobius/integrations/onnx_genai/workflow_metadata_test.py b/src/mobius/integrations/onnx_genai/workflow_metadata_test.py index 3e9c1c5ba..7f214898c 100644 --- a/src/mobius/integrations/onnx_genai/workflow_metadata_test.py +++ b/src/mobius/integrations/onnx_genai/workflow_metadata_test.py @@ -237,6 +237,12 @@ def collect_decoder_invokes(node): assert policy_invokes["decoder_step_update"]["inputs"]["logical_length"] == ( "cache_lengths.next" ) + assert workflow["inputs"]["request.image"]["required"] is True + assert workflow["inputs"]["request.has_media"]["default"] is False + media_branch = workflow["steps"][0]["setup"][0] + assert media_branch["kind"] == "branch" + assert set(media_branch["cases"]) == {"true", "false"} + assert media_branch["cases"]["false"]["component"] == "empty_image_features" kv_service = workflow["serving"]["kv_service"] assert kv_service["paging"] == "none" assert kv_service["compaction"] is False @@ -249,6 +255,7 @@ def collect_decoder_invokes(node): "output": "present.51.value", } 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: From afa65417970268223ba622007a235c5cc7ea6461 Mon Sep 17 00:00:00 2001 From: justinchuby Date: Thu, 13 Aug 2026 19:39:23 +0000 Subject: [PATCH 042/151] Refresh capture-stable workflow fixtures Regenerate all checked workflow packages after fixed-capacity state and alias-free YAML changes. Preserve the existing flattened vision setup for VLMs whose feature layout cannot represent an empty text-only feature matrix. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: justinchuby --- .../onnx_genai/workflow_metadata.py | 93 +++++++------ .../decoder/inference_metadata.yaml | 86 ++++++++---- .../policies/decoder_state_initializer.onnx | Bin 6641 -> 9625 bytes .../decoder/policies/decoder_step_update.onnx | Bin 2199 -> 3203 bytes .../masked/inference_metadata.yaml | 51 ++++++-- .../speculative/inference_metadata.yaml | 123 ++++++++++++++---- .../tts/inference_metadata.yaml | 110 ++++++++++++---- .../vlm/inference_metadata.yaml | 108 ++++++++++----- .../policies/decoder_state_initializer.onnx | Bin 7747 -> 10907 bytes .../vlm/policies/decoder_step_update.onnx | Bin 2199 -> 3203 bytes 10 files changed, 410 insertions(+), 161 deletions(-) diff --git a/src/mobius/integrations/onnx_genai/workflow_metadata.py b/src/mobius/integrations/onnx_genai/workflow_metadata.py index 0deb15122..53d2a0893 100644 --- a/src/mobius/integrations/onnx_genai/workflow_metadata.py +++ b/src/mobius/integrations/onnx_genai/workflow_metadata.py @@ -3004,60 +3004,59 @@ def build_vlm_workflow_metadata( if text_only_vision is not None: feature_name = text_only_vision.name - vision_setup: dict[str, Any] = { - "kind": "branch", - "predicate": "request.has_media", - "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}", - }, - ), - ], + vision_setup_nodes: list[dict[str, Any]] = [ + { + "kind": "branch", + "predicate": "request.has_media", + "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}"}, + ), }, - "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}", + "outputs": { + f"vision.{feature_name}": { + "cases": { + "true": f"vision.with_media.{feature_name}", + "false": f"vision.empty.{feature_name}", + } } - } - }, - } + }, + } + ] else: - vision_setup = { - "kind": "sequence", - "nodes": [ - _invoke( - "image_preprocess", - {"encoded": "request.image"}, - dict(preprocessing_values), - ), - _invoke("vision_encoder", vision_invoke_inputs, vision_outputs), - ], - } + 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, + *vision_setup_nodes, *audio_setup_nodes, _invoke( "decoder_state_initializer", diff --git a/tests/fixtures/onnx_genai_workflows/decoder/inference_metadata.yaml b/tests/fixtures/onnx_genai_workflows/decoder/inference_metadata.yaml index 7f620e841..76563a413 100644 --- a/tests/fixtures/onnx_genai_workflows/decoder/inference_metadata.yaml +++ b/tests/fixtures/onnx_genai_workflows/decoder/inference_metadata.yaml @@ -29,7 +29,7 @@ pipeline: kind: request required: true request.max_iterations: - contract: &id001 + contract: dtype: int64 rank: 1 shape: @@ -54,7 +54,11 @@ pipeline: required: true default: 127 package.one_token: - contract: *id001 + contract: + dtype: int64 + rank: 1 + shape: + - 1 role: kind: opaque source: @@ -62,7 +66,11 @@ pipeline: required: false default: 1 package.max_context: - contract: *id001 + contract: + dtype: int64 + rank: 1 + shape: + - 1 role: kind: opaque source: @@ -70,7 +78,7 @@ pipeline: required: false default: 8192 package.active: - contract: &id002 + contract: dtype: bool rank: 1 shape: @@ -82,7 +90,11 @@ pipeline: required: false default: true package.not_done: - contract: *id002 + contract: + dtype: bool + rank: 1 + shape: + - batch role: kind: opaque source: @@ -90,7 +102,7 @@ pipeline: required: false default: false package.slot_ids: - contract: &id003 + contract: dtype: int64 rank: 1 shape: @@ -102,7 +114,11 @@ pipeline: required: false default: 0 package.cache_lengths: - contract: *id003 + contract: + dtype: int64 + rank: 1 + shape: + - batch role: kind: opaque source: @@ -110,7 +126,11 @@ pipeline: required: false default: 0 package.zero_batch: - contract: *id003 + contract: + dtype: int64 + rank: 1 + shape: + - batch role: kind: opaque source: @@ -210,38 +230,58 @@ pipeline: recurrence: kind: invariant active: - contract: *id002 + contract: + dtype: bool + rank: 1 + shape: + - batch class: semantic scope: invocation initializer: package.active recurrence: kind: invariant done: - contract: *id002 + contract: + dtype: bool + rank: 1 + shape: + - batch class: semantic scope: invocation initializer: package.not_done recurrence: kind: invariant accepted_len: - contract: *id003 + contract: + dtype: int64 + rank: 1 + shape: + - batch class: semantic scope: invocation initializer: package.zero_batch recurrence: kind: invariant slot_ids: - contract: *id003 + contract: + dtype: int64 + rank: 1 + shape: + - batch class: semantic scope: invocation initializer: package.slot_ids recurrence: kind: invariant cache_lengths: - contract: *id003 + contract: + dtype: int64 + rank: 1 + shape: + - batch class: semantic scope: invocation - initializer: package.cache_lengths + initializer: initializer.cache_lengths recurrence: kind: invariant attention_mask: @@ -254,10 +294,7 @@ pipeline: scope: invocation initializer: initializer.body_attention_mask recurrence: - kind: growing - axis: 1 - increment: package.one_token - max: package.max_context + kind: invariant position_ids: contract: dtype: int64 @@ -281,9 +318,7 @@ pipeline: scope: invocation initializer: decoder.setup.present.0.key recurrence: - kind: bounded - axis: 2 - max: package.max_context + kind: invariant service_group: decoder_cache serving: active: active @@ -291,15 +326,15 @@ pipeline: accepted_len: accepted_len slot_ids: slot_ids kv_service: - paging: paged + paging: none allocation: runtime - compaction: true + compaction: false groups: decoder_cache: sequence_axis: 2 layout: bnsh logical_lengths: cache_lengths - storage: paged + storage: shared_buffer ports: model: cache_9: @@ -312,10 +347,12 @@ pipeline: component: decoder_state_initializer inputs: prompt_tokens: request.input_ids + max_iterations: request.max_iterations outputs: attention_mask: initializer.attention_mask body_attention_mask: initializer.body_attention_mask token_slot: initializer.token_slot + cache_lengths: initializer.cache_lengths position_ids: initializer.position_ids body_position_ids: initializer.body_position_ids past_key_values.0.key: initializer.past_key_values.0.key @@ -397,6 +434,7 @@ pipeline: component: decoder_step_update inputs: attention_mask: attention_mask + logical_length: cache_lengths.next position_ids: position_ids outputs: next_attention_mask: decoder_step.body_attention_mask diff --git a/tests/fixtures/onnx_genai_workflows/decoder/policies/decoder_state_initializer.onnx b/tests/fixtures/onnx_genai_workflows/decoder/policies/decoder_state_initializer.onnx index 3a042bd0b7d887ff7f856e266bb7a5bcf4b5fce7..b85fa5c58d6cb0ed70abd2aec96ad22f27d9a04c 100644 GIT binary patch literal 9625 zcmc&)TW{P%6ppu>>}IkF$)L8WK~-5*1r}*#&&=L<>P0*dkdQzXQmbNly>8-UV=rqb zN{Zm24}Adyh?@W*6(RltZ@lo!IKIz}yS*UU$xCKDo|(_*oH^$^XSTV9R}Q_gbL7AM z=somELF==*cQ^}-z&o&AAFm!8AMKel+tAcj+;F`K^xDf&Z)Y8?y5^zn&rHk4OFLK5 z<=Me**K^$y-<5WiP&>)5#ExgZTUFsS-c z){eAD+s*VrXA7ZYb9!VOjvM&bL2!+-Q2nH!%&G_y0OB<;w@;63`+;pRbt7#?U|GA7 z^veP!iK{1zt9K^i`b|L_hvtdl1aO1`2b3wutjEi&FUhQ5y3D>*5U(@7)5nctOFNUJD#z3(|Z-1*k2$Zh&V#8n;cGE3lHr(Co|KX7&^p* zq#0|q-(HFPWvTVsC>&(l^WlJnT=IPz$e=z@TTvqhHE7qvURjX7iyU0CG2SqPz;?qh zY8;yWf!exd`az<+VY?poI;*G_c-xHFCH z(bTmpo)fLqfc~{)PZ@;FmVp&Y#+!1fal7C8mNDy0{HM`o4)APl1@`!)Qj7pNCg#2?(}5@6D6 zXUSAD&{5HHL_0bgO=YvO(_Bja9xL^-oSWoGNWY-={vt8TLk|xlR=+l^1 z@M%n|z^BPcGShOOrsH+XwDxRcYP-9^o-cYerdO<2R;U^=zG@`u6)q1-7?m`;}CF|KU?J*2m{(VZr9WJ~g6T%kOGk=q$+6*r1rDT}ZPu z1X-H0QxB80@-j4hX<}P&V>&nB4lW39*MPs6Q|Ey_zjGvCwocHv*}$KmH}JKYH+8Jj z?!w94u^CuO!D%{FZVeGtE!d^jjkkgVhrMGxWAE#5cvP+okk+|_tSml+_Bj?u(kXdtxFL6{A zD*pLU`1yQ@)b(sLX@Fz{7x)_w+(awzg*CLaA0fy>h7MIR4#8zHfapEkg5d@%3N{Lp zBZS#KKsI#=h9b?S_4shjRd7A=fB=1zA;c)J`=;W|SAjw-zlsPOs*F5rUT|TZFZ&`{ za3oipwBUf7rR-l2%OM$LQQA-#NpW-1rrJVjLunw#rtC}7rsOHohBP^ynn^oJQL_Qi zk)&Qo%F4V&Bl0gz|ncK`qY literal 6641 zcmd5>PjBNy6pusNxUXqxXGOO)(yAM+7RX8=o}@|2Y1tJAdO>1^kXS92^x%4Z995UX(ZL$5xaO(tk4(!Zg@Xrh zZ*+20_gwefx5myWs9QtR_l<#LkIk_)xFE&nFAjD=lfGhm7{29=Y@dArSM!1B0uX%) zqz3Ss?e@2oE2UVgz>*)B<6yghqg#MT1&CCFSQH>Kz8|4F4DMY_3rIZ)^Ulf>ynXH>ZV7r0ixyGsK zpJ=O3O+PR~;H_$j4QeY;47}}fj@IbU0dfJNY^v$<(YZ@P=LomAk5WpG;zJMD z>0;NJ#&ZOoTR?v$slAy|dy`Z928J@0aah#MH6qE@*vi0aaj?W-j&$V~%zx&rk%C~2 zEtY(qeKE1^m$so}@`;}#b?;Tn(-(Oq8d;U?3}{(hFJ#|`aXArBDkIN# z!ZOwA`&uQfQmKApWx;oR3}M||h=O8=!oi_821l}cn#ll^&j6y`=UF$;PyDpWC4Slx z6F+U;XyV6)DAgcW1JD-P3@bdr;6&@+%iH>8)n+mRC*=E&FIT*T0Y3_#oBup2Yqqh0w>1GPfH94jL+7PE{%j)O z(Dk9GmdWSvHhDPmhK_YnpS!r;GXrY?lr$E8IlTq+ZOrzrW%nj>cS$TTQSOUt$koZa z3*`o0U=EE$>?sII6vaJ9p`Ttu=p*vZLZOMcuVGfFOa^d>QhV2s+9s`qQYDIaAMNVB z4SR+tW)#JN- zfd`->86})3@)J5qClr2y3KREd-zdgw3s5-LisPgoA6Q$>kiid`S4 zQEAguV{B7tQamEB8B8(hrlm{eEw~^3sZu|(UHs>RH33u0H!RN$#-;h!o`N>JF?{ueiGHU|Iz diff --git a/tests/fixtures/onnx_genai_workflows/decoder/policies/decoder_step_update.onnx b/tests/fixtures/onnx_genai_workflows/decoder/policies/decoder_step_update.onnx index 390f62573f15a272042ba9b3688c9fcdea316696..0333f2f9bf0c4417d7b0e6fd4e2772628899dcb9 100644 GIT binary patch literal 3203 zcmc(h-EPw`6vvxxT~n`v;%Tfz0*$DGsX|rLu3d`|Vw(#@;sc|x4Y)Zz5k;KK1eV>Mq9S! zQfyIFRjx*D%QgtAs&d#-E`Zz^4@Lx;7A>YurF3TjpmC} z@|K{QawTkX+?tvV?4&{$)3YAevwohQNyh3_0>U8!lOYmZbx-K$*y<6qrsVwwkKCG? z;Z{;{`^@DG{w&D3Fh(yeHxRU;6n+^TtfA(jo}}QuvV0iWJyXX6G$2-w_9qB6C=fQM z2%Gh{1{c?kdcHoIG3`c(Wa?QU_Rd{rR;GTZB{w-3Ri37+f1 zJE02az$}}U+FhP_h6xjvgjYs1UR`IoOW+mt%KH%n?ptNgKrFSh@>t`_KGJk67&e z?Oi+-{~#mp&fgc5cQ^M>rT_sY^4%>K&e%gDj2sjm?)yCd5&WCMzH&rUc zrQ#RBfeVLn;L1PXzp-A&ZQKM}4)$qwJUjdL&AeHC87~a{mOJFnJ{RE!f>kO6^Mvbr zWI%ZjFP)ON9U3t5*eJw?=R0y`&Z|`gV2*iC30&`^{#xoWP!l1B%O2h?U* zYi_~)pm$>Vo;Tul=mx^F`;>Fiaal-1yK{zf?>{ulU|freau9C&0pn2wd_}l$N)CNb z9_0yAHSS#63R%_E6Gx>8a7z0_MqE$ueTi`~D^wQTRkke7C2eX8M-@zN1 z=IB;wcJTJiK+6xy)z%?Qm*{^HckfQjd<>vVoOzH}=^Cltt)~8!Qd9V9SFXD`bG}_} z)<7F4P8z?K`nQofw2l`bF9B)Y*q`iaB>r>s05;>N-#THQe1s!$SuuYkDr!h=fpZ@} voudQTx-pWt4!$>z8W({MM{eJx;n{xm9Bo7SuQ^z|&1|9G_ZNB4)=GZ>i9y5+ diff --git a/tests/fixtures/onnx_genai_workflows/masked/inference_metadata.yaml b/tests/fixtures/onnx_genai_workflows/masked/inference_metadata.yaml index b1ca6edae..6a1480251 100644 --- a/tests/fixtures/onnx_genai_workflows/masked/inference_metadata.yaml +++ b/tests/fixtures/onnx_genai_workflows/masked/inference_metadata.yaml @@ -13,10 +13,10 @@ pipeline: - loop_induction_values inputs: request.input_ids: - contract: &id003 + contract: dtype: int64 rank: 2 - shape: &id001 + shape: - batch - sequence role: @@ -27,10 +27,12 @@ pipeline: kind: request required: true request.mask: - contract: &id004 + contract: dtype: bool rank: 2 - shape: *id001 + shape: + - batch + - sequence role: kind: opaque source: @@ -38,7 +40,7 @@ pipeline: name: masked_positions required: true request.seed: - contract: &id002 + contract: dtype: int64 rank: 1 shape: @@ -52,7 +54,11 @@ pipeline: required: false default: 0 request.rng_offset: - contract: *id002 + contract: + dtype: int64 + rank: 1 + shape: + - batch role: kind: opaque source: @@ -75,7 +81,11 @@ pipeline: required: false default: 8 package.num_steps: - contract: *id002 + contract: + dtype: int64 + rank: 1 + shape: + - batch role: kind: opaque source: @@ -96,7 +106,12 @@ pipeline: default: true outputs: tokens: - contract: *id003 + contract: + dtype: int64 + rank: 2 + shape: + - batch + - sequence role: tokens stage: pre_adapter components: @@ -124,19 +139,33 @@ pipeline: next_offset: next_offset state: tokens_state: - contract: *id003 + contract: + dtype: int64 + rank: 2 + shape: + - batch + - sequence scope: invocation initializer: request.input_ids recurrence: kind: invariant mask: - contract: *id004 + contract: + dtype: bool + rank: 2 + shape: + - batch + - sequence scope: invocation initializer: request.mask recurrence: kind: invariant rng_offset: - contract: *id002 + contract: + dtype: int64 + rank: 1 + shape: + - batch scope: invocation initializer: request.rng_offset recurrence: diff --git a/tests/fixtures/onnx_genai_workflows/speculative/inference_metadata.yaml b/tests/fixtures/onnx_genai_workflows/speculative/inference_metadata.yaml index cfff66939..5e34fd4a2 100644 --- a/tests/fixtures/onnx_genai_workflows/speculative/inference_metadata.yaml +++ b/tests/fixtures/onnx_genai_workflows/speculative/inference_metadata.yaml @@ -35,7 +35,7 @@ pipeline: kind: request required: true request.seed: - contract: &id001 + contract: dtype: int64 rank: 1 shape: @@ -49,7 +49,7 @@ pipeline: required: false default: 0 request.max_iterations: - contract: &id002 + contract: dtype: int64 rank: 1 shape: @@ -62,7 +62,11 @@ pipeline: kind: request required: true package.zero: - contract: *id001 + contract: + dtype: int64 + rank: 1 + shape: + - batch role: kind: opaque source: @@ -70,7 +74,11 @@ pipeline: required: false default: 0 package.one: - contract: *id001 + contract: + dtype: int64 + rank: 1 + shape: + - batch role: kind: opaque source: @@ -78,7 +86,11 @@ pipeline: required: false default: 1 package.max_context: - contract: *id002 + contract: + dtype: int64 + rank: 1 + shape: + - 1 role: kind: opaque source: @@ -86,7 +98,7 @@ pipeline: required: false default: 4096 package.false: - contract: &id003 + contract: dtype: bool rank: 1 shape: @@ -98,7 +110,11 @@ pipeline: required: false default: false package.active: - contract: *id003 + contract: + dtype: bool + rank: 1 + shape: + - batch role: kind: opaque source: @@ -106,7 +122,11 @@ pipeline: required: false default: true request.slot_ids: - contract: *id001 + contract: + dtype: int64 + rank: 1 + shape: + - batch role: kind: opaque source: @@ -114,7 +134,11 @@ pipeline: name: serving.slot_ids required: true request.cache_lengths: - contract: *id001 + contract: + dtype: int64 + rank: 1 + shape: + - batch role: kind: opaque source: @@ -123,7 +147,11 @@ pipeline: required: false default: 0 request.grammar_state: - contract: *id001 + contract: + dtype: int64 + rank: 1 + shape: + - batch role: kind: opaque source: @@ -144,7 +172,11 @@ pipeline: name: grammar.transition_table required: true request.adaptive_k: - contract: *id001 + contract: + dtype: int64 + rank: 1 + shape: + - batch role: kind: opaque source: @@ -153,7 +185,7 @@ pipeline: required: false default: 1 request.adaptive_estimates: - contract: &id004 + contract: dtype: float32 rank: 2 shape: @@ -515,63 +547,100 @@ pipeline: recurrence: kind: invariant rng_offset: - contract: *id001 + contract: + dtype: int64 + rank: 1 + shape: + - batch class: semantic scope: invocation initializer: package.zero recurrence: kind: invariant active: - contract: *id003 + contract: + dtype: bool + rank: 1 + shape: + - batch class: semantic scope: invocation initializer: package.active recurrence: kind: invariant done: - contract: *id003 + contract: + dtype: bool + rank: 1 + shape: + - batch class: semantic scope: invocation initializer: package.false recurrence: kind: invariant accepted_len: - contract: *id001 + contract: + dtype: int64 + rank: 1 + shape: + - batch class: semantic scope: invocation initializer: package.zero recurrence: kind: invariant slot_ids: - contract: *id001 + contract: + dtype: int64 + rank: 1 + shape: + - batch class: semantic scope: invocation initializer: request.slot_ids recurrence: kind: invariant cache_lengths: - contract: *id001 + contract: + dtype: int64 + rank: 1 + shape: + - batch class: semantic scope: invocation initializer: request.cache_lengths recurrence: kind: invariant grammar: - contract: *id001 + contract: + dtype: int64 + rank: 1 + shape: + - batch class: semantic scope: invocation initializer: request.grammar_state recurrence: kind: invariant proposal_k: - contract: *id001 + contract: + dtype: int64 + rank: 1 + shape: + - batch class: advisory scope: invocation initializer: request.adaptive_k recurrence: kind: invariant adaptive_estimates: - contract: *id004 + contract: + dtype: float32 + rank: 2 + shape: + - batch + - 24 class: advisory scope: invocation initializer: request.adaptive_estimates @@ -600,15 +669,15 @@ pipeline: accepted_len: accepted_len slot_ids: slot_ids kv_service: - paging: paged + paging: none allocation: runtime - compaction: true + compaction: false groups: verifier_cache: sequence_axis: 2 layout: bnsh logical_lengths: cache_lengths - storage: paged + storage: shared_buffer ports: verifier: cache_0: @@ -775,4 +844,8 @@ pipeline: next: verifier.present.0.key iteration: value: speculative.iteration - contract: *id001 + contract: + dtype: int64 + rank: 1 + shape: + - batch diff --git a/tests/fixtures/onnx_genai_workflows/tts/inference_metadata.yaml b/tests/fixtures/onnx_genai_workflows/tts/inference_metadata.yaml index cb0219254..7662cd563 100644 --- a/tests/fixtures/onnx_genai_workflows/tts/inference_metadata.yaml +++ b/tests/fixtures/onnx_genai_workflows/tts/inference_metadata.yaml @@ -29,7 +29,7 @@ pipeline: kind: request required: true request.max_iterations: - contract: &id001 + contract: dtype: int64 rank: 1 shape: @@ -42,7 +42,7 @@ pipeline: kind: request required: true package.false: - contract: &id003 + contract: dtype: bool rank: 1 shape: @@ -76,7 +76,11 @@ pipeline: required: false default: 1 package.remaining_groups: - contract: *id001 + contract: + dtype: int64 + rank: 1 + shape: + - 1 role: kind: opaque source: @@ -84,7 +88,11 @@ pipeline: required: false default: 2 package.predictor_context_limit: - contract: *id001 + contract: + dtype: int64 + rank: 1 + shape: + - 1 role: kind: opaque source: @@ -92,7 +100,11 @@ pipeline: required: false default: 4 package.predictor_mask_limit: - contract: *id001 + contract: + dtype: int64 + rank: 1 + shape: + - 1 role: kind: opaque source: @@ -100,7 +112,11 @@ pipeline: required: false default: 5 package.talker_context_limit: - contract: *id001 + contract: + dtype: int64 + rank: 1 + shape: + - 1 role: kind: opaque source: @@ -108,7 +124,11 @@ pipeline: required: false default: 128 package.one_control: - contract: *id001 + contract: + dtype: int64 + rank: 1 + shape: + - 1 role: kind: opaque source: @@ -116,7 +136,7 @@ pipeline: required: false default: 1 package.zero_batch: - contract: &id002 + contract: dtype: int64 rank: 1 shape: @@ -128,7 +148,11 @@ pipeline: required: false default: 0 package.one_batch: - contract: *id002 + contract: + dtype: int64 + rank: 1 + shape: + - batch role: kind: opaque source: @@ -136,7 +160,11 @@ pipeline: required: false default: 1 package.true: - contract: *id003 + contract: + dtype: bool + rank: 1 + shape: + - batch role: kind: opaque source: @@ -144,7 +172,11 @@ pipeline: required: false default: true package.slot_ids: - contract: *id002 + contract: + dtype: int64 + rank: 1 + shape: + - batch role: kind: opaque source: @@ -427,7 +459,11 @@ pipeline: recurrence: kind: invariant code_token: - contract: *id002 + contract: + dtype: int64 + rank: 1 + shape: + - batch scope: invocation initializer: frame.group1 recurrence: @@ -458,42 +494,66 @@ pipeline: recurrence: kind: invariant active: - contract: *id003 + contract: + dtype: bool + rank: 1 + shape: + - batch class: semantic scope: invocation initializer: package.true recurrence: kind: invariant done: - contract: *id003 + contract: + dtype: bool + rank: 1 + shape: + - batch class: semantic scope: invocation initializer: package.false recurrence: kind: invariant accepted_len: - contract: *id002 + contract: + dtype: int64 + rank: 1 + shape: + - batch class: semantic scope: invocation initializer: package.zero_batch recurrence: kind: invariant slot_ids: - contract: *id002 + contract: + dtype: int64 + rank: 1 + shape: + - batch class: semantic scope: invocation initializer: package.slot_ids recurrence: kind: invariant talker_cache_lengths: - contract: *id002 + contract: + dtype: int64 + rank: 1 + shape: + - batch class: semantic scope: invocation initializer: package.zero_batch recurrence: kind: invariant predictor_cache_lengths: - contract: *id002 + contract: + dtype: int64 + rank: 1 + shape: + - batch class: semantic scope: invocation initializer: package.zero_batch @@ -729,15 +789,15 @@ pipeline: accepted_len: accepted_len slot_ids: slot_ids kv_service: - paging: paged + paging: none allocation: runtime - compaction: true + compaction: false groups: talker_cache: sequence_axis: 2 layout: bnsh logical_lengths: talker_cache_lengths - storage: paged + storage: shared_buffer ports: talker: talker_cache_0: @@ -750,7 +810,7 @@ pipeline: sequence_axis: 2 layout: bnsh logical_lengths: predictor_cache_lengths - storage: paged + storage: shared_buffer ports: code_predictor: predictor_cache_0: @@ -1436,7 +1496,11 @@ pipeline: next: talker.continue iteration: value: talker.iteration - contract: *id002 + contract: + dtype: int64 + rank: 1 + shape: + - batch - kind: invoke component: codec_layout inputs: diff --git a/tests/fixtures/onnx_genai_workflows/vlm/inference_metadata.yaml b/tests/fixtures/onnx_genai_workflows/vlm/inference_metadata.yaml index 9f2c8e220..ee79a8a27 100644 --- a/tests/fixtures/onnx_genai_workflows/vlm/inference_metadata.yaml +++ b/tests/fixtures/onnx_genai_workflows/vlm/inference_metadata.yaml @@ -54,7 +54,7 @@ preprocessing: content: pixels dtype: float32 source: image.transform_4 - contract: &id004 + contract: dtype: float32 rank: 2 shape: @@ -64,7 +64,7 @@ preprocessing: content: grid_dimensions dtype: int64 source: image.output_grid_dimensions - contract: &id005 + contract: dtype: int64 rank: 2 shape: @@ -115,7 +115,7 @@ pipeline: kind: request required: true request.max_iterations: - contract: &id001 + contract: dtype: int64 rank: 1 shape: @@ -140,7 +140,11 @@ pipeline: required: false default: 2 package.max_context: - contract: *id001 + contract: + dtype: int64 + rank: 1 + shape: + - 1 role: kind: opaque source: @@ -148,7 +152,11 @@ pipeline: required: false default: 8192 package.one: - contract: *id001 + contract: + dtype: int64 + rank: 1 + shape: + - 1 role: kind: opaque source: @@ -156,7 +164,7 @@ pipeline: required: false default: 1 package.active: - contract: &id002 + contract: dtype: bool rank: 1 shape: @@ -168,7 +176,11 @@ pipeline: required: false default: true package.false: - contract: *id002 + contract: + dtype: bool + rank: 1 + shape: + - batch role: kind: opaque source: @@ -176,7 +188,7 @@ pipeline: required: false default: false package.zero_batch: - contract: &id003 + contract: dtype: int64 rank: 1 shape: @@ -188,7 +200,11 @@ pipeline: required: false default: 0 package.slot_ids: - contract: *id003 + contract: + dtype: int64 + rank: 1 + shape: + - batch role: kind: opaque source: @@ -243,8 +259,18 @@ pipeline: shape: - encoded_bytes outputs: - pixel_values: *id004 - grid_thw: *id005 + pixel_values: + dtype: float32 + rank: 2 + shape: + - 4 + - 1176 + grid_thw: + dtype: int64 + rank: 2 + shape: + - 1 + - 3 token_sampler: implementation: kind: onnx @@ -332,43 +358,60 @@ pipeline: scope: invocation initializer: initializer.body_attention_mask recurrence: - kind: growing - axis: 1 - increment: package.one - max: package.max_context + kind: invariant active: - contract: *id002 + contract: + dtype: bool + rank: 1 + shape: + - batch class: semantic scope: invocation initializer: package.active recurrence: kind: invariant done: - contract: *id002 + contract: + dtype: bool + rank: 1 + shape: + - batch class: semantic scope: invocation initializer: package.false recurrence: kind: invariant accepted_len: - contract: *id003 + contract: + dtype: int64 + rank: 1 + shape: + - batch class: semantic scope: invocation initializer: package.zero_batch recurrence: kind: invariant slot_ids: - contract: *id003 + contract: + dtype: int64 + rank: 1 + shape: + - batch class: semantic scope: invocation initializer: package.slot_ids recurrence: kind: invariant cache_lengths: - contract: *id003 + contract: + dtype: int64 + rank: 1 + shape: + - batch class: semantic scope: invocation - initializer: package.zero_batch + initializer: initializer.cache_lengths recurrence: kind: invariant position_ids: @@ -394,9 +437,7 @@ pipeline: scope: invocation initializer: decoder.setup.present.0.key recurrence: - kind: bounded - axis: 2 - max: package.max_context + kind: invariant service_group: decoder_cache cache_1: contract: @@ -410,9 +451,7 @@ pipeline: scope: invocation initializer: decoder.setup.present.0.value recurrence: - kind: bounded - axis: 2 - max: package.max_context + kind: invariant service_group: decoder_cache loop_0_active: contract: @@ -430,15 +469,15 @@ pipeline: accepted_len: accepted_len slot_ids: slot_ids kv_service: - paging: paged + paging: none allocation: runtime - compaction: true + compaction: false groups: decoder_cache: sequence_axis: 2 layout: bnsh logical_lengths: cache_lengths - storage: paged + storage: shared_buffer ports: decoder: cache_0: @@ -468,10 +507,12 @@ pipeline: component: decoder_state_initializer inputs: prompt_tokens: request.prompt_tokens + max_iterations: request.max_iterations outputs: attention_mask: initializer.attention_mask body_attention_mask: initializer.body_attention_mask token_slot: initializer.token_slot + cache_lengths: initializer.cache_lengths position_ids: initializer.position_ids body_position_ids: initializer.body_position_ids past_key_values.0.key: initializer.past_key_values.0.key @@ -572,6 +613,7 @@ pipeline: component: decoder_step_update inputs: attention_mask: attention_mask + logical_length: cache_lengths.next position_ids: position_ids outputs: next_attention_mask: decoder_step.body_attention_mask @@ -605,4 +647,8 @@ pipeline: next: loop.continue iteration: value: loop.iteration - contract: *id003 + contract: + dtype: int64 + rank: 1 + shape: + - batch diff --git a/tests/fixtures/onnx_genai_workflows/vlm/policies/decoder_state_initializer.onnx b/tests/fixtures/onnx_genai_workflows/vlm/policies/decoder_state_initializer.onnx index 4565c7336226903f9bd4b082c6dec88bd6d001b5..0f5286ad1769a18d5ea85794c2543660c2113cec 100644 GIT binary patch literal 10907 zcmc&)TW{P%6ppu>WHY&t3~HMiRFzd#V3Ah#%y@UhQ!nCy2nh*PRcci%uh&hSi+5Q& zQBnj#eLzTn0C5u_q$0#$;0+=1%Q(KzjO(;TvNtdB?0CkX&-FWJ+*-y9hvT7p+FnIVo^|N>6U%n+-1b#; zd9uIL9edu1Z%^Gx(6tXN-#2$%XKGFD-BaAScYpf|Ql-1XJxt#oPaHqI13fMWq6*5WiFK%CUKC?141Az@+L6 z*>$8v+HR#kbT$w=whoRQ)Aa)XItXsDDpWrzD6=Yp1b}!M%?ueqw?E-`g3+=0p5^U0rcTw>_$n4BZO=!;s?6>ONujJD zImSt@zT^a=<76U=k>NGoUv)#>$lNoIZn{x}$NmfnLBtV8-NZR{Ep*zqP9~N&GIhuY z$z`n7UV9-PSEbf(qI8fwo=*oX6%sQ!oI&W9!r2yQ{X#gaGH!e=C37KIY~fa% z&WWZjXZf6Hl_vDBBzwvrWX43G9>|S}=+zoi%w41B5gorc9FI=Ta*HCQHj9$tDWS4O zk#e#RCRG{PFVKt3&I#q{1p&67M6f6*5LPR__$xBA1OQ)0L#l~+=-DhoM*%p@oLDRN2D3rm<(He?_$B(jwgN~#Hj zl0wI*PC`iw6$0feg!)X5CzMnZ#6@+HP*PGLl$6v!D1VX&B_+9+@N|+=kxo)tHJt?R zma|z(@ycZEyAgSFM?TNpD61{ge#ys>kcLDIk{SXWqc%8FsHKGrK=}+H+Wl-2Kkade zpZ3MXPy4kderV0)8sy3-+Q%Caxf2XdD5mJ3YU`I(TgV8M&j_Mwt8C!@eny!;uu+4Y3L_bdHW0M09q8)ava2@Xl@!e6i^4k=K{FyFq8|?)cKkvEc)WsR>8!L3o%B z<%)ISK6IvckK}JFMrhbt#2=v7@U_YKz_m}irHi{mE3kJF!P-;rOupan9XM9>Z0GLF z^6@c7s0Gh$qDD&@j*-+6$M?}an7#8HsNTS@KNXd2!51QKa5_9hLJ?Qd;V0HnL;i>?iTT17Vk!LJpW8XhJWKLwt!Vj^T_kp9)`}C*``HJxrP) z*~A6@j1{-g0=!`Z&Fw`9GLou8AD%;SS~pJ>y$fOk zHiIlm8(Og`Zcf@%E0s1hBXexZo+NEb&X6`V=F+8^w3DV-HUT=4)C+B(bX7RX%L%$K z`BXz*CmJTCH{NO)Z=3QXndpV0n4y}?AFAxh`|_$Q>?^t>|0+jkCFFkG)`O*#T~R^zkF*2D&2QsG6pusNxUXqxXGK4ZwCYBy1+r2UkMju}mR)h67bI2)iPd5`j@!6R>=3(E zy9ljR1Sc+tMF>tSAugOaaRMR!GiL1Ao*BoAisY1>JRZ-y-~8U^yvXZlb?o--GwPVI?bH%Er&=|kI^8dGz4iHc929v*-qd`0)rJ=2|79{mDdW&=?m zR9Syyc)lJ0Z&pimP+Nnd@9vdze@TQ5YKZ}oc~;rUu*%N= z!YW_PwaN~X$)Ttnc3FGuF3lc)lR)qbh!LS%eRpuF3vR>NKrdwjc?PB*B(%XekJ?bS zB4-IrQf=Ib`~1SW4-XNo$5yMV$|mbTs+wXi&&>0O1S}Q=L9W)lk?SkF5sy)b4?JAO zi(O|H&k69{0{ROnYpG zD5R}dU5j7kMct1i(eXHvlBYp~^!pbR!x`wfp={0`18}%eT@A;1A^R?f%dvP=nYf-E zl&SW>Qz}W7iq)l+1>dqUgx=-?6qwvl7&v%oa0I)@sSJSG48Yrco^^9};>Q}3_;Hg@ z{J6Q)#7{0ni3XV(05{QgP~mX~$6J3ZZ|fITo5~27%?P||$(1;_EGpvah@lpDZY0!- zxKb$qXQu!;N51}HB8JLk0Nmz>9e&tl9i=BYO8S10XN64em`R0=J!R!zo-i6IY&wAq&I+~DlwB;t9zgb_>nrJueOMz2mc~k4Ay$S6PER!^?r#k6< z`t&+UiZWwlKeMKvo{9bL0rcfE`V8Jg_a^ShHZSX$i|c*EH-~^xW8s&xp1Nm|+SxI! z{&TTcgT`=5z~cjo_N5Unjh z;Y1PJ;;}&KWRGi2N2@55xf^$4*>vL_rr{bbrSD(o$|SWBFCx}-vd(#@;sc|x4Y)Zz5k;KK1eV>Mq9S! zQfyIFRjx*D%QgtAs&d#-E`Zz^4@Lx;7A>YurF3TjpmC} z@|K{QawTkX+?tvV?4&{$)3YAevwohQNyh3_0>U8!lOYmZbx-K$*y<6qrsVwwkKCG? z;Z{;{`^@DG{w&D3Fh(yeHxRU;6n+^TtfA(jo}}QuvV0iWJyXX6G$2-w_9qB6C=fQM z2%Gh{1{c?kdcHoIG3`c(Wa?QU_Rd{rR;GTZB{w-3Ri37+f1 zJE02az$}}U+FhP_h6xjvgjYs1UR`IoOW+mt%KH%n?ptNgKrFSh@>t`_KGJk67&e z?Oi+-{~#mp&fgc5cQ^M>rT_sY^4%>K&e%gDj2sjm?)yCd5&WCMzH&rUc zrQ#RBfeVLn;L1PXzp-A&ZQKM}4)$qwJUjdL&AeHC87~a{mOJFnJ{RE!f>kO6^Mvbr zWI%ZjFP)ON9U3t5*eJw?=R0y`&Z|`gV2*iC30&`^{#xoWP!l1B%O2h?U* zYi_~)pm$>Vo;Tul=mx^F`;>Fiaal-1yK{zf?>{ulU|freau9C&0pn2wd_}l$N)CNb z9_0yAHSS#63R%_E6Gx>8a7z0_MqE$ueTi`~D^wQTRkke7C2eX8M-@zN1 z=IB;wcJTJiK+6xy)z%?Qm*{^HckfQjd<>vVoOzH}=^Cltt)~8!Qd9V9SFXD`bG}_} z)<7F4P8z?K`nQofw2l`bF9B)Y*q`iaB>r>s05;>N-#THQe1s!$SuuYkDr!h=fpZ@} voudQTx-pWt4!$>z8W({MM{eJx;n{xm9Bo7SuQ^z|&1|9G_ZNB4)=GZ>i9y5+ From 9d781d86cff0826327e266d8e1ea92b241958e29 Mon Sep 17 00:00:00 2001 From: justinchuby Date: Thu, 13 Aug 2026 19:47:42 +0000 Subject: [PATCH 043/151] Fix shared-buffer cache recurrence metadata Keep cache state bounded in workflow IR so synthetic CPU components may grow present-cache tensors while runtime serving metadata still aliases fixed-capacity shared buffers. Refresh decoder and VLM conformance fixtures. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: justinchuby --- .../onnx_genai/workflow_metadata.py | 40 ++++++++----------- .../onnx_genai/workflow_metadata_test.py | 6 ++- .../decoder/inference_metadata.yaml | 4 +- .../vlm/inference_metadata.yaml | 8 +++- 4 files changed, 30 insertions(+), 28 deletions(-) diff --git a/src/mobius/integrations/onnx_genai/workflow_metadata.py b/src/mobius/integrations/onnx_genai/workflow_metadata.py index 53d2a0893..55423889a 100644 --- a/src/mobius/integrations/onnx_genai/workflow_metadata.py +++ b/src/mobius/integrations/onnx_genai/workflow_metadata.py @@ -2951,22 +2951,18 @@ def build_vlm_workflow_metadata( "contract": _contract(past), "scope": "invocation", "initializer": f"decoder.setup.{present.name}", - "recurrence": ( - {"kind": "invariant"} - if fixed_capacity - else { - "kind": "bounded", - "axis": next( - ( - axis - for axis, dimension in enumerate(_contract(past)["shape"]) - if "sequence" in str(dimension) - ), - 2, + "recurrence": { + "kind": "bounded", + "axis": next( + ( + axis + for axis, dimension in enumerate(_contract(past)["shape"]) + if "sequence" in str(dimension) ), - "max": "package.max_context", - } - ), + 2, + ), + "max": "package.max_context", + }, "service_group": "decoder_cache", } setup_decoder_outputs[present.name] = f"decoder.setup.{present.name}" @@ -4710,15 +4706,11 @@ def build_decoder_workflow_metadata( "contract": _contract(past), "scope": "invocation", "initializer": setup_value, - "recurrence": ( - {"kind": "invariant"} - if fixed_capacity - else { - "kind": "bounded", - "axis": decoder_kv_axis, - "max": "package.max_context", - } - ), + "recurrence": { + "kind": "bounded", + "axis": decoder_kv_axis, + "max": "package.max_context", + }, "service_group": "decoder_cache", } decoder_kv_ports[cell] = {"input": past.name, "output": present.name} diff --git a/src/mobius/integrations/onnx_genai/workflow_metadata_test.py b/src/mobius/integrations/onnx_genai/workflow_metadata_test.py index 7f214898c..a3e2a6467 100644 --- a/src/mobius/integrations/onnx_genai/workflow_metadata_test.py +++ b/src/mobius/integrations/onnx_genai/workflow_metadata_test.py @@ -228,7 +228,11 @@ def collect_decoder_invokes(node): "shape": ["batch", 202048], } assert workflow["state"]["attention_mask"]["recurrence"] == {"kind": "invariant"} - assert workflow["state"]["cache_103"]["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", diff --git a/tests/fixtures/onnx_genai_workflows/decoder/inference_metadata.yaml b/tests/fixtures/onnx_genai_workflows/decoder/inference_metadata.yaml index 76563a413..2ad814064 100644 --- a/tests/fixtures/onnx_genai_workflows/decoder/inference_metadata.yaml +++ b/tests/fixtures/onnx_genai_workflows/decoder/inference_metadata.yaml @@ -318,7 +318,9 @@ pipeline: scope: invocation initializer: decoder.setup.present.0.key recurrence: - kind: invariant + kind: bounded + axis: 2 + max: package.max_context service_group: decoder_cache serving: active: active diff --git a/tests/fixtures/onnx_genai_workflows/vlm/inference_metadata.yaml b/tests/fixtures/onnx_genai_workflows/vlm/inference_metadata.yaml index ee79a8a27..2399a345d 100644 --- a/tests/fixtures/onnx_genai_workflows/vlm/inference_metadata.yaml +++ b/tests/fixtures/onnx_genai_workflows/vlm/inference_metadata.yaml @@ -437,7 +437,9 @@ pipeline: scope: invocation initializer: decoder.setup.present.0.key recurrence: - kind: invariant + kind: bounded + axis: 2 + max: package.max_context service_group: decoder_cache cache_1: contract: @@ -451,7 +453,9 @@ pipeline: scope: invocation initializer: decoder.setup.present.0.value recurrence: - kind: invariant + kind: bounded + axis: 2 + max: package.max_context service_group: decoder_cache loop_0_active: contract: From e93d68e35185f147a1482968c10a0a634eba0d8d Mon Sep 17 00:00:00 2001 From: justinchuby Date: Thu, 13 Aug 2026 20:22:21 +0000 Subject: [PATCH 044/151] Make VLM media presence explicit Declare optional image requests with present_as and branch on runtime-observed presence. Text-only workflows now skip preprocessing and vision execution without sentinel tensors or application-provided flags, preserving parity with native text-only benchmarks. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: justinchuby --- benchmarks/README.md | 4 ++-- benchmarks/muse_workflow_h200.json | 2 +- .../onnx_genai/workflow_metadata.py | 18 ++++++------------ .../onnx_genai/workflow_metadata_test.py | 7 +++++-- 4 files changed, 14 insertions(+), 17 deletions(-) diff --git a/benchmarks/README.md b/benchmarks/README.md index f74fcee26..af718fa1c 100644 --- a/benchmarks/README.md +++ b/benchmarks/README.md @@ -6,8 +6,8 @@ paths use the same rendered 68-token prompt, greedy parameters, token budget, warmups, and steady-decode window. `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, bind `request.image` to an empty `uint8[0]` tensor -and `request.has_media` to `false`; the false branch supplies empty image features. +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: diff --git a/benchmarks/muse_workflow_h200.json b/benchmarks/muse_workflow_h200.json index 0629336d5..6aa330916 100644 --- a/benchmarks/muse_workflow_h200.json +++ b/benchmarks/muse_workflow_h200.json @@ -27,7 +27,7 @@ "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_tokens": 68, "image": null, - "workflow_media_binding": "empty_uint8_with_has_media_false", + "workflow_media_binding": "omit_optional_request_image", "max_new_tokens": 128, "request_max_length": 196, "model_max_context": 131072, diff --git a/src/mobius/integrations/onnx_genai/workflow_metadata.py b/src/mobius/integrations/onnx_genai/workflow_metadata.py index 55423889a..63358c058 100644 --- a/src/mobius/integrations/onnx_genai/workflow_metadata.py +++ b/src/mobius/integrations/onnx_genai/workflow_metadata.py @@ -2665,9 +2665,10 @@ def build_vlm_workflow_metadata( "contract": {"dtype": "uint8", "rank": 1, "shape": ["encoded_bytes"]}, "role": {"kind": "runtime", "version": "1.0", "role": "media"}, "source": {"kind": "request", "field": "media"}, - # The frozen contract forbids optional tensors without literal defaults. - # Text-only callers bind an empty uint8 tensor and leave has_media=false. - "required": True, + "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, @@ -2729,14 +2730,6 @@ def build_vlm_workflow_metadata( "default": 0, }, } - if text_only_vision is not None: - inputs["request.has_media"] = { - "contract": {"dtype": "bool", "rank": 1, "shape": [1]}, - "role": {"kind": "opaque"}, - "source": {"kind": "application", "name": "has_media"}, - "required": False, - "default": False, - } vision_invoke_inputs = { name: preprocessing_values[name] for name in vision_inputs @@ -3003,7 +2996,7 @@ def build_vlm_workflow_metadata( vision_setup_nodes: list[dict[str, Any]] = [ { "kind": "branch", - "predicate": "request.has_media", + "predicate": "request.image_present", "cases": { "true": { "kind": "sequence", @@ -3210,6 +3203,7 @@ def build_vlm_workflow_metadata( "nested_control_flow", "loop_induction_values", "typed_emit", + *(["input_presence"] if text_only_vision is not None else []), *( ["serving_service_contract", "bounded_state_recurrence"] if cache_pairs diff --git a/src/mobius/integrations/onnx_genai/workflow_metadata_test.py b/src/mobius/integrations/onnx_genai/workflow_metadata_test.py index a3e2a6467..6d7c22cd1 100644 --- a/src/mobius/integrations/onnx_genai/workflow_metadata_test.py +++ b/src/mobius/integrations/onnx_genai/workflow_metadata_test.py @@ -241,10 +241,13 @@ def collect_decoder_invokes(node): assert policy_invokes["decoder_step_update"]["inputs"]["logical_length"] == ( "cache_lengths.next" ) - assert workflow["inputs"]["request.image"]["required"] is True - assert workflow["inputs"]["request.has_media"]["default"] is False + 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"] 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" kv_service = workflow["serving"]["kv_service"] From 3934de7688d3962f2b2de9d3e5493c5a07e02bb9 Mon Sep 17 00:00:00 2001 From: justinchuby Date: Thu, 13 Aug 2026 20:28:50 +0000 Subject: [PATCH 045/151] Derive generation constants from package metadata Prefer packaged genai_config values for EOS and model context when producing VLM workflows. This prevents fallback or tiny build configs from silently changing real artifact termination and capacity contracts. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: justinchuby --- .../onnx_genai/workflow_metadata.py | 45 ++++++++++++++++--- .../onnx_genai/workflow_metadata_test.py | 11 ++++- 2 files changed, 49 insertions(+), 7 deletions(-) diff --git a/src/mobius/integrations/onnx_genai/workflow_metadata.py b/src/mobius/integrations/onnx_genai/workflow_metadata.py index 63358c058..5753502a9 100644 --- a/src/mobius/integrations/onnx_genai/workflow_metadata.py +++ b/src/mobius/integrations/onnx_genai/workflow_metadata.py @@ -55,6 +55,37 @@ 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( @@ -2651,9 +2682,7 @@ def build_vlm_workflow_metadata( batch_int = {"dtype": "int64", "rank": 1, "shape": [batch]} batch_bool = {"dtype": "bool", "rank": 1, "shape": [batch]} control_int = {"dtype": "int64", "rank": 1, "shape": [1]} - eos = getattr(config, "eos_token_id", 0) - if isinstance(eos, list): - eos = eos[0] if eos else 0 + 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), @@ -2685,14 +2714,20 @@ def build_vlm_workflow_metadata( "role": {"kind": "opaque"}, "source": {"kind": "literal"}, "required": False, - "default": int(eos or 0), + "default": eos, }, "package.max_context": { "contract": control_int, "role": {"kind": "opaque"}, "source": {"kind": "literal"}, "required": False, - "default": int(getattr(config, "max_position_embeddings", 4096)), + "default": int( + _source_model_value( + source, + "context_length", + getattr(config, "max_position_embeddings", 4096), + ) + ), }, "package.one": { "contract": control_int, diff --git a/src/mobius/integrations/onnx_genai/workflow_metadata_test.py b/src/mobius/integrations/onnx_genai/workflow_metadata_test.py index 6d7c22cd1..0f1c0c2cb 100644 --- a/src/mobius/integrations/onnx_genai/workflow_metadata_test.py +++ b/src/mobius/integrations/onnx_genai/workflow_metadata_test.py @@ -110,6 +110,10 @@ def test_vlm_writer_derives_real_decoder_contract_from_artifact(tmp_path): (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}}), + encoding="utf-8", + ) (source / "preprocessor_config.json").write_text( json.dumps( { @@ -177,10 +181,12 @@ def test_vlm_writer_derives_real_decoder_contract_from_artifact(tmp_path): ) # 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"), - _VlmConfig(), + config, source=str(source), ) serialized = Path(path).read_text(encoding="utf-8") @@ -216,7 +222,8 @@ def collect_decoder_invokes(node): len([name for name in workflow["state"] if name.removeprefix("cache_").isdigit()]) == 104 ) - assert workflow["inputs"]["package.max_context"]["default"] == 4096 + 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, From 32d1c46eebd92980f862dc4a133dfb11f2e8e38c Mon Sep 17 00:00:00 2001 From: justinchuby Date: Thu, 13 Aug 2026 20:36:23 +0000 Subject: [PATCH 046/151] Mark autoregressive EOS loop termination Emit the generation_eos loop contract for decoder and VLM workflows so stop_on_eos=false executes the fixed token count without inspecting device-resident predicates. Refresh the representative fixtures. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: justinchuby --- src/mobius/integrations/onnx_genai/workflow_metadata.py | 4 ++++ src/mobius/integrations/onnx_genai/workflow_metadata_test.py | 1 + .../onnx_genai_workflows/decoder/inference_metadata.yaml | 1 + .../fixtures/onnx_genai_workflows/vlm/inference_metadata.yaml | 1 + 4 files changed, 7 insertions(+) diff --git a/src/mobius/integrations/onnx_genai/workflow_metadata.py b/src/mobius/integrations/onnx_genai/workflow_metadata.py index 5753502a9..adf60c3d7 100644 --- a/src/mobius/integrations/onnx_genai/workflow_metadata.py +++ b/src/mobius/integrations/onnx_genai/workflow_metadata.py @@ -296,6 +296,8 @@ def convert(node: dict[str, Any]) -> dict[str, Any]: "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 @@ -3309,6 +3311,7 @@ def build_vlm_workflow_metadata( "setup": setup, "body": body, "condition": "loop.continue", + "termination": "generation_eos", "max_iterations": "request.max_iterations", "iteration": {"value": "loop.iteration", "contract": batch_int}, "carried": carried, @@ -4970,6 +4973,7 @@ def build_decoder_workflow_metadata( "setup": setup, "body": body, "condition": "loop.continue", + "termination": "generation_eos", **({"active_cell": "active"} if cache_pairs else {}), "max_iterations": "request.max_iterations", "iteration": { diff --git a/src/mobius/integrations/onnx_genai/workflow_metadata_test.py b/src/mobius/integrations/onnx_genai/workflow_metadata_test.py index 0f1c0c2cb..d12c214bf 100644 --- a/src/mobius/integrations/onnx_genai/workflow_metadata_test.py +++ b/src/mobius/integrations/onnx_genai/workflow_metadata_test.py @@ -252,6 +252,7 @@ def collect_decoder_invokes(node): 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" diff --git a/tests/fixtures/onnx_genai_workflows/decoder/inference_metadata.yaml b/tests/fixtures/onnx_genai_workflows/decoder/inference_metadata.yaml index 2ad814064..8cee222ff 100644 --- a/tests/fixtures/onnx_genai_workflows/decoder/inference_metadata.yaml +++ b/tests/fixtures/onnx_genai_workflows/decoder/inference_metadata.yaml @@ -464,6 +464,7 @@ pipeline: next: decoder_step.body_position_ids - cell: cache_9 next: decoder.body.present.0.key + termination: generation_eos iteration: value: loop.iteration contract: diff --git a/tests/fixtures/onnx_genai_workflows/vlm/inference_metadata.yaml b/tests/fixtures/onnx_genai_workflows/vlm/inference_metadata.yaml index 2399a345d..9bfd5ed06 100644 --- a/tests/fixtures/onnx_genai_workflows/vlm/inference_metadata.yaml +++ b/tests/fixtures/onnx_genai_workflows/vlm/inference_metadata.yaml @@ -649,6 +649,7 @@ pipeline: next: decoder.body.present.0.value - cell: loop_0_active next: loop.continue + termination: generation_eos iteration: value: loop.iteration contract: From f22c9df628c2aed9cee09e1a1ee566f21ab3aee9 Mon Sep 17 00:00:00 2001 From: justinchuby Date: Thu, 13 Aug 2026 20:40:08 +0000 Subject: [PATCH 047/151] Pin exact Muse benchmark prompt tokens Pass the canonical 68 prompt token IDs directly to the workflow profiler so tokenizer special-token insertion cannot drift the paired native/workflow request or KV capacity. Document and validate the fixed prompt contract. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: justinchuby --- benchmarks/README.md | 17 ++++---- benchmarks/muse_prompt_ids.json | 70 ++++++++++++++++++++++++++++++ benchmarks/muse_workflow_h200.json | 1 + scripts/benchmark_muse_workflow.py | 8 +++- 4 files changed, 86 insertions(+), 10 deletions(-) create mode 100644 benchmarks/muse_prompt_ids.json diff --git a/benchmarks/README.md b/benchmarks/README.md index af718fa1c..be55fc696 100644 --- a/benchmarks/README.md +++ b/benchmarks/README.md @@ -2,10 +2,11 @@ `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 same rendered 68-token prompt, greedy parameters, token budget, -warmups, and steady-decode window. `request_max_length` is prompt tokens plus -new tokens (68 + 128 = 196); `model_max_context` is the independent 131072-token -artifact admission ceiling. +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. +`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. @@ -27,7 +28,7 @@ python scripts/benchmark_muse_workflow.py \ --output artifacts/muse-int4-package/workflow-benchmark.json ``` -The workflow runner must support `--pipeline --backend ort --ep cuda` 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. +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_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 index 6aa330916..4215008bf 100644 --- a/benchmarks/muse_workflow_h200.json +++ b/benchmarks/muse_workflow_h200.json @@ -25,6 +25,7 @@ "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", diff --git a/scripts/benchmark_muse_workflow.py b/scripts/benchmark_muse_workflow.py index fde40f648..0f819ecf2 100644 --- a/scripts/benchmark_muse_workflow.py +++ b/scripts/benchmark_muse_workflow.py @@ -38,6 +38,10 @@ def main() -> int: config: dict[str, Any] = json.loads(args.config.read_text()) workload = config["workload"] sampling = config["sampling"] + 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"] ): @@ -70,8 +74,8 @@ def main() -> int: str(workload["runs"]), "--decode-skip", str(workload["decode_skip"]), - "--prompt", - workload["rendered_prompt"], + "--prompt-ids", + str(prompt_ids_path), ] if workload["image"]: command.extend(["--image", workload["image"]]) From b30e4db202c579f46d6dffdaff09b8b9a48ef622 Mon Sep 17 00:00:00 2001 From: justinchuby Date: Thu, 13 Aug 2026 20:41:47 +0000 Subject: [PATCH 048/151] Verify native Muse prompt identity Require ORT GenAI tokenization to match the canonical paired benchmark token IDs exactly, rather than accepting only the same token count. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: justinchuby --- scripts/benchmark_muse_native.py | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/scripts/benchmark_muse_native.py b/scripts/benchmark_muse_native.py index fc77ab79a..cdcd743cc 100644 --- a/scripts/benchmark_muse_native.py +++ b/scripts/benchmark_muse_native.py @@ -122,7 +122,11 @@ def main() -> int: if image is not None else workload["rendered_prompt"] ) - prompt_tokens = len(tokenizer.encode(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']}" From 8c936d613d8f4c94cb96837e5b879b86f380789d Mon Sep 17 00:00:00 2001 From: justinchuby Date: Thu, 13 Aug 2026 20:44:57 +0000 Subject: [PATCH 049/151] Pin paired Muse attention kernel Disable cuDNN Flash Attention explicitly for both native and workflow runners so the paired benchmark uses the established native GQA kernel and avoids unsupported cuDNN frontend shapes. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: justinchuby --- benchmarks/README.md | 2 ++ benchmarks/muse_workflow_h200.json | 1 + scripts/benchmark_muse_native.py | 4 +++- scripts/benchmark_muse_workflow.py | 13 ++++++++++++- 4 files changed, 18 insertions(+), 2 deletions(-) diff --git a/benchmarks/README.md b/benchmarks/README.md index be55fc696..4aadee135 100644 --- a/benchmarks/README.md +++ b/benchmarks/README.md @@ -5,6 +5,8 @@ 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 diff --git a/benchmarks/muse_workflow_h200.json b/benchmarks/muse_workflow_h200.json index 4215008bf..b13b6b94f 100644 --- a/benchmarks/muse_workflow_h200.json +++ b/benchmarks/muse_workflow_h200.json @@ -16,6 +16,7 @@ "onnxruntime_genai_commit": "ede24ecc6a254ef33354bc0eb7bc90d1daa540bd", "execution_provider": "CUDAExecutionProvider", "cuda_graph": true, + "cudnn_flash_attention": false, "shared_kv": true }, "hardware": { diff --git a/scripts/benchmark_muse_native.py b/scripts/benchmark_muse_native.py index cdcd743cc..fca1789a4 100644 --- a/scripts/benchmark_muse_native.py +++ b/scripts/benchmark_muse_native.py @@ -16,7 +16,7 @@ from pathlib import Path from typing import Any -os.environ.setdefault("ORT_ENABLE_CUDNN_FLASH_ATTENTION", "0") +os.environ["ORT_ENABLE_CUDNN_FLASH_ATTENTION"] = "0" import onnxruntime as ort import onnxruntime_genai as og @@ -103,6 +103,8 @@ def main() -> int: config = json.loads(args.config.read_text()) workload = config["workload"] sampling = config["sampling"] + if config["runtime"]["cudnn_flash_attention"]: + raise ValueError("paired Muse benchmark requires cuDNN Flash Attention disabled") if ort.__version__ != config["runtime"]["onnxruntime"]: raise RuntimeError( diff --git a/scripts/benchmark_muse_workflow.py b/scripts/benchmark_muse_workflow.py index 0f819ecf2..30bdcb32b 100644 --- a/scripts/benchmark_muse_workflow.py +++ b/scripts/benchmark_muse_workflow.py @@ -9,6 +9,7 @@ import argparse import hashlib import json +import os import re import subprocess from pathlib import Path @@ -38,6 +39,8 @@ def main() -> int: config: dict[str, Any] = json.loads(args.config.read_text()) workload = config["workload"] sampling = config["sampling"] + 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"]): @@ -79,7 +82,15 @@ def main() -> int: ] if workload["image"]: command.extend(["--image", workload["image"]]) - completed = subprocess.run(command, check=True, text=True, capture_output=True) + environment = os.environ.copy() + 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 " From 30cc0febd6b5329da42ba037adfe2bfa4218b403 Mon Sep 17 00:00:00 2001 From: justinchuby Date: Thu, 13 Aug 2026 20:51:43 +0000 Subject: [PATCH 050/151] Validate current workflow termination contract Pin metadata CI and the Muse benchmark contract to ONNX GenAI dbb5d0e5, which admits generation_eos loop termination and contains the device-resident fixed-count runtime support required by the real workflow. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: justinchuby --- .github/workflows/main.yml | 2 +- benchmarks/muse_workflow_h200.json | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/.github/workflows/main.yml b/.github/workflows/main.yml index 0fd3ec7f7..5828d9605 100644 --- a/.github/workflows/main.yml +++ b/.github/workflows/main.yml @@ -23,7 +23,7 @@ jobs: - uses: actions/checkout@v7 with: repository: justinchuby/onnx-genai - ref: 923530877b440145f5fc1b848d50cac73e02e2f6 + ref: dbb5d0e55fe113248a082b5ac23d997e94b2b534 path: validation/onnx-genai - uses: actions/setup-python@v7 with: diff --git a/benchmarks/muse_workflow_h200.json b/benchmarks/muse_workflow_h200.json index b13b6b94f..8599cc897 100644 --- a/benchmarks/muse_workflow_h200.json +++ b/benchmarks/muse_workflow_h200.json @@ -3,7 +3,7 @@ "package": { "repository": "justinchuby/Muse-Glimmer-30B-ONNX-INT4-CUDA", "weights_revision": "bf36a94a4519e14e3c48ad005c6ff1972ab44ccb", - "schema_head": "a341c463a0090298238102506796bc73d86b84fa", + "schema_head": "dbb5d0e55fe113248a082b5ac23d997e94b2b534", "artifacts": { "decoder": "decoder/model.onnx", "embedding": "embedding/model.onnx", From 6d51fa4af62710d2e9a7f2fabee64a811c519ec7 Mon Sep 17 00:00:00 2001 From: justinchuby Date: Thu, 13 Aug 2026 21:18:07 +0000 Subject: [PATCH 051/151] Derive dynamic KV storage from decoder ABI Treat differing past/present sequence dimensions as separate growing cache storage, initialize conventional decoder caches at zero length, and derive logical cache lengths from the prompt. Preserve fixed-capacity initialization only for artifacts whose admitted past/present shapes share the same capacity contract. This fixes real Muse prefill GQA and restores all seven ONNX GenAI workflow conformance fixtures under persistent execution islands. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: justinchuby --- src/mobius/generation/_policy_components.py | 13 ++-- .../generation/_policy_components_test.py | 8 ++- .../onnx_genai/workflow_metadata.py | 57 ++++++++---------- .../onnx_genai/workflow_metadata_test.py | 30 +++++++-- .../decoder/inference_metadata.yaml | 21 ++----- .../policies/decoder_state_initializer.onnx | Bin 9625 -> 7598 bytes .../decoder/policies/decoder_step_update.onnx | Bin 3203 -> 2199 bytes .../speculative/inference_metadata.yaml | 2 +- .../tts/inference_metadata.yaml | 4 +- .../vlm/inference_metadata.yaml | 9 +-- .../policies/decoder_state_initializer.onnx | Bin 10907 -> 8704 bytes .../vlm/policies/decoder_step_update.onnx | Bin 3203 -> 2199 bytes 12 files changed, 74 insertions(+), 70 deletions(-) diff --git a/src/mobius/generation/_policy_components.py b/src/mobius/generation/_policy_components.py index b69213b1d..caf40d21d 100644 --- a/src/mobius/generation/_policy_components.py +++ b/src/mobius/generation/_policy_components.py @@ -659,13 +659,12 @@ def build_decoder_state_initializer( if position_ids_input is not None: builder.add_output(body_position, "body_position_ids") builder.add_output(token_slot, "token_slot") - if fixed_capacity: - cache_lengths = op.Expand( - op.Unsqueeze(sequence_length, op.Constant(value_ints=[0])), - batch_shape, - ) - cache_lengths.shape = ir.Shape(["batch"]) - builder.add_output(cache_lengths, "cache_lengths") + cache_lengths = op.Expand( + op.Unsqueeze(sequence_length, op.Constant(value_ints=[0])), + batch_shape, + ) + cache_lengths.shape = ir.Shape(["batch"]) + builder.add_output(cache_lengths, "cache_lengths") for name in cache_inputs: value = decoder_inputs[name] diff --git a/src/mobius/generation/_policy_components_test.py b/src/mobius/generation/_policy_components_test.py index ae6699c17..35b5f01e9 100644 --- a/src/mobius/generation/_policy_components_test.py +++ b/src/mobius/generation/_policy_components_test.py @@ -220,7 +220,8 @@ def test_decoder_state_initializer_and_step_update_runtime(tmp_path): 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) + np.testing.assert_array_equal(outputs[5], [3]) + assert outputs[6].shape == (1, 2, 0, 4) updated = _run( build_decoder_step_update( @@ -344,7 +345,10 @@ def test_decoder_policy_chain_generates_multiple_tokens_from_prompt_only(tmp_pat tmp_path, {"prompt_tokens": prompt}, ) - attention, positions, body_attention, body_position, token, cache = initialized + attention, positions, body_attention, body_position, token, cache_lengths, cache = ( + initialized + ) + np.testing.assert_array_equal(cache_lengths, [2]) logits, cache = _run_model( decoder, tmp_path, diff --git a/src/mobius/integrations/onnx_genai/workflow_metadata.py b/src/mobius/integrations/onnx_genai/workflow_metadata.py index adf60c3d7..fcbfa2fdb 100644 --- a/src/mobius/integrations/onnx_genai/workflow_metadata.py +++ b/src/mobius/integrations/onnx_genai/workflow_metadata.py @@ -509,10 +509,28 @@ def _kv_storage_contract(model: ir.Model) -> dict[str, Any]: for name in input_names for marker in ("block_table", "block_tables", "page_table", "page_tables") ) + cache_pairs = _model_cache_pairs(model) + shared_buffer = bool(cache_pairs) and all( + past.shape is not None + and present.shape is not None + and len(past.shape) == len(present.shape) + and all( + str(getattr(past_dim, "value", past_dim)) + == str(getattr(present_dim, "value", present_dim)) + for past_dim, present_dim in zip(past.shape, present.shape) + ) + for past, present in cache_pairs + ) + if paged: + storage = "paged" + elif shared_buffer: + storage = "shared_buffer" + else: + storage = "separate" return { "paging": "paged" if paged else "none", "compaction": paged, - "storage": "paged" if paged else "shared_buffer", + "storage": storage, } @@ -2897,9 +2915,7 @@ def build_vlm_workflow_metadata( "contract": batch_int, "class": "semantic", "scope": "invocation", - "initializer": ( - "initializer.cache_lengths" if fixed_capacity else "package.zero_batch" - ), + "initializer": "initializer.cache_lengths", "recurrence": {"kind": "invariant"}, }, } @@ -2949,7 +2965,7 @@ def build_vlm_workflow_metadata( ), ( "cache_lengths", - "initializer.cache_lengths" if fixed_capacity else "package.zero_batch", + "initializer.cache_lengths", "state.cache_lengths.body", "cache_lengths.next", "state.cache_lengths.final", @@ -3094,11 +3110,7 @@ def build_vlm_workflow_metadata( attention_input.name: f"initializer.{attention_input.name}", "body_attention_mask": "initializer.body_attention_mask", "token_slot": "initializer.token_slot", - **( - {"cache_lengths": "initializer.cache_lengths"} - if fixed_capacity - else {} - ), + "cache_lengths": "initializer.cache_lengths", **( { position_input.name: f"initializer.{position_input.name}", @@ -4468,13 +4480,6 @@ def build_decoder_workflow_metadata( "required": False, "default": 0, }, - "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"}, @@ -4561,11 +4566,7 @@ def build_decoder_workflow_metadata( "contract": batch_int, "class": "semantic", "scope": "invocation", - "initializer": ( - "initializer.cache_lengths" - if fixed_capacity - else "package.cache_lengths" - ), + "initializer": "initializer.cache_lengths", "recurrence": {"kind": "invariant"}, }, } @@ -4621,11 +4622,7 @@ def build_decoder_workflow_metadata( }, { "cell": "cache_lengths", - "current": ( - "initializer.cache_lengths" - if fixed_capacity - else "package.cache_lengths" - ), + "current": "initializer.cache_lengths", "body_input": "state.cache_lengths.body", "body_output": "cache_lengths.next", "next": "state.cache_lengths.final", @@ -4773,11 +4770,7 @@ def build_decoder_workflow_metadata( attention_input.name: f"initializer.{attention_input.name}", "body_attention_mask": "initializer.body_attention_mask", "token_slot": "initializer.token_slot", - **( - {"cache_lengths": "initializer.cache_lengths"} - if fixed_capacity - else {} - ), + "cache_lengths": "initializer.cache_lengths", **( { position_input.name: f"initializer.{position_input.name}", diff --git a/src/mobius/integrations/onnx_genai/workflow_metadata_test.py b/src/mobius/integrations/onnx_genai/workflow_metadata_test.py index d12c214bf..fce12a4af 100644 --- a/src/mobius/integrations/onnx_genai/workflow_metadata_test.py +++ b/src/mobius/integrations/onnx_genai/workflow_metadata_test.py @@ -17,6 +17,7 @@ _VlmConfig, ) from mobius.integrations.onnx_genai.workflow_metadata import ( + _kv_storage_contract, build_language_diffusion_pipeline_metadata, build_speculative_workflow_metadata, build_vlm_workflow_metadata, @@ -25,6 +26,21 @@ ) +def test_kv_storage_is_derived_from_admitted_cache_shapes(): + dynamic = _model( + "dynamic", + [_value("past_key_values.0.key", ir.DataType.FLOAT, ["batch", 2, "past", 8])], + [("present.0.key", ir.DataType.FLOAT, ["batch", 2, "total", 8])], + ) + shared = _model( + "shared", + [_value("past_key_values.0.key", ir.DataType.FLOAT, ["batch", 2, "capacity", 8])], + [("present.0.key", ir.DataType.FLOAT, ["batch", 2, "capacity", 8])], + ) + assert _kv_storage_contract(dynamic)["storage"] == "separate" + assert _kv_storage_contract(shared)["storage"] == "shared_buffer" + + 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() @@ -234,7 +250,12 @@ def collect_decoder_invokes(node): "rank": 2, "shape": ["batch", 202048], } - assert workflow["state"]["attention_mask"]["recurrence"] == {"kind": "invariant"} + assert workflow["state"]["attention_mask"]["recurrence"] == { + "kind": "growing", + "axis": 1, + "increment": "package.one", + "max": "package.max_context", + } assert workflow["state"]["cache_103"]["recurrence"] == { "kind": "bounded", "axis": 2, @@ -243,11 +264,8 @@ def collect_decoder_invokes(node): assert workflow["state"]["cache_lengths"]["initializer"] == "initializer.cache_lengths" assert policy_invokes["decoder_state_initializer"]["inputs"] == { "prompt_tokens": "request.prompt_tokens", - "max_iterations": "request.max_iterations", } - assert policy_invokes["decoder_step_update"]["inputs"]["logical_length"] == ( - "cache_lengths.next" - ) + assert "logical_length" not in policy_invokes["decoder_step_update"]["inputs"] 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"] @@ -262,7 +280,7 @@ def collect_decoder_invokes(node): assert kv_service["paging"] == "none" assert kv_service["compaction"] is False decoder_cache = kv_service["groups"]["decoder_cache"] - assert decoder_cache["storage"] == "shared_buffer" + assert decoder_cache["storage"] == "separate" kv_ports = decoder_cache["ports"]["decoder"] assert len(kv_ports) == 104 assert kv_ports["cache_103"] == { diff --git a/tests/fixtures/onnx_genai_workflows/decoder/inference_metadata.yaml b/tests/fixtures/onnx_genai_workflows/decoder/inference_metadata.yaml index 8cee222ff..6efa79903 100644 --- a/tests/fixtures/onnx_genai_workflows/decoder/inference_metadata.yaml +++ b/tests/fixtures/onnx_genai_workflows/decoder/inference_metadata.yaml @@ -113,18 +113,6 @@ pipeline: kind: literal required: false default: 0 - package.cache_lengths: - contract: - dtype: int64 - rank: 1 - shape: - - batch - role: - kind: opaque - source: - kind: literal - required: false - default: 0 package.zero_batch: contract: dtype: int64 @@ -294,7 +282,10 @@ pipeline: scope: invocation initializer: initializer.body_attention_mask recurrence: - kind: invariant + kind: growing + axis: 1 + increment: package.one_token + max: package.max_context position_ids: contract: dtype: int64 @@ -336,7 +327,7 @@ pipeline: sequence_axis: 2 layout: bnsh logical_lengths: cache_lengths - storage: shared_buffer + storage: separate ports: model: cache_9: @@ -349,7 +340,6 @@ pipeline: component: decoder_state_initializer inputs: prompt_tokens: request.input_ids - max_iterations: request.max_iterations outputs: attention_mask: initializer.attention_mask body_attention_mask: initializer.body_attention_mask @@ -436,7 +426,6 @@ pipeline: component: decoder_step_update inputs: attention_mask: attention_mask - logical_length: cache_lengths.next position_ids: position_ids outputs: next_attention_mask: decoder_step.body_attention_mask diff --git a/tests/fixtures/onnx_genai_workflows/decoder/policies/decoder_state_initializer.onnx b/tests/fixtures/onnx_genai_workflows/decoder/policies/decoder_state_initializer.onnx index b85fa5c58d6cb0ed70abd2aec96ad22f27d9a04c..6fb3c64fe81b4f908fa8096554d13f5457539ba5 100644 GIT binary patch literal 7598 zcmd5>&2QsG6pusNq|+v}v)XQJ(5f4)7RXA$GtMVm_Cp-#1qpFx&g-yJE>UEs4|JE;W zqpuWH9FDxxVPFK_iS7D$$yT<|nje^>U@rrs7XXn85UB()D?ns`(}MV|g163$$DRw?a05eEUrwGw zE9&J!{Gzgh(3v?H+lJ!?{sT~4J}p!~C@5hSMFK#)33U5(Y}?Olqo$VP#}I6-oS%%- z0`?aNOv5*(ntSgkqJCZ7jrA433Zd1@iDD=3D=JzG)mbSOlK;AL14Rq7S^Zk4--?Lg0;ZjtwdsC>wZt`PEUeDO0hzkcx0YQllccnS_c& zpkfoF(EgSg7%jC3W@k^K(pu&9Xjs`qIrH4{Q*|?cD{wK{9S)JqV`X~*E8G7IR>Ezc z!^$=bEFljG1ePfwiMBFE{-Gf0W-u;HFyR{=uX|xgH>bqYfRGs;r0L7QS8y@pLFCH` zRks06Gze))=nJ_*zX?}Jgvi|aL}f%*MeXS`+DkK7|EVAz3r1mjWB3CvP;bS06nr|V zEx$e!pANuNbL!toX0Okfz0R3^8z>wH@l?H%%*MtHtOf^bb-Be)92+ST*w|pXP@GW^ z2te_3bbFeqrx&s*TQfRqaXJ%BnWkg;9a3%HvD8@{jCDhr;$(}3#rIA3*fvO8Ek;j~ z0^8-aaF|wlBTiJ&g3;nJD$q3cglT1wQN=_phiCFPCeRa&b1 z5kj4aW19A8%t^mfRUX+sSBw6W5_4@FU&)zKQ41JDNE z2~#|Qa3cGgX|rE8?My_Vd_)jUn`VXnYj`W;7*G>^T1)b2?OC7JUV%@8NH6-d_SB~l zT(nH9XBz|CJq~)l=+$&>#dg(+9IAA(FjPTzD7A7|;+sBGL+FcZh&1^oy{oCDOMk;T`IdTpf+5c&tx%15u^n?u(9Ik7K{uu|}=S``Kh ziXeSXXKOd^Y$51uDTg^$1f^F1;r*^{LC1MyK&czB?j)oK=D>Mok3JsD&FU`d6!Q2J z^g6yb^ahS~QC+yW+A#yGhbY#TemQA~`!-~8u4Q+|a^C=3$OO4huOXM7cx2ADY!w41N5IMFTKS8(&g`c3pgetU5bz$x%b2~{2f!9$E1ZDbn=5+(> z$;u4M^u=AF5u0FC`g=k%bUKrrzi~TR%r*{B9yj>nwQS-Q!M79fko%nfB>Wkf#>C1w zIi<8=BY^>I?C**NcCtv?t(8pd@&zU?pmoq>b_O;SoAH^Ll-p8zMDP|&=V4Qrg~j2q zriE9$RMZo#344mif;2x~#RT$tg6w)am6h~Mp^3nm^i>Jy)Mji~QKT*sQ<@a*CQb5| z6Tw1IoZ`t{DzBkC(LbWqW7~zlEm=7_xoC#vxxvV^0-XKq^GoF+y1VpC!|CGv$ybxx jYL3sHfn$y?9*|4r9y(aKhkbNpTS54b=%XCUZ14RG2ZU%- literal 9625 zcmc&)TW{P%6ppu>>}IkF$)L8WK~-5*1r}*#&&=L<>P0*dkdQzXQmbNly>8-UV=rqb zN{Zm24}Adyh?@W*6(RltZ@lo!IKIz}yS*UU$xCKDo|(_*oH^$^XSTV9R}Q_gbL7AM z=somELF==*cQ^}-z&o&AAFm!8AMKel+tAcj+;F`K^xDf&Z)Y8?y5^zn&rHk4OFLK5 z<=Me**K^$y-<5WiP&>)5#ExgZTUFsS-c z){eAD+s*VrXA7ZYb9!VOjvM&bL2!+-Q2nH!%&G_y0OB<;w@;63`+;pRbt7#?U|GA7 z^veP!iK{1zt9K^i`b|L_hvtdl1aO1`2b3wutjEi&FUhQ5y3D>*5U(@7)5nctOFNUJD#z3(|Z-1*k2$Zh&V#8n;cGE3lHr(Co|KX7&^p* zq#0|q-(HFPWvTVsC>&(l^WlJnT=IPz$e=z@TTvqhHE7qvURjX7iyU0CG2SqPz;?qh zY8;yWf!exd`az<+VY?poI;*G_c-xHFCH z(bTmpo)fLqfc~{)PZ@;FmVp&Y#+!1fal7C8mNDy0{HM`o4)APl1@`!)Qj7pNCg#2?(}5@6D6 zXUSAD&{5HHL_0bgO=YvO(_Bja9xL^-oSWoGNWY-={vt8TLk|xlR=+l^1 z@M%n|z^BPcGShOOrsH+XwDxRcYP-9^o-cYerdO<2R;U^=zG@`u6)q1-7?m`;}CF|KU?J*2m{(VZr9WJ~g6T%kOGk=q$+6*r1rDT}ZPu z1X-H0QxB80@-j4hX<}P&V>&nB4lW39*MPs6Q|Ey_zjGvCwocHv*}$KmH}JKYH+8Jj z?!w94u^CuO!D%{FZVeGtE!d^jjkkgVhrMGxWAE#5cvP+okk+|_tSml+_Bj?u(kXdtxFL6{A zD*pLU`1yQ@)b(sLX@Fz{7x)_w+(awzg*CLaA0fy>h7MIR4#8zHfapEkg5d@%3N{Lp zBZS#KKsI#=h9b?S_4shjRd7A=fB=1zA;c)J`=;W|SAjw-zlsPOs*F5rUT|TZFZ&`{ za3oipwBUf7rR-l2%OM$LQQA-#NpW-1rrJVjLunw#rtC}7rsOHohBP^ynn^oJQL_Qi zk)&Qo%F4V&Bl0gz|ncK`qY diff --git a/tests/fixtures/onnx_genai_workflows/decoder/policies/decoder_step_update.onnx b/tests/fixtures/onnx_genai_workflows/decoder/policies/decoder_step_update.onnx index 0333f2f9bf0c4417d7b0e6fd4e2772628899dcb9..390f62573f15a272042ba9b3688c9fcdea316696 100644 GIT binary patch literal 2199 zcmcIlO>fgc5cQ^M>rT_sY^4%>K&e%gDj2sjm?)yCd5&WCMzH&rUc zrQ#RBfeVLn;L1PXzp-A&ZQKM}4)$qwJUjdL&AeHC87~a{mOJFnJ{RE!f>kO6^Mvbr zWI%ZjFP)ON9U3t5*eJw?=R0y`&Z|`gV2*iC30&`^{#xoWP!l1B%O2h?U* zYi_~)pm$>Vo;Tul=mx^F`;>Fiaal-1yK{zf?>{ulU|freau9C&0pn2wd_}l$N)CNb z9_0yAHSS#63R%_E6Gx>8a7z0_MqE$ueTi`~D^wQTRkke7C2eX8M-@zN1 z=IB;wcJTJiK+6xy)z%?Qm*{^HckfQjd<>vVoOzH}=^Cltt)~8!Qd9V9SFXD`bG}_} z)<7F4P8z?K`nQofw2l`bF9B)Y*q`iaB>r>s05;>N-#THQe1s!$SuuYkDr!h=fpZ@} voudQTx-pWt4!$>z8W({MM{eJx;n{xm9Bo7SuQ^z|&1|9G_ZNB4)=GZ>i9y5+ literal 3203 zcmc(h-EPw`6vvxxT~n`v;%Tfz0*$DGsX|rLu3d`|Vw(#@;sc|x4Y)Zz5k;KK1eV>Mq9S! zQfyIFRjx*D%QgtAs&d#-E`Zz^4@Lx;7A>YurF3TjpmC} z@|K{QawTkX+?tvV?4&{$)3YAevwohQNyh3_0>U8!lOYmZbx-K$*y<6qrsVwwkKCG? z;Z{;{`^@DG{w&D3Fh(yeHxRU;6n+^TtfA(jo}}QuvV0iWJyXX6G$2-w_9qB6C=fQM z2%Gh{1{c?kdcHoIG3`c(Wa?QU_Rd{rR;GTZB{w-3Ri37+f1 zJE02az$}}U+FhP_h6xjvgjYs1UR`IoOW+mt%KH%n?ptNgKrFSh@>t`_KGJk67&e z?Oi+-{~#mp&`5Rrqc)AS#7s#(5f4)7RXAGXY4!-mwku>y&!=Q605~>9JfiE+97tU zb`e@>5uCUnUV_si#Emm25aJKwf8ZM@c4o%0vPH5>Po9itzTbTF`<`xf4KE-2-N8xl z=06Y7mkQb#js4?MXodcf>jilA)Ox>fk6cSr*C!92--A(YDH`oUCR=zE|IFSK;^ z<>WcKMZH#yf2dzU=+qvbxYoc6gL|O3mAp`Wub_lg6bS(F8qn>N6W4v>T86qAKSp4U z+RAjC7qCA#U>d$U)7-oJ6Y4kB?O0#&PZ6|cEm7?B_qvLfB6Zd`tIF2OUR5bg{+=t9 zyx{+&pm~__rX7Z^7Y=;SI<|u&b?u=YgjNK+RjaT;eHE2Me`lpQDz;@Rwlk>+29g^U z!ObL8Bmx!N7)AH**rC-?H{jaYQ>3&`ZFw@Ruc3;4HVATcGk+>zn9PPFB=cBlX0X!y zU$7FaeGV&47FZ%45(q3aLK3YoM*gNC$ubxhrkL>6uHQSiB+J?4X+TJd2YLGP?-aZd z@nGW1301cM&14YKlF&;rp)#c`W5Fm&Z>(VGhwAlM zkAhDp#$xlC`g8!En^XT*y7p%3+MC?9ZvlnlAfBsNlG)fw!D?}^ipwp2J-v{xvXjzThtrvGm3carUm?}z9ZQqN!B{t>Elze=SbWd+4qS_v z>c-@0qCm5@6beLx5ux21n+7DgsbG0*JiNGk4Ao{Itmhe%cZPKW!Bn_@O9@vpQ;U zIRI_pD^ZFk5Kd%&J8$;OSDT6ml#dAF)#h2D{~TFm90Lr|r*$NsW~P0bS%6Q2NH6*{ zbLP_s-f(QE?^;9GI|%!M=+$(s;Ct1H9IAAZ8LFT=lv}wg@l8+F5PEhEktg3|TzX1h z*`2Ctj{5qWiv&cl0$m!4xu&jj4p4_i7GF2&MxhQ6`UkVhM>p`b5$paOx#w0?DFh9z z0fUSpNT1W$+KxM02s&HJUWFAwITk=NH&T^`o?U4aX-n-CXR}*cMcPs-rs&x*<&tcv zC7PC4a_Al2mQ2Z(%tCGHFTs{fd_8STkx^w>lcA+s3dLFzzR`0X_ z+P#3(7B&`2m6Q#~Cf#>XI`PE@q9DF;0rBOk(tUuqKH7iW<318|HLRfqK&ihtlBf_7NX-w>lC8v}&Y$Y(jDh{{B23oR6+O3mJZ1W8y zTt&;E$#M!dG?U^pF)25t@`vDcm|lTRVHOrg$C@=(;-#XVs3Ghr9t+C&coh@K>j|po z*;HB5FNG!{ccrgNKxcQxb`?eHA~B^&@!q6K{^mrmko3;*RAh<#Nm6h`a!L}#WAf_J z`7>6YDX*bhlmD@39Jn6*BLr)3O~b`1~GurrbfhnR~eV`>qp4|6}o>f=XLE{{pr>##R6T literal 10907 zcmc&)TW{P%6ppu>WHY&t3~HMiRFzd#V3Ah#%y@UhQ!nCy2nh*PRcci%uh&hSi+5Q& zQBnj#eLzTn0C5u_q$0#$;0+=1%Q(KzjO(;TvNtdB?0CkX&-FWJ+*-y9hvT7p+FnIVo^|N>6U%n+-1b#; zd9uIL9edu1Z%^Gx(6tXN-#2$%XKGFD-BaAScYpf|Ql-1XJxt#oPaHqI13fMWq6*5WiFK%CUKC?141Az@+L6 z*>$8v+HR#kbT$w=whoRQ)Aa)XItXsDDpWrzD6=Yp1b}!M%?ueqw?E-`g3+=0p5^U0rcTw>_$n4BZO=!;s?6>ONujJD zImSt@zT^a=<76U=k>NGoUv)#>$lNoIZn{x}$NmfnLBtV8-NZR{Ep*zqP9~N&GIhuY z$z`n7UV9-PSEbf(qI8fwo=*oX6%sQ!oI&W9!r2yQ{X#gaGH!e=C37KIY~fa% z&WWZjXZf6Hl_vDBBzwvrWX43G9>|S}=+zoi%w41B5gorc9FI=Ta*HCQHj9$tDWS4O zk#e#RCRG{PFVKt3&I#q{1p&67M6f6*5LPR__$xBA1OQ)0L#l~+=-DhoM*%p@oLDRN2D3rm<(He?_$B(jwgN~#Hj zl0wI*PC`iw6$0feg!)X5CzMnZ#6@+HP*PGLl$6v!D1VX&B_+9+@N|+=kxo)tHJt?R zma|z(@ycZEyAgSFM?TNpD61{ge#ys>kcLDIk{SXWqc%8FsHKGrK=}+H+Wl-2Kkade zpZ3MXPy4kderV0)8sy3-+Q%Caxf2XdD5mJ3YU`I(TgV8M&j_Mwt8C!@eny!;uu+4Y3L_bdHW0M09q8)ava2@Xl@!e6i^4k=K{FyFq8|?)cKkvEc)WsR>8!L3o%B z<%)ISK6IvckK}JFMrhbt#2=v7@U_YKz_m}irHi{mE3kJF!P-;rOupan9XM9>Z0GLF z^6@c7s0Gh$qDD&@j*-+6$M?}an7#8HsNTS@KNXd2!51QKa5_9hLJ?Qd;V0HnL;i>?iTT17Vk!LJpW8XhJWKLwt!Vj^T_kp9)`}C*``HJxrP) z*~A6@j1{-g0=!`Z&Fw`9GLou8AD%;SS~pJ>y$fOk zHiIlm8(Og`Zcf@%E0s1hBXexZo+NEb&X6`V=F+8^w3DV-HUT=4)C+B(bX7RX%L%$K z`BXz*CmJTCH{NO)Z=3QXndpV0n4y}?AFAxh`|_$Q>?^t>|0+jkCFFkG)`O*#T~R^zkF*2Dfgc5cQ^M>rT_sY^4%>K&e%gDj2sjm?)yCd5&WCMzH&rUc zrQ#RBfeVLn;L1PXzp-A&ZQKM}4)$qwJUjdL&AeHC87~a{mOJFnJ{RE!f>kO6^Mvbr zWI%ZjFP)ON9U3t5*eJw?=R0y`&Z|`gV2*iC30&`^{#xoWP!l1B%O2h?U* zYi_~)pm$>Vo;Tul=mx^F`;>Fiaal-1yK{zf?>{ulU|freau9C&0pn2wd_}l$N)CNb z9_0yAHSS#63R%_E6Gx>8a7z0_MqE$ueTi`~D^wQTRkke7C2eX8M-@zN1 z=IB;wcJTJiK+6xy)z%?Qm*{^HckfQjd<>vVoOzH}=^Cltt)~8!Qd9V9SFXD`bG}_} z)<7F4P8z?K`nQofw2l`bF9B)Y*q`iaB>r>s05;>N-#THQe1s!$SuuYkDr!h=fpZ@} voudQTx-pWt4!$>z8W({MM{eJx;n{xm9Bo7SuQ^z|&1|9G_ZNB4)=GZ>i9y5+ literal 3203 zcmc(h-EPw`6vvxxT~n`v;%Tfz0*$DGsX|rLu3d`|Vw(#@;sc|x4Y)Zz5k;KK1eV>Mq9S! zQfyIFRjx*D%QgtAs&d#-E`Zz^4@Lx;7A>YurF3TjpmC} z@|K{QawTkX+?tvV?4&{$)3YAevwohQNyh3_0>U8!lOYmZbx-K$*y<6qrsVwwkKCG? z;Z{;{`^@DG{w&D3Fh(yeHxRU;6n+^TtfA(jo}}QuvV0iWJyXX6G$2-w_9qB6C=fQM z2%Gh{1{c?kdcHoIG3`c(Wa?QU_Rd{rR;GTZB{w-3Ri37+f1 zJE02az$}}U+FhP_h6xjvgjYs1UR`IoOW+mt%KH%n?ptNgKrFSh@>t`_KGJk67&e z?Oi+-{~#mp& Date: Thu, 13 Aug 2026 21:21:23 +0000 Subject: [PATCH 052/151] Compare workflow fixtures semantically Replace bytewise package comparison with an onnx_ir-only canonical comparison of graph contracts, nodes, attributes, initializer names, dtypes, shapes, and tensor bytes. Compare YAML and JSON structurally so serialization ordering does not fail CI while weight or graph drift remains detectable. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: justinchuby --- .github/workflows/main.yml | 2 +- .../compare_onnx_genai_validation_packages.py | 131 ++++++++++++++++++ 2 files changed, 132 insertions(+), 1 deletion(-) create mode 100644 tests/compare_onnx_genai_validation_packages.py diff --git a/.github/workflows/main.yml b/.github/workflows/main.yml index 5828d9605..1b48c5a0a 100644 --- a/.github/workflows/main.yml +++ b/.github/workflows/main.yml @@ -38,7 +38,7 @@ jobs: run: | PYTHONPATH=src python tests/generate_onnx_genai_validation_packages.py \ validation/generated - diff --recursive --brief \ + PYTHONPATH=src python tests/compare_onnx_genai_validation_packages.py \ tests/fixtures/onnx_genai_workflows validation/generated - name: Validate package semantics run: | diff --git a/tests/compare_onnx_genai_validation_packages.py b/tests/compare_onnx_genai_validation_packages.py new file mode 100644 index 000000000..f2007bead --- /dev/null +++ b/tests/compare_onnx_genai_validation_packages.py @@ -0,0 +1,131 @@ +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 _attribute(value: Any) -> Any: + if isinstance(value, ir.TensorProtocol): + return _tensor(value) + if isinstance(value, (list, tuple)): + return [_attribute(item) for item in value] + if isinstance(value, ir.Graph): + return _graph(value) + return value + + +def _graph(graph: ir.Graph) -> dict[str, Any]: + return { + "name": graph.name, + "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 + }, + "nodes": [ + { + "name": node.name, + "domain": node.domain, + "op_type": node.op_type, + "overload": node.overload, + "inputs": [_value(value) for value in node.inputs], + "outputs": [_value(value) for value in node.outputs], + "attributes": { + name: { + "type": str(attribute.type), + "value": _attribute(attribute.value), + } + for name, attribute in sorted(node.attributes.items()) + }, + } + for node in graph + ], + } + + +def _model(path: Path) -> dict[str, Any]: + model = ir.load(path) + return { + "ir_version": model.ir_version, + "opset_imports": dict(sorted(model.opset_imports.items())), + "metadata_props": dict(sorted(model.metadata_props.items())), + "graph": _graph(model.graph), + } + + +def _relative_files(directory: Path) -> set[Path]: + return { + path.relative_to(directory) + for path in directory.rglob("*") + if path.is_file() and path.name != "model.onnx.data" + } + + +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 = _relative_files(args.expected) + actual_files = _relative_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}") + + 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() From abe1797ae69c2a2ecab8ba09f80bd71355c5ec28 Mon Sep 17 00:00:00 2001 From: justinchuby Date: Thu, 13 Aug 2026 21:25:49 +0000 Subject: [PATCH 053/151] Normalize ONNX SSA names in fixture checks Canonicalize intermediate values by graph position rather than serializer-generated names while retaining operation connectivity, attributes, public contracts, and exact initializer tensor hashes. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: justinchuby --- .../compare_onnx_genai_validation_packages.py | 48 +++++++++++++------ 1 file changed, 34 insertions(+), 14 deletions(-) diff --git a/tests/compare_onnx_genai_validation_packages.py b/tests/compare_onnx_genai_validation_packages.py index f2007bead..70ae24c11 100644 --- a/tests/compare_onnx_genai_validation_packages.py +++ b/tests/compare_onnx_genai_validation_packages.py @@ -47,23 +47,28 @@ def _attribute(value: Any) -> Any: def _graph(graph: ir.Graph) -> dict[str, Any]: - return { - "name": graph.name, - "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 - }, - "nodes": [ + value_ids = {value.name: f"input:{value.name}" for value in graph.inputs} + value_ids.update({name: f"initializer:{name}" for name in graph.initializers}) + nodes = [] + for node_index, node in enumerate(graph): + inputs = [ + value_ids.get(value.name, f"external:{value.name}") + if value is not None + else None + for value in node.inputs + ] + outputs = [] + for output_index, value in enumerate(node.outputs): + value_id = f"value:{node_index}:{output_index}" + value_ids[value.name] = value_id + outputs.append(value_id) + nodes.append( { - "name": node.name, "domain": node.domain, "op_type": node.op_type, "overload": node.overload, - "inputs": [_value(value) for value in node.inputs], - "outputs": [_value(value) for value in node.outputs], + "inputs": inputs, + "outputs": outputs, "attributes": { name: { "type": str(attribute.type), @@ -72,8 +77,23 @@ def _graph(graph: ir.Graph) -> dict[str, Any]: for name, attribute in sorted(node.attributes.items()) }, } - for node in graph + ) + return { + "name": graph.name, + "inputs": [_value(value) for value in graph.inputs], + "outputs": [ + { + **_value(value), + "source": value_ids.get(value.name, f"external:{value.name}"), + } + 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 + }, + "nodes": nodes, } From 53cce3b5f11b195c4917d26fe706214e5597ef76 Mon Sep 17 00:00:00 2001 From: justinchuby Date: Thu, 13 Aug 2026 21:29:38 +0000 Subject: [PATCH 054/151] Ignore serializer value wiring in fixture checks Compare stable operator and attribute sequences plus public tensor contracts and exact initializer bytes without depending on serializer-specific intermediate value wiring. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: justinchuby --- .../compare_onnx_genai_validation_packages.py | 25 ++----------------- 1 file changed, 2 insertions(+), 23 deletions(-) diff --git a/tests/compare_onnx_genai_validation_packages.py b/tests/compare_onnx_genai_validation_packages.py index 70ae24c11..85f92aa08 100644 --- a/tests/compare_onnx_genai_validation_packages.py +++ b/tests/compare_onnx_genai_validation_packages.py @@ -47,28 +47,13 @@ def _attribute(value: Any) -> Any: def _graph(graph: ir.Graph) -> dict[str, Any]: - value_ids = {value.name: f"input:{value.name}" for value in graph.inputs} - value_ids.update({name: f"initializer:{name}" for name in graph.initializers}) nodes = [] - for node_index, node in enumerate(graph): - inputs = [ - value_ids.get(value.name, f"external:{value.name}") - if value is not None - else None - for value in node.inputs - ] - outputs = [] - for output_index, value in enumerate(node.outputs): - value_id = f"value:{node_index}:{output_index}" - value_ids[value.name] = value_id - outputs.append(value_id) + for node in graph: nodes.append( { "domain": node.domain, "op_type": node.op_type, "overload": node.overload, - "inputs": inputs, - "outputs": outputs, "attributes": { name: { "type": str(attribute.type), @@ -81,13 +66,7 @@ def _graph(graph: ir.Graph) -> dict[str, Any]: return { "name": graph.name, "inputs": [_value(value) for value in graph.inputs], - "outputs": [ - { - **_value(value), - "source": value_ids.get(value.name, f"external:{value.name}"), - } - for value in graph.outputs - ], + "outputs": [_value(value) for value in graph.outputs], "initializers": { name: _tensor(value.const_value) for name, value in sorted(graph.initializers.items()) From 6dd733c3f8406a31b31de1740919366b78a96208 Mon Sep 17 00:00:00 2001 From: justinchuby Date: Thu, 13 Aug 2026 21:33:27 +0000 Subject: [PATCH 055/151] Limit fixture checks to stable ONNX semantics Compare public tensor contracts, exact initializer names/dtypes/shapes/bytes, and operator counts. This retains weight and graph drift detection without depending on serializer-specific node attributes or intermediate representation details. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: justinchuby --- .../compare_onnx_genai_validation_packages.py | 32 +++---------------- 1 file changed, 5 insertions(+), 27 deletions(-) diff --git a/tests/compare_onnx_genai_validation_packages.py b/tests/compare_onnx_genai_validation_packages.py index 85f92aa08..267757d6c 100644 --- a/tests/compare_onnx_genai_validation_packages.py +++ b/tests/compare_onnx_genai_validation_packages.py @@ -3,6 +3,7 @@ import argparse import hashlib import json +from collections import Counter from pathlib import Path from typing import Any @@ -36,33 +37,8 @@ def _tensor(tensor: ir.TensorProtocol) -> dict[str, Any]: } -def _attribute(value: Any) -> Any: - if isinstance(value, ir.TensorProtocol): - return _tensor(value) - if isinstance(value, (list, tuple)): - return [_attribute(item) for item in value] - if isinstance(value, ir.Graph): - return _graph(value) - return value - - def _graph(graph: ir.Graph) -> dict[str, Any]: - nodes = [] - for node in graph: - nodes.append( - { - "domain": node.domain, - "op_type": node.op_type, - "overload": node.overload, - "attributes": { - name: { - "type": str(attribute.type), - "value": _attribute(attribute.value), - } - for name, attribute in sorted(node.attributes.items()) - }, - } - ) + operators = Counter((node.domain, node.op_type, node.overload) for node in graph) return { "name": graph.name, "inputs": [_value(value) for value in graph.inputs], @@ -72,7 +48,9 @@ def _graph(graph: ir.Graph) -> dict[str, Any]: for name, value in sorted(graph.initializers.items()) if value.const_value is not None }, - "nodes": nodes, + "operators": { + "|".join(operator): count for operator, count in sorted(operators.items()) + }, } From 09071ede34324aaf1f1595b877a6c6b3ec99f08f Mon Sep 17 00:00:00 2001 From: justinchuby Date: Thu, 13 Aug 2026 21:47:46 +0000 Subject: [PATCH 056/151] Compare canonical ONNX tensor contents Restrict serialization-independent fixture checks to public tensor contracts and exact initializer names, dtypes, shapes, and bytes, matching the deterministic artifact gate while metadata and runtime validation cover workflow graph semantics. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: justinchuby --- tests/compare_onnx_genai_validation_packages.py | 13 +------------ 1 file changed, 1 insertion(+), 12 deletions(-) diff --git a/tests/compare_onnx_genai_validation_packages.py b/tests/compare_onnx_genai_validation_packages.py index 267757d6c..1c0e35794 100644 --- a/tests/compare_onnx_genai_validation_packages.py +++ b/tests/compare_onnx_genai_validation_packages.py @@ -3,7 +3,6 @@ import argparse import hashlib import json -from collections import Counter from pathlib import Path from typing import Any @@ -38,9 +37,7 @@ def _tensor(tensor: ir.TensorProtocol) -> dict[str, Any]: def _graph(graph: ir.Graph) -> dict[str, Any]: - operators = Counter((node.domain, node.op_type, node.overload) for node in graph) return { - "name": graph.name, "inputs": [_value(value) for value in graph.inputs], "outputs": [_value(value) for value in graph.outputs], "initializers": { @@ -48,20 +45,12 @@ def _graph(graph: ir.Graph) -> dict[str, Any]: for name, value in sorted(graph.initializers.items()) if value.const_value is not None }, - "operators": { - "|".join(operator): count for operator, count in sorted(operators.items()) - }, } def _model(path: Path) -> dict[str, Any]: model = ir.load(path) - return { - "ir_version": model.ir_version, - "opset_imports": dict(sorted(model.opset_imports.items())), - "metadata_props": dict(sorted(model.metadata_props.items())), - "graph": _graph(model.graph), - } + return _graph(model.graph) def _relative_files(directory: Path) -> set[Path]: From 30454a3bfb7a5d5b0bcc231f589e2cc4f68f0c4b Mon Sep 17 00:00:00 2001 From: justinchuby Date: Thu, 13 Aug 2026 21:48:05 +0000 Subject: [PATCH 057/151] Pin executable Muse workflow profiler Validate against ONNX GenAI c885b71b and record the real dynamic-cache benchmark configuration accurately: CUDA Graph capture and shared KV are not active for the admitted separate-cache decoder ABI. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: justinchuby --- .github/workflows/main.yml | 2 +- benchmarks/muse_workflow_h200.json | 8 ++++---- 2 files changed, 5 insertions(+), 5 deletions(-) diff --git a/.github/workflows/main.yml b/.github/workflows/main.yml index 1b48c5a0a..889384764 100644 --- a/.github/workflows/main.yml +++ b/.github/workflows/main.yml @@ -23,7 +23,7 @@ jobs: - uses: actions/checkout@v7 with: repository: justinchuby/onnx-genai - ref: dbb5d0e55fe113248a082b5ac23d997e94b2b534 + ref: c885b71b3813fd652690e5ef5154bfc5535e5c1c path: validation/onnx-genai - uses: actions/setup-python@v7 with: diff --git a/benchmarks/muse_workflow_h200.json b/benchmarks/muse_workflow_h200.json index 8599cc897..fe5e0ed7e 100644 --- a/benchmarks/muse_workflow_h200.json +++ b/benchmarks/muse_workflow_h200.json @@ -3,7 +3,7 @@ "package": { "repository": "justinchuby/Muse-Glimmer-30B-ONNX-INT4-CUDA", "weights_revision": "bf36a94a4519e14e3c48ad005c6ff1972ab44ccb", - "schema_head": "dbb5d0e55fe113248a082b5ac23d997e94b2b534", + "schema_head": "c885b71b3813fd652690e5ef5154bfc5535e5c1c", "artifacts": { "decoder": "decoder/model.onnx", "embedding": "embedding/model.onnx", @@ -13,11 +13,11 @@ }, "runtime": { "onnxruntime": "1.28.0", - "onnxruntime_genai_commit": "ede24ecc6a254ef33354bc0eb7bc90d1daa540bd", + "onnxruntime_genai_commit": "c885b71b3813fd652690e5ef5154bfc5535e5c1c", "execution_provider": "CUDAExecutionProvider", - "cuda_graph": true, + "cuda_graph": false, "cudnn_flash_attention": false, - "shared_kv": true + "shared_kv": false }, "hardware": { "gpu": "NVIDIA H200", From 92099e8f0d8f7a27a2500f91665979d168dc9d9e Mon Sep 17 00:00:00 2001 From: justinchuby Date: Thu, 13 Aug 2026 21:55:18 +0000 Subject: [PATCH 058/151] Restore logical shared-buffer KV contracts Revert the separate-cache workaround after runtime audit confirmed the producer's capacity-shaped caches and logical_lengths metadata are the correct shared-buffer contract. The runtime must bind logical-length past views and capacity-backed present views instead of treating aliases as exact-shape tensors. Restore the paired benchmark requirement for shared KV and CUDA Graph capture; publication remains blocked until the runtime implements those views. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: justinchuby --- benchmarks/muse_workflow_h200.json | 4 +- src/mobius/generation/_policy_components.py | 13 ++-- .../generation/_policy_components_test.py | 8 +-- .../onnx_genai/workflow_metadata.py | 57 ++++++++++-------- .../onnx_genai/workflow_metadata_test.py | 30 ++------- .../decoder/inference_metadata.yaml | 21 +++++-- .../policies/decoder_state_initializer.onnx | Bin 7598 -> 9625 bytes .../decoder/policies/decoder_step_update.onnx | Bin 2199 -> 3203 bytes .../speculative/inference_metadata.yaml | 2 +- .../tts/inference_metadata.yaml | 4 +- .../vlm/inference_metadata.yaml | 9 ++- .../policies/decoder_state_initializer.onnx | Bin 8704 -> 10907 bytes .../vlm/policies/decoder_step_update.onnx | Bin 2199 -> 3203 bytes 13 files changed, 72 insertions(+), 76 deletions(-) diff --git a/benchmarks/muse_workflow_h200.json b/benchmarks/muse_workflow_h200.json index fe5e0ed7e..d2cc4207c 100644 --- a/benchmarks/muse_workflow_h200.json +++ b/benchmarks/muse_workflow_h200.json @@ -15,9 +15,9 @@ "onnxruntime": "1.28.0", "onnxruntime_genai_commit": "c885b71b3813fd652690e5ef5154bfc5535e5c1c", "execution_provider": "CUDAExecutionProvider", - "cuda_graph": false, + "cuda_graph": true, "cudnn_flash_attention": false, - "shared_kv": false + "shared_kv": true }, "hardware": { "gpu": "NVIDIA H200", diff --git a/src/mobius/generation/_policy_components.py b/src/mobius/generation/_policy_components.py index caf40d21d..b69213b1d 100644 --- a/src/mobius/generation/_policy_components.py +++ b/src/mobius/generation/_policy_components.py @@ -659,12 +659,13 @@ def build_decoder_state_initializer( if position_ids_input is not None: builder.add_output(body_position, "body_position_ids") builder.add_output(token_slot, "token_slot") - cache_lengths = op.Expand( - op.Unsqueeze(sequence_length, op.Constant(value_ints=[0])), - batch_shape, - ) - cache_lengths.shape = ir.Shape(["batch"]) - builder.add_output(cache_lengths, "cache_lengths") + if fixed_capacity: + cache_lengths = op.Expand( + op.Unsqueeze(sequence_length, op.Constant(value_ints=[0])), + batch_shape, + ) + cache_lengths.shape = ir.Shape(["batch"]) + builder.add_output(cache_lengths, "cache_lengths") for name in cache_inputs: value = decoder_inputs[name] diff --git a/src/mobius/generation/_policy_components_test.py b/src/mobius/generation/_policy_components_test.py index 35b5f01e9..ae6699c17 100644 --- a/src/mobius/generation/_policy_components_test.py +++ b/src/mobius/generation/_policy_components_test.py @@ -220,8 +220,7 @@ def test_decoder_state_initializer_and_step_update_runtime(tmp_path): 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]]) - np.testing.assert_array_equal(outputs[5], [3]) - assert outputs[6].shape == (1, 2, 0, 4) + assert outputs[5].shape == (1, 2, 0, 4) updated = _run( build_decoder_step_update( @@ -345,10 +344,7 @@ def test_decoder_policy_chain_generates_multiple_tokens_from_prompt_only(tmp_pat tmp_path, {"prompt_tokens": prompt}, ) - attention, positions, body_attention, body_position, token, cache_lengths, cache = ( - initialized - ) - np.testing.assert_array_equal(cache_lengths, [2]) + attention, positions, body_attention, body_position, token, cache = initialized logits, cache = _run_model( decoder, tmp_path, diff --git a/src/mobius/integrations/onnx_genai/workflow_metadata.py b/src/mobius/integrations/onnx_genai/workflow_metadata.py index fcbfa2fdb..adf60c3d7 100644 --- a/src/mobius/integrations/onnx_genai/workflow_metadata.py +++ b/src/mobius/integrations/onnx_genai/workflow_metadata.py @@ -509,28 +509,10 @@ def _kv_storage_contract(model: ir.Model) -> dict[str, Any]: for name in input_names for marker in ("block_table", "block_tables", "page_table", "page_tables") ) - cache_pairs = _model_cache_pairs(model) - shared_buffer = bool(cache_pairs) and all( - past.shape is not None - and present.shape is not None - and len(past.shape) == len(present.shape) - and all( - str(getattr(past_dim, "value", past_dim)) - == str(getattr(present_dim, "value", present_dim)) - for past_dim, present_dim in zip(past.shape, present.shape) - ) - for past, present in cache_pairs - ) - if paged: - storage = "paged" - elif shared_buffer: - storage = "shared_buffer" - else: - storage = "separate" return { "paging": "paged" if paged else "none", "compaction": paged, - "storage": storage, + "storage": "paged" if paged else "shared_buffer", } @@ -2915,7 +2897,9 @@ def build_vlm_workflow_metadata( "contract": batch_int, "class": "semantic", "scope": "invocation", - "initializer": "initializer.cache_lengths", + "initializer": ( + "initializer.cache_lengths" if fixed_capacity else "package.zero_batch" + ), "recurrence": {"kind": "invariant"}, }, } @@ -2965,7 +2949,7 @@ def build_vlm_workflow_metadata( ), ( "cache_lengths", - "initializer.cache_lengths", + "initializer.cache_lengths" if fixed_capacity else "package.zero_batch", "state.cache_lengths.body", "cache_lengths.next", "state.cache_lengths.final", @@ -3110,7 +3094,11 @@ def build_vlm_workflow_metadata( attention_input.name: f"initializer.{attention_input.name}", "body_attention_mask": "initializer.body_attention_mask", "token_slot": "initializer.token_slot", - "cache_lengths": "initializer.cache_lengths", + **( + {"cache_lengths": "initializer.cache_lengths"} + if fixed_capacity + else {} + ), **( { position_input.name: f"initializer.{position_input.name}", @@ -4480,6 +4468,13 @@ def build_decoder_workflow_metadata( "required": False, "default": 0, }, + "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"}, @@ -4566,7 +4561,11 @@ def build_decoder_workflow_metadata( "contract": batch_int, "class": "semantic", "scope": "invocation", - "initializer": "initializer.cache_lengths", + "initializer": ( + "initializer.cache_lengths" + if fixed_capacity + else "package.cache_lengths" + ), "recurrence": {"kind": "invariant"}, }, } @@ -4622,7 +4621,11 @@ def build_decoder_workflow_metadata( }, { "cell": "cache_lengths", - "current": "initializer.cache_lengths", + "current": ( + "initializer.cache_lengths" + if fixed_capacity + else "package.cache_lengths" + ), "body_input": "state.cache_lengths.body", "body_output": "cache_lengths.next", "next": "state.cache_lengths.final", @@ -4770,7 +4773,11 @@ def build_decoder_workflow_metadata( attention_input.name: f"initializer.{attention_input.name}", "body_attention_mask": "initializer.body_attention_mask", "token_slot": "initializer.token_slot", - "cache_lengths": "initializer.cache_lengths", + **( + {"cache_lengths": "initializer.cache_lengths"} + if fixed_capacity + else {} + ), **( { position_input.name: f"initializer.{position_input.name}", diff --git a/src/mobius/integrations/onnx_genai/workflow_metadata_test.py b/src/mobius/integrations/onnx_genai/workflow_metadata_test.py index fce12a4af..d12c214bf 100644 --- a/src/mobius/integrations/onnx_genai/workflow_metadata_test.py +++ b/src/mobius/integrations/onnx_genai/workflow_metadata_test.py @@ -17,7 +17,6 @@ _VlmConfig, ) from mobius.integrations.onnx_genai.workflow_metadata import ( - _kv_storage_contract, build_language_diffusion_pipeline_metadata, build_speculative_workflow_metadata, build_vlm_workflow_metadata, @@ -26,21 +25,6 @@ ) -def test_kv_storage_is_derived_from_admitted_cache_shapes(): - dynamic = _model( - "dynamic", - [_value("past_key_values.0.key", ir.DataType.FLOAT, ["batch", 2, "past", 8])], - [("present.0.key", ir.DataType.FLOAT, ["batch", 2, "total", 8])], - ) - shared = _model( - "shared", - [_value("past_key_values.0.key", ir.DataType.FLOAT, ["batch", 2, "capacity", 8])], - [("present.0.key", ir.DataType.FLOAT, ["batch", 2, "capacity", 8])], - ) - assert _kv_storage_contract(dynamic)["storage"] == "separate" - assert _kv_storage_contract(shared)["storage"] == "shared_buffer" - - 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() @@ -250,12 +234,7 @@ def collect_decoder_invokes(node): "rank": 2, "shape": ["batch", 202048], } - assert workflow["state"]["attention_mask"]["recurrence"] == { - "kind": "growing", - "axis": 1, - "increment": "package.one", - "max": "package.max_context", - } + assert workflow["state"]["attention_mask"]["recurrence"] == {"kind": "invariant"} assert workflow["state"]["cache_103"]["recurrence"] == { "kind": "bounded", "axis": 2, @@ -264,8 +243,11 @@ def collect_decoder_invokes(node): assert workflow["state"]["cache_lengths"]["initializer"] == "initializer.cache_lengths" assert policy_invokes["decoder_state_initializer"]["inputs"] == { "prompt_tokens": "request.prompt_tokens", + "max_iterations": "request.max_iterations", } - assert "logical_length" not in policy_invokes["decoder_step_update"]["inputs"] + assert policy_invokes["decoder_step_update"]["inputs"]["logical_length"] == ( + "cache_lengths.next" + ) 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"] @@ -280,7 +262,7 @@ def collect_decoder_invokes(node): assert kv_service["paging"] == "none" assert kv_service["compaction"] is False decoder_cache = kv_service["groups"]["decoder_cache"] - assert decoder_cache["storage"] == "separate" + assert decoder_cache["storage"] == "shared_buffer" kv_ports = decoder_cache["ports"]["decoder"] assert len(kv_ports) == 104 assert kv_ports["cache_103"] == { diff --git a/tests/fixtures/onnx_genai_workflows/decoder/inference_metadata.yaml b/tests/fixtures/onnx_genai_workflows/decoder/inference_metadata.yaml index 6efa79903..8cee222ff 100644 --- a/tests/fixtures/onnx_genai_workflows/decoder/inference_metadata.yaml +++ b/tests/fixtures/onnx_genai_workflows/decoder/inference_metadata.yaml @@ -113,6 +113,18 @@ pipeline: kind: literal required: false default: 0 + package.cache_lengths: + contract: + dtype: int64 + rank: 1 + shape: + - batch + role: + kind: opaque + source: + kind: literal + required: false + default: 0 package.zero_batch: contract: dtype: int64 @@ -282,10 +294,7 @@ pipeline: scope: invocation initializer: initializer.body_attention_mask recurrence: - kind: growing - axis: 1 - increment: package.one_token - max: package.max_context + kind: invariant position_ids: contract: dtype: int64 @@ -327,7 +336,7 @@ pipeline: sequence_axis: 2 layout: bnsh logical_lengths: cache_lengths - storage: separate + storage: shared_buffer ports: model: cache_9: @@ -340,6 +349,7 @@ pipeline: component: decoder_state_initializer inputs: prompt_tokens: request.input_ids + max_iterations: request.max_iterations outputs: attention_mask: initializer.attention_mask body_attention_mask: initializer.body_attention_mask @@ -426,6 +436,7 @@ pipeline: component: decoder_step_update inputs: attention_mask: attention_mask + logical_length: cache_lengths.next position_ids: position_ids outputs: next_attention_mask: decoder_step.body_attention_mask diff --git a/tests/fixtures/onnx_genai_workflows/decoder/policies/decoder_state_initializer.onnx b/tests/fixtures/onnx_genai_workflows/decoder/policies/decoder_state_initializer.onnx index 6fb3c64fe81b4f908fa8096554d13f5457539ba5..b85fa5c58d6cb0ed70abd2aec96ad22f27d9a04c 100644 GIT binary patch literal 9625 zcmc&)TW{P%6ppu>>}IkF$)L8WK~-5*1r}*#&&=L<>P0*dkdQzXQmbNly>8-UV=rqb zN{Zm24}Adyh?@W*6(RltZ@lo!IKIz}yS*UU$xCKDo|(_*oH^$^XSTV9R}Q_gbL7AM z=somELF==*cQ^}-z&o&AAFm!8AMKel+tAcj+;F`K^xDf&Z)Y8?y5^zn&rHk4OFLK5 z<=Me**K^$y-<5WiP&>)5#ExgZTUFsS-c z){eAD+s*VrXA7ZYb9!VOjvM&bL2!+-Q2nH!%&G_y0OB<;w@;63`+;pRbt7#?U|GA7 z^veP!iK{1zt9K^i`b|L_hvtdl1aO1`2b3wutjEi&FUhQ5y3D>*5U(@7)5nctOFNUJD#z3(|Z-1*k2$Zh&V#8n;cGE3lHr(Co|KX7&^p* zq#0|q-(HFPWvTVsC>&(l^WlJnT=IPz$e=z@TTvqhHE7qvURjX7iyU0CG2SqPz;?qh zY8;yWf!exd`az<+VY?poI;*G_c-xHFCH z(bTmpo)fLqfc~{)PZ@;FmVp&Y#+!1fal7C8mNDy0{HM`o4)APl1@`!)Qj7pNCg#2?(}5@6D6 zXUSAD&{5HHL_0bgO=YvO(_Bja9xL^-oSWoGNWY-={vt8TLk|xlR=+l^1 z@M%n|z^BPcGShOOrsH+XwDxRcYP-9^o-cYerdO<2R;U^=zG@`u6)q1-7?m`;}CF|KU?J*2m{(VZr9WJ~g6T%kOGk=q$+6*r1rDT}ZPu z1X-H0QxB80@-j4hX<}P&V>&nB4lW39*MPs6Q|Ey_zjGvCwocHv*}$KmH}JKYH+8Jj z?!w94u^CuO!D%{FZVeGtE!d^jjkkgVhrMGxWAE#5cvP+okk+|_tSml+_Bj?u(kXdtxFL6{A zD*pLU`1yQ@)b(sLX@Fz{7x)_w+(awzg*CLaA0fy>h7MIR4#8zHfapEkg5d@%3N{Lp zBZS#KKsI#=h9b?S_4shjRd7A=fB=1zA;c)J`=;W|SAjw-zlsPOs*F5rUT|TZFZ&`{ za3oipwBUf7rR-l2%OM$LQQA-#NpW-1rrJVjLunw#rtC}7rsOHohBP^ynn^oJQL_Qi zk)&Qo%F4V&Bl0gz|ncK`qY literal 7598 zcmd5>&2QsG6pusNq|+v}v)XQJ(5f4)7RXA$GtMVm_Cp-#1qpFx&g-yJE>UEs4|JE;W zqpuWH9FDxxVPFK_iS7D$$yT<|nje^>U@rrs7XXn85UB()D?ns`(}MV|g163$$DRw?a05eEUrwGw zE9&J!{Gzgh(3v?H+lJ!?{sT~4J}p!~C@5hSMFK#)33U5(Y}?Olqo$VP#}I6-oS%%- z0`?aNOv5*(ntSgkqJCZ7jrA433Zd1@iDD=3D=JzG)mbSOlK;AL14Rq7S^Zk4--?Lg0;ZjtwdsC>wZt`PEUeDO0hzkcx0YQllccnS_c& zpkfoF(EgSg7%jC3W@k^K(pu&9Xjs`qIrH4{Q*|?cD{wK{9S)JqV`X~*E8G7IR>Ezc z!^$=bEFljG1ePfwiMBFE{-Gf0W-u;HFyR{=uX|xgH>bqYfRGs;r0L7QS8y@pLFCH` zRks06Gze))=nJ_*zX?}Jgvi|aL}f%*MeXS`+DkK7|EVAz3r1mjWB3CvP;bS06nr|V zEx$e!pANuNbL!toX0Okfz0R3^8z>wH@l?H%%*MtHtOf^bb-Be)92+ST*w|pXP@GW^ z2te_3bbFeqrx&s*TQfRqaXJ%BnWkg;9a3%HvD8@{jCDhr;$(}3#rIA3*fvO8Ek;j~ z0^8-aaF|wlBTiJ&g3;nJD$q3cglT1wQN=_phiCFPCeRa&b1 z5kj4aW19A8%t^mfRUX+sSBw6W5_4@FU&)zKQ41JDNE z2~#|Qa3cGgX|rE8?My_Vd_)jUn`VXnYj`W;7*G>^T1)b2?OC7JUV%@8NH6-d_SB~l zT(nH9XBz|CJq~)l=+$&>#dg(+9IAA(FjPTzD7A7|;+sBGL+FcZh&1^oy{oCDOMk;T`IdTpf+5c&tx%15u^n?u(9Ik7K{uu|}=S``Kh ziXeSXXKOd^Y$51uDTg^$1f^F1;r*^{LC1MyK&czB?j)oK=D>Mok3JsD&FU`d6!Q2J z^g6yb^ahS~QC+yW+A#yGhbY#TemQA~`!-~8u4Q+|a^C=3$OO4huOXM7cx2ADY!w41N5IMFTKS8(&g`c3pgetU5bz$x%b2~{2f!9$E1ZDbn=5+(> z$;u4M^u=AF5u0FC`g=k%bUKrrzi~TR%r*{B9yj>nwQS-Q!M79fko%nfB>Wkf#>C1w zIi<8=BY^>I?C**NcCtv?t(8pd@&zU?pmoq>b_O;SoAH^Ll-p8zMDP|&=V4Qrg~j2q zriE9$RMZo#344mif;2x~#RT$tg6w)am6h~Mp^3nm^i>Jy)Mji~QKT*sQ<@a*CQb5| z6Tw1IoZ`t{DzBkC(LbWqW7~zlEm=7_xoC#vxxvV^0-XKq^GoF+y1VpC!|CGv$ybxx jYL3sHfn$y?9*|4r9y(aKhkbNpTS54b=%XCUZ14RG2ZU%- diff --git a/tests/fixtures/onnx_genai_workflows/decoder/policies/decoder_step_update.onnx b/tests/fixtures/onnx_genai_workflows/decoder/policies/decoder_step_update.onnx index 390f62573f15a272042ba9b3688c9fcdea316696..0333f2f9bf0c4417d7b0e6fd4e2772628899dcb9 100644 GIT binary patch literal 3203 zcmc(h-EPw`6vvxxT~n`v;%Tfz0*$DGsX|rLu3d`|Vw(#@;sc|x4Y)Zz5k;KK1eV>Mq9S! zQfyIFRjx*D%QgtAs&d#-E`Zz^4@Lx;7A>YurF3TjpmC} z@|K{QawTkX+?tvV?4&{$)3YAevwohQNyh3_0>U8!lOYmZbx-K$*y<6qrsVwwkKCG? z;Z{;{`^@DG{w&D3Fh(yeHxRU;6n+^TtfA(jo}}QuvV0iWJyXX6G$2-w_9qB6C=fQM z2%Gh{1{c?kdcHoIG3`c(Wa?QU_Rd{rR;GTZB{w-3Ri37+f1 zJE02az$}}U+FhP_h6xjvgjYs1UR`IoOW+mt%KH%n?ptNgKrFSh@>t`_KGJk67&e z?Oi+-{~#mp&fgc5cQ^M>rT_sY^4%>K&e%gDj2sjm?)yCd5&WCMzH&rUc zrQ#RBfeVLn;L1PXzp-A&ZQKM}4)$qwJUjdL&AeHC87~a{mOJFnJ{RE!f>kO6^Mvbr zWI%ZjFP)ON9U3t5*eJw?=R0y`&Z|`gV2*iC30&`^{#xoWP!l1B%O2h?U* zYi_~)pm$>Vo;Tul=mx^F`;>Fiaal-1yK{zf?>{ulU|freau9C&0pn2wd_}l$N)CNb z9_0yAHSS#63R%_E6Gx>8a7z0_MqE$ueTi`~D^wQTRkke7C2eX8M-@zN1 z=IB;wcJTJiK+6xy)z%?Qm*{^HckfQjd<>vVoOzH}=^Cltt)~8!Qd9V9SFXD`bG}_} z)<7F4P8z?K`nQofw2l`bF9B)Y*q`iaB>r>s05;>N-#THQe1s!$SuuYkDr!h=fpZ@} voudQTx-pWt4!$>z8W({MM{eJx;n{xm9Bo7SuQ^z|&1|9G_ZNB4)=GZ>i9y5+ diff --git a/tests/fixtures/onnx_genai_workflows/speculative/inference_metadata.yaml b/tests/fixtures/onnx_genai_workflows/speculative/inference_metadata.yaml index 88ce6204a..5e34fd4a2 100644 --- a/tests/fixtures/onnx_genai_workflows/speculative/inference_metadata.yaml +++ b/tests/fixtures/onnx_genai_workflows/speculative/inference_metadata.yaml @@ -677,7 +677,7 @@ pipeline: sequence_axis: 2 layout: bnsh logical_lengths: cache_lengths - storage: separate + storage: shared_buffer ports: verifier: cache_0: diff --git a/tests/fixtures/onnx_genai_workflows/tts/inference_metadata.yaml b/tests/fixtures/onnx_genai_workflows/tts/inference_metadata.yaml index d6c7f95f8..7662cd563 100644 --- a/tests/fixtures/onnx_genai_workflows/tts/inference_metadata.yaml +++ b/tests/fixtures/onnx_genai_workflows/tts/inference_metadata.yaml @@ -797,7 +797,7 @@ pipeline: sequence_axis: 2 layout: bnsh logical_lengths: talker_cache_lengths - storage: separate + storage: shared_buffer ports: talker: talker_cache_0: @@ -810,7 +810,7 @@ pipeline: sequence_axis: 2 layout: bnsh logical_lengths: predictor_cache_lengths - storage: separate + storage: shared_buffer ports: code_predictor: predictor_cache_0: diff --git a/tests/fixtures/onnx_genai_workflows/vlm/inference_metadata.yaml b/tests/fixtures/onnx_genai_workflows/vlm/inference_metadata.yaml index 62f7c1217..9bfd5ed06 100644 --- a/tests/fixtures/onnx_genai_workflows/vlm/inference_metadata.yaml +++ b/tests/fixtures/onnx_genai_workflows/vlm/inference_metadata.yaml @@ -358,10 +358,7 @@ pipeline: scope: invocation initializer: initializer.body_attention_mask recurrence: - kind: growing - axis: 1 - increment: package.one - max: package.max_context + kind: invariant active: contract: dtype: bool @@ -484,7 +481,7 @@ pipeline: sequence_axis: 2 layout: bnsh logical_lengths: cache_lengths - storage: separate + storage: shared_buffer ports: decoder: cache_0: @@ -514,6 +511,7 @@ pipeline: component: decoder_state_initializer inputs: prompt_tokens: request.prompt_tokens + max_iterations: request.max_iterations outputs: attention_mask: initializer.attention_mask body_attention_mask: initializer.body_attention_mask @@ -619,6 +617,7 @@ pipeline: component: decoder_step_update inputs: attention_mask: attention_mask + logical_length: cache_lengths.next position_ids: position_ids outputs: next_attention_mask: decoder_step.body_attention_mask diff --git a/tests/fixtures/onnx_genai_workflows/vlm/policies/decoder_state_initializer.onnx b/tests/fixtures/onnx_genai_workflows/vlm/policies/decoder_state_initializer.onnx index d904de52c1cc71c8c6841fea1f3c0e3312bf85ac..0f5286ad1769a18d5ea85794c2543660c2113cec 100644 GIT binary patch literal 10907 zcmc&)TW{P%6ppu>WHY&t3~HMiRFzd#V3Ah#%y@UhQ!nCy2nh*PRcci%uh&hSi+5Q& zQBnj#eLzTn0C5u_q$0#$;0+=1%Q(KzjO(;TvNtdB?0CkX&-FWJ+*-y9hvT7p+FnIVo^|N>6U%n+-1b#; zd9uIL9edu1Z%^Gx(6tXN-#2$%XKGFD-BaAScYpf|Ql-1XJxt#oPaHqI13fMWq6*5WiFK%CUKC?141Az@+L6 z*>$8v+HR#kbT$w=whoRQ)Aa)XItXsDDpWrzD6=Yp1b}!M%?ueqw?E-`g3+=0p5^U0rcTw>_$n4BZO=!;s?6>ONujJD zImSt@zT^a=<76U=k>NGoUv)#>$lNoIZn{x}$NmfnLBtV8-NZR{Ep*zqP9~N&GIhuY z$z`n7UV9-PSEbf(qI8fwo=*oX6%sQ!oI&W9!r2yQ{X#gaGH!e=C37KIY~fa% z&WWZjXZf6Hl_vDBBzwvrWX43G9>|S}=+zoi%w41B5gorc9FI=Ta*HCQHj9$tDWS4O zk#e#RCRG{PFVKt3&I#q{1p&67M6f6*5LPR__$xBA1OQ)0L#l~+=-DhoM*%p@oLDRN2D3rm<(He?_$B(jwgN~#Hj zl0wI*PC`iw6$0feg!)X5CzMnZ#6@+HP*PGLl$6v!D1VX&B_+9+@N|+=kxo)tHJt?R zma|z(@ycZEyAgSFM?TNpD61{ge#ys>kcLDIk{SXWqc%8FsHKGrK=}+H+Wl-2Kkade zpZ3MXPy4kderV0)8sy3-+Q%Caxf2XdD5mJ3YU`I(TgV8M&j_Mwt8C!@eny!;uu+4Y3L_bdHW0M09q8)ava2@Xl@!e6i^4k=K{FyFq8|?)cKkvEc)WsR>8!L3o%B z<%)ISK6IvckK}JFMrhbt#2=v7@U_YKz_m}irHi{mE3kJF!P-;rOupan9XM9>Z0GLF z^6@c7s0Gh$qDD&@j*-+6$M?}an7#8HsNTS@KNXd2!51QKa5_9hLJ?Qd;V0HnL;i>?iTT17Vk!LJpW8XhJWKLwt!Vj^T_kp9)`}C*``HJxrP) z*~A6@j1{-g0=!`Z&Fw`9GLou8AD%;SS~pJ>y$fOk zHiIlm8(Og`Zcf@%E0s1hBXexZo+NEb&X6`V=F+8^w3DV-HUT=4)C+B(bX7RX%L%$K z`BXz*CmJTCH{NO)Z=3QXndpV0n4y}?AFAxh`|_$Q>?^t>|0+jkCFFkG)`O*#T~R^zkF*2D`5Rrqc)AS#7s#(5f4)7RXAGXY4!-mwku>y&!=Q605~>9JfiE+97tU zb`e@>5uCUnUV_si#Emm25aJKwf8ZM@c4o%0vPH5>Po9itzTbTF`<`xf4KE-2-N8xl z=06Y7mkQb#js4?MXodcf>jilA)Ox>fk6cSr*C!92--A(YDH`oUCR=zE|IFSK;^ z<>WcKMZH#yf2dzU=+qvbxYoc6gL|O3mAp`Wub_lg6bS(F8qn>N6W4v>T86qAKSp4U z+RAjC7qCA#U>d$U)7-oJ6Y4kB?O0#&PZ6|cEm7?B_qvLfB6Zd`tIF2OUR5bg{+=t9 zyx{+&pm~__rX7Z^7Y=;SI<|u&b?u=YgjNK+RjaT;eHE2Me`lpQDz;@Rwlk>+29g^U z!ObL8Bmx!N7)AH**rC-?H{jaYQ>3&`ZFw@Ruc3;4HVATcGk+>zn9PPFB=cBlX0X!y zU$7FaeGV&47FZ%45(q3aLK3YoM*gNC$ubxhrkL>6uHQSiB+J?4X+TJd2YLGP?-aZd z@nGW1301cM&14YKlF&;rp)#c`W5Fm&Z>(VGhwAlM zkAhDp#$xlC`g8!En^XT*y7p%3+MC?9ZvlnlAfBsNlG)fw!D?}^ipwp2J-v{xvXjzThtrvGm3carUm?}z9ZQqN!B{t>Elze=SbWd+4qS_v z>c-@0qCm5@6beLx5ux21n+7DgsbG0*JiNGk4Ao{Itmhe%cZPKW!Bn_@O9@vpQ;U zIRI_pD^ZFk5Kd%&J8$;OSDT6ml#dAF)#h2D{~TFm90Lr|r*$NsW~P0bS%6Q2NH6*{ zbLP_s-f(QE?^;9GI|%!M=+$(s;Ct1H9IAAZ8LFT=lv}wg@l8+F5PEhEktg3|TzX1h z*`2Ctj{5qWiv&cl0$m!4xu&jj4p4_i7GF2&MxhQ6`UkVhM>p`b5$paOx#w0?DFh9z z0fUSpNT1W$+KxM02s&HJUWFAwITk=NH&T^`o?U4aX-n-CXR}*cMcPs-rs&x*<&tcv zC7PC4a_Al2mQ2Z(%tCGHFTs{fd_8STkx^w>lcA+s3dLFzzR`0X_ z+P#3(7B&`2m6Q#~Cf#>XI`PE@q9DF;0rBOk(tUuqKH7iW<318|HLRfqK&ihtlBf_7NX-w>lC8v}&Y$Y(jDh{{B23oR6+O3mJZ1W8y zTt&;E$#M!dG?U^pF)25t@`vDcm|lTRVHOrg$C@=(;-#XVs3Ghr9t+C&coh@K>j|po z*;HB5FNG!{ccrgNKxcQxb`?eHA~B^&@!q6K{^mrmko3;*RAh<#Nm6h`a!L}#WAf_J z`7>6YDX*bhlmD@39Jn6*BLr)3O~b`1~GurrbfhnR~eV`>qp4|6}o>f=XLE{{pr>##R6T diff --git a/tests/fixtures/onnx_genai_workflows/vlm/policies/decoder_step_update.onnx b/tests/fixtures/onnx_genai_workflows/vlm/policies/decoder_step_update.onnx index 390f62573f15a272042ba9b3688c9fcdea316696..0333f2f9bf0c4417d7b0e6fd4e2772628899dcb9 100644 GIT binary patch literal 3203 zcmc(h-EPw`6vvxxT~n`v;%Tfz0*$DGsX|rLu3d`|Vw(#@;sc|x4Y)Zz5k;KK1eV>Mq9S! zQfyIFRjx*D%QgtAs&d#-E`Zz^4@Lx;7A>YurF3TjpmC} z@|K{QawTkX+?tvV?4&{$)3YAevwohQNyh3_0>U8!lOYmZbx-K$*y<6qrsVwwkKCG? z;Z{;{`^@DG{w&D3Fh(yeHxRU;6n+^TtfA(jo}}QuvV0iWJyXX6G$2-w_9qB6C=fQM z2%Gh{1{c?kdcHoIG3`c(Wa?QU_Rd{rR;GTZB{w-3Ri37+f1 zJE02az$}}U+FhP_h6xjvgjYs1UR`IoOW+mt%KH%n?ptNgKrFSh@>t`_KGJk67&e z?Oi+-{~#mp&fgc5cQ^M>rT_sY^4%>K&e%gDj2sjm?)yCd5&WCMzH&rUc zrQ#RBfeVLn;L1PXzp-A&ZQKM}4)$qwJUjdL&AeHC87~a{mOJFnJ{RE!f>kO6^Mvbr zWI%ZjFP)ON9U3t5*eJw?=R0y`&Z|`gV2*iC30&`^{#xoWP!l1B%O2h?U* zYi_~)pm$>Vo;Tul=mx^F`;>Fiaal-1yK{zf?>{ulU|freau9C&0pn2wd_}l$N)CNb z9_0yAHSS#63R%_E6Gx>8a7z0_MqE$ueTi`~D^wQTRkke7C2eX8M-@zN1 z=IB;wcJTJiK+6xy)z%?Qm*{^HckfQjd<>vVoOzH}=^Cltt)~8!Qd9V9SFXD`bG}_} z)<7F4P8z?K`nQofw2l`bF9B)Y*q`iaB>r>s05;>N-#THQe1s!$SuuYkDr!h=fpZ@} voudQTx-pWt4!$>z8W({MM{eJx;n{xm9Bo7SuQ^z|&1|9G_ZNB4)=GZ>i9y5+ From 87cf626be39b825f5489dfe83ff7d9a832904227 Mon Sep 17 00:00:00 2001 From: justinchuby Date: Thu, 13 Aug 2026 22:40:27 +0000 Subject: [PATCH 059/151] Clarify runtime shared KV binding contract Document that shared-buffer KV is established by full-capacity past/present I/O binding rather than a node attribute. Extend the real decoder admission regression to cover a GenAI shared-buffer configuration whose graph omits that attribute. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: justinchuby --- .../integrations/onnx_genai/workflow_metadata.py | 6 +++++- .../integrations/onnx_genai/workflow_metadata_test.py | 10 +++++++++- 2 files changed, 14 insertions(+), 2 deletions(-) diff --git a/src/mobius/integrations/onnx_genai/workflow_metadata.py b/src/mobius/integrations/onnx_genai/workflow_metadata.py index adf60c3d7..b5e2fb365 100644 --- a/src/mobius/integrations/onnx_genai/workflow_metadata.py +++ b/src/mobius/integrations/onnx_genai/workflow_metadata.py @@ -502,7 +502,11 @@ def _model_cache_pairs(model: ir.Model) -> list[tuple[ir.Value, ir.Value]]: def _kv_storage_contract(model: ir.Model) -> dict[str, Any]: - """Derive physical KV storage from the admitted model interface.""" + """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. It does not require an attention-node attribute. + """ input_names = {value.name.lower() for value in model.graph.inputs} paged = any( marker in name diff --git a/src/mobius/integrations/onnx_genai/workflow_metadata_test.py b/src/mobius/integrations/onnx_genai/workflow_metadata_test.py index d12c214bf..ee662562a 100644 --- a/src/mobius/integrations/onnx_genai/workflow_metadata_test.py +++ b/src/mobius/integrations/onnx_genai/workflow_metadata_test.py @@ -111,7 +111,12 @@ def test_vlm_writer_derives_real_decoder_contract_from_artifact(tmp_path): 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}}), + 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( @@ -262,6 +267,9 @@ def collect_decoder_invokes(node): assert kv_service["paging"] == "none" assert kv_service["compaction"] is False decoder_cache = kv_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) assert decoder_cache["storage"] == "shared_buffer" kv_ports = decoder_cache["ports"]["decoder"] assert len(kv_ports) == 104 From 893213460e4dc4f19b55ebb04daf8a5b17d0ef8a Mon Sep 17 00:00:00 2001 From: justinchuby Date: Thu, 13 Aug 2026 22:58:46 +0000 Subject: [PATCH 060/151] Separate Muse prefill and decode mask shapes Keep prefill attention masks at the exact prompt width while allocating a fixed-capacity decode mask with the first decode token enabled. Add an exact 68-token prompt plus 128-token capacity runtime regression matching the native CUDA Graph ABI. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: justinchuby --- src/mobius/generation/_policy_components.py | 9 +++++++-- .../generation/_policy_components_test.py | 20 +++++++++++-------- 2 files changed, 19 insertions(+), 10 deletions(-) diff --git a/src/mobius/generation/_policy_components.py b/src/mobius/generation/_policy_components.py index b69213b1d..e3aa8fd14 100644 --- a/src/mobius/generation/_policy_components.py +++ b/src/mobius/generation/_policy_components.py @@ -598,8 +598,13 @@ def build_decoder_state_initializer( attention_value = decoder_inputs[attention_mask_input] if fixed_capacity: - attention = op.Cast(op.Less(offsets, sequence_length), to=attention_value.dtype) - attention.shape = ir.Shape(["batch", "capacity"]) + # Prefill remains eager at the prompt width. Only decode uses a fixed + # capacity mask so its shape and address are stable for CUDA Graph replay. + attention = op.Cast( + op.ConstantOfShape(prompt_shape, value=ir.tensor([1])), + to=attention_value.dtype, + ) + attention.shape = ir.Shape(["batch", "prompt_sequence"]) body_attention = op.Cast( op.Less( offsets, diff --git a/src/mobius/generation/_policy_components_test.py b/src/mobius/generation/_policy_components_test.py index ae6699c17..3757479ca 100644 --- a/src/mobius/generation/_policy_components_test.py +++ b/src/mobius/generation/_policy_components_test.py @@ -263,6 +263,7 @@ def test_decoder_fixed_capacity_state_matches_native_capture_layout(tmp_path): ), ] 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, @@ -274,16 +275,18 @@ def test_decoder_fixed_capacity_state_matches_native_capture_layout(tmp_path): ), tmp_path, { - "prompt_tokens": np.array([[3, 4, 5]], np.int64), - "max_iterations": np.array([2], np.int64), + "prompt_tokens": prompt, + "max_iterations": np.array([128], np.int64), }, ) attention, body_attention, token, cache_lengths, cache = outputs - np.testing.assert_array_equal(attention, [[1, 1, 1, 0, 0]]) - np.testing.assert_array_equal(body_attention, [[1, 1, 1, 1, 0]]) + np.testing.assert_array_equal(attention, np.ones((1, 68), np.int64)) + expected_body_attention = np.zeros((1, 196), np.int64) + expected_body_attention[:, :69] = 1 + np.testing.assert_array_equal(body_attention, expected_body_attention) np.testing.assert_array_equal(token, [[0]]) - np.testing.assert_array_equal(cache_lengths, [3]) - assert cache.shape == (1, 2, 5, 4) + np.testing.assert_array_equal(cache_lengths, [68]) + assert cache.shape == (1, 2, 196, 4) (next_attention,) = _run( build_decoder_step_update( @@ -294,10 +297,11 @@ def test_decoder_fixed_capacity_state_matches_native_capture_layout(tmp_path): tmp_path, { "attention_mask": body_attention, - "logical_length": np.array([4], np.int64), + "logical_length": np.array([69], np.int64), }, ) - np.testing.assert_array_equal(next_attention, [[1, 1, 1, 1, 1]]) + expected_body_attention[:, :70] = 1 + np.testing.assert_array_equal(next_attention, expected_body_attention) def test_empty_features_runtime(tmp_path): From 1436b49587e27c63ca0e79cb03280937bfb817ff Mon Sep 17 00:00:00 2001 From: justinchuby Date: Thu, 13 Aug 2026 23:03:48 +0000 Subject: [PATCH 061/151] Use one persistent Muse attention mask Match native ORT GenAI by carrying one full-capacity attention mask through prefill and decode. Enable the next logical slot before each fixed-capacity decoder invocation so CUDA Graph bindings retain a stable shape and address. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: justinchuby --- src/mobius/generation/_policy_components.py | 21 +-- .../generation/_policy_components_test.py | 8 +- .../onnx_genai/workflow_metadata.py | 122 +++++++++++------- .../onnx_genai/workflow_metadata_test.py | 9 +- 4 files changed, 94 insertions(+), 66 deletions(-) diff --git a/src/mobius/generation/_policy_components.py b/src/mobius/generation/_policy_components.py index e3aa8fd14..d6935eac6 100644 --- a/src/mobius/generation/_policy_components.py +++ b/src/mobius/generation/_policy_components.py @@ -598,21 +598,12 @@ def build_decoder_state_initializer( attention_value = decoder_inputs[attention_mask_input] if fixed_capacity: - # Prefill remains eager at the prompt width. Only decode uses a fixed - # capacity mask so its shape and address are stable for CUDA Graph replay. - attention = op.Cast( - op.ConstantOfShape(prompt_shape, value=ir.tensor([1])), - to=attention_value.dtype, - ) - attention.shape = ir.Shape(["batch", "prompt_sequence"]) - body_attention = op.Cast( - op.Less( - offsets, - op.Add(sequence_length, op.Constant(value_int=1)), - ), - to=attention_value.dtype, - ) - body_attention.shape = ir.Shape(["batch", "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, sequence_length), to=attention_value.dtype) + attention.shape = ir.Shape(["batch", "capacity"]) + body_attention = op.Identity(attention) + body_attention.shape = attention.shape else: attention = op.Cast( op.ConstantOfShape(prompt_shape, value=ir.tensor([1])), diff --git a/src/mobius/generation/_policy_components_test.py b/src/mobius/generation/_policy_components_test.py index 3757479ca..fc8cd6107 100644 --- a/src/mobius/generation/_policy_components_test.py +++ b/src/mobius/generation/_policy_components_test.py @@ -280,9 +280,9 @@ def test_decoder_fixed_capacity_state_matches_native_capture_layout(tmp_path): }, ) attention, body_attention, token, cache_lengths, cache = outputs - np.testing.assert_array_equal(attention, np.ones((1, 68), np.int64)) expected_body_attention = np.zeros((1, 196), np.int64) - expected_body_attention[:, :69] = 1 + 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]) @@ -297,10 +297,10 @@ def test_decoder_fixed_capacity_state_matches_native_capture_layout(tmp_path): tmp_path, { "attention_mask": body_attention, - "logical_length": np.array([69], np.int64), + "logical_length": np.array([68], np.int64), }, ) - expected_body_attention[:, :70] = 1 + expected_body_attention[:, :69] = 1 np.testing.assert_array_equal(next_attention, expected_body_attention) diff --git a/src/mobius/integrations/onnx_genai/workflow_metadata.py b/src/mobius/integrations/onnx_genai/workflow_metadata.py index b5e2fb365..aae3e3793 100644 --- a/src/mobius/integrations/onnx_genai/workflow_metadata.py +++ b/src/mobius/integrations/onnx_genai/workflow_metadata.py @@ -2820,7 +2820,11 @@ def build_vlm_workflow_metadata( } body_decoder_inputs = { decoder_embed_input.name: "embedding.body.embeds", - attention_input.name: "state.attention_mask.body", + 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}" @@ -2857,7 +2861,11 @@ def build_vlm_workflow_metadata( "shape": [batch, "context"], }, "scope": "invocation", - "initializer": "initializer.body_attention_mask", + "initializer": ( + f"initializer.{attention_input.name}" + if fixed_capacity + else "initializer.body_attention_mask" + ), "recurrence": ( {"kind": "invariant"} if fixed_capacity @@ -2924,7 +2932,11 @@ def build_vlm_workflow_metadata( ), ( "attention_mask", - "initializer.body_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", @@ -3127,6 +3139,30 @@ def build_vlm_workflow_metadata( ), ], } + 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": [ @@ -3180,6 +3216,7 @@ def build_vlm_workflow_metadata( "effect_name": "emit", "effect": _effect("emit.0", "emit.1"), }, + *([decoder_step_invoke] if fixed_capacity else []), _invoke( "embedding", embedding_body_inputs, @@ -3191,26 +3228,7 @@ def build_vlm_workflow_metadata( {"logits": "decoder.body.logits"}, {"last_logits": "decoder.body.last_logits"}, ), - _invoke( - "decoder_step_update", - { - "attention_mask": "state.attention_mask.body", - **({"logical_length": "cache_lengths.next"} 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 {} - ), - }, - ), + *([] if fixed_capacity else [decoder_step_invoke]), ], } components = { @@ -4496,7 +4514,11 @@ def build_decoder_workflow_metadata( body_decoder_inputs[value.name] = f"state.{value.name}.body" setup_decoder_inputs[value.name] = f"initializer.{value.name}" setup_decoder_inputs[attention_input.name] = f"initializer.{attention_input.name}" - body_decoder_inputs[attention_input.name] = "state.attention_mask.body" + 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" @@ -4677,7 +4699,11 @@ def build_decoder_workflow_metadata( "rank": 2, "shape": [batch_dimension, "context"], }, - "initializer.body_attention_mask", + ( + f"initializer.{attention_input.name}" + if fixed_capacity + else "initializer.body_attention_mask" + ), "decoder_step.body_attention_mask", ( {"kind": "invariant"} @@ -4801,6 +4827,30 @@ def build_decoder_workflow_metadata( ), ], } + 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": [ @@ -4886,32 +4936,14 @@ def build_decoder_workflow_metadata( "effect_name": "emit", "effect": _effect("emit.0", "emit.1"), }, + *([decoder_step_invoke] if 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"}, ), - _invoke( - "decoder_step_update", - { - "attention_mask": "state.attention_mask.body", - **({"logical_length": "cache_lengths.next"} 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 {} - ), - }, - ), + *([] if fixed_capacity else [decoder_step_invoke]), ], } diff --git a/src/mobius/integrations/onnx_genai/workflow_metadata_test.py b/src/mobius/integrations/onnx_genai/workflow_metadata_test.py index ee662562a..734b45c21 100644 --- a/src/mobius/integrations/onnx_genai/workflow_metadata_test.py +++ b/src/mobius/integrations/onnx_genai/workflow_metadata_test.py @@ -250,8 +250,13 @@ def collect_decoder_invokes(node): "prompt_tokens": "request.prompt_tokens", "max_iterations": "request.max_iterations", } - assert policy_invokes["decoder_step_update"]["inputs"]["logical_length"] == ( - "cache_lengths.next" + 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" From 448952f7bc7a02806ce01bdf87d0fdcee18235e6 Mon Sep 17 00:00:00 2001 From: justinchuby Date: Thu, 13 Aug 2026 23:06:37 +0000 Subject: [PATCH 062/151] Enforce fixed-count Muse benchmark semantics Reject paired workflow benchmark configurations that enable EOS stopping so the generation_eos loop always exercises the approved max-iteration device-carry path. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: justinchuby --- scripts/benchmark_muse_workflow.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/scripts/benchmark_muse_workflow.py b/scripts/benchmark_muse_workflow.py index 30bdcb32b..629dd6829 100644 --- a/scripts/benchmark_muse_workflow.py +++ b/scripts/benchmark_muse_workflow.py @@ -49,6 +49,8 @@ def main() -> 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, From 16dff437f257b5147c208a8e94772c8bef6dfa2a Mon Sep 17 00:00:00 2001 From: justinchuby Date: Thu, 13 Aug 2026 23:08:24 +0000 Subject: [PATCH 063/151] Verify paired benchmark runtime identity Require the workflow benchmark runner to come from the pinned c885b71b runtime checkout and reject binaries older than that checkout. Record the verified runtime source head in benchmark evidence. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: justinchuby --- scripts/benchmark_muse_workflow.py | 37 ++++++++++++++++++++++++++++++ 1 file changed, 37 insertions(+) diff --git a/scripts/benchmark_muse_workflow.py b/scripts/benchmark_muse_workflow.py index 629dd6829..f9f736c98 100644 --- a/scripts/benchmark_muse_workflow.py +++ b/scripts/benchmark_muse_workflow.py @@ -28,6 +28,11 @@ 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( "--config", type=Path, @@ -39,6 +44,37 @@ def main() -> int: config: dict[str, Any] = json.loads(args.config.read_text()) workload = config["workload"] sampling = config["sampling"] + 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}" + ) + 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"]) @@ -110,6 +146,7 @@ def main() -> int: "metadata_sha256": _sha256(args.model / "inference_metadata.yaml"), "genai_config_sha256": _sha256(args.model / "genai_config.json"), }, + "runtime_source_head": runtime_head, "metrics": { "ttft_ms": float(median.group(1)), "decode_ms_per_token": float(median.group(2)), From b4108e5b9647cc52f7730c67a95bffb08884ab5a Mon Sep 17 00:00:00 2001 From: justinchuby Date: Thu, 13 Aug 2026 23:09:20 +0000 Subject: [PATCH 064/151] Record paired benchmark runner identity Include the exact c885 source revision, runner path, and runner binary SHA-256 in workflow benchmark evidence so H200 results cannot be attributed to an obsolete build. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: justinchuby --- scripts/benchmark_muse_workflow.py | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/scripts/benchmark_muse_workflow.py b/scripts/benchmark_muse_workflow.py index f9f736c98..3e0ca9b07 100644 --- a/scripts/benchmark_muse_workflow.py +++ b/scripts/benchmark_muse_workflow.py @@ -146,7 +146,11 @@ def main() -> int: "metadata_sha256": _sha256(args.model / "inference_metadata.yaml"), "genai_config_sha256": _sha256(args.model / "genai_config.json"), }, - "runtime_source_head": runtime_head, + "runtime": { + "source_head": runtime_head, + "runner": str(args.runner.resolve()), + "runner_sha256": _sha256(args.runner), + }, "metrics": { "ttft_ms": float(median.group(1)), "decode_ms_per_token": float(median.group(2)), From 49dd6c6ce013d863bc8a771b71d35a2666ba0749 Mon Sep 17 00:00:00 2001 From: justinchuby Date: Thu, 13 Aug 2026 23:13:14 +0000 Subject: [PATCH 065/151] Pin exact prompt-ID runtime head Advance metadata validation and paired Muse benchmark identity to ONNX GenAI 56845704, which applies exact prompt IDs consistently across native, pipeline, and log-probability profiler paths. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: justinchuby --- .github/workflows/main.yml | 2 +- benchmarks/muse_workflow_h200.json | 4 ++-- 2 files changed, 3 insertions(+), 3 deletions(-) diff --git a/.github/workflows/main.yml b/.github/workflows/main.yml index 889384764..e921e687d 100644 --- a/.github/workflows/main.yml +++ b/.github/workflows/main.yml @@ -23,7 +23,7 @@ jobs: - uses: actions/checkout@v7 with: repository: justinchuby/onnx-genai - ref: c885b71b3813fd652690e5ef5154bfc5535e5c1c + ref: 568457045ea003e787816995a380bcd9d7443169 path: validation/onnx-genai - uses: actions/setup-python@v7 with: diff --git a/benchmarks/muse_workflow_h200.json b/benchmarks/muse_workflow_h200.json index d2cc4207c..92bc6924e 100644 --- a/benchmarks/muse_workflow_h200.json +++ b/benchmarks/muse_workflow_h200.json @@ -3,7 +3,7 @@ "package": { "repository": "justinchuby/Muse-Glimmer-30B-ONNX-INT4-CUDA", "weights_revision": "bf36a94a4519e14e3c48ad005c6ff1972ab44ccb", - "schema_head": "c885b71b3813fd652690e5ef5154bfc5535e5c1c", + "schema_head": "568457045ea003e787816995a380bcd9d7443169", "artifacts": { "decoder": "decoder/model.onnx", "embedding": "embedding/model.onnx", @@ -13,7 +13,7 @@ }, "runtime": { "onnxruntime": "1.28.0", - "onnxruntime_genai_commit": "c885b71b3813fd652690e5ef5154bfc5535e5c1c", + "onnxruntime_genai_commit": "568457045ea003e787816995a380bcd9d7443169", "execution_provider": "CUDAExecutionProvider", "cuda_graph": true, "cudnn_flash_attention": false, From 9f58c2bfecce803f1846747f54e472631fafa3e3 Mon Sep 17 00:00:00 2001 From: justinchuby Date: Thu, 13 Aug 2026 23:16:25 +0000 Subject: [PATCH 066/151] Format persistent-mask workflow changes Apply the repository Ruff formatter to the fixed-capacity decoder workflow and its real-artifact regression after the persistent-mask ordering update. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: justinchuby --- .../integrations/onnx_genai/workflow_metadata.py | 16 +++------------- .../onnx_genai/workflow_metadata_test.py | 4 +--- 2 files changed, 4 insertions(+), 16 deletions(-) diff --git a/src/mobius/integrations/onnx_genai/workflow_metadata.py b/src/mobius/integrations/onnx_genai/workflow_metadata.py index aae3e3793..abdea9ced 100644 --- a/src/mobius/integrations/onnx_genai/workflow_metadata.py +++ b/src/mobius/integrations/onnx_genai/workflow_metadata.py @@ -3143,11 +3143,7 @@ def build_vlm_workflow_metadata( "decoder_step_update", { "attention_mask": "state.attention_mask.body", - **( - {"logical_length": "state.cache_lengths.body"} - if fixed_capacity - else {} - ), + **({"logical_length": "state.cache_lengths.body"} if fixed_capacity else {}), **( {"position_ids": "state.position_ids.body"} if position_input is not None @@ -4515,9 +4511,7 @@ def build_decoder_workflow_metadata( setup_decoder_inputs[value.name] = f"initializer.{value.name}" 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" + "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}" @@ -4831,11 +4825,7 @@ def build_decoder_workflow_metadata( "decoder_step_update", { "attention_mask": "state.attention_mask.body", - **( - {"logical_length": "state.cache_lengths.body"} - if fixed_capacity - else {} - ), + **({"logical_length": "state.cache_lengths.body"} if fixed_capacity else {}), **( {"position_ids": "state.position_ids.body"} if position_input is not None diff --git a/src/mobius/integrations/onnx_genai/workflow_metadata_test.py b/src/mobius/integrations/onnx_genai/workflow_metadata_test.py index 734b45c21..7fa0f99b9 100644 --- a/src/mobius/integrations/onnx_genai/workflow_metadata_test.py +++ b/src/mobius/integrations/onnx_genai/workflow_metadata_test.py @@ -251,9 +251,7 @@ def collect_decoder_invokes(node): "max_iterations": "request.max_iterations", } assert policy_invokes["decoder_step_update"]["inputs"]["logical_length"] == "cache_lengths" - assert workflow["state"]["attention_mask"]["initializer"] == ( - "initializer.attention_mask" - ) + 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 From cf14bd8a72fca150b8934f93adf80f56b0fb4fc8 Mon Sep 17 00:00:00 2001 From: justinchuby Date: Thu, 13 Aug 2026 23:17:54 +0000 Subject: [PATCH 067/151] Run workflow conformance on ORT 1.28 Pin the ONNX GenAI execution job to onnxruntime 1.28.0 and explicitly pass that wheel's shared library to the 56845704 workflow conformance suite. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: justinchuby --- .github/workflows/main.yml | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/.github/workflows/main.yml b/.github/workflows/main.yml index e921e687d..d2a88fd52 100644 --- a/.github/workflows/main.yml +++ b/.github/workflows/main.yml @@ -32,7 +32,7 @@ jobs: run: | pip install torch --index-url https://download.pytorch.org/whl/cpu pip install -r requirements/ci/requirements.txt - pip install onnxruntime + pip install onnxruntime==1.28.0 pip install -e '.[testing]' - name: Generate representative packages run: | @@ -52,7 +52,10 @@ jobs: 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/tests/fixtures/onnx_genai_workflows" \ + ONNX_GENAI_ORT_LIB="$ORT_LIB" \ cargo test --manifest-path validation/onnx-genai/Cargo.toml \ -p onnx-genai-engine --test mobius_workflow_conformance -- --nocapture From 1b8157263e6ad5b42d9beb62fc2f03dd5ae43ae2 Mon Sep 17 00:00:00 2001 From: justinchuby Date: Thu, 13 Aug 2026 23:21:12 +0000 Subject: [PATCH 068/151] Pin image-enabled workflow runtime Advance metadata conformance and paired Muse benchmark identity to ONNX GenAI b4422f57, which preserves exact prompt IDs and supports real optional-image workflow profiling. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: justinchuby --- .github/workflows/main.yml | 2 +- benchmarks/muse_workflow_h200.json | 4 ++-- 2 files changed, 3 insertions(+), 3 deletions(-) diff --git a/.github/workflows/main.yml b/.github/workflows/main.yml index d2a88fd52..23ac9c84c 100644 --- a/.github/workflows/main.yml +++ b/.github/workflows/main.yml @@ -23,7 +23,7 @@ jobs: - uses: actions/checkout@v7 with: repository: justinchuby/onnx-genai - ref: 568457045ea003e787816995a380bcd9d7443169 + ref: b4422f5709fc41ea2545e1de46b849343ace9824 path: validation/onnx-genai - uses: actions/setup-python@v7 with: diff --git a/benchmarks/muse_workflow_h200.json b/benchmarks/muse_workflow_h200.json index 92bc6924e..1586d1a16 100644 --- a/benchmarks/muse_workflow_h200.json +++ b/benchmarks/muse_workflow_h200.json @@ -3,7 +3,7 @@ "package": { "repository": "justinchuby/Muse-Glimmer-30B-ONNX-INT4-CUDA", "weights_revision": "bf36a94a4519e14e3c48ad005c6ff1972ab44ccb", - "schema_head": "568457045ea003e787816995a380bcd9d7443169", + "schema_head": "b4422f5709fc41ea2545e1de46b849343ace9824", "artifacts": { "decoder": "decoder/model.onnx", "embedding": "embedding/model.onnx", @@ -13,7 +13,7 @@ }, "runtime": { "onnxruntime": "1.28.0", - "onnxruntime_genai_commit": "568457045ea003e787816995a380bcd9d7443169", + "onnxruntime_genai_commit": "b4422f5709fc41ea2545e1de46b849343ace9824", "execution_provider": "CUDAExecutionProvider", "cuda_graph": true, "cudnn_flash_attention": false, From 2c0580ab3aadd9881e1705af96ca09704f1f8f89 Mon Sep 17 00:00:00 2001 From: justinchuby Date: Thu, 13 Aug 2026 23:23:31 +0000 Subject: [PATCH 069/151] Require CUDA Graph in paired workflow runs Set ONNX_GENAI_CUDA_GRAPH=1 before profiler startup and record both CUDA Graph and cuDNN Flash environment controls in benchmark evidence. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: justinchuby --- scripts/benchmark_muse_workflow.py | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/scripts/benchmark_muse_workflow.py b/scripts/benchmark_muse_workflow.py index 3e0ca9b07..4a9a90048 100644 --- a/scripts/benchmark_muse_workflow.py +++ b/scripts/benchmark_muse_workflow.py @@ -121,6 +121,7 @@ def main() -> int: 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, @@ -150,6 +151,10 @@ def main() -> int: "source_head": runtime_head, "runner": str(args.runner.resolve()), "runner_sha256": _sha256(args.runner), + "environment": { + "ONNX_GENAI_CUDA_GRAPH": "1", + "ORT_ENABLE_CUDNN_FLASH_ATTENTION": "0", + }, }, "metrics": { "ttft_ms": float(median.group(1)), From 8a9aa994067b45661f30607000ae482527a9c343 Mon Sep 17 00:00:00 2001 From: justinchuby Date: Fri, 14 Aug 2026 01:57:45 +0000 Subject: [PATCH 070/151] Refresh persistent-mask workflow fixtures Regenerate decoder and VLM conformance fixtures after switching fixed-capacity generation to one persistent mask updated before each decode invocation. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: justinchuby --- .../decoder/inference_metadata.yaml | 22 +++++++++--------- .../policies/decoder_state_initializer.onnx | Bin 9625 -> 9116 bytes .../vlm/inference_metadata.yaml | 22 +++++++++--------- .../policies/decoder_state_initializer.onnx | Bin 10907 -> 10398 bytes 4 files changed, 22 insertions(+), 22 deletions(-) diff --git a/tests/fixtures/onnx_genai_workflows/decoder/inference_metadata.yaml b/tests/fixtures/onnx_genai_workflows/decoder/inference_metadata.yaml index 8cee222ff..76c27855e 100644 --- a/tests/fixtures/onnx_genai_workflows/decoder/inference_metadata.yaml +++ b/tests/fixtures/onnx_genai_workflows/decoder/inference_metadata.yaml @@ -292,7 +292,7 @@ pipeline: - batch - context scope: invocation - initializer: initializer.body_attention_mask + initializer: initializer.attention_mask recurrence: kind: invariant position_ids: @@ -416,12 +416,21 @@ pipeline: value: token.body output: tokens mode: append + - kind: invoke + component: decoder_step_update + inputs: + attention_mask: attention_mask + logical_length: cache_lengths + position_ids: position_ids + outputs: + next_attention_mask: decoder_step.body_attention_mask + next_position_ids: decoder_step.body_position_ids - kind: invoke component: model inputs: input_ids: token.body past_key_values.0.key: cache_9 - attention_mask: attention_mask + attention_mask: decoder_step.body_attention_mask position_ids: position_ids outputs: logits: decoder.body.logits @@ -432,15 +441,6 @@ pipeline: logits: decoder.body.logits outputs: last_logits: decoder.body.last_logits - - kind: invoke - component: decoder_step_update - inputs: - attention_mask: attention_mask - logical_length: cache_lengths.next - 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: diff --git a/tests/fixtures/onnx_genai_workflows/decoder/policies/decoder_state_initializer.onnx b/tests/fixtures/onnx_genai_workflows/decoder/policies/decoder_state_initializer.onnx index b85fa5c58d6cb0ed70abd2aec96ad22f27d9a04c..bb0bd7179f96f0326ad3a14f39becf05491f4658 100644 GIT binary patch delta 1108 zcmZWn&rj1}7^ZWK@v;7(+bC-%YVE-Q3GMfNY-}dRgcw3Rm`FUB7)xOb4p|3jsm_xJ zFJ8R7F;PQEJQ!o53C6SWZ}7%Vxq3janh^cg4Qbi!?fbp&)8~1g=jHXC1wK6@r<$#m z#m5_BsG}!eUHA&T&U%x$f`*y2d(LLv4i9+k;ilu(Y+|ZHuMEctJM{1} z?B)c{gbz85ozdnSfb)ug)y%x(^$a0!mv83?IZc(2o>31bNW>>X<314@`t)}!BVcr5 zCu?jg^%}>L=(l$zQouu5Sla!gxRoEYw60H0&ehI3|F64gNIgDJh3z!g~3HI1C`7(fH-N+(5O3J8W}L*ob>=&C@X|4hK|ixcqH5IaV&8MtbQ Z&@@CAs8_~W7N})oQiuv-clc$z`v)ysSSkPj delta 1198 zcmZWo&1(};5bws+rknjpU)rqMrnXs%R8-u3?`^j6P?RDHZT&>>P<^IN(1bQ=*)&>9 z{Xp;iA|874Wj9T;$?44u^L{h)`^{{fc|IjaqGaod6*{$ZvRaUO#3(U$wuBR$&!o9>`(#oPpjh)Kp88-E3 z{q=IN3CrPY;~Rn)t+CU0mU87hvoc~c74_EGbis8Qwsgre{BM>Ku6uA%y4wFt5cjfB ze6EsTW$;JT1q?6atDq^v@IzBzM!J2FAYO)dIiIUCOm(GU`j`~f^+h-@D*}b<@~B|L zBl#(OQ;)!;vJiLL$HNYM)sk?oZx$SNaKu4Uja^<|DHQG&7#`QvUE2pAhcEkboF}y( zcMyd=qM#Iuoy>><6jAot^-0B9B-YraT=`~!5mT3a%Y!GT(e1f)Lg0m>_7khUUt%@- zy+w%`045evYV1ncZFEnt-nZxZ1u$Iborf6|3lR(qX?;Xj8hZH&1b+42fbU2dB1l@Q zxW#I*Qf9@xt4otPx3Ln0!mD@@*5aCK2jVu5ufo59qgSIO(gTNz@*?{tcfCL6lo qz?hMO`-TEA4`LTtx6ibNV z;Kg7}%!`SMK;prJ2P4L#F~&c^8#nRlp%{-wqqExr-S)D1^XBb0-}k-uUd|oNa=w1C zXSFn6+$c}&&c5f0{nIZPmNzSvf>kM&EMwIyFN?wXQod%mlTzS%-ma{27FL3!1DjCn4X6j=!$kB~jhpKmg~C?BKp9CmjT_(@_;$gB z&%$M>_1>9z&9Iz9q`?zm*7dXdGApuG+mXAzJzqPk4n35t`!W&z-tfIOe=3-MdF*G zQ;w<8efQuX#4`%K4aKRA23G_Wj_^qGPn+dtLonRsY9Wl%8 zCMFpI$BAcfD?Z_{&9*yAvJTre2S}#_M041#57}%tuVmWDx+jO^TFL5dWOYh*%R_Y7 zA6Drc94B@@%#6J24#Zu8=~8aXF%R9jQYyhckVg&zE_jar$o0$ z=~yVu1>vw zgyylnb`46vlI!ZoBTn#C=$TuKnPQgeNuge|YhF93{-zJOtP@OnuLPcQ z!aWE=777YJ->IbFKoLZnZdFv2R)h+@lqufGQJm1kYGlLX#6Y`sgd>r;0M?tS zV!cLnj^Km?6YGjpp;JY(MjRup``}c?MsTrf7A6!Qj}cHs?Z%o^v#MHv;cwS<_^n70 ztaukoX1<&+6=^BPT3Zn z?eR*6H83^%SbC;s0H)QKwi|1`k!af47D=?UMG~wnlV=~S9_5K{_wZVh;YDy5_SAVl zZdQ?CTgCP8auheglU8_Qt7ZXy2@S(`urfxF;Lw^7hZV$V=$s^-os%T%e3eVWo_3N; z!LiU|cn}$~cq%R9L|{kNz@pitg3+piVW}_+jAp=S85=`~55dhbGUDWII7i6P8+aSY z$=g7-W8eBjo`t{qbPhA}cRp?Pzi&^gj{eyT(IL3pol09lX?Otkb85qq!{}fb)?*rv z2y8?{@O{t&TZ8hslp+XjhTsx-H&^I&?ou7yc`e(j{VRHfhW&!C(IC$jM~a25xgUy! zgSZagCbSi7LiT)PxEl|{=a>u+V Date: Fri, 14 Aug 2026 03:01:57 +0000 Subject: [PATCH 071/151] Manifest dirty benchmark runtimes Reject uncommitted runtime sources for release evidence by default. Diagnostic runs may opt in and now record git status, binary diff SHA-256, and per-file hashes alongside the runner identity. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: justinchuby --- scripts/benchmark_muse_workflow.py | 34 ++++++++++++++++++++++++++++++ 1 file changed, 34 insertions(+) diff --git a/scripts/benchmark_muse_workflow.py b/scripts/benchmark_muse_workflow.py index 4a9a90048..4ed238765 100644 --- a/scripts/benchmark_muse_workflow.py +++ b/scripts/benchmark_muse_workflow.py @@ -24,6 +24,32 @@ def _sha256(path: Path) -> str: 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) @@ -33,6 +59,7 @@ def main() -> int: type=Path, default=Path(".contract-schema-latest"), ) + parser.add_argument("--allow-dirty-runtime", action="store_true") parser.add_argument( "--config", type=Path, @@ -55,6 +82,12 @@ def main() -> int: 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( [ @@ -151,6 +184,7 @@ def main() -> int: "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", From a6ac2d7003a6c3fd58dd7857b236b52cf85be358 Mon Sep 17 00:00:00 2001 From: justinchuby Date: Fri, 14 Aug 2026 03:05:41 +0000 Subject: [PATCH 072/151] Record clean Muse release diagnostics Persist the exact c885 runner, 80ac overlay, paired metrics, transfer counts, and token-parity failure. Mark correctness, performance, and upload gates false. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: justinchuby --- benchmarks/muse_c885_release_evidence.json | 55 ++++++++++++++++++++++ 1 file changed, 55 insertions(+) create mode 100644 benchmarks/muse_c885_release_evidence.json diff --git a/benchmarks/muse_c885_release_evidence.json b/benchmarks/muse_c885_release_evidence.json new file mode 100644 index 000000000..a4c50f7c9 --- /dev/null +++ b/benchmarks/muse_c885_release_evidence.json @@ -0,0 +1,55 @@ +{ + "kind": "workflow_release_diagnostic", + "runtime_source_head": "c885b71b3813fd652690e5ef5154bfc5535e5c1c", + "runner_sha256": "69f22720f35a01b9df23e7e4c688e8a8e8d40b7113f2537e51dd2c3a13c1b015", + "overlay_manifest_sha256": "adbfcdd752d8aedc87d264c0ae0d1082f1f6a615e5dd96688500505456291cc6", + "metadata_sha256": "8727cb6ab44ff65a4e74b6281c5af9eadba336483f5fa1e6033996c3b99fcced", + "workload": { + "prompt_tokens": 68, + "generated_tokens": 128, + "warmups": 1, + "runs": 3, + "decode_skip": 8, + "stop_on_eos": false, + "image": null + }, + "metrics": { + "ttft_ms": 79.939, + "decode_ms_per_token": 28.974, + "throughput_tok_s": 34.51, + "native_ttft_ms": 49.02172100264579, + "native_throughput_tok_s": 63.36527373646786, + "throughput_ratio": 0.5446200728734305 + }, + "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.146 + }, + "token_parity": { + "exact": false, + "first_difference": { + "index": 38, + "runtime": 4243, + "native": 33386 + }, + "runtime_count": 128, + "native_count": 128 + }, + "gate": { + "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." + ] +} From 50129f2c974d4e07d508c3fbafb2842d866b08ce Mon Sep 17 00:00:00 2001 From: justinchuby Date: Fri, 14 Aug 2026 03:06:21 +0000 Subject: [PATCH 073/151] Pin measured Muse producer artifacts Add the exact 80ac producer revision and decoder-state initializer SHA-256 to the clean c885 release evidence so its current-overlay provenance is unambiguous. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: justinchuby --- benchmarks/muse_c885_release_evidence.json | 2 ++ 1 file changed, 2 insertions(+) diff --git a/benchmarks/muse_c885_release_evidence.json b/benchmarks/muse_c885_release_evidence.json index a4c50f7c9..524f639b4 100644 --- a/benchmarks/muse_c885_release_evidence.json +++ b/benchmarks/muse_c885_release_evidence.json @@ -1,9 +1,11 @@ { "kind": "workflow_release_diagnostic", + "producer_commit": "80acfc069f0959f9e6580785fad3172bcc4cc0aa", "runtime_source_head": "c885b71b3813fd652690e5ef5154bfc5535e5c1c", "runner_sha256": "69f22720f35a01b9df23e7e4c688e8a8e8d40b7113f2537e51dd2c3a13c1b015", "overlay_manifest_sha256": "adbfcdd752d8aedc87d264c0ae0d1082f1f6a615e5dd96688500505456291cc6", "metadata_sha256": "8727cb6ab44ff65a4e74b6281c5af9eadba336483f5fa1e6033996c3b99fcced", + "decoder_state_initializer_sha256": "77e4eae0e35cc618214476e0776af3c1b7930862924f431edb62724cadb97af7", "workload": { "prompt_tokens": 68, "generated_tokens": 128, From 80a5ae8f97d13bfc8f89eca37cb696eb28fd0895 Mon Sep 17 00:00:00 2001 From: justinchuby Date: Fri, 14 Aug 2026 03:09:17 +0000 Subject: [PATCH 074/151] Record SHA-named Muse release rerun Replace the prior release diagnostic metrics with the direct immutable 8727cb6a overlay rerun and its current manifest hash. Correctness, performance, and upload gates remain false. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: justinchuby --- benchmarks/muse_c885_release_evidence.json | 12 ++++++------ 1 file changed, 6 insertions(+), 6 deletions(-) diff --git a/benchmarks/muse_c885_release_evidence.json b/benchmarks/muse_c885_release_evidence.json index 524f639b4..f917ed221 100644 --- a/benchmarks/muse_c885_release_evidence.json +++ b/benchmarks/muse_c885_release_evidence.json @@ -3,7 +3,7 @@ "producer_commit": "80acfc069f0959f9e6580785fad3172bcc4cc0aa", "runtime_source_head": "c885b71b3813fd652690e5ef5154bfc5535e5c1c", "runner_sha256": "69f22720f35a01b9df23e7e4c688e8a8e8d40b7113f2537e51dd2c3a13c1b015", - "overlay_manifest_sha256": "adbfcdd752d8aedc87d264c0ae0d1082f1f6a615e5dd96688500505456291cc6", + "overlay_manifest_sha256": "fd77589ffb9333d69efe78909b7f4a489da1071a35affb49b243e87a76330c6a", "metadata_sha256": "8727cb6ab44ff65a4e74b6281c5af9eadba336483f5fa1e6033996c3b99fcced", "decoder_state_initializer_sha256": "77e4eae0e35cc618214476e0776af3c1b7930862924f431edb62724cadb97af7", "workload": { @@ -16,12 +16,12 @@ "image": null }, "metrics": { - "ttft_ms": 79.939, - "decode_ms_per_token": 28.974, - "throughput_tok_s": 34.51, + "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.5446200728734305 + "throughput_ratio": 0.5458825940507668 }, "diagnostics": { "captures": 1, @@ -33,7 +33,7 @@ "d2h_bytes": 4096, "d2d_calls": 3068, "d2d_bytes": 17376, - "island_elapsed_ms": 2904.146 + "island_elapsed_ms": 2904.668 }, "token_parity": { "exact": false, From 1e103e217d12628c53a5cc8246c2ae98efeb864b Mon Sep 17 00:00:00 2001 From: justinchuby Date: Fri, 14 Aug 2026 03:12:13 +0000 Subject: [PATCH 075/151] Record exact prompt-ID release evidence Persist the clean 568457 runtime rerun against the immutable 80ac overlay, including transfer diagnostics, paired performance failure, and token divergence at index 38. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: justinchuby --- benchmarks/muse_568457_release_evidence.json | 57 ++++++++++++++++++++ 1 file changed, 57 insertions(+) create mode 100644 benchmarks/muse_568457_release_evidence.json diff --git a/benchmarks/muse_568457_release_evidence.json b/benchmarks/muse_568457_release_evidence.json new file mode 100644 index 000000000..1a82caf81 --- /dev/null +++ b/benchmarks/muse_568457_release_evidence.json @@ -0,0 +1,57 @@ +{ + "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": { + "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." + ] +} From 6c22eb4124fd279b2422d62142a7d10fe7eeb5a9 Mon Sep 17 00:00:00 2001 From: justinchuby Date: Fri, 14 Aug 2026 04:08:30 +0000 Subject: [PATCH 076/151] Pin landed CUDA bridge runtime Advance workflow validation and benchmark identity to clean runtime 9e575719. Record the exact provenance, parity, transfer diagnostics, and failed performance gate for the bridge measurement now represented by that landed diff. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: justinchuby --- .github/workflows/main.yml | 2 +- .../muse_9e575_landed_bridge_evidence.json | 61 +++++++++++++++++++ benchmarks/muse_workflow_h200.json | 4 +- 3 files changed, 64 insertions(+), 3 deletions(-) create mode 100644 benchmarks/muse_9e575_landed_bridge_evidence.json diff --git a/.github/workflows/main.yml b/.github/workflows/main.yml index 23ac9c84c..1443e5db9 100644 --- a/.github/workflows/main.yml +++ b/.github/workflows/main.yml @@ -23,7 +23,7 @@ jobs: - uses: actions/checkout@v7 with: repository: justinchuby/onnx-genai - ref: b4422f5709fc41ea2545e1de46b849343ace9824 + ref: 9e5757196b98542390ce11f4ff966a58ab3ef578 path: validation/onnx-genai - uses: actions/setup-python@v7 with: diff --git a/benchmarks/muse_9e575_landed_bridge_evidence.json b/benchmarks/muse_9e575_landed_bridge_evidence.json new file mode 100644 index 000000000..f49559a70 --- /dev/null +++ b/benchmarks/muse_9e575_landed_bridge_evidence.json @@ -0,0 +1,61 @@ +{ + "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" + }, + "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 + }, + "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": { + "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_workflow_h200.json b/benchmarks/muse_workflow_h200.json index 1586d1a16..29a140b7d 100644 --- a/benchmarks/muse_workflow_h200.json +++ b/benchmarks/muse_workflow_h200.json @@ -3,7 +3,7 @@ "package": { "repository": "justinchuby/Muse-Glimmer-30B-ONNX-INT4-CUDA", "weights_revision": "bf36a94a4519e14e3c48ad005c6ff1972ab44ccb", - "schema_head": "b4422f5709fc41ea2545e1de46b849343ace9824", + "schema_head": "9e5757196b98542390ce11f4ff966a58ab3ef578", "artifacts": { "decoder": "decoder/model.onnx", "embedding": "embedding/model.onnx", @@ -13,7 +13,7 @@ }, "runtime": { "onnxruntime": "1.28.0", - "onnxruntime_genai_commit": "b4422f5709fc41ea2545e1de46b849343ace9824", + "onnxruntime_genai_commit": "9e5757196b98542390ce11f4ff966a58ab3ef578", "execution_provider": "CUDAExecutionProvider", "cuda_graph": true, "cudnn_flash_attention": false, From 7f66a6c5f8b59edaa57f14f755aa175a43203ee2 Mon Sep 17 00:00:00 2001 From: justinchuby Date: Fri, 14 Aug 2026 04:09:21 +0000 Subject: [PATCH 077/151] Pin dirty bridge evidence manifest Record the deterministic dirty-runtime manifest, five-file patch, ORT library, prompt, and environment hashes for the landed 9e575 bridge diagnostic. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: justinchuby --- benchmarks/muse_9e575_landed_bridge_evidence.json | 12 +++++++++++- 1 file changed, 11 insertions(+), 1 deletion(-) diff --git a/benchmarks/muse_9e575_landed_bridge_evidence.json b/benchmarks/muse_9e575_landed_bridge_evidence.json index f49559a70..3603fc9d9 100644 --- a/benchmarks/muse_9e575_landed_bridge_evidence.json +++ b/benchmarks/muse_9e575_landed_bridge_evidence.json @@ -7,7 +7,9 @@ "base_head": "b4422f5709fc41ea2545e1de46b849343ace9824", "dirty": true, "log_mtime": "2026-08-14T02:09:47Z", - "log_sha256": "d73e93f03086ff9408c12bb10c191678ed19a0f1650ecb738972911f0361d217" + "log_sha256": "d73e93f03086ff9408c12bb10c191678ed19a0f1650ecb738972911f0361d217", + "manifest_sha256": "940ccc975d87fe1920666e7ac86dad8d9ad4aafd03fa51ff36f066da3d200915", + "five_file_binary_patch_sha256": "53d5643611d760f7a685fa037e7ba7a99faeac35a719af46dbccec8ca6a9a03b" }, "landed_runtime": { "head": "9e5757196b98542390ce11f4ff966a58ab3ef578", @@ -24,6 +26,14 @@ "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, From 3fdbb1880753308ede71593671e972bddf846c2a Mon Sep 17 00:00:00 2001 From: justinchuby Date: Fri, 14 Aug 2026 04:15:37 +0000 Subject: [PATCH 078/151] Enforce 0.99x Muse release gate Raise the paired throughput threshold to 0.99x native and require clean landed runtime, exact token parity, and CUDA Graph controls in both benchmark harnesses and persisted evidence. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: justinchuby --- benchmarks/muse_568457_release_evidence.json | 1 + benchmarks/muse_9e575_landed_bridge_evidence.json | 1 + benchmarks/muse_c885_release_evidence.json | 1 + benchmarks/muse_workflow_h200.json | 6 ++++++ scripts/benchmark_muse_native.py | 7 +++++++ scripts/benchmark_muse_workflow.py | 7 +++++++ 6 files changed, 23 insertions(+) diff --git a/benchmarks/muse_568457_release_evidence.json b/benchmarks/muse_568457_release_evidence.json index 1a82caf81..b0a37de3d 100644 --- a/benchmarks/muse_568457_release_evidence.json +++ b/benchmarks/muse_568457_release_evidence.json @@ -46,6 +46,7 @@ "native_count": 128 }, "gate": { + "required_throughput_ratio": 0.99, "correctness": false, "paired_performance": false, "upload_allowed": false diff --git a/benchmarks/muse_9e575_landed_bridge_evidence.json b/benchmarks/muse_9e575_landed_bridge_evidence.json index 3603fc9d9..a311b3786 100644 --- a/benchmarks/muse_9e575_landed_bridge_evidence.json +++ b/benchmarks/muse_9e575_landed_bridge_evidence.json @@ -60,6 +60,7 @@ "native_count": 128 }, "gate": { + "required_throughput_ratio": 0.99, "real_multitoken_e2e": true, "paired_performance": false, "upload_allowed": false diff --git a/benchmarks/muse_c885_release_evidence.json b/benchmarks/muse_c885_release_evidence.json index f917ed221..3da2c22cd 100644 --- a/benchmarks/muse_c885_release_evidence.json +++ b/benchmarks/muse_c885_release_evidence.json @@ -46,6 +46,7 @@ "native_count": 128 }, "gate": { + "required_throughput_ratio": 0.99, "correctness": false, "paired_performance": false, "upload_allowed": false diff --git a/benchmarks/muse_workflow_h200.json b/benchmarks/muse_workflow_h200.json index 29a140b7d..a60b61e79 100644 --- a/benchmarks/muse_workflow_h200.json +++ b/benchmarks/muse_workflow_h200.json @@ -45,5 +45,11 @@ "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/scripts/benchmark_muse_native.py b/scripts/benchmark_muse_native.py index fca1789a4..6e4c550f4 100644 --- a/scripts/benchmark_muse_native.py +++ b/scripts/benchmark_muse_native.py @@ -103,6 +103,13 @@ def main() -> int: 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") diff --git a/scripts/benchmark_muse_workflow.py b/scripts/benchmark_muse_workflow.py index 4ed238765..ac3898c33 100644 --- a/scripts/benchmark_muse_workflow.py +++ b/scripts/benchmark_muse_workflow.py @@ -71,6 +71,13 @@ def main() -> int: 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, From 130e3934baee98e38f41c64effed2d81267bcfcd Mon Sep 17 00:00:00 2001 From: justinchuby Date: Fri, 14 Aug 2026 05:21:53 +0000 Subject: [PATCH 079/151] Add heterogeneous batched generation policies Make decoder and VLM policy graphs row-selective with independent counter-based RNG, per-row sampling controls, ragged prompt and output lengths, and inactive-row state preservation. Publish row-wise emit guards and regenerate schema-compliant workflow fixtures. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: justinchuby --- src/mobius/generation/__init__.py | 2 + src/mobius/generation/_policy_components.py | 224 ++++++++++++++---- .../generation/_policy_components_test.py | 164 +++++++++++++ .../onnx_genai/auto_export_test.py | 18 +- .../onnx_genai/workflow_metadata.py | 216 +++++++++++++++-- .../onnx_genai/workflow_metadata_test.py | 28 +++ .../decoder/inference_metadata.yaml | 59 ++++- .../decoder/policies/cache_length_update.onnx | Bin 467 -> 1044 bytes .../policies/decoder_state_initializer.onnx | Bin 9116 -> 9732 bytes .../decoder/policies/termination.onnx | Bin 3742 -> 4866 bytes .../decoder/policies/token_sampler.onnx | Bin 584 -> 1375 bytes .../decoder/policies/token_state_update.onnx | Bin 1304 -> 2029 bytes .../masked/inference_metadata.yaml | 1 + .../vlm/inference_metadata.yaml | 59 ++++- .../vlm/policies/cache_length_update.onnx | Bin 467 -> 1044 bytes .../policies/decoder_state_initializer.onnx | Bin 10398 -> 11015 bytes .../vlm/policies/termination.onnx | Bin 3742 -> 4866 bytes .../vlm/policies/token_sampler.onnx | Bin 584 -> 1375 bytes .../vlm/policies/token_state_update.onnx | Bin 1304 -> 2029 bytes 19 files changed, 700 insertions(+), 71 deletions(-) diff --git a/src/mobius/generation/__init__.py b/src/mobius/generation/__init__.py index 8d94f27b9..2cca64814 100644 --- a/src/mobius/generation/__init__.py +++ b/src/mobius/generation/__init__.py @@ -25,6 +25,7 @@ build_grammar_logits_processor, build_greedy_sampler, build_integer_add, + build_selective_integer_add, build_integer_minimum, build_iteration_cast, build_last_token_logits, @@ -65,6 +66,7 @@ "build_grammar_logits_processor", "build_greedy_sampler", "build_integer_add", + "build_selective_integer_add", "build_integer_minimum", "build_iteration_cast", "build_last_token_logits", diff --git a/src/mobius/generation/_policy_components.py b/src/mobius/generation/_policy_components.py index d6935eac6..d2f7cbefe 100644 --- a/src/mobius/generation/_policy_components.py +++ b/src/mobius/generation/_policy_components.py @@ -135,15 +135,27 @@ def _make_graph(name: str) -> tuple[ir.Graph, GraphBuilder]: return graph, GraphBuilder(graph) -def build_greedy_sampler(*, effect: str = "sample") -> PolicyComponent: - """Build ``logits -> token_ids`` greedy sampling over the final axis.""" +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"], ) - token_ids = builder.op.ArgMax(logits, axis=-1, keepdims=0) + 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@1", @@ -152,6 +164,7 @@ def build_greedy_sampler(*, effect: str = "sample") -> PolicyComponent: "role": "token_sampler", "mode": "greedy", "logits": "logits", + **({"active": "active", "done": "done"} if row_selective else {}), "token": "token", "effect": effect, }, @@ -213,6 +226,21 @@ def build_integer_add() -> PolicyComponent: 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_batch_minimum() -> PolicyComponent: """Synchronize a per-batch integer length to one conservative scalar.""" graph, builder = _make_graph("batch_minimum") @@ -552,6 +580,7 @@ def build_decoder_state_initializer( position_ids_input: str | None, cache_inputs: list[str], fixed_capacity: bool = False, + ragged: bool = False, ) -> PolicyComponent: """Build prompt-derived decoder state, optionally with capture-stable storage.""" graph, builder = _make_graph("decoder_state_initializer") @@ -570,6 +599,23 @@ def build_decoder_state_initializer( 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( @@ -579,7 +625,7 @@ def build_decoder_state_initializer( ) capacity = op.Add( sequence_length, - op.Squeeze(max_iterations, op.Constant(value_ints=[0])), + op.Squeeze(max_iterations, [0]), ) attention_shape = op.Concat( batch_shape, @@ -600,13 +646,22 @@ def build_decoder_state_initializer( 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, sequence_length), to=attention_value.dtype) + 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.ConstantOfShape(prompt_shape, value=ir.tensor([1])), + op.Less(offsets, op.Unsqueeze(prompt_lengths, [-1])), to=attention_value.dtype, ) attention.shape = attention_value.shape @@ -642,9 +697,9 @@ def build_decoder_state_initializer( ) positions = op.Cast(positions, to=position_value.dtype) positions.shape = position_value.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 = op.Unsqueeze( + op.Cast(prompt_lengths, to=position_value.dtype), + [-1], ) body_position.shape = ir.Shape(["batch", 1]) builder.add_output(attention, attention_mask_input) @@ -655,11 +710,15 @@ def build_decoder_state_initializer( if position_ids_input is not None: builder.add_output(body_position, "body_position_ids") builder.add_output(token_slot, "token_slot") - if fixed_capacity: - cache_lengths = op.Expand( - op.Unsqueeze(sequence_length, op.Constant(value_ints=[0])), + 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: + cache_lengths = op.Identity(prompt_lengths) cache_lengths.shape = ir.Shape(["batch"]) builder.add_output(cache_lengths, "cache_lengths") @@ -1052,13 +1111,16 @@ def build_seeded_categorical_sampler() -> PolicyComponent: 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, [1]) - top_k = builder.input("top_k", ir.DataType.INT64, [1]) - top_p = builder.input("top_p", ir.DataType.FLOAT, [1]) - min_p = builder.input("min_p", ir.DataType.FLOAT, [1]) + 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"]) grammar_mask = builder.input("grammar_mask", ir.DataType.BOOL, ["batch", "vocabulary"]) seed = builder.input("seed", ir.DataType.INT64, ["batch"]) offset = builder.input("offset", 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. @@ -1108,19 +1170,25 @@ def build_seeded_categorical_sampler() -> PolicyComponent: blocked = op.CastLike(op.Constant(value_float=-3.4028235e38), logits) constrained_logits = op.Where(grammar_mask, logits, blocked) - safe_temperature = op.Max(temperature, op.Constant(value_float=1e-6)) + safe_temperature = op.Unsqueeze( + op.Max(temperature, op.Constant(value_float=1e-6)), + [-1], + ) scaled_logits = op.Div(constrained_logits, safe_temperature) - safe_min_p = op.Clip( - min_p, - op.Constant(value_float=1e-20), - op.Constant(value_float=1.0), + 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.LessOrEqual(min_p, op.Constant(value_float=0.0)), + 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) @@ -1134,12 +1202,19 @@ def build_seeded_categorical_sampler() -> PolicyComponent: effective_k = op.Min(op.Max(requested_k, op.Constant(value_int=1)), vocabulary) _, top_indices = op.TopK( scaled_logits, - effective_k, + vocabulary, axis=-1, largest=1, - sorted=0, + 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( @@ -1147,10 +1222,7 @@ def build_seeded_categorical_sampler() -> PolicyComponent: value=ir.tensor([0], dtype=ir.DataType.INT64), ), top_indices, - op.ConstantOfShape( - op.Shape(top_indices), - value=ir.tensor([1], dtype=ir.DataType.INT64), - ), + op.Cast(keep_top_k, to=ir.DataType.INT64), axis=1, ), op.Constant(value_int=0), @@ -1171,10 +1243,13 @@ def build_seeded_categorical_sampler() -> PolicyComponent: sorted_probabilities, op.Constant(value_int=1), ) - safe_top_p = op.Clip( - top_p, - op.Constant(value_float=1e-6), - op.Constant(value_float=1.0), + 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), @@ -1223,7 +1298,12 @@ def build_seeded_categorical_sampler() -> PolicyComponent: token_ids, op.Constant(value_int=-1), ) - next_offset = op.Add(offset, op.Constant(value_int=1)) + token_ids = op.Where(enabled, token_ids, op.Constant(value_int=-1)) + next_offset = op.Where( + enabled, + op.Add(offset, op.Constant(value_int=1)), + offset, + ) builder.add_output(token_ids, "token") builder.add_output(next_offset, "next_offset") return _component( @@ -1239,6 +1319,8 @@ def build_seeded_categorical_sampler() -> PolicyComponent: "top_p": "top_p", "min_p": "min_p", "grammar_mask": "grammar_mask", + "active": "active", + "done": "done", "rng": { "rng_seed": "seed", "rng_offset": "offset", @@ -1250,14 +1332,22 @@ def build_seeded_categorical_sampler() -> PolicyComponent: ) -def build_eos_termination() -> PolicyComponent: +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("token_ids", ir.DataType.INT64, ["batch"]) eos_ids = builder.input("eos_ids", ir.DataType.INT64, ["num_eos"]) - iteration = builder.input("iteration", ir.DataType.INT64, ["batch"]) - max_iterations = builder.input("max_iterations", ir.DataType.INT64, ["batch"]) + iteration = builder.input( + "iteration", + ir.DataType.INT64, + [1] if row_selective else ["batch"], + ) + max_iterations = builder.input( + "max_iterations", + ir.DataType.INT64, + [1] if row_selective else ["batch"], + ) tokens = op.Unsqueeze(token_ids, op.Constant(value_ints=[-1])) eos = op.Unsqueeze(eos_ids, op.Constant(value_ints=[0])) matches = op.Equal(tokens, eos) @@ -1271,13 +1361,24 @@ def build_eos_termination() -> PolicyComponent: op.Add(iteration, op.Constant(value_int=1)), max_iterations, ) - done = op.Or(hit_eos, hit_limit) - continued = op.Equal( - op.ReduceMax(op.Cast(done, to=ir.DataType.INT64), keepdims=1), + if row_selective: + active = builder.input("active", ir.DataType.BOOL, ["batch"]) + previous_done = builder.input("previous_done", ir.DataType.BOOL, ["batch"]) + enabled = op.And(active, op.Not(previous_done)) + newly_done = op.And(enabled, op.Or(hit_eos, hit_limit)) + done = op.Or(previous_done, newly_done) + next_active = op.And(enabled, op.Not(newly_done)) + 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), ) continued.shape = ir.Shape([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@1", @@ -1288,7 +1389,11 @@ def build_eos_termination() -> PolicyComponent: "eos_ids": "eos_ids", "iteration": "iteration", "max_iterations": "max_iterations", + **( + {"active": "active", "previous_done": "previous_done"} if row_selective else {} + ), "done": "done", + **({"next_active": "next_active"} if row_selective else {}), "continue": "continue", "effect": "termination", }, @@ -1623,17 +1728,33 @@ def build_token_block_identity() -> PolicyComponent: return _component("mobius.policy.auxiliary@1", graph, {}, "state") -def build_token_state_update() -> PolicyComponent: - """Build explicit token-history append and sequence-length update math.""" +def build_token_state_update(*, row_selective: bool = False) -> PolicyComponent: + """Selectively update token state and per-row generated lengths.""" 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"]) - next_state = op.Add( - op.Mul(current, op.Constant(value_int=0)), - op.Unsqueeze(update, op.Constant(value_ints=[-1])), - ) + if row_selective: + lengths = builder.input("lengths", 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)) + next_state = op.Where( + op.Unsqueeze(enabled, [-1]), + op.Unsqueeze(update, [-1]), + current, + ) + emitted_length = op.Cast(enabled, to=ir.DataType.INT64) + next_lengths = op.Add(lengths, emitted_length) + else: + next_state = op.Unsqueeze(update, [-1]) + next_state.shape = ir.Shape(["batch", 1]) builder.add_output(next_state, "next") + if row_selective: + next_lengths.shape = ir.Shape(["batch"]) + emitted_length.shape = ir.Shape(["batch"]) + builder.add_output(next_lengths, "next_lengths") + builder.add_output(emitted_length, "emitted_length") return _component( "onnx-genai.state-update@1", graph, @@ -1641,7 +1762,20 @@ def build_token_state_update() -> PolicyComponent: "role": "state_update", "current": "current", "update": "update", + **( + {"lengths": "lengths", "active": "active", "done": "done"} + if row_selective + else {} + ), "next": "next", + **( + { + "next_lengths": "next_lengths", + "emitted_length": "emitted_length", + } + if row_selective + else {} + ), "effect": "state", }, "state", diff --git a/src/mobius/generation/_policy_components_test.py b/src/mobius/generation/_policy_components_test.py index fc8cd6107..2781ff94b 100644 --- a/src/mobius/generation/_policy_components_test.py +++ b/src/mobius/generation/_policy_components_test.py @@ -158,6 +158,19 @@ def test_greedy_sampler_runtime(tmp_path): 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}) @@ -304,6 +317,76 @@ def test_decoder_fixed_capacity_state_matches_native_capture_layout(tmp_path): 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), @@ -419,6 +502,8 @@ def test_seeded_sampler_is_counter_based_and_reproducible(tmp_path): "grammar_mask": np.array([[True, True, True, True]], np.bool_), "seed": np.array([7], np.int64), "offset": 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) @@ -439,6 +524,8 @@ def test_seeded_sampler_applies_request_top_k_and_grammar_mask(tmp_path): "grammar_mask": np.array([[False, True, True, True]], np.bool_), "seed": np.array([17], np.int64), "offset": np.array([0], np.int64), + "active": np.array([True], np.bool_), + "done": np.array([False], np.bool_), }, ) np.testing.assert_array_equal(token, [1]) @@ -457,6 +544,8 @@ def test_seeded_sampler_rejects_empty_grammar_vocabulary(tmp_path): "grammar_mask": np.array([[False, False, False]], np.bool_), "seed": np.array([1], np.int64), "offset": np.array([0], np.int64), + "active": np.array([True], np.bool_), + "done": np.array([False], np.bool_), }, ) np.testing.assert_array_equal(token, [-1]) @@ -475,12 +564,87 @@ def test_seeded_sampler_applies_request_min_p_in_logit_space(tmp_path): "grammar_mask": np.array([[True, True, True]], np.bool_), "seed": np.array([23], np.int64), "offset": 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_offset, [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), + "grammar_mask": np.array( + [ + [True, True, True, True, True], + [True, False, True, True, True], + [True, True, True, False, False], + [True, True, True, True, True], + ], + np.bool_, + ), + "seed": np.array([3, 7, 11, 13], np.int64), + "offset": 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, next_lengths, emitted = _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), + "lengths": np.array([2, 4, 6], 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]]) + np.testing.assert_array_equal(next_lengths, [3, 4, 6]) + np.testing.assert_array_equal(emitted, [1, 0, 0]) + + done, next_active, continued = _run( + build_eos_termination(row_selective=True), + tmp_path, + { + "token_ids": np.array([2, 8, 9], np.int64), + "eos_ids": np.array([2, 9], np.int64), + "iteration": np.array([0], np.int64), + "max_iterations": np.array([5], np.int64), + "active": np.array([True, False, True], np.bool_), + "previous_done": np.array([False, False, True], np.bool_), + }, + ) + np.testing.assert_array_equal(done, [True, False, True]) + np.testing.assert_array_equal(next_active, [False, False, False]) + np.testing.assert_array_equal(continued, [False]) + + def test_eos_termination_runtime(tmp_path): terminated, continued = _run( build_eos_termination(), diff --git a/src/mobius/integrations/onnx_genai/auto_export_test.py b/src/mobius/integrations/onnx_genai/auto_export_test.py index 1b3c369fa..26b29a606 100644 --- a/src/mobius/integrations/onnx_genai/auto_export_test.py +++ b/src/mobius/integrations/onnx_genai/auto_export_test.py @@ -168,9 +168,13 @@ def test_dispatch_decoder(tmp_path): "onnx-genai.termination-predicate" ) assert workflow["steps"][0]["kind"] == "loop" - assert all( - value["source"]["kind"] != "application" for value in workflow["inputs"].values() - ) + application_inputs = { + name + for name, value in workflow["inputs"].items() + if value["source"]["kind"] == "application" + } + assert application_inputs == {"request.prompt_lengths"} + assert workflow["inputs"]["request.prompt_lengths"]["default"] == -1 assert [node["component"] for node in workflow["steps"][0]["setup"]] == [ "decoder_state_initializer", "model", @@ -221,6 +225,8 @@ def test_seeded_decoder_sampler_uses_request_controls_and_direct_kv_carry(): "rng_seed": "seed", "rng_offset": "offset", "rng_next_offset": "next_offset", + "active": "active", + "done": "done", } sampler_step = next( step @@ -236,7 +242,13 @@ def test_seeded_decoder_sampler_uses_request_controls_and_direct_kv_carry(): "grammar_mask": "request.grammar_mask", "seed": "request.seed", "offset": "rng_offset", + "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["state"]["rng_offset"]["class"] == "semantic" assert workflow["state"]["rng_offset"]["initializer"] == "request.rng_offset" assert not any("kv_update" in name for name in workflow["components"]) diff --git a/src/mobius/integrations/onnx_genai/workflow_metadata.py b/src/mobius/integrations/onnx_genai/workflow_metadata.py index abdea9ced..a50c25e5d 100644 --- a/src/mobius/integrations/onnx_genai/workflow_metadata.py +++ b/src/mobius/integrations/onnx_genai/workflow_metadata.py @@ -22,18 +22,22 @@ build_decoder_state_initializer, build_decoder_step_update, build_empty_features, + build_eos_termination, build_euler_model_input, build_euler_solver_step, build_greedy_sampler, build_integer_add, + build_selective_integer_add, build_integer_minimum, build_last_token_logits, build_model_token_cast, build_proposal_metrics, build_schedule_constant, build_schedule_lookup, + build_seeded_categorical_sampler, build_sequence_length, build_token_to_slot, + build_token_state_update, build_tts_decoder_state_initializer, build_tts_decoder_step_update, build_tts_state_initializer, @@ -228,6 +232,8 @@ def convert(node: dict[str, Any]) -> dict[str, Any]: } 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 = { @@ -2671,6 +2677,7 @@ def build_vlm_workflow_metadata( 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, ), ) pkg.add_policy_component( @@ -2681,8 +2688,14 @@ def build_vlm_workflow_metadata( fixed_capacity=fixed_capacity, ), ) + pkg.add_policy_component("token_sampler", build_greedy_sampler(row_selective=True)) + pkg.add_policy_component("termination", build_eos_termination(row_selective=True)) + pkg.add_policy_component( + "token_state_update", + build_token_state_update(row_selective=True), + ) if cache_pairs: - pkg.add_policy_component("cache_length_update", build_integer_add()) + pkg.add_policy_component("cache_length_update", build_selective_integer_add()) batch = _contract(token_input)["shape"][0] batch_int = {"dtype": "int64", "rank": 1, "shape": [batch]} @@ -2715,6 +2728,13 @@ def build_vlm_workflow_metadata( "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, + }, "package.eos_ids": { "contract": {"dtype": "int64", "rank": 1, "shape": [1]}, "role": {"kind": "opaque"}, @@ -2736,7 +2756,7 @@ def build_vlm_workflow_metadata( ), }, "package.one": { - "contract": control_int, + "contract": batch_int, "role": {"kind": "opaque"}, "source": {"kind": "literal"}, "required": False, @@ -2854,6 +2874,13 @@ def build_vlm_workflow_metadata( "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"], @@ -2941,11 +2968,18 @@ def build_vlm_workflow_metadata( "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", + ), ( "active", "package.active", "state.active.body", - "loop.continue", + "loop.next_active", "state.active.final", ), ("done", "package.false", "state.done.body", "loop.done", "state.done.final"), @@ -3104,12 +3138,14 @@ def build_vlm_workflow_metadata( "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 fixed_capacity @@ -3164,7 +3200,11 @@ def build_vlm_workflow_metadata( "nodes": [ _invoke( "token_sampler", - {"logits": "state.logits.body"}, + { + "logits": "state.logits.body", + "active": "state.active.body", + "done": "state.done.body", + }, {"token": "sample.body"}, {"sample": _effect("sample.0", "sample.1")}, ), @@ -3175,8 +3215,14 @@ def build_vlm_workflow_metadata( "eos_ids": "package.eos_ids", "iteration": "loop.iteration", "max_iterations": "request.max_iterations", + "active": "state.active.body", + "previous_done": "state.done.body", + }, + { + "done": "loop.done", + "next_active": "loop.next_active", + "continue": "loop.continue", }, - {"done": "loop.done", "continue": "loop.continue"}, {"termination": _effect("termination.0", "termination.1")}, ), *( @@ -3186,12 +3232,19 @@ def build_vlm_workflow_metadata( { "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"}, + { + "left": "package.zero_batch", + "right": "package.one", + "active": "state.active.body", + "done": "state.done.body", + }, {"total": "accepted_len.next"}, ), ] @@ -3200,8 +3253,18 @@ def build_vlm_workflow_metadata( ), _invoke( "token_state_update", - {"current": "state.token.body", "update": "sample.body"}, - {"next": "token.body"}, + { + "current": "state.token.body", + "update": "sample.body", + "lengths": "state.generated_lengths.body", + "active": "state.active.body", + "done": "state.done.body", + }, + { + "next": "token.body", + "next_lengths": "token.next_lengths", + "emitted_length": "token.emitted_length", + }, {"state": _effect("state.0", "state.1")}, ), { @@ -3209,6 +3272,8 @@ def build_vlm_workflow_metadata( "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"), }, @@ -3258,6 +3323,7 @@ def build_vlm_workflow_metadata( "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"] @@ -4292,6 +4358,7 @@ def build_decoder_workflow_metadata( 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), ), ) pkg.add_policy_component( @@ -4360,7 +4427,7 @@ def build_decoder_workflow_metadata( "default": eos_token_id, }, "package.one_token": { - "contract": control_int, + "contract": batch_int, "role": {"kind": "opaque"}, "source": {"kind": "literal"}, "required": False, @@ -4375,12 +4442,24 @@ def build_decoder_workflow_metadata( }, } ) + if cache_pairs: + workflow_inputs["request.prompt_lengths"] = { + "contract": batch_int, + "role": {"kind": "opaque"}, + "source": {"kind": "application", "name": "prompt_lengths"}, + "required": False, + "default": -1, + } stochastic_sampler = sampler != "greedy" if stochastic_sampler: workflow_inputs.update( { "request.temperature": { - "contract": {"dtype": "float32", "rank": 1, "shape": [1]}, + "contract": { + "dtype": "float32", + "rank": 1, + "shape": [batch_dimension], + }, "role": { "kind": "runtime", "version": "1.0", @@ -4394,7 +4473,7 @@ def build_decoder_workflow_metadata( "default": 1.0, }, "request.top_k": { - "contract": {"dtype": "int64", "rank": 1, "shape": [1]}, + "contract": batch_int, "role": { "kind": "runtime", "version": "1.0", @@ -4405,7 +4484,11 @@ def build_decoder_workflow_metadata( "default": 0, }, "request.top_p": { - "contract": {"dtype": "float32", "rank": 1, "shape": [1]}, + "contract": { + "dtype": "float32", + "rank": 1, + "shape": [batch_dimension], + }, "role": { "kind": "runtime", "version": "1.0", @@ -4416,7 +4499,11 @@ def build_decoder_workflow_metadata( "default": 1.0, }, "request.min_p": { - "contract": {"dtype": "float32", "rank": 1, "shape": [1]}, + "contract": { + "dtype": "float32", + "rank": 1, + "shape": [batch_dimension], + }, "role": { "kind": "runtime", "version": "1.0", @@ -4462,7 +4549,20 @@ def build_decoder_workflow_metadata( } ) if cache_pairs: - pkg.add_policy_component("cache_length_update", build_integer_add()) + pkg.add_policy_component("cache_length_update", build_selective_integer_add()) + pkg.add_policy_component( + "token_sampler", + ( + build_greedy_sampler(row_selective=True) + if not stochastic_sampler + else build_seeded_categorical_sampler() + ), + ) + pkg.add_policy_component("termination", build_eos_termination(row_selective=True)) + pkg.add_policy_component( + "token_state_update", + build_token_state_update(row_selective=True), + ) workflow_inputs.update( { "package.active": { @@ -4549,6 +4649,13 @@ def build_decoder_workflow_metadata( 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", @@ -4625,11 +4732,18 @@ def build_decoder_workflow_metadata( 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.continue", + "body_output": "loop.next_active", "next": "state.active.final", }, { @@ -4745,7 +4859,9 @@ def build_decoder_workflow_metadata( decoder_kv_ports: dict[str, Any] = {} decoder_kv_axis = 2 for past, present in cache_pairs: - cell = f"cache_{len(carried)}" + # Generated-length state is orthogonal to the admitted cache ABI and + # must not renumber stable cache service cells. + cell = f"cache_{len(carried) - 1}" setup_value = f"decoder.setup.{present.name}" body_value = f"decoder.body.{present.name}" setup_decoder_outputs[present.name] = setup_value @@ -4791,12 +4907,18 @@ def build_decoder_workflow_metadata( "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", "token_slot": "initializer.token_slot", + **( + {"generated_lengths": "initializer.generated_lengths"} + if cache_pairs + else {} + ), **( {"cache_lengths": "initializer.cache_lengths"} if fixed_capacity @@ -4861,6 +4983,14 @@ def build_decoder_workflow_metadata( if stochastic_sampler else {} ), + **( + { + "active": "state.active.body", + "done": "state.done.body", + } + if cache_pairs + else {} + ), }, { "token": "sample.body", @@ -4870,8 +5000,30 @@ def build_decoder_workflow_metadata( ), _invoke( "token_state_update", - {"current": "state.token.body", "update": "sample.body"}, - {"next": "token.body"}, + { + "current": "state.token.body", + "update": "sample.body", + **( + { + "lengths": "state.generated_lengths.body", + "active": "state.active.body", + "done": "state.done.body", + } + if cache_pairs + else {} + ), + }, + { + "next": "token.body", + **( + { + "next_lengths": "token.next_lengths", + "emitted_length": "token.emitted_length", + } + if cache_pairs + else {} + ), + }, {"state": _effect("state.0", "state.1")}, ), *( @@ -4892,8 +5044,20 @@ def build_decoder_workflow_metadata( "eos_ids": "package.eos_ids", "iteration": "loop.iteration", "max_iterations": "request.max_iterations", + **( + { + "active": "state.active.body", + "previous_done": "state.done.body", + } + if cache_pairs + else {} + ), + }, + { + "done": "loop.done", + "continue": "loop.continue", + **({"next_active": "loop.next_active"} if cache_pairs else {}), }, - {"done": "loop.done", "continue": "loop.continue"}, {"termination": _effect("termination.0", "termination.1")}, ), *( @@ -4903,6 +5067,8 @@ def build_decoder_workflow_metadata( { "left": "state.cache_lengths.body", "right": "package.one_token", + "active": "state.active.body", + "done": "state.done.body", }, {"total": "cache_lengths.next"}, ), @@ -4911,6 +5077,8 @@ def build_decoder_workflow_metadata( { "left": "package.zero_batch", "right": "package.one_token", + "active": "state.active.body", + "done": "state.done.body", }, {"total": "accepted_len.next"}, ), @@ -4923,6 +5091,14 @@ def build_decoder_workflow_metadata( "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"), }, @@ -4948,6 +5124,7 @@ def build_decoder_workflow_metadata( "linear_effects", "nested_control_flow", "typed_emit", + "emit_valid_length", "loop_induction_values", *(["serving_service_contract"] if cache_pairs else []), *(["bounded_state_recurrence"] if cache_pairs else []), @@ -5243,6 +5420,7 @@ def update_invoke( "linear_effects", "nested_control_flow", "typed_emit", + "emit_valid_length", "loop_induction_values", ], }, diff --git a/src/mobius/integrations/onnx_genai/workflow_metadata_test.py b/src/mobius/integrations/onnx_genai/workflow_metadata_test.py index 7fa0f99b9..f4429585c 100644 --- a/src/mobius/integrations/onnx_genai/workflow_metadata_test.py +++ b/src/mobius/integrations/onnx_genai/workflow_metadata_test.py @@ -248,8 +248,36 @@ def collect_decoder_invokes(node): 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["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 policy_invokes["token_state_update"]["inputs"]["lengths"] == ("generated_lengths") + assert policy_invokes["token_state_update"]["outputs"]["next_lengths"] == ( + "token.next_lengths" + ) + + 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 policy_invokes["decoder_step_update"]["inputs"]["logical_length"] == "cache_lengths" assert workflow["state"]["attention_mask"]["initializer"] == ("initializer.attention_mask") assert any( diff --git a/tests/fixtures/onnx_genai_workflows/decoder/inference_metadata.yaml b/tests/fixtures/onnx_genai_workflows/decoder/inference_metadata.yaml index 76c27855e..e2af713d3 100644 --- a/tests/fixtures/onnx_genai_workflows/decoder/inference_metadata.yaml +++ b/tests/fixtures/onnx_genai_workflows/decoder/inference_metadata.yaml @@ -10,6 +10,7 @@ pipeline: - linear_effects - nested_control_flow - typed_emit + - emit_valid_length - loop_induction_values - serving_service_contract - bounded_state_recurrence @@ -58,7 +59,7 @@ pipeline: dtype: int64 rank: 1 shape: - - 1 + - batch role: kind: opaque source: @@ -77,6 +78,19 @@ pipeline: kind: literal required: false default: 8192 + request.prompt_lengths: + contract: + dtype: int64 + rank: 1 + shape: + - batch + role: + kind: opaque + source: + kind: application + name: prompt_lengths + required: false + default: -1 package.active: contract: dtype: bool @@ -161,6 +175,8 @@ pipeline: version: '1' bindings: logits: logits + active: active + done: done token: token parameters: mode: greedy @@ -177,7 +193,10 @@ pipeline: eos_ids: eos_ids iteration: iteration max_iterations: max_iterations + active: active + previous_done: previous_done done: done + next_active: next_active continue: continue token_state_update: implementation: @@ -189,7 +208,12 @@ pipeline: bindings: current: current update: update + lengths: lengths + active: active + done: done next: next + next_lengths: next_lengths + emitted_length: emitted_length last_token_logits: implementation: kind: onnx @@ -229,6 +253,17 @@ pipeline: initializer: decoder.setup.last_logits recurrence: kind: invariant + generated_lengths: + contract: + dtype: int64 + rank: 1 + shape: + - batch + class: semantic + scope: invocation + initializer: initializer.generated_lengths + recurrence: + kind: invariant active: contract: dtype: bool @@ -349,11 +384,13 @@ pipeline: component: decoder_state_initializer inputs: prompt_tokens: request.input_ids + prompt_lengths: request.prompt_lengths max_iterations: request.max_iterations outputs: attention_mask: initializer.attention_mask body_attention_mask: initializer.body_attention_mask token_slot: initializer.token_slot + generated_lengths: initializer.generated_lengths cache_lengths: initializer.cache_lengths position_ids: initializer.position_ids body_position_ids: initializer.body_position_ids @@ -379,6 +416,8 @@ pipeline: component: token_sampler inputs: logits: logits + active: active + done: done outputs: token: sample.body - kind: invoke @@ -386,8 +425,13 @@ pipeline: inputs: current: token update: sample.body + lengths: generated_lengths + active: active + done: done outputs: next: token.body + next_lengths: token.next_lengths + emitted_length: token.emitted_length - kind: invoke component: termination inputs: @@ -395,14 +439,19 @@ pipeline: eos_ids: package.eos_ids iteration: loop.iteration max_iterations: request.max_iterations + active: active + previous_done: done 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 @@ -410,12 +459,16 @@ pipeline: 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: decoder_step_update inputs: @@ -448,8 +501,10 @@ pipeline: next: token.body - cell: logits next: decoder.body.last_logits + - cell: generated_lengths + next: token.next_lengths - cell: active - next: loop.continue + next: loop.next_active - cell: done next: loop.done - cell: cache_lengths diff --git a/tests/fixtures/onnx_genai_workflows/decoder/policies/cache_length_update.onnx b/tests/fixtures/onnx_genai_workflows/decoder/policies/cache_length_update.onnx index 45764e632822b23fc759eb608a1f8079786bf246..980ab9a3d96fd6966551637a8856faffc0bf4bb9 100644 GIT binary patch delta 330 zcmcc2JcUD_gIkC#H$N$}wAiYVgR7j2B_%&ERfxST-Y>r--av{AMCIkDq{bU4F$38X zO~Q4#ni#p*5|c|Z%Tl>u3PCy?^HSmsrMN&8Scjn!vt!;w;|LY5Ixd!+)U*;V)}qYx zj1sWgl$3ZQpjsdWR%neqzIx_h_xiYBr!*dCp;syC>5f|Scw(P zn`jcjEn1wKlL~fH{KVc8fhZH03x%|}G&ndX3oxqjMFGtM`wlKNnUT@dK1&$k3l1(e d4o)E^E(VAo*d;>ff*@Cc1=&F=Cz~JeSQ&p$|2?>NMwJNLidYd?xwefDE zqzE4Rz!M-r%S|9C0{#Loh!_45zW~n6UOY2ft4*cdHZPg+dS<@;=A1d-IpfMAp4%Vx z1_!~L|GterGtlyAJlr3Jb~xPi`~WW;+V5{WBhPM_OHsog_Mumu6}`=6wBS4YUNCZ8 z56^5~L>ET8TlJyu9|i7sFbeDLo)ZN2_P`rEV|V))mu}tOd>NVAxA+`(;0{M#z`ubW zP6t91AS4D-L+F&}_t%YwMya-hW`ocfhwC%Ydj=4x0+DJEOA15|I4_9z4ZL({Umg0u zh9BC@Je{`1F0yJR`Ji?Np+jfyz_SN_7+eP8=JP`Jt$|XjqDTRV7s1>c zNmGEes`GI_FJM0zNM0WymosORzTLpfhxP}45JPL4tCNlZ*Q_q2pX3Ga0r-T1VnQRq z=+M4)G;;jD-7=TsmO$33&PCn4V12C`Sz9-6t{mD|`hA-abCLg#K#}Sx z=;TG$HwHf;lEc&{)LczF(fU+!gW}HGpAD3KR&!i_}O__V5XQNOl`%FfXs$n`vVpzSY2fRxGzPgC08!j3r|Rx1Om zH4Rpty5m<7sZxSSbD9yhc-(r=@wYsiw9VzH74g}w&WdhP9^1)K*O90z&SN1qNSvrX z+tGZsoAudl0Y3W!#0KT;=<;wDdnd8D*TxEjHp~@SjL=5jY6LXgBw|V-i>TC5lA(^0 z0;sz$0$0k8$|EH_+P~vLzCSplu~1I%kYXoGDwGr*8T~)lIa+lhz4OiWp!b=Z_Q< zql$K7P;n|!J294_@7YwAWbMSLN*7ZF0JAazfLR3yfbTTz#H@*zG$iH3nyPYQ&AgQp zh=7T3V@R&!h&2`R$b3Y>OfaDEYRVhoD7JR_uu!kC^; zvrYl#^_O($5`1RM^LhEk>)UY^Y)XV>-RC00exYplvg^3pUW!~3dRC)wp=V{nge8Xw zRp>#LKR2PLg?uYhZL*%TYLlf^o2+cL$tt4Sq!VZ9fy;vQA<>}Pl}IJW4gE)UTP1E; zT2ahm>SK{@$70%om=QNcwq`-hi28-p$o`gvz175LM?7!b^^R?!)Ps7X4xK3{kk!hb zK<5&Tk3i!an`JKX^V9eEeBX27ZgXtIZCS|g!65oMdxN{)xMuPsV?(+Rnj{>VXJoVT zU;KCTK)+7yqh4hme~4bgmqx?AfqPt^x@f)Ugzh$?xbfI;#y1^-2lrFH>)qVYZxnC? z)nVP*ldpU8CB6QS5vsr(WmKvdy&}%kQ&LKVv0}E zqYo48A-7^+`$XiFPGost?vb5}MgLr>xUXCtFjlE5(2XG2pCJ zY}MyU@>bGF=fRVEfb=r~GzqzIai#!;Lr8uLqQ3fhwd|~NEEgO|GhJ>?37@2uJ`kKl z=`?9e@1o)BDEKaD_@J->DCS{O@IfF{AWW;h>ChQ?qz#4IZ!&a-500!PDarCzTY#gN z4^#m`WpDD{Mzmn80eEdKwMTAtG9|Q=zPUDGR~ltRCXs9=;f(%cz5o2ZYX|dwE-rzy8b_7(hgSu literal 9116 zcmc&)TW{P%6ppu>>}I+N$)L8WK^3g30*ka-&vHO<$E9vFJx4I zA$^XtNSpQKMQamayVQK){@P-;~KNdWPAFt<;R9p{0=OnozXjKG@B zTHG!Qm?EyODz5HQ#PyqoHjk`R=7w;DLKl=N$*e2Ntf$DVSGdf+(vYY#vC+e|6Lw=f zW_=y;e~}jZ&2#Wvk_HE=1_w*g;IA4=Bq6D8fHzLqhh7jH{7~P@8j-Mu&Gqz6Q6l?O zIE@C90ddm;2Vq+wGdlWa{1`zqI?c6tyLBE_ty4EB)Y*QoaakqKiW!E!l{V&h3_-P9 z;9wG8Rw^u#gGCm>D$?)s=qX7nvp?xQ!RUm&Z+UwTGpN2Xf0~PvHdmu&QD*mpqEIs@ zIa82adBF)pC&)w=qlq`g^J;YU?euugy_?Z3;l%y|0YS#Wy>5IsjSfGwZ=X&rZ_Esc z2l3-vtG(vhyj_%9zlp*@y1W<;OwJ|Wb$|>SeZ3Jq%t7^=Ro*HJ(zlU=ixFJ_{EVwK4bDree3p^gALDOgu`Y% zc~g|Ql!Hk+>cJ$Pg$EOST4Ag)7pGGHPa>0#?VlKOoFYc4aq92L{!=U%MQIcZCVr@2 zN)W$yU&y2gX{M|^?r7;~HHu%Y&APT#+pVv}e~+|kQEFELETK8Ii^)Q10amrfT$mus z2@~Lhg&alI|uSCY99NFtKw$lH-_;gnoeWiO}7o^XOCRTfD~{*p-vp_~gRvLRh5 z)RJDtOh{j1CZu03Ga-YV_?DAJh|GlaWpNE;WhoVAQbjew#HJV>%C!hn?ly_wp-n1#Xp@%eq5lz6Gn3*AQ=a)BI%h1(wSso1 z4UGj^nrY^0mpLK7rUCw9P2300?9Q=zYd1!t zdJTV!UdPv_{=~J<+6xD_M^BZ50m{W8883>+v}d$x1uW%XhLBUFcVtEf`fMt-iu zaSz>v-jALG)tmT@r=qegxO;>W$><14tPU5mT(bL7h#pt6DQMMnlebZxPGA!e5N!K7 z1j`k46A9Ic|2fAEBu7BpL2ZE2dUhzI&(Y&gB6g5n*D!nbsle!I>&Xi20N^vm4gUD8Zr!1SVlXzA_owo ziyMHV29tt~!r%yDT2M}hu0U6$xumemp|6AM;T`gVP#?_prOhVBIWOlZ0@ zWkPue7v(}mI!Klilr$0r3UVU-g>oXY8WlN_;S=Np32w6ER8DbnnRbA56*)mVm5kz@ zqG>=DlXX?paZs!g4go|3(cWlZ>Xju)Wowi>ft5XBUQ<+^x)tCZ;!trb5c#vN>O<`{ z^y>WgFYP_Yi*MN5^Pt0Q-wS7!9pcKTpFPxWpsP>)lDT7iIsPh+T-NcaJ8`Yq*>&Tg Yb`8CH>fAgCf`5SGUA+0fc+h{rnVn5`QanxGd*6KXz4yL+x-hr4TEb)Rk!U*Kv!#MS9$JE&w#Q~;)wMYZ<^Tb}~wFDJ^kcK&iC6ry- ziUQB>nmRlArLMI|19nOJK=##EL(RQMBhTJ2Aw4)|gG|(?_8WW!ObxYE9 zdaJ6vBymxV^&-4;lf850g@7o;_ryMXWuNX_QBw&^^ZKKW1gNzQ=ro!U%;z|W1FY8<0wh}Xtou{ zCbzZp-bV=@qksl|v4A=@0)se8!zjb%$(<4QQ#_l#<5((BkoW2y*@5rwuRtxNbw3T` z2tf#difT{>(!qJSbV&7;#iUYg7oI>CYyyp}>KBcJuCN}mrC%d=^>)Lq&ZdB`ejn<0 zq;t#D1qIYe9g%-{iTtQ9zhrteu4mX_NN!XFphEL9T;N&mBO!Z6@YH{xt^l3W_Jr(N z$abjkC6eED^SV>jw{m>s3aru5EsUjKGPz;2c(Ma`sfw`>Us-%mYD#Hxs&8UNyrGGdAc6-` zXfpXGce9A0sg!_oeqM1&VqOW<9@EJMJSC3P7`f1l;Nk_^2vcfkD#VeTpI4HZSDGrt z>snZvm;=>rro;;7P7dS_pWMJZfk&TiSkpGxXQU$ zQu6ash1kpD{qjrV4W+n1R9=2cYP_KmGmt$oHe8RZ31~oKa!F=cDi=&4NQYxyO1zO2 z7l;DuFj8U$vVl4*xn?s8@s!0o=jRodB<7XG8%y!S*75IuPf G;~W5%T7g0U delta 69 zcmV-L0J{I*3djT>2n!MhZEs>}b#pqk1PY`93I=R%XK8eE5(RW`Yh`YcB$boy0%Mb} b11|=G0g3?%A(QX~Ap#O*lVSuxlg9*@_&ydt diff --git a/tests/fixtures/onnx_genai_workflows/decoder/policies/token_state_update.onnx b/tests/fixtures/onnx_genai_workflows/decoder/policies/token_state_update.onnx index 9a8aafaf3172c300217724e5f79d073252f30bbb..5ff11b46b3a74fa10132a719a2dede5cd17cbe8c 100644 GIT binary patch literal 2029 zcmcgt&u`N(6mGK8q;DNm=29X7U6_~{H0qBH!MLz->wz0WDneG8cx|@U35i=a5SQM# zz!A9dcX8zp0ozI1rPV}xD7o10bDrPl_dY-Co2W4KyWS`|{B;A~L*Df%LE0F<^ck)j zAO`BY1lP?CK5uV9k=jEN1-3&_uDt_yg0r6GQ+ggbp%*aA8Q4*T`yL7H(CJ?w_08LM z6^yyR#2z?u{D4H_511%WupQ=&38b+Fi6iRbh5=%bBpT)hpRbF=cSQu3F##N(MhrJx zTt{1D{E9~JMudDIxM`GUh9uQAi&JD>)V}fji~`qYl>RdY+hekzymWcjJ+bL z%&%xa6?I*NzH&lw{0QaoRO0&-Q2}ytqq*glT-MU%iSk*6?fAK|dW7;ff_Xl* zV)y37GUHfs(FB-IIT{N?fXR$w z;>{C?9JwuSqOj$d*y+to@u|(%v#mw^F^7nMPELp;ca@5IVyMDU> diff --git a/tests/fixtures/onnx_genai_workflows/masked/inference_metadata.yaml b/tests/fixtures/onnx_genai_workflows/masked/inference_metadata.yaml index 6a1480251..0fbcfe3de 100644 --- a/tests/fixtures/onnx_genai_workflows/masked/inference_metadata.yaml +++ b/tests/fixtures/onnx_genai_workflows/masked/inference_metadata.yaml @@ -10,6 +10,7 @@ pipeline: - linear_effects - nested_control_flow - typed_emit + - emit_valid_length - loop_induction_values inputs: request.input_ids: diff --git a/tests/fixtures/onnx_genai_workflows/vlm/inference_metadata.yaml b/tests/fixtures/onnx_genai_workflows/vlm/inference_metadata.yaml index a31090a15..ae709cc1a 100644 --- a/tests/fixtures/onnx_genai_workflows/vlm/inference_metadata.yaml +++ b/tests/fixtures/onnx_genai_workflows/vlm/inference_metadata.yaml @@ -84,6 +84,7 @@ pipeline: - nested_control_flow - loop_induction_values - typed_emit + - emit_valid_length - serving_service_contract - bounded_state_recurrence inputs: @@ -127,6 +128,19 @@ pipeline: source: kind: request required: true + request.prompt_lengths: + contract: + dtype: int64 + rank: 1 + shape: + - batch + role: + kind: opaque + source: + kind: application + name: prompt_lengths + required: false + default: -1 package.eos_ids: contract: dtype: int64 @@ -156,7 +170,7 @@ pipeline: dtype: int64 rank: 1 shape: - - 1 + - batch role: kind: opaque source: @@ -280,6 +294,8 @@ pipeline: version: '1' bindings: logits: logits + active: active + done: done token: token parameters: mode: greedy @@ -296,7 +312,10 @@ pipeline: eos_ids: eos_ids iteration: iteration max_iterations: max_iterations + active: active + previous_done: previous_done done: done + next_active: next_active continue: continue token_state_update: implementation: @@ -308,7 +327,12 @@ pipeline: bindings: current: current update: update + lengths: lengths + active: active + done: done next: next + next_lengths: next_lengths + emitted_length: emitted_length last_token_logits: implementation: kind: onnx @@ -348,6 +372,17 @@ pipeline: initializer: decoder.setup.last_logits recurrence: kind: invariant + generated_lengths: + contract: + dtype: int64 + rank: 1 + shape: + - batch + class: semantic + scope: invocation + initializer: initializer.generated_lengths + recurrence: + kind: invariant attention_mask: contract: dtype: int64 @@ -511,11 +546,13 @@ pipeline: component: decoder_state_initializer inputs: prompt_tokens: request.prompt_tokens + prompt_lengths: request.prompt_lengths max_iterations: request.max_iterations outputs: attention_mask: initializer.attention_mask body_attention_mask: initializer.body_attention_mask token_slot: initializer.token_slot + generated_lengths: initializer.generated_lengths cache_lengths: initializer.cache_lengths position_ids: initializer.position_ids body_position_ids: initializer.body_position_ids @@ -551,6 +588,8 @@ pipeline: component: token_sampler inputs: logits: logits + active: active + done: done outputs: token: sample.body - kind: invoke @@ -560,14 +599,19 @@ pipeline: eos_ids: package.eos_ids iteration: loop.iteration max_iterations: request.max_iterations + active: active + previous_done: done 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 @@ -575,6 +619,8 @@ pipeline: inputs: left: package.zero_batch right: package.one + active: active + done: done outputs: total: accepted_len.next - kind: invoke @@ -582,12 +628,19 @@ pipeline: inputs: current: token update: sample.body + lengths: generated_lengths + active: active + done: done outputs: next: token.body + next_lengths: token.next_lengths + emitted_length: token.emitted_length - kind: emit value: token.body output: tokens mode: append + valid_length: token.emitted_length + when: active - kind: invoke component: decoder_step_update inputs: @@ -631,8 +684,10 @@ pipeline: next: decoder.body.last_logits - cell: attention_mask next: decoder_step.body_attention_mask + - cell: generated_lengths + next: token.next_lengths - cell: active - next: loop.continue + next: loop.next_active - cell: done next: loop.done - cell: accepted_len diff --git a/tests/fixtures/onnx_genai_workflows/vlm/policies/cache_length_update.onnx b/tests/fixtures/onnx_genai_workflows/vlm/policies/cache_length_update.onnx index 45764e632822b23fc759eb608a1f8079786bf246..980ab9a3d96fd6966551637a8856faffc0bf4bb9 100644 GIT binary patch delta 330 zcmcc2JcUD_gIkC#H$N$}wAiYVgR7j2B_%&ERfxST-Y>r--av{AMCIkDq{bU4F$38X zO~Q4#ni#p*5|c|Z%Tl>u3PCy?^HSmsrMN&8Scjn!vt!;w;|LY5Ixd!+)U*;V)}qYx zj1sWgl$3ZQpjsdWR%neqzIx_h_xiYBr!*dCp;syC>5f|Scw(P zn`jcjEn1wKlL~fH{KVc8fhZH03x%|}G&ndX3oxqjMFGtM`wlKNnUT@dK1&$k3l1(e d4o)E^E(VAo*d;>ff*@Cc1=&F=Cz~WHY%mHENR@RKcn$ut=+&@pvyhwW%u92nh*PRcci%uh(rh7w@um zqNE6f`oI$)LCZ}bC<6WhF9;#=n|OyaeXH#P( zX>zc%J|Fjs0`{YZ#QhO+ndEFduv_}_fql&nVrZ7JnsqohtG@IpFJC-o&$e#!}SaAba&$-YW_s z6X2e6Xb(X`PB;WZkjw$BHZe}zGEOFK4h^>ox705bOO@d0wFCRoU|?fxEVBO*C|o}Q zoub6`wZ=R|;!JH!jMcOg?N0Sl>%^2==m(>+j~y1V=P^E zfCWcD%zBme8Vjfrj@IW3%w4_%C6j^ZFZ4%)BU^rGtX>^3u;EdRkk)I-lcL0>1bdPT z21E=7vg-i1TUcWCs}xJcXiqF5l#;d_N(rP0C5F;JxCzAGf2W0C`V~xRc0m@WF{SDk-X`5nKzLxOIkBINo2s}Ie)~M zOd@e76Dm$+awpTwiSKwSOConNi7YNtCIBYQ5df2xKmfi~xRXgUD`^SL$+TqVWLiZt zCm;bCaWgG_X*>#soCORAfk*<(c4;KwGm*_Pa>S{K*^%IMN`vz=hZAFv$mJ9pA`LB0MIoG#W+-SK=@zVQZjTm_pFVa?ui5n(@NrxoQ)A&Z6}dh9Rbo6W$3`)J?wZfvMGhs0M3P zQKhE!N4a9$Ep!t`uRr+?=neh#$G(~Cz=tGI%Go3zr}9*q z2kKtyf1;sj{;VVj+v`Gp*%ED08i68EW~JL}Vl5#XsE7koXAo&AL;7CAl*Ud=pULQR zs0s=@v;88ARoNY~CCxd*ZnffjkG5Vzb1+N~X@R$&r>Mw?ORG7Qp+ni83I# zOvv6_i#Cio2`lHP`pDEy6bb33Z>}uZ6J}|VN+d^VW=AQn2KCQG@dP=VWU3rZ`6Lx# zm^2keUXD-=OJ0kKgR5X%1)91Yl?_dPW5|b-1)1O#oue{RWh#}XiV6uiSGHv{%t9(Wj=(xGn`(jLnD=*zU*=CypEHl@C9< zr(HtlANxytIMC0=zlwQ)0lhqR4e$1C&kgyfrW>H*tLy&*q0*x& literal 10398 zcmc&)TW{P%6ppu>WHa4_WKi4GpbA!1fkj%aXJ&6a^&%b`At68rsZ|+yy>8-U<6YKH zloWwbA9yJbAnqWfBE(d2c;> z4}Gbj_33PKIQ6+dIj~(1uO0Ib_ROix+j@O|;Z8;{YAlDNoprS4nuoSGH7y%2?Oa7y zrU$#NiR+$t*36muEo*Fg9^Z59nK`rePI2Y#y`5K(uKX%^2KTJV)b@g3pvU<@3<-!4 z0cj$1&UQ!J+GDNKtfOVmH)sC#5{woAkw_4U0rjidQHyUhw-7ou$454IT;ID6f?F*L)ejmENiSr zMTTDziMdX&^+Oe4=&_8piD_-9Z_aoMP}W?W%fuzqRzxY7q1@k8>12L z=_vRYYO&Wi2iGNOu&-*czYq=nqM=w4lIr?+{g{8~dXd2o^v!e-3Tx0{A^Ff9lWmg#e&t#3q^AvB}iSecKT=TX%>alArBS zlpmVjfnL95dVZw5exn+Unro=yPqtSJEX;la70N*Hi^Iw2l*>2u&D)~@tbZCI95ib2 zoub60I3{VUj!D|NjtMcXu&ObWr;_+jLX(f|pBQqUB1Wls>aWoL6D$~oG>UuUiLYOZ z5x;w1U{ZuMlBztKXzFM=%wNronzp&RQ(KDu9&6R2)UG60LNjU?$wFxfR<*`Vnjniw z6A*(14kIR2!UE0}GZz9RBEAlz9qANalB=rhW>nb~JXlg?kwwX$a#2EPPz%ShAss2z zl5Sd=ke*bTkY2f!3F&9Vw-^>7Rwkq;i>ohJCS;JQOh_4A;zy-2A%oOQ#JYq6nGov| z3S^?#x&(@Vv@D`+T%S%nC#(XTktY@b)F`b8_)3f+L2(q8A{3fOX-z^&X>gQKOEVdO ziWxwz$QP6NDHRARg?>?#cf&wSS*eL1IDV=@p&q9UZ-teeWN@p7c-=CsJ)4hhch}$Zl&ID(Td&YTifha{ zVW2=5;yX%*7rM}e)#{49u?&Y``5dB1d;3jj#EG`X+On-NBV%jKD8bfLxn(WfKc4BeAIZ z%Mx@Y0X!Sd zr(Dav^Rjvli4m&7d#k8Y(}t5wNAn)K3$q_R2dX#m8&5@LneeFrbVP@TNFs5#nAgR- zANuG?9omB4nr!km%CZN{5COrqo?-JaOa>OzhF@qXFBp^rVgFLfAliZtRm2HHn9#%tII~N=ZeN!Ul&Lg# zK|Xz1K~+%IlCq?ezod)fC0oh)VRB1;abJ4MOF?rgJsBHFla~c@42MWK~?n|3L5YPnY=tNYkx@AI>ohc8LcW_Z|ohFN9IYAROR-hm! z(#w?-k=3Zki44w=6SUOg9Vc>%+F8j2NLP^)G;rco!K7#!&^Jk^D(X6D4+~BJiAl6K z+LLBwNmAJwHF5o!Y5BPF>1PkM8|dm&f61K@z8w83O6KO#i8FT0 d+39uTp>_?ulKTw%{(al>!!H1Df?sWJ{|8o9-dg|w diff --git a/tests/fixtures/onnx_genai_workflows/vlm/policies/termination.onnx b/tests/fixtures/onnx_genai_workflows/vlm/policies/termination.onnx index 72763e63ade7f01e72cff99022dea83883b2f75b..e8cf5a74afb81d1e02e4e6c620c8dbdce8d1e7aa 100644 GIT binary patch delta 1001 zcmZuw&ui2`6mHTrnM^iS$F;?vb`4v9;1)O8#&s(SrD6|K1wm+w5aMPIS!g1gAM_%^ ziXJ>H>fAgCf`5SGUA+0fc+h{rnVn5`QanxGd*6KXz4yL+x-hr4TEb)Rk!U*Kv!#MS9$JE&w#Q~;)wMYZ<^Tb}~wFDJ^kcK&iC6ry- ziUQB>nmRlArLMI|19nOJK=##EL(RQMBhTJ2Aw4)|gG|(?_8WW!ObxYE9 zdaJ6vBymxV^&-4;lf850g@7o;_ryMXWuNX_QBw&^^ZKKW1gNzQ=ro!U%;z|W1FY8<0wh}Xtou{ zCbzZp-bV=@qksl|v4A=@0)se8!zjb%$(<4QQ#_l#<5((BkoW2y*@5rwuRtxNbw3T` z2tf#difT{>(!qJSbV&7;#iUYg7oI>CYyyp}>KBcJuCN}mrC%d=^>)Lq&ZdB`ejn<0 zq;t#D1qIYe9g%-{iTtQ9zhrteu4mX_NN!XFphEL9T;N&mBO!Z6@YH{xt^l3W_Jr(N z$abjkC6eED^SV>jw{m>s3aru5EsUjKGPz;2c(Ma`sfw`>Us-%mYD#Hxs&8UNyrGGdAc6-` zXfpXGce9A0sg!_oeqM1&VqOW<9@EJMJSC3P7`f1l;Nk_^2vcfkD#VeTpI4HZSDGrt z>snZvm;=>rro;;7P7dS_pWMJZfk&TiSkpGxXQU$ zQu6ash1kpD{qjrV4W+n1R9=2cYP_KmGmt$oHe8RZ31~oKa!F=cDi=&4NQYxyO1zO2 z7l;DuFj8U$vVl4*xn?s8@s!0o=jRodB<7XG8%y!S*75IuPf G;~W5%T7g0U delta 69 zcmV-L0J{I*3djT>2n!MhZEs>}b#pqk1PY`93I=R%XK8eE5(RW`Yh`YcB$boy0%Mb} b11|=G0g3?%A(QX~Ap#O*lVSuxlg9*@_&ydt diff --git a/tests/fixtures/onnx_genai_workflows/vlm/policies/token_state_update.onnx b/tests/fixtures/onnx_genai_workflows/vlm/policies/token_state_update.onnx index 9a8aafaf3172c300217724e5f79d073252f30bbb..5ff11b46b3a74fa10132a719a2dede5cd17cbe8c 100644 GIT binary patch literal 2029 zcmcgt&u`N(6mGK8q;DNm=29X7U6_~{H0qBH!MLz->wz0WDneG8cx|@U35i=a5SQM# zz!A9dcX8zp0ozI1rPV}xD7o10bDrPl_dY-Co2W4KyWS`|{B;A~L*Df%LE0F<^ck)j zAO`BY1lP?CK5uV9k=jEN1-3&_uDt_yg0r6GQ+ggbp%*aA8Q4*T`yL7H(CJ?w_08LM z6^yyR#2z?u{D4H_511%WupQ=&38b+Fi6iRbh5=%bBpT)hpRbF=cSQu3F##N(MhrJx zTt{1D{E9~JMudDIxM`GUh9uQAi&JD>)V}fji~`qYl>RdY+hekzymWcjJ+bL z%&%xa6?I*NzH&lw{0QaoRO0&-Q2}ytqq*glT-MU%iSk*6?fAK|dW7;ff_Xl* zV)y37GUHfs(FB-IIT{N?fXR$w z;>{C?9JwuSqOj$d*y+to@u|(%v#mw^F^7nMPELp;ca@5IVyMDU> From c3028e070ab1a20e3629f49a325e248b336e60e8 Mon Sep 17 00:00:00 2001 From: justinchuby Date: Fri, 14 Aug 2026 06:44:42 +0000 Subject: [PATCH 080/151] Upgrade batched generation policy contracts to v2 Publish per-row sampler, ragged termination, and selective state-update contracts for decoder and VLM workflows. Add dynamic EOS and iteration initialization artifacts plus runtime parity coverage for heterogeneous rows. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: justinchuby --- src/mobius/generation/__init__.py | 6 +- src/mobius/generation/_policy_components.py | 92 +++++- .../generation/_policy_components_test.py | 63 +++- .../onnx_genai/auto_export_test.py | 28 +- .../onnx_genai/inference_metadata.py | 20 +- .../onnx_genai/workflow_metadata.py | 274 ++++++++++++++++-- .../decoder/inference_metadata.yaml | 227 ++++++++++++++- .../decoder/policies/iteration_broadcast.onnx | Bin 0 -> 611 bytes .../decoder/policies/termination.onnx | Bin 4866 -> 6378 bytes .../termination_batch_initializer.onnx | Bin 0 -> 2003 bytes .../decoder/policies/token_sampler.onnx | Bin 1375 -> 59345 bytes .../decoder/policies/token_state_update.onnx | Bin 2029 -> 2081 bytes .../vlm/inference_metadata.yaml | 219 +++++++++++++- .../vlm/policies/iteration_broadcast.onnx | Bin 0 -> 611 bytes .../vlm/policies/termination.onnx | Bin 4866 -> 6378 bytes .../termination_batch_initializer.onnx | Bin 0 -> 2003 bytes .../vlm/policies/token_sampler.onnx | Bin 1375 -> 59345 bytes .../vlm/policies/token_state_update.onnx | Bin 2029 -> 2081 bytes 18 files changed, 862 insertions(+), 67 deletions(-) create mode 100644 tests/fixtures/onnx_genai_workflows/decoder/policies/iteration_broadcast.onnx create mode 100644 tests/fixtures/onnx_genai_workflows/decoder/policies/termination_batch_initializer.onnx create mode 100644 tests/fixtures/onnx_genai_workflows/vlm/policies/iteration_broadcast.onnx create mode 100644 tests/fixtures/onnx_genai_workflows/vlm/policies/termination_batch_initializer.onnx diff --git a/src/mobius/generation/__init__.py b/src/mobius/generation/__init__.py index 2cca64814..f34a6c97b 100644 --- a/src/mobius/generation/__init__.py +++ b/src/mobius/generation/__init__.py @@ -25,8 +25,8 @@ build_grammar_logits_processor, build_greedy_sampler, build_integer_add, - build_selective_integer_add, build_integer_minimum, + build_integer_row_broadcast, build_iteration_cast, build_last_token_logits, build_masked_token_update, @@ -35,9 +35,11 @@ build_schedule_constant, build_schedule_lookup, build_seeded_categorical_sampler, + build_selective_integer_add, build_sequence_length, build_speculative_acceptance, build_speculative_state_rollback, + build_termination_batch_initializer, build_token_block_identity, build_token_state_update, build_token_to_slot, @@ -66,6 +68,7 @@ "build_grammar_logits_processor", "build_greedy_sampler", "build_integer_add", + "build_integer_row_broadcast", "build_selective_integer_add", "build_integer_minimum", "build_iteration_cast", @@ -82,6 +85,7 @@ "build_token_block_identity", "build_token_state_update", "build_token_to_slot", + "build_termination_batch_initializer", "build_tts_decoder_state_initializer", "build_tts_decoder_step_update", "build_tts_state_initializer", diff --git a/src/mobius/generation/_policy_components.py b/src/mobius/generation/_policy_components.py index d2f7cbefe..b15876f2b 100644 --- a/src/mobius/generation/_policy_components.py +++ b/src/mobius/generation/_policy_components.py @@ -158,11 +158,12 @@ def build_greedy_sampler( token_ids = sampled builder.add_output(token_ids, "token") return _component( - "onnx-genai.token-sampler@1", + "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", @@ -241,6 +242,45 @@ def build_selective_integer_add() -> PolicyComponent: 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") @@ -1307,11 +1347,13 @@ def build_seeded_categorical_sampler() -> PolicyComponent: builder.add_output(token_ids, "token") builder.add_output(next_offset, "next_offset") return _component( - "onnx-genai.token-sampler@1", + "onnx-genai.token-sampler@2", graph, { "role": "token_sampler", "mode": "seeded_stochastic", + "batching": "per_row", + "inactive_rows": "preserve", "logits": "logits", "token": "token", "temperature": "temperature", @@ -1337,20 +1379,39 @@ def build_eos_termination(*, row_selective: bool = False) -> PolicyComponent: graph, builder = _make_graph("eos_termination") op = builder.op token_ids = builder.input("token_ids", ir.DataType.INT64, ["batch"]) - eos_ids = builder.input("eos_ids", ir.DataType.INT64, ["num_eos"]) + 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"], + ["batch"], ) max_iterations = builder.input( "max_iterations", ir.DataType.INT64, - [1] if row_selective else ["batch"], + ["batch"], ) tokens = op.Unsqueeze(token_ids, op.Constant(value_ints=[-1])) - eos = op.Unsqueeze(eos_ids, op.Constant(value_ints=[0])) + 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], @@ -1381,7 +1442,11 @@ def build_eos_termination(*, row_selective: bool = False) -> PolicyComponent: builder.add_output(next_active, "next_active") builder.add_output(continued, "continue") return _component( - "onnx-genai.termination-predicate@1", + ( + "onnx-genai.termination-predicate@2" + if row_selective + else "onnx-genai.termination-predicate@1" + ), graph, { "role": "termination_predicate", @@ -1390,7 +1455,15 @@ def build_eos_termination(*, row_selective: bool = False) -> PolicyComponent: "iteration": "iteration", "max_iterations": "max_iterations", **( - {"active": "active", "previous_done": "previous_done"} if row_selective else {} + { + "eos_lengths": "eos_lengths", + "active": "active", + "previous_done": "previous_done", + "batching": "per_row", + "inactive_rows": "preserve", + } + if row_selective + else {} ), "done": "done", **({"next_active": "next_active"} if row_selective else {}), @@ -1756,12 +1829,13 @@ def build_token_state_update(*, row_selective: bool = False) -> PolicyComponent: builder.add_output(next_lengths, "next_lengths") builder.add_output(emitted_length, "emitted_length") return _component( - "onnx-genai.state-update@1", + "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 {}), **( {"lengths": "lengths", "active": "active", "done": "done"} if row_selective diff --git a/src/mobius/generation/_policy_components_test.py b/src/mobius/generation/_policy_components_test.py index 2781ff94b..dc628ccfa 100644 --- a/src/mobius/generation/_policy_components_test.py +++ b/src/mobius/generation/_policy_components_test.py @@ -24,6 +24,7 @@ build_grammar_logits_processor, build_greedy_sampler, build_integer_minimum, + build_integer_row_broadcast, build_last_token_logits, build_masked_token_update, build_model_token_cast, @@ -31,6 +32,7 @@ build_seeded_categorical_sampler, build_speculative_acceptance, build_speculative_state_rollback, + build_termination_batch_initializer, build_token_state_update, ) from mobius.generation._policy_components import _make_graph @@ -633,9 +635,10 @@ def test_row_selective_state_and_termination_preserve_inactive_rows(tmp_path): tmp_path, { "token_ids": np.array([2, 8, 9], np.int64), - "eos_ids": np.array([2, 9], np.int64), - "iteration": np.array([0], np.int64), - "max_iterations": np.array([5], 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, 0, 0], np.int64), + "max_iterations": np.array([5, 5, 5], np.int64), "active": np.array([True, False, True], np.bool_), "previous_done": np.array([False, False, True], np.bool_), }, @@ -645,6 +648,60 @@ def test_row_selective_state_and_termination_preserve_inactive_rows(tmp_path): 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 = { + "token_ids": 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([0, 1, 4, 2], np.int64), + "max_iterations": np.array([5, 5, 10, 3], np.int64), + "active": np.array([True, True, True, True], np.bool_), + "previous_done": np.array([False, False, False, False], 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[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(), diff --git a/src/mobius/integrations/onnx_genai/auto_export_test.py b/src/mobius/integrations/onnx_genai/auto_export_test.py index 26b29a606..46b44d863 100644 --- a/src/mobius/integrations/onnx_genai/auto_export_test.py +++ b/src/mobius/integrations/onnx_genai/auto_export_test.py @@ -173,11 +173,19 @@ def test_dispatch_decoder(tmp_path): for name, value in workflow["inputs"].items() if value["source"]["kind"] == "application" } - assert application_inputs == {"request.prompt_lengths"} + assert application_inputs == { + "request.prompt_lengths", + "request.eos_ids", + "request.eos_lengths", + "request.row_max_iterations", + "request.grammar_mask", + "request.rng_offset", + } 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"] @@ -214,6 +222,12 @@ def test_seeded_decoder_sampler_uses_request_controls_and_direct_kv_carry(): )["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", @@ -249,8 +263,18 @@ def test_seeded_decoder_sampler_uses_request_controls_and_direct_kv_carry(): 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_offset"]["class"] == "semantic" assert workflow["state"]["rng_offset"]["initializer"] == "request.rng_offset" + assert workflow["components"]["termination"]["contract"]["version"] == "2" + assert workflow["components"]["termination"]["contract"]["parameters"] == { + "batching": "per_row", + "inactive_rows": "preserve", + } + assert workflow["components"]["token_state_update"]["contract"]["version"] == "2" assert not any("kv_update" in name for name in workflow["components"]) @@ -451,7 +475,7 @@ def test_dispatch_vision_multimodal_pipeline(tmp_path): 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"][3]["component"] == "embedding" + assert workflow["steps"][0]["setup"][4]["component"] == "embedding" assert workflow["steps"][0]["iteration"]["value"] == "loop.iteration" assert workflow["state"]["logits"]["contract"] == { "dtype": "float32", diff --git a/src/mobius/integrations/onnx_genai/inference_metadata.py b/src/mobius/integrations/onnx_genai/inference_metadata.py index 432cdb200..dc251637c 100644 --- a/src/mobius/integrations/onnx_genai/inference_metadata.py +++ b/src/mobius/integrations/onnx_genai/inference_metadata.py @@ -1413,7 +1413,16 @@ def semantic_contract(component: Any) -> dict[str, Any]: bindings = { key: value for key, value in contract.items() - if key not in {"role", "mode", "effect", "rng", "state_class"} + if key + not in { + "role", + "mode", + "effect", + "rng", + "state_class", + "batching", + "inactive_rows", + } and isinstance(value, str) } rng = contract.get("rng") @@ -1426,8 +1435,13 @@ def semantic_contract(component: Any) -> dict[str, Any]: "version": version, "bindings": bindings, } - if "mode" in contract: - declaration["parameters"] = {"mode": contract["mode"]} + parameters = { + key: contract[key] + for key in ("mode", "batching", "inactive_rows") + if key in contract + } + if parameters: + declaration["parameters"] = parameters return declaration for name, component in policy_components.items(): diff --git a/src/mobius/integrations/onnx_genai/workflow_metadata.py b/src/mobius/integrations/onnx_genai/workflow_metadata.py index a50c25e5d..8878e99ec 100644 --- a/src/mobius/integrations/onnx_genai/workflow_metadata.py +++ b/src/mobius/integrations/onnx_genai/workflow_metadata.py @@ -27,17 +27,19 @@ build_euler_solver_step, build_greedy_sampler, build_integer_add, - build_selective_integer_add, build_integer_minimum, + build_integer_row_broadcast, build_last_token_logits, build_model_token_cast, build_proposal_metrics, build_schedule_constant, build_schedule_lookup, build_seeded_categorical_sampler, + build_selective_integer_add, build_sequence_length, - build_token_to_slot, + build_termination_batch_initializer, build_token_state_update, + build_token_to_slot, build_tts_decoder_state_initializer, build_tts_decoder_step_update, build_tts_state_initializer, @@ -2688,8 +2690,13 @@ def build_vlm_workflow_metadata( fixed_capacity=fixed_capacity, ), ) - pkg.add_policy_component("token_sampler", build_greedy_sampler(row_selective=True)) + 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("iteration_broadcast", build_integer_row_broadcast()) pkg.add_policy_component( "token_state_update", build_token_state_update(row_selective=True), @@ -2735,6 +2742,27 @@ def build_vlm_workflow_metadata( "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"}, @@ -2791,6 +2819,79 @@ def build_vlm_workflow_metadata( "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.grammar_mask": { + "contract": { + "dtype": "bool", + "rank": 2, + "shape": [batch, _contract(logits_output)["shape"][-1]], + }, + "role": {"kind": "opaque"}, + "source": {"kind": "application", "name": "grammar_mask"}, + "required": False, + "default": True, + }, + "request.rng_offset": { + "contract": batch_int, + "role": {"kind": "opaque"}, + "source": {"kind": "application", "name": "rng_offset"}, + "required": False, + "default": 0, + }, + } + ) vision_invoke_inputs = { name: preprocessing_values[name] for name in vision_inputs @@ -2941,6 +3042,13 @@ def build_vlm_workflow_metadata( ), "recurrence": {"kind": "invariant"}, }, + "rng_offset": { + "contract": batch_int, + "class": "semantic", + "scope": "invocation", + "initializer": "request.rng_offset", + "recurrence": {"kind": "invariant"}, + }, } state_specs = [ ( @@ -2975,6 +3083,13 @@ def build_vlm_workflow_metadata( "token.next_lengths", "state.generated_lengths.final", ), + ( + "rng_offset", + "request.rng_offset", + "state.rng_offset.body", + "sample.next_offset", + "state.rng_offset.final", + ), ( "active", "package.active", @@ -3162,6 +3277,21 @@ def build_vlm_workflow_metadata( **{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, @@ -3198,23 +3328,36 @@ def build_vlm_workflow_metadata( body = { "kind": "sequence", "nodes": [ + _invoke( + "iteration_broadcast", + {"value": "loop.iteration", "active": "state.active.body"}, + {"rows": "loop.iteration_rows"}, + ), _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", + "grammar_mask": "request.grammar_mask", + "seed": "request.seed", + "offset": "state.rng_offset.body", "active": "state.active.body", "done": "state.done.body", }, - {"token": "sample.body"}, + {"token": "sample.body", "next_offset": "sample.next_offset"}, {"sample": _effect("sample.0", "sample.1")}, ), _invoke( "termination", { "token_ids": "sample.body", - "eos_ids": "package.eos_ids", - "iteration": "loop.iteration", - "max_iterations": "request.max_iterations", + "eos_ids": "termination.eos_ids", + "eos_lengths": "termination.eos_lengths", + "iteration": "loop.iteration_rows", + "max_iterations": "termination.max_iterations", "active": "state.active.body", "previous_done": "state.done.body", }, @@ -4443,15 +4586,48 @@ def build_decoder_workflow_metadata( } ) if cache_pairs: - workflow_inputs["request.prompt_lengths"] = { - "contract": batch_int, - "role": {"kind": "opaque"}, - "source": {"kind": "application", "name": "prompt_lengths"}, - "required": False, - "default": -1, - } + 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" - if stochastic_sampler: + sampler_with_rng = stochastic_sampler or bool(cache_pairs) + if sampler_with_rng: workflow_inputs.update( { "request.temperature": { @@ -4481,7 +4657,7 @@ def build_decoder_workflow_metadata( }, "source": {"kind": "request", "field": "sampling_top_k"}, "required": False, - "default": 0, + "default": 0 if stochastic_sampler else 1, }, "request.top_p": { "contract": { @@ -4524,7 +4700,8 @@ def build_decoder_workflow_metadata( "role": "seed", }, "source": {"kind": "request", "field": "seed"}, - "required": True, + "required": False, + "default": 0, }, "request.grammar_mask": { "contract": { @@ -4537,7 +4714,8 @@ def build_decoder_workflow_metadata( }, "role": {"kind": "opaque"}, "source": {"kind": "application", "name": "grammar_mask"}, - "required": True, + "required": False, + "default": True, }, "request.rng_offset": { "contract": batch_int, @@ -4552,13 +4730,14 @@ def build_decoder_workflow_metadata( pkg.add_policy_component("cache_length_update", build_selective_integer_add()) pkg.add_policy_component( "token_sampler", - ( - build_greedy_sampler(row_selective=True) - if not stochastic_sampler - else build_seeded_categorical_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("iteration_broadcast", build_integer_row_broadcast()) pkg.add_policy_component( "token_state_update", build_token_state_update(row_selective=True), @@ -4780,7 +4959,7 @@ def build_decoder_workflow_metadata( }, ] ) - if stochastic_sampler: + if sampler_with_rng: state["rng_offset"] = { "contract": batch_int, "scope": "invocation", @@ -4936,6 +5115,27 @@ def build_decoder_workflow_metadata( }, ), _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"}, @@ -4966,6 +5166,17 @@ def build_decoder_workflow_metadata( body = { "kind": "sequence", "nodes": [ + *( + [ + _invoke( + "iteration_broadcast", + {"value": "loop.iteration", "active": "state.active.body"}, + {"rows": "loop.iteration_rows"}, + ) + ] + if cache_pairs + else [] + ), _invoke( "token_sampler", { @@ -4980,7 +5191,7 @@ def build_decoder_workflow_metadata( "seed": "request.seed", "offset": "state.rng_offset.body", } - if stochastic_sampler + if sampler_with_rng else {} ), **( @@ -4994,7 +5205,7 @@ def build_decoder_workflow_metadata( }, { "token": "sample.body", - **({"next_offset": "sample.next_offset"} if stochastic_sampler else {}), + **({"next_offset": "sample.next_offset"} if sampler_with_rng else {}), }, {"sample": _effect("sample.0", "sample.1")}, ), @@ -5041,9 +5252,14 @@ def build_decoder_workflow_metadata( "termination", { "token_ids": "sample.body", - "eos_ids": "package.eos_ids", - "iteration": "loop.iteration", - "max_iterations": "request.max_iterations", + "eos_ids": ("termination.eos_ids" if cache_pairs else "package.eos_ids"), + **({"eos_lengths": "termination.eos_lengths"} if cache_pairs else {}), + "iteration": ("loop.iteration_rows" if cache_pairs else "loop.iteration"), + "max_iterations": ( + "termination.max_iterations" + if cache_pairs + else "request.max_iterations" + ), **( { "active": "state.active.body", diff --git a/tests/fixtures/onnx_genai_workflows/decoder/inference_metadata.yaml b/tests/fixtures/onnx_genai_workflows/decoder/inference_metadata.yaml index e2af713d3..0c03c60c9 100644 --- a/tests/fixtures/onnx_genai_workflows/decoder/inference_metadata.yaml +++ b/tests/fixtures/onnx_genai_workflows/decoder/inference_metadata.yaml @@ -91,6 +91,143 @@ pipeline: 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: 127 + request.eos_lengths: + contract: + dtype: int64 + rank: 1 + shape: + - batch + role: + kind: opaque + source: + kind: application + name: eos_lengths + required: false + default: 1 + request.row_max_iterations: + contract: + dtype: int64 + rank: 1 + shape: + - batch + role: + kind: opaque + source: + kind: application + name: row_max_iterations + required: false + default: -1 + request.temperature: + contract: + dtype: float32 + rank: 1 + shape: + - batch + 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 + 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 + 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 + 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 + role: + kind: runtime + version: '1.0' + role: seed + source: + kind: request + required: false + default: 0 + request.grammar_mask: + contract: + dtype: bool + rank: 2 + shape: + - batch + - 128 + role: + kind: opaque + source: + kind: application + name: grammar_mask + required: false + default: true + request.rng_offset: + contract: + dtype: int64 + rank: 1 + shape: + - batch + role: + kind: opaque + source: + kind: application + name: rng_offset + required: false + default: 0 package.active: contract: dtype: bool @@ -172,14 +309,24 @@ pipeline: artifact: policies/token_sampler.onnx contract: id: onnx-genai.token-sampler - version: '1' + version: '2' bindings: logits: logits + token: token + temperature: temperature + top_k: top_k + top_p: top_p + min_p: min_p + grammar_mask: grammar_mask active: active done: done - token: token + rng_seed: seed + rng_offset: offset + rng_next_offset: next_offset parameters: - mode: greedy + mode: seeded_stochastic + batching: per_row + inactive_rows: preserve application_overridable: true termination: implementation: @@ -187,24 +334,28 @@ pipeline: artifact: policies/termination.onnx contract: id: onnx-genai.termination-predicate - version: '1' + version: '2' bindings: tokens: token_ids eos_ids: eos_ids iteration: iteration max_iterations: max_iterations + eos_lengths: eos_lengths active: active previous_done: previous_done 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 contract: id: onnx-genai.state-update - version: '1' + version: '2' bindings: current: current update: update @@ -214,6 +365,9 @@ pipeline: next: next next_lengths: next_lengths emitted_length: emitted_length + parameters: + batching: per_row + inactive_rows: preserve last_token_logits: implementation: kind: onnx @@ -230,6 +384,14 @@ pipeline: implementation: kind: onnx artifact: policies/cache_length_update.onnx + termination_batch_initializer: + implementation: + kind: onnx + artifact: policies/termination_batch_initializer.onnx + iteration_broadcast: + implementation: + kind: onnx + artifact: policies/iteration_broadcast.onnx state: token: contract: @@ -319,6 +481,17 @@ pipeline: initializer: initializer.cache_lengths recurrence: kind: invariant + rng_offset: + contract: + dtype: int64 + rank: 1 + shape: + - batch + scope: invocation + class: semantic + initializer: request.rng_offset + recurrence: + kind: invariant attention_mask: contract: dtype: int64 @@ -341,7 +514,7 @@ pipeline: initializer: initializer.body_position_ids recurrence: kind: invariant - cache_9: + cache_10: contract: dtype: float32 rank: 4 @@ -374,7 +547,7 @@ pipeline: storage: shared_buffer ports: model: - cache_9: + cache_10: input: past_key_values.0.key output: present.0.key steps: @@ -405,6 +578,18 @@ pipeline: 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: @@ -412,14 +597,29 @@ pipeline: outputs: last_logits: decoder.setup.last_logits steps: + - kind: invoke + component: iteration_broadcast + inputs: + value: loop.iteration + active: active + outputs: + rows: loop.iteration_rows - 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 + grammar_mask: request.grammar_mask + seed: request.seed + offset: rng_offset active: active done: done outputs: token: sample.body + next_offset: sample.next_offset - kind: invoke component: token_state_update inputs: @@ -436,9 +636,10 @@ pipeline: component: termination inputs: token_ids: sample.body - eos_ids: package.eos_ids - iteration: loop.iteration - max_iterations: request.max_iterations + eos_ids: termination.eos_ids + eos_lengths: termination.eos_lengths + iteration: loop.iteration_rows + max_iterations: termination.max_iterations active: active previous_done: done outputs: @@ -482,7 +683,7 @@ pipeline: component: model inputs: input_ids: token.body - past_key_values.0.key: cache_9 + past_key_values.0.key: cache_10 attention_mask: decoder_step.body_attention_mask position_ids: position_ids outputs: @@ -513,11 +714,13 @@ pipeline: next: accepted_len.next - cell: slot_ids next: slot_ids + - cell: rng_offset + next: sample.next_offset - cell: attention_mask next: decoder_step.body_attention_mask - cell: position_ids next: decoder_step.body_position_ids - - cell: cache_9 + - cell: cache_10 next: decoder.body.present.0.key termination: generation_eos iteration: diff --git a/tests/fixtures/onnx_genai_workflows/decoder/policies/iteration_broadcast.onnx b/tests/fixtures/onnx_genai_workflows/decoder/policies/iteration_broadcast.onnx new file mode 100644 index 0000000000000000000000000000000000000000..6618ad65560e12af5e2f95f56cf1d059c7b64cda GIT binary patch literal 611 zcmcIh!AiqG6l_A8rcVp9B5Ds(OHuIedT1XhBnu z2{VdNhH8Tz9H^&pD213As=2ZuO^Go$;Y<@vlR0wl9|j$0e`75*u%S{hvt+=#1NDrg zQwA#o$kX!Kv|G0`MMR7XAN{Uh!hbS?cDS%?%rwq#<48*qQ)29JyRwG>f}nvsa06#} z@s-*soCHnOgD&bQnQ}T0xy=KxT#A)Q7!ga8=p2f_&~m5ox1Gn`^;p&{ESzIFENcwM zOjIYs(gUVaSe^HZ+?Usda}B*;nV83DzhYI7km-!4oap&gf8iWK=SK}T8Zl~%I&VCv HZS8#mtB=JC literal 0 HcmV?d00001 diff --git a/tests/fixtures/onnx_genai_workflows/decoder/policies/termination.onnx b/tests/fixtures/onnx_genai_workflows/decoder/policies/termination.onnx index e8cf5a74afb81d1e02e4e6c620c8dbdce8d1e7aa..46ce9432722602b317df0a22c323084f5f063c85 100644 GIT binary patch literal 6378 zcmc&(+in|081^~`+Y>hlQ`%aq0%28gVYR_Kb{xAGsI)>z;ZUg+H$`Z)Sx=Kq<6U=m zZCX%oB7p=5kf1jpgu(;B>+m2v0Q@tvJD%C3G&kG&W@r96ee?Y@|6yET!OP=lH<%{Z z|8~)D7H-byttd=VFHC92UQ^{T8Zp{&>ban8BQ*04r;G++n%wv||5?XYZRZ?ngx;7X z6VGS3vU3H!K6$XW8HM4a#E*kXy6GQ!NkaDn7JIS3{}@+4{cL9g*@a$w4x0GUgeANO zz(b8RdcZ;&j1puZ#ph=9lQ4NWW$X(^U3-0A5y)Mqp(JI&`_V!s3RGiJ!m;AUjK249 z>K#&IH)TPfBF>Vim4)ahQ=i?UKD(Me*HJ_1vsF2^s%3%t)iQXwyYEdH?b%IP5WMu9 zB~fc5w7|kqs|s-1O|%5ZGEQ4?#7YDA6pq-8W|$E2`Bk9qh<5M_h_mChWv}K%0e8!( zW%aT+D=4yGRAm2IiY$|>W6QXMU8GZFH1+ufxZ{O;jBeXaSrE)^J4>Qg7H@FY<|`ZO zG;ZqPguNKN0DA%EY=s+$Ll*9(`-uk6HO<-Z1iax3fCpE>B+0A}xlTW1V8-2!-4q3( zQn%wQ@>*Fc{YI!Xv!-0T83tUR(RafUbzOUfe#9~(`5c6U0aZr6TgUXviZPD6>>YpDK!vB@YH zC#5=qrxrivc|bwkp1rOr5?1e7X01&B^t*)$eDQ;?khwh`BfTFpFJ&=xx9qjNEHQ&K z$OL7bx*~7=6TH0?ckn$NzW_M#(jW>=%KGARj^J``!{vmpq-q;3=cu&H`Io$$VpX&{ z_8w7%F2U!O3U9~Kmbe4^Vy3uUBvB5WRa0J;UXQ@wPGU9-qG>`$QOG!JA4gzOhV}~o zN<2en0V-vI_{u`do}UIY#!Z`CJvXhSqnTFHDb=+8EP{*tYPnnJG>zffjNXnZacu-2 zK@V}A>g~8JcD_~ypL9^ZxT1&X^pH~N0U+jC$sG|Pc#)qSjgfAdj6Bg6&JGOeX~&iH zNIzZoB}ERmu+#&}j>+h=BCe> ztGN>tYc!riO1EH(&oXGYjf9;(Q(quTw<~-H&%v~N5vlWu5voJ~MO3X@2Q4iz=!llu zDJ`ziVjIc4IqEetD|B%aR;%GXbF9wJU z#`A*eZ_d+K@(WtxkLzd|Ca=P>1-JtR7q~guJXFovZt{FHXI1V7e2>gNja`b}WEq*- z4Ig07t&7abG>n%FO9S~v&P|JRHV=u!kcQY+VaS|@SA}~PF|fA68>dH8oF+M+7?uZr zrj+lJbhNe4Yk|-yU5VCJO-p*3RH#aj1WszMs>^eVq{C7K@^ftN7|d}-(vkJc0(w(E zX>abaFv~|ar5U2!X2%d0;D4%LI&pN!oa-BoX67ga;G@8Y+|IcQ4g9e=5nb|UXD-fO zb-6q*f>hZhE}vJRWo+agytyFB=FA>~@_F+rw8}_^H)JJ{8L0wqG6m2>nCu|j6OE9p z#?j{-J~#~#;PblJ91KiYJQG{yu_Ny#* literal 4866 zcmcgwPj4GV6!*FzUXRlxokD9?p-NYQL)04YI*#L>N)bY`1d%{og3xHQo;I7Oj+An%iqIj$XVMqrDEG6Waf>goCRT;Jp5PwjF(1jaGlgcZ_Ja4 z=W|*aY>_*YlP9ex3{Mk34kl^KA9+c_jshNgv48Z8R=@scut}`KEc_3a_|b$XI0N93 zdK#VZkOjkpXh`Yx8T%?so=!RcmNUm%UkC!ZW7l&{S@3?**>j8Rxe9wXH1;YblJ>hn zYCKBLQReEDb|RFW)R(Qy4a!h8T;hv)EluKYbART=G?L05>i>F{1H4Y-UK-W1!2D^D904IdMkjNItc?czf)A%Y|8hEJyVm+b z5D0hedal`CBjyR`lVLDUn&57wfqkI|L9ZADJtYX%fn=FNN(BWmpmU`qj~;Ez*rze~ zQXaECYgN7q68G%dTrSJJ1;+YC#`>2r2C23*GyVht5C9)u8bqPSRv*LRqv4P_ZL5J# zL6g(AuYph&M}LVB$aW>GWABvZ;1YT>x88o7$sIOA6r^m)K=L#eow#0ICKS8*61>+Aen`w?*A?AGlt|8Al6r6yShC{?jUy+ii^-V%)99%{K33+fjv=;sT> z-dT5+?~UiYdvU3kuyM-Ud8&7njf7Uzz>cE5snlsWM22ZGwW~~B{B2LszJuh+9J=5E zd|S__+97Tiv&EdIGvn_nxto1FX+7cLyjgB#(Zw*O#89^AtLk@l9F4gBVAED2ltFtq z41BnB*!N%nKWijniu}NV>iNYT_3{<8dGCw*g#aUZXNgR1fiM zi?|fu*P6+KI|FI8S7NKzf;uN->bu0GsMOnr+wRFocQ~K;rpH{Nn>O)%Y%%H)&$Onf z{PtaM3$&b96a)LF<5EBJfKvrELwx9PUl1;EM_W5v$N4*Lb-uB?4)u`O0mY1q&C++J zTolv~1+Fvh1_F=C1T!IjNic~B7f?|n#4kNcj;C-!?u42Odb~Q1J0{p2!R-grnX86$RMYXl&8;g1r$G8XWmDnnZ z_NxM|a+n!*f31)X>(B29ygXrU}i1wmRP^3g}dUYv({V{>hsM zDM~$jLAtUs-97Z{%8b6)IBpl01;@O-O~R)V&NPc_7=7qWxO5 zKU40(on&n4Bg)f(peJRDO@|9l&QVnU^m(Q{h3&t3$_|5Wre0~SlVHMZM!e~h&P>^X Ut%Wh@a7axNkI!cnD6Mb)2KtqxhyVZp literal 0 HcmV?d00001 diff --git a/tests/fixtures/onnx_genai_workflows/decoder/policies/token_sampler.onnx b/tests/fixtures/onnx_genai_workflows/decoder/policies/token_sampler.onnx index 29209ca573fdea2268d08475ddd411a15201d9b7..af9f73753fbbd81fe648083be389f699c689b6e1 100644 GIT binary patch literal 59345 zcmd5_TaP5kRi2t_PtDcky^Pz&UKbWZjXkQoRc2a9c-Jfoi|rL&BcrvfQcZWwOz%!t zwYzG@Ymfz6ghAXfmH`{^U~{#&%SP~E3Gu+I9}zzQiAUafSiXqNh&bm&=6D~7sOguk zjI4@0aZa2Q=MrD!&ZW`$qt&g$`|A&W=euWypB#7o$@zzHj*M@6Ffi+1c_J>;2X7Vy(V0d}A;?Ki}URK3ELRvlS7~E)QnO?oj>HL?q)g z_#BFeUm6V0t{01g(awX}hv)0fEZx1N|LbR`6*pcfa|+|Nh_ajK{3PGHXhiHI2xup`FwAQ$D)-VD{YM=Cgu8b$B}S69 z{gf)Io;`rAoa|mw{~KcolZ}f|V*AoybcI-{B-y=gtWx{@gs5ONJN4UxAt!5ev9;z& zx_jA~^CX@<{}m^*9G0kATM{RvVUEqhN4 zFCHA8EZ|_RmgBP@dH%;f@o;b^mSLzo$yR-aVG>cYm6L4EBpc6||81vaYbM!rA|;zn zh-9Cnl0EnFpLpTNLyGnr+PIo&g;}R;t-VYqWVwFnWR@m&xyF^Am1#1*i0>1dfGcB+Ph7=oMOq~4qJ*P!?f_g@~Mad zE}eiYCg3WXfWQ4H0askWwI|@32zaaJIh}IzTyyiBO+53QO}1d3vxza!X(r5bmiXp5 zOX8a6G!xpMCB{5wNDx_nIZI>u%ULG1H|tFcU5=-{S!T62%c=He@iIf6dCn4Rp0k{3 zZ-%JDx-WY1&k%K(maV$Y8CDKsj2pis!^&Y=xOvVa3V7-STrvTd+cVEpz$F(jf)69k zmWv2JOv~O@&2v^T^PE+}JlDQ?uD4;Hb)HF?xhxgtIiL9EIiJKe&si$Wb3QTVIZudr z&QSXZfl;UcSono~eL&+A8uZPh&IIWy0B$<1@_Y7r*6 zkyl$V&v|9M=RA?-x$ZE}bqw?D(#(0H$kS=}+}P!URw;_mRSLO0)S}u{S+eMIQAAxX z3U9e6dY6kqJZn<8<)UEAMGm@U2NDlMN!(cY&ZNXD&}8I6YgJ3 zQ!o6Rrd#kYrm68SibD7o)6DlTrdbUC!WB=7LagsoyEIJc(l8CXG>GEK)Qb+Lz0tul z@z(RHiw>qNI+(_b4!rPhn%VGg%A-DmtJkVd~9mlC9iFlqPAW6we4yfwe89#3QAeqt~zSlRpi=s zDQnx6t!-Ddwp~T7ZC74xyXviNKN(MZD_7gDSZ%x7Vr{!}NwJF8wkxl;T|F+f?aIqn zRlWJDABreom#?Z=zN(6sukyTSDqxqSl1bgVaD)8(R$x?I%Wa#8m# z7pLP{liDp8HCry~c*}(=p485buWiYqrX`CyY{|kEPih;o)*2Y6g4Y_rIQ3ue!K$Y* zizl@gE7ZNQLd4=p?P7(R#R_%2Siu!f>cU2WH4hSMFG#3qkPtOV_vL}gp_}Rm&~S(X z?3JEcXxc{kh-i^Tlco^ST(wL@ zae7Ndlt1SPl-gADEwQ9Vfqfat#MLOwHVw2zsH6NDT{Rm@%(7ljN7oq&C7I;AZwEsS z61U}Y9;1}j2Bwo$?Umhvx z#MCHSNduCt5dBor$@a?Me_>F0b{@9DjipwYZs4^IHqV-Fpti%~2eHPufJd3`L_5!| z9h9~^$vSLz5~U3O#ewq5TDqhSNOb(6ZzhO8obgV=mcpg+P6|(@lip%@NT`Zb>7*bk zJ>@DLuO99j?<8}^J259<7$s>A>0m0owVJpmX{;>5cqhG8@sOAmPtp~Ubj>9luQ2W# zZzgF|!DuB7+_yB|TLgYf;|=6@b7qUBRv2$!zdhp}1^n$9?^I{VTRSL?H}K!Sd*f81|AZ~B2^kVa2v4#2u}U6h)8hXc&Di| z-oS!$B|#V?{IPUv6?0G0SXqSe?#F|NM6`I4E{UWoF6nrwFyDAHNt^nHF^5fkgD%r4 z&5FgX%}C=7TsWVXwJ_d5hI__43N+j^-oS<%4;SEYk)8lL+&A8F;NidOJa6E`4Z_74 zZ6L#m(GCm^M+(Q|jHOBMDbfoGAmfy#usg1#c>x++nioLAdv>y)30S3(<^_mwyO015 zP8Skj!DAiy`0H*yWyJl@^oHF=df2e5J~T-YV8KJ8S$~a?<6_@T*T)-bP#-7MltHQS z=8_+0WW8}l$dZjWmq?tE@i-%k7iS=?Z)|2Zb>2uO0t|q+r8ffXw#(~pH4g~rwlo(2 zaQB#tkiZuyVt_2Ixd6zUnhPMjW0{M8>=rTL-E3h8#@iDyO5?@{jf{VhZU|t(f?DB* z0RHW{Apn4HApr<1uvMv2x*>pq`)&x(;1P~aB{5`HN+%$9i-*F3J97;zc*w;=nrnc( zJzegtu!jUoc)Cm?3&6RVF2{>3{1`GP)!1A_7P%i;zL4xo1|tL?{-RMLe30JbqCFyYd8D^U33V!TD+9Z+#zGXY@SGv0uWZ^3vM8epc~ z*4(SS13&H?Z{Wz|8t>dIy#rTnR-A&B-UB1crSS%|+*9e^(tAj(j5OW=ms>{+=yGbj zBf`sl;|(w}p^N)t0GfL${n(V=0X6p|-CKGO37YXFO-k=Tn^Q?ggq!=un@SqEbLPte zcrGQqMdZ2kya7EI)C%Jb@VRHaqoB{d1R&t&&JIe?8vu0Qc*jAY6VDqEbdi1o2Hg-1 z5IXU^1EbKT@dgZCs&qdNJtUGwsx*M;)<^=0&QyAvrFTG;30*>^{nGpP;OL&D{ghZg zAU!0c#*;Kji2;(%Bpoj$<{NJ&X^|NNEL|G!E#lIp@dhxR8E+udx#taNx@WwjVADP0 z4RE@%gVJ~dp6(m(IQVo@y#YYoRBxQe4QM(s+JUj@zvoqNfXHnmZ1w}|Qyp{w1LQiT zEbdM&fYRl10%Uqmcz-rx4F|wl^jP6Nso|6)2q<0E1h8}}y|7giQo{kv z*XUuXVI|F60<6R(1pz^inG^)J+?z|i(MHS~4v^$_E&)YO=MoU)vE~w~;gqHGNh%R$ z07EWUDnQ6xj%cgdKp@Dah=C*TF&8mwIKX;35ocV)@eVd5H5?}5jJutHntLLqhv05^ zD+$oa%oza&ojW5yp?l5hS6%@gFm~T~1I`}TcvoHl9`JUv;sD=HjCWvwyG#xN;qIw)Zvj3g0tXm(YrFw+ zXDYqT0z8F71Mbe1#4o@68QBak)j~O(YFMn50u5_>LyB zOmElo7RYyy#*KH{fqX~X7K}HZYCUfc`kr<{-_aP^c-OK34-$;=P8s$cS1~>BbmA4@ z(bT8X5fM0`f~Yhrz^9WfV&5S+n=b(&LZ!U|d^&kt3h?Q~ONpVWPtp+)IH1Bwnq%M7 zZDQYvil9VCeaTU~bdpI1h!d90oHYf!k-o*mA3QCiR1=7S{PLkf`TN|1k6J}fQ45TBob!yIwiuFUIJncKGg5*MPMdd@|LEZZ z2p6_reb6XAT#e7*DioJ+a3_|?_PK;XtU{gwf==}_dn%&f#XQf(Ug)-$xE#TR?`Y7GkbiL8@9v+Z0Nge=-4^I+~x% zj^^vv3~|c$G0#@SjV?Wyee&+&L{VDrUTr?MWY*(z`Z^SiQo!mJ8NOW1w4ML@;dcj|X&!an&Z$(rXxyj4!*fT6%h|EefLvTRyQ=OT9z)D(%}}4(iKo^KT!xx+pir%(?V}&16=jE3lo7Pz z@7&s|0Q+2RKfXA)zrSd7y4oGN%a(h1Ppoa*yEfc;ZLv5$I6PY42)X!vL7!?zRaXZ! zA6eajc=56*^`WJ&9-qNgDB&VZh7GX{z|$+`r?iCFtSj*8BuoZ0J>-4|sjxo&JZ1QE zPhn4;3aGGu+O4p_qQ4D%TddbFojm`U`}2FV6p-{MXhS1VAk#0A)=+}Q{$HOSsT*Vh z0-o+IX?PKN_$BSHnet-pseLsmUL60#gV_rwi}_}8;y9Y`J}uf>rsm_T{6;9g{(;oo zJR5dWs@`Ob^D0vy(l41KGjTCK{nAM&+8*lTgPVuiXag3p-zfzaojIjIqla}$|7c)b zW#t;be0P2fcM72NtIfw&n&tSMzTO)S&n%V)H`JQ~_AU?4!96_L1U1YuMW@~@Zajn5 zBIC{C&RMt)KI>5XaJjhwvvWF>Y=0kSXR|t;RJpm_i0DP+lhT3|3p=}98 z8oXvsgQ1j4m_$NArn5;0SbEq=Cxr^sxfCjp=RHCV=?17E$OedbkF`=%8w3G^Zh|1d z&q)vjW2Tt%PNQcVwh)l|KJ{G<5Ls4I}UFYVa<@0oWp=}8@A3U!4i7SB4 zd$<--oZl2VD_~JMvxI(^$|in%zPy7a5NPzP&BvCH)%cvg4#h?I3iRnr#5}72EV}0z z;>G9TJ;Pt@@)ZC&XTAbZ=Z&uric!c!N`4R%1gtrX&hxJy&zJC8fi%C?es07FocZ~t zAryIS%Br!-M2a$HrBLL~-EXOjk+K4$j&VqEbE@$lDj)*;fd9~+TJnJZP?w=t*FC>- zF~9P_^XuQJt;KtTpL9OJOSfAS4D9#~;~1F6mPf_-nz*yKJ3Oa-(!HxYgX4aLge(m5g{MOWV7c=A-d+WbO5e;HmwHz(C>v@S#9ibQzdIezf+{(OUo^Z9#= zqs4Ng=482hLpHXoFUQx!oxMw~`w()Ok?S8`e=fb*ez%LPjm2l=`;2ul%KoQjs38WFku@?eonv06)J?|~chdnmfBf9^fAz*@O^6zRPE^e| zsPZ|h0nqx>FYvZjKp7HM?8a7X2B7QGX7s;xJ)|+i^j6t3#t}^HCTC?sVZ4kTEdTKR zqnGa=;cl^et^M3GwHTjo8bYmaG9?XQyT?Deb4#S81J_bL>imd^y>eYpcj{E|3fBgi zD^cOXOT$^Z17LVtX2J{EX@yI_kv&p{3$Ra7KyqGl3ZT5_MB_E=;X@&0otXl3D~-y} zynO%GECp8nvii`d2cYt2a1}~7zNGU-?Q?~9-vIafHkyAK$wN44y@gz`%!% z9+`I!tb3PFidU@nc=veK`s3!^N9<-2)6SSDOC!vv0w>2a=zeae(BhXAPjd@a_SY5AEH{_0R{*r|RwOCT{k3 zd;~2Ze?z1Rt4gN{om0RF5g zJbl<02?w;$r^^Mr+1;RBM*!;E9U$tm;U*CE?G6yMMcV-)=O|G0QqO^;@13K8_e_16 zYzt`_$JHF*dQ&6Ar>^O0xX_~RY1lE4yztY?)iZWd;j|`7Iaf>Uw??Q??Nc= z7S`JOJI%YA%(5{{0h#Be0pRAn_yDhn4<8?Vo5cqJ&`P?xcGKW(@!h*=fT54Ko2Fy) zOzi+H8Bp}Z$HXPuARp4l#HwN>kH;!adG_Rv+DDiHG=HuAtd|fpjL$a>Z2eH~Sll!m zw)3%l>DE`z0H9}l!6-W)o~xBPb4B2lj0}W_s10Orhq@qO=#EnO_|YFMxZff9s@w!yPb2aIqR; z|JrMdrQly^6pU15VCJPN12^yZ$5c6_UxsZ~=U2CHuNNE6D1hhTuc=M28F2G%*EOz! z?z%>u-ipx?U}zj&S}tDS%o>8Z2myr;8s3d|*Ner$;$Vi-?BdSqxC>fYi0=Wgn0YMyZO)_k*n zclh4XdphqvSnbbm-NzQP&%Kf{K-1`b!!i1O_oMrL)K{U2002XUcmOx#y^?WveoHFnx~pr_~Z|mbDoU1u9XiEuaNX5n7awi>B{?OIo>|l(UF_wqN6@_h6I5(aMXGOq@%Y3#QWadPh4yN zu<`mbpW4C=+x0<5>{|N;jrUY_^xP-tf{v9-rf5`~&lv0lRlO_aqIG_9Z!r!~jb=csan5j@}Ni zq{HTPVJC<`K)iCGdw@9Z-_a#H>N09&pHV{{+jS%ch;!iGJvl&>OWjR7Koo*qE2pdK zGHR$^b`P-XpX(^Z`JBRetvwylxh@At2QLR$)6v@j;(TB?K-bzoY#ev%Q(L=XyFTcM zU2DIfk)-YDxlhnYJvzem;_O@yBmqrNPuSF8f|Du@OAX)Jmqy|gNY@wT8|-K@p>zCQ z@~Cm7aokUFwq+9#COGuVHB0K8yeV}~;F3BjaoLwMZ=^4{XFVf@cjoAlBZZgdZh9k4 zC?DKN6UqQVUE+)!rOsJ7F;aMKjsV%0vvagBXXr#ZlsXzMLM=zMLY+zNjY1kupcl8!7YT+(?-#N1Evj5Jb1$oLXlI(QxI= zL?qqNv6JRU2p^tYheG4s?520Y66U}OO5w(J^woNB$|aC}(NRnGGFQ%HFT6Q7&D==w zK2_H&sp}LzoVmV{QsIlf;Lml86lleszTnO|y5vaV*14PBNb%xTch`-ypbU`Ib(+DV z*9o*OB1^#paq*CVNLWWiFrS3yxpMNU>X_uP>_q za-{J89MRrLaiYJw>tu>ktzEOEuG0+SRQW#Y`puv^Mcp5%I=8&LFLfrkqf3kwr*e~< za-{0`Yxh$-QgsBZYvve2+9HrzMTM|p!d!j?S98)>&Uhyx_a zj~oNw2=+1nN0@dnfF$rpH;)*5(m8gLz~h>$4vcgpJ!PW=96_ktFX^YMe2Y# zIqIN7)p`l7npC~F39g^&zUhW`-%R>JXZMZJRrk$<&2f3^d^+>scaR-)DYjoA2|Yl1 z`V{ibf?BwK=0}>Y=47R2*v<3bBX*y_@4Ji+2l(A7o$!9z?npuk?@^CxSl}o(nQR=` zWr=9G(4J1Exzj@hP)LjlO{U#jpiVv#Ur?n~dlC}=t-hsF%j%#SeoL^c^$XP)zOUm% z9!+qQ*OzQVqEGdF(XwsabI^?9;Kf#0PpQnlN+iQcr0rC`a1_-oBVT>XV*T|kSD9dY zds}}|Ex0=WV@?{uu_I(0Pm`(C=KgfqL{6mEV^qDliPSoa%@3O92HLD2&@}J5P2T#{ z2Rayj&s9@8CNjK8HIdcxbqdBDWRQ|9p%7;4Mm9geK12)NC)7zXB2}owcmWSXu?PA@ z`3+xEQQz?UG-3qoLvUbndrwg>`-MMg`faod1*YZzDg~t0KQ)trJu$bvp58hA>i74c zWbKz4#+i*!K=86V`sW>pPxd9jMB;0{lWxOE;chBsFsfTjCj3n*W&Jw>He@%+6wqix z`a^z?Pme~uJR}6Spx>j*j8fJkd;(MzuT7DNR0S}gbNidZpF8Ta`537D==7;>5 zTg;j}cA0BES;cQ^2I0o4?u;2E%R#&2Suj>Suk~r0Ua?TAZ_XI8 znFMc^68SE+S43?TGpG@{#{Jg!+mvAZ#Owy-ci5A;kdJt)n`@n;j-D{psv^Y z-pOLUIC-#8H}C0#kyUH^32pkR_VMalEaU9S(-^fU|JYY#HI0!Arics4UM`-pc<21Q;*!L?l6Yew){^|})I2Gk z@Ql==RERkylRxyv2?+6}7p11ARK^!4<`(3n7ER7R95UJBXdH(w7evqGhey=xvc!SUaJ_NlTI#9AdC u1tmzJDOtf-lg~fXPy)LbBnDvsm8GVor6!kvgdk3x9R1vl(Qxvj=lTHs^nwup diff --git a/tests/fixtures/onnx_genai_workflows/decoder/policies/token_state_update.onnx b/tests/fixtures/onnx_genai_workflows/decoder/policies/token_state_update.onnx index 5ff11b46b3a74fa10132a719a2dede5cd17cbe8c..268d221d754c99df43c4bc05049adc3ae7c5e2bb 100644 GIT binary patch delta 93 zcmaFMzfgdegIkC#H$N$}wAiXoU?Xo58{^)|1#GHJ)0rkuXLHd{N-Rmv$jnPuvQkhg sNG*yl$}d;aQBcavOH3}wEK3FPi@|b5sl};9Ws|S5J!dqUypLTU0BL+5I{*Lx delta 45 zcmZ1|@RpyKgIkC#H$N$}wAktv|3=;fF B4oCn1 diff --git a/tests/fixtures/onnx_genai_workflows/vlm/inference_metadata.yaml b/tests/fixtures/onnx_genai_workflows/vlm/inference_metadata.yaml index ae709cc1a..e2ebfb1c4 100644 --- a/tests/fixtures/onnx_genai_workflows/vlm/inference_metadata.yaml +++ b/tests/fixtures/onnx_genai_workflows/vlm/inference_metadata.yaml @@ -141,6 +141,46 @@ pipeline: 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: 2 + request.eos_lengths: + contract: + dtype: int64 + rank: 1 + shape: + - batch + role: + kind: opaque + source: + kind: application + name: eos_lengths + required: false + default: 1 + request.row_max_iterations: + contract: + dtype: int64 + rank: 1 + shape: + - batch + role: + kind: opaque + source: + kind: application + name: row_max_iterations + required: false + default: -1 package.eos_ids: contract: dtype: int64 @@ -225,6 +265,103 @@ pipeline: kind: literal required: false default: 0 + request.temperature: + contract: + dtype: float32 + rank: 1 + shape: + - batch + 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 + 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 + 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 + 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 + role: + kind: runtime + version: '1.0' + role: seed + source: + kind: request + required: false + default: 0 + request.grammar_mask: + contract: + dtype: bool + rank: 2 + shape: + - batch + - 128 + role: + kind: opaque + source: + kind: application + name: grammar_mask + required: false + default: true + request.rng_offset: + contract: + dtype: int64 + rank: 1 + shape: + - batch + role: + kind: opaque + source: + kind: application + name: rng_offset + required: false + default: 0 package.loop_0_active: contract: dtype: bool @@ -291,14 +428,24 @@ pipeline: artifact: policies/token_sampler.onnx contract: id: onnx-genai.token-sampler - version: '1' + version: '2' bindings: logits: logits + token: token + temperature: temperature + top_k: top_k + top_p: top_p + min_p: min_p + grammar_mask: grammar_mask active: active done: done - token: token + rng_seed: seed + rng_offset: offset + rng_next_offset: next_offset parameters: - mode: greedy + mode: seeded_stochastic + batching: per_row + inactive_rows: preserve application_overridable: true termination: implementation: @@ -306,24 +453,28 @@ pipeline: artifact: policies/termination.onnx contract: id: onnx-genai.termination-predicate - version: '1' + version: '2' bindings: tokens: token_ids eos_ids: eos_ids iteration: iteration max_iterations: max_iterations + eos_lengths: eos_lengths active: active previous_done: previous_done 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 contract: id: onnx-genai.state-update - version: '1' + version: '2' bindings: current: current update: update @@ -333,6 +484,9 @@ pipeline: next: next next_lengths: next_lengths emitted_length: emitted_length + parameters: + batching: per_row + inactive_rows: preserve last_token_logits: implementation: kind: onnx @@ -345,6 +499,14 @@ pipeline: implementation: kind: onnx artifact: policies/decoder_step_update.onnx + termination_batch_initializer: + implementation: + kind: onnx + artifact: policies/termination_batch_initializer.onnx + iteration_broadcast: + implementation: + kind: onnx + artifact: policies/iteration_broadcast.onnx cache_length_update: implementation: kind: onnx @@ -449,6 +611,17 @@ pipeline: initializer: initializer.cache_lengths recurrence: kind: invariant + rng_offset: + contract: + dtype: int64 + rank: 1 + shape: + - batch + class: semantic + scope: invocation + initializer: request.rng_offset + recurrence: + kind: invariant position_ids: contract: dtype: int64 @@ -558,6 +731,18 @@ pipeline: 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: @@ -584,21 +769,37 @@ pipeline: outputs: last_logits: decoder.setup.last_logits steps: + - kind: invoke + component: iteration_broadcast + inputs: + value: loop.iteration + active: active + outputs: + rows: loop.iteration_rows - 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 + grammar_mask: request.grammar_mask + seed: request.seed + offset: rng_offset active: active done: done outputs: token: sample.body + next_offset: sample.next_offset - kind: invoke component: termination inputs: token_ids: sample.body - eos_ids: package.eos_ids - iteration: loop.iteration - max_iterations: request.max_iterations + eos_ids: termination.eos_ids + eos_lengths: termination.eos_lengths + iteration: loop.iteration_rows + max_iterations: termination.max_iterations active: active previous_done: done outputs: @@ -686,6 +887,8 @@ pipeline: next: decoder_step.body_attention_mask - cell: generated_lengths next: token.next_lengths + - cell: rng_offset + next: sample.next_offset - cell: active next: loop.next_active - cell: done diff --git a/tests/fixtures/onnx_genai_workflows/vlm/policies/iteration_broadcast.onnx b/tests/fixtures/onnx_genai_workflows/vlm/policies/iteration_broadcast.onnx new file mode 100644 index 0000000000000000000000000000000000000000..6618ad65560e12af5e2f95f56cf1d059c7b64cda GIT binary patch literal 611 zcmcIh!AiqG6l_A8rcVp9B5Ds(OHuIedT1XhBnu z2{VdNhH8Tz9H^&pD213As=2ZuO^Go$;Y<@vlR0wl9|j$0e`75*u%S{hvt+=#1NDrg zQwA#o$kX!Kv|G0`MMR7XAN{Uh!hbS?cDS%?%rwq#<48*qQ)29JyRwG>f}nvsa06#} z@s-*soCHnOgD&bQnQ}T0xy=KxT#A)Q7!ga8=p2f_&~m5ox1Gn`^;p&{ESzIFENcwM zOjIYs(gUVaSe^HZ+?Usda}B*;nV83DzhYI7km-!4oap&gf8iWK=SK}T8Zl~%I&VCv HZS8#mtB=JC literal 0 HcmV?d00001 diff --git a/tests/fixtures/onnx_genai_workflows/vlm/policies/termination.onnx b/tests/fixtures/onnx_genai_workflows/vlm/policies/termination.onnx index e8cf5a74afb81d1e02e4e6c620c8dbdce8d1e7aa..46ce9432722602b317df0a22c323084f5f063c85 100644 GIT binary patch literal 6378 zcmc&(+in|081^~`+Y>hlQ`%aq0%28gVYR_Kb{xAGsI)>z;ZUg+H$`Z)Sx=Kq<6U=m zZCX%oB7p=5kf1jpgu(;B>+m2v0Q@tvJD%C3G&kG&W@r96ee?Y@|6yET!OP=lH<%{Z z|8~)D7H-byttd=VFHC92UQ^{T8Zp{&>ban8BQ*04r;G++n%wv||5?XYZRZ?ngx;7X z6VGS3vU3H!K6$XW8HM4a#E*kXy6GQ!NkaDn7JIS3{}@+4{cL9g*@a$w4x0GUgeANO zz(b8RdcZ;&j1puZ#ph=9lQ4NWW$X(^U3-0A5y)Mqp(JI&`_V!s3RGiJ!m;AUjK249 z>K#&IH)TPfBF>Vim4)ahQ=i?UKD(Me*HJ_1vsF2^s%3%t)iQXwyYEdH?b%IP5WMu9 zB~fc5w7|kqs|s-1O|%5ZGEQ4?#7YDA6pq-8W|$E2`Bk9qh<5M_h_mChWv}K%0e8!( zW%aT+D=4yGRAm2IiY$|>W6QXMU8GZFH1+ufxZ{O;jBeXaSrE)^J4>Qg7H@FY<|`ZO zG;ZqPguNKN0DA%EY=s+$Ll*9(`-uk6HO<-Z1iax3fCpE>B+0A}xlTW1V8-2!-4q3( zQn%wQ@>*Fc{YI!Xv!-0T83tUR(RafUbzOUfe#9~(`5c6U0aZr6TgUXviZPD6>>YpDK!vB@YH zC#5=qrxrivc|bwkp1rOr5?1e7X01&B^t*)$eDQ;?khwh`BfTFpFJ&=xx9qjNEHQ&K z$OL7bx*~7=6TH0?ckn$NzW_M#(jW>=%KGARj^J``!{vmpq-q;3=cu&H`Io$$VpX&{ z_8w7%F2U!O3U9~Kmbe4^Vy3uUBvB5WRa0J;UXQ@wPGU9-qG>`$QOG!JA4gzOhV}~o zN<2en0V-vI_{u`do}UIY#!Z`CJvXhSqnTFHDb=+8EP{*tYPnnJG>zffjNXnZacu-2 zK@V}A>g~8JcD_~ypL9^ZxT1&X^pH~N0U+jC$sG|Pc#)qSjgfAdj6Bg6&JGOeX~&iH zNIzZoB}ERmu+#&}j>+h=BCe> ztGN>tYc!riO1EH(&oXGYjf9;(Q(quTw<~-H&%v~N5vlWu5voJ~MO3X@2Q4iz=!llu zDJ`ziVjIc4IqEetD|B%aR;%GXbF9wJU z#`A*eZ_d+K@(WtxkLzd|Ca=P>1-JtR7q~guJXFovZt{FHXI1V7e2>gNja`b}WEq*- z4Ig07t&7abG>n%FO9S~v&P|JRHV=u!kcQY+VaS|@SA}~PF|fA68>dH8oF+M+7?uZr zrj+lJbhNe4Yk|-yU5VCJO-p*3RH#aj1WszMs>^eVq{C7K@^ftN7|d}-(vkJc0(w(E zX>abaFv~|ar5U2!X2%d0;D4%LI&pN!oa-BoX67ga;G@8Y+|IcQ4g9e=5nb|UXD-fO zb-6q*f>hZhE}vJRWo+agytyFB=FA>~@_F+rw8}_^H)JJ{8L0wqG6m2>nCu|j6OE9p z#?j{-J~#~#;PblJ91KiYJQG{yu_Ny#* literal 4866 zcmcgwPj4GV6!*FzUXRlxokD9?p-NYQL)04YI*#L>N)bY`1d%{og3xHQo;I7Oj+An%iqIj$XVMqrDEG6Waf>goCRT;Jp5PwjF(1jaGlgcZ_Ja4 z=W|*aY>_*YlP9ex3{Mk34kl^KA9+c_jshNgv48Z8R=@scut}`KEc_3a_|b$XI0N93 zdK#VZkOjkpXh`Yx8T%?so=!RcmNUm%UkC!ZW7l&{S@3?**>j8Rxe9wXH1;YblJ>hn zYCKBLQReEDb|RFW)R(Qy4a!h8T;hv)EluKYbART=G?L05>i>F{1H4Y-UK-W1!2D^D904IdMkjNItc?czf)A%Y|8hEJyVm+b z5D0hedal`CBjyR`lVLDUn&57wfqkI|L9ZADJtYX%fn=FNN(BWmpmU`qj~;Ez*rze~ zQXaECYgN7q68G%dTrSJJ1;+YC#`>2r2C23*GyVht5C9)u8bqPSRv*LRqv4P_ZL5J# zL6g(AuYph&M}LVB$aW>GWABvZ;1YT>x88o7$sIOA6r^m)K=L#eow#0ICKS8*61>+Aen`w?*A?AGlt|8Al6r6yShC{?jUy+ii^-V%)99%{K33+fjv=;sT> z-dT5+?~UiYdvU3kuyM-Ud8&7njf7Uzz>cE5snlsWM22ZGwW~~B{B2LszJuh+9J=5E zd|S__+97Tiv&EdIGvn_nxto1FX+7cLyjgB#(Zw*O#89^AtLk@l9F4gBVAED2ltFtq z41BnB*!N%nKWijniu}NV>iNYT_3{<8dGCw*g#aUZXNgR1fiM zi?|fu*P6+KI|FI8S7NKzf;uN->bu0GsMOnr+wRFocQ~K;rpH{Nn>O)%Y%%H)&$Onf z{PtaM3$&b96a)LF<5EBJfKvrELwx9PUl1;EM_W5v$N4*Lb-uB?4)u`O0mY1q&C++J zTolv~1+Fvh1_F=C1T!IjNic~B7f?|n#4kNcj;C-!?u42Odb~Q1J0{p2!R-grnX86$RMYXl&8;g1r$G8XWmDnnZ z_NxM|a+n!*f31)X>(B29ygXrU}i1wmRP^3g}dUYv({V{>hsM zDM~$jLAtUs-97Z{%8b6)IBpl01;@O-O~R)V&NPc_7=7qWxO5 zKU40(on&n4Bg)f(peJRDO@|9l&QVnU^m(Q{h3&t3$_|5Wre0~SlVHMZM!e~h&P>^X Ut%Wh@a7axNkI!cnD6Mb)2KtqxhyVZp literal 0 HcmV?d00001 diff --git a/tests/fixtures/onnx_genai_workflows/vlm/policies/token_sampler.onnx b/tests/fixtures/onnx_genai_workflows/vlm/policies/token_sampler.onnx index 29209ca573fdea2268d08475ddd411a15201d9b7..af9f73753fbbd81fe648083be389f699c689b6e1 100644 GIT binary patch literal 59345 zcmd5_TaP5kRi2t_PtDcky^Pz&UKbWZjXkQoRc2a9c-Jfoi|rL&BcrvfQcZWwOz%!t zwYzG@Ymfz6ghAXfmH`{^U~{#&%SP~E3Gu+I9}zzQiAUafSiXqNh&bm&=6D~7sOguk zjI4@0aZa2Q=MrD!&ZW`$qt&g$`|A&W=euWypB#7o$@zzHj*M@6Ffi+1c_J>;2X7Vy(V0d}A;?Ki}URK3ELRvlS7~E)QnO?oj>HL?q)g z_#BFeUm6V0t{01g(awX}hv)0fEZx1N|LbR`6*pcfa|+|Nh_ajK{3PGHXhiHI2xup`FwAQ$D)-VD{YM=Cgu8b$B}S69 z{gf)Io;`rAoa|mw{~KcolZ}f|V*AoybcI-{B-y=gtWx{@gs5ONJN4UxAt!5ev9;z& zx_jA~^CX@<{}m^*9G0kATM{RvVUEqhN4 zFCHA8EZ|_RmgBP@dH%;f@o;b^mSLzo$yR-aVG>cYm6L4EBpc6||81vaYbM!rA|;zn zh-9Cnl0EnFpLpTNLyGnr+PIo&g;}R;t-VYqWVwFnWR@m&xyF^Am1#1*i0>1dfGcB+Ph7=oMOq~4qJ*P!?f_g@~Mad zE}eiYCg3WXfWQ4H0askWwI|@32zaaJIh}IzTyyiBO+53QO}1d3vxza!X(r5bmiXp5 zOX8a6G!xpMCB{5wNDx_nIZI>u%ULG1H|tFcU5=-{S!T62%c=He@iIf6dCn4Rp0k{3 zZ-%JDx-WY1&k%K(maV$Y8CDKsj2pis!^&Y=xOvVa3V7-STrvTd+cVEpz$F(jf)69k zmWv2JOv~O@&2v^T^PE+}JlDQ?uD4;Hb)HF?xhxgtIiL9EIiJKe&si$Wb3QTVIZudr z&QSXZfl;UcSono~eL&+A8uZPh&IIWy0B$<1@_Y7r*6 zkyl$V&v|9M=RA?-x$ZE}bqw?D(#(0H$kS=}+}P!URw;_mRSLO0)S}u{S+eMIQAAxX z3U9e6dY6kqJZn<8<)UEAMGm@U2NDlMN!(cY&ZNXD&}8I6YgJ3 zQ!o6Rrd#kYrm68SibD7o)6DlTrdbUC!WB=7LagsoyEIJc(l8CXG>GEK)Qb+Lz0tul z@z(RHiw>qNI+(_b4!rPhn%VGg%A-DmtJkVd~9mlC9iFlqPAW6we4yfwe89#3QAeqt~zSlRpi=s zDQnx6t!-Ddwp~T7ZC74xyXviNKN(MZD_7gDSZ%x7Vr{!}NwJF8wkxl;T|F+f?aIqn zRlWJDABreom#?Z=zN(6sukyTSDqxqSl1bgVaD)8(R$x?I%Wa#8m# z7pLP{liDp8HCry~c*}(=p485buWiYqrX`CyY{|kEPih;o)*2Y6g4Y_rIQ3ue!K$Y* zizl@gE7ZNQLd4=p?P7(R#R_%2Siu!f>cU2WH4hSMFG#3qkPtOV_vL}gp_}Rm&~S(X z?3JEcXxc{kh-i^Tlco^ST(wL@ zae7Ndlt1SPl-gADEwQ9Vfqfat#MLOwHVw2zsH6NDT{Rm@%(7ljN7oq&C7I;AZwEsS z61U}Y9;1}j2Bwo$?Umhvx z#MCHSNduCt5dBor$@a?Me_>F0b{@9DjipwYZs4^IHqV-Fpti%~2eHPufJd3`L_5!| z9h9~^$vSLz5~U3O#ewq5TDqhSNOb(6ZzhO8obgV=mcpg+P6|(@lip%@NT`Zb>7*bk zJ>@DLuO99j?<8}^J259<7$s>A>0m0owVJpmX{;>5cqhG8@sOAmPtp~Ubj>9luQ2W# zZzgF|!DuB7+_yB|TLgYf;|=6@b7qUBRv2$!zdhp}1^n$9?^I{VTRSL?H}K!Sd*f81|AZ~B2^kVa2v4#2u}U6h)8hXc&Di| z-oS!$B|#V?{IPUv6?0G0SXqSe?#F|NM6`I4E{UWoF6nrwFyDAHNt^nHF^5fkgD%r4 z&5FgX%}C=7TsWVXwJ_d5hI__43N+j^-oS<%4;SEYk)8lL+&A8F;NidOJa6E`4Z_74 zZ6L#m(GCm^M+(Q|jHOBMDbfoGAmfy#usg1#c>x++nioLAdv>y)30S3(<^_mwyO015 zP8Skj!DAiy`0H*yWyJl@^oHF=df2e5J~T-YV8KJ8S$~a?<6_@T*T)-bP#-7MltHQS z=8_+0WW8}l$dZjWmq?tE@i-%k7iS=?Z)|2Zb>2uO0t|q+r8ffXw#(~pH4g~rwlo(2 zaQB#tkiZuyVt_2Ixd6zUnhPMjW0{M8>=rTL-E3h8#@iDyO5?@{jf{VhZU|t(f?DB* z0RHW{Apn4HApr<1uvMv2x*>pq`)&x(;1P~aB{5`HN+%$9i-*F3J97;zc*w;=nrnc( zJzegtu!jUoc)Cm?3&6RVF2{>3{1`GP)!1A_7P%i;zL4xo1|tL?{-RMLe30JbqCFyYd8D^U33V!TD+9Z+#zGXY@SGv0uWZ^3vM8epc~ z*4(SS13&H?Z{Wz|8t>dIy#rTnR-A&B-UB1crSS%|+*9e^(tAj(j5OW=ms>{+=yGbj zBf`sl;|(w}p^N)t0GfL${n(V=0X6p|-CKGO37YXFO-k=Tn^Q?ggq!=un@SqEbLPte zcrGQqMdZ2kya7EI)C%Jb@VRHaqoB{d1R&t&&JIe?8vu0Qc*jAY6VDqEbdi1o2Hg-1 z5IXU^1EbKT@dgZCs&qdNJtUGwsx*M;)<^=0&QyAvrFTG;30*>^{nGpP;OL&D{ghZg zAU!0c#*;Kji2;(%Bpoj$<{NJ&X^|NNEL|G!E#lIp@dhxR8E+udx#taNx@WwjVADP0 z4RE@%gVJ~dp6(m(IQVo@y#YYoRBxQe4QM(s+JUj@zvoqNfXHnmZ1w}|Qyp{w1LQiT zEbdM&fYRl10%Uqmcz-rx4F|wl^jP6Nso|6)2q<0E1h8}}y|7giQo{kv z*XUuXVI|F60<6R(1pz^inG^)J+?z|i(MHS~4v^$_E&)YO=MoU)vE~w~;gqHGNh%R$ z07EWUDnQ6xj%cgdKp@Dah=C*TF&8mwIKX;35ocV)@eVd5H5?}5jJutHntLLqhv05^ zD+$oa%oza&ojW5yp?l5hS6%@gFm~T~1I`}TcvoHl9`JUv;sD=HjCWvwyG#xN;qIw)Zvj3g0tXm(YrFw+ zXDYqT0z8F71Mbe1#4o@68QBak)j~O(YFMn50u5_>LyB zOmElo7RYyy#*KH{fqX~X7K}HZYCUfc`kr<{-_aP^c-OK34-$;=P8s$cS1~>BbmA4@ z(bT8X5fM0`f~Yhrz^9WfV&5S+n=b(&LZ!U|d^&kt3h?Q~ONpVWPtp+)IH1Bwnq%M7 zZDQYvil9VCeaTU~bdpI1h!d90oHYf!k-o*mA3QCiR1=7S{PLkf`TN|1k6J}fQ45TBob!yIwiuFUIJncKGg5*MPMdd@|LEZZ z2p6_reb6XAT#e7*DioJ+a3_|?_PK;XtU{gwf==}_dn%&f#XQf(Ug)-$xE#TR?`Y7GkbiL8@9v+Z0Nge=-4^I+~x% zj^^vv3~|c$G0#@SjV?Wyee&+&L{VDrUTr?MWY*(z`Z^SiQo!mJ8NOW1w4ML@;dcj|X&!an&Z$(rXxyj4!*fT6%h|EefLvTRyQ=OT9z)D(%}}4(iKo^KT!xx+pir%(?V}&16=jE3lo7Pz z@7&s|0Q+2RKfXA)zrSd7y4oGN%a(h1Ppoa*yEfc;ZLv5$I6PY42)X!vL7!?zRaXZ! zA6eajc=56*^`WJ&9-qNgDB&VZh7GX{z|$+`r?iCFtSj*8BuoZ0J>-4|sjxo&JZ1QE zPhn4;3aGGu+O4p_qQ4D%TddbFojm`U`}2FV6p-{MXhS1VAk#0A)=+}Q{$HOSsT*Vh z0-o+IX?PKN_$BSHnet-pseLsmUL60#gV_rwi}_}8;y9Y`J}uf>rsm_T{6;9g{(;oo zJR5dWs@`Ob^D0vy(l41KGjTCK{nAM&+8*lTgPVuiXag3p-zfzaojIjIqla}$|7c)b zW#t;be0P2fcM72NtIfw&n&tSMzTO)S&n%V)H`JQ~_AU?4!96_L1U1YuMW@~@Zajn5 zBIC{C&RMt)KI>5XaJjhwvvWF>Y=0kSXR|t;RJpm_i0DP+lhT3|3p=}98 z8oXvsgQ1j4m_$NArn5;0SbEq=Cxr^sxfCjp=RHCV=?17E$OedbkF`=%8w3G^Zh|1d z&q)vjW2Tt%PNQcVwh)l|KJ{G<5Ls4I}UFYVa<@0oWp=}8@A3U!4i7SB4 zd$<--oZl2VD_~JMvxI(^$|in%zPy7a5NPzP&BvCH)%cvg4#h?I3iRnr#5}72EV}0z z;>G9TJ;Pt@@)ZC&XTAbZ=Z&uric!c!N`4R%1gtrX&hxJy&zJC8fi%C?es07FocZ~t zAryIS%Br!-M2a$HrBLL~-EXOjk+K4$j&VqEbE@$lDj)*;fd9~+TJnJZP?w=t*FC>- zF~9P_^XuQJt;KtTpL9OJOSfAS4D9#~;~1F6mPf_-nz*yKJ3Oa-(!HxYgX4aLge(m5g{MOWV7c=A-d+WbO5e;HmwHz(C>v@S#9ibQzdIezf+{(OUo^Z9#= zqs4Ng=482hLpHXoFUQx!oxMw~`w()Ok?S8`e=fb*ez%LPjm2l=`;2ul%KoQjs38WFku@?eonv06)J?|~chdnmfBf9^fAz*@O^6zRPE^e| zsPZ|h0nqx>FYvZjKp7HM?8a7X2B7QGX7s;xJ)|+i^j6t3#t}^HCTC?sVZ4kTEdTKR zqnGa=;cl^et^M3GwHTjo8bYmaG9?XQyT?Deb4#S81J_bL>imd^y>eYpcj{E|3fBgi zD^cOXOT$^Z17LVtX2J{EX@yI_kv&p{3$Ra7KyqGl3ZT5_MB_E=;X@&0otXl3D~-y} zynO%GECp8nvii`d2cYt2a1}~7zNGU-?Q?~9-vIafHkyAK$wN44y@gz`%!% z9+`I!tb3PFidU@nc=veK`s3!^N9<-2)6SSDOC!vv0w>2a=zeae(BhXAPjd@a_SY5AEH{_0R{*r|RwOCT{k3 zd;~2Ze?z1Rt4gN{om0RF5g zJbl<02?w;$r^^Mr+1;RBM*!;E9U$tm;U*CE?G6yMMcV-)=O|G0QqO^;@13K8_e_16 zYzt`_$JHF*dQ&6Ar>^O0xX_~RY1lE4yztY?)iZWd;j|`7Iaf>Uw??Q??Nc= z7S`JOJI%YA%(5{{0h#Be0pRAn_yDhn4<8?Vo5cqJ&`P?xcGKW(@!h*=fT54Ko2Fy) zOzi+H8Bp}Z$HXPuARp4l#HwN>kH;!adG_Rv+DDiHG=HuAtd|fpjL$a>Z2eH~Sll!m zw)3%l>DE`z0H9}l!6-W)o~xBPb4B2lj0}W_s10Orhq@qO=#EnO_|YFMxZff9s@w!yPb2aIqR; z|JrMdrQly^6pU15VCJPN12^yZ$5c6_UxsZ~=U2CHuNNE6D1hhTuc=M28F2G%*EOz! z?z%>u-ipx?U}zj&S}tDS%o>8Z2myr;8s3d|*Ner$;$Vi-?BdSqxC>fYi0=Wgn0YMyZO)_k*n zclh4XdphqvSnbbm-NzQP&%Kf{K-1`b!!i1O_oMrL)K{U2002XUcmOx#y^?WveoHFnx~pr_~Z|mbDoU1u9XiEuaNX5n7awi>B{?OIo>|l(UF_wqN6@_h6I5(aMXGOq@%Y3#QWadPh4yN zu<`mbpW4C=+x0<5>{|N;jrUY_^xP-tf{v9-rf5`~&lv0lRlO_aqIG_9Z!r!~jb=csan5j@}Ni zq{HTPVJC<`K)iCGdw@9Z-_a#H>N09&pHV{{+jS%ch;!iGJvl&>OWjR7Koo*qE2pdK zGHR$^b`P-XpX(^Z`JBRetvwylxh@At2QLR$)6v@j;(TB?K-bzoY#ev%Q(L=XyFTcM zU2DIfk)-YDxlhnYJvzem;_O@yBmqrNPuSF8f|Du@OAX)Jmqy|gNY@wT8|-K@p>zCQ z@~Cm7aokUFwq+9#COGuVHB0K8yeV}~;F3BjaoLwMZ=^4{XFVf@cjoAlBZZgdZh9k4 zC?DKN6UqQVUE+)!rOsJ7F;aMKjsV%0vvagBXXr#ZlsXzMLM=zMLY+zNjY1kupcl8!7YT+(?-#N1Evj5Jb1$oLXlI(QxI= zL?qqNv6JRU2p^tYheG4s?520Y66U}OO5w(J^woNB$|aC}(NRnGGFQ%HFT6Q7&D==w zK2_H&sp}LzoVmV{QsIlf;Lml86lleszTnO|y5vaV*14PBNb%xTch`-ypbU`Ib(+DV z*9o*OB1^#paq*CVNLWWiFrS3yxpMNU>X_uP>_q za-{J89MRrLaiYJw>tu>ktzEOEuG0+SRQW#Y`puv^Mcp5%I=8&LFLfrkqf3kwr*e~< za-{0`Yxh$-QgsBZYvve2+9HrzMTM|p!d!j?S98)>&Uhyx_a zj~oNw2=+1nN0@dnfF$rpH;)*5(m8gLz~h>$4vcgpJ!PW=96_ktFX^YMe2Y# zIqIN7)p`l7npC~F39g^&zUhW`-%R>JXZMZJRrk$<&2f3^d^+>scaR-)DYjoA2|Yl1 z`V{ibf?BwK=0}>Y=47R2*v<3bBX*y_@4Ji+2l(A7o$!9z?npuk?@^CxSl}o(nQR=` zWr=9G(4J1Exzj@hP)LjlO{U#jpiVv#Ur?n~dlC}=t-hsF%j%#SeoL^c^$XP)zOUm% z9!+qQ*OzQVqEGdF(XwsabI^?9;Kf#0PpQnlN+iQcr0rC`a1_-oBVT>XV*T|kSD9dY zds}}|Ex0=WV@?{uu_I(0Pm`(C=KgfqL{6mEV^qDliPSoa%@3O92HLD2&@}J5P2T#{ z2Rayj&s9@8CNjK8HIdcxbqdBDWRQ|9p%7;4Mm9geK12)NC)7zXB2}owcmWSXu?PA@ z`3+xEQQz?UG-3qoLvUbndrwg>`-MMg`faod1*YZzDg~t0KQ)trJu$bvp58hA>i74c zWbKz4#+i*!K=86V`sW>pPxd9jMB;0{lWxOE;chBsFsfTjCj3n*W&Jw>He@%+6wqix z`a^z?Pme~uJR}6Spx>j*j8fJkd;(MzuT7DNR0S}gbNidZpF8Ta`537D==7;>5 zTg;j}cA0BES;cQ^2I0o4?u;2E%R#&2Suj>Suk~r0Ua?TAZ_XI8 znFMc^68SE+S43?TGpG@{#{Jg!+mvAZ#Owy-ci5A;kdJt)n`@n;j-D{psv^Y z-pOLUIC-#8H}C0#kyUH^32pkR_VMalEaU9S(-^fU|JYY#HI0!Arics4UM`-pc<21Q;*!L?l6Yew){^|})I2Gk z@Ql==RERkylRxyv2?+6}7p11ARK^!4<`(3n7ER7R95UJBXdH(w7evqGhey=xvc!SUaJ_NlTI#9AdC u1tmzJDOtf-lg~fXPy)LbBnDvsm8GVor6!kvgdk3x9R1vl(Qxvj=lTHs^nwup diff --git a/tests/fixtures/onnx_genai_workflows/vlm/policies/token_state_update.onnx b/tests/fixtures/onnx_genai_workflows/vlm/policies/token_state_update.onnx index 5ff11b46b3a74fa10132a719a2dede5cd17cbe8c..268d221d754c99df43c4bc05049adc3ae7c5e2bb 100644 GIT binary patch delta 93 zcmaFMzfgdegIkC#H$N$}wAiXoU?Xo58{^)|1#GHJ)0rkuXLHd{N-Rmv$jnPuvQkhg sNG*yl$}d;aQBcavOH3}wEK3FPi@|b5sl};9Ws|S5J!dqUypLTU0BL+5I{*Lx delta 45 zcmZ1|@RpyKgIkC#H$N$}wAktv|3=;fF B4oCn1 From c34cde4e9849735e4d0f2fcaa4d9af44f4ceb449 Mon Sep 17 00:00:00 2001 From: justinchuby Date: Fri, 14 Aug 2026 07:54:07 +0000 Subject: [PATCH 081/151] Bind row-wise emits to semantic request identities Carry explicit stable row IDs through decoder, VLM, and speculative loops and bind every ragged emit to them. Keep reusable serving slot IDs separate and stabilize decoder KV service cell numbering. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: justinchuby --- .../onnx_genai/auto_export_test.py | 8 ++ .../onnx_genai/workflow_metadata.py | 73 ++++++++++++++++++- .../onnx_genai/workflow_metadata_test.py | 7 ++ .../decoder/inference_metadata.yaml | 35 ++++++++- .../speculative/inference_metadata.yaml | 28 +++++++ .../vlm/inference_metadata.yaml | 27 +++++++ 6 files changed, 172 insertions(+), 6 deletions(-) diff --git a/src/mobius/integrations/onnx_genai/auto_export_test.py b/src/mobius/integrations/onnx_genai/auto_export_test.py index 46b44d863..a41e4b1db 100644 --- a/src/mobius/integrations/onnx_genai/auto_export_test.py +++ b/src/mobius/integrations/onnx_genai/auto_export_test.py @@ -178,6 +178,7 @@ def test_dispatch_decoder(tmp_path): "request.eos_ids", "request.eos_lengths", "request.row_max_iterations", + "request.row_ids", "request.grammar_mask", "request.rng_offset", } @@ -269,6 +270,13 @@ def test_seeded_decoder_sampler_uses_request_controls_and_direct_kv_carry(): ] assert workflow["state"]["rng_offset"]["class"] == "semantic" assert workflow["state"]["rng_offset"]["initializer"] == "request.rng_offset" + assert workflow["state"]["row_ids"]["initializer"] == "request.row_ids" + emit = next(node for node in workflow["steps"][0]["steps"] if node["kind"] == "emit") + assert emit["row_ids"] == "row_ids" + assert "emit_row_identity" in workflow["manifest"]["capabilities"] + assert set( + workflow["serving"]["kv_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", diff --git a/src/mobius/integrations/onnx_genai/workflow_metadata.py b/src/mobius/integrations/onnx_genai/workflow_metadata.py index 8878e99ec..fb8932084 100644 --- a/src/mobius/integrations/onnx_genai/workflow_metadata.py +++ b/src/mobius/integrations/onnx_genai/workflow_metadata.py @@ -236,6 +236,8 @@ def convert(node: dict[str, Any]) -> dict[str, Any]: result["valid_length"] = rewrite(node["valid_length"]) if "when" in node: result["when"] = rewrite(node["when"]) + if "row_ids" in node: + result["row_ids"] = rewrite(node["row_ids"]) return result if kind == "branch": result = { @@ -2763,6 +2765,12 @@ def build_vlm_workflow_metadata( "required": False, "default": -1, }, + "request.row_ids": { + "contract": batch_int, + "role": {"kind": "opaque"}, + "source": {"kind": "application", "name": "row_ids"}, + "required": True, + }, "package.eos_ids": { "contract": {"dtype": "int64", "rank": 1, "shape": [1]}, "role": {"kind": "opaque"}, @@ -3033,6 +3041,13 @@ def build_vlm_workflow_metadata( "initializer": "package.slot_ids", "recurrence": {"kind": "invariant"}, }, + "row_ids": { + "contract": batch_int, + "class": "semantic", + "scope": "invocation", + "initializer": "request.row_ids", + "recurrence": {"kind": "invariant"}, + }, "cache_lengths": { "contract": batch_int, "class": "semantic", @@ -3112,6 +3127,13 @@ def build_vlm_workflow_metadata( "state.slot_ids.body", "state.slot_ids.final", ), + ( + "row_ids", + "request.row_ids", + "state.row_ids.body", + "state.row_ids.body", + "state.row_ids.final", + ), ( "cache_lengths", "initializer.cache_lengths" if fixed_capacity else "package.zero_batch", @@ -3417,6 +3439,7 @@ def build_vlm_workflow_metadata( "mode": "append", "when": "state.active.body", "valid_length": "token.emitted_length", + "row_ids": "state.row_ids.body", "effect_name": "emit", "effect": _effect("emit.0", "emit.1"), }, @@ -3467,6 +3490,7 @@ def build_vlm_workflow_metadata( "loop_induction_values", "typed_emit", "emit_valid_length", + "emit_row_identity", *(["input_presence"] if text_only_vision is not None else []), *( ["serving_service_contract", "bounded_state_recurrence"] @@ -3716,6 +3740,12 @@ def build_speculative_workflow_metadata( "source": {"kind": "application", "name": "serving.slot_ids"}, "required": True, }, + "request.row_ids": { + "contract": batch_int, + "role": {"kind": "opaque"}, + "source": {"kind": "application", "name": "serving.row_ids"}, + "required": True, + }, "request.cache_lengths": { "contract": batch_int, "role": {"kind": "opaque"}, @@ -4037,6 +4067,7 @@ def build_speculative_workflow_metadata( "valid_length": emit_length, "output": "tokens", "mode": "append", + "row_ids": "state.row_ids.body", "effect_name": "emit", "effect": _effect("emit.0", "emit.1"), }, @@ -4049,6 +4080,7 @@ def build_speculative_workflow_metadata( "valid_length": "grammar.forced_length", "output": "tokens", "mode": "append", + "row_ids": "state.row_ids.body", "effect_name": "emit", "effect": _effect("emit.1", "emit.2"), } @@ -4096,6 +4128,13 @@ def build_speculative_workflow_metadata( "initializer": "request.slot_ids", "recurrence": {"kind": "invariant"}, }, + "row_ids": { + "contract": batch_int, + "class": "semantic", + "scope": "invocation", + "initializer": "request.row_ids", + "recurrence": {"kind": "invariant"}, + }, "cache_lengths": { "contract": batch_int, "class": "semantic", @@ -4147,6 +4186,13 @@ def build_speculative_workflow_metadata( "state.slot_ids.body", "state.slot_ids.final", ), + ( + "row_ids", + "request.row_ids", + "state.row_ids.body", + "state.row_ids.body", + "state.row_ids.final", + ), ( "cache_lengths", "request.cache_lengths", @@ -4278,6 +4324,7 @@ def build_speculative_workflow_metadata( "loop_induction_values", "typed_emit", "emit_valid_length", + "emit_row_identity", "bounded_state_recurrence", "serving_service_contract", ], @@ -4623,6 +4670,12 @@ def build_decoder_workflow_metadata( "required": False, "default": -1, }, + "request.row_ids": { + "contract": batch_int, + "role": {"kind": "opaque"}, + "source": {"kind": "application", "name": "row_ids"}, + "required": True, + }, } ) stochastic_sampler = sampler != "greedy" @@ -4863,6 +4916,13 @@ def build_decoder_workflow_metadata( "initializer": "package.slot_ids", "recurrence": {"kind": "invariant"}, }, + "row_ids": { + "contract": batch_int, + "class": "semantic", + "scope": "invocation", + "initializer": "request.row_ids", + "recurrence": {"kind": "invariant"}, + }, "cache_lengths": { "contract": batch_int, "class": "semantic", @@ -4957,6 +5017,13 @@ def build_decoder_workflow_metadata( "body_output": "state.slot_ids.body", "next": "state.slot_ids.final", }, + { + "cell": "row_ids", + "current": "request.row_ids", + "body_input": "state.row_ids.body", + "body_output": "state.row_ids.body", + "next": "state.row_ids.final", + }, ] ) if sampler_with_rng: @@ -5037,10 +5104,10 @@ def build_decoder_workflow_metadata( ) decoder_kv_ports: dict[str, Any] = {} decoder_kv_axis = 2 - for past, present in cache_pairs: + 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_{len(carried) - 1}" + 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 @@ -5311,6 +5378,7 @@ def build_decoder_workflow_metadata( { "when": "state.active.body", "valid_length": "token.emitted_length", + "row_ids": "state.row_ids.body", } if cache_pairs else {} @@ -5341,6 +5409,7 @@ def build_decoder_workflow_metadata( "nested_control_flow", "typed_emit", "emit_valid_length", + *(["emit_row_identity"] if cache_pairs else []), "loop_induction_values", *(["serving_service_contract"] if cache_pairs else []), *(["bounded_state_recurrence"] if cache_pairs else []), diff --git a/src/mobius/integrations/onnx_genai/workflow_metadata_test.py b/src/mobius/integrations/onnx_genai/workflow_metadata_test.py index f4429585c..0651a354a 100644 --- a/src/mobius/integrations/onnx_genai/workflow_metadata_test.py +++ b/src/mobius/integrations/onnx_genai/workflow_metadata_test.py @@ -48,7 +48,10 @@ def test_speculative_emit_uses_accepted_prefix_length(): ] emit = next(node for node in workflow["steps"][0]["steps"] if node["kind"] == "emit") assert emit["valid_length"] == "acceptance.length" + assert emit["row_ids"] == "row_ids" assert "emit_valid_length" in workflow["manifest"]["capabilities"] + assert "emit_row_identity" in workflow["manifest"]["capabilities"] + assert workflow["inputs"]["request.row_ids"]["source"]["name"] == "serving.row_ids" assert workflow["outputs"]["tokens"]["contract"]["shape"][-1] == "accepted_sequence" assert workflow["state"]["cache_0"]["recurrence"] == { "kind": "bounded", @@ -278,6 +281,9 @@ def collect_emits(node): ) assert emit["when"] == "active" assert emit["valid_length"] == "token.emitted_length" + assert emit["row_ids"] == "row_ids" + assert workflow["state"]["row_ids"]["initializer"] == "request.row_ids" + assert "emit_row_identity" 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( @@ -503,6 +509,7 @@ def test_speculative_workflow_uses_per_row_ragged_state_and_rng(): 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 emit["row_ids"] == "row_ids" assert not any(node["kind"] == "branch" for node in body) assert workflow["serving"]["active"] == "active" assert workflow["serving"]["done"] == "done" diff --git a/tests/fixtures/onnx_genai_workflows/decoder/inference_metadata.yaml b/tests/fixtures/onnx_genai_workflows/decoder/inference_metadata.yaml index 0c03c60c9..2906c593a 100644 --- a/tests/fixtures/onnx_genai_workflows/decoder/inference_metadata.yaml +++ b/tests/fixtures/onnx_genai_workflows/decoder/inference_metadata.yaml @@ -11,6 +11,7 @@ pipeline: - nested_control_flow - typed_emit - emit_valid_length + - emit_row_identity - loop_induction_values - serving_service_contract - bounded_state_recurrence @@ -131,6 +132,18 @@ pipeline: name: row_max_iterations required: false default: -1 + request.row_ids: + contract: + dtype: int64 + rank: 1 + shape: + - batch + role: + kind: opaque + source: + kind: application + name: row_ids + required: true request.temperature: contract: dtype: float32 @@ -470,6 +483,17 @@ pipeline: initializer: package.slot_ids recurrence: kind: invariant + row_ids: + contract: + dtype: int64 + rank: 1 + shape: + - batch + class: semantic + scope: invocation + initializer: request.row_ids + recurrence: + kind: invariant cache_lengths: contract: dtype: int64 @@ -514,7 +538,7 @@ pipeline: initializer: initializer.body_position_ids recurrence: kind: invariant - cache_10: + cache_0: contract: dtype: float32 rank: 4 @@ -547,7 +571,7 @@ pipeline: storage: shared_buffer ports: model: - cache_10: + cache_0: input: past_key_values.0.key output: present.0.key steps: @@ -670,6 +694,7 @@ pipeline: mode: append valid_length: token.emitted_length when: active + row_ids: row_ids - kind: invoke component: decoder_step_update inputs: @@ -683,7 +708,7 @@ pipeline: component: model inputs: input_ids: token.body - past_key_values.0.key: cache_10 + past_key_values.0.key: cache_0 attention_mask: decoder_step.body_attention_mask position_ids: position_ids outputs: @@ -714,13 +739,15 @@ pipeline: next: accepted_len.next - cell: slot_ids next: slot_ids + - cell: row_ids + next: row_ids - cell: rng_offset next: sample.next_offset - cell: attention_mask next: decoder_step.body_attention_mask - cell: position_ids next: decoder_step.body_position_ids - - cell: cache_10 + - cell: cache_0 next: decoder.body.present.0.key termination: generation_eos iteration: diff --git a/tests/fixtures/onnx_genai_workflows/speculative/inference_metadata.yaml b/tests/fixtures/onnx_genai_workflows/speculative/inference_metadata.yaml index 5e34fd4a2..21b54c37b 100644 --- a/tests/fixtures/onnx_genai_workflows/speculative/inference_metadata.yaml +++ b/tests/fixtures/onnx_genai_workflows/speculative/inference_metadata.yaml @@ -12,6 +12,7 @@ pipeline: - loop_induction_values - typed_emit - emit_valid_length + - emit_row_identity - bounded_state_recurrence - serving_service_contract - grammar_guidance_adapter @@ -133,6 +134,18 @@ pipeline: kind: application name: serving.slot_ids required: true + request.row_ids: + contract: + dtype: int64 + rank: 1 + shape: + - batch + role: + kind: opaque + source: + kind: application + name: serving.row_ids + required: true request.cache_lengths: contract: dtype: int64 @@ -601,6 +614,17 @@ pipeline: initializer: request.slot_ids recurrence: kind: invariant + row_ids: + contract: + dtype: int64 + rank: 1 + shape: + - batch + class: semantic + scope: invocation + initializer: request.row_ids + recurrence: + kind: invariant cache_lengths: contract: dtype: int64 @@ -812,11 +836,13 @@ pipeline: output: tokens mode: append valid_length: grammar.committed_length + row_ids: row_ids - kind: emit value: grammar.token output: tokens mode: append valid_length: grammar.forced_length + row_ids: row_ids continue_when: active max_iterations: request.max_iterations carried: @@ -832,6 +858,8 @@ pipeline: next: grammar.committed_length - cell: slot_ids next: slot_ids + - cell: row_ids + next: row_ids - cell: cache_lengths next: cache_lengths.next - cell: grammar diff --git a/tests/fixtures/onnx_genai_workflows/vlm/inference_metadata.yaml b/tests/fixtures/onnx_genai_workflows/vlm/inference_metadata.yaml index e2ebfb1c4..3ba37a0ad 100644 --- a/tests/fixtures/onnx_genai_workflows/vlm/inference_metadata.yaml +++ b/tests/fixtures/onnx_genai_workflows/vlm/inference_metadata.yaml @@ -85,6 +85,7 @@ pipeline: - loop_induction_values - typed_emit - emit_valid_length + - emit_row_identity - serving_service_contract - bounded_state_recurrence inputs: @@ -181,6 +182,18 @@ pipeline: name: row_max_iterations required: false default: -1 + request.row_ids: + contract: + dtype: int64 + rank: 1 + shape: + - batch + role: + kind: opaque + source: + kind: application + name: row_ids + required: true package.eos_ids: contract: dtype: int64 @@ -600,6 +613,17 @@ pipeline: initializer: package.slot_ids recurrence: kind: invariant + row_ids: + contract: + dtype: int64 + rank: 1 + shape: + - batch + class: semantic + scope: invocation + initializer: request.row_ids + recurrence: + kind: invariant cache_lengths: contract: dtype: int64 @@ -842,6 +866,7 @@ pipeline: mode: append valid_length: token.emitted_length when: active + row_ids: row_ids - kind: invoke component: decoder_step_update inputs: @@ -897,6 +922,8 @@ pipeline: next: accepted_len.next - cell: slot_ids next: slot_ids + - cell: row_ids + next: row_ids - cell: cache_lengths next: cache_lengths.next - cell: position_ids From 8784ad49500086fb444230ff9df6ecb93d2438db Mon Sep 17 00:00:00 2001 From: justinchuby Date: Fri, 14 Aug 2026 07:57:57 +0000 Subject: [PATCH 082/151] Validate workflows against semantic row runtime Pin ONNX GenAI validation to the row-identity runtime commit and sync its authoritative all-family/B>1 conformance harness. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: justinchuby --- .github/workflows/main.yml | 2 +- tests/onnx_genai_workflow_conformance.rs | 204 ++++++++++++++++++++--- 2 files changed, 186 insertions(+), 20 deletions(-) diff --git a/.github/workflows/main.yml b/.github/workflows/main.yml index 1443e5db9..8e3542a26 100644 --- a/.github/workflows/main.yml +++ b/.github/workflows/main.yml @@ -23,7 +23,7 @@ jobs: - uses: actions/checkout@v7 with: repository: justinchuby/onnx-genai - ref: 9e5757196b98542390ce11f4ff966a58ab3ef578 + ref: 7c64e1d3230e63b8f3cc50cc5d30929bd76d3a10 path: validation/onnx-genai - uses: actions/setup-python@v7 with: diff --git a/tests/onnx_genai_workflow_conformance.rs b/tests/onnx_genai_workflow_conformance.rs index 9b32f3e1f..dbcd8ed33 100644 --- a/tests/onnx_genai_workflow_conformance.rs +++ b/tests/onnx_genai_workflow_conformance.rs @@ -3,16 +3,23 @@ //! 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::{ - Engine, EngineConfig, GenerateOptions, GeneratePrompt, GenerateRequest, PipelineGenerateRequest, + Engine, EngineConfig, GenerateOptions, GeneratePrompt, GenerateRequest, + PipelineGenerateRequest, pipeline::WorkflowOutputRole, }; use onnx_genai_ort::{DataType, Value}; use std::path::PathBuf; fn root(name: &str) -> anyhow::Result { let root = std::env::var_os("MOBIUS_WORKFLOW_CONFORMANCE_DIR") - .ok_or_else(|| anyhow::anyhow!("MOBIUS_WORKFLOW_CONFORMANCE_DIR must be set"))?; - Ok(PathBuf::from(root).join(name)) + .map(PathBuf::from) + .unwrap_or_else(|| { + PathBuf::from(env!("CARGO_MANIFEST_DIR")) + .join("../../tests/fixtures/onnx_genai_workflows") + }); + Ok(root.join(name)) } fn options(max_new_tokens: usize) -> GenerateOptions { @@ -22,14 +29,163 @@ fn options(max_new_tokens: usize) -> GenerateOptions { options } +fn decoder_batch_request( + input_ids: &[i64], + batch: i64, + sequence: i64, + prompt_lengths: &[i64], + active: &[bool], + 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 slot_ids = (0..batch).collect::>(); + let row_ids = (100..100 + batch).collect::>(); + 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( + "package.slot_ids", + Value::from_slice_i64(&slot_ids, &[batch])?, + ) + .with_input( + "request.row_ids", + Value::from_slice_i64(&row_ids, &[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(PipelineGenerateRequest::new(GenerateRequest { - prompt: GeneratePrompt::TokenIds(vec![4, 5]), - options: options(3), - }))?; - assert_eq!(output["tokens"].to_vec_i64()?.len(), 3); + let output = engine.run_pipeline_outputs( + PipelineGenerateRequest::new(GenerateRequest { + prompt: GeneratePrompt::TokenIds(vec![4, 5]), + options: options(3), + }) + .with_input("request.row_ids", Value::from_slice_i64(&[0], &[1])?), + )?; + assert_eq!( + engine + .structured_output_for_role(&output, WorkflowOutputRole::Tokens) + .expect("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") + ); + + 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, 100); + 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, 101); + assert_eq!(rows[1].1.to_vec_i64()?, second_tokens); + + let inactive = decoder_batch_request(&[4, 5, 6, 0], 2, 2, &[2, 1], &[true, false], 3)?; + let inactive_output = engine.run_pipeline_outputs(inactive)?; + let inactive_rows = engine.output_rows_for_role(&inactive_output, WorkflowOutputRole::Tokens); + assert_eq!(inactive_rows.len(), 1); + assert_eq!(inactive_rows[0].0, 100); + assert_eq!(inactive_rows[0].1.to_vec_i64()?, first_tokens); + + let first_inactive = decoder_batch_request(&[4, 5, 6, 0], 2, 2, &[2, 1], &[false, true], 3)?; + let first_inactive_output = engine.run_pipeline_outputs(first_inactive)?; + let first_inactive_rows = + engine.output_rows_for_role(&first_inactive_output, WorkflowOutputRole::Tokens); + assert_eq!(first_inactive_rows.len(), 1); + assert_eq!(first_inactive_rows[0].0, 101); + assert_eq!(first_inactive_rows[0].1.to_vec_i64()?, second_tokens); + assert_eq!( + engine + .structured_output_for_role(&first_inactive_output, WorkflowOutputRole::Tokens) + .expect("semantic lookup must return the first emitted row") + .to_vec_i64()?, + second_tokens + ); + + let replay = decoder_batch_request(&[4, 5], 1, 2, &[2], &[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()?, + first_tokens + ); Ok(()) } @@ -49,9 +205,16 @@ fn mobius_vlm_workflow_executes_complete_image_path() -> anyhow::Result<()> { .with_input( "request.image", Value::from_raw_bytes(png, &[png_len], DataType::Uint8)?, + ) + .with_input("request.row_ids", Value::from_slice_i64(&[0], &[1])?); + 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] ); - let output = engine.run_pipeline(request)?; - assert_eq!(output["tokens"].shape(), [1, 2]); Ok(()) } @@ -63,12 +226,14 @@ fn mobius_euler_diffusion_workflow_executes_complete_path() -> anyhow::Result<() options: options(2), }) .with_input("latent", Value::from_slice_f32(&[1.0; 64], &[1, 4, 4, 4])?); - let output = engine.run_pipeline(request)?; + 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())); + assert!( + output["image"] + .to_vec_f32()? + .iter() + .all(|value| value.is_finite()) + ); Ok(()) } @@ -84,7 +249,7 @@ fn mobius_masked_diffusion_workflow_executes() -> anyhow::Result<()> { 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(request)?; + let output = engine.run_pipeline_outputs(request)?; assert_eq!(output["tokens"].shape(), [1, 2]); Ok(()) } @@ -98,7 +263,7 @@ fn mobius_codec_workflow_executes() -> anyhow::Result<()> { "request.waveform", Value::from_slice_f32(&[0.25, -0.5], &[1, 1, 2])?, ); - let output = engine.run_pipeline(request)?; + let output = engine.run_pipeline_outputs(request)?; assert_eq!(output["waveform"].to_vec_f32()?, [0.25, -0.5]); Ok(()) } @@ -106,7 +271,7 @@ fn mobius_codec_workflow_executes() -> anyhow::Result<()> { #[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(PipelineGenerateRequest::new(GenerateRequest { + let output = engine.run_pipeline_outputs(PipelineGenerateRequest::new(GenerateRequest { prompt: GeneratePrompt::TokenIds(vec![1, 2]), options: options(1), }))?; @@ -123,6 +288,7 @@ fn mobius_speculative_workflow_executes_rejection_and_correction() -> anyhow::Re options: options(1), }) .with_input("serving.slot_ids", Value::from_slice_i64(&[0], &[1])?) + .with_input("serving.row_ids", Value::from_slice_i64(&[0], &[1])?) .with_input( "verifier.past_key_values.0.key", Value::from_slice_f32(&[], &[1, 2, 0, 8])?, @@ -139,7 +305,7 @@ fn mobius_speculative_workflow_executes_rejection_and_correction() -> anyhow::Re ) .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(request)?; + let output = engine.run_pipeline_outputs(request)?; assert_eq!(output["tokens.row.0"].to_vec_i64()?, [1, 31]); Ok(()) } From 509c049ef38ca79c89288386b6f80980c5c810c2 Mon Sep 17 00:00:00 2001 From: justinchuby Date: Fri, 14 Aug 2026 08:03:17 +0000 Subject: [PATCH 083/151] Expose row-output conformance errors Include the complete multi-row generation error in the cross-repository assertion so semantic output failures remain diagnosable. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: justinchuby --- tests/onnx_genai_workflow_conformance.rs | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/tests/onnx_genai_workflow_conformance.rs b/tests/onnx_genai_workflow_conformance.rs index dbcd8ed33..4a3fd5741 100644 --- a/tests/onnx_genai_workflow_conformance.rs +++ b/tests/onnx_genai_workflow_conformance.rs @@ -135,7 +135,8 @@ fn mobius_decoder_rows_match_independent_runs_and_dynamic_batch_replay() -> anyh assert!( multi_row_error .to_string() - .contains("multi-row ragged output") + .contains("multi-row ragged output"), + "{multi_row_error:#}" ); let batched = decoder_batch_request(&[4, 5, 6, 0], 2, 2, &[2, 1], &[true, true], 3)?; From bdc10f030c65d193391b6ea656f7cbef68946c48 Mon Sep 17 00:00:00 2001 From: justinchuby Date: Fri, 14 Aug 2026 08:08:13 +0000 Subject: [PATCH 084/151] Exercise explicit batched policy controls Supply all per-row sampler, EOS, RNG, and grammar tensors in cross-repository B>1 conformance so symbolic batch defaults cannot collapse heterogeneous runtime requests. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: justinchuby --- tests/onnx_genai_workflow_conformance.rs | 38 ++++++++++++++++++++++++ 1 file changed, 38 insertions(+) diff --git a/tests/onnx_genai_workflow_conformance.rs b/tests/onnx_genai_workflow_conformance.rs index 4a3fd5741..dce086ddc 100644 --- a/tests/onnx_genai_workflow_conformance.rs +++ b/tests/onnx_genai_workflow_conformance.rs @@ -40,8 +40,12 @@ fn decoder_batch_request( 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)?]; let slot_ids = (0..batch).collect::>(); let row_ids = (100..100 + batch).collect::>(); + let grammar_mask = vec![1_u8; usize::try_from(batch * 128)?]; Ok(PipelineGenerateRequest::new(GenerateRequest { prompt: GeneratePrompt::TokenIds(vec![0]), options: options(max_new_tokens), @@ -71,6 +75,40 @@ fn decoder_batch_request( "request.row_ids", Value::from_slice_i64(&row_ids, &[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(&slot_ids, &[batch])?) + .with_input( + "request.grammar_mask", + Value::from_raw_bytes(grammar_mask, &[batch, 128], DataType::Bool)?, + ) + .with_input( + "request.rng_offset", + Value::from_slice_i64(&zeros, &[batch])?, + ) .with_input( "package.cache_lengths", Value::from_slice_i64(&zeros, &[batch])?, From f597878b18a83c06a964402cdf2e07ed5367dfc9 Mon Sep 17 00:00:00 2001 From: justinchuby Date: Fri, 14 Aug 2026 08:18:39 +0000 Subject: [PATCH 085/151] Bind row emits to carried serving slots Adopt the final runtime contract where explicit serving slot IDs provide semantic row identity, are preserved through loop compaction, and bind structured ragged emits directly. Remove the superseded duplicate row-ID state. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: justinchuby --- .github/workflows/main.yml | 2 +- .../onnx_genai/auto_export_test.py | 5 +- .../onnx_genai/workflow_metadata.py | 78 ++----------------- .../onnx_genai/workflow_metadata_test.py | 9 +-- .../decoder/inference_metadata.yaml | 33 +------- .../speculative/inference_metadata.yaml | 29 +------ .../vlm/inference_metadata.yaml | 33 +------- tests/onnx_genai_workflow_conformance.rs | 18 ++--- 8 files changed, 31 insertions(+), 176 deletions(-) diff --git a/.github/workflows/main.yml b/.github/workflows/main.yml index 8e3542a26..7ef53f219 100644 --- a/.github/workflows/main.yml +++ b/.github/workflows/main.yml @@ -23,7 +23,7 @@ jobs: - uses: actions/checkout@v7 with: repository: justinchuby/onnx-genai - ref: 7c64e1d3230e63b8f3cc50cc5d30929bd76d3a10 + ref: 5b51d1f67f3d3c3a66aa60af91d5110f0725b722 path: validation/onnx-genai - uses: actions/setup-python@v7 with: diff --git a/src/mobius/integrations/onnx_genai/auto_export_test.py b/src/mobius/integrations/onnx_genai/auto_export_test.py index a41e4b1db..097b64ad7 100644 --- a/src/mobius/integrations/onnx_genai/auto_export_test.py +++ b/src/mobius/integrations/onnx_genai/auto_export_test.py @@ -178,7 +178,7 @@ def test_dispatch_decoder(tmp_path): "request.eos_ids", "request.eos_lengths", "request.row_max_iterations", - "request.row_ids", + "package.slot_ids", "request.grammar_mask", "request.rng_offset", } @@ -270,9 +270,8 @@ def test_seeded_decoder_sampler_uses_request_controls_and_direct_kv_carry(): ] assert workflow["state"]["rng_offset"]["class"] == "semantic" assert workflow["state"]["rng_offset"]["initializer"] == "request.rng_offset" - assert workflow["state"]["row_ids"]["initializer"] == "request.row_ids" emit = next(node for node in workflow["steps"][0]["steps"] if node["kind"] == "emit") - assert emit["row_ids"] == "row_ids" + assert emit["row_ids"] == "slot_ids" assert "emit_row_identity" in workflow["manifest"]["capabilities"] assert set( workflow["serving"]["kv_service"]["groups"]["decoder_cache"]["ports"]["model"] diff --git a/src/mobius/integrations/onnx_genai/workflow_metadata.py b/src/mobius/integrations/onnx_genai/workflow_metadata.py index fb8932084..933ca5d83 100644 --- a/src/mobius/integrations/onnx_genai/workflow_metadata.py +++ b/src/mobius/integrations/onnx_genai/workflow_metadata.py @@ -2765,12 +2765,6 @@ def build_vlm_workflow_metadata( "required": False, "default": -1, }, - "request.row_ids": { - "contract": batch_int, - "role": {"kind": "opaque"}, - "source": {"kind": "application", "name": "row_ids"}, - "required": True, - }, "package.eos_ids": { "contract": {"dtype": "int64", "rank": 1, "shape": [1]}, "role": {"kind": "opaque"}, @@ -2822,9 +2816,8 @@ def build_vlm_workflow_metadata( "package.slot_ids": { "contract": batch_int, "role": {"kind": "opaque"}, - "source": {"kind": "literal"}, - "required": False, - "default": 0, + "source": {"kind": "application", "name": "serving.slot_ids"}, + "required": True, }, } inputs.update( @@ -3041,13 +3034,6 @@ def build_vlm_workflow_metadata( "initializer": "package.slot_ids", "recurrence": {"kind": "invariant"}, }, - "row_ids": { - "contract": batch_int, - "class": "semantic", - "scope": "invocation", - "initializer": "request.row_ids", - "recurrence": {"kind": "invariant"}, - }, "cache_lengths": { "contract": batch_int, "class": "semantic", @@ -3127,13 +3113,6 @@ def build_vlm_workflow_metadata( "state.slot_ids.body", "state.slot_ids.final", ), - ( - "row_ids", - "request.row_ids", - "state.row_ids.body", - "state.row_ids.body", - "state.row_ids.final", - ), ( "cache_lengths", "initializer.cache_lengths" if fixed_capacity else "package.zero_batch", @@ -3439,7 +3418,7 @@ def build_vlm_workflow_metadata( "mode": "append", "when": "state.active.body", "valid_length": "token.emitted_length", - "row_ids": "state.row_ids.body", + "row_ids": "state.slot_ids.body", "effect_name": "emit", "effect": _effect("emit.0", "emit.1"), }, @@ -3740,12 +3719,6 @@ def build_speculative_workflow_metadata( "source": {"kind": "application", "name": "serving.slot_ids"}, "required": True, }, - "request.row_ids": { - "contract": batch_int, - "role": {"kind": "opaque"}, - "source": {"kind": "application", "name": "serving.row_ids"}, - "required": True, - }, "request.cache_lengths": { "contract": batch_int, "role": {"kind": "opaque"}, @@ -4067,7 +4040,7 @@ def build_speculative_workflow_metadata( "valid_length": emit_length, "output": "tokens", "mode": "append", - "row_ids": "state.row_ids.body", + "row_ids": "state.slot_ids.body", "effect_name": "emit", "effect": _effect("emit.0", "emit.1"), }, @@ -4080,7 +4053,7 @@ def build_speculative_workflow_metadata( "valid_length": "grammar.forced_length", "output": "tokens", "mode": "append", - "row_ids": "state.row_ids.body", + "row_ids": "state.slot_ids.body", "effect_name": "emit", "effect": _effect("emit.1", "emit.2"), } @@ -4128,13 +4101,6 @@ def build_speculative_workflow_metadata( "initializer": "request.slot_ids", "recurrence": {"kind": "invariant"}, }, - "row_ids": { - "contract": batch_int, - "class": "semantic", - "scope": "invocation", - "initializer": "request.row_ids", - "recurrence": {"kind": "invariant"}, - }, "cache_lengths": { "contract": batch_int, "class": "semantic", @@ -4186,13 +4152,6 @@ def build_speculative_workflow_metadata( "state.slot_ids.body", "state.slot_ids.final", ), - ( - "row_ids", - "request.row_ids", - "state.row_ids.body", - "state.row_ids.body", - "state.row_ids.final", - ), ( "cache_lengths", "request.cache_lengths", @@ -4670,12 +4629,6 @@ def build_decoder_workflow_metadata( "required": False, "default": -1, }, - "request.row_ids": { - "contract": batch_int, - "role": {"kind": "opaque"}, - "source": {"kind": "application", "name": "row_ids"}, - "required": True, - }, } ) stochastic_sampler = sampler != "greedy" @@ -4814,9 +4767,8 @@ def build_decoder_workflow_metadata( "package.slot_ids": { "contract": batch_int, "role": {"kind": "opaque"}, - "source": {"kind": "literal"}, - "required": False, - "default": 0, + "source": {"kind": "application", "name": "serving.slot_ids"}, + "required": True, }, "package.cache_lengths": { "contract": batch_int, @@ -4916,13 +4868,6 @@ def build_decoder_workflow_metadata( "initializer": "package.slot_ids", "recurrence": {"kind": "invariant"}, }, - "row_ids": { - "contract": batch_int, - "class": "semantic", - "scope": "invocation", - "initializer": "request.row_ids", - "recurrence": {"kind": "invariant"}, - }, "cache_lengths": { "contract": batch_int, "class": "semantic", @@ -5017,13 +4962,6 @@ def build_decoder_workflow_metadata( "body_output": "state.slot_ids.body", "next": "state.slot_ids.final", }, - { - "cell": "row_ids", - "current": "request.row_ids", - "body_input": "state.row_ids.body", - "body_output": "state.row_ids.body", - "next": "state.row_ids.final", - }, ] ) if sampler_with_rng: @@ -5378,7 +5316,7 @@ def build_decoder_workflow_metadata( { "when": "state.active.body", "valid_length": "token.emitted_length", - "row_ids": "state.row_ids.body", + "row_ids": "state.slot_ids.body", } if cache_pairs else {} diff --git a/src/mobius/integrations/onnx_genai/workflow_metadata_test.py b/src/mobius/integrations/onnx_genai/workflow_metadata_test.py index 0651a354a..bf07eec9d 100644 --- a/src/mobius/integrations/onnx_genai/workflow_metadata_test.py +++ b/src/mobius/integrations/onnx_genai/workflow_metadata_test.py @@ -48,10 +48,10 @@ def test_speculative_emit_uses_accepted_prefix_length(): ] emit = next(node for node in workflow["steps"][0]["steps"] if node["kind"] == "emit") assert emit["valid_length"] == "acceptance.length" - assert emit["row_ids"] == "row_ids" + assert emit["row_ids"] == "slot_ids" assert "emit_valid_length" in workflow["manifest"]["capabilities"] assert "emit_row_identity" in workflow["manifest"]["capabilities"] - assert workflow["inputs"]["request.row_ids"]["source"]["name"] == "serving.row_ids" + assert workflow["inputs"]["request.slot_ids"]["source"]["name"] == "serving.slot_ids" assert workflow["outputs"]["tokens"]["contract"]["shape"][-1] == "accepted_sequence" assert workflow["state"]["cache_0"]["recurrence"] == { "kind": "bounded", @@ -281,8 +281,7 @@ def collect_emits(node): ) assert emit["when"] == "active" assert emit["valid_length"] == "token.emitted_length" - assert emit["row_ids"] == "row_ids" - assert workflow["state"]["row_ids"]["initializer"] == "request.row_ids" + assert emit["row_ids"] == "slot_ids" assert "emit_row_identity" in workflow["manifest"]["capabilities"] assert policy_invokes["decoder_step_update"]["inputs"]["logical_length"] == "cache_lengths" assert workflow["state"]["attention_mask"]["initializer"] == ("initializer.attention_mask") @@ -509,7 +508,7 @@ def test_speculative_workflow_uses_per_row_ragged_state_and_rng(): 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 emit["row_ids"] == "row_ids" + assert emit["row_ids"] == "slot_ids" assert not any(node["kind"] == "branch" for node in body) assert workflow["serving"]["active"] == "active" assert workflow["serving"]["done"] == "done" diff --git a/tests/fixtures/onnx_genai_workflows/decoder/inference_metadata.yaml b/tests/fixtures/onnx_genai_workflows/decoder/inference_metadata.yaml index 2906c593a..051536f0b 100644 --- a/tests/fixtures/onnx_genai_workflows/decoder/inference_metadata.yaml +++ b/tests/fixtures/onnx_genai_workflows/decoder/inference_metadata.yaml @@ -132,18 +132,6 @@ pipeline: name: row_max_iterations required: false default: -1 - request.row_ids: - contract: - dtype: int64 - rank: 1 - shape: - - batch - role: - kind: opaque - source: - kind: application - name: row_ids - required: true request.temperature: contract: dtype: float32 @@ -274,9 +262,9 @@ pipeline: role: kind: opaque source: - kind: literal - required: false - default: 0 + kind: application + name: serving.slot_ids + required: true package.cache_lengths: contract: dtype: int64 @@ -483,17 +471,6 @@ pipeline: initializer: package.slot_ids recurrence: kind: invariant - row_ids: - contract: - dtype: int64 - rank: 1 - shape: - - batch - class: semantic - scope: invocation - initializer: request.row_ids - recurrence: - kind: invariant cache_lengths: contract: dtype: int64 @@ -694,7 +671,7 @@ pipeline: mode: append valid_length: token.emitted_length when: active - row_ids: row_ids + row_ids: slot_ids - kind: invoke component: decoder_step_update inputs: @@ -739,8 +716,6 @@ pipeline: next: accepted_len.next - cell: slot_ids next: slot_ids - - cell: row_ids - next: row_ids - cell: rng_offset next: sample.next_offset - cell: attention_mask diff --git a/tests/fixtures/onnx_genai_workflows/speculative/inference_metadata.yaml b/tests/fixtures/onnx_genai_workflows/speculative/inference_metadata.yaml index 21b54c37b..08d50b25b 100644 --- a/tests/fixtures/onnx_genai_workflows/speculative/inference_metadata.yaml +++ b/tests/fixtures/onnx_genai_workflows/speculative/inference_metadata.yaml @@ -134,18 +134,6 @@ pipeline: kind: application name: serving.slot_ids required: true - request.row_ids: - contract: - dtype: int64 - rank: 1 - shape: - - batch - role: - kind: opaque - source: - kind: application - name: serving.row_ids - required: true request.cache_lengths: contract: dtype: int64 @@ -614,17 +602,6 @@ pipeline: initializer: request.slot_ids recurrence: kind: invariant - row_ids: - contract: - dtype: int64 - rank: 1 - shape: - - batch - class: semantic - scope: invocation - initializer: request.row_ids - recurrence: - kind: invariant cache_lengths: contract: dtype: int64 @@ -836,13 +813,13 @@ pipeline: output: tokens mode: append valid_length: grammar.committed_length - row_ids: row_ids + row_ids: slot_ids - kind: emit value: grammar.token output: tokens mode: append valid_length: grammar.forced_length - row_ids: row_ids + row_ids: slot_ids continue_when: active max_iterations: request.max_iterations carried: @@ -858,8 +835,6 @@ pipeline: next: grammar.committed_length - cell: slot_ids next: slot_ids - - cell: row_ids - next: row_ids - cell: cache_lengths next: cache_lengths.next - cell: grammar diff --git a/tests/fixtures/onnx_genai_workflows/vlm/inference_metadata.yaml b/tests/fixtures/onnx_genai_workflows/vlm/inference_metadata.yaml index 3ba37a0ad..7d52809f1 100644 --- a/tests/fixtures/onnx_genai_workflows/vlm/inference_metadata.yaml +++ b/tests/fixtures/onnx_genai_workflows/vlm/inference_metadata.yaml @@ -182,18 +182,6 @@ pipeline: name: row_max_iterations required: false default: -1 - request.row_ids: - contract: - dtype: int64 - rank: 1 - shape: - - batch - role: - kind: opaque - source: - kind: application - name: row_ids - required: true package.eos_ids: contract: dtype: int64 @@ -275,9 +263,9 @@ pipeline: role: kind: opaque source: - kind: literal - required: false - default: 0 + kind: application + name: serving.slot_ids + required: true request.temperature: contract: dtype: float32 @@ -613,17 +601,6 @@ pipeline: initializer: package.slot_ids recurrence: kind: invariant - row_ids: - contract: - dtype: int64 - rank: 1 - shape: - - batch - class: semantic - scope: invocation - initializer: request.row_ids - recurrence: - kind: invariant cache_lengths: contract: dtype: int64 @@ -866,7 +843,7 @@ pipeline: mode: append valid_length: token.emitted_length when: active - row_ids: row_ids + row_ids: slot_ids - kind: invoke component: decoder_step_update inputs: @@ -922,8 +899,6 @@ pipeline: next: accepted_len.next - cell: slot_ids next: slot_ids - - cell: row_ids - next: row_ids - cell: cache_lengths next: cache_lengths.next - cell: position_ids diff --git a/tests/onnx_genai_workflow_conformance.rs b/tests/onnx_genai_workflow_conformance.rs index dce086ddc..7ed6d35a0 100644 --- a/tests/onnx_genai_workflow_conformance.rs +++ b/tests/onnx_genai_workflow_conformance.rs @@ -44,7 +44,6 @@ fn decoder_batch_request( let floats_zero = vec![0.0_f32; usize::try_from(batch)?]; let floats_one = vec![1.0_f32; usize::try_from(batch)?]; let slot_ids = (0..batch).collect::>(); - let row_ids = (100..100 + batch).collect::>(); let grammar_mask = vec![1_u8; usize::try_from(batch * 128)?]; Ok(PipelineGenerateRequest::new(GenerateRequest { prompt: GeneratePrompt::TokenIds(vec![0]), @@ -71,10 +70,6 @@ fn decoder_batch_request( "package.slot_ids", Value::from_slice_i64(&slot_ids, &[batch])?, ) - .with_input( - "request.row_ids", - Value::from_slice_i64(&row_ids, &[batch])?, - ) .with_input( "request.eos_ids", Value::from_slice_i64(&vec![2_i64; usize::try_from(batch)?], &[batch, 1])?, @@ -127,7 +122,7 @@ fn mobius_decoder_workflow_executes() -> anyhow::Result<()> { prompt: GeneratePrompt::TokenIds(vec![4, 5]), options: options(3), }) - .with_input("request.row_ids", Value::from_slice_i64(&[0], &[1])?), + .with_input("package.slot_ids", Value::from_slice_i64(&[0], &[1])?), )?; assert_eq!( engine @@ -181,7 +176,7 @@ fn mobius_decoder_rows_match_independent_runs_and_dynamic_batch_replay() -> anyh 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, 100); + 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())?; @@ -191,14 +186,14 @@ fn mobius_decoder_rows_match_independent_runs_and_dynamic_batch_replay() -> anyh .structured_output_for_role(&second_output, WorkflowOutputRole::Tokens) .expect("independent second row must emit tokens") .to_vec_i64()?; - assert_eq!(rows[1].0, 101); + assert_eq!(rows[1].0, 1); assert_eq!(rows[1].1.to_vec_i64()?, second_tokens); let inactive = decoder_batch_request(&[4, 5, 6, 0], 2, 2, &[2, 1], &[true, false], 3)?; let inactive_output = engine.run_pipeline_outputs(inactive)?; let inactive_rows = engine.output_rows_for_role(&inactive_output, WorkflowOutputRole::Tokens); assert_eq!(inactive_rows.len(), 1); - assert_eq!(inactive_rows[0].0, 100); + assert_eq!(inactive_rows[0].0, 0); assert_eq!(inactive_rows[0].1.to_vec_i64()?, first_tokens); let first_inactive = decoder_batch_request(&[4, 5, 6, 0], 2, 2, &[2, 1], &[false, true], 3)?; @@ -206,7 +201,7 @@ fn mobius_decoder_rows_match_independent_runs_and_dynamic_batch_replay() -> anyh let first_inactive_rows = engine.output_rows_for_role(&first_inactive_output, WorkflowOutputRole::Tokens); assert_eq!(first_inactive_rows.len(), 1); - assert_eq!(first_inactive_rows[0].0, 101); + assert_eq!(first_inactive_rows[0].0, 1); assert_eq!(first_inactive_rows[0].1.to_vec_i64()?, second_tokens); assert_eq!( engine @@ -245,7 +240,7 @@ fn mobius_vlm_workflow_executes_complete_image_path() -> anyhow::Result<()> { "request.image", Value::from_raw_bytes(png, &[png_len], DataType::Uint8)?, ) - .with_input("request.row_ids", Value::from_slice_i64(&[0], &[1])?); + .with_input("package.slot_ids", Value::from_slice_i64(&[0], &[1])?); let output = engine.run_pipeline_outputs(request)?; assert_eq!( engine @@ -327,7 +322,6 @@ fn mobius_speculative_workflow_executes_rejection_and_correction() -> anyhow::Re options: options(1), }) .with_input("serving.slot_ids", Value::from_slice_i64(&[0], &[1])?) - .with_input("serving.row_ids", Value::from_slice_i64(&[0], &[1])?) .with_input( "verifier.past_key_values.0.key", Value::from_slice_f32(&[], &[1, 2, 0, 8])?, From 6b9bb7bfe63777282c2749b5ac90d6357849124b Mon Sep 17 00:00:00 2001 From: justinchuby Date: Fri, 14 Aug 2026 08:33:19 +0000 Subject: [PATCH 086/151] Match the exact batched policy v2 ABI Align sampler counters, singleton termination iteration, ragged EOS ports, and shape-identical selective state update to the authoritative fusion gate. Split token shaping and generated-length math into auxiliary ONNX components. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: justinchuby --- .github/workflows/main.yml | 2 +- src/mobius/generation/_policy_components.py | 84 +++----- .../generation/_policy_components_test.py | 76 ++----- .../onnx_genai/auto_export_test.py | 17 +- .../onnx_genai/workflow_metadata.py | 194 ++++++++---------- .../onnx_genai/workflow_metadata_test.py | 14 +- .../decoder/inference_metadata.yaml | 93 ++++----- .../policies/generated_length_update.onnx | Bin 0 -> 1064 bytes .../decoder/policies/iteration_broadcast.onnx | Bin 611 -> 0 bytes .../decoder/policies/termination.onnx | Bin 6378 -> 5932 bytes .../decoder/policies/token_sampler.onnx | Bin 59345 -> 57881 bytes .../decoder/policies/token_state_update.onnx | Bin 2081 -> 1338 bytes .../decoder/policies/token_to_slot.onnx | Bin 0 -> 480 bytes .../vlm/inference_metadata.yaml | 93 ++++----- .../vlm/policies/generated_length_update.onnx | Bin 0 -> 1064 bytes .../vlm/policies/iteration_broadcast.onnx | Bin 611 -> 0 bytes .../vlm/policies/termination.onnx | Bin 6378 -> 5932 bytes .../vlm/policies/token_sampler.onnx | Bin 59345 -> 57881 bytes .../vlm/policies/token_state_update.onnx | Bin 2081 -> 1338 bytes .../vlm/policies/token_to_slot.onnx | Bin 0 -> 480 bytes tests/onnx_genai_workflow_conformance.rs | 7 +- 21 files changed, 245 insertions(+), 335 deletions(-) create mode 100644 tests/fixtures/onnx_genai_workflows/decoder/policies/generated_length_update.onnx delete mode 100644 tests/fixtures/onnx_genai_workflows/decoder/policies/iteration_broadcast.onnx create mode 100644 tests/fixtures/onnx_genai_workflows/decoder/policies/token_to_slot.onnx create mode 100644 tests/fixtures/onnx_genai_workflows/vlm/policies/generated_length_update.onnx delete mode 100644 tests/fixtures/onnx_genai_workflows/vlm/policies/iteration_broadcast.onnx create mode 100644 tests/fixtures/onnx_genai_workflows/vlm/policies/token_to_slot.onnx diff --git a/.github/workflows/main.yml b/.github/workflows/main.yml index 7ef53f219..e8c6dbdd1 100644 --- a/.github/workflows/main.yml +++ b/.github/workflows/main.yml @@ -23,7 +23,7 @@ jobs: - uses: actions/checkout@v7 with: repository: justinchuby/onnx-genai - ref: 5b51d1f67f3d3c3a66aa60af91d5110f0725b722 + ref: f0046a8d89553bde8bc301071f5d0f09f6b7254e path: validation/onnx-genai - uses: actions/setup-python@v7 with: diff --git a/src/mobius/generation/_policy_components.py b/src/mobius/generation/_policy_components.py index b15876f2b..9c1c317c4 100644 --- a/src/mobius/generation/_policy_components.py +++ b/src/mobius/generation/_policy_components.py @@ -1145,7 +1145,7 @@ 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 grammar constraints remain request inputs; + 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") @@ -1155,9 +1155,8 @@ def build_seeded_categorical_sampler() -> PolicyComponent: 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"]) - grammar_mask = builder.input("grammar_mask", ir.DataType.BOOL, ["batch", "vocabulary"]) seed = builder.input("seed", ir.DataType.INT64, ["batch"]) - offset = builder.input("offset", 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)) @@ -1170,7 +1169,7 @@ def build_seeded_categorical_sampler() -> PolicyComponent: 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(offset, to=ir.DataType.UINT64), k0) + 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): @@ -1209,12 +1208,11 @@ def build_seeded_categorical_sampler() -> PolicyComponent: uniform = op.Cast(uniform, to=ir.DataType.FLOAT) blocked = op.CastLike(op.Constant(value_float=-3.4028235e38), logits) - constrained_logits = op.Where(grammar_mask, logits, blocked) safe_temperature = op.Unsqueeze( op.Max(temperature, op.Constant(value_float=1e-6)), [-1], ) - scaled_logits = op.Div(constrained_logits, safe_temperature) + scaled_logits = op.Div(logits, safe_temperature) safe_min_p = op.Unsqueeze( op.Clip( min_p, @@ -1328,24 +1326,14 @@ def build_seeded_categorical_sampler() -> PolicyComponent: axis=-1, keepdims=0, ) - has_allowed_token = op.ReduceMax( - op.Cast(grammar_mask, to=ir.DataType.INT64), - axes=[-1], - keepdims=0, - ) - token_ids = op.Where( - op.Greater(has_allowed_token, op.Constant(value_int=0)), - token_ids, - op.Constant(value_int=-1), - ) token_ids = op.Where(enabled, token_ids, op.Constant(value_int=-1)) - next_offset = op.Where( + next_counter = op.Where( enabled, - op.Add(offset, op.Constant(value_int=1)), - offset, + op.Add(counter, op.Constant(value_int=1)), + counter, ) builder.add_output(token_ids, "token") - builder.add_output(next_offset, "next_offset") + builder.add_output(next_counter, "next_counter") return _component( "onnx-genai.token-sampler@2", graph, @@ -1360,14 +1348,11 @@ def build_seeded_categorical_sampler() -> PolicyComponent: "top_k": "top_k", "top_p": "top_p", "min_p": "min_p", - "grammar_mask": "grammar_mask", "active": "active", "done": "done", - "rng": { - "rng_seed": "seed", - "rng_offset": "offset", - "rng_next_offset": "next_offset", - }, + "seed": "seed", + "counter": "counter", + "next_counter": "next_counter", "effect": "rng", }, "rng", @@ -1378,7 +1363,7 @@ 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("token_ids", ir.DataType.INT64, ["batch"]) + token_ids = builder.input("tokens", ir.DataType.INT64, ["batch"]) eos_ids = builder.input( "eos_ids", ir.DataType.INT64, @@ -1390,7 +1375,7 @@ def build_eos_termination(*, row_selective: bool = False) -> PolicyComponent: iteration = builder.input( "iteration", ir.DataType.INT64, - ["batch"], + [1] if row_selective else ["batch"], ) max_iterations = builder.input( "max_iterations", @@ -1424,11 +1409,9 @@ def build_eos_termination(*, row_selective: bool = False) -> PolicyComponent: ) if row_selective: active = builder.input("active", ir.DataType.BOOL, ["batch"]) - previous_done = builder.input("previous_done", ir.DataType.BOOL, ["batch"]) - enabled = op.And(active, op.Not(previous_done)) - newly_done = op.And(enabled, op.Or(hit_eos, hit_limit)) - done = op.Or(previous_done, newly_done) - next_active = op.And(enabled, op.Not(newly_done)) + 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) @@ -1450,7 +1433,7 @@ def build_eos_termination(*, row_selective: bool = False) -> PolicyComponent: graph, { "role": "termination_predicate", - "tokens": "token_ids", + "tokens": "tokens", "eos_ids": "eos_ids", "iteration": "iteration", "max_iterations": "max_iterations", @@ -1458,7 +1441,6 @@ def build_eos_termination(*, row_selective: bool = False) -> PolicyComponent: { "eos_lengths": "eos_lengths", "active": "active", - "previous_done": "previous_done", "batching": "per_row", "inactive_rows": "preserve", } @@ -1802,32 +1784,28 @@ def build_token_block_identity() -> PolicyComponent: def build_token_state_update(*, row_selective: bool = False) -> PolicyComponent: - """Selectively update token state and per-row generated lengths.""" + """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"]) + update = builder.input( + "update", + ir.DataType.INT64, + ["batch", 1] if row_selective else ["batch"], + ) if row_selective: - lengths = builder.input("lengths", 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)) next_state = op.Where( op.Unsqueeze(enabled, [-1]), - op.Unsqueeze(update, [-1]), + update, current, ) - emitted_length = op.Cast(enabled, to=ir.DataType.INT64) - next_lengths = op.Add(lengths, emitted_length) else: next_state = op.Unsqueeze(update, [-1]) next_state.shape = ir.Shape(["batch", 1]) builder.add_output(next_state, "next") - if row_selective: - next_lengths.shape = ir.Shape(["batch"]) - emitted_length.shape = ir.Shape(["batch"]) - builder.add_output(next_lengths, "next_lengths") - builder.add_output(emitted_length, "emitted_length") return _component( "onnx-genai.state-update@2" if row_selective else "onnx-genai.state-update@1", graph, @@ -1836,20 +1814,8 @@ def build_token_state_update(*, row_selective: bool = False) -> PolicyComponent: "current": "current", "update": "update", **({"batching": "per_row", "inactive_rows": "preserve"} if row_selective else {}), - **( - {"lengths": "lengths", "active": "active", "done": "done"} - if row_selective - else {} - ), + **({"active": "active", "done": "done"} if row_selective else {}), "next": "next", - **( - { - "next_lengths": "next_lengths", - "emitted_length": "emitted_length", - } - if row_selective - else {} - ), "effect": "state", }, "state", diff --git a/src/mobius/generation/_policy_components_test.py b/src/mobius/generation/_policy_components_test.py index dc628ccfa..159c54f11 100644 --- a/src/mobius/generation/_policy_components_test.py +++ b/src/mobius/generation/_policy_components_test.py @@ -454,7 +454,7 @@ def test_decoder_policy_chain_generates_multiple_tokens_from_prompt_only(tmp_pat build_eos_termination(), tmp_path, { - "token_ids": sample, + "tokens": sample, "eos_ids": np.array([7], np.int64), "iteration": np.array([iteration], np.int64), "max_iterations": max_output_tokens, @@ -501,9 +501,8 @@ def test_seeded_sampler_is_counter_based_and_reproducible(tmp_path): "top_k": np.array([0], np.int64), "top_p": np.array([1.0], np.float32), "min_p": np.array([0.0], np.float32), - "grammar_mask": np.array([[True, True, True, True]], np.bool_), "seed": np.array([7], np.int64), - "offset": np.array([11], np.int64), + "counter": np.array([11], np.int64), "active": np.array([True], np.bool_), "done": np.array([False], np.bool_), } @@ -513,7 +512,7 @@ def test_seeded_sampler_is_counter_based_and_reproducible(tmp_path): np.testing.assert_array_equal(first[1], [12]) -def test_seeded_sampler_applies_request_top_k_and_grammar_mask(tmp_path): +def test_seeded_sampler_applies_request_top_k(tmp_path): (token, _) = _run( build_seeded_categorical_sampler(), tmp_path, @@ -523,38 +522,17 @@ def test_seeded_sampler_applies_request_top_k_and_grammar_mask(tmp_path): "top_k": np.array([1], np.int64), "top_p": np.array([0.5], np.float32), "min_p": np.array([0.0], np.float32), - "grammar_mask": np.array([[False, True, True, True]], np.bool_), "seed": np.array([17], np.int64), - "offset": np.array([0], 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, [1]) - - -def test_seeded_sampler_rejects_empty_grammar_vocabulary(tmp_path): - (token, _) = _run( - build_seeded_categorical_sampler(), - tmp_path, - { - "logits": np.array([[1.0, 2.0, 3.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), - "grammar_mask": np.array([[False, False, False]], np.bool_), - "seed": np.array([1], np.int64), - "offset": np.array([0], np.int64), - "active": np.array([True], np.bool_), - "done": np.array([False], np.bool_), - }, - ) - np.testing.assert_array_equal(token, [-1]) + np.testing.assert_array_equal(token, [0]) def test_seeded_sampler_applies_request_min_p_in_logit_space(tmp_path): - (token, next_offset) = _run( + (token, next_counter) = _run( build_seeded_categorical_sampler(), tmp_path, { @@ -563,15 +541,14 @@ def test_seeded_sampler_applies_request_min_p_in_logit_space(tmp_path): "top_k": np.array([0], np.int64), "top_p": np.array([1.0], np.float32), "min_p": np.array([0.95], np.float32), - "grammar_mask": np.array([[True, True, True]], np.bool_), "seed": np.array([23], np.int64), - "offset": np.array([4], 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_offset, [5]) + np.testing.assert_array_equal(next_counter, [5]) def test_seeded_sampler_heterogeneous_batch_matches_independent_rows(tmp_path): @@ -590,17 +567,8 @@ def test_seeded_sampler_heterogeneous_batch_matches_independent_rows(tmp_path): "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), - "grammar_mask": np.array( - [ - [True, True, True, True, True], - [True, False, True, True, True], - [True, True, True, False, False], - [True, True, True, True, True], - ], - np.bool_, - ), "seed": np.array([3, 7, 11, 13], np.int64), - "offset": np.array([0, 5, 9, 12], 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_), } @@ -615,35 +583,31 @@ def test_seeded_sampler_heterogeneous_batch_matches_independent_rows(tmp_path): def test_row_selective_state_and_termination_preserve_inactive_rows(tmp_path): - next_state, next_lengths, emitted = _run( + (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), - "lengths": np.array([2, 4, 6], 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]]) - np.testing.assert_array_equal(next_lengths, [3, 4, 6]) - np.testing.assert_array_equal(emitted, [1, 0, 0]) done, next_active, continued = _run( build_eos_termination(row_selective=True), tmp_path, { - "token_ids": np.array([2, 8, 9], np.int64), + "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, 0, 0], 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_), - "previous_done": np.array([False, False, True], np.bool_), }, ) - np.testing.assert_array_equal(done, [True, False, True]) + 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]) @@ -653,13 +617,12 @@ def test_row_selective_termination_heterogeneous_batch_matches_independent_rows( ): component = build_eos_termination(row_selective=True) feeds = { - "token_ids": np.array([2, 9, 5, 7], np.int64), + "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([0, 1, 4, 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_), - "previous_done": np.array([False, False, False, False], np.bool_), } done, next_active, continued = _run(component, tmp_path, feeds) np.testing.assert_array_equal(done, [True, False, True, True]) @@ -669,7 +632,10 @@ def test_row_selective_termination_heterogeneous_batch_matches_independent_rows( row_outputs = _run( component, tmp_path, - {name: value[row : row + 1] for name, value in feeds.items()}, + { + 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]) @@ -707,7 +673,7 @@ def test_eos_termination_runtime(tmp_path): build_eos_termination(), tmp_path, { - "token_ids": np.array([2, 8, 9], np.int64), + "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), diff --git a/src/mobius/integrations/onnx_genai/auto_export_test.py b/src/mobius/integrations/onnx_genai/auto_export_test.py index 097b64ad7..a0eae019b 100644 --- a/src/mobius/integrations/onnx_genai/auto_export_test.py +++ b/src/mobius/integrations/onnx_genai/auto_export_test.py @@ -179,8 +179,7 @@ def test_dispatch_decoder(tmp_path): "request.eos_lengths", "request.row_max_iterations", "package.slot_ids", - "request.grammar_mask", - "request.rng_offset", + "request.rng_counter", } assert workflow["inputs"]["request.prompt_lengths"]["default"] == -1 assert [node["component"] for node in workflow["steps"][0]["setup"]] == [ @@ -236,10 +235,9 @@ def test_seeded_decoder_sampler_uses_request_controls_and_direct_kv_carry(): "top_k": "top_k", "top_p": "top_p", "min_p": "min_p", - "grammar_mask": "grammar_mask", - "rng_seed": "seed", - "rng_offset": "offset", - "rng_next_offset": "next_offset", + "seed": "seed", + "counter": "counter", + "next_counter": "next_counter", "active": "active", "done": "done", } @@ -254,9 +252,8 @@ def test_seeded_decoder_sampler_uses_request_controls_and_direct_kv_carry(): "top_k": "request.top_k", "top_p": "request.top_p", "min_p": "request.min_p", - "grammar_mask": "request.grammar_mask", "seed": "request.seed", - "offset": "rng_offset", + "counter": "rng_counter", "active": "active", "done": "done", } @@ -268,8 +265,8 @@ def test_seeded_decoder_sampler_uses_request_controls_and_direct_kv_carry(): "batch", "num_eos", ] - assert workflow["state"]["rng_offset"]["class"] == "semantic" - assert workflow["state"]["rng_offset"]["initializer"] == "request.rng_offset" + 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 emit["row_ids"] == "slot_ids" assert "emit_row_identity" in workflow["manifest"]["capabilities"] diff --git a/src/mobius/integrations/onnx_genai/workflow_metadata.py b/src/mobius/integrations/onnx_genai/workflow_metadata.py index 933ca5d83..ea5ce7862 100644 --- a/src/mobius/integrations/onnx_genai/workflow_metadata.py +++ b/src/mobius/integrations/onnx_genai/workflow_metadata.py @@ -28,7 +28,6 @@ build_greedy_sampler, build_integer_add, build_integer_minimum, - build_integer_row_broadcast, build_last_token_logits, build_model_token_cast, build_proposal_metrics, @@ -2698,11 +2697,12 @@ def build_vlm_workflow_metadata( "termination_batch_initializer", build_termination_batch_initializer(), ) - pkg.add_policy_component("iteration_broadcast", build_integer_row_broadcast()) 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()) @@ -2873,21 +2873,10 @@ def build_vlm_workflow_metadata( "required": False, "default": 0, }, - "request.grammar_mask": { - "contract": { - "dtype": "bool", - "rank": 2, - "shape": [batch, _contract(logits_output)["shape"][-1]], - }, - "role": {"kind": "opaque"}, - "source": {"kind": "application", "name": "grammar_mask"}, - "required": False, - "default": True, - }, - "request.rng_offset": { + "request.rng_counter": { "contract": batch_int, "role": {"kind": "opaque"}, - "source": {"kind": "application", "name": "rng_offset"}, + "source": {"kind": "application", "name": "rng_counter"}, "required": False, "default": 0, }, @@ -3043,11 +3032,11 @@ def build_vlm_workflow_metadata( ), "recurrence": {"kind": "invariant"}, }, - "rng_offset": { + "rng_counter": { "contract": batch_int, "class": "semantic", "scope": "invocation", - "initializer": "request.rng_offset", + "initializer": "request.rng_counter", "recurrence": {"kind": "invariant"}, }, } @@ -3085,11 +3074,11 @@ def build_vlm_workflow_metadata( "state.generated_lengths.final", ), ( - "rng_offset", - "request.rng_offset", - "state.rng_offset.body", - "sample.next_offset", - "state.rng_offset.final", + "rng_counter", + "request.rng_counter", + "state.rng_counter.body", + "sample.next_counter", + "state.rng_counter.final", ), ( "active", @@ -3329,11 +3318,6 @@ def build_vlm_workflow_metadata( body = { "kind": "sequence", "nodes": [ - _invoke( - "iteration_broadcast", - {"value": "loop.iteration", "active": "state.active.body"}, - {"rows": "loop.iteration_rows"}, - ), _invoke( "token_sampler", { @@ -3342,25 +3326,23 @@ def build_vlm_workflow_metadata( "top_k": "request.top_k", "top_p": "request.top_p", "min_p": "request.min_p", - "grammar_mask": "request.grammar_mask", "seed": "request.seed", - "offset": "state.rng_offset.body", + "counter": "state.rng_counter.body", "active": "state.active.body", "done": "state.done.body", }, - {"token": "sample.body", "next_offset": "sample.next_offset"}, + {"token": "sample.body", "next_counter": "sample.next_counter"}, {"sample": _effect("sample.0", "sample.1")}, ), _invoke( "termination", { - "token_ids": "sample.body", + "tokens": "sample.body", "eos_ids": "termination.eos_ids", "eos_lengths": "termination.eos_lengths", - "iteration": "loop.iteration_rows", + "iteration": "loop.iteration", "max_iterations": "termination.max_iterations", "active": "state.active.body", - "previous_done": "state.done.body", }, { "done": "loop.done", @@ -3395,20 +3377,36 @@ def build_vlm_workflow_metadata( if cache_pairs else [] ), + _invoke("token_to_slot", {"token": "sample.body"}, {"slot": "sample.slot"}), _invoke( - "token_state_update", + "generated_length_update", { - "current": "state.token.body", - "update": "sample.body", - "lengths": "state.generated_lengths.body", + "left": "state.generated_lengths.body", + "right": "package.one", "active": "state.active.body", "done": "state.done.body", }, + {"total": "token.next_lengths"}, + ), + _invoke( + "generated_length_update", { - "next": "token.body", - "next_lengths": "token.next_lengths", - "emitted_length": "token.emitted_length", + "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")}, ), { @@ -4709,24 +4707,10 @@ def build_decoder_workflow_metadata( "required": False, "default": 0, }, - "request.grammar_mask": { - "contract": { - "dtype": "bool", - "rank": 2, - "shape": [ - batch_dimension, - _contract(logits_output)["shape"][-1], - ], - }, - "role": {"kind": "opaque"}, - "source": {"kind": "application", "name": "grammar_mask"}, - "required": False, - "default": True, - }, - "request.rng_offset": { + "request.rng_counter": { "contract": batch_int, "role": {"kind": "opaque"}, - "source": {"kind": "application", "name": "rng_offset"}, + "source": {"kind": "application", "name": "rng_counter"}, "required": False, "default": 0, }, @@ -4743,11 +4727,12 @@ def build_decoder_workflow_metadata( "termination_batch_initializer", build_termination_batch_initializer(), ) - pkg.add_policy_component("iteration_broadcast", build_integer_row_broadcast()) 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": { @@ -4965,23 +4950,23 @@ def build_decoder_workflow_metadata( ] ) if sampler_with_rng: - state["rng_offset"] = { + state["rng_counter"] = { "contract": batch_int, "scope": "invocation", "class": "semantic", - "initializer": "request.rng_offset", + "initializer": "request.rng_counter", "recurrence": {"kind": "invariant"}, } - initial_effects["state:rng_offset"] = "state:rng_offset.0" + initial_effects["state:rng_counter"] = "state:rng_counter.0" carried.append( { - "cell": "rng_offset", - "current": "request.rng_offset", - "body_input": "state.rng_offset.body", - "body_output": "sample.next_offset", - "next": "state.rng_offset.final", - "read_effect": _effect("state:rng_offset.0", "state:rng_offset.read"), - "write_effect": _effect("state:rng_offset.read", "state:rng_offset.1"), + "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 = { @@ -5171,17 +5156,6 @@ def build_decoder_workflow_metadata( body = { "kind": "sequence", "nodes": [ - *( - [ - _invoke( - "iteration_broadcast", - {"value": "loop.iteration", "active": "state.active.body"}, - {"rows": "loop.iteration_rows"}, - ) - ] - if cache_pairs - else [] - ), _invoke( "token_sampler", { @@ -5192,9 +5166,8 @@ def build_decoder_workflow_metadata( "top_k": "request.top_k", "top_p": "request.top_p", "min_p": "request.min_p", - "grammar_mask": "request.grammar_mask", "seed": "request.seed", - "offset": "state.rng_offset.body", + "counter": "state.rng_counter.body", } if sampler_with_rng else {} @@ -5210,18 +5183,48 @@ def build_decoder_workflow_metadata( }, { "token": "sample.body", - **({"next_offset": "sample.next_offset"} if sampler_with_rng else {}), + **({"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.body", + "update": "sample.slot" if cache_pairs else "sample.body", **( { - "lengths": "state.generated_lengths.body", "active": "state.active.body", "done": "state.done.body", } @@ -5229,17 +5232,7 @@ def build_decoder_workflow_metadata( else {} ), }, - { - "next": "token.body", - **( - { - "next_lengths": "token.next_lengths", - "emitted_length": "token.emitted_length", - } - if cache_pairs - else {} - ), - }, + {"next": "token.body"}, {"state": _effect("state.0", "state.1")}, ), *( @@ -5256,23 +5249,16 @@ def build_decoder_workflow_metadata( _invoke( "termination", { - "token_ids": "sample.body", + "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_rows" if cache_pairs else "loop.iteration"), + "iteration": "loop.iteration", "max_iterations": ( "termination.max_iterations" if cache_pairs else "request.max_iterations" ), - **( - { - "active": "state.active.body", - "previous_done": "state.done.body", - } - if cache_pairs - else {} - ), + **({"active": "state.active.body"} if cache_pairs else {}), }, { "done": "loop.done", diff --git a/src/mobius/integrations/onnx_genai/workflow_metadata_test.py b/src/mobius/integrations/onnx_genai/workflow_metadata_test.py index bf07eec9d..e7af9fb00 100644 --- a/src/mobius/integrations/onnx_genai/workflow_metadata_test.py +++ b/src/mobius/integrations/onnx_genai/workflow_metadata_test.py @@ -262,10 +262,16 @@ def collect_decoder_invokes(node): ) assert policy_invokes["token_sampler"]["inputs"]["active"] == "active" assert policy_invokes["token_sampler"]["inputs"]["done"] == "done" - assert policy_invokes["token_state_update"]["inputs"]["lengths"] == ("generated_lengths") - assert policy_invokes["token_state_update"]["outputs"]["next_lengths"] == ( - "token.next_lengths" - ) + 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): diff --git a/tests/fixtures/onnx_genai_workflows/decoder/inference_metadata.yaml b/tests/fixtures/onnx_genai_workflows/decoder/inference_metadata.yaml index 051536f0b..7075a0095 100644 --- a/tests/fixtures/onnx_genai_workflows/decoder/inference_metadata.yaml +++ b/tests/fixtures/onnx_genai_workflows/decoder/inference_metadata.yaml @@ -202,21 +202,7 @@ pipeline: kind: request required: false default: 0 - request.grammar_mask: - contract: - dtype: bool - rank: 2 - shape: - - batch - - 128 - role: - kind: opaque - source: - kind: application - name: grammar_mask - required: false - default: true - request.rng_offset: + request.rng_counter: contract: dtype: int64 rank: 1 @@ -226,7 +212,7 @@ pipeline: kind: opaque source: kind: application - name: rng_offset + name: rng_counter required: false default: 0 package.active: @@ -318,12 +304,11 @@ pipeline: top_k: top_k top_p: top_p min_p: min_p - grammar_mask: grammar_mask active: active done: done - rng_seed: seed - rng_offset: offset - rng_next_offset: next_offset + seed: seed + counter: counter + next_counter: next_counter parameters: mode: seeded_stochastic batching: per_row @@ -337,13 +322,12 @@ pipeline: id: onnx-genai.termination-predicate version: '2' bindings: - tokens: token_ids + tokens: tokens eos_ids: eos_ids iteration: iteration max_iterations: max_iterations eos_lengths: eos_lengths active: active - previous_done: previous_done done: done next_active: next_active continue: continue @@ -360,12 +344,9 @@ pipeline: bindings: current: current update: update - lengths: lengths active: active done: done next: next - next_lengths: next_lengths - emitted_length: emitted_length parameters: batching: per_row inactive_rows: preserve @@ -389,10 +370,14 @@ pipeline: implementation: kind: onnx artifact: policies/termination_batch_initializer.onnx - iteration_broadcast: + token_to_slot: + implementation: + kind: onnx + artifact: policies/token_to_slot.onnx + generated_length_update: implementation: kind: onnx - artifact: policies/iteration_broadcast.onnx + artifact: policies/generated_length_update.onnx state: token: contract: @@ -482,7 +467,7 @@ pipeline: initializer: initializer.cache_lengths recurrence: kind: invariant - rng_offset: + rng_counter: contract: dtype: int64 rank: 1 @@ -490,7 +475,7 @@ pipeline: - batch scope: invocation class: semantic - initializer: request.rng_offset + initializer: request.rng_counter recurrence: kind: invariant attention_mask: @@ -598,13 +583,6 @@ pipeline: outputs: last_logits: decoder.setup.last_logits steps: - - kind: invoke - component: iteration_broadcast - inputs: - value: loop.iteration - active: active - outputs: - rows: loop.iteration_rows - kind: invoke component: token_sampler inputs: @@ -613,36 +591,55 @@ pipeline: top_k: request.top_k top_p: request.top_p min_p: request.min_p - grammar_mask: request.grammar_mask seed: request.seed - offset: rng_offset + counter: rng_counter active: active done: done outputs: token: sample.body - next_offset: sample.next_offset + 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.body - lengths: generated_lengths + update: sample.slot active: active done: done outputs: next: token.body - next_lengths: token.next_lengths - emitted_length: token.emitted_length - kind: invoke component: termination inputs: - token_ids: sample.body + tokens: sample.body eos_ids: termination.eos_ids eos_lengths: termination.eos_lengths - iteration: loop.iteration_rows + iteration: loop.iteration max_iterations: termination.max_iterations active: active - previous_done: done outputs: done: loop.done continue: loop.continue @@ -716,8 +713,8 @@ pipeline: next: accepted_len.next - cell: slot_ids next: slot_ids - - cell: rng_offset - next: sample.next_offset + - cell: rng_counter + next: sample.next_counter - cell: attention_mask next: decoder_step.body_attention_mask - cell: position_ids diff --git a/tests/fixtures/onnx_genai_workflows/decoder/policies/generated_length_update.onnx b/tests/fixtures/onnx_genai_workflows/decoder/policies/generated_length_update.onnx new file mode 100644 index 0000000000000000000000000000000000000000..81bc28f0bf8faea2948bc40eb1e892115033e890 GIT binary patch literal 1064 zcmcIjO-lnY5bgHcQ9&Yt_8?k{dXU8`sCcp-d+;hK64GS48(ouSH&u)1p;yoTUz2Q$ zAc*wP(@ch$yf^b^n#-jAqz7`8Uwk#;1*$?TB#mc!tqnbDfip@Aq(|+F_4})^q}UVY zDdR{={SIuWk5Qo1DL2UHUr8mMJ5ijBt?f*xCZYdrq3#TN_1{|sF7zGv?^ z@PZPl8Esh7TDgfa!-C66yXIdDng-d(F+nd0mC9g*8D&D;o=zEsaM)~+1K1;-R40;; z!w8iX-JqZes>sB2lnVO{2h8yJ_G~uRc?-{Hn?YRV(j1VMs^8{0>jsvpAGZgbqsNE! zpXF^r-Eu8MsR(*e$?u>O*asUfa}w;sZgGpVSxT+KV!D)TWioprr2O_i33|}^HCX!|2%kCuDgPlMBz~Rtp?{D*}3Z=EpZ+d!DCjbBd literal 0 HcmV?d00001 diff --git a/tests/fixtures/onnx_genai_workflows/decoder/policies/iteration_broadcast.onnx b/tests/fixtures/onnx_genai_workflows/decoder/policies/iteration_broadcast.onnx deleted file mode 100644 index 6618ad65560e12af5e2f95f56cf1d059c7b64cda..0000000000000000000000000000000000000000 GIT binary patch literal 0 HcmV?d00001 literal 611 zcmcIh!AiqG6l_A8rcVp9B5Ds(OHuIedT1XhBnu z2{VdNhH8Tz9H^&pD213As=2ZuO^Go$;Y<@vlR0wl9|j$0e`75*u%S{hvt+=#1NDrg zQwA#o$kX!Kv|G0`MMR7XAN{Uh!hbS?cDS%?%rwq#<48*qQ)29JyRwG>f}nvsa06#} z@s-*soCHnOgD&bQnQ}T0xy=KxT#A)Q7!ga8=p2f_&~m5ox1Gn`^;p&{ESzIFENcwM zOjIYs(gUVaSe^HZ+?Usda}B*;nV83DzhYI7km-!4oap&gf8iWK=SK}T8Zl~%I&VCv HZS8#mtB=JC diff --git a/tests/fixtures/onnx_genai_workflows/decoder/policies/termination.onnx b/tests/fixtures/onnx_genai_workflows/decoder/policies/termination.onnx index 46ce9432722602b317df0a22c323084f5f063c85..3af2ca2e7abf2118f6f33cbf12082031af251140 100644 GIT binary patch delta 515 zcmY+Aze~eV5XVg{O>%AegFHNMT&!)e}Rj$;Oyj}i;E^!t4{as?t7oR_kGrTIa!R5@Sd|(Z?|gQ+@*XabvX#T z&MvdBZyx}=(JaxcDy%iB?O2TJRgGV3;zD#Y(eZ`mwp;fY2p;u|T6BN@7(|aQ*cLVN z#DA?iE;V#D62z~0L)S+ABmm5u{Q`zk60zBVO9xzBl^b(M@V7)lmSeNOz`_uSEAfeV zQH?H|t)Fd0Y{%QPMH2JqIK~?d_>A)9nPpME@=x%V)`Zlzx6{h3HWugxjuH{VqC|KYX{15&?+U(b6!Ait pS2H)-)Ag&qX0coLyZx7iUk;;tl{Z7*Jm ziXW%)L=Zf9(u&x5D9I z<+HnIyw+bKTutITJYr$6|Nav`w~MfGug}vk9u8O#N1UiByA~(RUA8r8qd4Hgfse*U z?IFU7-%Y}l;~WJG`Zhe#E)eS4c)g#4sq2*1`+L7`o_>TR?s60`x`-<&6A0?c23d&0 zn$dFw=CFWq2_RM1P8q~67+#r|U)(IUQ+uzV0(lUlMN=|%XW zY{0yEe8#i&RevxVo_erZ3E_jOH$7se>;?~pU4GTyVYF@6CSFd`hM(qx*^X_k#?c`0 zqyHow_=C3L?~G}+v5~TImUiHTWus-dUfqNj))j4Gav}I;z0rjqCcN(_VI0AC>nyyj zb&XrC!q&hf=MOY;yVR&>dkF&Owuy=rwL(gG3Le)THOq*5!hXH!%@&R!P*sWQ<+buw zG^;H{z9jn=ey WM0nnKcIwa*IxAH!}q$Y!emfd+)xFdqIqdq$UNob~RdqOsBge-DP+7 zYjz7n6KcYTRoIqeD^%Np0u`Z6D`K=1T$GwdQ;jk4kB0q#EofpI%Rf#0r}5nT-n@^Q z<)4{*-@EsobARV|&iTE+|E=M~^@g^l7J23H$m0W(6B`a(f1^!V(Q-H@4D~CpQ4OAwij;Vsb zNGBS%^u>WGw4^JoeBfHrVgC=lg}%n(q|Q}B&{YN2{ieVdz@ML44|^{3!T+9(uOm|X zRKdHMP7GPX`Ov^98ZebMF>)=Kka=}5N=)IqwS?~$_Z=9t>bfNLakL6re?gV;3Xhcf zC;LhcLbf9NQ}=8eo`i#|EPlCdBfHQxRa*JLeNFQXIO8xvfbEOMOa}V|ZR^Vxnh!>)2#*sGu8~5)VFb4mC9R z_j7yUD8|mv>j5y}_tX6kUJ!A1L&P0*y_4fpg$Kt=MdX!lhA7FJXRajzn|{16qQjcG zM25&sQ-z*!^vY0L@ZVixXrtQ(>$%y1F;09gW@cmaZvSBLiUw@)BUL}}e!iPr+ z@c!|xn5DG$K2aR?-dNCb;_0|u50(vMqbc}nWC!G4Nabz0&ASM?M3qkW235Z(NUs#( zw=Z;RR9aCP87=IVTBZt)Dc!JSoH!oo2CH!n{(Iyx*#Cu&xNg=54T?JyIR48Wn~p{J zf!-Nm2Aisvi>4OLMN{)~ksk4L(bU4;A%5Ej$_3rT^x^{tY*WY017*R$(^NXE6Zbqb z-`ERRUfsz->(WiD{$NWSjIhDX%dt&{6h5#9I}w~misfBc-%ML+cZRMl8`@6}M{M}c z$!vnwv)`m5{X0)3(k8|=NG+4)SBSa5@A*@*K?C@ij3c&n&uFoXX>6%VhdXxcT55i! zcmQjf9&L52sE800pU-w#OH@%yR8jcn&pJM%))Y`n;{}eVHk#6vHZgJy>G0ue6KRWt zowrEP`Hz0y9wjn7d?pbwAbUEEzt-tvSzVOPB33-z;$@SM+*3^`zcm2c&!jSj6k>+N zE3LZ)_Mhp>SW2q5uUNqXv3004Tq;*4s%Kd6#+gG)q^9Mk=FHiYZiT~Xg;Uk4OI3@f z>R7u7Mm3#GUN&M@O=em=Gu6|<#Iu+n=g7A!!@Tj?Dw5NkI5)Rd*!qh@Obvt@v)_Zj z)`F>mvv8`|8m0=s@pGM>s*f9ACc{%V<5(KsJ(t4YJLe8A(eZ5J(tbkK!=TuNyRc<) zxOZ|Gy*88%ckDWCz{c}qakDOwY!h3?({SND@10MsGiy$@{Jv@1p}uL`*f)1$-b_zt zy6hzq$957Y^I8XzY5&c{DxFub-pR+&J?v?G;QGe$4LF0lsc(u7=q{;YxJ`YAQ#HhJ zsv^THDJ@KJss#w%`3DhWw^~yK`{6V3DAkecDgI}jel|x{pY4={*iMPd_GFIg;jsC1 zrknacMv_1+HHV>smE1Ef*B-98;ftsP-M6KH1L@m*qxGJFajgH!` zDi|fUsFOaoWv-GiSGNu9#?nYB(xW`q$&1wQBYW|Eeo-U!eKXniQj=`LRoF zQOe*;m$C`9s*^ZV9V<+Qy_XKewK{e(taw7JmE_rRwFc-5llm7)}$eW4d;v9 zCnEy)s_xcVeSc#XB!Bmln8ihAKqoV}cI*>Th55oVv|ei2RHR@=srG$?Hrswaw$ZIN(IZrsj@CAn^gbb@>pyb2U~UOrmNW zy-)rWKSWSHVpG(3Ci!QD0?b#>&vaoW=A@9Q>H&q7YB|0a-|q?Jw0aBV&j zHFSv0b*G6eNli`7@^WdJv}~};*Wu7Jr$kH`oKu>yHQ<^?c{yC0O%PjHXv90Dd9A2w zXB5)T>2eegN(R!=VDRPJ=>}B9aCa`A(!{bxvCKiiQLB)fPv*G#3Ey%Cp0nI~4#gqGN$7}9?7>5~6vX1tt+nXOg;6Jx=tNqA??@t{THN)LoN%Z>j zi|G^sII8vcw-%hxjyDb@;2$93J}D!T$q8+P#wi delta 5066 zcmZ`->vI&x5ob@J$Ldax)xESQPH+KYV*|T8`*h+5We8Ocwv6SJQ>N_oAss@`lTIgv zK%5k*3fPHpz-`{fJZy{!juQ-l#3>tnNUHKBMf}M^2ni~cRE1QE_#a56XLj#)_fGzz zo0*;I?&iMsEO1!3bn)eZT%DDk9>CHQbb(B@8kH|{PykRx#DGI$*$Lu zXE;9GSI7A=;1Dt_@2wGXZz zSO?D@7>#PxW@>1r2!FfS1EE80YlzA>P1@zFC`1!@oDcPnqI6x1u#qcXhl2+PIo+|| zSipLNS|9lh$9wj{=_~P+;|)_3p3C)4^yGX#G2zBrdqsm@Xmh>~Ra#=04qS~E5{=!cxb?yHjO`>KTA7oIsBju2|1YnK46N1~hxKRTSmzlRR*;&9)Q zn5YW%lXh1D)oEfQ{<>PUObEEgUxZ^v8pE33bC`jGss}Sr)uHLU4c|T*Z#BfG?j70D zQtGG%^c?+p)U3X}3EZ(aGNurjv^N!U*9gEPn`zYqC$nk+qdC_HIDpCA1$o>FZ|@Xus1!oIj-RCwt_ z0K808tHz54doE2hA7_Xx8w6jS6x(#Mr8|ILO^5Wmg{V;#0?l9!NuFxmpz#=n(lJHT zE8N4J{#9W}DjvX^}8P=Q`!knf>Okbdg>8yys z+Oy3>t|(o)LaPgf8VV+2gpFLWB(s=KSv0AkYK5UH+&??4s6xyuS6$oPi0GP_EY+~m zqAON|H{W}!gRrWWMIR|Tt!i}SZqVV!pZ3GK_mhgo2a<%Rh0r}hwRZPpiwK^LL%HGH zNRi&afcx*aa|Se=+qVGlbPBv8BtJ+hMlj5J5R7^i7kCPojBJvWH@<#pC$X^jyNmd-%$LJ!C4S8hJuiJ zFCu#m!KWSVw|RGE{2LlM}<;6T^4| zO>B0@u0;6PwJhs(XD+UTfA6%zxAo zv#M*qK@n95X*{%Ch;|ZM)w(CP*@}s6i4GhiO(|1@Gg#jF;-j>*yIU$yVOSn2Omwsl z6{hS%#g5Z4!Bm2@wgIHNj3k@pG+(=pQ)c$SokIOBO8Wm@md8OTW*4G!U3iMg5FE(HD^EX^O zeWfENm)H+YXGf=nyyVKyqDqx_EQMX<*#kEG{ddA@Ip`pZ@juolE!q#u4AUaed39iw zCDBxKzthBIRMArXRMAoyRSaKiU4SB58pzjPmnxW{rTOt_=@odibhvkIY=P!y8Ib?| ziB(2LW3rGAdZgBMhzZgk3T=7?J}i?`spgNG>EJqK0j)V=&C5|U8g0}1Dh7drZ?s3?(4s_(HA9AAeDkX^9$ z;|RdB6^cw@Nijc|8>vu%Z?Jrv$W;$8T%xUn>@UCYg7RecFryru$oMcTza4N!-wxk8 z;uFedeFZd5 zt!-Fak_q>wRzd$%p^5N*>6hUZR=cjmURC(XtraQ9w{@-q+0hdHgwJoSz*r8>bc86YoK0F- z=e{NoNd4%iQXoBk~K)^&_N=``L)@VNaa7CUD~deugpB@5ij}`Qn&eURyYai z(*lWylxgd9Dun!~>vp?g295FggEEiG8@T>JPIKz^iWG8o+^~Q@keSTgUcqV1AIKz@ z&n4@S%28=BHWLpa4|wCtWIJ~1l70R}#eV$yTp9y*V>+F{Uh3`fx+dJ76Y=ei;1jX! z(H%5<0)D}Mw|0Ih24(YnnloW=ekDG#4$dcPXrs9}pThgmtaftK3nn&)z_n6UB_4CW z9lMkljsRM)Iy-0q`LAvG&n7R0eaNA87^OUp1-EQ`FnCxj%U>9(+;ZpH{8R2KblzDR zI#2`Ox%1Bs|NqD`k1Wgh=kUPY2c9pC^x^vM4F1~Vxm+*Z#C{)M#@3w5Nmt#GyY4Hi N^rL_JD!|>n{|8l9IeGv9 diff --git a/tests/fixtures/onnx_genai_workflows/decoder/policies/token_state_update.onnx b/tests/fixtures/onnx_genai_workflows/decoder/policies/token_state_update.onnx index 268d221d754c99df43c4bc05049adc3ae7c5e2bb..cd2f0e0103b06a711fdd3203a42a0923d036cab2 100644 GIT binary patch delta 118 zcmZ1|u#1bAgIkC#H$N$}wAgAn*GAr1jFOWWxp>RsL-UFYOH)&;Qsa%d*h&jh5=&Ai zCo;(}8c(iciktkOIh)fQE^ftQ%$k^7l36xcfK_KQCu`Q^*R07bO^lBjCl|7*FzscW QJe{p#vMRgDU`#L?p9Vk*g(U|+z<1hmMi zAE+6kwwMb|vk*^SYDEduq+%&9$CMO^NoGpSK=wra@X6Lp+5E;Z{|aexsdKPTR%BKc zj52|niV&TAk7>SGk|h_4v2bZMacc|*zzv*S%bcv5C5sjaLOfht9PC1DTnt=H9E?qj zdzrWtCkwDjPX5o#$uym5G8apQ5Y%rfnYq*e!f5N*C}sg^*u*%mbEoZdF6CM>ZDPp+ zzx6c<+tAW{A#KWqB(-n=y>vdMM(d>&nMxfM3AdKbl+1W0W-C&EcntgC{h&$>W`#+m zEeX=znqkhc(!!;OQKvcnW-Y`Y)w{o{oRb0gxJi0&O1#u0O04Lf2#vTCvnYUVDslnC zg7*KWp^nj0cdruwTo1DL2UHUr8mMJ5ijBt?f*xCZYdrq3#TN_1{|sF7zGv?^ z@PZPl8Esh7TDgfa!-C66yXIdDng-d(F+nd0mC9g*8D&D;o=zEsaM)~+1K1;-R40;; z!w8iX-JqZes>sB2lnVO{2h8yJ_G~uRc?-{Hn?YRV(j1VMs^8{0>jsvpAGZgbqsNE! zpXF^r-Eu8MsR(*e$?u>O*asUfa}w;sZgGpVSxT+KV!D)TWioprr2O_i33|}^HCX!|2%kCuDgPlMBz~Rtp?{D*}3Z=EpZ+d!DCjbBd literal 0 HcmV?d00001 diff --git a/tests/fixtures/onnx_genai_workflows/vlm/policies/iteration_broadcast.onnx b/tests/fixtures/onnx_genai_workflows/vlm/policies/iteration_broadcast.onnx deleted file mode 100644 index 6618ad65560e12af5e2f95f56cf1d059c7b64cda..0000000000000000000000000000000000000000 GIT binary patch literal 0 HcmV?d00001 literal 611 zcmcIh!AiqG6l_A8rcVp9B5Ds(OHuIedT1XhBnu z2{VdNhH8Tz9H^&pD213As=2ZuO^Go$;Y<@vlR0wl9|j$0e`75*u%S{hvt+=#1NDrg zQwA#o$kX!Kv|G0`MMR7XAN{Uh!hbS?cDS%?%rwq#<48*qQ)29JyRwG>f}nvsa06#} z@s-*soCHnOgD&bQnQ}T0xy=KxT#A)Q7!ga8=p2f_&~m5ox1Gn`^;p&{ESzIFENcwM zOjIYs(gUVaSe^HZ+?Usda}B*;nV83DzhYI7km-!4oap&gf8iWK=SK}T8Zl~%I&VCv HZS8#mtB=JC diff --git a/tests/fixtures/onnx_genai_workflows/vlm/policies/termination.onnx b/tests/fixtures/onnx_genai_workflows/vlm/policies/termination.onnx index 46ce9432722602b317df0a22c323084f5f063c85..3af2ca2e7abf2118f6f33cbf12082031af251140 100644 GIT binary patch delta 515 zcmY+Aze~eV5XVg{O>%AegFHNMT&!)e}Rj$;Oyj}i;E^!t4{as?t7oR_kGrTIa!R5@Sd|(Z?|gQ+@*XabvX#T z&MvdBZyx}=(JaxcDy%iB?O2TJRgGV3;zD#Y(eZ`mwp;fY2p;u|T6BN@7(|aQ*cLVN z#DA?iE;V#D62z~0L)S+ABmm5u{Q`zk60zBVO9xzBl^b(M@V7)lmSeNOz`_uSEAfeV zQH?H|t)Fd0Y{%QPMH2JqIK~?d_>A)9nPpME@=x%V)`Zlzx6{h3HWugxjuH{VqC|KYX{15&?+U(b6!Ait pS2H)-)Ag&qX0coLyZx7iUk;;tl{Z7*Jm ziXW%)L=Zf9(u&x5D9I z<+HnIyw+bKTutITJYr$6|Nav`w~MfGug}vk9u8O#N1UiByA~(RUA8r8qd4Hgfse*U z?IFU7-%Y}l;~WJG`Zhe#E)eS4c)g#4sq2*1`+L7`o_>TR?s60`x`-<&6A0?c23d&0 zn$dFw=CFWq2_RM1P8q~67+#r|U)(IUQ+uzV0(lUlMN=|%XW zY{0yEe8#i&RevxVo_erZ3E_jOH$7se>;?~pU4GTyVYF@6CSFd`hM(qx*^X_k#?c`0 zqyHow_=C3L?~G}+v5~TImUiHTWus-dUfqNj))j4Gav}I;z0rjqCcN(_VI0AC>nyyj zb&XrC!q&hf=MOY;yVR&>dkF&Owuy=rwL(gG3Le)THOq*5!hXH!%@&R!P*sWQ<+buw zG^;H{z9jn=ey WM0nnKcIwa*IxAH!}q$Y!emfd+)xFdqIqdq$UNob~RdqOsBge-DP+7 zYjz7n6KcYTRoIqeD^%Np0u`Z6D`K=1T$GwdQ;jk4kB0q#EofpI%Rf#0r}5nT-n@^Q z<)4{*-@EsobARV|&iTE+|E=M~^@g^l7J23H$m0W(6B`a(f1^!V(Q-H@4D~CpQ4OAwij;Vsb zNGBS%^u>WGw4^JoeBfHrVgC=lg}%n(q|Q}B&{YN2{ieVdz@ML44|^{3!T+9(uOm|X zRKdHMP7GPX`Ov^98ZebMF>)=Kka=}5N=)IqwS?~$_Z=9t>bfNLakL6re?gV;3Xhcf zC;LhcLbf9NQ}=8eo`i#|EPlCdBfHQxRa*JLeNFQXIO8xvfbEOMOa}V|ZR^Vxnh!>)2#*sGu8~5)VFb4mC9R z_j7yUD8|mv>j5y}_tX6kUJ!A1L&P0*y_4fpg$Kt=MdX!lhA7FJXRajzn|{16qQjcG zM25&sQ-z*!^vY0L@ZVixXrtQ(>$%y1F;09gW@cmaZvSBLiUw@)BUL}}e!iPr+ z@c!|xn5DG$K2aR?-dNCb;_0|u50(vMqbc}nWC!G4Nabz0&ASM?M3qkW235Z(NUs#( zw=Z;RR9aCP87=IVTBZt)Dc!JSoH!oo2CH!n{(Iyx*#Cu&xNg=54T?JyIR48Wn~p{J zf!-Nm2Aisvi>4OLMN{)~ksk4L(bU4;A%5Ej$_3rT^x^{tY*WY017*R$(^NXE6Zbqb z-`ERRUfsz->(WiD{$NWSjIhDX%dt&{6h5#9I}w~misfBc-%ML+cZRMl8`@6}M{M}c z$!vnwv)`m5{X0)3(k8|=NG+4)SBSa5@A*@*K?C@ij3c&n&uFoXX>6%VhdXxcT55i! zcmQjf9&L52sE800pU-w#OH@%yR8jcn&pJM%))Y`n;{}eVHk#6vHZgJy>G0ue6KRWt zowrEP`Hz0y9wjn7d?pbwAbUEEzt-tvSzVOPB33-z;$@SM+*3^`zcm2c&!jSj6k>+N zE3LZ)_Mhp>SW2q5uUNqXv3004Tq;*4s%Kd6#+gG)q^9Mk=FHiYZiT~Xg;Uk4OI3@f z>R7u7Mm3#GUN&M@O=em=Gu6|<#Iu+n=g7A!!@Tj?Dw5NkI5)Rd*!qh@Obvt@v)_Zj z)`F>mvv8`|8m0=s@pGM>s*f9ACc{%V<5(KsJ(t4YJLe8A(eZ5J(tbkK!=TuNyRc<) zxOZ|Gy*88%ckDWCz{c}qakDOwY!h3?({SND@10MsGiy$@{Jv@1p}uL`*f)1$-b_zt zy6hzq$957Y^I8XzY5&c{DxFub-pR+&J?v?G;QGe$4LF0lsc(u7=q{;YxJ`YAQ#HhJ zsv^THDJ@KJss#w%`3DhWw^~yK`{6V3DAkecDgI}jel|x{pY4={*iMPd_GFIg;jsC1 zrknacMv_1+HHV>smE1Ef*B-98;ftsP-M6KH1L@m*qxGJFajgH!` zDi|fUsFOaoWv-GiSGNu9#?nYB(xW`q$&1wQBYW|Eeo-U!eKXniQj=`LRoF zQOe*;m$C`9s*^ZV9V<+Qy_XKewK{e(taw7JmE_rRwFc-5llm7)}$eW4d;v9 zCnEy)s_xcVeSc#XB!Bmln8ihAKqoV}cI*>Th55oVv|ei2RHR@=srG$?Hrswaw$ZIN(IZrsj@CAn^gbb@>pyb2U~UOrmNW zy-)rWKSWSHVpG(3Ci!QD0?b#>&vaoW=A@9Q>H&q7YB|0a-|q?Jw0aBV&j zHFSv0b*G6eNli`7@^WdJv}~};*Wu7Jr$kH`oKu>yHQ<^?c{yC0O%PjHXv90Dd9A2w zXB5)T>2eegN(R!=VDRPJ=>}B9aCa`A(!{bxvCKiiQLB)fPv*G#3Ey%Cp0nI~4#gqGN$7}9?7>5~6vX1tt+nXOg;6Jx=tNqA??@t{THN)LoN%Z>j zi|G^sII8vcw-%hxjyDb@;2$93J}D!T$q8+P#wi delta 5066 zcmZ`->vI&x5ob@J$Ldax)xESQPH+KYV*|T8`*h+5We8Ocwv6SJQ>N_oAss@`lTIgv zK%5k*3fPHpz-`{fJZy{!juQ-l#3>tnNUHKBMf}M^2ni~cRE1QE_#a56XLj#)_fGzz zo0*;I?&iMsEO1!3bn)eZT%DDk9>CHQbb(B@8kH|{PykRx#DGI$*$Lu zXE;9GSI7A=;1Dt_@2wGXZz zSO?D@7>#PxW@>1r2!FfS1EE80YlzA>P1@zFC`1!@oDcPnqI6x1u#qcXhl2+PIo+|| zSipLNS|9lh$9wj{=_~P+;|)_3p3C)4^yGX#G2zBrdqsm@Xmh>~Ra#=04qS~E5{=!cxb?yHjO`>KTA7oIsBju2|1YnK46N1~hxKRTSmzlRR*;&9)Q zn5YW%lXh1D)oEfQ{<>PUObEEgUxZ^v8pE33bC`jGss}Sr)uHLU4c|T*Z#BfG?j70D zQtGG%^c?+p)U3X}3EZ(aGNurjv^N!U*9gEPn`zYqC$nk+qdC_HIDpCA1$o>FZ|@Xus1!oIj-RCwt_ z0K808tHz54doE2hA7_Xx8w6jS6x(#Mr8|ILO^5Wmg{V;#0?l9!NuFxmpz#=n(lJHT zE8N4J{#9W}DjvX^}8P=Q`!knf>Okbdg>8yys z+Oy3>t|(o)LaPgf8VV+2gpFLWB(s=KSv0AkYK5UH+&??4s6xyuS6$oPi0GP_EY+~m zqAON|H{W}!gRrWWMIR|Tt!i}SZqVV!pZ3GK_mhgo2a<%Rh0r}hwRZPpiwK^LL%HGH zNRi&afcx*aa|Se=+qVGlbPBv8BtJ+hMlj5J5R7^i7kCPojBJvWH@<#pC$X^jyNmd-%$LJ!C4S8hJuiJ zFCu#m!KWSVw|RGE{2LlM}<;6T^4| zO>B0@u0;6PwJhs(XD+UTfA6%zxAo zv#M*qK@n95X*{%Ch;|ZM)w(CP*@}s6i4GhiO(|1@Gg#jF;-j>*yIU$yVOSn2Omwsl z6{hS%#g5Z4!Bm2@wgIHNj3k@pG+(=pQ)c$SokIOBO8Wm@md8OTW*4G!U3iMg5FE(HD^EX^O zeWfENm)H+YXGf=nyyVKyqDqx_EQMX<*#kEG{ddA@Ip`pZ@juolE!q#u4AUaed39iw zCDBxKzthBIRMArXRMAoyRSaKiU4SB58pzjPmnxW{rTOt_=@odibhvkIY=P!y8Ib?| ziB(2LW3rGAdZgBMhzZgk3T=7?J}i?`spgNG>EJqK0j)V=&C5|U8g0}1Dh7drZ?s3?(4s_(HA9AAeDkX^9$ z;|RdB6^cw@Nijc|8>vu%Z?Jrv$W;$8T%xUn>@UCYg7RecFryru$oMcTza4N!-wxk8 z;uFedeFZd5 zt!-Fak_q>wRzd$%p^5N*>6hUZR=cjmURC(XtraQ9w{@-q+0hdHgwJoSz*r8>bc86YoK0F- z=e{NoNd4%iQXoBk~K)^&_N=``L)@VNaa7CUD~deugpB@5ij}`Qn&eURyYai z(*lWylxgd9Dun!~>vp?g295FggEEiG8@T>JPIKz^iWG8o+^~Q@keSTgUcqV1AIKz@ z&n4@S%28=BHWLpa4|wCtWIJ~1l70R}#eV$yTp9y*V>+F{Uh3`fx+dJ76Y=ei;1jX! z(H%5<0)D}Mw|0Ih24(YnnloW=ekDG#4$dcPXrs9}pThgmtaftK3nn&)z_n6UB_4CW z9lMkljsRM)Iy-0q`LAvG&n7R0eaNA87^OUp1-EQ`FnCxj%U>9(+;ZpH{8R2KblzDR zI#2`Ox%1Bs|NqD`k1Wgh=kUPY2c9pC^x^vM4F1~Vxm+*Z#C{)M#@3w5Nmt#GyY4Hi N^rL_JD!|>n{|8l9IeGv9 diff --git a/tests/fixtures/onnx_genai_workflows/vlm/policies/token_state_update.onnx b/tests/fixtures/onnx_genai_workflows/vlm/policies/token_state_update.onnx index 268d221d754c99df43c4bc05049adc3ae7c5e2bb..cd2f0e0103b06a711fdd3203a42a0923d036cab2 100644 GIT binary patch delta 118 zcmZ1|u#1bAgIkC#H$N$}wAgAn*GAr1jFOWWxp>RsL-UFYOH)&;Qsa%d*h&jh5=&Ai zCo;(}8c(iciktkOIh)fQE^ftQ%$k^7l36xcfK_KQCu`Q^*R07bO^lBjCl|7*FzscW QJe{p#vMRgDU`#L?p9Vk*g(U|+z<1hmMi zAE+6kwwMb|vk*^SYDEduq+%&9$CMO^NoGpSK=wra@X6Lp+5E;Z{|aexsdKPTR%BKc zj52|niV&TAk7>SGk|h_4v2bZMacc|*zzv*S%bcv5C5sjaLOfht9PC1DTnt=H9E?qj zdzrWtCkwDjPX5o#$uym5G8apQ5Y%rfnYq*e!f5N*C}sg^*u*%mbEoZdF6CM>ZDPp+ zzx6c<+tAW{A#KWqB(-n=y>vdMM(d>&nMxfM3AdKbl+1W0W-C&EcntgC{h&$>W`#+m zEeX=znqkhc(!!;OQKvcnW-Y`Y)w{o{oRb0gxJi0&O1#u0O04Lf2#vTCvnYUVDslnC zg7*KWp^nj0cdruwT>(); - let grammar_mask = vec![1_u8; usize::try_from(batch * 128)?]; Ok(PipelineGenerateRequest::new(GenerateRequest { prompt: GeneratePrompt::TokenIds(vec![0]), options: options(max_new_tokens), @@ -97,11 +96,7 @@ fn decoder_batch_request( ) .with_input("request.seed", Value::from_slice_i64(&slot_ids, &[batch])?) .with_input( - "request.grammar_mask", - Value::from_raw_bytes(grammar_mask, &[batch, 128], DataType::Bool)?, - ) - .with_input( - "request.rng_offset", + "request.rng_counter", Value::from_slice_i64(&zeros, &[batch])?, ) .with_input( From 96580bb1d076cb964a065b41150ce6c79f92fd79 Mon Sep 17 00:00:00 2001 From: justinchuby Date: Fri, 14 Aug 2026 08:38:21 +0000 Subject: [PATCH 087/151] Lock exact batched policy bindings Assert the authoritative sampler, termination, and selective state-update v2 semantic role sets so extra or missing bindings fail producer tests. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: justinchuby --- .../onnx_genai/auto_export_test.py | 27 ++++++++++++++++++- 1 file changed, 26 insertions(+), 1 deletion(-) diff --git a/src/mobius/integrations/onnx_genai/auto_export_test.py b/src/mobius/integrations/onnx_genai/auto_export_test.py index a0eae019b..408cd68f3 100644 --- a/src/mobius/integrations/onnx_genai/auto_export_test.py +++ b/src/mobius/integrations/onnx_genai/auto_export_test.py @@ -278,7 +278,32 @@ def test_seeded_decoder_sampler_uses_request_controls_and_direct_kv_carry(): "batching": "per_row", "inactive_rows": "preserve", } - assert workflow["components"]["token_state_update"]["contract"]["version"] == "2" + assert set(workflow["components"]["termination"]["contract"]["bindings"]) == { + "tokens", + "active", + "eos_ids", + "eos_lengths", + "iteration", + "max_iterations", + "done", + "next_active", + "continue", + } + 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"]) From 28f8129a9ef4ab0c4783ed0226e3d9970d474457 Mon Sep 17 00:00:00 2001 From: justinchuby Date: Fri, 14 Aug 2026 08:55:23 +0000 Subject: [PATCH 088/151] Require complete diffusion workflow packages Update the newly landed flat-model regression to enforce the generic workflow contract: diffusion exports require separate denoiser and VAE decoder components rather than reviving legacy pipeline metadata. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: justinchuby --- .../integrations/onnx_genai/auto_export_test.py | 11 ++++++----- 1 file changed, 6 insertions(+), 5 deletions(-) diff --git a/src/mobius/integrations/onnx_genai/auto_export_test.py b/src/mobius/integrations/onnx_genai/auto_export_test.py index 408cd68f3..b7354b6c8 100644 --- a/src/mobius/integrations/onnx_genai/auto_export_test.py +++ b/src/mobius/integrations/onnx_genai/auto_export_test.py @@ -352,12 +352,13 @@ def test_dispatch_diffusion(tmp_path): 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 test_rejects_unsupported_qwen_image_edit_runtime_export(tmp_path): From faf86865ca4fea70d32a3a87020b7ea1cb71af66 Mon Sep 17 00:00:00 2001 From: justinchuby Date: Fri, 14 Aug 2026 09:12:01 +0000 Subject: [PATCH 089/151] Assert batched policy super-island admission Publish the VLM loop induction value as the singleton required by the termination v2 ABI, regenerate its fixture, and make cross-repository conformance prove sampler, termination, and state update execute in one fused island without fallback. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: justinchuby --- src/mobius/_model_package.py | 9 +- src/mobius/_model_package_test.py | 6 +- src/mobius/generation/_policy_components.py | 29 +- .../generation/_policy_components_test.py | 46 ++ .../onnx_genai/auto_export_test.py | 21 +- .../onnx_genai/inference_metadata.py | 23 + .../onnx_genai/inference_metadata_test.py | 17 +- .../onnx_genai/workflow_metadata.py | 5 +- .../onnx_genai/workflow_metadata_test.py | 9 +- .../decoder/inference_metadata.yaml | 369 +++++++++++++++ .../decoder/policies/cache_length_update.onnx | Bin 1044 -> 930 bytes .../policies/decoder_state_initializer.onnx | Bin 9732 -> 9244 bytes .../decoder/policies/decoder_step_update.onnx | Bin 3203 -> 3028 bytes .../policies/generated_length_update.onnx | Bin 1064 -> 930 bytes .../decoder/policies/last_token_logits.onnx | Bin 756 -> 631 bytes .../decoder/policies/termination.onnx | Bin 5932 -> 5880 bytes .../termination_batch_initializer.onnx | Bin 2003 -> 1677 bytes .../decoder/policies/token_sampler.onnx | Bin 57881 -> 58720 bytes .../decoder/policies/token_state_update.onnx | Bin 1338 -> 1229 bytes .../decoder/policies/token_to_slot.onnx | Bin 480 -> 438 bytes .../diffusion/inference_metadata.yaml | 113 +++++ .../policies/continue_predicate.onnx | Bin 929 -> 910 bytes .../diffusion/policies/euler_model_input.onnx | Bin 2142 -> 1926 bytes .../diffusion/policies/schedule_lookup.onnx | Bin 669 -> 607 bytes .../diffusion/policies/solver_step.onnx | Bin 3248 -> 3060 bytes .../masked/inference_metadata.yaml | 75 +++ .../masked/policies/masked_update.onnx | Bin 13264 -> 13182 bytes .../speculative/inference_metadata.yaml | 220 +++++++++ .../speculative/policies/adaptive_k.onnx | Bin 33930 -> 34702 bytes .../policies/cache_length_update.onnx | Bin 467 -> 386 bytes .../policies/grammar_guidance.onnx | Bin 2054 -> 1926 bytes .../speculative/policies/grammar_length.onnx | Bin 460 -> 394 bytes .../policies/grammar_sampler_logits.onnx | Bin 784 -> 631 bytes .../policies/proposal_metrics.onnx | Bin 1146 -> 1026 bytes .../policies/speculative_acceptance.onnx | Bin 8835 -> 8736 bytes .../tts/code_predictor/model.onnx | Bin 108810 -> 106338 bytes .../tts/embedding/model.onnx | Bin 5622 -> 5313 bytes .../tts/inference_metadata.yaml | 445 ++++++++++++++++++ .../tts/policies/cache_length_update.onnx | Bin 467 -> 386 bytes .../tts/policies/code_frame_update.onnx | Bin 1648 -> 1580 bytes .../tts/policies/code_history_append.onnx | Bin 981 -> 853 bytes .../tts/policies/codec_layout.onnx | Bin 508 -> 428 bytes .../tts/policies/continue_predicate.onnx | Bin 929 -> 910 bytes .../tts/policies/last_token_logits.onnx | Bin 756 -> 631 bytes .../tts/policies/predictor_body_sampler.onnx | Bin 626 -> 536 bytes .../policies/predictor_prefill_sampler.onnx | Bin 638 -> 539 bytes .../policies/predictor_state_initializer.onnx | Bin 15717 -> 14690 bytes .../tts/policies/predictor_step_update.onnx | Bin 2213 -> 2039 bytes .../tts/policies/setup_predictor_sampler.onnx | Bin 630 -> 537 bytes .../tts/policies/setup_talker_sampler.onnx | Bin 618 -> 534 bytes .../tts/policies/talker_sampler.onnx | Bin 594 -> 528 bytes .../policies/talker_state_initializer.onnx | Bin 7875 -> 7457 bytes .../tts/policies/talker_step_update.onnx | Bin 2200 -> 2044 bytes .../tts/policies/token_to_slot.onnx | Bin 480 -> 438 bytes .../tts/policies/tts_state_initializer.onnx | Bin 2339 -> 2194 bytes .../tts/talker/model.onnx | Bin 32118 -> 31685 bytes .../tts/talker_prefill_embedder/model.onnx | Bin 22703 -> 21848 bytes .../vlm/inference_metadata.yaml | 379 ++++++++++++++- .../vlm/policies/cache_length_update.onnx | Bin 1044 -> 930 bytes .../policies/decoder_state_initializer.onnx | Bin 11015 -> 10460 bytes .../vlm/policies/decoder_step_update.onnx | Bin 3203 -> 3028 bytes .../vlm/policies/generated_length_update.onnx | Bin 1064 -> 930 bytes .../vlm/policies/last_token_logits.onnx | Bin 756 -> 631 bytes .../vlm/policies/termination.onnx | Bin 5932 -> 5880 bytes .../termination_batch_initializer.onnx | Bin 2003 -> 1677 bytes .../vlm/policies/token_sampler.onnx | Bin 57881 -> 58720 bytes .../vlm/policies/token_state_update.onnx | Bin 1338 -> 1229 bytes .../vlm/policies/token_to_slot.onnx | Bin 480 -> 438 bytes tests/onnx_genai_workflow_conformance.rs | 35 +- 69 files changed, 1784 insertions(+), 17 deletions(-) diff --git a/src/mobius/_model_package.py b/src/mobius/_model_package.py index 8fa6df5e1..89166c784 100644 --- a/src/mobius/_model_package.py +++ b/src/mobius/_model_package.py @@ -211,11 +211,10 @@ def save_policy_components( if check_weights: _check_weights(name, component.model) relative_path = f"policies/{name}.onnx" - with _namespaced_symbolic_dimensions( - component.model, - f"policy.{name}", - ) as saved_model: - ir.save(saved_model, os.path.join(directory, relative_path)) + # 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 diff --git a/src/mobius/_model_package_test.py b/src/mobius/_model_package_test.py index d9e076f03..ada31473d 100644 --- a/src/mobius/_model_package_test.py +++ b/src/mobius/_model_package_test.py @@ -477,7 +477,7 @@ def test_save_anonymizes_nested_graph_intermediate_symbols(self): assert str(branch_output.shape[1]) == "branch_sequence" - def test_save_namespaces_policy_symbols(self, tmp_path): + 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) @@ -485,8 +485,8 @@ def test_save_namespaces_policy_symbols(self, tmp_path): pkg.save(str(tmp_path)) saved = ir.load(tmp_path / "policies" / "sample.onnx") - assert str(saved.graph.inputs[0].shape[0]) == "policy.sample.batch" - assert str(saved.graph.inputs[0].shape[1]) == "policy.sample.vocabulary" + 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" diff --git a/src/mobius/generation/_policy_components.py b/src/mobius/generation/_policy_components.py index 9c1c317c4..90587a639 100644 --- a/src/mobius/generation/_policy_components.py +++ b/src/mobius/generation/_policy_components.py @@ -135,6 +135,11 @@ def _make_graph(name: str) -> tuple[ir.Graph, GraphBuilder]: 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", @@ -1332,6 +1337,11 @@ def build_seeded_categorical_sampler() -> PolicyComponent: 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( @@ -1419,7 +1429,17 @@ def build_eos_termination(*, row_selective: bool = False) -> PolicyComponent: op.ReduceMax(op.Cast(next_active, to=ir.DataType.INT64), keepdims=1), op.Constant(value_int=0), ) - continued.shape = ir.Shape([1]) + 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") @@ -1804,7 +1824,12 @@ def build_token_state_update(*, row_selective: bool = False) -> PolicyComponent: ) else: next_state = op.Unsqueeze(update, [-1]) - next_state.shape = ir.Shape(["batch", 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", diff --git a/src/mobius/generation/_policy_components_test.py b/src/mobius/generation/_policy_components_test.py index 159c54f11..363a21e57 100644 --- a/src/mobius/generation/_policy_components_test.py +++ b/src/mobius/generation/_policy_components_test.py @@ -512,6 +512,52 @@ def test_seeded_sampler_is_counter_based_and_reproducible(tmp_path): 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(), diff --git a/src/mobius/integrations/onnx_genai/auto_export_test.py b/src/mobius/integrations/onnx_genai/auto_export_test.py index b7354b6c8..1c2bef966 100644 --- a/src/mobius/integrations/onnx_genai/auto_export_test.py +++ b/src/mobius/integrations/onnx_genai/auto_export_test.py @@ -241,6 +241,16 @@ def test_seeded_decoder_sampler_uses_request_controls_and_direct_kv_carry(): "active": "active", "done": "done", } + assert sampler["ports"]["inputs"]["logits"] == { + "dtype": "float32", + "rank": 2, + "shape": ["batch", "vocabulary"], + } + assert sampler["ports"]["outputs"]["token"] == { + "dtype": "int64", + "rank": 1, + "shape": ["batch"], + } sampler_step = next( step for step in workflow["steps"][0]["steps"] @@ -289,6 +299,11 @@ def test_seeded_decoder_sampler_uses_request_controls_and_direct_kv_carry(): "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", @@ -487,7 +502,11 @@ def test_dispatch_diffusion_auto_reads_scheduler_from_source(tmp_path): with open(arts["inference_metadata"]) as handle: meta = yaml.safe_load(handle) components = meta["pipeline"]["workflow"]["components"] - assert "ports" not in components["diffusion_schedule"] + 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] diff --git a/src/mobius/integrations/onnx_genai/inference_metadata.py b/src/mobius/integrations/onnx_genai/inference_metadata.py index dc251637c..995b82ccc 100644 --- a/src/mobius/integrations/onnx_genai/inference_metadata.py +++ b/src/mobius/integrations/onnx_genai/inference_metadata.py @@ -1444,12 +1444,35 @@ def semantic_contract(component: Any) -> dict[str, Any]: 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) + return { + "dtype": dtype, + "rank": port.rank, + "shape": _shape_metadata(port), + } + for name, component in policy_components.items(): 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) diff --git a/src/mobius/integrations/onnx_genai/inference_metadata_test.py b/src/mobius/integrations/onnx_genai/inference_metadata_test.py index e763a69a3..2c7f7312f 100644 --- a/src/mobius/integrations/onnx_genai/inference_metadata_test.py +++ b/src/mobius/integrations/onnx_genai/inference_metadata_test.py @@ -150,7 +150,22 @@ def test_workflow_policy_components_reference_saved_onnx_artifacts(tmp_path): "kind": "onnx", "artifact": "policies/sample.onnx", } - assert "ports" not in component + assert component["ports"] == { + "inputs": { + "logits": { + "dtype": "float32", + "rank": 2, + "shape": ["batch", "vocabulary"], + } + }, + "outputs": { + "token": { + "dtype": "int64", + "rank": 1, + "shape": ["batch"], + } + }, + } assert component["contract"] == { "id": "onnx-genai.token-sampler", "version": "1", diff --git a/src/mobius/integrations/onnx_genai/workflow_metadata.py b/src/mobius/integrations/onnx_genai/workflow_metadata.py index ea5ce7862..99195459e 100644 --- a/src/mobius/integrations/onnx_genai/workflow_metadata.py +++ b/src/mobius/integrations/onnx_genai/workflow_metadata.py @@ -3541,7 +3541,10 @@ def build_vlm_workflow_metadata( "condition": "loop.continue", "termination": "generation_eos", "max_iterations": "request.max_iterations", - "iteration": {"value": "loop.iteration", "contract": batch_int}, + "iteration": { + "value": "loop.iteration", + "contract": {"dtype": "int64", "rank": 1, "shape": [1]}, + }, "carried": carried, }, } diff --git a/src/mobius/integrations/onnx_genai/workflow_metadata_test.py b/src/mobius/integrations/onnx_genai/workflow_metadata_test.py index e7af9fb00..f3d2154ec 100644 --- a/src/mobius/integrations/onnx_genai/workflow_metadata_test.py +++ b/src/mobius/integrations/onnx_genai/workflow_metadata_test.py @@ -256,6 +256,11 @@ def collect_decoder_invokes(node): } 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" @@ -485,9 +490,9 @@ def test_speculative_grammar_and_adaptive_k_use_typed_state_contracts(): assert workflow["state"]["proposal_k"]["class"] == "advisory" assert workflow["state"]["adaptive_estimates"]["class"] == "advisory" assert all( - "ports" not in component and "effects" not in component + "ports" in component and "effects" not in component for component in workflow["components"].values() - if component["implementation"]["kind"] == "onnx" + if component["implementation"]["kind"] == "onnx" and "contract" in component ) assert "ports" in workflow["components"]["grammar_commit"] proposer = next( diff --git a/tests/fixtures/onnx_genai_workflows/decoder/inference_metadata.yaml b/tests/fixtures/onnx_genai_workflows/decoder/inference_metadata.yaml index 7075a0095..b715dd625 100644 --- a/tests/fixtures/onnx_genai_workflows/decoder/inference_metadata.yaml +++ b/tests/fixtures/onnx_genai_workflows/decoder/inference_metadata.yaml @@ -294,6 +294,65 @@ pipeline: implementation: kind: onnx artifact: policies/token_sampler.onnx + ports: + inputs: + logits: + dtype: float32 + rank: 2 + shape: + - batch + - vocabulary + temperature: + dtype: float32 + rank: 1 + shape: + - batch + top_k: + dtype: int64 + rank: 1 + shape: + - batch + top_p: + dtype: float32 + rank: 1 + shape: + - batch + min_p: + dtype: float32 + rank: 1 + shape: + - batch + seed: + dtype: int64 + rank: 1 + shape: + - batch + counter: + dtype: int64 + rank: 1 + shape: + - batch + active: + dtype: bool + rank: 1 + shape: + - batch + done: + dtype: bool + rank: 1 + shape: + - batch + outputs: + token: + dtype: int64 + rank: 1 + shape: + - batch + next_counter: + dtype: int64 + rank: 1 + shape: + - batch contract: id: onnx-genai.token-sampler version: '2' @@ -318,6 +377,55 @@ pipeline: implementation: kind: onnx artifact: policies/termination.onnx + ports: + inputs: + tokens: + dtype: int64 + rank: 1 + shape: + - batch + eos_ids: + dtype: int64 + rank: 2 + shape: + - batch + - num_eos + eos_lengths: + dtype: int64 + rank: 1 + shape: + - batch + iteration: + dtype: int64 + rank: 1 + shape: + - 1 + max_iterations: + dtype: int64 + rank: 1 + shape: + - batch + active: + dtype: bool + rank: 1 + shape: + - batch + outputs: + done: + dtype: bool + rank: 1 + shape: + - batch + next_active: + dtype: bool + rank: 1 + shape: + - batch + continue: + dtype: bool + rank: 1 + shape: + - 1 contract: id: onnx-genai.termination-predicate version: '2' @@ -338,6 +446,37 @@ pipeline: implementation: kind: onnx artifact: policies/token_state_update.onnx + ports: + inputs: + current: + dtype: int64 + rank: 2 + shape: + - batch + - 1 + update: + dtype: int64 + rank: 2 + shape: + - batch + - 1 + active: + dtype: bool + rank: 1 + shape: + - batch + done: + dtype: bool + rank: 1 + shape: + - batch + outputs: + next: + dtype: int64 + rank: 2 + shape: + - batch + - 1 contract: id: onnx-genai.state-update version: '2' @@ -354,30 +493,260 @@ pipeline: implementation: kind: onnx artifact: policies/last_token_logits.onnx + ports: + inputs: + logits: + dtype: float32 + rank: 3 + shape: + - batch + - sequence + - vocabulary + outputs: + last_logits: + dtype: float32 + rank: 2 + shape: + - batch + - vocabulary decoder_state_initializer: implementation: kind: onnx artifact: policies/decoder_state_initializer.onnx + ports: + inputs: + prompt_tokens: + dtype: int64 + rank: 2 + shape: + - batch + - prompt_sequence + prompt_lengths: + dtype: int64 + rank: 1 + shape: + - batch + max_iterations: + dtype: int64 + rank: 1 + shape: + - 1 + outputs: + attention_mask: + dtype: int64 + rank: 2 + shape: + - batch + - capacity + position_ids: + dtype: int64 + rank: 2 + shape: + - batch + - sequence + body_attention_mask: + dtype: int64 + rank: 2 + shape: + - batch + - capacity + body_position_ids: + dtype: int64 + rank: 2 + shape: + - batch + - 1 + token_slot: + dtype: int64 + rank: 2 + shape: + - batch + - 1 + generated_lengths: + dtype: int64 + rank: 1 + shape: + - batch + cache_lengths: + dtype: int64 + rank: 1 + shape: + - batch + past_key_values.0.key: + dtype: float32 + rank: 4 + shape: + - batch + - 2 + - capacity + - 8 decoder_step_update: implementation: kind: onnx artifact: policies/decoder_step_update.onnx + ports: + inputs: + attention_mask: + dtype: int64 + rank: 2 + shape: + - batch + - context + logical_length: + dtype: int64 + rank: 1 + shape: + - batch + position_ids: + dtype: int64 + rank: 2 + shape: + - batch + - 1 + outputs: + next_attention_mask: + dtype: int64 + rank: 2 + shape: + - batch + - context + next_position_ids: + dtype: int64 + rank: 2 + shape: + - batch + - 1 cache_length_update: implementation: kind: onnx artifact: policies/cache_length_update.onnx + ports: + inputs: + left: + dtype: int64 + rank: 1 + shape: + - batch + right: + dtype: int64 + rank: 1 + shape: + - batch + active: + dtype: bool + rank: 1 + shape: + - batch + done: + dtype: bool + rank: 1 + shape: + - batch + outputs: + total: + dtype: int64 + rank: 1 + shape: + - batch 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 + input_eos_lengths: + dtype: int64 + rank: 1 + shape: + - batch + input_max_iterations: + dtype: int64 + rank: 1 + shape: + - batch + fallback_max_iterations: + dtype: int64 + rank: 1 + shape: + - 1 + active: + dtype: bool + rank: 1 + shape: + - batch + outputs: + row_eos_ids: + dtype: int64 + rank: 2 + shape: + - batch + - num_eos + eos_lengths: + dtype: int64 + rank: 1 + shape: + - batch + max_iterations: + dtype: int64 + rank: 1 + shape: + - batch token_to_slot: implementation: kind: onnx artifact: policies/token_to_slot.onnx + ports: + inputs: + token: + dtype: int64 + rank: 1 + shape: + - batch + outputs: + slot: + dtype: int64 + rank: 2 + shape: + - batch + - 1 generated_length_update: implementation: kind: onnx artifact: policies/generated_length_update.onnx + ports: + inputs: + left: + dtype: int64 + rank: 1 + shape: + - batch + right: + dtype: int64 + rank: 1 + shape: + - batch + active: + dtype: bool + rank: 1 + shape: + - batch + done: + dtype: bool + rank: 1 + shape: + - batch + outputs: + total: + dtype: int64 + rank: 1 + shape: + - batch state: token: contract: diff --git a/tests/fixtures/onnx_genai_workflows/decoder/policies/cache_length_update.onnx b/tests/fixtures/onnx_genai_workflows/decoder/policies/cache_length_update.onnx index 980ab9a3d96fd6966551637a8856faffc0bf4bb9..af3bc3fca3b73acdb986806a33ddcf35d9032ffc 100644 GIT binary patch delta 240 zcmbQjv51|QgIkC#H$N$}wAiYjeIu_5Q@uDBOHOK9i4Z>*F9*92Cl|XAYf@rKaz>N{ z7i&>wdIlC5NiMd;SVo9oAQesJRMwBrZYf)x;25Fj1xY!btOESw+!B%lXt*R$v z7tnnv`FW_C2%u9(k U#1KqLi8lfZvV&AkW@VlW0M1xrx&QzG diff --git a/tests/fixtures/onnx_genai_workflows/decoder/policies/decoder_state_initializer.onnx b/tests/fixtures/onnx_genai_workflows/decoder/policies/decoder_state_initializer.onnx index 32de3cf6278cf495e54d6c4ad13379d9de0f7d5a..6e93156a3863f12fdf03a7fe95b4529558b3278e 100644 GIT binary patch delta 496 zcmaKpKTiTN7>9Eh<#^#BSHeIP<*Sgs57N7QGrNk%w$Lo zp{;gcel5pVoff1IK6n3=0uM+(0X3Qqd3>bk3q2<>-8* s599C9Jdo@{cjyCFh34c!6CXy*oK05b08{u@gRbtAxXnl);6_C2fcXm4R-Y* zls$yc;~Q9Jr&|%lWB&ZU^L+YrVVMP#_at*7b@+b%YF$=0Bd~%{dVwZd_JyyI4a{6X z_ppt&0_i$D#=~NDwTv0nya#R6oVxDgl&i!;Bz!JzI#5bauJF6Mr%)4cqF2x88F=_o zt<@xysMf;Qj`WE~RUh5KZ34f}z$f+%v@>PDi#ZJ_cXZ4=C5G^Qs2d4|n>cwO0o*hY$&{&lh0DS02N`%&UoRRf b%TO4QlTaj)Lvnyj$eW8O2XmvVjkoy^AYr`| diff --git a/tests/fixtures/onnx_genai_workflows/decoder/policies/decoder_step_update.onnx b/tests/fixtures/onnx_genai_workflows/decoder/policies/decoder_step_update.onnx index 0333f2f9bf0c4417d7b0e6fd4e2772628899dcb9..d61b374c0aba21bd49e09839e485c098162ccaa4 100644 GIT binary patch delta 302 zcmZpcyduua!7ap=o1c_fT5NStd?W8Oj(Qy~zQmG})Vz|+{Ji+w#NuorDJ}^Pb|Dci zb|KcJ#FFF;E>0o#60tDt~jf3@s`Dh t<`ox~rlwY<#+zZaK1+^^vn<}VuoUPGb0I-4ehy9{Zny(CH*qsE0svs0P8t9J delta 479 zcmca2-Ym(>!7ap=o1c_fT5R<}Vk7S}j`}1nzQmG})Vz|+{Ji+w#NuqB04_fcb|Eh= zB_V}^{G80>O1+fSqWI#H)PnfZf|SIPRK29clH?376(J=&YLfHwN>VFIqHMVM za`MwNlM{2|b5irtOEQGCxHN!vsBkIO6SgGEj*F)tzc>@@!pxLnAw4b~pk@uC9LU7M zn3T>XoCkD9JT`yUQ`4hK&Rl|En}~NVCp*F9*92Cl|XAYf@rKaz>N{ z7i&>wdIlC5NiMd;CMwnPsV9E9*H) zuoLLkl>9u1MiTTT0o_`XUy_)E#gj33eVHZ9#aS`b|KcJ#FFF; zE-oRC;?%;@)V$GSvi#)4q|%(kqRJ#4F7BMf;*xl%PFXH#piXgwPN3?^rcC<* DBibI7 delta 241 zcmey)@`aU`gIkC#H$N$}wAkt)>qg#S#`-!gww(O*%#vcE60RZ+Mxi_|1tGbD{G80> zO1+%K;*$81{Or`cc!(Okq{Ncs3@&9MMVu;%QwvK|^O94!RE3mrDk{rQPE0DzNi3=i dON!&-2HOL5g+G@s&=sCUxB{OAlXo!g2LSQqRMh|g diff --git a/tests/fixtures/onnx_genai_workflows/decoder/policies/termination.onnx b/tests/fixtures/onnx_genai_workflows/decoder/policies/termination.onnx index 3af2ca2e7abf2118f6f33cbf12082031af251140..3183438a88f3ecce329e4afc84272274af35cdc8 100644 GIT binary patch delta 727 zcmZ{i%Syvg5Qb~fCYc&*PwTxZf|rFNf0;Py3fLw za3em7&)_?_^fcl%>H3@bXa0ZYY2&@Bt9g>^xJ^6o_ui_{`s?`R4CI7+!?;fhpsNbe zKp`VdQ&{cB3MkC=(Qf%<7N#XDe>2dC!h;S<-dKQqzwDOrYtjBA6Fn3cdrL6ZG4Ie8 z%rk}UrVWxin3k|Z7Vn__+MziG!!4Jm06CSjJAs25CqcDbF6_u7SI5hshj`BYZot?* zLrO|elz#LdKUsCiSpzM^gKl8nq7ggR(j7-y>9O}uayMOap)YyFd>?6z6rrGMB>yus zzRH8nprN a*cE8dB&<~W`NaVkTQUCqpf5uc>)8)%Jf+(J delta 772 zcma)(&o2W}5XY-&y4#_`Q$NyFh`10MZ&j3t5D^DLLR?%X+xCTB?OXcZZpx)sXFZ7j z!A1NVZu~Kh#%@qegX?$ZJM;a_yv;x4(`GMA*1VGAtJVAbWBO&FIR{DYoeQbh6ik{1 z%R!cnU3gWee3jRtUUQ_a9Z%*gH^w;bLI+QGo^_MA1O^kL~JAnupQ@=BoW+74WWY^yH&V!VQ zL5haPf6Z?_56b0nNC$W<)uk`QwLs4L0O*iOh(X-!aYrE);P$0&SJ7frF;gr-rC95R z-8DPBE8V<%LMh}6?Ot~T%m6oK1v&qrhs(BXrIAO=jA)sgh5i5!Ma3_RBl;F!WF{od zqQT-2qliSXP&NzMXgPY^-O9fQ6Ack05X_aI*iJ*f}NZN_R&z|dFmpdJzO;&+cb r1V(@dWL#+zGoQ(=g-&(f(`fNFb@=0SA&pfCmjnm9kO&vM z5NlFmNpc1krx1HyX)Z`UN`*@hp*SZsFTEtAScspC7pQ>~t|3Z&vOKFeHz0A+c!7ap=o1c_fT5NTIegf>B3m9Tspc~yAuBQ_i7($nkr1@%u=L_bc7lrZKOHUIF>QTvt)OK zR0`vnvGP==_|cgcNUNz&&*;UBSuX04JtWMdTfFEq?WEBx%xGt=m1bmR2HR6@oX+Jn zWT}npD(U(JJnebXJYFd0O}z-$4v?nIxS6Zicr79XKX0Ot2Mej}m<4`T&x7N}4wBt@ z{er0>&tt9Ej-KKu746-EsYR{5;`0n^Z7@HfDR8ZWgeuxlretc=S_})I?2~J!tB1(5 z743M=JfAD+9~6pO3WOPQNE(Ltk5M;R9Gu{MHhh5L*!7`KYqu2;2fNu<*MB(HL-qt9 zDHLOPhtL>2q%=2hsGEeXLq@Z?4|zJ0FzUCtfvOaR$Pv>qT-0Jvlp*uR?8=D3ecqTT zrpX=%c4}S*jwK-QxAQV^3@MHbjA)X9>IaJm7?=rq|A}pZ|m?$4UR;g#p#y4uuH2k+OWYAm<_hC9vfMITuMt5*R+7i=@~r zjuJT+NjLyLZ&DQF7H{&V44i~zVL#oJfm32)9U7_IQivuQIE9GvvY!&wxWFfF$-pUO zo-f@3z(|`&z-dI244g(pdBAB=jSJj=TLw-ejZ?fW1E*o*$SsjZG|9j!BFY0+MKxeN zR8fvI9^ExzJX9#eZ#NYw%WO@~g$gCGxtijcll{t8q7WtUs^#;cM8$S-TC$hy4cLbg z6^FnH>Le`{&2}kQ5b?h`Smut><93RR18c*6yaT8eF}=>i#QM~!6+{`JRuE}Ewc=_L zso!1%)Cw}$mlpvwMLI^Nrid~?jd`_C?U}Vmowy6AF}wcwE~JKy1UdvWY(R}UwomO@ z_RD^U7{W3v6iY5CuL==xEX4gtiN06@IYCUCtuKL`#32pH3F6fT$VuFt3siL6l_Xa3 z;xeR$C1SN@K%GD)Lhio=q6|%IH1Nn+;1Px z#6{|n6+n%-_;)LS8nbbk8uM{LjTyO5?KwF~M;?cNZbT42$0Ox6>p5@a zDjKCWi5Ovs2iS%n-rTU+N5fU?fRmXE8g-mlj4)oR|LVV&FP$sv`mD}ll*oRvkKg}< zz55@sBh+BLQvY1DFt{0SexGEt*Y7xn3i*M+;1KrY@xkF2$v%FFUMdY2-#k~&Olj2l z5f5Aj$@z;@>$Da*e+jntavRS>N2q1Pi#VgdetX2fPbvRC*dp{WX{~7QXZ0el8UIk& zq9O%C)(s^~{&$Wi*LPt3^wP{Cjvy$8T l`_q5z=$D_TZqQLZynN5G&;LH}&wb2!8HyUNcCJ5d_#eJ~A8G&q delta 4740 zcma)=O>7%Q6vyKvb(~28t^1KCF%5}8fJzeY?D{JQP-v@OXrqD<5FxD`XKB`TY}a19 zDH5`T6bV&a+L5R@ap1%S3H3w|RX;+~rXT4?nl?aCQ$8eaz@R_T2i@E z)-npu7%)iu$&_W5jfxY{`qXqe&qOXO#C+=P@ol!2ymr9=@dsLhjig~*BkfsZ+R_nc zSN*bZ-PqoX!*6aUEoJ?1&eFrqQmytBl{+4N-FpI_qxq=>Z$J2N$4#wOS6Zq@M1m9dBiG5CA+PBl0?~G{feD?RHZ)p#C z-a|kEV!9SaRS>pvd8oUKbU0z}PMRO{w4z}6wJSp=8itTCt%)l&$QK60)27+Q34U*QB(EuV_o5e61)LZ#+sNz*d+o<4N84ZF;Fp#tc z6RC?BfryQBLwnu%fX>x70tHOZ;glkPi=Cgt`bNQSo*w2YMM34HVn;!bxiC?H$g=Yo zCJF>(7v{0PF(?Qsb_{}IePfWcrdJFicG^c^UoKm;l37cSX+b|}4zSO!cA>)&gNNht zRS6MR5+TMRK1Pf~{u*K&f_o9;szChW8b*x6iqBoch;dMHT+%qi$A}5YUqejP1y2BQ zc6I@DX#&u>jgf$|{pYko45%hvoYmWg*j2;Ja8KDJ% z7Ep!iTZHi*UQ`A^2~JBpM}QNu5Dj}T^j@K3A#20Z?5Sj zDL3d%?D~>2jMR?9vf|vJ8>K_+baDHO_HT9ZE6~M`12i^sUt)mf!yew37!t7YFhc?Y zhS-ti`V1ZuIXQ zw7%z1kLT3NM->5D?Yf1B$}TVlA$PARs<(Zp$5(a{vi`cNI!c@KzB;v{3_45hH_JMu zVmoIj4MW6j=W{}ytc>n=LJlRTby{@>XDHokzmwj}7LS(<;}e4qD#3O__~YGA8c1Wn zm+rKG2sm4*lqruxZ2qU0D`KOJndbolV1w17Xfwzh|T_jB9G3uGG)Fp9;|!hz$ZWJager;B*F;`}t~tlC=3ub*u_+xZ4h zInWwx<~cWT75TH4{}iFzX5jZO(pJ{rpEL?wChjrzkWYlLFMo-u&I{;Z!ASDI1o#9| zY5~y+BgN5iUg)f2wMST~BS!T>;|z&Q=@9?_f>v!tM5I=qgm6y$j<;*pV&`=3(_ioT zG*F-FO;^7@{JyJI`~-4t614E_NZ-q{Ncs z3@#=P#wb}Xw$g%>#FA9((vn5 LPd49Sj$;G>XrVTR delta 367 zcmX@hxr>XJgIkC#H$N$}wAgAn*GArWrg{r5_TW z2d9uS5mrVSalizP3$0r--av?hiw)>D vCN737u%KgJO1vSGpez?}S$t?-abam{YE^2ykq{3T7f?ML7sKX!<~T+GnXqSb diff --git a/tests/fixtures/onnx_genai_workflows/decoder/policies/token_to_slot.onnx b/tests/fixtures/onnx_genai_workflows/decoder/policies/token_to_slot.onnx index 2ac337e271fb50fb203de6db46cf55db8ee7532b..eb8e4b6599f5d6095617f9fe0d294eb1d098d216 100644 GIT binary patch delta 70 zcmaFByp5TMgIkC#H$N$}wAgC#M4r965?rh$`Pr#?Li}939PC1zT3Z6j|Ildw1!OG*P61ax&6f b9A)v&iNz)H24FdMh+I-)Npi;I^Gufj$pI1- delta 123 zcmeBUU&zkO!7ap=o1c_fT5Q$FzL7VHDZ!A7B_%&ERY;Rdor6IB{Mm(Bvmgdu_QSoNsNmFCMU$j#R1gD!o|eFm?gr+Q5Nrfs TAjHAN#=$Pc#Kkaq57Q+84z3^m diff --git a/tests/fixtures/onnx_genai_workflows/diffusion/policies/euler_model_input.onnx b/tests/fixtures/onnx_genai_workflows/diffusion/policies/euler_model_input.onnx index 4dfc8aa475265467b9bacd3f9c003cd9079f4924..cd1ecb2db40cadbc9aa5d62530db9fa3a9eefe0d 100644 GIT binary patch delta 420 zcmca7(8kZp!7ap=o1c_fT5R>1ZzHcEYrO>*TXAA;K~AcW9+wUWqmTv{yAW$qVo7oa z7ncx6az)Kb5e`qbMsSDbK*1e3Q9}#k`hajGq{w66mhCZ&PdG5OU)_fQWTQMsVE~g zGd-iE9iy^z6SsR diff --git a/tests/fixtures/onnx_genai_workflows/diffusion/policies/schedule_lookup.onnx b/tests/fixtures/onnx_genai_workflows/diffusion/policies/schedule_lookup.onnx index c1adcd7639612c95d4d727a7b9db456e3135df41..c74e15c7c69c47028a7e86fee9534309449589b8 100644 GIT binary patch delta 154 zcmbQsdY^@tgIkC#H$N$}wAktl%SK)+#(GsQj^gBu)RfYkR3S+&aSlcyVJ<--ewbK% zPHJ9yNk)`77fW$TYJm_x7cU395GNPA5NlFmNpeP#G#5unW^O86k`YNVOO}hPEZ#k_ WBqOyb-T&)0EwTPmtK+)rN_lmT#{NKq{gMf!7il0B_||26SG%>1kDNg>#s5torV>tjk9YVqY diff --git a/tests/fixtures/onnx_genai_workflows/diffusion/policies/solver_step.onnx b/tests/fixtures/onnx_genai_workflows/diffusion/policies/solver_step.onnx index 9851bf0d28a982d71d6b4f569a45e7e2e1c0d7c3..a3c4c2f2a45887656936909d57d822c7b08adc05 100644 GIT binary patch delta 595 zcmb7?ze)o^5XL#tWHTWKSLBF6MMSXBNJLExHez9wD%b>J*vk#sJ1%R??yX{TU~iQd zQG5g6#CPxkyhIWLSEZO{e$0H|{Q0Q-x+PZRxxu3s#Jct4ewOxF(uABAgJDFx2FJ|d zN08@ho-jQLz*$cUr6|(CcurE#mu7;W1K-H*HGof}zzP%C<*Qv-)rN)%FGGptd0{Tw zJp`tElKeRJQ~@{jv}Wqi?B86*N!8spDlanH|N`d{r|72^-0kvz4+IA@xb{5 D@3@cv delta 776 zcmchVJxjwt9L7y8?fKVsIE&gBQHpdZ5{&p#-NeDElely_@^bbP(o5`JV(pT}!QI*K zpt$<6`~Y543MmA~60-cBhll_3aKF2s^$Oa=rBPtHk+}Y-e^m}p@~S~en4{3AcnTAw zVF?Yq7Y4p-zvv?HV>%{6(onZdX*1Zz%^b~UCg;=_(8To|nNfGh@^h0=e|82&MMMbOF}96!u+C&NFZaASQRFWa?d9 zq+CnQS$>(}X~=zKkzNMY0LA3i6p~z}^l_figo?reS9@?tj1vri)Gb9QpmnbU)tEeR q@f1-yrNl6B4X~QK3TsN&7Qfe+W%nbW7^h0xQVygX>G%BuJ?#(Ji{O+1 diff --git a/tests/fixtures/onnx_genai_workflows/masked/inference_metadata.yaml b/tests/fixtures/onnx_genai_workflows/masked/inference_metadata.yaml index 0fbcfe3de..8c1a28821 100644 --- a/tests/fixtures/onnx_genai_workflows/masked/inference_metadata.yaml +++ b/tests/fixtures/onnx_genai_workflows/masked/inference_metadata.yaml @@ -124,6 +124,81 @@ pipeline: implementation: kind: onnx artifact: policies/masked_update.onnx + ports: + inputs: + current_tokens: + dtype: int64 + rank: 2 + shape: + - batch + - sequence + proposed_tokens: + dtype: int64 + rank: 2 + shape: + - batch + - sequence + logits: + dtype: float32 + rank: 3 + shape: + - batch + - sequence + - vocabulary + masked: + dtype: bool + rank: 2 + shape: + - batch + - sequence + step: + dtype: int64 + rank: 1 + shape: + - batch + total_steps: + dtype: int64 + rank: 1 + shape: + - batch + seed: + dtype: int64 + rank: 1 + shape: + - batch + offset: + dtype: int64 + rank: 1 + shape: + - batch + outputs: + next_state: + dtype: int64 + rank: 2 + shape: + - batch + - sequence + next_mask: + dtype: bool + rank: 2 + shape: + - batch + - sequence + next_offset: + dtype: int64 + rank: 1 + shape: + - batch + done: + dtype: bool + rank: 1 + shape: + - batch + continue: + dtype: bool + rank: 1 + shape: + - 1 contract: id: onnx-genai.masked-update version: '1' diff --git a/tests/fixtures/onnx_genai_workflows/masked/policies/masked_update.onnx b/tests/fixtures/onnx_genai_workflows/masked/policies/masked_update.onnx index e6c8140fc61aa97b93aafd5e911c6db03eeaa1fe..4e0bc858a5bf66d50ba9b15fb31731206717335d 100644 GIT binary patch delta 1863 zcma)6%TC)s6vZtK_Jxo#h*BO2DHH($0)8dXT}7p;Qt6_os>Ehw#}gnVHuX#NA^9Wx zB^5u=UAycrbj?3h9gkBdjvd7Mo_prrbI+OIXMbBoei9YBmVV(mFK%1ki+h{FE=)H) z+a{)qUF(vV4qAc*ege(G1bU=vZu1-f@ecXq5wl57Yf!qftt-nR#=Y<|6gt+~h3lX# z*x)&|Hbfi5yfB5px0;&nbu|0)={gj;nsZ4EHjFT8n0Xj?TyljNlq0R*;Sgn>r+)R_c1z?hEmvU+@xl<*569fj0mbh~TU8bWl=%`oe zsD@<{HjWVQc=L6;zy*si;p3Nvfe$jHwx7V1kB=;qj!N%(hzc-H1sZEF0q^5uuY<+? zY!fR+grb;NEDkV=5(;m3rjFKO+Q%PEhaQo9Ay|~rEX>fCipeS90nfF|(X)(MK{HUI zfXM_9?(8lH3~+4Gr|TI2jI}Eu(8o2$#bSfixd8%NpNG_N2W(%C?#p;zUHDi$%l2D8JYsE5_kT$$mx5x8)xw)$q9 zjn>muQkjXk5jp7VE5-5+$`gL8sg)G~AAdY2Ho;O>iSSpgy1J3xMHcB}vox@!WS8*a z-t%A>CTC#Ln)D`X(E_h+LbGCHSz*Qy28}UBLDl52vYNF{gjkEldQS|mNlv^j4&+>e z(7g2wA~)_#X8T#1PtFNX+D&pQJpuKLr(S?j1Imgr>ZrjipE$SMp9>KiuU|`pw?_FwE zY2_VEv~8BdDd5f_{6fJ_idAh53POJgchdZd0>QD|?U% zb>ISJ{#~FY7lr$}FZeiNk2yx#3W72c^cVf!qy$oQ|6NZp57`FhYcZ^tP`HxsLJ~mPxr9nQI(6 z`}r)DoV*Jo9(iRLWL-!@3N0OiL?>IAElfh%BRg)Bh^tYj!K=Y>QH-h`Qpl@$8289q z)4`WSKTsmd!WaPVOF;sL(#giXU>~1U1Z-eu7;ic+J%6c6Y7d4kQ7NPoqW-w7PR{^%{XpIe?yqf%skVu%aPFB4?v2^5n6esk21%o6=-fUv+$t@* zEg}zc0lz#(L>Qp5s}to?tai{dB^Imr zl0Dp_NCURWQ4FJSDoTGf$cmyy)y|KVyx$jDWi2L5;4%;38|pYjS?3mKf&e1(&|=dd zN@!zGeb-_{6?@J{Nmgav;1cd$)=-TD@6}LMyX>dBUaD&`O21cITiy3riVM@bsibr_ Pbsu>trp?(`@jUqlT8!Wf diff --git a/tests/fixtures/onnx_genai_workflows/speculative/inference_metadata.yaml b/tests/fixtures/onnx_genai_workflows/speculative/inference_metadata.yaml index 08d50b25b..bfc4b955e 100644 --- a/tests/fixtures/onnx_genai_workflows/speculative/inference_metadata.yaml +++ b/tests/fixtures/onnx_genai_workflows/speculative/inference_metadata.yaml @@ -483,6 +483,63 @@ pipeline: implementation: kind: onnx artifact: policies/speculative_acceptance.onnx + ports: + inputs: + target_scores: + dtype: float32 + rank: 3 + shape: + - batch + - draft_sequence + - vocabulary + proposed_tokens: + dtype: int64 + rank: 2 + shape: + - batch + - draft_sequence + seed: + dtype: int64 + rank: 1 + shape: + - batch + offset: + dtype: int64 + rank: 1 + shape: + - batch + outputs: + accepted_tokens: + dtype: int64 + rank: 2 + shape: + - batch + - draft_sequence + accepted_len: + dtype: int64 + rank: 1 + shape: + - batch + done: + dtype: bool + rank: 1 + shape: + - batch + next_offset: + dtype: int64 + rank: 1 + shape: + - batch + rollback_len: + dtype: int64 + rank: 1 + shape: + - batch + continue: + dtype: bool + rank: 1 + shape: + - batch contract: id: onnx-genai.speculative-verifier version: '1' @@ -500,10 +557,97 @@ pipeline: implementation: kind: onnx artifact: policies/grammar_guidance.onnx + ports: + inputs: + logits: + dtype: float32 + rank: 2 + shape: + - batch + - vocabulary + logits_mask: + dtype: bool + rank: 2 + shape: + - batch + - vocabulary + forced_tokens: + dtype: int64 + rank: 2 + shape: + - batch + - 1 + forced_length: + dtype: int64 + rank: 1 + shape: + - batch + outputs: + token: + dtype: int64 + rank: 2 + shape: + - batch + - 1 adaptive_k: implementation: kind: onnx artifact: policies/adaptive_k.onnx + ports: + inputs: + current_k: + dtype: int64 + rank: 1 + shape: + - batch + accepted: + dtype: int64 + rank: 1 + shape: + - batch + evaluated: + dtype: int64 + rank: 1 + shape: + - batch + committed_tokens: + dtype: int64 + rank: 1 + shape: + - batch + filled_proposal_budget: + dtype: bool + rank: 1 + shape: + - batch + draft_ms: + dtype: float32 + rank: 1 + shape: + - batch + target_ms: + dtype: float32 + rank: 1 + shape: + - batch + estimates: + dtype: float32 + rank: 2 + shape: + - batch + - 24 + outputs: + next_k: + dtype: int64 + rank: 1 + shape: + - batch + next_estimates: + dtype: float32 + rank: 2 + shape: + - batch + - 24 contract: id: onnx-genai.adaptive-proposal-budget version: '1' @@ -522,18 +666,94 @@ pipeline: implementation: kind: onnx artifact: policies/grammar_length.onnx + ports: + inputs: + left: + dtype: int64 + rank: 1 + shape: + - batch + right: + dtype: int64 + rank: 1 + shape: + - batch + outputs: + minimum: + dtype: int64 + rank: 1 + shape: + - batch grammar_sampler_logits: implementation: kind: onnx artifact: policies/grammar_sampler_logits.onnx + ports: + inputs: + logits: + dtype: float32 + rank: 3 + shape: + - batch + - sequence + - vocabulary + outputs: + last_logits: + dtype: float32 + rank: 2 + shape: + - batch + - vocabulary proposal_metrics: implementation: kind: onnx artifact: policies/proposal_metrics.onnx + ports: + inputs: + proposed_tokens: + dtype: int64 + rank: 2 + shape: + - batch + - proposal + requested_k: + dtype: int64 + rank: 1 + shape: + - batch + outputs: + evaluated: + dtype: int64 + rank: 1 + shape: + - batch + filled_proposal_budget: + dtype: bool + rank: 1 + shape: + - batch cache_length_update: implementation: kind: onnx artifact: policies/cache_length_update.onnx + ports: + inputs: + left: + dtype: int64 + rank: 1 + shape: + - batch + right: + dtype: int64 + rank: 1 + shape: + - batch + outputs: + total: + dtype: int64 + rank: 1 + shape: + - batch state: tokens_state: contract: diff --git a/tests/fixtures/onnx_genai_workflows/speculative/policies/adaptive_k.onnx b/tests/fixtures/onnx_genai_workflows/speculative/policies/adaptive_k.onnx index 2dd29eba585333888f1074c9e30ed84e086eb349..58de18d70f6992b9663deaece49f999960fe8032 100644 GIT binary patch delta 5043 zcma)A&2Jk;6vt7TIOC=eJ4w?*la!=Vn}Fo;`m++CDj=w!DN;lts*2IZ-cq}cop{$K zZ4rau1cxGJk4RiOSKRmq_!HpTUXb_)5IuEf-o`sSA2#Rr+xPK%zxU?ZKmIWB(A{tL{p&BDJe~ckP2WH>^?tACbprAr`zo4EPiJS)boR`?8`SULc^#$Qdfht= zyoPRMpLzqg-FJ=FHS}7&dvMSUn3V+G2VTe5U0j<$i{X3UFH&KeQ9?I*?!AB+LGyqm=4W$gE}hEG4r@Xa>81S(XsY8K$0c8lYu>+t&Iy~*Mfhju z#UAxH}jP^2g!b1f--xt{>g4%U)_4gq1qk zk|2e=36g4wdbV6YAO%QeTH-5w_AnRzd$djg)ovb=LQ(jUQ7lv{5lT;3K%`LR@lIs! zC`eNEQXTKbNQ_*y1H5SUFeMDgXC42j?|F|FShFqc%=qt6HYo#gqtPHvNoZM?YwHqk z_Dz4=-6l?1llrOoX6GJpD!Tb1njVnbJ>pb#n+Utz4R(oB)2%D$?0|gO^VlHtNHISS zfHgq)kAIOSy*L`E1+{k?q*xqf1Y+bNm%cR661jbs`8BYTXpTpUAcPr7xk4%%;nTNs z7m5}QaO?b%Wh^$rXTjp4VSl7r8zIgwQ1djg-83L#no7V`LJCUSB^V`TL8V%f6N8?> zO8LZ=fVe~qLEQVz0Vx4*38R_-@TMi8!VCn4&mU&5mGjvtG|4vXD=_{a9QEBcDZ`}5 z-QatKmJ37l3LluA{yr(|J5+Q|0B-e}GYdxoz(=lZSvF$z?cMi!9x01$L)Dic+F21s zUN_RVi8Hlu6=!N-i7&>}Oev{Yv?irm(VqjshYTs0r8ms9V&;@WCY;pDI8rO)Nv*&N z*R&&Nn&1>j_pxH3Px4o_@V6Pd#|ukF=FYUhXygox=S@={eV5gV3V%Ib z-h!sP$E0duK>+v^RudeE!%A9BRg0*^;>uNZN*I0fG@iO)9yMe2k@$EXH6sJ5)|&iE z)Id4%KzAPU)!fQVz`278xqUf)M50?0Y_vCYq46=I=Agkf(m3j8zLyHoDN}t|9>S$?ik=IxvpuIyjFiIly^j&~KHj5WL7$076`cS2X-S0wDoqRz<)T?jv> z$V-w})F)9*XY8h8H@liJ5B57nFwME>CRa)Mk7m}zgAX~E^r}QEd?TEvY6K-qc7TsV zj1qx(I!YEU_|lQ_^OP!V@Kf*gb9xmB0X} zBZIL^eY_uP*4FIAt`IbY!79QTzWOx~vfmJd5LH ntJ|r&?5*dQ{GQ&_=noh1`wyE>`AN9AhK}@i^|!CS%T2rl*tG?T delta 4287 zcma)9QEwYX5RRh|+v}txj*~QPny5;PimF6j+vl^3kO+mMN;IX45F(AVI@=d&$FUvz zd`eT$0TL1*1P}L)5Wj$mcYZ@3`5ArUncbc9?cDY*%IkhRGyBbaGdK6^Z__{jG(Dfm zc1X549uDoE!;bSe$z~>G-F&`dwXA{Dc5QaN z)3BW8vx9Atv6@YL;MlFpT<7j>*XoX~g!`WE-yjRk{>e$(;RejaxbT^_A^N<7^@)Z$hdyw2mBh0>M&=_vLj?%CUyr`6cg{coIfV`-CT9 zl)p{3GO7G#kYg_{$z(cnwQ-ls_UyCERi$B)EAnzun4~Bjf4ld3WsPKAw$JNrvvPip z5dOYdGUeave~BRzL4WlRR<98N)>4(A?`61@rYg`G32p&ff*vMo7JS*VSKnKIZU7o+?3e{ zx#F_Vd!y5_ZJ+ZhLsm_Y1!o$&?wa3a;d7j1R z4|_+fY@+iel6KkSAuHF=p>)01cUZZOPOp+Vmwhs{xon55a=VxaQHA)gM%TQkVOgOa z;LvNaN+tFbm;nQI7xPl5m{lZy1)Ei%Fj=7xS+gzw<@Qp(qL)M*7Z!D;iPPl;O@3QG zEiQZ7;j-u)sYwH9LQx{mv6%s)lwBoXgUoO((J@d4Z<6vyZI>D51DUn>AV66K(I0`h z`_}1_|CyVEKK=M~%*SOFG$&Ig@n75=uO!#p{WH)D#ZpArlhY{8 z;53IQ`oBNlFrWt38LR0uBqEesLmXr(p&x5{Ipy(O4RP|HoNs^(AMprvyOf+6kJJ$% zC9_ZM(TLF^6iGczewB9ljHE@>lF|%8KsUx}hZc3R(oBqZS_HPr?3Vw**Q*CY5$X&@ zfr7?Dh_?PVG<^NVZwl`g>tRZ6l0}yZ;*WnNJsmfQH?P(bTGqv&)LV{KgP7lxuTaT+`5qMm26E-gR5olEkq!UVx4D#*(1(0` z3L#A>6{K<6X(DfryEqq1PHI|-5I+|$2fGj_7rPK^QesJRMwA2> ZYf)x;2Bu6BP^KimBrykB&tyZ!MF3MT6yyK^ delta 178 zcmZo-zRb+a!7ap=o1c_fT5NTIc_MF5y%85nPHI|-kQSE)2fL68my(b|L4HnVa;08! zVsb`md`@a!dPzonX+cV2Nvd8_Vo7pFlra};QD%AuNt%*?no9CZ5_3>(lEi7%#0QH2 DWYjl; diff --git a/tests/fixtures/onnx_genai_workflows/speculative/policies/grammar_guidance.onnx b/tests/fixtures/onnx_genai_workflows/speculative/policies/grammar_guidance.onnx index e79efb9b0951cf011dd6e1582fe7c860ff045b82..cb4ad9693ecdab7f45483069d859c941cafc93d6 100644 GIT binary patch delta 356 zcmZn@XyfPQ;1*)b%}>fKEw=j1w~<$swO*ZzEhj%cv!qx^mP?w0QAnJNU5GU)u_QTz zi${p7EI&ChsWd0Cs4_~2iyNXaJ~y#A8?2fWw`vtG-n9IpAr~-@y5s|fjv4|kX_bMj*Gi2-n}R_u_U!9-W2Q(P9%3^0X8&r& zXYc`hBo97;FW^$~BvSD@|9sy+%%Z;7a9I`YSbAZWo`2L|-Bx+t2Ro8Ms8j60H7j5j zwy`;tQRvUwf#Pw@6$M#1*M{uU~B#PMOGM`ebe_8i6qb?+=rIIch=%w2g<9)zZ5Rt%Qy1| zI63u4LJ2xT0A`+C0t>A17C=syD(LYiI>wrbDrlR;95~*DT25CoNi|P2b5X(I0=PjN1moYb@uA$~4i4t60S`b|KcJ#FFF; zE-oRC;?%;@)V$GSvi#)4q|%(kqRJ#4F7BMf;*xl%PFXH#piXgwPN3?^flT`W D_oW_A delta 282 zcmey)GJ%bkgIkC#H$N$}wAkt|>qg#S#`=CnF1DQf^vsfCp%zB2Mn(=sp;|6gA?1Sn zoXq4(z4W5Q+}y;X_~OLef}GT%c!*NHq{Ncs3@!~JHT=qoQwvK|^O94!w1hOUtE`8a vSeBohm{gjRSX7ym!Nr}ESX=^?5sKsr2RbU4SV!Ts3e%=6F|Nr+88-m{TAO4u diff --git a/tests/fixtures/onnx_genai_workflows/speculative/policies/proposal_metrics.onnx b/tests/fixtures/onnx_genai_workflows/speculative/policies/proposal_metrics.onnx index a8ad0be82b3c1f4cc5281fb4a4f14ec3a89ed4ca..0084a107f34b0d3c63c7a853108d6ea7a925af8d 100644 GIT binary patch delta 175 zcmeyx(Zs>a!7ap=o1c_fT5R>6eIsulQ@tJ+e?d`xL4I*+N_4Vk z)VWlFW-G$YuE%9~k_i`QYFT1VX(CV~hBlnK9J$2OGIMf(cEMZ_pH!NXo>~Gnf)mMR Qh!MQ`#Yo0O~Vhi2wiq diff --git a/tests/fixtures/onnx_genai_workflows/speculative/policies/speculative_acceptance.onnx b/tests/fixtures/onnx_genai_workflows/speculative/policies/speculative_acceptance.onnx index e436431ea73e454fc01b3fc829f0187e769ea357..912d023e99436d4787610497e7782c1fee02fd73 100644 GIT binary patch delta 1283 zcmb7Ezi!h&7$>MC&X=U2heTyZf~NbuX{vNP5V%6IZ zNu47mQRv)Zf;6D6q{vgqkb4@H`Z3%mtA;}dA%LB^%+?v?33yzmnrl(#W4{B%k;6w^ zFasq&Wws!Xp{is^X#q!n9aaQmhEIx62(j4&+uwjR9}EPOGe=q|MRlEx!5hJONGSt zGHxD1K7?~yNNP*$kTk4>6TM?bv3A|zdv}Q41w{V!VwuxS+X>S|-WT%{+#X?lbtawTV%>s~OC#b*`zNZJ-*r`VWn6<#i*<9sAl|B6sk)QffxN4v(3` zP+cZ{@T*cf4A=B(LDZQ?V3x<*7P)A(rXJ6D!_DYw&oHpno;eJ*s+W6DVAaD9wg^8D zcEPaSArD|JT#!QY`d>1Rh-f~AymE|$v~;LD z2C^cS71_+CYAtu^ip(c4>C$6?`#nYDG~A{3nd6wmI{X{x4bUy-inbf{!EgTgapUD; zIhq3qal_x+Cdc?80?_4$k%knc4rU=2;;o~AIM`4Jx!yqmvLW7K0#}E|GoXiflW>7+ zLpn2%mviQF`DY1SM;S;!5>IE?8OZuM8lq{lDHC!7TNv7U<%uCPa2DDcPVdC$l;=GdF-KRdq#t>4xv`~D0y<3%P z&_Ft5^c4cDU=LMw0V+iT@gI78{1cE>i!BSvCM7De&$qG2g;wBi7 z0R%#hZ~=LQ8Wq%00u4k<35lx%V~GZ5i*?o2L>E^Xg;-Owvb!St-TQZUI*heNRl4u( zd(J)geCK@k+{4z762I(7%uOC-Nqn^K{)g7AvL-mjI>$Q3o$5$5eW!`uN63C+;^>gJ z(IA#?BoCGHkx7N7`M-H+)q^$efOQnt&d%kYP(qfPSy}hMa!e0PPz(5{&<#F3WvS*4 zGEMd3LpUqR#;4y5ai`zR=8iKwHTd5MSGqm zlQ_PEl8hd=)e|$hb*71r&q^aVPpY^-Ty$pv8Qhj)zT~qeRzC1RjmN_^vrD)l_@Dgn zoZ;l*A1t~2)$GUb)~gf+4*N1NfEZmvZQybj=ftNT>3@*)rv@e*Q|Hdr@v_ z;lhI9>{yn{#fy?e%@dTwiwg206BYmvwRroYB~#l?%HjNBON zcgv*z9<_bY#lZ)HWH4sc4(>GNEB_M6d*cpi3@+bFIT~DC49@VpB?p?EyMG% zCS9~`rUWh@o+Pc0mkU=fnN8eTB3H#Une%eh@W2z4&g({~;|S9R%GJD{mfFd{u43+l za6|2K;vGxv;_Mcj+PrKU-&$(tL7x_IW$o|y;K)7@AahCd&rM}xOW#6x) z5%zG^@-5U-7H5!;JHMRG4Q0Ap>mJQC1gt^M+>)yRegBccnm52dlsOD*|IRNv)ea1GfO^BDKFu4H@otqTkMn->6(<+3oC{RyS1!cg(e78`^sEOWQG{Z3Zj*4{OG9|}<2JFhiW1rV_;iEE zI#YRV_v4NTgxYU*$Z+4aHbCBf9*}8g?nvS4e>~3#>t;qb?tIcL>JEYjd|t%^dEEJA zuE7TekWe(QH}dhdChl36M!xX1b)zW4omPFBrIhEbi!!Ls4+hnx;u+NEk22`**Wci| z4NEm&fKN1dMR_F^DEib*JEIsOeZzg4$HZN4OG6UH(#I%Sv~HL71uR_MG}h<~f>5Kj zeb}HAsxLTzPzNf^)Cm=aMhMmC*9kSYLB}CV(tJT~Y_fx|M=153P3hzaf7k?GmDyUz z)tjs!PsRC?SLGg(S0#<5Z%E?(PrJ#XgDLeMehhS~2-fe6IW{s4mQ6Oe;#^$4$;36A z(~#)b*4@t^ZoW!xc|l1h4-kn94ZjSH63}wr>aW?lbvV9B>Rr*IQ}2qF5=4Wf-gmc- zB2W1G)(#mHGOFv3>P+j81V`a6ooPeOI@9W+yIM;5&<+b^PVDe<=#-6%51V+<&TPa% z@y-{CUwy{Te{LSa-#)WJv1X`QvL@oIC4`-KwFJ0r#|Z9gN#`p&O8H-HSzNfYK5pf3 z(YgvJ1h2|3S>jp)nfTr2Ql2Ojr!wt&iaUkQrTsQ_@5$r~EhD*m57U6{2K6Z>>^YPw>RZ5|$i{sx`~`!ybeOr@pRGu==M_bwJ%`do$6iY1f8IM0-j_=i z+1|EPQE9)*J*%ix9C`v`pi$9t7iDp^JyQ#~x$VGY03ms_{NQq2$XHn-9IGRQ+jGbc z*R&ldk$kEJ6yY)!Dq5BPYiq57xa8H?W?p5K_ZRF})T*l-Ac4CP$52ds+c6~n#wm{S zFQ&NYrA&h#{Cnx0Tm`h`-+*88@BgQxpTna%B!>^s0T%x1b{&eQBNDjroGsCe*FerY z!{z6vNW@Qcm!Ra$jiSDCL5R!GONtsCcKk7&fPK$iJW*;cWNAP>^9Ym{U&s>X21?c( zUYtJ&cCMd4Gb)98`6AojUxD10s>d&o2a1)`{mSQFj42m&dnrrv>DooxUb1NpKQxy} z!iKbefd07`rT$WGxDqLf$6c_if#hzi#e2_pxbt!bIl|qS^N33?`=#0PSEmfI z$gfWPI*Sk~vhppykS?y@LxXts$06Lklhvxt_5y>`>5_!a4Ye^3@$}MDu-C$ABT{(s zC-VA`Dmp*O;?_@O3QOq{PW!Y#PvHTwvP5qDYzac)^Uo%PpSz5BeeiRMJch{Fe)GJ# z@!P+2Ar?BnY|@A+^uQBX&ozUs{fqtK4PPB1 z?pcV7`a&wFU9pf~4}Up~IJx;sqvi~-xa=ehIS zIa0uUb)3N&j00`rH%Z(VNHREGG7n4m=+(jSb3brW(#ALbjEv|vA|jLkAtVvKhb;Uq zk(|4~pc2jHk7yS+eP5zERk<7B-tT7su`2jHm#1EH!n5*%J$Y>>@l#iA zytuCjffcLEcHYbj_w{8;(v{gS`VOK{zWx`Jta4KU1F}Clmj%bm^uMpq%#MmwoW_G! z?o5UX?&~{m*bFj3AsSn5@K{FhrN0yo(%dD5hP!?&AmcHENf{VvX}}f}&pVqZmE;&*ev4gN&!7J_5pS)L z(sJJVzb4HU;Eeye3093d=4X>=Sc7{lwXcY;V}LRi+shoaa>1>1etBjFd=&{7)WYGN zH>_6h@mVBMlC71h^KBjvLDrXnPPUlhT%hY|=R=#mU;6%Hfn z&*IHwGCJLE73Fr(y_S-6s&#urXCeg=D~Ruy1_2Ehg*Ta2;E^Ra()rdMDXNh~v8C-b zDRG;1=K#>=WXja=-~y?8VCcLtWD|8oD1a61x_tw;g1p+QR5TueRk9x;s*Hr1<_ONe z!H7$uy&y-NHBytLTH#2i#c^j#)z@uKq>)MiIuJ?XY$A)S|h4a*?9r*cuSBTvCM=QP3pjEvDh3I-6#L zlqah-toYXBjPh;FLy~Z%)rxV&T{+~^941k@1b4}jY}DGVRC2)(fB`RigdO+!p7icKle0>7a(Cd&%Cb!KH4^O0Bf=6K%NXV{( z=pIE60m>=~Law-bG&zxlVB+?-N7GK3OEAGIR*az{aj*mh0Ej^x4pQEkG6EZ`&GW_a zUnBa##Dg@iOHpy1pAyBba>@cVgNwz(W00XBWAR=o-dEErbRB!jsGp8wBIrbp5WVFT zBV@0~I$CUS(?CY{dg2&anj|S%f|?+fP9USOx@no_^@;{Jplm)@^ikZ@KdVjdl3yFugvG8RIYh_qSZ@olh+qz=k`l(gsh@xWp7I40`Yo; z`6xMzUZ2ITeh|3c67dilsZJsG`%!p%{i4W4IzJz-p&VfwL&f6wQ5p>9IhP3LC>xIl zCP-OxdiCIe>)UZokTQe;dcL?Gpgh@7lUxiXW|1LKMS5i`J6HT_JQc|Apk#UG9+MidoKqUmtVgOg}%BsTsuiE?FZd{;q5_(%;=uS~`e zMsAT_32vDZgn4&X(kBq?I;Q5QPo@#%X|J1%1WNUJhb7N^VEM|`53PPe3=Pq%?c=7> zBOycZ_8xi4CYmO@^3|Kb{tnT-5d~1oCdw9HY$6{v&3T%piLrR4Bs<3Ht!vTzG*!ci zy-x#ApXmN4DpP&JYj-K?PeLwsZNSw~y$xI$x(DdD0uLl=*JhfIL=&$yL?l#)S5D)t z8g224!@pI`)3^ng9caXV6~vmcZqc_7lNDPCy%K>13q`{gfO2aenEm@LI8?KPInqb}R$ID<$?K6PTD`szm1Ui8^K1+rvEtxDSroY2KLvmy8 z!_xBunrZO7TpZV}LyJ%15H~AdZ(`YYz%gk%*)_jIT#Wu6Yn5q!r$kAYQnRvC2I!*Z zCTf1S1Z@oZTdJ91ddUSNr8p!efGP(m^;1b;deE@}Q|C^kG;=_$JRzA7Kx!nMG z$59;Pwnw+P^Ds`YPvX@r3SN*Vy_1Mnqf5*YG!JIB3YrB_zY{bQShlrdTf9iKWY*!` z?h`j7GxCipF7BmSnj@el)r&`e;)u_WU?2J9G5@;>sh|+-*eCBQ3EVyKjGFrOeyWsP z?nktH#5emv8D2e|qpv3wH@YhQF*WIVU9-6P9PDuhBqZ6Q@Bjtn)9Ky=u-c)Y>+tvJ zBZ-d>(0$nMRh%^Xf`~ps#QKADM~>%joj*>m-=p_~3#H|5>4xYdK&(56Ju~27Rnv1k z*}as-vFbUX5C$lB!cjnChO7Vq0z$)xrK%ErpcPz_c7#Utmzl4&Li*Yf6m($bC>Sj zIt=hll{HipH7H+#Ju9vbl8g@c%J#f~_6JxUvFU+bz;BE!dTgTTAXzl8Q$jfd)&RUR z5i3G++`kCp?f1OGQ3*MY;=eRx#UN{qn=jH7Sumh>R5~%~mjKSdWvy~%z-=wkjw#b7 ZAESxzNaHcwvo3M`7{c2vdXLel{{vX{l|ldj delta 8864 zcmZ{p3v`r4mdCfg?xd5ZlQan=olfTgq!YkE`s;TKO5%%$Vv_ZU>nQSA(LoUsaYlzh z2`Y#PNyrt-RU^r|!#blvFtcPvP=aSXxT^@P<2s{Gz;PUPRBSyv>Z~IBzg1s%I_R97 zll~r6w{G2gtA2Iw_s8e6H~)8bY0j9!>|0hXUH+rh-mJiMb9&(1w~uE#KQ^4Qf2OG3 zd*0Lr`(~Tdq`vbfcmKp;55)>6qr}n2jZ2MifLhjiSP@!zL5T=c!2utAnQTxk`=^P07`T{*27R*q22s}i;aJMP zD9pY3e(#m*){6Yx`Acr&LBgi2a*C9s9Qx>Umk~Ax`Y!Sd<6muZ>;~bbPtVVm?mLCu z2qoyVb!Yd7@_t)zME5ai-NQj@YrkZ)@;A@EQH1E7nHBWE`%8^*h};LRNhzI6?Q>3v zpxk!5@a0Ek?N>|^i40MG z`Mg~`qXs^+pqPHY**799ZX^onvui^}1fIMy#y^KzR%UisFHC4S1OG>q^Hck)DS>(9 zJCaMj$_$}36j{cz?ed1wjvQB!w zx&lqlm8&imdE`Ch#R&Pb_lF{jn*C$xXG?16(bvl9gC&pC!r%VV>Xj$k3q&pz*ZO!A zNPg!b55_N&Lk%KZ329j77g1_iR86luQVz%W6`e=N7A>Qm$tf8Hbk`dWS}}#EG&118 z^~*L3ik`#MVF|mZ=!wKLCKCVtiF{kf6B$ppEk8l-=|BGNOsxh2YrpRXG|iLscZ&jh zG~i7rBbDDY-iQY1vuTyIdH8$!KVk9!s@GV<871?IY8txP&3(-6UfI4*6jAF*2WHz6 z#4=pbgorq_Vu3K_-ZJszclGI7c?%rPvEAI}pc#<@BN`kPLa#(dLZ~Z>7-@3Rp4Kro zb5TG=xg5G%6dFMlKz!|>WZYpyLn@4#+s0wla~O>llS{p~PJ)h5ncQtazL z+FAYV(J1z#vGct$Fk9qz|0Ks2Q(KY;lJY4q%Vk7kYF|2AoF>btj-ZQmZ%-p8suAO% z6Kk?cEbJ%_M9(G+51c2prJ zQ$Bt2)1^j?krTlNm_m&;#eJ}ri=JEOWZYHX%f&*>sO9cXsJ$!0?1+GSUxP8tZ z3fwab@y>BKbdQ^Qo^#q`Oq!I%sOuglkduQ=i;3fbnEdW}E+eL?llCr)s@gEa7RIsR z?7kL_=((qO8SdRsjsK0S|4auqa>#=|^17QuDacbw`JL5@J~!T{=+oVy;gjsP;Iqv| z4G#cPO;r?7*M;=Mt$unM&$f&&pvSk~O@S>(=r6@S7KMmDegb=^)$-=fTGhM#fgJFM zdyS(vJ68<^w#1o0e%iXVgr3{BUSp)8XB<_PJHVX=Ek~Zb346)n&M|Z^)ygfMBA5_LC=>L@FPv6WeO$9?rYL&7Ag3|6?{t3 zzi7uMkpu6%GczSZLuhVyhelblyNncnEW7Z`YUE5>vanQ*dAcVbJ_lh2s&-@vBd!^= z_;Jl3@$-qqEfj+uKPd%$R?Wk@Wb3>Xc5riSlPHCsS{7*~03IW*H<@FL8dwuXP_nyl z1V0!+;isB99&x1c!}iOku188RUoC8^MIcD2rJa+Xp@yDo>FFu81Oe2yEIkMEO}&(= zb{u5{WlCWV@no1osmUk~Jzmu-$eeZ~-sZ{X&BAR<04Lm};@^rXIc*0wP{_I4o?z1G zSL2Pu2x)k{UzolO5ui)oesA~vPT2&J+}7M)BLg=J2m6WnBHh;j(YNoKYa~p1Vb_)T z(U@Ujf7(PB;_zR6a9M^=y-%#-vv^Yym^v`wFNv zS(>`Zkw>=(r_Brmq&Nq3<{7u{E1|pt6`YP(v)HGsvS)Ct*hGg0=ZJ``>H#1uff+KA zbs3LcR*GT6snHZh(aTAP&R8QJ++WP-30m{(-0!5o zS2;!i96`o*YQGyC!AulLteJ4)v&XWLgJwPFqVo?$o5K2ql z_vB|pt{G5yPC9fC4NbN{Gp|Efc~04(^R0d8)f@Ksl>Z6rn9f z^GDDB%{Ju;$2oZt2xi5Mo164t;6Uv|jFIUepfc zY9vLKxZ?9K(8OvEEc~s&iJ8V8uA!B+E_tc|>!F>|efhFt42uK8>ldadq~^uB5WBUh z$3+|7xB%i_dSj-D>EeK;FVB{d1a;+}B7xdMrbKQZAJrZMDJ%^%zbTkYm($L-kUt~m zWr1J6Rl-P8FBOyK<8{;s1{5B2i()6*;oH~vppEJT)ODp5(7poIqGkls^moT_ zD-JHcmqV2z#|ZK+pd)(GyAE0&!!pJ*IC&uXUa>GDTCe8!d?KV)soY68$Aqy={rJxn zh~_sfsx1Fta*8VTpEFe)^vVabu`#wy!k$$3VMzOP;L#5|1=g$iMFq|KsEpo9`sA=* zG(eDrhvX-R<+n8mmHN-8TN=nfi{K>3g^s)p@vUr_qK;H8`$}omCl!h@O9`Egtv0>Q zq+8j8vBD|TmaNE26{yT$+*^k|!9#&hvpktu3ErGl_+rHrVh;XwjPca>$aX5(>4XOo z3T5kd;TH+D9k2Q%d!&GbCB)UQOqn3j5+h{M294>3qzBtAa*}~BlQ)2`3N2LiRkrMJ z5jk@Ay^JSxV6>zkr;!k1wCI`H>D|xA8zG&DSoYmVJ5Rajz{O5umrbtPE^3T~Q^6%I zX-o=IAydtx+6ZZ`K(`F8gmft^V`(UuFqE4WIVNaJE>_c;HussDb5Tqr^G# zjNgblY5nQAJnay*JFdX;>UuFphH^lS;nQlfymNH7g2P_fdIff@omlhC25F8H6Xj_k zEIP>MY_W{DbKZ+HhVaWfr9}lLp-52yRZq}cy0e9gWmsGY2ezAH|^nLc%O;5NQD@jTt5il&Ut?palR5 zEMOd}W0NvRed!H71I>pZL<0+hCkw?`-mKDZ6%B$$SV2JI^--#P%ctgU(V&!X^5oFN zVk~n(`>5X|6d6YMNF!pTkDIHD$R*O1fIArwHo=^?qh>;*%46(Nbq zDHBab#FST;i3Ru_M4{|GVUK_uIy6pIB6d(2lS6({F_R)tV#cf*+;w$dz2|FEkh|fjF{H2roouxh%k#5m1LM?1L=`(!bNeovb7VA~B}6nk>e1~xnnW2Iley-K^kWdp_^|Qlz=4C@ z^zDAP9Eb}=imPQHBo3<-BrrwjN}d^2jL^QhEGSSs);S2bkTb(5&7zzNR3+smvqSehb`jt%Kpd1I5~VP4n~I}Vk%DLA4kM=-Z}=a zieg*Bw$K_AmH6q3iFfenLq3%d6*wGr$00k;d#gMa7jL5DlL@qo%a;?PHa9UbTf8_EM4j~+=F*M03^d^k zC3kE@cm6Vjk099Q$d#aQ{bO=6<{EFl~3x0?Q(Up#cbC|4cZ z9uh_J5BI}D)^~ylvSzb-r@uA8`p~CHuHK9$nbmwX3wK))b~7g$?OlUOpg501t+Gj~ zI4|2^A*B1uva5mVAR>llw4`>6vqk+oOwB- zFxUAH__O7-du~#TVt+;2sEM>u%T#qj{^M2zePxf(w!b~dF^(DCYK%bE*S{Qf7*)agyhN83(BZR?0p2jSM>;4(IHx3l-3S$DU5QU1D)dX z?QYCeXX|pC3?@a6A0Gt6Ky)D}1vXhpHf3zDUj>5kn`Q zfR(DrCrJDNLo_~$N2xSQ7H>!6>mGz>uECXUrcPnfA61~LishBtMJ>iX`k=UkGnA6; z;zFboY|S{50=D6_j3Yc7^Ad%VThw<(COMP??Vdtg`DgbqLeOHTk+kn-dl;%Q)z< zsZD)kRCR;!h}FxqgqZx*HfX}f*|hQ1tJH~0xq;uS2X4Iy7O^`JlVvz-I{WeDF?AS>4Ov-D~w6POLduN9k`1NQI|akRRFwF`C5ei+E-=V)C%#4 za@#by_*Ri`rHo)aJZL@Y>OuY=wN}PRx2X225+p6^>P3=?UDg_BhKPLlF>xN${M}<> RI+llHr!cWR5$T5M-!{5JUmtERDVS$t}4 zQff*{W?s5paCv5NM!boX7+gpnB9xb(k{WNK#0ut3j$!}6Xgc{Iquk_d4nZL6BV+93 zTO89Qell@kGmdMrB8N7!sgTCx`yz6adAK>5&7`y^`*X-m{=><`Xg0Z@*&C#N@+&TF zW-~6O$qz+jCYx~EFq&^JK_KfRW9;PBT+_t9GjU-W#5Gxv zU7OiVNMo|0h{7Z;c1H8b=QwnLWIwYvNY~^$oZ8IhTuPH0gybe`a9J^0OcrGI0LuJk zwA{R$JAuPi9ha;5xOg}ig*dnvfRKrUBTE^J^^;{q?0`;e6gHatkCPK6X(DfryEqq1PHI|-5I+|$2fGj_7rPK^QesJRMwA2> ZYf)x;2Bu6BP^KimBrykB&tyZ!MF3MT6yyK^ delta 178 zcmZo-zRb+a!7ap=o1c_fT5NTIc_MF5y%85nPHI|-kQSE)2fL68my(b|L4HnVa;08! zVsb`md`@a!dPzonX+cV2Nvd8_Vo7pFlra};QD%AuNt%*?no9CZ5_3>(lEi7%#0QH2 DWYjl; diff --git a/tests/fixtures/onnx_genai_workflows/tts/policies/code_frame_update.onnx b/tests/fixtures/onnx_genai_workflows/tts/policies/code_frame_update.onnx index 102314d5d32107012502c49a271bf2dd75ed6280..76ba875d751abeb777abedb6463134ac1114ac56 100644 GIT binary patch delta 191 zcmeysvxbLNCJaedTwHmn6(tZmh;qQ>8!VHgWVv|D;zRR_3rkZ|t5V}lg?PBQfCjR0 KF--1eT?hbe=|uwo diff --git a/tests/fixtures/onnx_genai_workflows/tts/policies/code_history_append.onnx b/tests/fixtures/onnx_genai_workflows/tts/policies/code_history_append.onnx index 0b703d45166e2e22d6297661bfbe906c54bd5cd3..41f2f9fe6724529e403e852d55a18346e830dbbc 100644 GIT binary patch delta 192 zcmcc0ewB@vgIkC#H$N$}wAktp+eY5kjP)8^>=~KGCHX~_Lb6=a9PC2kTx$&iaDFSVj19%h&#mpo9BG(r)N z5EsNS1#JaGkbzlJTs&oy|1n7@YH;zE#fRn<7nY``R;9)p0^J}A)GP?sJlTMGDFB}F BE&c!i delta 313 zcmcc0c9or%gIkC#H$N$}wAktp`$pc^jP-F`>=~KGCHX~_Lcv^t9PC2ATuMR;1^GFd z$(4G^`6;RKQ0e%@f`Zh%6uqRxlH?37Wg$g8YSM}lb5o1Cm^fIX%(+;>JRvWaFSVisra&m1D->uzAesR%X;kN_3#k%xo`SZ5A=l)5CIu-)F5a^E d(7fWp($v(d)ObT70WLnERa{&QlQ%Fe1pqcFWJ&-4 diff --git a/tests/fixtures/onnx_genai_workflows/tts/policies/codec_layout.onnx b/tests/fixtures/onnx_genai_workflows/tts/policies/codec_layout.onnx index dfca6fee18ae6ad030cbe3c796112a7eafca77fc..31ce6b3a9b69ef63df9a7b26d8b599edb29c19b0 100644 GIT binary patch delta 100 zcmeyvyoQ;FgIkC#H$N$}wAgCqM4nCF8eHrdnZ+gfMU_IbT+$ruLgHNPLaa%NCCM3F l971eqMTxno#av7rEJ3Z6j|Ildw1!OG*P61ax&6f b9A)v&iNz)H24FdMh+I-)Npi;I^Gufj$pI1- delta 123 zcmeBUU&zkO!7ap=o1c_fT5Q$FzL7VHDZ!A7B_%&ERY;Rdor6IB{Mm(Bvmgdu_QSoNsNmFCMU$j#R1gD!o|eFm?gr+Q5Nrfs TAjHAN#=$Pc#Kkaq57Q+84z3^m diff --git a/tests/fixtures/onnx_genai_workflows/tts/policies/last_token_logits.onnx b/tests/fixtures/onnx_genai_workflows/tts/policies/last_token_logits.onnx index 4142d9050c9fdc97b12ded4b6b105e73487316e5..f674183330680e7f8e3468568778f80f8327ea0b 100644 GIT binary patch delta 115 zcmeyu`kjTBgIkC#H$N$}wAkti%SPT{#t36Bww(O*%#vauO)hl~Mj>S`b|KcJ#FFF; zE-oRC;?%;@)V$GSvi#)4q|%(kqRJ#4F7BMf;*xl%PFXH#piXgwPN3?^rcC<* DBibI7 delta 241 zcmey)@`aU`gIkC#H$N$}wAkt)>qg#S#`-!gww(O*%#vcE60RZ+Mxi_|1tGbD{G80> zO1+%K;*$81{Or`cc!(Okq{Ncs3@&9MMVu;%QwvK|^O94!RE3mrDk{rQPE0DzNi3=i dON!&-2HOL5g+G@s&=sCUxB{OAlXo!g2LSQqRMh|g diff --git a/tests/fixtures/onnx_genai_workflows/tts/policies/predictor_body_sampler.onnx b/tests/fixtures/onnx_genai_workflows/tts/policies/predictor_body_sampler.onnx index 5f78f3304d884cd9404c5d2beb36397e77ca81e9..e2685bfe5abf6f9a185f76a54debcdd4d02e72e5 100644 GIT binary patch delta 87 zcmeywGJ}POgIkC#H$N$}wAiY3BF|1+buPA?{PfI{Vj)>BX%0pqaV~Zt)}+LePx# delta 192 zcmbQi@`;6qgIkC#H$N$}wAkwYM4p}X$y{tX`RSP@#X^x>;T()Y!Cb0B$_4p3naP!U z1x2YTnaL&jMe#}bDV6cXiMa(isYQB8i6zMyTv|dJ_?4CACnqMA<|GzXCYf=umgHxr i<_YO?X>+g(sUz%A#%oCv<4i^_#mVm&B`2#fegXhhEJ1bv diff --git a/tests/fixtures/onnx_genai_workflows/tts/policies/predictor_prefill_sampler.onnx b/tests/fixtures/onnx_genai_workflows/tts/policies/predictor_prefill_sampler.onnx index b5ab9658c9b24c00658df8c476b003a4f910a928..d67d86e339afd9306caf255dc6666a3de6190f98 100644 GIT binary patch delta 87 zcmeyzGMj~mgIkC#H$N$}wAiX+BF|1+buPA?{PfI{Vj)>BX%0pqaV~Zt)}+Le(9Gb4 Un$g5KhmlKh@<&F=$!d&W0TN;>4gdfE diff --git a/tests/fixtures/onnx_genai_workflows/tts/policies/predictor_state_initializer.onnx b/tests/fixtures/onnx_genai_workflows/tts/policies/predictor_state_initializer.onnx index 153847646c60f9b86eb52e6d191dd3ebe3dfe75f..5f30409921a77135af5e7155745f39eabe49d79d 100644 GIT binary patch delta 785 zcmaD_^{9xKgIkC#H$N$}wAku=(Zr3i8?}P<>MgnW3W`$GGIMg`Q*)D2Q;LPuxl}nA zg%r8ig;Qu9hO^Yh|!6N|HjG`Z9{ z*oBl4S|x^w0c5p?GP?J-Q!nA~U}A|}iw2sD-#VJs5|W0KqC^9B+Ix?I{o zkE$a}ac~I>3Bp_q_9@6}E@qAf#w3r)afVWcBx-Wwnp`L*W++6WHjl}5MpC3Y&}}lm mF&Vy^XH14o`X*%9w9kYLoBU15u<4$ulmxCI-yCMf&kg`Y5XKDv delta 1763 zcmc(dyGjE=6o#3I$uerxkz5pn*hz%l+>AsY!9v7RaM<09I%IFznFvX<2$mrd=1GFR zU?bK(f?#WJAvjq`p;+t|S-UCD@crL^&YS-}SW%Wx+I9`vV~4MU@nm~aK2w7LW<8%= zQp>``Hb|Qx8+H_cu3#DMdagyyzA6XW)a0&@8P_?%)S;Z}799{@HFR!X!+lhbFgpV= z((Mt)Brv4}ig62IR_B~JoVpISb=E;Hw3Q5^u_7FyR+I=j<5=>sMV!)JiHQ$26ZGJ+ z=Q8Oe`BXMQpegY<7QTwK2>6-#6mQ%CxM8^MzF17Jo`cYiWdxa<)zp$zbmkb8c||DX@y_}V8NU2tV8q>&e^RSO3knk-Y6rEdks71Fe0gzX1kiZ)gAj diff --git a/tests/fixtures/onnx_genai_workflows/tts/policies/predictor_step_update.onnx b/tests/fixtures/onnx_genai_workflows/tts/policies/predictor_step_update.onnx index 306634d708b9cb702261aafc596e47aabdfa9891..5c2e73d7a71c4f373a8eaac8f1919cc9f7b06eb7 100644 GIT binary patch delta 199 zcmZ1~_?@4ZgIkC#H$N$}wAkti|3=_Q@3 z>_V(bi6zMyT%1De$@zIDsTCzr%3M4J`Nf%F#hEF^Lc&~vK()LGwM-n0Nv2%Fc|gVS nXy(ds$pF<#Ak^{-aVKj-3|7!qFig_m5(MkRWy5Ai_7p|{x=%G< delta 374 zcmey)zf_QygIkC#H$N$}wAgBr;6~m%to3PJe2FC`sd*)t`FZiViN)DMAzVQm>_UEA zDnd#H`8k=%m3jq5sVSMsCHY11#U-f)@udYRi6yCeNr@%N8C+^YDtJ{T=jWBAR+L0J zaPbu67iWU4$xJC0GT_n!+M-oYq%}+&j7j-i!g)ZW;?dk2!4(EHDTqjew1hN}-5Z7B cEvSPPv=t1KT)70nHez!*LHjmOW=mlN0BsL|U;qFB diff --git a/tests/fixtures/onnx_genai_workflows/tts/policies/setup_predictor_sampler.onnx b/tests/fixtures/onnx_genai_workflows/tts/policies/setup_predictor_sampler.onnx index 6a5ca34f02a182e41dfb043ee54e9a92c401c67f..a8bbba82c4edcee282d0dba1e7a6fabbd0cdd57d 100644 GIT binary patch delta 87 zcmeyyGLwaegIkC#H$N$}wAiX`BF|1+buPA?{PfI{Vj)>BX%0pqaV~Zt)}+Lee(0PzwO^8f$< delta 195 zcmbQq@{NUugIkC#H$N$}wAkwLM4p}Xsa$M1`RSP@#X`|sksORdpzngrC9BX%0pqaV~Zt)}+LeY-0P6r0)c^nh delta 172 zcmbQn@`{CrgIkC#H$N$}wAkwUM4p}X@my><`RSP@#X_N6!5oZ2{#?pJiUs*OnaP!U z#i=Ew1@R?`IoYX2@x_U`1v#ljdP#{T$r)T4LTdOFmE|WVCY9zS7F8yhaIu!;XQ$=~ WX>)0EunVc8nlFTD!Dc2#5k>${;5ion diff --git a/tests/fixtures/onnx_genai_workflows/tts/policies/talker_sampler.onnx b/tests/fixtures/onnx_genai_workflows/tts/policies/talker_sampler.onnx index 9ea6d838f9d735672f69e84200f381839be31248..83f3e3cdb794a651f74c3f9b7191e0fe327971f0 100644 GIT binary patch delta 88 zcmcb_GJ%DcgIkC#H$N$}wAkt|^F-bqw(4ALIr-_CCB;IrT+$qjLgHNPLaa%NCCM3F pJVIP$`N@e%r8$X3l}Qp@tR?x`sd+;DT)Z6YLYyFVn_U?B7y$_V6&(No delta 169 zcmbQha*2hPgIkC#H$N$}wAgAt%S7HC^}$?hIr-_CCB;IXT<#o7#{%ux#u=C diff --git a/tests/fixtures/onnx_genai_workflows/tts/policies/talker_state_initializer.onnx b/tests/fixtures/onnx_genai_workflows/tts/policies/talker_state_initializer.onnx index 75a589d22319597d99196ab56d2af2e5ff653506..d16e72d0909679244091d68a91161a4a10207b2e 100644 GIT binary patch delta 350 zcmX?XyU>c4gIkC#H$N$}wAiZ8YT`!Ojaqu*^_E_W;2tr9|F1&PHa@j!j?IjMOH+6qusk}(%gL4I+j0a$Y;&{}0KMW9|;pvlY#{d__^ zFe^d&lkB)Av5AO+oyd;tL@^| delta 536 zcmZ2zb=a1dgIkC#H$N$}wAgBc?Zl0;8@2SrRVumo3W`$GGIMg`Q*)D2Q;LPsxl%b8 zg%Y{cg;Wdjb25`F^(OxplMvKPN-Rmv;4%_2KoYgmgJ~{KEi6sVOHSot;^0WC=Hg2% zDM`&M$;{7-&rK}O7RuyG=U^8~#$l+1kQtJpc6tSg#U=4T6XSDI^Axldpsb`~E}nw? z;!LoqnLw8&b0q=|jRjiCjKg3(Asr-x1@vI91R0!E$2EyfL>TN}cAWk-6EZ_Q@3 z>_V(bi6zMyT%1De$@zIDsTCzr>Rdbp`Nf%F#hEF^LgHMaK(&HgOdQMz^-LU$Nv2%F qc|Z;EXlBcC$pF<$Ak-T33UNaXSI|~4Ow!>J1na|R%Vrby6h;67o;9KX delta 347 zcmeyvKSPk0gIkC#H$N$}wAiXua3k+c*7^i4zQmG})Vz|+{Ji+w#NuorKQ12*b|DWg zMIrfu{G80>O1+ZAob1%1_~Me(g80&cl*E!$y`;pF<~Z1rheyah%1xdkQhCHdK@dBsAZT)`adLjGJT zLP`bsIho0odL<>r@x>*HC8_b5d6^}di8+~7sYQB8i6zMyTcug;;JXk8YT{wBs(rHuoL5pbMi~z1`w;CG0BCC2V{4A cMrLtIeo>{636~KEyRwiTL7#5cWlv=U0Ns^nxc~qF diff --git a/tests/fixtures/onnx_genai_workflows/tts/talker/model.onnx b/tests/fixtures/onnx_genai_workflows/tts/talker/model.onnx index 299fde68df9c8a4e9740c378dac63d201f3eebb5..5dcc0471fb86dba5fd296d8834e70a112d589e62 100644 GIT binary patch delta 542 zcmezNi}C1pMiCBfA-3H7q|DM{DOLkLLp=kl>EAbsRQWL7VVb`isO{KIznkOF)6lFBsd=p5rd|~1O8_y0kesV*#29v4Kqu<-f z$swVP=92?F#U`H*kpi+BJ#!|9h28_oCRRyI_6-wdG8dXW(MKG}>_nv(2^%3#6D0(8AH6+7_rauiA>ybW*iG zFr02y+XKT*vc_ieg-UIZ-dqrzx7`A&^8qtT>VIiPdoYjuDo`_yUy01Upl Ao&W#< delta 680 zcmX^5o$=c*MiCBfA-3H7q|DM{DOLkLLp=kl6Tdc!RQWJIn%wWB%cwB9Fi@O3Gd(vy zGbP^CSW0=aqo0UGFjOEfKP43?s>BWznEWU31>+B(Mja+DG@U{+K$FBKUkGGnHj$DC znKJoJAQz*_*-T0cs9I*SR{$@g z+2)70F2?lZDXai9X_!&xLSJo>RraXgQfbRF~0mvu$WRBakyUOhVqwKuWqKF(*5< zNIxgBGPS5!&p_WXB_$pl1ZD=4|GQqByeeFZ+01}za$vgfWNr^yBg$O{5Gn$+oc}>b#N*pTZoS#=*l9*Qlwas|) z2PVnMB~jaKjJQPLCc+h&84K~lMT1L|;>}E?gkZJ+xez-{CNrump3E7eB4o&gu6(k6 zj3zKDGGksbnoSOgRRog2Xqh}Q_8HJOHPv!J{`5FCM)S@4;uc6T8$sfl6P`*Y|ESal z#!G#rEzpp?l~zFIAQAJ)8ddhd2*{|i0S4;QDlH~pxLvF=0h-BHt-@$DS-V=9(Ri{y Q(0r50dDS{VZL_Pj02x!|LI3~& diff --git a/tests/fixtures/onnx_genai_workflows/tts/talker_prefill_embedder/model.onnx b/tests/fixtures/onnx_genai_workflows/tts/talker_prefill_embedder/model.onnx index 99da132888d817c6caa93b98716b22bc23104a02..0bdbaa3e8a3442fc45a82ff4d89aaf795b580be9 100644 GIT binary patch delta 1691 zcmZuxT})eL7|wT2*OtTD(Zjkyfy(I*494jBKRsu_N?@=kW3U3PFIy9V!9iNiNqzu-T1T2WG*o=S)yB}yO^5H_d8t+S-d&t`#JCX zKJW89@42)_e!58N-POGFrTOP)mhyoeVyDz8S{E-8?o*b7Hx75M4f`d-bCe8$3FwSbc4wAPMCX+5!515tM*=8@%VZkv;W8rI$iM!lT2Ru6BgqT-l zi~y`P)WNW$7S4LBiBY`Z?IDc8!*xv{aS9U-CaQ52srk)umEj|Bjty{q^K}a-dbZPc+Ev2~Ns3s7rw}^raDDXkO%Blh5F;lzW3vfh*3*HlO z-LXnG;g?7ZMvlAiY{6=B*<4E+2xdd=mTHn#mW{N212>v-NbYTa2NjHwkoA&c>@D;{ z&T1?PCTSeh1T99VfDK|G)`R#gEA)NG-*v6)%hEt^`9UN8n_S70TR z2cf?K&KNF6jzhM)g}#=%A5vsh6=ZNn8errSVfdniwJ1?++&2-X=Q%|{zlLpFc)K{K z7-D;;rhp3krW!^z(GIXncY-80LRW$#%F1AZV04h4GEy0Y8&hVxTj5IY&(LGgz-+n^ z!MPcd) zTt6j-hxfa`EuJ|1po&#W|FrVoarM{`wsv_q0Vk7H2>fqc?TS(lERU$W;ylN>-7en2 zF>Z2Ne)5m==ZJ=Y+C~p!>Gg3xyfxZ|TC9&gjXcx+nPhpg>F!!GSefq9*D!@XA;#R{3yD!zw?F4WM34<0dwd8joYHv*V(j z$E2S8cRWs=rcW3cKR02*Q)!t|WVksYW3V@6VDLan!{EC~Kb%X2Fu~_36+TPV!au1n a6{R^Xd&F%turPWo2(P8x*z?=zw*LUFlmR9H delta 1814 zcmZWpT}&KR6!y*x>@v*qGYD*lELhROo_npfrZ21#4TM zh;0*}Ol$+^#WzdhL$zt_W=%{QjeTkp8dDk*qefEei_x?KiJF*H^xR>AMKj6d-h1x( z&N<(A?)mm(@=1aCz13>vskv8Y&YjUJ^nTi}CodL=d{qQ85-S{S^av{gETs1@22@zt zS4m8`!uD5ZXI?)wH#41QS|?2T{mt9^*}3WCc^34P_OM{y7Jm7U|BznhueQO0sCiQ7 zPUbBw$gyhq+QqSgBCa_GeB^yYw6y`W_rj%?`r?(ARw8{T%C2{X>kUhzNzqnqcdw)<_W1w0IQN>tanUBhZ^jX-sHRQni0on3rM>Pqc? zhlo@)utHs;hGwBt09MszViazA2s~*xEt=XvJdMJnP>0kmwY&){(luCaw1~0TbO4Xq z1O4!*Sr(a=_c8|~U&GQqTd-6p zq*``_qqYVeiho~hk8;%BBOQ6nT@mU=YQM65k=klI*UZb)pm8n4si@n{DZZ=Px^`N0 z+yc33D$={u5+Ks5D7y%iTL_Cy{cvwo5vYU3I>k|fhKlzyV9CO5x8zOZKh!8 zSY(uIhRfo78;XWUpSKUguc1HTWXR@QJ$O_!v2zL)8=(W(k7dE)Zm=(VyOE%>U`|neMj_CV?h+lfgVS7l$rZItzyj19w2GyrWKvS3m-@}M3OOzhWO z2s?C3&!Oq*ylHIdST3?sMyP;6yTxRc!SAwHSE1f}c5V%)19-64T34 zV!BJ*`ar}bBPcOivf_LY7}8pzBPgRmTVA3Nn6J8lQh_;Qs?~89e_eT}>9@+xiuaT)Wvf?GGSAGgV!RtLi z_YB3$qoc^t52GVE^g2_+$oHuf^CeVS_If?4gdPOi)9&p0CS~F5QpS#8ifdzy@a}jF z)8yc3@8I8p&$4j@TFY8ke*F9*92Cl|XAYf@rKaz>N{ z7i&>wdIlC5NiMd;SVo9oAQesJRMwBrZYf)x;25Fj1xY!btOESw+!B%lXt*R$v z7tnnv`FW_C2%u9(k U#1KqLi8lfZvV&AkW@VlW0M1xrx&QzG diff --git a/tests/fixtures/onnx_genai_workflows/vlm/policies/decoder_state_initializer.onnx b/tests/fixtures/onnx_genai_workflows/vlm/policies/decoder_state_initializer.onnx index 72b444b6458851dce1be15f7f1ec7b95ae60f57b..0747740a619a10328d92eb1609bcf9f926eb64e6 100644 GIT binary patch delta 561 zcmaKp%}T>S6oqO0>GZZ~$&@0+zou3Z#3H&85y6$~x=?U4%p{jKw9UlKq|jBDZru4I zN}s}a@L^2HN-IW}GkoWsbMJY&`{>xDNY{r_1c|zM@4VU9?(G@aQ7ng%W?J?!RI~{V zlBZjcr)z?1uLlmT%n5~$35Fg%xKN&d4KVELo}v}7Q&nN6dRb|4~Lr(kLzHO)j?;B+y>m}~S|#tDS+7Y(pI?)5MenI#mLP>Z7KegX`5l5+q6 delta 756 zcmb7B&q@M87;hzP`*Z~d%(`Vtl*M3D5Ya^;qGN{=j2dRfK^Jvr+nE*i5?!JP@CAB> z@O65OjvZQTHQl1{I3K_N^LF`sZJ8;Q>`P|*>g4_E)v6VWK3J|N`>rNhc7>x*3z)fr z?qCOPyRv8V01rPlH?o*f%{x#ItWPLUNmRIG*}&AJUvl8snPHF>J0u2vTflAk805( zf3IBrqOlrSobrzN2Ohbw2gtcpY0?z|8PJ|D6t3X#fdp{VKqOWsHw~A7`42jsjspqP sWRC{qe+rKs7M3>fs98)K>v>2G$hjxNH9T^ROh}rmC;>~Or&4426XYJ|_W%F@ diff --git a/tests/fixtures/onnx_genai_workflows/vlm/policies/decoder_step_update.onnx b/tests/fixtures/onnx_genai_workflows/vlm/policies/decoder_step_update.onnx index 0333f2f9bf0c4417d7b0e6fd4e2772628899dcb9..d61b374c0aba21bd49e09839e485c098162ccaa4 100644 GIT binary patch delta 302 zcmZpcyduua!7ap=o1c_fT5NStd?W8Oj(Qy~zQmG})Vz|+{Ji+w#NuorDJ}^Pb|Dci zb|KcJ#FFF;E>0o#60tDt~jf3@s`Dh t<`ox~rlwY<#+zZaK1+^^vn<}VuoUPGb0I-4ehy9{Zny(CH*qsE0svs0P8t9J delta 479 zcmca2-Ym(>!7ap=o1c_fT5R<}Vk7S}j`}1nzQmG})Vz|+{Ji+w#NuqB04_fcb|Eh= zB_V}^{G80>O1+fSqWI#H)PnfZf|SIPRK29clH?376(J=&YLfHwN>VFIqHMVM za`MwNlM{2|b5irtOEQGCxHN!vsBkIO6SgGEj*F)tzc>@@!pxLnAw4b~pk@uC9LU7M zn3T>XoCkD9JT`yUQ`4hK&Rl|En}~NVCp*F9*92Cl|XAYf@rKaz>N{ z7i&>wdIlC5NiMd;CMwnPsV9E9*H) zuoLLkl>9u1MiTTT0o_`XUy_)E#gj33eVHZ9#aS`b|KcJ#FFF; zE-oRC;?%;@)V$GSvi#)4q|%(kqRJ#4F7BMf;*xl%PFXH#piXgwPN3?^rcC<* DBibI7 delta 241 zcmey)@`aU`gIkC#H$N$}wAkt)>qg#S#`-!gww(O*%#vcE60RZ+Mxi_|1tGbD{G80> zO1+%K;*$81{Or`cc!(Okq{Ncs3@&9MMVu;%QwvK|^O94!RE3mrDk{rQPE0DzNi3=i dON!&-2HOL5g+G@s&=sCUxB{OAlXo!g2LSQqRMh|g diff --git a/tests/fixtures/onnx_genai_workflows/vlm/policies/termination.onnx b/tests/fixtures/onnx_genai_workflows/vlm/policies/termination.onnx index 3af2ca2e7abf2118f6f33cbf12082031af251140..3183438a88f3ecce329e4afc84272274af35cdc8 100644 GIT binary patch delta 727 zcmZ{i%Syvg5Qb~fCYc&*PwTxZf|rFNf0;Py3fLw za3em7&)_?_^fcl%>H3@bXa0ZYY2&@Bt9g>^xJ^6o_ui_{`s?`R4CI7+!?;fhpsNbe zKp`VdQ&{cB3MkC=(Qf%<7N#XDe>2dC!h;S<-dKQqzwDOrYtjBA6Fn3cdrL6ZG4Ie8 z%rk}UrVWxin3k|Z7Vn__+MziG!!4Jm06CSjJAs25CqcDbF6_u7SI5hshj`BYZot?* zLrO|elz#LdKUsCiSpzM^gKl8nq7ggR(j7-y>9O}uayMOap)YyFd>?6z6rrGMB>yus zzRH8nprN a*cE8dB&<~W`NaVkTQUCqpf5uc>)8)%Jf+(J delta 772 zcma)(&o2W}5XY-&y4#_`Q$NyFh`10MZ&j3t5D^DLLR?%X+xCTB?OXcZZpx)sXFZ7j z!A1NVZu~Kh#%@qegX?$ZJM;a_yv;x4(`GMA*1VGAtJVAbWBO&FIR{DYoeQbh6ik{1 z%R!cnU3gWee3jRtUUQ_a9Z%*gH^w;bLI+QGo^_MA1O^kL~JAnupQ@=BoW+74WWY^yH&V!VQ zL5haPf6Z?_56b0nNC$W<)uk`QwLs4L0O*iOh(X-!aYrE);P$0&SJ7frF;gr-rC95R z-8DPBE8V<%LMh}6?Ot~T%m6oK1v&qrhs(BXrIAO=jA)sgh5i5!Ma3_RBl;F!WF{od zqQT-2qliSXP&NzMXgPY^-O9fQ6Ack05X_aI*iJ*f}NZN_R&z|dFmpdJzO;&+cb r1V(@dWL#+zGoQ(=g-&(f(`fNFb@=0SA&pfCmjnm9kO&vM z5NlFmNpc1krx1HyX)Z`UN`*@hp*SZsFTEtAScspC7pQ>~t|3Z&vOKFeHz0A+c!7ap=o1c_fT5NTIegf>B3m9Tspc~yAuBQ_i7($nkr1@%u=L_bc7lrZKOHUIF>QTvt)OK zR0`vnvGP==_|cgcNUNz&&*;UBSuX04JtWMdTfFEq?WEBx%xGt=m1bmR2HR6@oX+Jn zWT}npD(U(JJnebXJYFd0O}z-$4v?nIxS6Zicr79XKX0Ot2Mej}m<4`T&x7N}4wBt@ z{er0>&tt9Ej-KKu746-EsYR{5;`0n^Z7@HfDR8ZWgeuxlretc=S_})I?2~J!tB1(5 z743M=JfAD+9~6pO3WOPQNE(Ltk5M;R9Gu{MHhh5L*!7`KYqu2;2fNu<*MB(HL-qt9 zDHLOPhtL>2q%=2hsGEeXLq@Z?4|zJ0FzUCtfvOaR$Pv>qT-0Jvlp*uR?8=D3ecqTT zrpX=%c4}S*jwK-QxAQV^3@MHbjA)X9>IaJm7?=rq|A}pZ|m?$4UR;g#p#y4uuH2k+OWYAm<_hC9vfMITuMt5*R+7i=@~r zjuJT+NjLyLZ&DQF7H{&V44i~zVL#oJfm32)9U7_IQivuQIE9GvvY!&wxWFfF$-pUO zo-f@3z(|`&z-dI244g(pdBAB=jSJj=TLw-ejZ?fW1E*o*$SsjZG|9j!BFY0+MKxeN zR8fvI9^ExzJX9#eZ#NYw%WO@~g$gCGxtijcll{t8q7WtUs^#;cM8$S-TC$hy4cLbg z6^FnH>Le`{&2}kQ5b?h`Smut><93RR18c*6yaT8eF}=>i#QM~!6+{`JRuE}Ewc=_L zso!1%)Cw}$mlpvwMLI^Nrid~?jd`_C?U}Vmowy6AF}wcwE~JKy1UdvWY(R}UwomO@ z_RD^U7{W3v6iY5CuL==xEX4gtiN06@IYCUCtuKL`#32pH3F6fT$VuFt3siL6l_Xa3 z;xeR$C1SN@K%GD)Lhio=q6|%IH1Nn+;1Px z#6{|n6+n%-_;)LS8nbbk8uM{LjTyO5?KwF~M;?cNZbT42$0Ox6>p5@a zDjKCWi5Ovs2iS%n-rTU+N5fU?fRmXE8g-mlj4)oR|LVV&FP$sv`mD}ll*oRvkKg}< zz55@sBh+BLQvY1DFt{0SexGEt*Y7xn3i*M+;1KrY@xkF2$v%FFUMdY2-#k~&Olj2l z5f5Aj$@z;@>$Da*e+jntavRS>N2q1Pi#VgdetX2fPbvRC*dp{WX{~7QXZ0el8UIk& zq9O%C)(s^~{&$Wi*LPt3^wP{Cjvy$8T l`_q5z=$D_TZqQLZynN5G&;LH}&wb2!8HyUNcCJ5d_#eJ~A8G&q delta 4740 zcma)=O>7%Q6vyKvb(~28t^1KCF%5}8fJzeY?D{JQP-v@OXrqD<5FxD`XKB`TY}a19 zDH5`T6bV&a+L5R@ap1%S3H3w|RX;+~rXT4?nl?aCQ$8eaz@R_T2i@E z)-npu7%)iu$&_W5jfxY{`qXqe&qOXO#C+=P@ol!2ymr9=@dsLhjig~*BkfsZ+R_nc zSN*bZ-PqoX!*6aUEoJ?1&eFrqQmytBl{+4N-FpI_qxq=>Z$J2N$4#wOS6Zq@M1m9dBiG5CA+PBl0?~G{feD?RHZ)p#C z-a|kEV!9SaRS>pvd8oUKbU0z}PMRO{w4z}6wJSp=8itTCt%)l&$QK60)27+Q34U*QB(EuV_o5e61)LZ#+sNz*d+o<4N84ZF;Fp#tc z6RC?BfryQBLwnu%fX>x70tHOZ;glkPi=Cgt`bNQSo*w2YMM34HVn;!bxiC?H$g=Yo zCJF>(7v{0PF(?Qsb_{}IePfWcrdJFicG^c^UoKm;l37cSX+b|}4zSO!cA>)&gNNht zRS6MR5+TMRK1Pf~{u*K&f_o9;szChW8b*x6iqBoch;dMHT+%qi$A}5YUqejP1y2BQ zc6I@DX#&u>jgf$|{pYko45%hvoYmWg*j2;Ja8KDJ% z7Ep!iTZHi*UQ`A^2~JBpM}QNu5Dj}T^j@K3A#20Z?5Sj zDL3d%?D~>2jMR?9vf|vJ8>K_+baDHO_HT9ZE6~M`12i^sUt)mf!yew37!t7YFhc?Y zhS-ti`V1ZuIXQ zw7%z1kLT3NM->5D?Yf1B$}TVlA$PARs<(Zp$5(a{vi`cNI!c@KzB;v{3_45hH_JMu zVmoIj4MW6j=W{}ytc>n=LJlRTby{@>XDHokzmwj}7LS(<;}e4qD#3O__~YGA8c1Wn zm+rKG2sm4*lqruxZ2qU0D`KOJndbolV1w17Xfwzh|T_jB9G3uGG)Fp9;|!hz$ZWJager;B*F;`}t~tlC=3ub*u_+xZ4h zInWwx<~cWT75TH4{}iFzX5jZO(pJ{rpEL?wChjrzkWYlLFMo-u&I{;Z!ASDI1o#9| zY5~y+BgN5iUg)f2wMST~BS!T>;|z&Q=@9?_f>v!tM5I=qgm6y$j<;*pV&`=3(_ioT zG*F-FO;^7@{JyJI`~-4t614E_NZ-q{Ncs z3@#=P#wb}Xw$g%>#FA9((vn5 LPd49Sj$;G>XrVTR delta 367 zcmX@hxr>XJgIkC#H$N$}wAgAn*GArWrg{r5_TW z2d9uS5mrVSalizP3$0r--av?hiw)>D vCN737u%KgJO1vSGpez?}S$t?-abam{YE^2ykq{3T7f?ML7sKX!<~T+GnXqSb diff --git a/tests/fixtures/onnx_genai_workflows/vlm/policies/token_to_slot.onnx b/tests/fixtures/onnx_genai_workflows/vlm/policies/token_to_slot.onnx index 2ac337e271fb50fb203de6db46cf55db8ee7532b..eb8e4b6599f5d6095617f9fe0d294eb1d098d216 100644 GIT binary patch delta 70 zcmaFByp5TMgIkC#H$N$}wAgC#M4r965?rh$`Pr#?Li}939PC1zT GenerateOptions { options } +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, @@ -127,6 +158,7 @@ fn mobius_decoder_workflow_executes() -> anyhow::Result<()> { .len(), 3 ); + assert_batched_policy_super_island(&engine); Ok(()) } @@ -244,6 +276,7 @@ fn mobius_vlm_workflow_executes_complete_image_path() -> anyhow::Result<()> { .shape(), [1, 2] ); + assert_batched_policy_super_island(&engine); Ok(()) } From 00202cd2bb168af1aba068a6f426a839c0c0f3e4 Mon Sep 17 00:00:00 2001 From: justinchuby Date: Fri, 14 Aug 2026 10:09:21 +0000 Subject: [PATCH 090/151] Declare coordinated row compaction semantics Mark every batched KV workflow as compactable regardless of paging layout, require carried row identity and mutable state coverage, and exercise same-shape row permutation plus request-epoch slot reuse without changing stable island bindings. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: justinchuby --- .../onnx_genai/workflow_metadata.py | 10 ++- .../onnx_genai/workflow_metadata_test.py | 16 ++++- .../decoder/inference_metadata.yaml | 2 +- .../speculative/inference_metadata.yaml | 2 +- .../vlm/inference_metadata.yaml | 2 +- tests/onnx_genai_workflow_conformance.rs | 68 +++++++++++++++++-- 6 files changed, 89 insertions(+), 11 deletions(-) diff --git a/src/mobius/integrations/onnx_genai/workflow_metadata.py b/src/mobius/integrations/onnx_genai/workflow_metadata.py index 99195459e..a83fb5535 100644 --- a/src/mobius/integrations/onnx_genai/workflow_metadata.py +++ b/src/mobius/integrations/onnx_genai/workflow_metadata.py @@ -522,9 +522,12 @@ def _kv_storage_contract(model: ir.Model) -> dict[str, Any]: for name in input_names for marker in ("block_table", "block_tables", "page_table", "page_tables") ) + has_cache = bool(_model_cache_pairs(model)) return { "paging": "paged" if paged else "none", - "compaction": paged, + # 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": "paged" if paged else "shared_buffer", } @@ -1630,7 +1633,10 @@ def frame_generation_nodes(prefix: str, hidden: str, logits: str) -> list[dict[s else "none" ), "allocation": "runtime", - "compaction": talker_kv["compaction"] or predictor_kv["compaction"], + # The current workflow schema cannot represent an invariant + # serving slot identity through the nested predictor loop + # without redefining the outer SSA value. + "compaction": False, "groups": { **( { diff --git a/src/mobius/integrations/onnx_genai/workflow_metadata_test.py b/src/mobius/integrations/onnx_genai/workflow_metadata_test.py index f3d2154ec..1777ddfa4 100644 --- a/src/mobius/integrations/onnx_genai/workflow_metadata_test.py +++ b/src/mobius/integrations/onnx_genai/workflow_metadata_test.py @@ -312,7 +312,21 @@ def collect_emits(node): assert media_branch["cases"]["false"]["component"] == "empty_image_features" kv_service = workflow["serving"]["kv_service"] assert kv_service["paging"] == "none" - assert kv_service["compaction"] is False + assert kv_service["compaction"] is True + carried = {item["cell"] for item in workflow["steps"][0]["carried"]} + assert { + "slot_ids", + "token", + "logits", + "generated_lengths", + "rng_counter", + "active", + "done", + "accepted_len", + "cache_lengths", + "attention_mask", + "cache_0", + } <= carried decoder_cache = kv_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. diff --git a/tests/fixtures/onnx_genai_workflows/decoder/inference_metadata.yaml b/tests/fixtures/onnx_genai_workflows/decoder/inference_metadata.yaml index b715dd625..7712409f3 100644 --- a/tests/fixtures/onnx_genai_workflows/decoder/inference_metadata.yaml +++ b/tests/fixtures/onnx_genai_workflows/decoder/inference_metadata.yaml @@ -893,7 +893,7 @@ pipeline: kv_service: paging: none allocation: runtime - compaction: false + compaction: true groups: decoder_cache: sequence_axis: 2 diff --git a/tests/fixtures/onnx_genai_workflows/speculative/inference_metadata.yaml b/tests/fixtures/onnx_genai_workflows/speculative/inference_metadata.yaml index bfc4b955e..d8db9b97f 100644 --- a/tests/fixtures/onnx_genai_workflows/speculative/inference_metadata.yaml +++ b/tests/fixtures/onnx_genai_workflows/speculative/inference_metadata.yaml @@ -892,7 +892,7 @@ pipeline: kv_service: paging: none allocation: runtime - compaction: false + compaction: true groups: verifier_cache: sequence_axis: 2 diff --git a/tests/fixtures/onnx_genai_workflows/vlm/inference_metadata.yaml b/tests/fixtures/onnx_genai_workflows/vlm/inference_metadata.yaml index 6267da169..fb62f2767 100644 --- a/tests/fixtures/onnx_genai_workflows/vlm/inference_metadata.yaml +++ b/tests/fixtures/onnx_genai_workflows/vlm/inference_metadata.yaml @@ -1046,7 +1046,7 @@ pipeline: kv_service: paging: none allocation: runtime - compaction: false + compaction: true groups: decoder_cache: sequence_axis: 2 diff --git a/tests/onnx_genai_workflow_conformance.rs b/tests/onnx_genai_workflow_conformance.rs index 2f3f5c8a2..9a6f658a5 100644 --- a/tests/onnx_genai_workflow_conformance.rs +++ b/tests/onnx_genai_workflow_conformance.rs @@ -67,6 +67,27 @@ fn decoder_batch_request( prompt_lengths: &[i64], active: &[bool], max_new_tokens: usize, +) -> anyhow::Result { + let slot_ids = (0..batch).collect::>(); + decoder_batch_request_with_slots( + input_ids, + batch, + sequence, + prompt_lengths, + active, + &slot_ids, + max_new_tokens, + ) +} + +fn decoder_batch_request_with_slots( + input_ids: &[i64], + batch: i64, + sequence: i64, + prompt_lengths: &[i64], + active: &[bool], + slot_ids: &[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)?]; @@ -74,7 +95,7 @@ fn decoder_batch_request( 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)?]; - let slot_ids = (0..batch).collect::>(); + assert_eq!(slot_ids.len(), usize::try_from(batch)?); Ok(PipelineGenerateRequest::new(GenerateRequest { prompt: GeneratePrompt::TokenIds(vec![0]), options: options(max_new_tokens), @@ -98,7 +119,7 @@ fn decoder_batch_request( .with_input("package.one_token", Value::from_slice_i64(&ones, &[batch])?) .with_input( "package.slot_ids", - Value::from_slice_i64(&slot_ids, &[batch])?, + Value::from_slice_i64(slot_ids, &[batch])?, ) .with_input( "request.eos_ids", @@ -125,7 +146,7 @@ fn decoder_batch_request( "request.min_p", Value::from_slice_f32(&floats_zero, &[batch])?, ) - .with_input("request.seed", Value::from_slice_i64(&slot_ids, &[batch])?) + .with_input("request.seed", Value::from_slice_i64(slot_ids, &[batch])?) .with_input( "request.rng_counter", Value::from_slice_i64(&zeros, &[batch])?, @@ -216,6 +237,43 @@ fn mobius_decoder_rows_match_independent_runs_and_dynamic_batch_replay() -> anyh 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_slots( + &[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() + }; + assert_eq!(compacted_row(0)?, first_tokens); + assert_eq!(compacted_row(1)?, second_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" + ); + let inactive = decoder_batch_request(&[4, 5, 6, 0], 2, 2, &[2, 1], &[true, false], 3)?; let inactive_output = engine.run_pipeline_outputs(inactive)?; let inactive_rows = engine.output_rows_for_role(&inactive_output, WorkflowOutputRole::Tokens); @@ -238,14 +296,14 @@ fn mobius_decoder_rows_match_independent_runs_and_dynamic_batch_replay() -> anyh second_tokens ); - let replay = decoder_batch_request(&[4, 5], 1, 2, &[2], &[true], 3)?; + 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()?, - first_tokens + second_tokens ); Ok(()) } From 817578080de4d48bcb8379954ec6ec63612e4892 Mon Sep 17 00:00:00 2001 From: justinchuby Date: Fri, 14 Aug 2026 11:00:49 +0000 Subject: [PATCH 091/151] Enable nested TTS row compaction Adopt ONNX GenAI's lexical slot-provenance semantics for nested loops, enable coordinated compaction for the TTS talker and predictor caches, and pin cross-repository validation to the fixed runtime. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: justinchuby --- .github/workflows/main.yml | 2 +- .../integrations/onnx_genai/codec_workflow_metadata_test.py | 2 ++ src/mobius/integrations/onnx_genai/workflow_metadata.py | 5 +---- .../onnx_genai_workflows/tts/inference_metadata.yaml | 2 +- 4 files changed, 5 insertions(+), 6 deletions(-) diff --git a/.github/workflows/main.yml b/.github/workflows/main.yml index e8c6dbdd1..1066a35c2 100644 --- a/.github/workflows/main.yml +++ b/.github/workflows/main.yml @@ -23,7 +23,7 @@ jobs: - uses: actions/checkout@v7 with: repository: justinchuby/onnx-genai - ref: f0046a8d89553bde8bc301071f5d0f09f6b7254e + ref: 8e17b988834503c0d4aa9ef776040984ac03e2d9 path: validation/onnx-genai - uses: actions/setup-python@v7 with: diff --git a/src/mobius/integrations/onnx_genai/codec_workflow_metadata_test.py b/src/mobius/integrations/onnx_genai/codec_workflow_metadata_test.py index 2b37ceb00..c39293146 100644 --- a/src/mobius/integrations/onnx_genai/codec_workflow_metadata_test.py +++ b/src/mobius/integrations/onnx_genai/codec_workflow_metadata_test.py @@ -177,6 +177,7 @@ def test_real_qwen3_tts_workflow_carries_trained_transitions_and_kv_state(): 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"]["kv_service"]["compaction"] is True assert workflow["serving"]["kv_service"]["groups"]["talker_cache"]["ports"]["talker"][ "talker_cache_0" ]["input"].startswith("past_key_values.") @@ -194,6 +195,7 @@ def test_real_qwen3_tts_workflow_carries_trained_transitions_and_kv_state(): 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/workflow_metadata.py b/src/mobius/integrations/onnx_genai/workflow_metadata.py index a83fb5535..09fbf100b 100644 --- a/src/mobius/integrations/onnx_genai/workflow_metadata.py +++ b/src/mobius/integrations/onnx_genai/workflow_metadata.py @@ -1633,10 +1633,7 @@ def frame_generation_nodes(prefix: str, hidden: str, logits: str) -> list[dict[s else "none" ), "allocation": "runtime", - # The current workflow schema cannot represent an invariant - # serving slot identity through the nested predictor loop - # without redefining the outer SSA value. - "compaction": False, + "compaction": (talker_kv["compaction"] or predictor_kv["compaction"]), "groups": { **( { diff --git a/tests/fixtures/onnx_genai_workflows/tts/inference_metadata.yaml b/tests/fixtures/onnx_genai_workflows/tts/inference_metadata.yaml index 9fc0f594a..5e579b293 100644 --- a/tests/fixtures/onnx_genai_workflows/tts/inference_metadata.yaml +++ b/tests/fixtures/onnx_genai_workflows/tts/inference_metadata.yaml @@ -1236,7 +1236,7 @@ pipeline: kv_service: paging: none allocation: runtime - compaction: false + compaction: true groups: talker_cache: sequence_axis: 2 From 3dc1f9623e07f2fec53adbce3c2f29a7d47b9d96 Mon Sep 17 00:00:00 2001 From: justinchuby Date: Fri, 14 Aug 2026 11:13:10 +0000 Subject: [PATCH 092/151] Exercise nested TTS batch compaction Require explicit TTS serving slot identities and validate heterogeneous B>1 row permutation, stable binding reuse, and request-epoch slot reuse through the nested predictor loop. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: justinchuby --- .../codec_workflow_metadata_test.py | 6 ++ .../onnx_genai/workflow_metadata.py | 5 +- .../tts/inference_metadata.yaml | 6 +- tests/onnx_genai_workflow_conformance.rs | 81 +++++++++++++++++-- 4 files changed, 87 insertions(+), 11 deletions(-) diff --git a/src/mobius/integrations/onnx_genai/codec_workflow_metadata_test.py b/src/mobius/integrations/onnx_genai/codec_workflow_metadata_test.py index c39293146..436ca93be 100644 --- a/src/mobius/integrations/onnx_genai/codec_workflow_metadata_test.py +++ b/src/mobius/integrations/onnx_genai/codec_workflow_metadata_test.py @@ -178,6 +178,12 @@ def test_real_qwen3_tts_workflow_carries_trained_transitions_and_kv_state(): assert workflow["state"]["talker_cache_0"]["service_group"] == "talker_cache" assert workflow["state"]["predictor_cache_0"]["service_group"] == "predictor_cache" assert workflow["serving"]["kv_service"]["compaction"] is True + assert workflow["inputs"]["package.slot_ids"] == { + "contract": {"dtype": "int64", "rank": 1, "shape": ["batch"]}, + "role": {"kind": "opaque"}, + "source": {"kind": "application", "name": "serving.slot_ids"}, + "required": True, + } assert workflow["serving"]["kv_service"]["groups"]["talker_cache"]["ports"]["talker"][ "talker_cache_0" ]["input"].startswith("past_key_values.") diff --git a/src/mobius/integrations/onnx_genai/workflow_metadata.py b/src/mobius/integrations/onnx_genai/workflow_metadata.py index 09fbf100b..15ff8d516 100644 --- a/src/mobius/integrations/onnx_genai/workflow_metadata.py +++ b/src/mobius/integrations/onnx_genai/workflow_metadata.py @@ -768,9 +768,8 @@ def _build_real_tts_workflow_metadata(pkg: Any, config: Any) -> dict[str, Any]: "package.slot_ids": { "contract": batch_int, "role": {"kind": "opaque"}, - "source": {"kind": "literal"}, - "required": False, - "default": 0, + "source": {"kind": "application", "name": "serving.slot_ids"}, + "required": True, }, } for iteration in range(num_groups - 2): diff --git a/tests/fixtures/onnx_genai_workflows/tts/inference_metadata.yaml b/tests/fixtures/onnx_genai_workflows/tts/inference_metadata.yaml index 5e579b293..37a9faad7 100644 --- a/tests/fixtures/onnx_genai_workflows/tts/inference_metadata.yaml +++ b/tests/fixtures/onnx_genai_workflows/tts/inference_metadata.yaml @@ -180,9 +180,9 @@ pipeline: role: kind: opaque source: - kind: literal - required: false - default: 0 + kind: application + name: serving.slot_ids + required: true package.setup_predictor_iteration_0: contract: dtype: int64 diff --git a/tests/onnx_genai_workflow_conformance.rs b/tests/onnx_genai_workflow_conformance.rs index 9a6f658a5..a941addad 100644 --- a/tests/onnx_genai_workflow_conformance.rs +++ b/tests/onnx_genai_workflow_conformance.rs @@ -388,15 +388,86 @@ fn mobius_codec_workflow_executes() -> anyhow::Result<()> { Ok(()) } +fn tts_request( + prompt_tokens: &[i64], + batch: i64, + slot_ids: &[i64], +) -> anyhow::Result { + let rows = usize::try_from(batch)?; + assert_eq!(prompt_tokens.len(), rows * 2); + assert_eq!(slot_ids.len(), rows); + 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)?, + ) + .with_input( + "package.slot_ids", + Value::from_slice_i64(slot_ids, &[batch])?, + ), + ) +} + #[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(PipelineGenerateRequest::new(GenerateRequest { - prompt: GeneratePrompt::TokenIds(vec![1, 2]), - options: options(1), - }))?; + let output = engine.run_pipeline_outputs(tts_request(&[1, 2], 1, &[0])?)?; assert_eq!(output["waveform"].shape()[..2], [1, 1]); - assert!(!output["waveform"].to_vec_f32()?.is_empty()); + 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, &[1])?)?; + let second = second_output["waveform"].to_vec_f32()?; + + let batched = engine.run_pipeline_outputs(tts_request(&[1, 2, 3, 4], 2, &[0, 1])?)?; + 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, &[1, 0])?)?; + 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, &[0])?)?; + assert_eq!(reused["waveform"].to_vec_f32()?, second); Ok(()) } From 7abe41676f88f879508c6a7f460a3b53472b35c3 Mon Sep 17 00:00:00 2001 From: justinchuby Date: Sat, 15 Aug 2026 05:44:41 +0000 Subject: [PATCH 093/151] Add generic runtime adapter artifact model Represent low-rank adapter factors independently from model-family wiring, validate exact base fingerprints and target shapes, and model heterogeneous per-request composition with compaction-safe semantic row state. Keep persistence and ONNX GenAI metadata emission gated on the final runtime schema contract. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: justinchuby --- src/mobius/__init__.py | 18 +++ src/mobius/_model_package.py | 24 +++ src/mobius/adapters.py | 285 +++++++++++++++++++++++++++++++++++ src/mobius/adapters_test.py | 202 +++++++++++++++++++++++++ 4 files changed, 529 insertions(+) create mode 100644 src/mobius/adapters.py create mode 100644 src/mobius/adapters_test.py diff --git a/src/mobius/__init__.py b/src/mobius/__init__.py index cf3237219..6b7fd9500 100644 --- a/src/mobius/__init__.py +++ b/src/mobius/__init__.py @@ -5,6 +5,12 @@ __all__ = [ "ArchitectureConfig", + "AdapterApplication", + "AdapterArtifact", + "AdapterBatchSelection", + "AdapterRowSelection", + "AdapterTarget", + "AdapterWeights", "AudioConfig", "BaseModelConfig", "CausalLMConfig", @@ -44,9 +50,11 @@ "build_from_gguf", "build_from_module", "build_from_nemo", + "compose_adapter_deltas", "components", "ep_capabilities", "ep_registry", + "fingerprint_model_weights", "generation", "get_build_dtype", "get_ep", @@ -61,6 +69,16 @@ __version__ = "0.1.0" from mobius import components, generation, models, tasks +from mobius.adapters import ( + AdapterApplication, + AdapterArtifact, + AdapterBatchSelection, + AdapterRowSelection, + AdapterTarget, + AdapterWeights, + compose_adapter_deltas, + fingerprint_model_weights, +) from mobius._build_context import build_context, ep_capabilities, get_build_dtype from mobius._builder import build_from_module from mobius._configs import ( diff --git a/src/mobius/_model_package.py b/src/mobius/_model_package.py index 89166c784..66597f6cb 100644 --- a/src/mobius/_model_package.py +++ b/src/mobius/_model_package.py @@ -34,6 +34,7 @@ import tqdm from mobius._optimizations import fold_initializers_after_weights +from mobius.adapters import AdapterArtifact from mobius.generation import PolicyComponent from mobius.integrations._weight_loading import _assign_weight @@ -53,10 +54,19 @@ def __init__( 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, ) -> None: super().__init__(models or {}) self.config = config self.policy_components = dict(policy_components or {}) + 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) @@ -195,6 +205,20 @@ def add_policy_component(self, name: str, component: PolicyComponent) -> None: 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) + self.adapter_artifacts[artifact.name] = artifact + def save_policy_components( self, directory: str, diff --git a/src/mobius/adapters.py b/src/mobius/adapters.py new file mode 100644 index 000000000..1ebf408e1 --- /dev/null +++ b/src/mobius/adapters.py @@ -0,0 +1,285 @@ +# 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", + "AdapterRowSelection", + "AdapterTarget", + "AdapterWeights", + "compose_adapter_deltas", + "fingerprint_model_weights", +] + +import dataclasses +import hashlib +import json +import math +from collections.abc import Mapping, Sequence +from typing import Any + +import numpy as np +import onnx_ir as ir + + +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()) + 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 fingerprint_model_weights(models: Mapping[str, ir.Model]) -> str: + """Return a deterministic SHA-256 fingerprint of loaded base-model weights. + + Component and initializer names are part of the digest, so an adapter cannot + silently bind to an equal-shaped parameter in a different model component. + """ + digest = hashlib.sha256() + for component_name, model in sorted(models.items()): + digest.update(component_name.encode()) + digest.update(b"\0") + for parameter_name, initializer in sorted(model.graph.initializers.items()): + if initializer.const_value is None: + raise ValueError( + f"cannot fingerprint unloaded initializer " + f"{component_name!r}/{parameter_name!r}" + ) + digest.update(parameter_name.encode()) + digest.update(b"\0") + _update_tensor_hash(digest, initializer.const_value) + digest.update(b"\0") + return f"sha256:{digest.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 AdapterWeights: + """LoRA factors for one target, computing ``B @ A * alpha / rank``.""" + + target: AdapterTarget + a: ir.Tensor + b: ir.Tensor + alpha: float + + 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): + raise ValueError("adapter alpha must be finite") + + @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, ...] + + 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 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: + """Return a 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()}" + + def validate_base(self, models: Mapping[str, ir.Model]) -> None: + """Validate fingerprint, target existence, dtype, and matrix dimensions.""" + actual_fingerprint = fingerprint_model_weights(models) + if actual_fingerprint != self.base_fingerprint: + raise ValueError( + f"adapter {self.name!r} base fingerprint mismatch: " + f"expected {self.base_fingerprint}, got {actual_fingerprint}" + ) + 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] + if 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 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}" + ) + + 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): + raise ValueError("adapter application scale must be finite") + + +@dataclasses.dataclass(frozen=True) +class AdapterRowSelection: + """Adapter composition for one stable semantic request row.""" + + row_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("an adapter may appear at most once in a row composition") + + +@dataclasses.dataclass(frozen=True) +class AdapterBatchSelection: + """Fixed-shape, compaction-safe adapter state for a heterogeneous batch.""" + + rows: tuple[AdapterRowSelection, ...] + + def __post_init__(self) -> None: + row_ids = [row.row_id for row in self.rows] + if len(row_ids) != len(set(row_ids)): + raise ValueError("adapter batch row IDs must be unique") + + def validate_catalog(self, artifacts: Mapping[str, AdapterArtifact]) -> None: + for row in self.rows: + for application in row.adapters: + if application.adapter not in artifacts: + raise ValueError( + f"row {row.row_id} selects unknown adapter {application.adapter!r}" + ) + + def compact(self, permutation: Sequence[int]) -> AdapterBatchSelection: + """Apply the same physical-row permutation used for all workflow state.""" + if sorted(permutation) != list(range(len(self.rows))): + raise ValueError("adapter compaction must be a permutation of all batch rows") + return AdapterBatchSelection(tuple(self.rows[index] for index in permutation)) + + +def compose_adapter_deltas( + row: AdapterRowSelection, + artifacts: Mapping[str, AdapterArtifact], +) -> dict[AdapterTarget, np.ndarray]: + """Compose a row's selected adapters into reference parameter updates.""" + deltas: dict[AdapterTarget, np.ndarray] = {} + for application in row.adapters: + try: + artifact = artifacts[application.adapter] + except KeyError as error: + raise ValueError( + f"row {row.row_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..5249df0fc --- /dev/null +++ b/src/mobius/adapters_test.py @@ -0,0 +1,202 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT License. + +"""Tests for generic low-rank adapter artifacts and request state.""" + +from __future__ import annotations + +import numpy as np +import onnx_ir as ir +import pytest + +from mobius import ( + AdapterApplication, + AdapterArtifact, + AdapterBatchSelection, + AdapterRowSelection, + AdapterTarget, + AdapterWeights, + ModelPackage, + compose_adapter_deltas, + fingerprint_model_weights, +) + + +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), + ) + graph = ir.Graph([], [], nodes=[], 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, +) -> 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, + ) + + +def _artifact(model: ir.Model, name: str = "style") -> AdapterArtifact: + models = {"decoder": model} + return AdapterArtifact( + name=name, + base_fingerprint=fingerprint_model_weights(models), + weights=(_weights(),), + ) + + +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}) + + +@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 = AdapterRowSelection( + 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( + ( + AdapterRowSelection(row_id=100, request_epoch=4), + AdapterRowSelection( + row_id=101, + request_epoch=7, + adapters=(AdapterApplication("style", 0.5),), + ), + AdapterRowSelection( + row_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.rows[0].adapters == () + assert [item.adapter for item in batch.rows[2].adapters] == ["style", "speaker"] + + +def test_compaction_preserves_semantic_rows_and_slot_reuse_uses_epoch() -> None: + original = AdapterBatchSelection( + ( + AdapterRowSelection(100, 1, (AdapterApplication("style"),)), + AdapterRowSelection(101, 5, (AdapterApplication("speaker", 0.25),)), + ) + ) + compacted = original.compact([1, 0]) + assert [row.row_id for row in compacted.rows] == [101, 100] + assert compacted.rows[0].request_epoch == 5 + assert compacted.compact([1, 0]) == original + + reused_slot = AdapterRowSelection( + row_id=200, + request_epoch=2, + adapters=(AdapterApplication("speaker"),), + ) + assert reused_slot != original.rows[0] + + +def test_model_package_catalog_validates_and_rejects_duplicates() -> None: + model = _model() + artifact = _artifact(model) + package = ModelPackage({"decoder": model}) + package.add_adapter_artifact(artifact) + assert package.adapter_artifacts["style"] is artifact + + 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_selection_rejects_unknown_adapter_and_invalid_permutation() -> None: + batch = AdapterBatchSelection( + (AdapterRowSelection(100, 0, (AdapterApplication("missing"),)),) + ) + with pytest.raises(ValueError, match="unknown adapter"): + batch.validate_catalog({}) + with pytest.raises(ValueError, match="permutation"): + batch.compact([1]) From 33c99291b0553b6456c95398b8f69a4745658299 Mon Sep 17 00:00:00 2001 From: justinchuby Date: Sat, 15 Aug 2026 05:53:50 +0000 Subject: [PATCH 094/151] Fix adapter framework lint diagnostics Sort the public adapter imports and use a property-style checksum docstring so the repository's current Ruff rules accept the new framework. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: justinchuby --- src/mobius/__init__.py | 20 ++++++++++---------- src/mobius/adapters.py | 2 +- 2 files changed, 11 insertions(+), 11 deletions(-) diff --git a/src/mobius/__init__.py b/src/mobius/__init__.py index 6b7fd9500..470749ceb 100644 --- a/src/mobius/__init__.py +++ b/src/mobius/__init__.py @@ -69,16 +69,6 @@ __version__ = "0.1.0" from mobius import components, generation, models, tasks -from mobius.adapters import ( - AdapterApplication, - AdapterArtifact, - AdapterBatchSelection, - AdapterRowSelection, - AdapterTarget, - AdapterWeights, - compose_adapter_deltas, - fingerprint_model_weights, -) from mobius._build_context import build_context, ep_capabilities, get_build_dtype from mobius._builder import build_from_module from mobius._configs import ( @@ -116,6 +106,16 @@ ModelRegistry, registry, ) +from mobius.adapters import ( + AdapterApplication, + AdapterArtifact, + AdapterBatchSelection, + AdapterRowSelection, + AdapterTarget, + 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/adapters.py b/src/mobius/adapters.py index 1ebf408e1..a761ae0cb 100644 --- a/src/mobius/adapters.py +++ b/src/mobius/adapters.py @@ -143,7 +143,7 @@ def __post_init__(self) -> None: @property def checksum(self) -> str: - """Return a deterministic checksum covering bindings, metadata, and tensors.""" + """Deterministic checksum covering bindings, metadata, and tensors.""" digest = hashlib.sha256() digest.update(self.name.encode()) digest.update(b"\0") From 8f1efca6189cb501f3ce782611c72c549c62ef07 Mon Sep 17 00:00:00 2001 From: justinchuby Date: Sat, 15 Aug 2026 06:12:07 +0000 Subject: [PATCH 095/151] Align adapter groundwork with native LoRA contracts Add an authoritative graph target manifest, PEFT safetensors ingestion with rank and alpha patterns, optional provenance-preserving .onnx_adapter declarations, aligned N-adapter catalogs, and paged-lifecycle reference accounting. Cover PR #318 and #374 migration semantics without freezing the pending ONNX GenAI #828 schema field names. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: justinchuby --- src/mobius/__init__.py | 11 +++ src/mobius/_model_package.py | 20 +++- src/mobius/adapter_io.py | 151 +++++++++++++++++++++++++++++ src/mobius/adapters.py | 170 ++++++++++++++++++++++++++++++++- src/mobius/adapters_test.py | 179 ++++++++++++++++++++++++++++++++++- 5 files changed, 527 insertions(+), 4 deletions(-) create mode 100644 src/mobius/adapter_io.py diff --git a/src/mobius/__init__.py b/src/mobius/__init__.py index 470749ceb..d4233e39a 100644 --- a/src/mobius/__init__.py +++ b/src/mobius/__init__.py @@ -9,7 +9,11 @@ "AdapterArtifact", "AdapterBatchSelection", "AdapterRowSelection", + "AdapterSource", "AdapterTarget", + "AdapterTargetDescriptor", + "AdapterTargetManifest", + "AdapterTargetSlice", "AdapterWeights", "AudioConfig", "BaseModelConfig", @@ -44,6 +48,7 @@ "WorldModelTask", "YolosConfig", "apply_weights", + "adapter_source_from_onnx_adapter", "build", "build_context", "build_diffusers_pipeline", @@ -55,6 +60,7 @@ "ep_capabilities", "ep_registry", "fingerprint_model_weights", + "load_peft_adapter", "generation", "get_build_dtype", "get_ep", @@ -106,12 +112,17 @@ ModelRegistry, registry, ) +from mobius.adapter_io import adapter_source_from_onnx_adapter, load_peft_adapter from mobius.adapters import ( AdapterApplication, AdapterArtifact, AdapterBatchSelection, AdapterRowSelection, + AdapterSource, AdapterTarget, + AdapterTargetDescriptor, + AdapterTargetManifest, + AdapterTargetSlice, AdapterWeights, compose_adapter_deltas, fingerprint_model_weights, diff --git a/src/mobius/_model_package.py b/src/mobius/_model_package.py index 66597f6cb..b81bf6fba 100644 --- a/src/mobius/_model_package.py +++ b/src/mobius/_model_package.py @@ -34,7 +34,7 @@ import tqdm from mobius._optimizations import fold_initializers_after_weights -from mobius.adapters import AdapterArtifact +from mobius.adapters import AdapterArtifact, AdapterTargetManifest from mobius.generation import PolicyComponent from mobius.integrations._weight_loading import _assign_weight @@ -55,10 +55,14 @@ def __init__( config: object | None = None, policy_components: dict[str, PolicyComponent] | None = None, adapter_artifacts: dict[str, AdapterArtifact] | None = None, + adapter_target_manifest: AdapterTargetManifest | None = None, ) -> None: super().__init__(models or {}) self.config = config self.policy_components = dict(policy_components or {}) + self.adapter_target_manifest = adapter_target_manifest + 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: @@ -217,6 +221,20 @@ def add_adapter_artifact( raise ValueError(f"adapter artifact {artifact.name!r} is already attached") if validate_base: artifact.validate_base(self.data) + if self.adapter_target_manifest is not None: + 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))}" + ) + if self.adapter_artifacts: + expected = next(iter(self.adapter_artifacts.values())).target_bindings + if artifact.target_bindings != expected: + raise ValueError( + f"adapter artifact {artifact.name!r} target set does not align " + "with the existing N-adapter catalog" + ) self.adapter_artifacts[artifact.name] = artifact def save_policy_components( diff --git a/src/mobius/adapter_io.py b/src/mobius/adapter_io.py new file mode 100644 index 000000000..ce0ce9165 --- /dev/null +++ b/src/mobius/adapter_io.py @@ -0,0 +1,151 @@ +# 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]) -> AdapterTarget: + if module_key in targets: + return 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][1] + + +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}" + ) + loaded_weights.append( + AdapterWeights( + _resolve_target(module_key, target_bindings), + ir.tensor(a), + ir.tensor(b), + alpha, + ) + ) + + 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 index a761ae0cb..c43958861 100644 --- a/src/mobius/adapters.py +++ b/src/mobius/adapters.py @@ -10,7 +10,11 @@ "AdapterArtifact", "AdapterBatchSelection", "AdapterRowSelection", + "AdapterSource", "AdapterTarget", + "AdapterTargetDescriptor", + "AdapterTargetManifest", + "AdapterTargetSlice", "AdapterWeights", "compose_adapter_deltas", "fingerprint_model_weights", @@ -21,7 +25,7 @@ import json import math from collections.abc import Mapping, Sequence -from typing import Any +from typing import Any, Literal import numpy as np import onnx_ir as ir @@ -86,6 +90,147 @@ def __post_init__(self) -> None: 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, ...] = () + + 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 zip(ordered, ordered[1:]): + 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") + + +@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: + actual_fingerprint = fingerprint_model_weights(models) + if actual_fingerprint != self.base_fingerprint: + raise ValueError( + "adapter target manifest base fingerprint mismatch: " + f"expected {self.base_fingerprint}, got {actual_fingerprint}" + ) + 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 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}" + ) + + @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 + + def __post_init__(self) -> None: + 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``.""" @@ -130,6 +275,9 @@ class AdapterArtifact: name: str base_fingerprint: str weights: tuple[AdapterWeights, ...] + source: AdapterSource = dataclasses.field( + default_factory=lambda: AdapterSource("in_memory") + ) def __post_init__(self) -> None: _validate_identifier(self.name, "adapter name") @@ -163,6 +311,13 @@ def checksum(self) -> str: _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]) -> None: """Validate fingerprint, target existence, dtype, and matrix dimensions.""" actual_fingerprint = fingerprint_model_weights(models) @@ -171,6 +326,7 @@ def validate_base(self, models: Mapping[str, ir.Model]) -> None: f"adapter {self.name!r} base fingerprint mismatch: " f"expected {self.base_fingerprint}, got {actual_fingerprint}" ) + for weight in self.weights: model = models.get(weight.target.component) if model is None: @@ -199,6 +355,11 @@ def validate_base(self, models: Mapping[str, ir.Model]) -> None: f"but adapter factors have dtype {weight.dtype.name}" ) + @property + def target_bindings(self) -> frozenset[AdapterTarget]: + """Exact base parameters modified by this artifact.""" + return frozenset(weight.target for weight in self.weights) + def validate_checksum(self, expected: str) -> None: """Reject corrupted or substituted adapter tensor data.""" if self.checksum != expected: @@ -262,6 +423,13 @@ def compact(self, permutation: Sequence[int]) -> AdapterBatchSelection: raise ValueError("adapter compaction must be a permutation of all batch rows") return AdapterBatchSelection(tuple(self.rows[index] for index in permutation)) + @property + def referenced_adapters(self) -> frozenset[str]: + """Live adapter set that a paged runtime must pin against eviction.""" + return frozenset( + application.adapter for row in self.rows for application in row.adapters + ) + def compose_adapter_deltas( row: AdapterRowSelection, diff --git a/src/mobius/adapters_test.py b/src/mobius/adapters_test.py index 5249df0fc..8c5dae6b3 100644 --- a/src/mobius/adapters_test.py +++ b/src/mobius/adapters_test.py @@ -5,20 +5,32 @@ from __future__ import annotations +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 save_file from mobius import ( AdapterApplication, AdapterArtifact, AdapterBatchSelection, AdapterRowSelection, + AdapterSource, AdapterTarget, + AdapterTargetDescriptor, + AdapterTargetManifest, + AdapterTargetSlice, AdapterWeights, ModelPackage, + adapter_source_from_onnx_adapter, compose_adapter_deltas, fingerprint_model_weights, + load_peft_adapter, ) @@ -30,7 +42,21 @@ def _model(weight: np.ndarray | None = None) -> ir.Model: type=ir.TensorType(ir.DataType.FLOAT), shape=ir.Shape(values.shape), ) - graph = ir.Graph([], [], nodes=[], initializers=[initializer], name="adapter_test_model") + 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) @@ -61,6 +87,33 @@ def _artifact(model: ir.Model, name: str = "style") -> AdapterArtifact: ) +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),), + ) + ] + 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)) + + def test_artifact_validates_base_target_shape_dtype_and_checksum() -> None: model = _model() artifact = _artifact(model) @@ -74,6 +127,29 @@ def test_artifact_validates_base_target_shape_dtype_and_checksum() -> None: artifact.validate_base({"decoder": changed}) +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"), [ @@ -169,6 +245,7 @@ def test_compaction_preserves_semantic_rows_and_slot_reuse_uses_epoch() -> None: assert [row.row_id for row in compacted.rows] == [101, 100] assert compacted.rows[0].request_epoch == 5 assert compacted.compact([1, 0]) == original + assert compacted.referenced_adapters == {"style", "speaker"} reused_slot = AdapterRowSelection( row_id=200, @@ -181,9 +258,10 @@ def test_compaction_preserves_semantic_rows_and_slot_reuse_uses_epoch() -> None: def test_model_package_catalog_validates_and_rejects_duplicates() -> None: model = _model() artifact = _artifact(model) - package = ModelPackage({"decoder": model}) + package = ModelPackage({"decoder": model}, adapter_target_manifest=_manifest(model)) 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) @@ -192,6 +270,103 @@ def test_model_package_catalog_validates_and_rejects_duplicates() -> None: artifact.validate_checksum("sha256:" + "0" * 64) +def test_model_package_requires_n_adapter_target_alignment() -> 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) + package = ModelPackage( + {"decoder": model}, adapter_target_manifest=_manifest(model, include_second=True) + ) + package.add_adapter_artifact(_artifact(model, "style")) + second = AdapterArtifact( + "speaker", + fingerprint_model_weights({"decoder": model}), + (_weights(parameter="other.weight"),), + ) + with pytest.raises(ValueError, match="does not align"): + package.add_adapter_artifact(second) + + +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() + artifact = load_peft_adapter( + directory, + name="peft-style", + base_fingerprint=fingerprint_model_weights({"decoder": model}), + target_bindings={ + "layers.0.self_attn.q_proj": AdapterTarget("decoder", "projection.weight") + }, + ) + artifact.validate_base({"decoder": model}) + assert artifact.weights[0].rank == 2 + assert artifact.weights[0].alpha == 6.0 + assert artifact.source.format == "peft_safetensors" + assert artifact.source.base_model == "synthetic/base" + assert artifact.source.revision == "producer-fixture" + np.testing.assert_allclose( + artifact.weights[0].delta(), b @ a * 3.0, rtol=1e-6, atol=1e-6 + ) + 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_adapter_source_rejects_unprovenanced_external_artifact() -> None: + with pytest.raises(ValueError, match="requires a path"): + AdapterSource("onnx_adapter") + + def test_selection_rejects_unknown_adapter_and_invalid_permutation() -> None: batch = AdapterBatchSelection( (AdapterRowSelection(100, 0, (AdapterApplication("missing"),)),) From 20f8df10da77251acf05d7cb5973403b688a1856 Mon Sep 17 00:00:00 2001 From: justinchuby Date: Sat, 15 Aug 2026 06:16:36 +0000 Subject: [PATCH 096/151] Validate adapter target parameter consumption Require each authoritative manifest node to consume the declared base parameter, preventing a same-named but semantically stale node/value declaration from passing producer validation. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: justinchuby --- src/mobius/adapters.py | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/src/mobius/adapters.py b/src/mobius/adapters.py index c43958861..37a21cd37 100644 --- a/src/mobius/adapters.py +++ b/src/mobius/adapters.py @@ -200,6 +200,11 @@ def validate(self, models: Mapping[str, ir.Model]) -> None: 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 }: From 7fdd437c39a2c2b9be9ccfea44d5645c90401194 Mon Sep 17 00:00:00 2001 From: justinchuby Date: Sat, 15 Aug 2026 07:10:17 +0000 Subject: [PATCH 097/151] Emit generic parameter adapter workflow artifacts Serialize exact ONNX GenAI adapter catalogs and checksummed portable bundles from ModelPackage, preserve PEFT and optional native adapter provenance, and emit row identity, request epoch, cache, planning, and capability contracts without model-family discovery. Add an executable heterogeneous adapter fixture covering zero/one/ordered composition, inactive rows, compaction, replay, and slot reuse against ONNX GenAI 8549e425. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: justinchuby --- .github/workflows/main.yml | 2 +- src/mobius/__init__.py | 2 + src/mobius/_model_package.py | 161 +++++++++++- src/mobius/adapters.py | 47 +++- src/mobius/adapters_test.py | 184 ++++++++++++++ .../integrations/onnx_genai/auto_export.py | 2 + .../onnx_genai/inference_metadata.py | 88 +++++++ .../onnx_genai/workflow_metadata.py | 8 + tests/fixtures/onnx_genai_workflows/README.md | 6 +- .../adapter/adapters/blue.json | 1 + .../adapter/adapters/red.json | 1 + .../adapter/inference_metadata.yaml | 171 +++++++++++++ .../onnx_genai_workflows/adapter/model.onnx | Bin 0 -> 356 bytes ...generate_onnx_genai_validation_packages.py | 229 +++++++++++++++++- tests/onnx_genai_workflow_conformance.rs | 91 ++++++- 15 files changed, 981 insertions(+), 12 deletions(-) create mode 100644 tests/fixtures/onnx_genai_workflows/adapter/adapters/blue.json create mode 100644 tests/fixtures/onnx_genai_workflows/adapter/adapters/red.json create mode 100644 tests/fixtures/onnx_genai_workflows/adapter/inference_metadata.yaml create mode 100644 tests/fixtures/onnx_genai_workflows/adapter/model.onnx diff --git a/.github/workflows/main.yml b/.github/workflows/main.yml index 1066a35c2..1eeed08d9 100644 --- a/.github/workflows/main.yml +++ b/.github/workflows/main.yml @@ -23,7 +23,7 @@ jobs: - uses: actions/checkout@v7 with: repository: justinchuby/onnx-genai - ref: 8e17b988834503c0d4aa9ef776040984ac03e2d9 + ref: 8549e42529eda6b9a135f1f5efd93416dd5d63aa path: validation/onnx-genai - uses: actions/setup-python@v7 with: diff --git a/src/mobius/__init__.py b/src/mobius/__init__.py index d4233e39a..322f88bce 100644 --- a/src/mobius/__init__.py +++ b/src/mobius/__init__.py @@ -9,6 +9,7 @@ "AdapterArtifact", "AdapterBatchSelection", "AdapterRowSelection", + "AdapterServiceOptions", "AdapterSource", "AdapterTarget", "AdapterTargetDescriptor", @@ -118,6 +119,7 @@ AdapterArtifact, AdapterBatchSelection, AdapterRowSelection, + AdapterServiceOptions, AdapterSource, AdapterTarget, AdapterTargetDescriptor, diff --git a/src/mobius/_model_package.py b/src/mobius/_model_package.py index b81bf6fba..cefc41b97 100644 --- a/src/mobius/_model_package.py +++ b/src/mobius/_model_package.py @@ -20,9 +20,12 @@ __all__ = ["ModelPackage"] +import hashlib import inspect +import json import logging import os +import shutil import threading from collections import UserDict from collections.abc import Callable, Iterator @@ -34,13 +37,49 @@ import tqdm from mobius._optimizations import fold_initializers_after_weights -from mobius.adapters import AdapterArtifact, AdapterTargetManifest +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) -> tuple[str, 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, "ort_genai" + if artifact.source.format == "peft_safetensors": + path = artifact.source.path + if os.path.isdir(path): + path = os.path.join(path, "adapter_model.safetensors") + return path, "safetensors" + 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. @@ -56,11 +95,13 @@ def __init__( 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] = {} @@ -97,6 +138,7 @@ def save( 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. @@ -150,6 +192,8 @@ def save( 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 @@ -202,6 +246,8 @@ def save( 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.""" @@ -237,6 +283,119 @@ def add_adapter_artifact( ) 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 + } + catalog: dict[str, dict[str, object]] = {} + identities: set[tuple[str, str]] = set() + for alias, artifact in 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) + ranks = {weight.rank for weight in artifact.weights} + alphas = {weight.alpha for weight in artifact.weights} + dtypes = {weight.dtype for weight in artifact.weights} + if len(ranks) != 1 or len(alphas) != 1 or len(dtypes) != 1: + raise ValueError( + f"adapter {alias!r} has heterogeneous target rank/alpha/dtype, " + "which the ONNX GenAI artifact contract cannot represent" + ) + rank = ranks.pop() + alpha = alphas.pop() + dtype = _adapter_dtype_name(dtypes.pop()) + if ( + self.adapter_service_options.preserve_source_format + and artifact.source.format == "onnx_adapter" + and alpha != rank + ): + raise ValueError( + f"adapter {alias!r} imports .onnx_adapter weights with baked scale; " + "alpha must equal rank" + ) + targets: list[dict[str, object]] = [] + portable_targets: dict[str, dict[str, list[float]]] = {} + for weight in sorted( + artifact.weights, + key=lambda item: (item.target.component, item.target.parameter), + ): + descriptor = descriptors[weight.target] + weight_key = descriptor.semantic_name + targets.append( + { + "component": weight.target.component, + "parameter": weight.target.parameter, + "weight_key": weight_key, + "input_features": descriptor.input_size, + "output_features": descriptor.output_size, + } + ) + portable_targets[weight_key] = { + "a": weight.a.numpy().reshape(-1).astype("float32").tolist(), + "b": weight.b.numpy().reshape(-1).astype("float32").tolist(), + } + + if self.adapter_service_options.preserve_source_format: + source_path, weight_format = _adapter_source_file(artifact) + extension = ".onnx_adapter" if weight_format == "ort_genai" else ".safetensors" + relative_location = f"adapters/{alias}{extension}" + destination = os.path.join(directory, relative_location) + shutil.copyfile(source_path, destination) + with open(destination, "rb") as handle: + payload = handle.read() + else: + relative_location = f"adapters/{alias}.json" + destination = os.path.join(directory, relative_location) + payload = json.dumps( + {"targets": portable_targets}, + sort_keys=True, + separators=(",", ":"), + allow_nan=False, + ).encode("utf-8") + with open(destination, "wb") as handle: + handle.write(payload) + weight_format = "json" + provenance_parts = [artifact.source.format] + if artifact.source.base_model: + provenance_parts.append(f"base={artifact.source.base_model}") + if artifact.source.revision: + provenance_parts.append(f"revision={artifact.source.revision}") + if artifact.source.checksum: + provenance_parts.append(f"source_{artifact.source.checksum}") + catalog[alias] = { + "identity": artifact.stable_identity, + "version": artifact.version, + "base_model_fingerprint": artifact.base_fingerprint, + "rank": rank, + "alpha": alpha, + "dtype": dtype, + "provenance": ";".join(provenance_parts), + "weights": [ + { + "location": relative_location, + "sha256": hashlib.sha256(payload).hexdigest(), + "format": weight_format, + } + ], + "targets": targets, + } + return catalog + def save_policy_components( self, directory: str, diff --git a/src/mobius/adapters.py b/src/mobius/adapters.py index 37a21cd37..96b2ed866 100644 --- a/src/mobius/adapters.py +++ b/src/mobius/adapters.py @@ -10,6 +10,7 @@ "AdapterArtifact", "AdapterBatchSelection", "AdapterRowSelection", + "AdapterServiceOptions", "AdapterSource", "AdapterTarget", "AdapterTargetDescriptor", @@ -257,8 +258,8 @@ def __post_init__(self) -> None: ) 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): - raise ValueError("adapter alpha must be finite") + if not math.isfinite(self.alpha) or self.alpha <= 0.0: + raise ValueError("adapter alpha must be finite and greater than zero") @property def rank(self) -> int: @@ -283,11 +284,17 @@ class AdapterArtifact: 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] @@ -365,6 +372,11 @@ 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: @@ -383,8 +395,35 @@ class AdapterApplication: def __post_init__(self) -> None: _validate_identifier(self.adapter, "adapter application name") - if not math.isfinite(self.scale): - raise ValueError("adapter application scale must be finite") + 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.""" + + row_ids: str | None = None + request_epochs: str | None = None + active: str | None = None + application_capability: str = "onnx-genai.adapters" + 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 + + 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.portable_fallback and self.preserve_source_format: + raise ValueError( + "portable fallback requires portable JSON artifacts; " + "source-format preservation requires a native adapter capability" + ) @dataclasses.dataclass(frozen=True) diff --git a/src/mobius/adapters_test.py b/src/mobius/adapters_test.py index 8c5dae6b3..53b21b4d7 100644 --- a/src/mobius/adapters_test.py +++ b/src/mobius/adapters_test.py @@ -6,6 +6,7 @@ from __future__ import annotations import json +import hashlib import shutil import uuid from pathlib import Path @@ -20,6 +21,7 @@ AdapterArtifact, AdapterBatchSelection, AdapterRowSelection, + AdapterServiceOptions, AdapterSource, AdapterTarget, AdapterTargetDescriptor, @@ -32,6 +34,9 @@ fingerprint_model_weights, load_peft_adapter, ) +from mobius.integrations.onnx_genai.inference_metadata import ( + add_adapter_service_to_workflow, +) def _model(weight: np.ndarray | None = None) -> ir.Model: @@ -362,11 +367,190 @@ def test_onnx_adapter_migration_source_is_optional_and_checksummed() -> None: 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() + package = ModelPackage( + {"decoder": model}, + adapter_target_manifest=_manifest(model), + adapter_service_options=AdapterServiceOptions( + portable_fallback=False, + preserve_source_format=True, + ), + ) + package.add_adapter_artifact( + AdapterArtifact( + "style", + fingerprint_model_weights({"decoder": model}), + (_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["location"] == "adapters/style.onnx_adapter" + 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() + package = ModelPackage( + {"decoder": model}, + adapter_target_manifest=_manifest(model), + adapter_service_options=AdapterServiceOptions( + row_ids="request.row_ids", + active="request.active", + cache_max_entries=2, + ), + ) + package.add_adapter_artifact( + AdapterArtifact( + "red", + fingerprint_model_weights({"decoder": model}), + (_weights(alpha=2.0),), + identity="style-red", + version="2026.08", + ) + ) + metadata = { + "pipeline": { + "workflow": { + "manifest": {"capabilities": []}, + "inputs": { + "request.row_ids": { + "contract": { + "dtype": "int64", + "rank": 1, + "shape": ["batch"], + } + }, + "request.active": { + "contract": { + "dtype": "bool", + "rank": 1, + "shape": ["batch"], + } + }, + }, + "components": {"decoder": {"implementation": {"kind": "binding"}}}, + "steps": [], + } + } + } + add_adapter_service_to_workflow(metadata, package, str(directory)) + service = metadata["pipeline"]["workflow"]["adapters"] + assert service["base_model_fingerprint"].startswith("sha256:") + assert service["row_ids"] == "request.row_ids" + assert service["request_epochs"] == "request.request_epochs" + assert service["active"] == "request.active" + assert service["application_capability"] == "onnx-genai.adapters" + 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 metadata["pipeline"]["workflow"]["inputs"]["request.request_epochs"] == { + "contract": {"dtype": "int64", "rank": 1, "shape": ["batch"]}, + "role": { + "kind": "runtime", + "version": "1.0", + "role": "request_epochs", + }, + "source": {"kind": "request"}, + } + artifact = service["artifacts"]["red"] + assert artifact["identity"] == "style-red" + assert artifact["version"] == "2026.08" + assert artifact["rank"] == 2 + assert artifact["alpha"] == 2.0 + assert artifact["dtype"] == "float32" + assert artifact["targets"] == [ + { + "component": "decoder", + "parameter": "projection.weight", + "weight_key": "layers.0.self_attn.q_proj", + "input_features": 4, + "output_features": 3, + } + ] + weight = artifact["weights"][0] + payload = (directory / weight["location"]).read_bytes() + assert weight["format"] == "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_wire_contract_rejects_heterogeneous_target_rank() -> 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) + package = ModelPackage( + {"decoder": model}, + adapter_target_manifest=_manifest(model, include_second=True), + ) + package.add_adapter_artifact( + AdapterArtifact( + "mixed-rank", + fingerprint_model_weights({"decoder": model}), + ( + _weights(), + _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: + with pytest.raises(ValueError, match="heterogeneous target rank"): + 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_unknown_adapter_and_invalid_permutation() -> None: batch = AdapterBatchSelection( (AdapterRowSelection(100, 0, (AdapterApplication("missing"),)),) diff --git a/src/mobius/integrations/onnx_genai/auto_export.py b/src/mobius/integrations/onnx_genai/auto_export.py index 37152ac15..45a29ab73 100644 --- a/src/mobius/integrations/onnx_genai/auto_export.py +++ b/src/mobius/integrations/onnx_genai/auto_export.py @@ -26,6 +26,7 @@ from mobius.integrations.onnx_genai.inference_metadata import ( SchedulerConfig, add_explicit_package_io, + add_adapter_service_to_workflow, add_policy_components_to_workflow, load_diffusers_scheduler_config, write_speech_to_text_pipeline_metadata, @@ -108,6 +109,7 @@ def _add_explicit_io_to_file(path: str, pkg: Any, config: Any) -> None: metadata = yaml.safe_load(handle) add_explicit_package_io(metadata, pkg, config) add_policy_components_to_workflow(metadata, pkg) + add_adapter_service_to_workflow(metadata, pkg, os.path.dirname(path)) with open(path, "w", encoding="utf-8") as handle: yaml.safe_dump(metadata, handle, sort_keys=False) diff --git a/src/mobius/integrations/onnx_genai/inference_metadata.py b/src/mobius/integrations/onnx_genai/inference_metadata.py index 995b82ccc..5f223fcee 100644 --- a/src/mobius/integrations/onnx_genai/inference_metadata.py +++ b/src/mobius/integrations/onnx_genai/inference_metadata.py @@ -1482,6 +1482,94 @@ def tensor_contract(value: Any) -> dict[str, Any]: return metadata +def add_adapter_service_to_workflow( + 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 = metadata.get("pipeline", {}).get("workflow") + if not isinstance(workflow, dict): + raise ValueError("parameter adapters require pipeline.workflow metadata") + 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.get("inputs", {}) + + def compatible_input(name: str, *, dtype: str) -> bool: + declaration = inputs.get(name) + if not isinstance(declaration, dict): + return False + contract = declaration.get("contract", {}) + return ( + contract.get("dtype") == dtype + and contract.get("rank") == 1 + and contract.get("shape") == ["batch"] + ) + + row_ids = options.row_ids + if row_ids is None: + row_ids = next( + ( + candidate + for candidate in ("package.slot_ids", "request.row_ids") + if compatible_input(candidate, dtype="int64") + ), + None, + ) + if row_ids is None or not compatible_input(row_ids, dtype="int64"): + raise ValueError("adapter row_ids must reference an int64[batch] workflow input") + request_epochs = options.request_epochs or "request.request_epochs" + if request_epochs not in inputs: + inputs[request_epochs] = { + "contract": {"dtype": "int64", "rank": 1, "shape": ["batch"]}, + "role": { + "kind": "runtime", + "version": "1.0", + "role": "request_epochs", + }, + "source": {"kind": "request"}, + } + if not compatible_input(request_epochs, dtype="int64"): + raise ValueError( + "adapter request_epochs must reference an int64[batch] workflow input" + ) + active = options.active + if active is None and compatible_input("package.active", dtype="bool"): + active = "package.active" + if active is not None and not compatible_input(active, dtype="bool"): + raise ValueError("adapter active must reference a bool[batch] workflow input") + + catalog = pkg.save_adapter_artifacts(output_dir) + workflow["adapters"] = { + "base_model_fingerprint": manifest.base_fingerprint, + "row_ids": row_ids, + "request_epochs": request_epochs, + **({"active": active} if active is not None else {}), + "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, + } + capabilities = workflow.setdefault("manifest", {}).setdefault("capabilities", []) + for capability in ("parameter_adapters", "heterogeneous_adapter_batching"): + if capability not in capabilities: + capabilities.append(capability) + return metadata + + def _topological_order( names: Iterable[str], edges: list[dict[str, Any]], diff --git a/src/mobius/integrations/onnx_genai/workflow_metadata.py b/src/mobius/integrations/onnx_genai/workflow_metadata.py index 15ff8d516..22627b104 100644 --- a/src/mobius/integrations/onnx_genai/workflow_metadata.py +++ b/src/mobius/integrations/onnx_genai/workflow_metadata.py @@ -46,6 +46,7 @@ from mobius.integrations.onnx_genai.inference_metadata import ( _port, _shape_metadata, + add_adapter_service_to_workflow, add_policy_components_to_workflow, build_native_vlm_package_metadata, ) @@ -484,6 +485,7 @@ 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_workflow(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) @@ -2181,6 +2183,7 @@ def write_tts_workflow_metadata(pkg: Any, output_dir: str, config: Any) -> str: metadata = build_tts_workflow_metadata(pkg, config) os.makedirs(output_dir, exist_ok=True) pkg.save_policy_components(output_dir) + add_adapter_service_to_workflow(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) @@ -2527,6 +2530,7 @@ def write_diffusion_workflow_metadata( timesteps=timesteps, ) pkg.save_policy_components(output_dir) + add_adapter_service_to_workflow(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) @@ -3569,6 +3573,7 @@ def write_vlm_workflow_metadata( 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_workflow(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) @@ -4377,6 +4382,7 @@ def write_speculative_workflow_metadata( adaptive_k_max=adaptive_k_max, ) pkg.save_policy_components(output_dir) + add_adapter_service_to_workflow(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) @@ -5681,6 +5687,7 @@ def write_decoder_workflow_metadata( 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_workflow(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) @@ -5700,6 +5707,7 @@ def write_language_diffusion_workflow_metadata( num_inference_steps=num_inference_steps, ) pkg.save_policy_components(output_dir) + add_adapter_service_to_workflow(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) diff --git a/tests/fixtures/onnx_genai_workflows/README.md b/tests/fixtures/onnx_genai_workflows/README.md index 9e396e658..574b2ad16 100644 --- a/tests/fixtures/onnx_genai_workflows/README.md +++ b/tests/fixtures/onnx_genai_workflows/README.md @@ -1,9 +1,11 @@ # ONNX GenAI workflow conformance fixtures Generated by `tests/generate_onnx_genai_validation_packages.py` for semantic -validation and runtime conformance against `justinchuby/onnx-genai@c9bddd6e`. +validation and runtime conformance against `justinchuby/onnx-genai@8549e425`. 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. No downloaded -model weights are included. +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.json b/tests/fixtures/onnx_genai_workflows/adapter/adapters/blue.json new file mode 100644 index 000000000..a822ff25f --- /dev/null +++ b/tests/fixtures/onnx_genai_workflows/adapter/adapters/blue.json @@ -0,0 +1 @@ +{"targets":{"projection":{"a":[0.0,1.0],"b":[3.0,4.0]}}} \ No newline at end of file diff --git a/tests/fixtures/onnx_genai_workflows/adapter/adapters/red.json b/tests/fixtures/onnx_genai_workflows/adapter/adapters/red.json new file mode 100644 index 000000000..f1d1e9700 --- /dev/null +++ b/tests/fixtures/onnx_genai_workflows/adapter/adapters/red.json @@ -0,0 +1 @@ +{"targets":{"projection":{"a":[1.0,0.0],"b":[1.0,2.0]}}} \ 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..b3984a591 --- /dev/null +++ b/tests/fixtures/onnx_genai_workflows/adapter/inference_metadata.yaml @@ -0,0 +1,171 @@ +schema_version: v1 +pipeline: + workflow: + manifest: + ir_version: '1.0' + onnx_opsets: + ai.onnx: 24 + adapter_abis: + onnx-genai.parameter-overlay: '1' + capabilities: + - workflow_ssa + - typed_emit + - parameter_adapters + - heterogeneous_adapter_batching + inputs: + request.row_ids: + contract: + dtype: int64 + rank: 1 + shape: + - batch + role: + kind: runtime + version: '1.0' + role: row_ids + source: + kind: request + request.active: + contract: + dtype: bool + rank: 1 + shape: + - batch + role: + kind: opaque + source: + kind: application + name: active + request.request_epochs: + contract: + dtype: int64 + rank: 1 + shape: + - batch + role: + kind: runtime + version: '1.0' + role: request_epochs + 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 + ports: {} + 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 + adapters: + base_model_fingerprint: sha256:3fd36990826adeab952ece1ca18d2b2c721e47a21a40cd0d32e74c2af46bcd96 + row_ids: request.row_ids + request_epochs: request.request_epochs + active: request.active + application_capability: onnx-genai.adapters + 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: + identity: blue + version: '1' + base_model_fingerprint: sha256:3fd36990826adeab952ece1ca18d2b2c721e47a21a40cd0d32e74c2af46bcd96 + rank: 1 + alpha: 1.0 + dtype: float32 + provenance: in_memory + weights: + - location: adapters/blue.json + sha256: 66ecbb05ef164997eb5d21cd7ced595ee7457daa67ac6c4b26bdba27d7d238e7 + format: json + targets: + - component: decoder + parameter: projection + weight_key: projection + input_features: 2 + output_features: 2 + red: + identity: red + version: '1' + base_model_fingerprint: sha256:3fd36990826adeab952ece1ca18d2b2c721e47a21a40cd0d32e74c2af46bcd96 + rank: 1 + alpha: 1.0 + dtype: float32 + provenance: in_memory + weights: + - location: adapters/red.json + sha256: 7de6cf124e348f4e8fce31694559fb73644a463ee6b43931bf66cc617e157783 + format: json + targets: + - component: decoder + parameter: projection + weight_key: projection + input_features: 2 + output_features: 2 diff --git a/tests/fixtures/onnx_genai_workflows/adapter/model.onnx b/tests/fixtures/onnx_genai_workflows/adapter/model.onnx new file mode 100644 index 0000000000000000000000000000000000000000..caf3d7ffbe1b82ff0fdad6034b7d4248cc7bcbf7 GIT binary patch literal 356 zcmdsh0KPwd|ke??ch+ymGmzETimPjE ModelPackage: ) +def _adapter_package() -> 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) + fingerprint = fingerprint_model_weights({"decoder": model}) + target = AdapterTarget("decoder", "projection") + manifest = AdapterTargetManifest( + fingerprint, + ( + AdapterTargetDescriptor( + target, + semantic_name="projection", + node_name="projection", + output_name="projection.output", + input_size=2, + output_size=2, + ), + ), + ) + package = ModelPackage( + {"decoder": model}, + adapter_target_manifest=manifest, + adapter_service_options=AdapterServiceOptions( + row_ids="request.row_ids", + request_epochs="request.request_epochs", + active="request.active", + cache_max_entries=2, + ), + ) + 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), + ), + ): + package.add_adapter_artifact( + AdapterArtifact( + name, + fingerprint, + (AdapterWeights(target, ir.tensor(a), ir.tensor(b), 1.0),), + identity=name, + version="1", + ) + ) + return package + + +def _write_adapter_metadata(package: ModelPackage, directory: Path) -> None: + metadata = { + "schema_version": "v1", + "pipeline": { + "workflow": { + "manifest": { + "ir_version": "1.0", + "onnx_opsets": {"ai.onnx": 24}, + "adapter_abis": {"onnx-genai.parameter-overlay": "1"}, + "capabilities": [ + "workflow_ssa", + "typed_emit", + "parameter_adapters", + "heterogeneous_adapter_batching", + ], + }, + "inputs": { + "request.row_ids": { + "contract": { + "dtype": "int64", + "rank": 1, + "shape": ["batch"], + }, + "role": { + "kind": "runtime", + "version": "1.0", + "role": "row_ids", + }, + "source": {"kind": "request"}, + }, + "request.active": { + "contract": { + "dtype": "bool", + "rank": 1, + "shape": ["batch"], + }, + "role": {"kind": "opaque"}, + "source": {"kind": "application", "name": "active"}, + }, + "request.request_epochs": { + "contract": { + "dtype": "int64", + "rank": 1, + "shape": ["batch"], + }, + "role": { + "kind": "runtime", + "version": "1.0", + "role": "request_epochs", + }, + "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", + }, + "ports": {}, + }, + "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_workflow(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 main() -> None: parser = argparse.ArgumentParser() parser.add_argument("output", type=Path) @@ -504,16 +722,23 @@ def main() -> None: codec.save(str(directory), progress_bar=False, check_weights=False) write_audio_codec_workflow_metadata(codec, str(directory)) + adapter = _adapter_package() + directory = args.output / "adapter" + adapter.save(str(directory), progress_bar=False) + _write_adapter_metadata(adapter, directory) + (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 against `justinchuby/onnx-genai@c9bddd6e`. +validation and runtime conformance against `justinchuby/onnx-genai@8549e425`. 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. No downloaded -model weights are included. +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", ) diff --git a/tests/onnx_genai_workflow_conformance.rs b/tests/onnx_genai_workflow_conformance.rs index a941addad..6ed619756 100644 --- a/tests/onnx_genai_workflow_conformance.rs +++ b/tests/onnx_genai_workflow_conformance.rs @@ -6,8 +6,9 @@ #![allow(clippy::field_reassign_with_default)] use onnx_genai_engine::{ - Engine, EngineConfig, GenerateOptions, GeneratePrompt, GenerateRequest, - PipelineGenerateRequest, pipeline::{PipelineEngine, WorkflowOutputRole}, + AdapterActivation, AdapterSelection, Engine, EngineConfig, GenerateOptions, + GeneratePrompt, GenerateRequest, PipelineGenerateRequest, + pipeline::{PipelineEngine, WorkflowOutputRole}, }; use onnx_genai_ort::{DataType, Value}; use std::path::PathBuf; @@ -29,6 +30,92 @@ fn options(max_new_tokens: usize) -> GenerateOptions { options } +fn adapter_request( + row_ids: &[i64], + request_epochs: &[i64], + active: &[bool], + values: &[f32], + selection: AdapterSelection, +) -> anyhow::Result { + let batch = i64::try_from(row_ids.len())?; + Ok(PipelineGenerateRequest::new(GenerateRequest { + prompt: GeneratePrompt::TokenIds(vec![]), + options: Default::default(), + }) + .with_input("request.row_ids", Value::from_slice_i64(row_ids, &[batch])?) + .with_input( + "request.request_epochs", + Value::from_slice_i64(request_epochs, &[batch])?, + ) + .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])?, + ) + .with_adapters(selection)) +} + +#[test] +fn mobius_parameter_adapters_preserve_order_rows_compaction_and_epochs() -> anyhow::Result<()> { + let mut engine = Engine::from_pipeline_dir(&root("adapter")?, EngineConfig::default())?; + let selection = AdapterSelection::default() + .with_row(10, 0, [AdapterActivation::new("red", 1.0)]) + .with_row( + 30, + 0, + [ + AdapterActivation::new("red", 0.5), + AdapterActivation::new("blue", 1.0), + ], + ); + let output = engine.run_pipeline(adapter_request( + &[10, 20, 30], + &[0, 0, 0], + &[true, false, true], + &[1.0, 2.0, 3.0, 4.0, 5.0, 6.0], + selection.clone(), + )?)?; + assert_eq!( + output["result"].to_vec_f32()?, + vec![2.0, 4.0, 3.0, 4.0, 25.5, 35.0] + ); + let compacted = engine.run_pipeline(adapter_request( + &[30, 10], + &[0, 0], + &[true, true], + &[5.0, 6.0, 1.0, 2.0], + selection, + )?)?; + assert_eq!( + compacted["result"].to_vec_f32()?, + vec![25.5, 35.0, 2.0, 4.0] + ); + let reused = + AdapterSelection::default().with_row(10, 1, [AdapterActivation::new("blue", 1.0)]); + for _ in 0..2 { + let output = engine.run_pipeline(adapter_request( + &[10], + &[1], + &[true], + &[1.0, 2.0], + reused.clone(), + )?)?; + assert_eq!(output["result"].to_vec_f32()?, vec![7.0, 10.0]); + } + let diagnostic = engine.adapter_lifecycle_diagnostic(); + assert_eq!(diagnostic.loads, 2); + assert!(diagnostic.cache_hits > 0); + assert!(diagnostic.replayed_plans > 0); + Ok(()) +} + fn assert_batched_policy_super_island(engine: &PipelineEngine) { let diagnostics = engine.execution_island_diagnostics(); let island = diagnostics From d25779aa71f15cc6073a60b8e4ad936731bb4e4a Mon Sep 17 00:00:00 2001 From: justinchuby Date: Sat, 15 Aug 2026 07:34:55 +0000 Subject: [PATCH 098/151] Fix adapter metadata lint findings Use the type-specific exception required by Ruff and normalize fixture generator import ordering so the adapter producer passes CI lint. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: justinchuby --- src/mobius/integrations/onnx_genai/inference_metadata.py | 2 +- tests/generate_onnx_genai_validation_packages.py | 8 ++++---- 2 files changed, 5 insertions(+), 5 deletions(-) diff --git a/src/mobius/integrations/onnx_genai/inference_metadata.py b/src/mobius/integrations/onnx_genai/inference_metadata.py index 5f223fcee..0e0a1237d 100644 --- a/src/mobius/integrations/onnx_genai/inference_metadata.py +++ b/src/mobius/integrations/onnx_genai/inference_metadata.py @@ -1493,7 +1493,7 @@ def add_adapter_service_to_workflow( return metadata workflow = metadata.get("pipeline", {}).get("workflow") if not isinstance(workflow, dict): - raise ValueError("parameter adapters require pipeline.workflow metadata") + raise TypeError("parameter adapters require pipeline.workflow metadata") manifest = getattr(pkg, "adapter_target_manifest", None) if manifest is None: raise ValueError("parameter adapters require an authoritative adapter target manifest") diff --git a/tests/generate_onnx_genai_validation_packages.py b/tests/generate_onnx_genai_validation_packages.py index 2429cbb21..4b25bc571 100644 --- a/tests/generate_onnx_genai_validation_packages.py +++ b/tests/generate_onnx_genai_validation_packages.py @@ -8,6 +8,7 @@ import yaml from onnxscript import GraphBuilder +from mobius._model_package import ModelPackage from mobius.adapters import ( AdapterArtifact, AdapterServiceOptions, @@ -17,20 +18,19 @@ AdapterWeights, fingerprint_model_weights, ) -from mobius._model_package import ModelPackage 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_workflow, +) from mobius.integrations.onnx_genai.workflow_metadata import ( write_audio_codec_workflow_metadata, write_language_diffusion_workflow_metadata, write_speculative_workflow_metadata, ) -from mobius.integrations.onnx_genai.inference_metadata import ( - add_adapter_service_to_workflow, -) from mobius.models.qwen3_tts import Qwen3TTSForConditionalGeneration from mobius.models.qwen3_tts_test import _TINY_CONFIG from mobius.tasks import TTSTask From f72825535a059ff10583edb07f41f6543cbfe678 Mon Sep 17 00:00:00 2001 From: justinchuby Date: Sat, 15 Aug 2026 07:42:36 +0000 Subject: [PATCH 099/151] Resolve remaining adapter lint checks Adopt pairwise iteration, robust floating-point assertions, and normalized import ordering across the adapter producer changes. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: justinchuby --- src/mobius/adapters.py | 3 ++- src/mobius/adapters_test.py | 6 +++--- src/mobius/integrations/onnx_genai/auto_export.py | 2 +- 3 files changed, 6 insertions(+), 5 deletions(-) diff --git a/src/mobius/adapters.py b/src/mobius/adapters.py index 96b2ed866..bb131ab48 100644 --- a/src/mobius/adapters.py +++ b/src/mobius/adapters.py @@ -23,6 +23,7 @@ import dataclasses import hashlib +import itertools import json import math from collections.abc import Mapping, Sequence @@ -142,7 +143,7 @@ def __post_init__(self) -> None: 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 zip(ordered, ordered[1:]): + 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: diff --git a/src/mobius/adapters_test.py b/src/mobius/adapters_test.py index 53b21b4d7..4b5c485ad 100644 --- a/src/mobius/adapters_test.py +++ b/src/mobius/adapters_test.py @@ -5,8 +5,8 @@ from __future__ import annotations -import json import hashlib +import json import shutil import uuid from pathlib import Path @@ -336,7 +336,7 @@ def test_peft_migration_source_preserves_rank_alpha_and_provenance() -> None: ) artifact.validate_base({"decoder": model}) assert artifact.weights[0].rank == 2 - assert artifact.weights[0].alpha == 6.0 + 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" @@ -483,7 +483,7 @@ def test_exact_onnx_genai_catalog_and_portable_bundle_serialization() -> None: assert artifact["identity"] == "style-red" assert artifact["version"] == "2026.08" assert artifact["rank"] == 2 - assert artifact["alpha"] == 2.0 + assert artifact["alpha"] == pytest.approx(2.0) assert artifact["dtype"] == "float32" assert artifact["targets"] == [ { diff --git a/src/mobius/integrations/onnx_genai/auto_export.py b/src/mobius/integrations/onnx_genai/auto_export.py index 45a29ab73..7133ceb39 100644 --- a/src/mobius/integrations/onnx_genai/auto_export.py +++ b/src/mobius/integrations/onnx_genai/auto_export.py @@ -25,8 +25,8 @@ ) from mobius.integrations.onnx_genai.inference_metadata import ( SchedulerConfig, - add_explicit_package_io, add_adapter_service_to_workflow, + add_explicit_package_io, add_policy_components_to_workflow, load_diffusers_scheduler_config, write_speech_to_text_pipeline_metadata, From d6e48f1944ad512fc14a6cdbc734915bd7c9719e Mon Sep 17 00:00:00 2001 From: justinchuby Date: Sat, 15 Aug 2026 07:52:21 +0000 Subject: [PATCH 100/151] Pin finalized adapter runtime contract Validate the persisted adapter catalog and fixtures against the immutable ONNX GenAI request-epoch contract at 9ec89afe. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: justinchuby --- .github/workflows/main.yml | 2 +- tests/generate_onnx_genai_validation_packages.py | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/.github/workflows/main.yml b/.github/workflows/main.yml index 1eeed08d9..a4e88cd32 100644 --- a/.github/workflows/main.yml +++ b/.github/workflows/main.yml @@ -23,7 +23,7 @@ jobs: - uses: actions/checkout@v7 with: repository: justinchuby/onnx-genai - ref: 8549e42529eda6b9a135f1f5efd93416dd5d63aa + ref: 9ec89afef244b68e8724f3eacb49568b3083c47b path: validation/onnx-genai - uses: actions/setup-python@v7 with: diff --git a/tests/generate_onnx_genai_validation_packages.py b/tests/generate_onnx_genai_validation_packages.py index 4b25bc571..217a13162 100644 --- a/tests/generate_onnx_genai_validation_packages.py +++ b/tests/generate_onnx_genai_validation_packages.py @@ -731,7 +731,7 @@ def main() -> None: """# ONNX GenAI workflow conformance fixtures Generated by `tests/generate_onnx_genai_validation_packages.py` for semantic -validation and runtime conformance against `justinchuby/onnx-genai@8549e425`. +validation and runtime conformance against `justinchuby/onnx-genai@9ec89afe`. The decoder, VLM, diffusion, masked diffusion, speculative, and codec packages contain executable synthetic models. The TTS fixture uses the real tiny From 7809f21dfd960cd0d240e8fbfa725a8daf0d55f8 Mon Sep 17 00:00:00 2001 From: justinchuby Date: Sat, 15 Aug 2026 07:56:20 +0000 Subject: [PATCH 101/151] Enforce final adapter selection contract Pin the runtime validation commit that rejects duplicate aliases and align producer-side selection validation with its stable error contract. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: justinchuby --- .github/workflows/main.yml | 2 +- src/mobius/adapters.py | 2 +- src/mobius/adapters_test.py | 6 ++++++ tests/generate_onnx_genai_validation_packages.py | 2 +- 4 files changed, 9 insertions(+), 3 deletions(-) diff --git a/.github/workflows/main.yml b/.github/workflows/main.yml index a4e88cd32..1eeed08d9 100644 --- a/.github/workflows/main.yml +++ b/.github/workflows/main.yml @@ -23,7 +23,7 @@ jobs: - uses: actions/checkout@v7 with: repository: justinchuby/onnx-genai - ref: 9ec89afef244b68e8724f3eacb49568b3083c47b + ref: 8549e42529eda6b9a135f1f5efd93416dd5d63aa path: validation/onnx-genai - uses: actions/setup-python@v7 with: diff --git a/src/mobius/adapters.py b/src/mobius/adapters.py index bb131ab48..960731324 100644 --- a/src/mobius/adapters.py +++ b/src/mobius/adapters.py @@ -440,7 +440,7 @@ def __post_init__(self) -> None: 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("an adapter may appear at most once in a row composition") + raise ValueError("adapter row contains duplicate adapter") @dataclasses.dataclass(frozen=True) diff --git a/src/mobius/adapters_test.py b/src/mobius/adapters_test.py index 4b5c485ad..0b78f1485 100644 --- a/src/mobius/adapters_test.py +++ b/src/mobius/adapters_test.py @@ -551,6 +551,12 @@ def test_application_scale_matches_runtime_bound() -> None: AdapterApplication("style", 16.1) +def test_selection_rejects_duplicate_adapter() -> None: + application = AdapterApplication("style") + with pytest.raises(ValueError, match="contains duplicate adapter"): + AdapterRowSelection(100, 0, (application, application)) + + def test_selection_rejects_unknown_adapter_and_invalid_permutation() -> None: batch = AdapterBatchSelection( (AdapterRowSelection(100, 0, (AdapterApplication("missing"),)),) diff --git a/tests/generate_onnx_genai_validation_packages.py b/tests/generate_onnx_genai_validation_packages.py index 217a13162..4b25bc571 100644 --- a/tests/generate_onnx_genai_validation_packages.py +++ b/tests/generate_onnx_genai_validation_packages.py @@ -731,7 +731,7 @@ def main() -> None: """# ONNX GenAI workflow conformance fixtures Generated by `tests/generate_onnx_genai_validation_packages.py` for semantic -validation and runtime conformance against `justinchuby/onnx-genai@9ec89afe`. +validation and runtime conformance against `justinchuby/onnx-genai@8549e425`. The decoder, VLM, diffusion, masked diffusion, speculative, and codec packages contain executable synthetic models. The TTS fixture uses the real tiny From 5c9841a9579a6a8b391bce30c8a94a47bb80e453 Mon Sep 17 00:00:00 2001 From: justinchuby Date: Sat, 15 Aug 2026 08:30:14 +0000 Subject: [PATCH 102/151] Freeze generic adapter wire ABI Align ModelPackage persistence and workflow metadata with ONNX GenAI adapters@1, including target-scoped base fingerprints, RFC 8785 artifacts, stable catalog indices, fixed-shape SSA selection tensors, canonical safetensors, and native parameter bindings. Regenerate executable heterogeneous batching fixtures and cover ordered composition, inactive rows, compaction, request-epoch slot reuse, eviction, reload, capture invalidation, and replay against ONNX GenAI 21a935c2. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: justinchuby --- .github/workflows/main.yml | 2 +- pyproject.toml | 1 + src/mobius/__init__.py | 2 + src/mobius/_model_package.py | 105 +++++--- src/mobius/adapter_io.py | 13 +- src/mobius/adapters.py | 238 +++++++++++++++--- src/mobius/adapters_test.py | 180 +++++++++++-- .../onnx_genai/inference_metadata.py | 110 +++++--- tests/fixtures/onnx_genai_workflows/README.md | 2 +- .../adapter/adapters/blue.json | 1 - .../adapter/adapters/blue/adapter.json | 1 + .../adapter/adapters/green/adapter.json | 1 + .../adapter/adapters/red.json | 1 - .../adapter/adapters/red/adapter.json | 1 + .../adapter/inference_metadata.yaml | 93 +++++-- ...generate_onnx_genai_validation_packages.py | 18 +- tests/onnx_genai_workflow_conformance.rs | 80 +++++- 17 files changed, 691 insertions(+), 158 deletions(-) delete mode 100644 tests/fixtures/onnx_genai_workflows/adapter/adapters/blue.json create mode 100644 tests/fixtures/onnx_genai_workflows/adapter/adapters/blue/adapter.json create mode 100644 tests/fixtures/onnx_genai_workflows/adapter/adapters/green/adapter.json delete mode 100644 tests/fixtures/onnx_genai_workflows/adapter/adapters/red.json create mode 100644 tests/fixtures/onnx_genai_workflows/adapter/adapters/red/adapter.json diff --git a/.github/workflows/main.yml b/.github/workflows/main.yml index 1eeed08d9..631f62bc8 100644 --- a/.github/workflows/main.yml +++ b/.github/workflows/main.yml @@ -23,7 +23,7 @@ jobs: - uses: actions/checkout@v7 with: repository: justinchuby/onnx-genai - ref: 8549e42529eda6b9a135f1f5efd93416dd5d63aa + ref: 21a935c241f9fb8bb4b77e4df39b2716b0f70a26 path: validation/onnx-genai - uses: actions/setup-python@v7 with: 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/src/mobius/__init__.py b/src/mobius/__init__.py index 322f88bce..6157e6e90 100644 --- a/src/mobius/__init__.py +++ b/src/mobius/__init__.py @@ -9,6 +9,7 @@ "AdapterArtifact", "AdapterBatchSelection", "AdapterRowSelection", + "AdapterSelectionTensors", "AdapterServiceOptions", "AdapterSource", "AdapterTarget", @@ -119,6 +120,7 @@ AdapterArtifact, AdapterBatchSelection, AdapterRowSelection, + AdapterSelectionTensors, AdapterServiceOptions, AdapterSource, AdapterTarget, diff --git a/src/mobius/_model_package.py b/src/mobius/_model_package.py index cefc41b97..836ac3e1f 100644 --- a/src/mobius/_model_package.py +++ b/src/mobius/_model_package.py @@ -22,7 +22,6 @@ import hashlib import inspect -import json import logging import os import shutil @@ -32,9 +31,12 @@ from contextlib import contextmanager from typing import Any +import numpy as np import onnx_ir as ir +import rfc8785 import torch import tqdm +from safetensors.numpy import save_file as save_safetensors_file from mobius._optimizations import fold_initializers_after_weights from mobius.adapters import ( @@ -62,18 +64,13 @@ def _adapter_dtype_name(dtype: ir.DataType) -> str: ) from error -def _adapter_source_file(artifact: AdapterArtifact) -> tuple[str, str]: +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, "ort_genai" - if artifact.source.format == "peft_safetensors": - path = artifact.source.path - if os.path.isdir(path): - path = os.path.join(path, "adapter_model.safetensors") - return path, "safetensors" + return artifact.source.path raise ValueError( f"adapter {artifact.name!r} source format {artifact.source.format!r} " "cannot be preserved" @@ -266,21 +263,26 @@ def add_adapter_artifact( 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) + artifact.validate_base( + self.data, + fingerprint_targets=( + tuple(self.adapter_target_manifest.bindings) + 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" + ) 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))}" ) - if self.adapter_artifacts: - expected = next(iter(self.adapter_artifacts.values())).target_bindings - if artifact.target_bindings != expected: - raise ValueError( - f"adapter artifact {artifact.name!r} target set does not align " - "with the existing N-adapter catalog" - ) self.adapter_artifacts[artifact.name] = artifact def save_adapter_artifacts(self, directory: str) -> dict[str, dict[str, object]]: @@ -300,7 +302,9 @@ def save_adapter_artifacts(self, directory: str) -> dict[str, dict[str, object]] } catalog: dict[str, dict[str, object]] = {} identities: set[tuple[str, str]] = set() - for alias, artifact in sorted(self.adapter_artifacts.items()): + 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( @@ -330,43 +334,67 @@ def save_adapter_artifacts(self, directory: str) -> dict[str, dict[str, object]] ) targets: list[dict[str, object]] = [] portable_targets: dict[str, dict[str, list[float]]] = {} + portable_safetensors: dict[str, np.ndarray] = {} + native_parameters = { + target: {"a": a, "b": b} for target, a, b in artifact.source.native_parameters + } for weight in sorted( artifact.weights, key=lambda item: (item.target.component, item.target.parameter), ): descriptor = descriptors[weight.target] weight_key = descriptor.semantic_name - targets.append( - { - "component": weight.target.component, - "parameter": weight.target.parameter, - "weight_key": weight_key, - "input_features": descriptor.input_size, - "output_features": descriptor.output_size, - } - ) + target_entry: dict[str, object] = { + "component": weight.target.component, + "parameter": weight.target.parameter, + "weight_key": weight_key, + "input_features": descriptor.input_size, + "output_features": descriptor.output_size, + } + if ( + self.adapter_service_options.preserve_source_format + and artifact.source.format == "onnx_adapter" + ): + try: + target_entry["native_parameters"] = native_parameters[weight.target] + except KeyError as error: + raise ValueError( + f"adapter {alias!r} must declare exact native A/B parameter " + f"names for {weight.target.component}.{weight.target.parameter}" + ) from error + targets.append(target_entry) portable_targets[weight_key] = { "a": weight.a.numpy().reshape(-1).astype("float32").tolist(), "b": weight.b.numpy().reshape(-1).astype("float32").tolist(), } + portable_safetensors[f"{weight_key}.a"] = weight.a.numpy() + portable_safetensors[f"{weight_key}.b"] = weight.b.numpy() + artifact_dir = os.path.join(directory, "adapters", alias) + os.makedirs(artifact_dir, exist_ok=True) if self.adapter_service_options.preserve_source_format: - source_path, weight_format = _adapter_source_file(artifact) - extension = ".onnx_adapter" if weight_format == "ort_genai" else ".safetensors" - relative_location = f"adapters/{alias}{extension}" - destination = os.path.join(directory, relative_location) - shutil.copyfile(source_path, destination) + if artifact.source.format == "onnx_adapter": + source_path = _adapter_source_file(artifact) + weight_format = "ort_genai" + relative_location = f"adapters/{alias}/adapter.onnx_adapter" + destination = os.path.join(directory, relative_location) + shutil.copyfile(source_path, destination) + elif artifact.source.format == "peft_safetensors": + weight_format = "safetensors" + relative_location = f"adapters/{alias}/adapter.safetensors" + destination = os.path.join(directory, relative_location) + save_safetensors_file(portable_safetensors, destination) + else: + raise ValueError( + f"adapter {alias!r} source format {artifact.source.format!r} " + "cannot be preserved" + ) with open(destination, "rb") as handle: payload = handle.read() else: - relative_location = f"adapters/{alias}.json" + relative_location = f"adapters/{alias}/adapter.json" destination = os.path.join(directory, relative_location) - payload = json.dumps( - {"targets": portable_targets}, - sort_keys=True, - separators=(",", ":"), - allow_nan=False, - ).encode("utf-8") + payload = rfc8785.dumps({"targets": portable_targets}) with open(destination, "wb") as handle: handle.write(payload) weight_format = "json" @@ -378,6 +406,7 @@ def save_adapter_artifacts(self, directory: str) -> dict[str, dict[str, object]] if artifact.source.checksum: provenance_parts.append(f"source_{artifact.source.checksum}") catalog[alias] = { + "index": artifact_index, "identity": artifact.stable_identity, "version": artifact.version, "base_model_fingerprint": artifact.base_fingerprint, diff --git a/src/mobius/adapter_io.py b/src/mobius/adapter_io.py index ce0ce9165..6c502f441 100644 --- a/src/mobius/adapter_io.py +++ b/src/mobius/adapter_io.py @@ -138,7 +138,11 @@ def load_peft_adapter( ) -def adapter_source_from_onnx_adapter(path: str | Path) -> AdapterSource: +def adapter_source_from_onnx_adapter( + path: str | Path, + *, + native_parameters: Mapping[AdapterTarget, tuple[str, str]] | None = None, +) -> AdapterSource: """Declare an ORT FlatBuffers adapter source without making it mandatory.""" path = Path(path) payload = path.read_bytes() @@ -148,4 +152,11 @@ def adapter_source_from_onnx_adapter(path: str | Path) -> AdapterSource: "onnx_adapter", path=str(path), checksum=f"sha256:{hashlib.sha256(payload).hexdigest()}", + native_parameters=tuple( + (target, names[0], names[1]) + for target, names in sorted( + (native_parameters or {}).items(), + key=lambda item: (item[0].component, item[0].parameter), + ) + ), ) diff --git a/src/mobius/adapters.py b/src/mobius/adapters.py index 960731324..f1d1ffa61 100644 --- a/src/mobius/adapters.py +++ b/src/mobius/adapters.py @@ -10,6 +10,7 @@ "AdapterArtifact", "AdapterBatchSelection", "AdapterRowSelection", + "AdapterSelectionTensors", "AdapterServiceOptions", "AdapterSource", "AdapterTarget", @@ -26,11 +27,13 @@ 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: @@ -40,6 +43,10 @@ def _validate_identifier(value: str, description: str) -> None: 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") @@ -56,27 +63,108 @@ def _update_tensor_hash(digest: Any, tensor: ir.Tensor) -> None: digest.update(_tensor_bytes(tensor)) -def fingerprint_model_weights(models: Mapping[str, ir.Model]) -> str: - """Return a deterministic SHA-256 fingerprint of loaded base-model weights. +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}" + ) - Component and initializer names are part of the digest, so an adapter cannot - silently bind to an equal-shaped parameter in a different model component. - """ - digest = hashlib.sha256() - for component_name, model in sorted(models.items()): - digest.update(component_name.encode()) - digest.update(b"\0") - for parameter_name, initializer in sorted(model.graph.initializers.items()): - if initializer.const_value is None: - raise ValueError( - f"cannot fingerprint unloaded initializer " - f"{component_name!r}/{parameter_name!r}" - ) - digest.update(parameter_name.encode()) - digest.update(b"\0") - _update_tensor_hash(digest, initializer.const_value) - digest.update(b"\0") - return f"sha256:{digest.hexdigest()}" + +def _target_fingerprint_record( + models: Mapping[str, ir.Model], + target: AdapterTarget, +) -> dict[str, object]: + 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, + } + ) + return { + "component": target.component, + "consumers": consumers, + "dtype": int(initializer.dtype), + "parameter": target.parameter, + "shape": [int(dimension) for dimension in initializer.shape], + "tensor_sha256": hashlib.sha256(_tensor_bytes(initializer.const_value)).hexdigest(), + } + + +def fingerprint_model_weights( + models: Mapping[str, ir.Model], + targets: Sequence[AdapterTarget] | 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(set(targets), key=lambda item: (item.component, item.parameter)) + if not unique_targets: + raise ValueError("adapter base fingerprint requires at least one target") + canonical = rfc8785.dumps( + { + "schema": "onnx-genai-targeted-base-v1", + "targets": [ + _target_fingerprint_record(models, target) for target in unique_targets + ], + } + ) + return f"onnx-genai-targeted-base-v1:sha256:{hashlib.sha256(canonical).hexdigest()}" @dataclasses.dataclass(frozen=True) @@ -170,12 +258,6 @@ def __post_init__(self) -> None: raise ValueError("adapter target manifest contains duplicate semantic names") def validate(self, models: Mapping[str, ir.Model]) -> None: - actual_fingerprint = fingerprint_model_weights(models) - if actual_fingerprint != self.base_fingerprint: - raise ValueError( - "adapter target manifest base fingerprint mismatch: " - f"expected {self.base_fingerprint}, got {actual_fingerprint}" - ) for descriptor in self.targets: model = models.get(descriptor.target.component) if model is None: @@ -214,6 +296,14 @@ def validate(self, models: Mapping[str, ir.Model]) -> None: f"adapter manifest node {descriptor.node_name!r} does not produce " f"{descriptor.output_name!r}" ) + actual_fingerprint = fingerprint_model_weights( + models, tuple(descriptor.target for descriptor in 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]: @@ -230,12 +320,20 @@ class AdapterSource: checksum: str | None = None base_model: str | None = None revision: str | None = None + native_parameters: tuple[tuple[AdapterTarget, str, str], ...] = () def __post_init__(self) -> None: 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") + targets = [target for target, _, _ in self.native_parameters] + if len(targets) != len(set(targets)): + raise ValueError("adapter source contains duplicate native parameter bindings") + if any(not a or not b or a == b for _, a, b in self.native_parameters): + raise ValueError( + "adapter native parameters must contain distinct non-empty A/B names" + ) @dataclasses.dataclass(frozen=True) @@ -331,15 +429,13 @@ def nbytes(self) -> int: weight.a.numpy().nbytes + weight.b.numpy().nbytes for weight in self.weights ) - def validate_base(self, models: Mapping[str, ir.Model]) -> None: + def validate_base( + self, + models: Mapping[str, ir.Model], + *, + fingerprint_targets: Sequence[AdapterTarget] | None = None, + ) -> None: """Validate fingerprint, target existence, dtype, and matrix dimensions.""" - actual_fingerprint = fingerprint_model_weights(models) - if actual_fingerprint != self.base_fingerprint: - raise ValueError( - f"adapter {self.name!r} base fingerprint mismatch: " - f"expected {self.base_fingerprint}, got {actual_fingerprint}" - ) - for weight in self.weights: model = models.get(weight.target.component) if model is None: @@ -367,6 +463,15 @@ def validate_base(self, models: Mapping[str, ir.Model]) -> None: 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]: @@ -406,8 +511,12 @@ class AdapterServiceOptions: row_ids: str | None = None request_epochs: str | None = None + adapter_ids: str = "request.adapter_ids" + adapter_counts: str = "request.adapter_counts" + scales: str = "request.adapter_scales" active: str | None = None - application_capability: str = "onnx-genai.adapters" + 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 @@ -420,6 +529,8 @@ def __post_init__(self) -> None: 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.portable_fallback and self.preserve_source_format: raise ValueError( "portable fallback requires portable JSON artifacts; " @@ -468,6 +579,50 @@ def compact(self, permutation: Sequence[int]) -> AdapterBatchSelection: raise ValueError("adapter compaction must be a permutation of all batch rows") return AdapterBatchSelection(tuple(self.rows[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.rows) + if len(active) != len(self.rows): + raise ValueError("adapter active rows must match the selection batch size") + self.validate_catalog(artifacts) + aliases = tuple(sorted(artifacts)) + indices = {alias: index for index, alias in enumerate(aliases)} + adapter_ids = np.full((len(self.rows), max_adapters), -1, dtype=np.int64) + scales = np.zeros((len(self.rows), max_adapters), dtype=np.float32) + counts = np.zeros((len(self.rows),), dtype=np.int64) + for row_index, (row, is_active) in enumerate(zip(self.rows, active)): + if not is_active: + continue + if len(row.adapters) > max_adapters: + raise ValueError( + f"adapter row {row.row_id} selects {len(row.adapters)} adapters, " + f"exceeding max_adapters {max_adapters}" + ) + counts[row_index] = len(row.adapters) + for slot, application in enumerate(row.adapters): + adapter_ids[row_index, slot] = indices[application.adapter] + scales[row_index, slot] = application.scale + return AdapterSelectionTensors( + row_ids=np.asarray([row.row_id for row in self.rows], dtype=np.int64), + request_epochs=np.asarray( + [row.request_epoch for row in self.rows], dtype=np.int64 + ), + adapter_ids=adapter_ids, + 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.""" @@ -476,6 +631,19 @@ def referenced_adapters(self) -> frozenset[str]: ) +@dataclasses.dataclass(frozen=True) +class AdapterSelectionTensors: + """Fixed-shape SSA request buffers for the ``onnx-genai.adapters@1`` ABI.""" + + row_ids: np.ndarray + request_epochs: np.ndarray + adapter_ids: np.ndarray + adapter_counts: np.ndarray + scales: np.ndarray + active: np.ndarray + aliases: tuple[str, ...] + + def compose_adapter_deltas( row: AdapterRowSelection, artifacts: Mapping[str, AdapterArtifact], diff --git a/src/mobius/adapters_test.py b/src/mobius/adapters_test.py index 0b78f1485..144817e97 100644 --- a/src/mobius/adapters_test.py +++ b/src/mobius/adapters_test.py @@ -14,7 +14,7 @@ import numpy as np import onnx_ir as ir import pytest -from safetensors.numpy import save_file +from safetensors.numpy import load_file, save_file from mobius import ( AdapterApplication, @@ -85,10 +85,11 @@ def _weights( 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=(_weights(),), + base_fingerprint=fingerprint_model_weights(models, (weights.target,)), + weights=(weights,), ) @@ -116,7 +117,12 @@ def _manifest(model: ir.Model, *, include_second: bool = False) -> AdapterTarget output_size=3, ) ) - return AdapterTargetManifest(fingerprint_model_weights({"decoder": model}), tuple(targets)) + return AdapterTargetManifest( + fingerprint_model_weights( + {"decoder": model}, tuple(descriptor.target for descriptor in targets) + ), + tuple(targets), + ) def test_artifact_validates_base_target_shape_dtype_and_checksum() -> None: @@ -132,6 +138,27 @@ def test_artifact_validates_base_target_shape_dtype_and_checksum() -> None: 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) @@ -260,6 +287,46 @@ def test_compaction_preserves_semantic_rows_and_slot_reuse_uses_epoch() -> None: assert reused_slot != original.rows[0] +def test_selection_lowers_to_stable_fixed_shape_request_tensors() -> None: + model = _model() + artifact = _artifact(model) + batch = AdapterBatchSelection( + ( + AdapterRowSelection( + 100, + 4, + ( + AdapterApplication("red", 0.5), + AdapterApplication("blue", -0.25), + ), + ), + AdapterRowSelection(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.row_ids, [100, 101]) + np.testing.assert_array_equal(tensors.request_epochs, [4, 5]) + np.testing.assert_array_equal(tensors.adapter_ids, [[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.adapter_ids.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.row_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() artifact = _artifact(model) @@ -275,7 +342,7 @@ def test_model_package_catalog_validates_and_rejects_duplicates() -> None: artifact.validate_checksum("sha256:" + "0" * 64) -def test_model_package_requires_n_adapter_target_alignment() -> None: +def test_model_package_allows_distinct_manifest_targets_per_adapter() -> None: model = _model() other = ir.Value( name="other.weight", @@ -288,17 +355,23 @@ def test_model_package_requires_n_adapter_target_alignment() -> None: node.outputs[0].name = "other.output" model.graph.append(node) model.graph.initializers.add(other) - package = ModelPackage( - {"decoder": model}, adapter_target_manifest=_manifest(model, include_second=True) + manifest = _manifest(model, include_second=True) + package = ModelPackage({"decoder": model}, adapter_target_manifest=manifest) + package.add_adapter_artifact( + AdapterArtifact("style", manifest.base_fingerprint, (_weights(),)) ) - package.add_adapter_artifact(_artifact(model, "style")) second = AdapterArtifact( "speaker", - fingerprint_model_weights({"decoder": model}), + manifest.base_fingerprint, (_weights(parameter="other.weight"),), ) - with pytest.raises(ValueError, match="does not align"): - package.add_adapter_artifact(second) + 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: @@ -340,6 +413,22 @@ def test_peft_migration_source_preserves_rank_alpha_and_provenance() -> None: 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(model), + 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)) + saved = load_file(output / catalog["peft-style"]["weights"][0]["location"]) + assert set(saved) == { + "layers.0.self_attn.q_proj.a", + "layers.0.self_attn.q_proj.b", + } np.testing.assert_allclose( artifact.weights[0].delta(), b @ a * 3.0, rtol=1e-6, atol=1e-6 ) @@ -388,15 +477,29 @@ def test_onnx_adapter_source_can_be_declared_for_native_capability() -> None: package.add_adapter_artifact( AdapterArtifact( "style", - fingerprint_model_weights({"decoder": model}), + fingerprint_model_weights( + {"decoder": model}, (AdapterTarget("decoder", "projection.weight"),) + ), (_weights(alpha=2.0),), - source=adapter_source_from_onnx_adapter(source_path), + source=adapter_source_from_onnx_adapter( + source_path, + native_parameters={ + AdapterTarget("decoder", "projection.weight"): ( + "style.projection.a", + "style.projection.b", + ) + }, + ), ) ) catalog = package.save_adapter_artifacts(str(output_directory)) declared = catalog["style"]["weights"][0] assert declared["format"] == "ort_genai" - assert declared["location"] == "adapters/style.onnx_adapter" + assert declared["location"] == "adapters/style/adapter.onnx_adapter" + assert catalog["style"]["targets"][0]["native_parameters"] == { + "a": "style.projection.a", + "b": "style.projection.b", + } copied = output_directory / declared["location"] assert copied.read_bytes() == source_path.read_bytes() assert declared["sha256"] == hashlib.sha256(copied.read_bytes()).hexdigest() @@ -420,13 +523,16 @@ def test_exact_onnx_genai_catalog_and_portable_bundle_serialization() -> None: adapter_service_options=AdapterServiceOptions( row_ids="request.row_ids", active="request.active", + max_adapters=2, cache_max_entries=2, ), ) package.add_adapter_artifact( AdapterArtifact( "red", - fingerprint_model_weights({"decoder": model}), + fingerprint_model_weights( + {"decoder": model}, (AdapterTarget("decoder", "projection.weight"),) + ), (_weights(alpha=2.0),), identity="style-red", version="2026.08", @@ -442,14 +548,26 @@ def test_exact_onnx_genai_catalog_and_portable_bundle_serialization() -> None: "dtype": "int64", "rank": 1, "shape": ["batch"], - } + }, + "role": { + "kind": "runtime", + "version": "1.0", + "role": "row_ids", + }, + "source": {"kind": "request"}, }, "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"}}}, @@ -459,11 +577,19 @@ def test_exact_onnx_genai_catalog_and_portable_bundle_serialization() -> None: } add_adapter_service_to_workflow(metadata, package, str(directory)) service = metadata["pipeline"]["workflow"]["adapters"] - assert service["base_model_fingerprint"].startswith("sha256:") - assert service["row_ids"] == "request.row_ids" - assert service["request_epochs"] == "request.request_epochs" - assert service["active"] == "request.active" - assert service["application_capability"] == "onnx-genai.adapters" + assert service["base_model_fingerprint"].startswith( + "onnx-genai-targeted-base-v1:sha256:" + ) + assert service["selection"] == { + "row_ids": "request.row_ids", + "request_epochs": "request.request_epochs", + "adapter_ids": "request.adapter_ids", + "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["cache"] == {"max_entries": 2, "eviction": "lru"} assert service["planning"] == { "bucket_by_adapter_set": True, @@ -480,6 +606,7 @@ def test_exact_onnx_genai_catalog_and_portable_bundle_serialization() -> None: "source": {"kind": "request"}, } artifact = service["artifacts"]["red"] + assert artifact["index"] == 0 assert artifact["identity"] == "style-red" assert artifact["version"] == "2026.08" assert artifact["rank"] == 2 @@ -497,6 +624,7 @@ def test_exact_onnx_genai_catalog_and_portable_bundle_serialization() -> None: weight = artifact["weights"][0] payload = (directory / weight["location"]).read_bytes() assert weight["format"] == "json" + assert weight["location"] == "adapters/red/adapter.json" assert len(weight["sha256"]) == 64 assert weight["sha256"] == hashlib.sha256(payload).hexdigest() bundle = json.loads(payload) @@ -526,7 +654,13 @@ def test_wire_contract_rejects_heterogeneous_target_rank() -> None: package.add_adapter_artifact( AdapterArtifact( "mixed-rank", - fingerprint_model_weights({"decoder": model}), + fingerprint_model_weights( + {"decoder": model}, + ( + AdapterTarget("decoder", "projection.weight"), + AdapterTarget("decoder", "other.weight"), + ), + ), ( _weights(), _weights( diff --git a/src/mobius/integrations/onnx_genai/inference_metadata.py b/src/mobius/integrations/onnx_genai/inference_metadata.py index 0e0a1237d..58cde78b4 100644 --- a/src/mobius/integrations/onnx_genai/inference_metadata.py +++ b/src/mobius/integrations/onnx_genai/inference_metadata.py @@ -1500,56 +1500,94 @@ def add_adapter_service_to_workflow( options = pkg.adapter_service_options inputs = workflow.get("inputs", {}) - def compatible_input(name: str, *, dtype: str) -> bool: + def compatible_input( + name: str, + *, + dtype: str, + shape: list[str | int], + role: str, + ) -> 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") == 1 - and contract.get("shape") == ["batch"] + and contract.get("rank") == len(shape) + and contract.get("shape") == shape + and declaration.get("required", True) + and declaration.get("source") == {"kind": "request"} + and semantic_role == {"kind": "runtime", "version": "1.0", "role": role} ) - row_ids = options.row_ids - if row_ids is None: - row_ids = next( - ( - candidate - for candidate in ("package.slot_ids", "request.row_ids") - if compatible_input(candidate, dtype="int64") - ), - None, - ) - if row_ids is None or not compatible_input(row_ids, dtype="int64"): - raise ValueError("adapter row_ids must reference an int64[batch] workflow input") + def ensure_input( + name: str, + *, + dtype: str, + shape: list[str | int], + role: str, + ) -> None: + if name not in inputs: + inputs[name] = { + "contract": {"dtype": dtype, "rank": len(shape), "shape": shape}, + "role": {"kind": "runtime", "version": "1.0", "role": role}, + "source": {"kind": "request"}, + } + if not compatible_input(name, dtype=dtype, shape=shape, role=role): + raise ValueError( + f"adapter {role} must reference a required request-sourced " + f"{dtype}{shape} workflow input" + ) + + row_ids = options.row_ids or "request.row_ids" + ensure_input(row_ids, dtype="int64", shape=["batch"], role="row_ids") request_epochs = options.request_epochs or "request.request_epochs" - if request_epochs not in inputs: - inputs[request_epochs] = { - "contract": {"dtype": "int64", "rank": 1, "shape": ["batch"]}, - "role": { - "kind": "runtime", - "version": "1.0", - "role": "request_epochs", - }, - "source": {"kind": "request"}, - } - if not compatible_input(request_epochs, dtype="int64"): - raise ValueError( - "adapter request_epochs must reference an int64[batch] workflow input" - ) + ensure_input( + request_epochs, + dtype="int64", + shape=["batch"], + role="request_epochs", + ) + ensure_input( + options.adapter_ids, + dtype="int64", + shape=["batch", options.max_adapters], + role="adapter_ids", + ) + 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", + ) active = options.active - if active is None and compatible_input("package.active", dtype="bool"): - active = "package.active" - if active is not None and not compatible_input(active, dtype="bool"): - raise ValueError("adapter active must reference a bool[batch] workflow input") + if active is not None: + ensure_input( + active, + dtype="bool", + shape=["batch"], + role="adapter_active", + ) catalog = pkg.save_adapter_artifacts(output_dir) workflow["adapters"] = { "base_model_fingerprint": manifest.base_fingerprint, - "row_ids": row_ids, - "request_epochs": request_epochs, - **({"active": active} if active is not None else {}), + "selection": { + "row_ids": row_ids, + "request_epochs": request_epochs, + "adapter_ids": options.adapter_ids, + "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": { diff --git a/tests/fixtures/onnx_genai_workflows/README.md b/tests/fixtures/onnx_genai_workflows/README.md index 574b2ad16..6c5c62c73 100644 --- a/tests/fixtures/onnx_genai_workflows/README.md +++ b/tests/fixtures/onnx_genai_workflows/README.md @@ -1,7 +1,7 @@ # ONNX GenAI workflow conformance fixtures Generated by `tests/generate_onnx_genai_validation_packages.py` for semantic -validation and runtime conformance against `justinchuby/onnx-genai@8549e425`. +validation and runtime conformance against `justinchuby/onnx-genai@21a935c2`. The decoder, VLM, diffusion, masked diffusion, speculative, and codec packages contain executable synthetic models. The TTS fixture uses the real tiny diff --git a/tests/fixtures/onnx_genai_workflows/adapter/adapters/blue.json b/tests/fixtures/onnx_genai_workflows/adapter/adapters/blue.json deleted file mode 100644 index a822ff25f..000000000 --- a/tests/fixtures/onnx_genai_workflows/adapter/adapters/blue.json +++ /dev/null @@ -1 +0,0 @@ -{"targets":{"projection":{"a":[0.0,1.0],"b":[3.0,4.0]}}} \ No newline at end of file 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/red.json b/tests/fixtures/onnx_genai_workflows/adapter/adapters/red.json deleted file mode 100644 index f1d1e9700..000000000 --- a/tests/fixtures/onnx_genai_workflows/adapter/adapters/red.json +++ /dev/null @@ -1 +0,0 @@ -{"targets":{"projection":{"a":[1.0,0.0],"b":[1.0,2.0]}}} \ 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 index b3984a591..e6f6933b5 100644 --- a/tests/fixtures/onnx_genai_workflows/adapter/inference_metadata.yaml +++ b/tests/fixtures/onnx_genai_workflows/adapter/inference_metadata.yaml @@ -32,10 +32,11 @@ pipeline: shape: - batch role: - kind: opaque + kind: runtime + version: '1.0' + role: adapter_active source: - kind: application - name: active + kind: request request.request_epochs: contract: dtype: int64 @@ -60,6 +61,44 @@ pipeline: source: kind: application name: activations + request.adapter_ids: + contract: + dtype: int64 + rank: 2 + shape: + - batch + - 2 + role: + kind: runtime + version: '1.0' + role: adapter_ids + source: + kind: request + request.adapter_counts: + contract: + dtype: int64 + rank: 1 + shape: + - batch + role: + kind: runtime + version: '1.0' + role: adapter_counts + source: + kind: request + request.adapter_scales: + contract: + dtype: float32 + rank: 2 + shape: + - batch + - 2 + role: + kind: runtime + version: '1.0' + role: adapter_scales + source: + kind: request outputs: result: contract: @@ -119,11 +158,16 @@ pipeline: output: result mode: replace adapters: - base_model_fingerprint: sha256:3fd36990826adeab952ece1ca18d2b2c721e47a21a40cd0d32e74c2af46bcd96 - row_ids: request.row_ids - request_epochs: request.request_epochs - active: request.active - application_capability: onnx-genai.adapters + base_model_fingerprint: onnx-genai-targeted-base-v1:sha256:3bc8619d443995ec33644536e15ad6e777609c63eb078273e098fb733d2c415f + selection: + row_ids: request.row_ids + request_epochs: request.request_epochs + adapter_ids: request.adapter_ids + 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 @@ -134,16 +178,36 @@ pipeline: invalidate_capture_on_eviction: true artifacts: blue: + index: 0 identity: blue version: '1' - base_model_fingerprint: sha256:3fd36990826adeab952ece1ca18d2b2c721e47a21a40cd0d32e74c2af46bcd96 + base_model_fingerprint: onnx-genai-targeted-base-v1:sha256:3bc8619d443995ec33644536e15ad6e777609c63eb078273e098fb733d2c415f + rank: 1 + alpha: 1.0 + dtype: float32 + provenance: in_memory + weights: + - location: adapters/blue/adapter.json + sha256: d9e67499ad74d4c2d62c45a79ceeea8dfa94893afde9615ea6030b79d052ddbe + format: json + targets: + - component: decoder + parameter: projection + weight_key: projection + input_features: 2 + output_features: 2 + green: + index: 1 + identity: green + version: '1' + base_model_fingerprint: onnx-genai-targeted-base-v1:sha256:3bc8619d443995ec33644536e15ad6e777609c63eb078273e098fb733d2c415f rank: 1 alpha: 1.0 dtype: float32 provenance: in_memory weights: - - location: adapters/blue.json - sha256: 66ecbb05ef164997eb5d21cd7ced595ee7457daa67ac6c4b26bdba27d7d238e7 + - location: adapters/green/adapter.json + sha256: d49ae9a3336884fa1ed8503caca1380699788add893870a670f6290c473bc935 format: json targets: - component: decoder @@ -152,16 +216,17 @@ pipeline: input_features: 2 output_features: 2 red: + index: 2 identity: red version: '1' - base_model_fingerprint: sha256:3fd36990826adeab952ece1ca18d2b2c721e47a21a40cd0d32e74c2af46bcd96 + base_model_fingerprint: onnx-genai-targeted-base-v1:sha256:3bc8619d443995ec33644536e15ad6e777609c63eb078273e098fb733d2c415f rank: 1 alpha: 1.0 dtype: float32 provenance: in_memory weights: - - location: adapters/red.json - sha256: 7de6cf124e348f4e8fce31694559fb73644a463ee6b43931bf66cc617e157783 + - location: adapters/red/adapter.json + sha256: b4bd656305aa2fea1d518df351af5fd3bf99b3d0a6c8bd8501dfc5107ac278a1 format: json targets: - component: decoder diff --git a/tests/generate_onnx_genai_validation_packages.py b/tests/generate_onnx_genai_validation_packages.py index 4b25bc571..7fcd1e866 100644 --- a/tests/generate_onnx_genai_validation_packages.py +++ b/tests/generate_onnx_genai_validation_packages.py @@ -492,8 +492,8 @@ def _adapter_package() -> ModelPackage: "projection.output", ) model = ir.Model(graph, ir_version=11) - fingerprint = fingerprint_model_weights({"decoder": model}) target = AdapterTarget("decoder", "projection") + fingerprint = fingerprint_model_weights({"decoder": model}, (target,)) manifest = AdapterTargetManifest( fingerprint, ( @@ -514,6 +514,7 @@ def _adapter_package() -> ModelPackage: row_ids="request.row_ids", request_epochs="request.request_epochs", active="request.active", + max_adapters=2, cache_max_entries=2, ), ) @@ -528,6 +529,11 @@ def _adapter_package() -> ModelPackage: 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( @@ -577,8 +583,12 @@ def _write_adapter_metadata(package: ModelPackage, directory: Path) -> None: "rank": 1, "shape": ["batch"], }, - "role": {"kind": "opaque"}, - "source": {"kind": "application", "name": "active"}, + "role": { + "kind": "runtime", + "version": "1.0", + "role": "adapter_active", + }, + "source": {"kind": "request"}, }, "request.request_epochs": { "contract": { @@ -731,7 +741,7 @@ def main() -> None: """# ONNX GenAI workflow conformance fixtures Generated by `tests/generate_onnx_genai_validation_packages.py` for semantic -validation and runtime conformance against `justinchuby/onnx-genai@8549e425`. +validation and runtime conformance against `justinchuby/onnx-genai@21a935c2`. The decoder, VLM, diffusion, masked diffusion, speculative, and codec packages contain executable synthetic models. The TTS fixture uses the real tiny diff --git a/tests/onnx_genai_workflow_conformance.rs b/tests/onnx_genai_workflow_conformance.rs index 6ed619756..34d6f4801 100644 --- a/tests/onnx_genai_workflow_conformance.rs +++ b/tests/onnx_genai_workflow_conformance.rs @@ -38,6 +38,27 @@ fn adapter_request( selection: AdapterSelection, ) -> anyhow::Result { let batch = i64::try_from(row_ids.len())?; + let mut adapter_ids = vec![-1i64; row_ids.len() * 2]; + let mut adapter_counts = vec![0i64; row_ids.len()]; + let mut adapter_scales = vec![0.0f32; row_ids.len() * 2]; + for (row, (&row_id, &request_epoch)) in row_ids.iter().zip(request_epochs).enumerate() { + let identity = onnx_genai_engine::AdapterRowIdentity { + row_id, + request_epoch, + }; + if let Some(activations) = selection.rows.get(&identity) { + adapter_counts[row] = i64::try_from(activations.len())?; + for (slot, activation) in activations.iter().enumerate() { + adapter_ids[row * 2 + slot] = match activation.adapter.as_str() { + "blue" => 0, + "green" => 1, + "red" => 2, + 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(), @@ -47,6 +68,18 @@ fn adapter_request( "request.request_epochs", Value::from_slice_i64(request_epochs, &[batch])?, ) + .with_input( + "request.adapter_ids", + Value::from_slice_i64(&adapter_ids, &[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( @@ -58,8 +91,7 @@ fn adapter_request( .with_input( "activations", Value::from_slice_f32(values, &[batch, 2])?, - ) - .with_adapters(selection)) + )) } #[test] @@ -67,6 +99,11 @@ fn mobius_parameter_adapters_preserve_order_rows_compaction_and_epochs() -> anyh let mut engine = Engine::from_pipeline_dir(&root("adapter")?, EngineConfig::default())?; let selection = AdapterSelection::default() .with_row(10, 0, [AdapterActivation::new("red", 1.0)]) + .with_row( + 20, + 0, + [AdapterActivation::new("blue", 1.0)], + ) .with_row( 30, 0, @@ -99,6 +136,18 @@ fn mobius_parameter_adapters_preserve_order_rows_compaction_and_epochs() -> anyh ); let reused = AdapterSelection::default().with_row(10, 1, [AdapterActivation::new("blue", 1.0)]); + let stale = engine.run_pipeline(adapter_request( + &[10], + &[1], + &[true], + &[1.0, 2.0], + AdapterSelection::default().with_row( + 10, + 0, + [AdapterActivation::new("red", 1.0)], + ), + )?)?; + assert_eq!(stale["result"].to_vec_f32()?, vec![1.0, 2.0]); for _ in 0..2 { let output = engine.run_pipeline(adapter_request( &[10], @@ -109,9 +158,34 @@ fn mobius_parameter_adapters_preserve_order_rows_compaction_and_epochs() -> anyh )?)?; assert_eq!(output["result"].to_vec_f32()?, vec![7.0, 10.0]); } + let green = + AdapterSelection::default().with_row(40, 0, [AdapterActivation::new("green", 1.0)]); + let output = engine.run_pipeline(adapter_request( + &[40], + &[0], + &[true], + &[1.0, 2.0], + green, + )?)?; + assert_eq!(output["result"].to_vec_f32()?, vec![4.0, 5.0]); + let red = + AdapterSelection::default().with_row(50, 0, [AdapterActivation::new("red", 1.0)]); + for _ in 0..2 { + let output = engine.run_pipeline(adapter_request( + &[50], + &[0], + &[true], + &[1.0, 2.0], + red.clone(), + )?)?; + assert_eq!(output["result"].to_vec_f32()?, vec![2.0, 4.0]); + } let diagnostic = engine.adapter_lifecycle_diagnostic(); - assert_eq!(diagnostic.loads, 2); + 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(()) } From 08ce7aef6c8d6bd0c6994629dc7916420e9f5fab Mon Sep 17 00:00:00 2001 From: justinchuby Date: Sat, 15 Aug 2026 09:04:09 +0000 Subject: [PATCH 103/151] Integrate canonical LoRA adapter metadata ABI Replace the superseded workflow-local adapter catalog with the canonical top-level target manifest, PEFT/ORT-compatible sources, graph-input bindings, and fixed-shape segment routing from ONNX GenAI 903a2d1a. Preserve targeted base fingerprints and add executable heterogeneous slot/epoch fixtures. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: justinchuby --- .github/workflows/main.yml | 2 +- src/mobius/__init__.py | 4 +- src/mobius/_model_package.py | 203 +++++++++------ src/mobius/adapter_io.py | 21 +- src/mobius/adapters.py | 229 ++++++++++++----- src/mobius/adapters_test.py | 231 +++++++++++------- .../integrations/onnx_genai/auto_export.py | 4 +- .../onnx_genai/inference_metadata.py | 120 +++++---- .../onnx_genai/workflow_metadata.py | 16 +- tests/fixtures/onnx_genai_workflows/README.md | 2 +- .../adapter/adapters/peft/adapter.json | 1 + .../adapter/adapters/peft/adapter_config.json | 1 + .../adapters/peft/adapter_model.safetensors | Bin 0 -> 208 bytes .../adapter/inference_metadata.yaml | 198 ++++++++------- ...generate_onnx_genai_validation_packages.py | 87 +++++-- tests/onnx_genai_workflow_conformance.rs | 40 +-- 16 files changed, 732 insertions(+), 427 deletions(-) create mode 100644 tests/fixtures/onnx_genai_workflows/adapter/adapters/peft/adapter.json create mode 100644 tests/fixtures/onnx_genai_workflows/adapter/adapters/peft/adapter_config.json create mode 100644 tests/fixtures/onnx_genai_workflows/adapter/adapters/peft/adapter_model.safetensors diff --git a/.github/workflows/main.yml b/.github/workflows/main.yml index 631f62bc8..45e13f46c 100644 --- a/.github/workflows/main.yml +++ b/.github/workflows/main.yml @@ -23,7 +23,7 @@ jobs: - uses: actions/checkout@v7 with: repository: justinchuby/onnx-genai - ref: 21a935c241f9fb8bb4b77e4df39b2716b0f70a26 + ref: 903a2d1aad6e58ecd966db70e8b2fd310fce146a path: validation/onnx-genai - uses: actions/setup-python@v7 with: diff --git a/src/mobius/__init__.py b/src/mobius/__init__.py index 6157e6e90..edaec18a4 100644 --- a/src/mobius/__init__.py +++ b/src/mobius/__init__.py @@ -8,7 +8,7 @@ "AdapterApplication", "AdapterArtifact", "AdapterBatchSelection", - "AdapterRowSelection", + "AdapterSlotSelection", "AdapterSelectionTensors", "AdapterServiceOptions", "AdapterSource", @@ -119,9 +119,9 @@ AdapterApplication, AdapterArtifact, AdapterBatchSelection, - AdapterRowSelection, AdapterSelectionTensors, AdapterServiceOptions, + AdapterSlotSelection, AdapterSource, AdapterTarget, AdapterTargetDescriptor, diff --git a/src/mobius/_model_package.py b/src/mobius/_model_package.py index 836ac3e1f..94232a8f7 100644 --- a/src/mobius/_model_package.py +++ b/src/mobius/_model_package.py @@ -31,12 +31,10 @@ from contextlib import contextmanager from typing import Any -import numpy as np import onnx_ir as ir import rfc8785 import torch import tqdm -from safetensors.numpy import save_file as save_safetensors_file from mobius._optimizations import fold_initializers_after_weights from mobius.adapters import ( @@ -266,7 +264,7 @@ def add_adapter_artifact( artifact.validate_base( self.data, fingerprint_targets=( - tuple(self.adapter_target_manifest.bindings) + self.adapter_target_manifest.targets if self.adapter_target_manifest is not None else None ), @@ -300,6 +298,9 @@ def save_adapter_artifacts(self, directory: str) -> dict[str, dict[str, object]] descriptor.target: descriptor for descriptor in self.adapter_target_manifest.targets } + manifest_target_ids = { + target["id"] 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( @@ -312,92 +313,119 @@ def save_adapter_artifacts(self, directory: str) -> dict[str, dict[str, object]] f"{identity_version[1]} must be unique" ) identities.add(identity_version) - ranks = {weight.rank for weight in artifact.weights} - alphas = {weight.alpha for weight in artifact.weights} + 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(ranks) != 1 or len(alphas) != 1 or len(dtypes) != 1: + if len(dtypes) != 1: raise ValueError( - f"adapter {alias!r} has heterogeneous target rank/alpha/dtype, " + f"adapter {alias!r} has heterogeneous target dtypes, " "which the ONNX GenAI artifact contract cannot represent" ) - rank = ranks.pop() - alpha = alphas.pop() + rank = ordered_weights[0].rank + alpha = ordered_weights[0].alpha dtype = _adapter_dtype_name(dtypes.pop()) - if ( - self.adapter_service_options.preserve_source_format - and artifact.source.format == "onnx_adapter" - and alpha != rank - ): - raise ValueError( - f"adapter {alias!r} imports .onnx_adapter weights with baked scale; " - "alpha must equal rank" - ) - targets: list[dict[str, object]] = [] + bindings: list[dict[str, object]] = [] portable_targets: dict[str, dict[str, list[float]]] = {} - portable_safetensors: dict[str, np.ndarray] = {} - native_parameters = { - target: {"a": a, "b": b} for target, a, b in artifact.source.native_parameters - } - for weight in sorted( - artifact.weights, - key=lambda item: (item.target.component, item.target.parameter), - ): + for weight in ordered_weights: descriptor = descriptors[weight.target] - weight_key = descriptor.semantic_name - target_entry: dict[str, object] = { - "component": weight.target.component, - "parameter": weight.target.parameter, + target_id = weight.target_id or descriptor.semantic_name + if target_id not in manifest_target_ids: + raise ValueError( + f"adapter {alias!r} references target ID {target_id!r} " + "outside the authoritative manifest" + ) + weight_key = weight.weight_key or descriptor.semantic_name + binding: dict[str, object] = { + "target": target_id, "weight_key": weight_key, - "input_features": descriptor.input_size, - "output_features": descriptor.output_size, } - if ( - self.adapter_service_options.preserve_source_format - and artifact.source.format == "onnx_adapter" - ): - try: - target_entry["native_parameters"] = native_parameters[weight.target] - except KeyError as error: - raise ValueError( - f"adapter {alias!r} must declare exact native A/B parameter " - f"names for {weight.target.component}.{weight.target.parameter}" - ) from error - targets.append(target_entry) + 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(), } - portable_safetensors[f"{weight_key}.a"] = weight.a.numpy() - portable_safetensors[f"{weight_key}.b"] = weight.b.numpy() artifact_dir = os.path.join(directory, "adapters", alias) os.makedirs(artifact_dir, exist_ok=True) - if self.adapter_service_options.preserve_source_format: + 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(), + "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) - weight_format = "ort_genai" 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(), + "format": "ort_genai", + } + ) elif artifact.source.format == "peft_safetensors": - weight_format = "safetensors" - relative_location = f"adapters/{alias}/adapter.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) - save_safetensors_file(portable_safetensors, destination) + 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(), + "format": "hf_peft", + } + ) else: raise ValueError( f"adapter {alias!r} source format {artifact.source.format!r} " "cannot be preserved" ) - with open(destination, "rb") as handle: - payload = handle.read() - else: - 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_format = "json" + if not weight_artifacts: + raise ValueError( + f"adapter {alias!r} must emit a portable or preserved source artifact" + ) provenance_parts = [artifact.source.format] if artifact.source.base_model: provenance_parts.append(f"base={artifact.source.base_model}") @@ -414,17 +442,54 @@ def save_adapter_artifacts(self, directory: str) -> dict[str, dict[str, object]] "alpha": alpha, "dtype": dtype, "provenance": ";".join(provenance_parts), - "weights": [ - { - "location": relative_location, - "sha256": hashlib.sha256(payload).hexdigest(), - "format": weight_format, - } - ], - "targets": targets, + "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, + "parameter": descriptor.target.parameter, + "output_value": descriptor.output_name, + "activation_dtype": activation_dtype, + "input_features": descriptor.input_size, + "output_features": descriptor.output_size, + } + if descriptor.graph_input_a is not None: + graph_inputs = { + "a": descriptor.graph_input_a, + "b": descriptor.graph_input_b, + } + 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}" + sliced["output_slice"] = { + "offset": target_slice.offset, + "width": target_slice.width, + } + targets.append(sliced) + return {"targets": targets} + def save_policy_components( self, directory: str, diff --git a/src/mobius/adapter_io.py b/src/mobius/adapter_io.py index 6c502f441..19a621885 100644 --- a/src/mobius/adapter_io.py +++ b/src/mobius/adapter_io.py @@ -47,16 +47,18 @@ def _pattern_value(patterns: Mapping[str, object], module_key: str, default: obj return max(matches, key=lambda item: len(item[0]))[1] -def _resolve_target(module_key: str, targets: Mapping[str, AdapterTarget]) -> AdapterTarget: +def _resolve_target( + module_key: str, targets: Mapping[str, AdapterTarget] +) -> tuple[str, AdapterTarget]: if module_key in targets: - return targets[module_key] + 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][1] + return matches[0] def load_peft_adapter( @@ -114,12 +116,14 @@ def load_peft_adapter( 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( - _resolve_target(module_key, target_bindings), + target, ir.tensor(a), ir.tensor(b), alpha, + weight_key=weight_key, ) ) @@ -140,8 +144,6 @@ def load_peft_adapter( def adapter_source_from_onnx_adapter( path: str | Path, - *, - native_parameters: Mapping[AdapterTarget, tuple[str, str]] | None = None, ) -> AdapterSource: """Declare an ORT FlatBuffers adapter source without making it mandatory.""" path = Path(path) @@ -152,11 +154,4 @@ def adapter_source_from_onnx_adapter( "onnx_adapter", path=str(path), checksum=f"sha256:{hashlib.sha256(payload).hexdigest()}", - native_parameters=tuple( - (target, names[0], names[1]) - for target, names in sorted( - (native_parameters or {}).items(), - key=lambda item: (item[0].component, item[0].parameter), - ) - ), ) diff --git a/src/mobius/adapters.py b/src/mobius/adapters.py index f1d1ffa61..9a7623821 100644 --- a/src/mobius/adapters.py +++ b/src/mobius/adapters.py @@ -9,7 +9,7 @@ "AdapterApplication", "AdapterArtifact", "AdapterBatchSelection", - "AdapterRowSelection", + "AdapterSlotSelection", "AdapterSelectionTensors", "AdapterServiceOptions", "AdapterSource", @@ -100,8 +100,13 @@ def _canonical_attribute_value(attribute: ir.Attr) -> object: def _target_fingerprint_record( models: Mapping[str, ir.Model], - target: AdapterTarget, + 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}") @@ -132,7 +137,7 @@ def _target_fingerprint_record( "op_type": node.op_type, } ) - return { + record: dict[str, object] = { "component": target.component, "consumers": consumers, "dtype": int(initializer.dtype), @@ -140,11 +145,26 @@ def _target_fingerprint_record( "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["output_value"] = 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] | None = None, + targets: Sequence[AdapterTarget | AdapterTargetDescriptor] | None = None, ) -> str: """Fingerprint the exact immutable base parameters targeted by adapters.""" if targets is None: @@ -153,15 +173,37 @@ def fingerprint_model_weights( for component, model in sorted(models.items()) for parameter in sorted(model.graph.initializers) ) - unique_targets = sorted(set(targets), key=lambda item: (item.component, item.parameter)) + 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"] = { + "offset": target_slice.offset, + "width": target_slice.width, + } + records.append(sliced) canonical = rfc8785.dumps( { "schema": "onnx-genai-targeted-base-v1", - "targets": [ - _target_fingerprint_record(models, target) for target in unique_targets - ], + "targets": records, } ) return f"onnx-genai-targeted-base-v1:sha256:{hashlib.sha256(canonical).hexdigest()}" @@ -215,6 +257,10 @@ class AdapterTargetDescriptor: 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: @@ -236,6 +282,19 @@ def __post_init__(self) -> None: 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) @@ -296,9 +355,7 @@ def validate(self, models: Mapping[str, ir.Model]) -> None: f"adapter manifest node {descriptor.node_name!r} does not produce " f"{descriptor.output_name!r}" ) - actual_fingerprint = fingerprint_model_weights( - models, tuple(descriptor.target for descriptor in self.targets) - ) + actual_fingerprint = fingerprint_model_weights(models, self.targets) if actual_fingerprint != self.base_fingerprint: raise ValueError( "adapter target manifest base fingerprint mismatch: " @@ -320,20 +377,12 @@ class AdapterSource: checksum: str | None = None base_model: str | None = None revision: str | None = None - native_parameters: tuple[tuple[AdapterTarget, str, str], ...] = () def __post_init__(self) -> None: 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") - targets = [target for target, _, _ in self.native_parameters] - if len(targets) != len(set(targets)): - raise ValueError("adapter source contains duplicate native parameter bindings") - if any(not a or not b or a == b for _, a, b in self.native_parameters): - raise ValueError( - "adapter native parameters must contain distinct non-empty A/B names" - ) @dataclasses.dataclass(frozen=True) @@ -344,6 +393,8 @@ class AdapterWeights: 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: @@ -359,6 +410,10 @@ def __post_init__(self) -> None: 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: @@ -433,9 +488,24 @@ def validate_base( self, models: Mapping[str, ir.Model], *, - fingerprint_targets: Sequence[AdapterTarget] | None = None, + 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: @@ -451,12 +521,41 @@ def validate_base( ) expected_shape = [int(weight.b.shape[0]), int(weight.a.shape[1])] actual_shape = [int(dimension) for dimension in initializer.shape] - if actual_shape != expected_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}/" @@ -509,9 +608,9 @@ def __post_init__(self) -> None: class AdapterServiceOptions: """Producer-neutral runtime lifecycle, planning, and artifact format options.""" - row_ids: str | None = None + slot_ids: str | None = None request_epochs: str | None = None - adapter_ids: str = "request.adapter_ids" + segments: str = "request.adapter_segments" adapter_counts: str = "request.adapter_counts" scales: str = "request.adapter_scales" active: str | None = None @@ -523,6 +622,7 @@ class AdapterServiceOptions: 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: @@ -531,18 +631,19 @@ def __post_init__(self) -> None: 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.portable_fallback and self.preserve_source_format: + 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( - "portable fallback requires portable JSON artifacts; " - "source-format preservation requires a native adapter capability" + "adapter service must emit a portable fallback or preserved source format" ) @dataclasses.dataclass(frozen=True) -class AdapterRowSelection: - """Adapter composition for one stable semantic request row.""" +class AdapterSlotSelection: + """Adapter composition for one stable semantic serving slot.""" - row_id: int + slot_id: int request_epoch: int adapters: tuple[AdapterApplication, ...] = () @@ -551,33 +652,33 @@ def __post_init__(self) -> None: 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 row contains duplicate adapter") + raise ValueError("adapter slot contains duplicate adapter") @dataclasses.dataclass(frozen=True) class AdapterBatchSelection: """Fixed-shape, compaction-safe adapter state for a heterogeneous batch.""" - rows: tuple[AdapterRowSelection, ...] + slots: tuple[AdapterSlotSelection, ...] def __post_init__(self) -> None: - row_ids = [row.row_id for row in self.rows] - if len(row_ids) != len(set(row_ids)): - raise ValueError("adapter batch row IDs must be unique") + 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 row in self.rows: - for application in row.adapters: + for slot in self.slots: + for application in slot.adapters: if application.adapter not in artifacts: raise ValueError( - f"row {row.row_id} selects unknown adapter {application.adapter!r}" + f"slot {slot.slot_id} selects unknown adapter {application.adapter!r}" ) def compact(self, permutation: Sequence[int]) -> AdapterBatchSelection: - """Apply the same physical-row permutation used for all workflow state.""" - if sorted(permutation) != list(range(len(self.rows))): - raise ValueError("adapter compaction must be a permutation of all batch rows") - return AdapterBatchSelection(tuple(self.rows[index] for index in permutation)) + """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, @@ -590,33 +691,33 @@ def to_tensors( if max_adapters <= 0: raise ValueError("adapter max_adapters must be greater than zero") if active is None: - active = [True] * len(self.rows) - if len(active) != len(self.rows): - raise ValueError("adapter active rows must match the selection batch size") + 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)} - adapter_ids = np.full((len(self.rows), max_adapters), -1, dtype=np.int64) - scales = np.zeros((len(self.rows), max_adapters), dtype=np.float32) - counts = np.zeros((len(self.rows),), dtype=np.int64) - for row_index, (row, is_active) in enumerate(zip(self.rows, active)): + 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(row.adapters) > max_adapters: + if len(slot.adapters) > max_adapters: raise ValueError( - f"adapter row {row.row_id} selects {len(row.adapters)} adapters, " + f"adapter slot {slot.slot_id} selects {len(slot.adapters)} adapters, " f"exceeding max_adapters {max_adapters}" ) - counts[row_index] = len(row.adapters) - for slot, application in enumerate(row.adapters): - adapter_ids[row_index, slot] = indices[application.adapter] - scales[row_index, slot] = application.scale + 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( - row_ids=np.asarray([row.row_id for row in self.rows], dtype=np.int64), + slot_ids=np.asarray([slot.slot_id for slot in self.slots], dtype=np.int64), request_epochs=np.asarray( - [row.request_epoch for row in self.rows], dtype=np.int64 + [slot.request_epoch for slot in self.slots], dtype=np.int64 ), - adapter_ids=adapter_ids, + segments=segments, adapter_counts=counts, scales=scales, active=np.asarray(active, dtype=np.bool_), @@ -627,7 +728,7 @@ def to_tensors( def referenced_adapters(self) -> frozenset[str]: """Live adapter set that a paged runtime must pin against eviction.""" return frozenset( - application.adapter for row in self.rows for application in row.adapters + application.adapter for slot in self.slots for application in slot.adapters ) @@ -635,9 +736,9 @@ def referenced_adapters(self) -> frozenset[str]: class AdapterSelectionTensors: """Fixed-shape SSA request buffers for the ``onnx-genai.adapters@1`` ABI.""" - row_ids: np.ndarray + slot_ids: np.ndarray request_epochs: np.ndarray - adapter_ids: np.ndarray + segments: np.ndarray adapter_counts: np.ndarray scales: np.ndarray active: np.ndarray @@ -645,17 +746,17 @@ class AdapterSelectionTensors: def compose_adapter_deltas( - row: AdapterRowSelection, + slot: AdapterSlotSelection, artifacts: Mapping[str, AdapterArtifact], ) -> dict[AdapterTarget, np.ndarray]: - """Compose a row's selected adapters into reference parameter updates.""" + """Compose a slot's selected adapters into reference parameter updates.""" deltas: dict[AdapterTarget, np.ndarray] = {} - for application in row.adapters: + for application in slot.adapters: try: artifact = artifacts[application.adapter] except KeyError as error: raise ValueError( - f"row {row.row_id} selects unknown adapter {application.adapter!r}" + f"slot {slot.slot_id} selects unknown adapter {application.adapter!r}" ) from error for weight in artifact.weights: update = weight.delta() * application.scale diff --git a/src/mobius/adapters_test.py b/src/mobius/adapters_test.py index 144817e97..265b581d3 100644 --- a/src/mobius/adapters_test.py +++ b/src/mobius/adapters_test.py @@ -20,8 +20,8 @@ AdapterApplication, AdapterArtifact, AdapterBatchSelection, - AdapterRowSelection, AdapterServiceOptions, + AdapterSlotSelection, AdapterSource, AdapterTarget, AdapterTargetDescriptor, @@ -35,7 +35,7 @@ load_peft_adapter, ) from mobius.integrations.onnx_genai.inference_metadata import ( - add_adapter_service_to_workflow, + add_adapter_service_to_metadata, ) @@ -72,6 +72,8 @@ def _weights( 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 @@ -80,6 +82,8 @@ def _weights( ir.tensor(a_values), ir.tensor(b_values), alpha, + weight_key=weight_key, + target_id=target_id, ) @@ -104,6 +108,10 @@ def _manifest(model: ir.Model, *, include_second: bool = False) -> AdapterTarget 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: @@ -118,9 +126,7 @@ def _manifest(model: ir.Model, *, include_second: bool = False) -> AdapterTarget ) ) return AdapterTargetManifest( - fingerprint_model_weights( - {"decoder": model}, tuple(descriptor.target for descriptor in targets) - ), + fingerprint_model_weights({"decoder": model}, tuple(targets)), tuple(targets), ) @@ -222,7 +228,7 @@ def test_composed_delta_matches_scaled_sum() -> None: style.base_fingerprint, (_weights(alpha=2.0),), ) - row = AdapterRowSelection( + row = AdapterSlotSelection( 100, 1, ( @@ -239,14 +245,14 @@ def test_composed_delta_matches_scaled_sum() -> None: def test_zero_one_and_composed_per_row_adapters() -> None: batch = AdapterBatchSelection( ( - AdapterRowSelection(row_id=100, request_epoch=4), - AdapterRowSelection( - row_id=101, + AdapterSlotSelection(slot_id=100, request_epoch=4), + AdapterSlotSelection( + slot_id=101, request_epoch=7, adapters=(AdapterApplication("style", 0.5),), ), - AdapterRowSelection( - row_id=102, + AdapterSlotSelection( + slot_id=102, request_epoch=2, adapters=( AdapterApplication("style", 0.25), @@ -262,29 +268,29 @@ def test_zero_one_and_composed_per_row_adapters() -> None: } batch.validate_catalog(catalog) - assert batch.rows[0].adapters == () - assert [item.adapter for item in batch.rows[2].adapters] == ["style", "speaker"] + 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( ( - AdapterRowSelection(100, 1, (AdapterApplication("style"),)), - AdapterRowSelection(101, 5, (AdapterApplication("speaker", 0.25),)), + AdapterSlotSelection(100, 1, (AdapterApplication("style"),)), + AdapterSlotSelection(101, 5, (AdapterApplication("speaker", 0.25),)), ) ) compacted = original.compact([1, 0]) - assert [row.row_id for row in compacted.rows] == [101, 100] - assert compacted.rows[0].request_epoch == 5 + 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 = AdapterRowSelection( - row_id=200, + reused_slot = AdapterSlotSelection( + slot_id=200, request_epoch=2, adapters=(AdapterApplication("speaker"),), ) - assert reused_slot != original.rows[0] + assert reused_slot != original.slots[0] def test_selection_lowers_to_stable_fixed_shape_request_tensors() -> None: @@ -292,7 +298,7 @@ def test_selection_lowers_to_stable_fixed_shape_request_tensors() -> None: artifact = _artifact(model) batch = AdapterBatchSelection( ( - AdapterRowSelection( + AdapterSlotSelection( 100, 4, ( @@ -300,7 +306,7 @@ def test_selection_lowers_to_stable_fixed_shape_request_tensors() -> None: AdapterApplication("blue", -0.25), ), ), - AdapterRowSelection(101, 5, (AdapterApplication("blue", 1.0),)), + AdapterSlotSelection(101, 5, (AdapterApplication("blue", 1.0),)), ) ) tensors = batch.to_tensors( @@ -309,19 +315,19 @@ def test_selection_lowers_to_stable_fixed_shape_request_tensors() -> None: active=[True, False], ) assert tensors.aliases == ("blue", "red") - np.testing.assert_array_equal(tensors.row_ids, [100, 101]) + 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.adapter_ids, [[1, 0, -1], [-1, -1, -1]]) + 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.adapter_ids.dtype == np.int64 + 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.row_ids, [101, 100]) + 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) @@ -329,8 +335,9 @@ def test_selection_lowers_to_stable_fixed_shape_request_tensors() -> None: def test_model_package_catalog_validates_and_rejects_duplicates() -> None: model = _model() - artifact = _artifact(model) - package = ModelPackage({"decoder": model}, adapter_target_manifest=_manifest(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 @@ -399,15 +406,16 @@ def test_peft_migration_source_preserves_rank_alpha_and_provenance() -> None: directory / "adapter_model.safetensors", ) model = _model() + manifest = _manifest(model) artifact = load_peft_adapter( directory, name="peft-style", - base_fingerprint=fingerprint_model_weights({"decoder": model}), + base_fingerprint=manifest.base_fingerprint, target_bindings={ "layers.0.self_attn.q_proj": AdapterTarget("decoder", "projection.weight") }, ) - artifact.validate_base({"decoder": model}) + 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" @@ -415,7 +423,7 @@ def test_peft_migration_source_preserves_rank_alpha_and_provenance() -> None: assert artifact.source.revision == "producer-fixture" package = ModelPackage( {"decoder": model}, - adapter_target_manifest=_manifest(model), + adapter_target_manifest=manifest, adapter_service_options=AdapterServiceOptions( portable_fallback=False, preserve_source_format=True, @@ -424,11 +432,17 @@ def test_peft_migration_source_preserves_rank_alpha_and_provenance() -> None: package.add_adapter_artifact(artifact) output = directory / "package" catalog = package.save_adapter_artifacts(str(output)) - saved = load_file(output / catalog["peft-style"]["weights"][0]["location"]) + declaration = catalog["peft-style"]["weights"][0] + assert declaration["format"] == "hf_peft" + assert declaration["loader_capability"] == "onnx-genai.adapters.hf-peft@1" + saved = load_file(output / declaration["location"]) assert set(saved) == { - "layers.0.self_attn.q_proj.a", - "layers.0.self_attn.q_proj.b", + 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 ) @@ -466,9 +480,10 @@ def test_onnx_adapter_source_can_be_declared_for_native_capability() -> None: 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(model), + adapter_target_manifest=manifest, adapter_service_options=AdapterServiceOptions( portable_fallback=False, preserve_source_format=True, @@ -477,29 +492,22 @@ def test_onnx_adapter_source_can_be_declared_for_native_capability() -> None: package.add_adapter_artifact( AdapterArtifact( "style", - fingerprint_model_weights( - {"decoder": model}, (AdapterTarget("decoder", "projection.weight"),) - ), + manifest.base_fingerprint, (_weights(alpha=2.0),), - source=adapter_source_from_onnx_adapter( - source_path, - native_parameters={ - AdapterTarget("decoder", "projection.weight"): ( - "style.projection.a", - "style.projection.b", - ) - }, - ), + 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["location"] == "adapters/style/adapter.onnx_adapter" - assert catalog["style"]["targets"][0]["native_parameters"] == { - "a": "style.projection.a", - "b": "style.projection.b", - } + 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() @@ -517,11 +525,12 @@ def test_exact_onnx_genai_catalog_and_portable_bundle_serialization() -> None: directory.mkdir(parents=True) try: model = _model() + manifest = _manifest(model) package = ModelPackage( {"decoder": model}, - adapter_target_manifest=_manifest(model), + adapter_target_manifest=manifest, adapter_service_options=AdapterServiceOptions( - row_ids="request.row_ids", + slot_ids="request.slot_ids", active="request.active", max_adapters=2, cache_max_entries=2, @@ -530,9 +539,7 @@ def test_exact_onnx_genai_catalog_and_portable_bundle_serialization() -> None: package.add_adapter_artifact( AdapterArtifact( "red", - fingerprint_model_weights( - {"decoder": model}, (AdapterTarget("decoder", "projection.weight"),) - ), + manifest.base_fingerprint, (_weights(alpha=2.0),), identity="style-red", version="2026.08", @@ -543,18 +550,17 @@ def test_exact_onnx_genai_catalog_and_portable_bundle_serialization() -> None: "workflow": { "manifest": {"capabilities": []}, "inputs": { - "request.row_ids": { + "request.slot_ids": { "contract": { "dtype": "int64", "rank": 1, "shape": ["batch"], }, - "role": { - "kind": "runtime", - "version": "1.0", - "role": "row_ids", + "role": {"kind": "opaque"}, + "source": { + "kind": "application", + "name": "serving.slot_ids", }, - "source": {"kind": "request"}, }, "request.active": { "contract": { @@ -575,21 +581,42 @@ def test_exact_onnx_genai_catalog_and_portable_bundle_serialization() -> None: } } } - add_adapter_service_to_workflow(metadata, package, str(directory)) - service = metadata["pipeline"]["workflow"]["adapters"] + 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"] == { - "row_ids": "request.row_ids", + "slot_ids": "request.slot_ids", "request_epochs": "request.request_epochs", - "adapter_ids": "request.adapter_ids", + "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", + "parameter": "projection.weight", + "output_value": "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"] == { + "offset": 0, + "width": 3, + } assert service["cache"] == {"max_entries": 2, "eviction": "lru"} assert service["planning"] == { "bucket_by_adapter_set": True, @@ -612,18 +639,16 @@ def test_exact_onnx_genai_catalog_and_portable_bundle_serialization() -> None: assert artifact["rank"] == 2 assert artifact["alpha"] == pytest.approx(2.0) assert artifact["dtype"] == "float32" - assert artifact["targets"] == [ + assert artifact["bindings"] == [ { - "component": "decoder", - "parameter": "projection.weight", + "target": "layers.0.self_attn.q_proj", "weight_key": "layers.0.self_attn.q_proj", - "input_features": 4, - "output_features": 3, } ] 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["location"] == "adapters/red/adapter.json" assert len(weight["sha256"]) == 64 assert weight["sha256"] == hashlib.sha256(payload).hexdigest() @@ -635,7 +660,42 @@ def test_exact_onnx_genai_catalog_and_portable_bundle_serialization() -> None: shutil.rmtree(directory) -def test_wire_contract_rejects_heterogeneous_target_rank() -> None: +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", @@ -647,22 +707,14 @@ def test_wire_contract_rejects_heterogeneous_target_rank() -> None: node.outputs[0].name = "other.output" model.graph.append(node) model.graph.initializers.add(other) - package = ModelPackage( - {"decoder": model}, - adapter_target_manifest=_manifest(model, include_second=True), - ) + manifest = _manifest(model, include_second=True) + package = ModelPackage({"decoder": model}, adapter_target_manifest=manifest) package.add_adapter_artifact( AdapterArtifact( "mixed-rank", - fingerprint_model_weights( - {"decoder": model}, - ( - AdapterTarget("decoder", "projection.weight"), - AdapterTarget("decoder", "other.weight"), - ), - ), + manifest.base_fingerprint, ( - _weights(), + _weights(target_id="layers.0.self_attn.q_proj.q"), _weights( parameter="other.weight", a=np.ones((1, 4), dtype=np.float32), @@ -674,8 +726,11 @@ def test_wire_contract_rejects_heterogeneous_target_rank() -> None: directory = Path("artifacts") / f"adapter-rank-test-{uuid.uuid4().hex}" directory.mkdir(parents=True) try: - with pytest.raises(ValueError, match="heterogeneous target rank"): - package.save_adapter_artifacts(str(directory)) + 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) @@ -688,12 +743,12 @@ def test_application_scale_matches_runtime_bound() -> None: def test_selection_rejects_duplicate_adapter() -> None: application = AdapterApplication("style") with pytest.raises(ValueError, match="contains duplicate adapter"): - AdapterRowSelection(100, 0, (application, application)) + AdapterSlotSelection(100, 0, (application, application)) def test_selection_rejects_unknown_adapter_and_invalid_permutation() -> None: batch = AdapterBatchSelection( - (AdapterRowSelection(100, 0, (AdapterApplication("missing"),)),) + (AdapterSlotSelection(100, 0, (AdapterApplication("missing"),)),) ) with pytest.raises(ValueError, match="unknown adapter"): batch.validate_catalog({}) diff --git a/src/mobius/integrations/onnx_genai/auto_export.py b/src/mobius/integrations/onnx_genai/auto_export.py index 7133ceb39..5d10d9e0b 100644 --- a/src/mobius/integrations/onnx_genai/auto_export.py +++ b/src/mobius/integrations/onnx_genai/auto_export.py @@ -25,7 +25,7 @@ ) from mobius.integrations.onnx_genai.inference_metadata import ( SchedulerConfig, - add_adapter_service_to_workflow, + add_adapter_service_to_metadata, add_explicit_package_io, add_policy_components_to_workflow, load_diffusers_scheduler_config, @@ -109,7 +109,7 @@ def _add_explicit_io_to_file(path: str, pkg: Any, config: Any) -> None: metadata = yaml.safe_load(handle) add_explicit_package_io(metadata, pkg, config) add_policy_components_to_workflow(metadata, pkg) - add_adapter_service_to_workflow(metadata, pkg, os.path.dirname(path)) + add_adapter_service_to_metadata(metadata, pkg, os.path.dirname(path)) with open(path, "w", encoding="utf-8") as handle: yaml.safe_dump(metadata, handle, sort_keys=False) diff --git a/src/mobius/integrations/onnx_genai/inference_metadata.py b/src/mobius/integrations/onnx_genai/inference_metadata.py index 58cde78b4..9ed0b8e2d 100644 --- a/src/mobius/integrations/onnx_genai/inference_metadata.py +++ b/src/mobius/integrations/onnx_genai/inference_metadata.py @@ -1482,7 +1482,7 @@ def tensor_contract(value: Any) -> dict[str, Any]: return metadata -def add_adapter_service_to_workflow( +def add_adapter_service_to_metadata( metadata: dict[str, Any], pkg: Any, output_dir: str, @@ -1491,21 +1491,20 @@ def add_adapter_service_to_workflow( artifacts = getattr(pkg, "adapter_artifacts", {}) if not artifacts: return metadata - workflow = metadata.get("pipeline", {}).get("workflow") - if not isinstance(workflow, dict): - raise TypeError("parameter adapters require pipeline.workflow 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.get("inputs", {}) + inputs = workflow.setdefault("inputs", {}) if workflow is not None else {} def compatible_input( name: str, *, dtype: str, shape: list[str | int], - role: str, + role: str | None, ) -> bool: declaration = inputs.get(name) if not isinstance(declaration, dict): @@ -1517,8 +1516,11 @@ def compatible_input( and contract.get("rank") == len(shape) and contract.get("shape") == shape and declaration.get("required", True) - and declaration.get("source") == {"kind": "request"} - and semantic_role == {"kind": "runtime", "version": "1.0", "role": role} + 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( @@ -1526,63 +1528,82 @@ def ensure_input( *, dtype: str, shape: list[str | int], - role: str, + 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}, - "source": {"kind": "request"}, + "role": ( + {"kind": "runtime", "version": "1.0", "role": role} + if role is not None + else {"kind": "opaque"} + ), + "source": source or {"kind": "request"}, } if not compatible_input(name, dtype=dtype, shape=shape, role=role): raise ValueError( - f"adapter {role} must reference a required request-sourced " + f"adapter {role or 'slot_ids'} must reference a required " + "request/application-sourced " f"{dtype}{shape} workflow input" ) - row_ids = options.row_ids or "request.row_ids" - ensure_input(row_ids, dtype="int64", shape=["batch"], role="row_ids") + serving = workflow.get("serving") if workflow is not None else None + serving_slot_ids = serving.get("slot_ids") if isinstance(serving, dict) else None + slot_ids = options.slot_ids or serving_slot_ids or "request.slot_ids" request_epochs = options.request_epochs or "request.request_epochs" - ensure_input( - request_epochs, - dtype="int64", - shape=["batch"], - role="request_epochs", - ) - ensure_input( - options.adapter_ids, - dtype="int64", - shape=["batch", options.max_adapters], - role="adapter_ids", - ) - 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", - ) active = options.active - if active is not None: + if workflow is not None: + ensure_input( + slot_ids, + dtype="int64", + shape=["batch"], + role=None, + source={"kind": "application", "name": "serving.slot_ids"}, + ) ensure_input( - active, - dtype="bool", + request_epochs, + dtype="int64", shape=["batch"], - role="adapter_active", + role="request_epochs", ) + 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) - workflow["adapters"] = { + 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": { - "row_ids": row_ids, + "slot_ids": slot_ids, "request_epochs": request_epochs, - "adapter_ids": options.adapter_ids, + "segments": options.segments, "adapter_counts": options.adapter_counts, "scales": options.scales, **({"active": active} if active is not None else {}), @@ -1601,10 +1622,11 @@ def ensure_input( }, "artifacts": catalog, } - capabilities = workflow.setdefault("manifest", {}).setdefault("capabilities", []) - for capability in ("parameter_adapters", "heterogeneous_adapter_batching"): - if capability not in capabilities: - capabilities.append(capability) + 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) return metadata diff --git a/src/mobius/integrations/onnx_genai/workflow_metadata.py b/src/mobius/integrations/onnx_genai/workflow_metadata.py index 22627b104..6169bbe95 100644 --- a/src/mobius/integrations/onnx_genai/workflow_metadata.py +++ b/src/mobius/integrations/onnx_genai/workflow_metadata.py @@ -46,7 +46,7 @@ from mobius.integrations.onnx_genai.inference_metadata import ( _port, _shape_metadata, - add_adapter_service_to_workflow, + add_adapter_service_to_metadata, add_policy_components_to_workflow, build_native_vlm_package_metadata, ) @@ -485,7 +485,7 @@ 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_workflow(metadata, pkg, 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) @@ -2183,7 +2183,7 @@ def write_tts_workflow_metadata(pkg: Any, output_dir: str, config: Any) -> str: metadata = build_tts_workflow_metadata(pkg, config) os.makedirs(output_dir, exist_ok=True) pkg.save_policy_components(output_dir) - add_adapter_service_to_workflow(metadata, pkg, 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) @@ -2530,7 +2530,7 @@ def write_diffusion_workflow_metadata( timesteps=timesteps, ) pkg.save_policy_components(output_dir) - add_adapter_service_to_workflow(metadata, pkg, 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) @@ -3573,7 +3573,7 @@ def write_vlm_workflow_metadata( 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_workflow(metadata, pkg, 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) @@ -4382,7 +4382,7 @@ def write_speculative_workflow_metadata( adaptive_k_max=adaptive_k_max, ) pkg.save_policy_components(output_dir) - add_adapter_service_to_workflow(metadata, pkg, 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) @@ -5687,7 +5687,7 @@ def write_decoder_workflow_metadata( 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_workflow(metadata, pkg, 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) @@ -5707,7 +5707,7 @@ def write_language_diffusion_workflow_metadata( num_inference_steps=num_inference_steps, ) pkg.save_policy_components(output_dir) - add_adapter_service_to_workflow(metadata, pkg, 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) diff --git a/tests/fixtures/onnx_genai_workflows/README.md b/tests/fixtures/onnx_genai_workflows/README.md index 6c5c62c73..b66e98e54 100644 --- a/tests/fixtures/onnx_genai_workflows/README.md +++ b/tests/fixtures/onnx_genai_workflows/README.md @@ -1,7 +1,7 @@ # ONNX GenAI workflow conformance fixtures Generated by `tests/generate_onnx_genai_validation_packages.py` for semantic -validation and runtime conformance against `justinchuby/onnx-genai@21a935c2`. +validation and runtime conformance against `justinchuby/onnx-genai@903a2d1a`. The decoder, VLM, diffusion, masked diffusion, speculative, and codec packages contain executable synthetic models. The TTS fixture uses the real tiny 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..761796c59 --- /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,"target_modules":["projection"]} \ No newline at end of file diff --git a/tests/fixtures/onnx_genai_workflows/adapter/adapters/peft/adapter_model.safetensors b/tests/fixtures/onnx_genai_workflows/adapter/adapters/peft/adapter_model.safetensors new file mode 100644 index 0000000000000000000000000000000000000000..ee02117f4d5e1d232a04c0481fd4b240f961c302 GIT binary patch literal 208 zcmdnN00GrXNr}a&@wxdasX2NDMfq8&$t9Wjd3rhdMTzl_dgZB^=@}(TR@F)=C6xuK zN>)m4#zsmyO2rw8AYQbgj!~?RQc7Y;VtjsDT5)PgF;LJz$0D{?2cN-C1PwOQF~l_3 ZLdVc7wzgIQ2pAX|?16Yc&=(-u4ggvtK*#_9 literal 0 HcmV?d00001 diff --git a/tests/fixtures/onnx_genai_workflows/adapter/inference_metadata.yaml b/tests/fixtures/onnx_genai_workflows/adapter/inference_metadata.yaml index e6f6933b5..3cdd8fa51 100644 --- a/tests/fixtures/onnx_genai_workflows/adapter/inference_metadata.yaml +++ b/tests/fixtures/onnx_genai_workflows/adapter/inference_metadata.yaml @@ -13,18 +13,17 @@ pipeline: - parameter_adapters - heterogeneous_adapter_batching inputs: - request.row_ids: + request.slot_ids: contract: dtype: int64 rank: 1 shape: - batch role: - kind: runtime - version: '1.0' - role: row_ids + kind: opaque source: - kind: request + kind: application + name: serving.slot_ids request.active: contract: dtype: bool @@ -61,7 +60,7 @@ pipeline: source: kind: application name: activations - request.adapter_ids: + request.adapter_segments: contract: dtype: int64 rank: 2 @@ -71,7 +70,7 @@ pipeline: role: kind: runtime version: '1.0' - role: adapter_ids + role: adapter_segments source: kind: request request.adapter_counts: @@ -157,80 +156,111 @@ pipeline: value: adapted output: result mode: replace - adapters: - base_model_fingerprint: onnx-genai-targeted-base-v1:sha256:3bc8619d443995ec33644536e15ad6e777609c63eb078273e098fb733d2c415f - selection: - row_ids: request.row_ids - request_epochs: request.request_epochs - adapter_ids: request.adapter_ids - 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:3bc8619d443995ec33644536e15ad6e777609c63eb078273e098fb733d2c415f - rank: 1 - alpha: 1.0 - dtype: float32 - provenance: in_memory - weights: - - location: adapters/blue/adapter.json - sha256: d9e67499ad74d4c2d62c45a79ceeea8dfa94893afde9615ea6030b79d052ddbe - format: json - targets: - - component: decoder - parameter: projection - weight_key: projection - input_features: 2 - output_features: 2 - green: - index: 1 - identity: green - version: '1' - base_model_fingerprint: onnx-genai-targeted-base-v1:sha256:3bc8619d443995ec33644536e15ad6e777609c63eb078273e098fb733d2c415f - rank: 1 - alpha: 1.0 - dtype: float32 - provenance: in_memory - weights: - - location: adapters/green/adapter.json - sha256: d49ae9a3336884fa1ed8503caca1380699788add893870a670f6290c473bc935 - format: json - targets: - - component: decoder - parameter: projection - weight_key: projection - input_features: 2 - output_features: 2 - red: - index: 2 - identity: red - version: '1' - base_model_fingerprint: onnx-genai-targeted-base-v1:sha256:3bc8619d443995ec33644536e15ad6e777609c63eb078273e098fb733d2c415f - rank: 1 - alpha: 1.0 - dtype: float32 - provenance: in_memory - weights: - - location: adapters/red/adapter.json - sha256: b4bd656305aa2fea1d518df351af5fd3bf99b3d0a6c8bd8501dfc5107ac278a1 - format: json - targets: - - component: decoder - parameter: projection - weight_key: projection - input_features: 2 - output_features: 2 +adapters: + base_model_fingerprint: onnx-genai-targeted-base-v1:sha256:ece1984372c107109d5e02743e04b6eefd452bfb8551a7b1e65bd7216871fcda + target_manifest: + targets: + - id: projection + component: decoder + parameter: projection + output_value: projection.output + activation_dtype: float32 + input_features: 2 + output_features: 2 + graph_inputs: + a: lora.projection.a + b: lora.projection.b + scale: lora.projection.scale + discovery_fallback: disabled + selection: + slot_ids: request.slot_ids + request_epochs: request.request_epochs + 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:ece1984372c107109d5e02743e04b6eefd452bfb8551a7b1e65bd7216871fcda + rank: 1 + alpha: 1.0 + dtype: float32 + provenance: in_memory + weights: + - location: adapters/blue/adapter.json + loader_capability: onnx-genai.adapters.json@1 + sha256: d9e67499ad74d4c2d62c45a79ceeea8dfa94893afde9615ea6030b79d052ddbe + format: json + bindings: + - target: projection + weight_key: projection + green: + index: 1 + identity: green + version: '1' + base_model_fingerprint: onnx-genai-targeted-base-v1:sha256:ece1984372c107109d5e02743e04b6eefd452bfb8551a7b1e65bd7216871fcda + rank: 1 + alpha: 1.0 + dtype: float32 + provenance: in_memory + weights: + - location: adapters/green/adapter.json + loader_capability: onnx-genai.adapters.json@1 + sha256: d49ae9a3336884fa1ed8503caca1380699788add893870a670f6290c473bc935 + format: json + bindings: + - target: projection + weight_key: projection + peft: + index: 2 + identity: peft + version: '1' + base_model_fingerprint: onnx-genai-targeted-base-v1:sha256:ece1984372c107109d5e02743e04b6eefd452bfb8551a7b1e65bd7216871fcda + rank: 1 + alpha: 1.0 + dtype: float32 + provenance: peft_safetensors;base=synthetic/adapter-base;source_sha256:1be2eedb9402accd480873439cd9d4cb2b5542de9369985b541c5c488c6be935 + weights: + - location: adapters/peft/adapter.json + loader_capability: onnx-genai.adapters.json@1 + sha256: 4351d2c682e0d1666ca3fdd6e967baaab6ab93c33cf5ace7bf3d6269713e85ce + 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: e3ce70ff5b09c7ddcbcd9426fd0e7a867c3c8433cb15b6c5dd5f18a3c5630a55 + 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:ece1984372c107109d5e02743e04b6eefd452bfb8551a7b1e65bd7216871fcda + rank: 1 + alpha: 1.0 + dtype: float32 + provenance: in_memory + weights: + - location: adapters/red/adapter.json + loader_capability: onnx-genai.adapters.json@1 + sha256: b4bd656305aa2fea1d518df351af5fd3bf99b3d0a6c8bd8501dfc5107ac278a1 + format: json + bindings: + - target: projection + weight_key: projection diff --git a/tests/generate_onnx_genai_validation_packages.py b/tests/generate_onnx_genai_validation_packages.py index 7fcd1e866..dad99306a 100644 --- a/tests/generate_onnx_genai_validation_packages.py +++ b/tests/generate_onnx_genai_validation_packages.py @@ -1,14 +1,18 @@ 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.adapter_io import load_peft_adapter from mobius.adapters import ( AdapterArtifact, AdapterServiceOptions, @@ -24,7 +28,7 @@ _VlmCfg, ) from mobius.integrations.onnx_genai.inference_metadata import ( - add_adapter_service_to_workflow, + add_adapter_service_to_metadata, ) from mobius.integrations.onnx_genai.workflow_metadata import ( write_audio_codec_workflow_metadata, @@ -474,7 +478,7 @@ def _executable_codec_package() -> ModelPackage: ) -def _adapter_package() -> ModelPackage: +def _adapter_package(source_root: Path) -> ModelPackage: graph, builder = _graph("decoder") activations = builder.input("activations", ir.DataType.FLOAT, ["batch", 2]) weight = ir.Value( @@ -493,29 +497,30 @@ def _adapter_package() -> ModelPackage: ) model = ir.Model(graph, ir_version=11) target = AdapterTarget("decoder", "projection") - fingerprint = fingerprint_model_weights({"decoder": model}, (target,)) - manifest = AdapterTargetManifest( - fingerprint, - ( - AdapterTargetDescriptor( - target, - semantic_name="projection", - node_name="projection", - output_name="projection.output", - input_size=2, - output_size=2, - ), - ), - ) + descriptor = AdapterTargetDescriptor( + target, + semantic_name="projection", + node_name="projection", + output_name="projection.output", + input_size=2, + output_size=2, + 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( - row_ids="request.row_ids", + slot_ids="request.slot_ids", request_epochs="request.request_epochs", active="request.active", max_adapters=2, cache_max_entries=2, + preserve_source_format=True, ), ) for name, a, b in ( @@ -544,6 +549,35 @@ def _adapter_package() -> ModelPackage: 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"], + }, + 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 @@ -564,18 +598,17 @@ def _write_adapter_metadata(package: ModelPackage, directory: Path) -> None: ], }, "inputs": { - "request.row_ids": { + "request.slot_ids": { "contract": { "dtype": "int64", "rank": 1, "shape": ["batch"], }, - "role": { - "kind": "runtime", - "version": "1.0", - "role": "row_ids", + "role": {"kind": "opaque"}, + "source": { + "kind": "application", + "name": "serving.slot_ids", }, - "source": {"kind": "request"}, }, "request.active": { "contract": { @@ -684,7 +717,7 @@ def _write_adapter_metadata(package: ModelPackage, directory: Path) -> None: } }, } - add_adapter_service_to_workflow(metadata, package, str(directory)) + 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) @@ -732,16 +765,18 @@ def main() -> None: codec.save(str(directory), progress_bar=False, check_weights=False) write_audio_codec_workflow_metadata(codec, str(directory)) - adapter = _adapter_package() 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 against `justinchuby/onnx-genai@21a935c2`. +validation and runtime conformance against `justinchuby/onnx-genai@903a2d1a`. The decoder, VLM, diffusion, masked diffusion, speculative, and codec packages contain executable synthetic models. The TTS fixture uses the real tiny diff --git a/tests/onnx_genai_workflow_conformance.rs b/tests/onnx_genai_workflow_conformance.rs index 34d6f4801..0b5db4e06 100644 --- a/tests/onnx_genai_workflow_conformance.rs +++ b/tests/onnx_genai_workflow_conformance.rs @@ -31,28 +31,28 @@ fn options(max_new_tokens: usize) -> GenerateOptions { } fn adapter_request( - row_ids: &[i64], + slot_ids: &[i64], request_epochs: &[i64], active: &[bool], values: &[f32], selection: AdapterSelection, ) -> anyhow::Result { - let batch = i64::try_from(row_ids.len())?; - let mut adapter_ids = vec![-1i64; row_ids.len() * 2]; - let mut adapter_counts = vec![0i64; row_ids.len()]; - let mut adapter_scales = vec![0.0f32; row_ids.len() * 2]; - for (row, (&row_id, &request_epoch)) in row_ids.iter().zip(request_epochs).enumerate() { - let identity = onnx_genai_engine::AdapterRowIdentity { - row_id, + let batch = i64::try_from(slot_ids.len())?; + let mut segments = vec![-1i64; slot_ids.len() * 2]; + let mut adapter_counts = vec![0i64; slot_ids.len()]; + let mut adapter_scales = vec![0.0f32; slot_ids.len() * 2]; + for (row, (&slot_id, &request_epoch)) in slot_ids.iter().zip(request_epochs).enumerate() { + let identity = onnx_genai_engine::AdapterSlotIdentity { + slot_id, request_epoch, }; if let Some(activations) = selection.rows.get(&identity) { adapter_counts[row] = i64::try_from(activations.len())?; for (slot, activation) in activations.iter().enumerate() { - adapter_ids[row * 2 + slot] = match activation.adapter.as_str() { + segments[row * 2 + slot] = match activation.adapter.as_str() { "blue" => 0, "green" => 1, - "red" => 2, + "red" => 3, other => anyhow::bail!("unknown test adapter {other}"), }; adapter_scales[row * 2 + slot] = activation.scale; @@ -63,14 +63,14 @@ fn adapter_request( prompt: GeneratePrompt::TokenIds(vec![]), options: Default::default(), }) - .with_input("request.row_ids", Value::from_slice_i64(row_ids, &[batch])?) + .with_input("request.slot_ids", Value::from_slice_i64(slot_ids, &[batch])?) .with_input( "request.request_epochs", Value::from_slice_i64(request_epochs, &[batch])?, ) .with_input( - "request.adapter_ids", - Value::from_slice_i64(&adapter_ids, &[batch, 2])?, + "request.adapter_segments", + Value::from_slice_i64(&segments, &[batch, 2])?, ) .with_input( "request.adapter_counts", @@ -98,13 +98,13 @@ fn adapter_request( fn mobius_parameter_adapters_preserve_order_rows_compaction_and_epochs() -> anyhow::Result<()> { let mut engine = Engine::from_pipeline_dir(&root("adapter")?, EngineConfig::default())?; let selection = AdapterSelection::default() - .with_row(10, 0, [AdapterActivation::new("red", 1.0)]) - .with_row( + .with_slot(10, 0, [AdapterActivation::new("red", 1.0)]) + .with_slot( 20, 0, [AdapterActivation::new("blue", 1.0)], ) - .with_row( + .with_slot( 30, 0, [ @@ -135,13 +135,13 @@ fn mobius_parameter_adapters_preserve_order_rows_compaction_and_epochs() -> anyh vec![25.5, 35.0, 2.0, 4.0] ); let reused = - AdapterSelection::default().with_row(10, 1, [AdapterActivation::new("blue", 1.0)]); + AdapterSelection::default().with_slot(10, 1, [AdapterActivation::new("blue", 1.0)]); let stale = engine.run_pipeline(adapter_request( &[10], &[1], &[true], &[1.0, 2.0], - AdapterSelection::default().with_row( + AdapterSelection::default().with_slot( 10, 0, [AdapterActivation::new("red", 1.0)], @@ -159,7 +159,7 @@ fn mobius_parameter_adapters_preserve_order_rows_compaction_and_epochs() -> anyh assert_eq!(output["result"].to_vec_f32()?, vec![7.0, 10.0]); } let green = - AdapterSelection::default().with_row(40, 0, [AdapterActivation::new("green", 1.0)]); + AdapterSelection::default().with_slot(40, 0, [AdapterActivation::new("green", 1.0)]); let output = engine.run_pipeline(adapter_request( &[40], &[0], @@ -169,7 +169,7 @@ fn mobius_parameter_adapters_preserve_order_rows_compaction_and_epochs() -> anyh )?)?; assert_eq!(output["result"].to_vec_f32()?, vec![4.0, 5.0]); let red = - AdapterSelection::default().with_row(50, 0, [AdapterActivation::new("red", 1.0)]); + AdapterSelection::default().with_slot(50, 0, [AdapterActivation::new("red", 1.0)]); for _ in 0..2 { let output = engine.run_pipeline(adapter_request( &[50], From 53e1271bba6087ed4e077730e29214bcb4c70138 Mon Sep 17 00:00:00 2001 From: justinchuby Date: Sat, 15 Aug 2026 09:31:44 +0000 Subject: [PATCH 104/151] Freeze final LoRA adapter producer ABI Pin ONNX GenAI d9482bca and emit exact target node/output names, labeled slices, structured provenance, and source-specific scale encoding so PEFT applies alpha/rank while TORT avoids double scaling. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: justinchuby --- .github/workflows/main.yml | 2 +- src/mobius/_model_package.py | 17 ++++++---- src/mobius/adapters.py | 7 +++- src/mobius/adapters_test.py | 13 ++++++- tests/fixtures/onnx_genai_workflows/README.md | 2 +- .../adapter/adapters/peft/adapter_config.json | 2 +- .../adapter/inference_metadata.yaml | 34 +++++++++++++------ ...generate_onnx_genai_validation_packages.py | 3 +- 8 files changed, 56 insertions(+), 24 deletions(-) diff --git a/.github/workflows/main.yml b/.github/workflows/main.yml index 45e13f46c..0cff255a6 100644 --- a/.github/workflows/main.yml +++ b/.github/workflows/main.yml @@ -23,7 +23,7 @@ jobs: - uses: actions/checkout@v7 with: repository: justinchuby/onnx-genai - ref: 903a2d1aad6e58ecd966db70e8b2fd310fce146a + ref: d9482bca0ad6ccec907a2b49faf033c2006a9fa9 path: validation/onnx-genai - uses: actions/setup-python@v7 with: diff --git a/src/mobius/_model_package.py b/src/mobius/_model_package.py index 94232a8f7..fd4f55ec3 100644 --- a/src/mobius/_model_package.py +++ b/src/mobius/_model_package.py @@ -369,6 +369,7 @@ def save_adapter_artifacts(self, directory: str) -> dict[str, dict[str, object]] "location": relative_location, "loader_capability": "onnx-genai.adapters.json@1", "sha256": hashlib.sha256(payload).hexdigest(), + "scale_encoding": "alpha_over_rank", "format": "json", } ) @@ -388,6 +389,7 @@ def save_adapter_artifacts(self, directory: str) -> dict[str, dict[str, object]] "location": relative_location, "loader_capability": "onnxruntime.lora-adapter@1", "sha256": hashlib.sha256(payload).hexdigest(), + "scale_encoding": "baked", "format": "ort_genai", } ) @@ -414,6 +416,7 @@ def save_adapter_artifacts(self, directory: str) -> dict[str, dict[str, object]] "sha256": hashlib.sha256(payload).hexdigest(), "config_location": relative_config, "config_sha256": hashlib.sha256(config_payload).hexdigest(), + "scale_encoding": "alpha_over_rank", "format": "hf_peft", } ) @@ -426,13 +429,11 @@ def save_adapter_artifacts(self, directory: str) -> dict[str, dict[str, object]] raise ValueError( f"adapter {alias!r} must emit a portable or preserved source artifact" ) - provenance_parts = [artifact.source.format] + provenance = {"producer": artifact.source.producer} if artifact.source.base_model: - provenance_parts.append(f"base={artifact.source.base_model}") + provenance["source"] = artifact.source.base_model if artifact.source.revision: - provenance_parts.append(f"revision={artifact.source.revision}") - if artifact.source.checksum: - provenance_parts.append(f"source_{artifact.source.checksum}") + provenance["revision"] = artifact.source.revision catalog[alias] = { "index": artifact_index, "identity": artifact.stable_identity, @@ -441,7 +442,7 @@ def save_adapter_artifacts(self, directory: str) -> dict[str, dict[str, object]] "rank": rank, "alpha": alpha, "dtype": dtype, - "provenance": ";".join(provenance_parts), + "provenance": provenance, "weights": weight_artifacts, "bindings": bindings, } @@ -466,7 +467,8 @@ def adapter_target_manifest_metadata(self) -> dict[str, object]: "id": descriptor.semantic_name, "component": descriptor.target.component, "parameter": descriptor.target.parameter, - "output_value": descriptor.output_name, + "node_name": descriptor.node_name, + "output_name": descriptor.output_name, "activation_dtype": activation_dtype, "input_features": descriptor.input_size, "output_features": descriptor.output_size, @@ -484,6 +486,7 @@ def adapter_target_manifest_metadata(self) -> dict[str, object]: sliced = dict(base) sliced["id"] = f"{descriptor.semantic_name}.{target_slice.role}" sliced["output_slice"] = { + "role": target_slice.role, "offset": target_slice.offset, "width": target_slice.width, } diff --git a/src/mobius/adapters.py b/src/mobius/adapters.py index 9a7623821..470ecc4d5 100644 --- a/src/mobius/adapters.py +++ b/src/mobius/adapters.py @@ -147,7 +147,8 @@ def _target_fingerprint_record( } if isinstance(resolved_target, AdapterTargetDescriptor): record["id"] = resolved_target.semantic_name - record["output_value"] = resolved_target.output_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"] = { @@ -196,6 +197,7 @@ def fingerprint_model_weights( 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, } @@ -377,8 +379,11 @@ class AdapterSource: 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:"): diff --git a/src/mobius/adapters_test.py b/src/mobius/adapters_test.py index 265b581d3..d5a646a90 100644 --- a/src/mobius/adapters_test.py +++ b/src/mobius/adapters_test.py @@ -435,6 +435,12 @@ def test_peft_migration_source_preserves_rank_alpha_and_provenance() -> None: 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", @@ -501,6 +507,7 @@ def test_onnx_adapter_source_can_be_declared_for_native_capability() -> None: 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"] == [ { @@ -603,7 +610,8 @@ def test_exact_onnx_genai_catalog_and_portable_bundle_serialization() -> None: "id": "layers.0.self_attn.q_proj", "component": "decoder", "parameter": "projection.weight", - "output_value": "projection.output", + "node_name": "projection", + "output_name": "projection.output", "activation_dtype": "float32", "input_features": 4, "output_features": 3, @@ -614,6 +622,7 @@ def test_exact_onnx_genai_catalog_and_portable_bundle_serialization() -> None: }, } assert service["target_manifest"]["targets"][1]["output_slice"] == { + "role": "q", "offset": 0, "width": 3, } @@ -639,6 +648,7 @@ def test_exact_onnx_genai_catalog_and_portable_bundle_serialization() -> None: 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", @@ -649,6 +659,7 @@ def test_exact_onnx_genai_catalog_and_portable_bundle_serialization() -> None: 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() diff --git a/tests/fixtures/onnx_genai_workflows/README.md b/tests/fixtures/onnx_genai_workflows/README.md index b66e98e54..66f4b998f 100644 --- a/tests/fixtures/onnx_genai_workflows/README.md +++ b/tests/fixtures/onnx_genai_workflows/README.md @@ -1,7 +1,7 @@ # ONNX GenAI workflow conformance fixtures Generated by `tests/generate_onnx_genai_validation_packages.py` for semantic -validation and runtime conformance against `justinchuby/onnx-genai@903a2d1a`. +validation and runtime conformance against `justinchuby/onnx-genai@d9482bca`. The decoder, VLM, diffusion, masked diffusion, speculative, and codec packages contain executable synthetic models. The TTS fixture uses the real tiny 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 index 761796c59..c55844d20 100644 --- a/tests/fixtures/onnx_genai_workflows/adapter/adapters/peft/adapter_config.json +++ b/tests/fixtures/onnx_genai_workflows/adapter/adapters/peft/adapter_config.json @@ -1 +1 @@ -{"base_model_name_or_path":"synthetic/adapter-base","lora_alpha":1.0,"r":1,"target_modules":["projection"]} \ No newline at end of file +{"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/inference_metadata.yaml b/tests/fixtures/onnx_genai_workflows/adapter/inference_metadata.yaml index 3cdd8fa51..9f3f070c1 100644 --- a/tests/fixtures/onnx_genai_workflows/adapter/inference_metadata.yaml +++ b/tests/fixtures/onnx_genai_workflows/adapter/inference_metadata.yaml @@ -157,13 +157,14 @@ pipeline: output: result mode: replace adapters: - base_model_fingerprint: onnx-genai-targeted-base-v1:sha256:ece1984372c107109d5e02743e04b6eefd452bfb8551a7b1e65bd7216871fcda + base_model_fingerprint: onnx-genai-targeted-base-v1:sha256:e7001822e692c6c2e51b0665bc105db3a4ae9a517d90c116e6f3dca1ee992316 target_manifest: targets: - id: projection component: decoder parameter: projection - output_value: projection.output + node_name: projection + output_name: projection.output activation_dtype: float32 input_features: 2 output_features: 2 @@ -194,15 +195,17 @@ adapters: index: 0 identity: blue version: '1' - base_model_fingerprint: onnx-genai-targeted-base-v1:sha256:ece1984372c107109d5e02743e04b6eefd452bfb8551a7b1e65bd7216871fcda + base_model_fingerprint: onnx-genai-targeted-base-v1:sha256:e7001822e692c6c2e51b0665bc105db3a4ae9a517d90c116e6f3dca1ee992316 rank: 1 alpha: 1.0 dtype: float32 - provenance: in_memory + 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 @@ -211,15 +214,17 @@ adapters: index: 1 identity: green version: '1' - base_model_fingerprint: onnx-genai-targeted-base-v1:sha256:ece1984372c107109d5e02743e04b6eefd452bfb8551a7b1e65bd7216871fcda + base_model_fingerprint: onnx-genai-targeted-base-v1:sha256:e7001822e692c6c2e51b0665bc105db3a4ae9a517d90c116e6f3dca1ee992316 rank: 1 alpha: 1.0 dtype: float32 - provenance: in_memory + 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 @@ -228,21 +233,26 @@ adapters: index: 2 identity: peft version: '1' - base_model_fingerprint: onnx-genai-targeted-base-v1:sha256:ece1984372c107109d5e02743e04b6eefd452bfb8551a7b1e65bd7216871fcda + base_model_fingerprint: onnx-genai-targeted-base-v1:sha256:e7001822e692c6c2e51b0665bc105db3a4ae9a517d90c116e6f3dca1ee992316 rank: 1 alpha: 1.0 dtype: float32 - provenance: peft_safetensors;base=synthetic/adapter-base;source_sha256:1be2eedb9402accd480873439cd9d4cb2b5542de9369985b541c5c488c6be935 + 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: e3ce70ff5b09c7ddcbcd9426fd0e7a867c3c8433cb15b6c5dd5f18a3c5630a55 + config_sha256: b8419ce55415ad1be8844bd6eaef9a873881f172f427694b1d8e46049a8784df + scale_encoding: alpha_over_rank format: hf_peft bindings: - target: projection @@ -251,15 +261,17 @@ adapters: index: 3 identity: red version: '1' - base_model_fingerprint: onnx-genai-targeted-base-v1:sha256:ece1984372c107109d5e02743e04b6eefd452bfb8551a7b1e65bd7216871fcda + base_model_fingerprint: onnx-genai-targeted-base-v1:sha256:e7001822e692c6c2e51b0665bc105db3a4ae9a517d90c116e6f3dca1ee992316 rank: 1 alpha: 1.0 dtype: float32 - provenance: in_memory + 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 diff --git a/tests/generate_onnx_genai_validation_packages.py b/tests/generate_onnx_genai_validation_packages.py index dad99306a..b32d695a3 100644 --- a/tests/generate_onnx_genai_validation_packages.py +++ b/tests/generate_onnx_genai_validation_packages.py @@ -558,6 +558,7 @@ def _adapter_package(source_root: Path) -> ModelPackage: "r": 1, "lora_alpha": 1.0, "target_modules": ["projection"], + "revision": "synthetic-revision", }, sort_keys=True, separators=(",", ":"), @@ -776,7 +777,7 @@ def main() -> None: """# ONNX GenAI workflow conformance fixtures Generated by `tests/generate_onnx_genai_validation_packages.py` for semantic -validation and runtime conformance against `justinchuby/onnx-genai@903a2d1a`. +validation and runtime conformance against `justinchuby/onnx-genai@d9482bca`. The decoder, VLM, diffusion, masked diffusion, speculative, and codec packages contain executable synthetic models. The TTS fixture uses the real tiny From 32eaca6beec7492519792451177ce1a9a8a07790 Mon Sep 17 00:00:00 2001 From: justinchuby Date: Sat, 15 Aug 2026 09:47:25 +0000 Subject: [PATCH 105/151] Align adapter targets with published ABI Pin ONNX GenAI 793bfe9b, serialize exact initializer and layer metadata, preserve target and fused-slice rank/alpha policies, and reject artifact bindings that violate the authoritative manifest. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: justinchuby --- .github/workflows/main.yml | 2 +- src/mobius/_model_package.py | 45 ++++++++++++++++--- src/mobius/adapters.py | 2 +- src/mobius/adapters_test.py | 34 +++++++++++++- tests/fixtures/onnx_genai_workflows/README.md | 2 +- .../adapter/inference_metadata.yaml | 15 ++++--- ...generate_onnx_genai_validation_packages.py | 5 ++- 7 files changed, 89 insertions(+), 16 deletions(-) diff --git a/.github/workflows/main.yml b/.github/workflows/main.yml index 0cff255a6..f52c039cf 100644 --- a/.github/workflows/main.yml +++ b/.github/workflows/main.yml @@ -23,7 +23,7 @@ jobs: - uses: actions/checkout@v7 with: repository: justinchuby/onnx-genai - ref: d9482bca0ad6ccec907a2b49faf033c2006a9fa9 + ref: 793bfe9b5489257a8ae4126ba66b760e3b2cfe19 path: validation/onnx-genai - uses: actions/setup-python@v7 with: diff --git a/src/mobius/_model_package.py b/src/mobius/_model_package.py index fd4f55ec3..7f386ae52 100644 --- a/src/mobius/_model_package.py +++ b/src/mobius/_model_package.py @@ -298,8 +298,9 @@ def save_adapter_artifacts(self, directory: str) -> dict[str, dict[str, object]] descriptor.target: descriptor for descriptor in self.adapter_target_manifest.targets } - manifest_target_ids = { - target["id"] for target in self.adapter_target_manifest_metadata()["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() @@ -335,11 +336,34 @@ def save_adapter_artifacts(self, directory: str) -> dict[str, dict[str, object]] for weight in ordered_weights: descriptor = descriptors[weight.target] target_id = weight.target_id or descriptor.semantic_name - if target_id not in manifest_target_ids: + 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, @@ -466,13 +490,19 @@ def adapter_target_manifest_metadata(self) -> dict[str, object]: base: dict[str, object] = { "id": descriptor.semantic_name, "component": descriptor.target.component, - "parameter": descriptor.target.parameter, + "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, @@ -485,11 +515,16 @@ def adapter_target_manifest_metadata(self) -> dict[str, object]: for target_slice in descriptor.slices: sliced = dict(base) sliced["id"] = f"{descriptor.semantic_name}.{target_slice.role}" - sliced["output_slice"] = { + 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} diff --git a/src/mobius/adapters.py b/src/mobius/adapters.py index 470ecc4d5..3fa8aa7ff 100644 --- a/src/mobius/adapters.py +++ b/src/mobius/adapters.py @@ -141,7 +141,7 @@ def _target_fingerprint_record( "component": target.component, "consumers": consumers, "dtype": int(initializer.dtype), - "parameter": target.parameter, + "initializer": target.parameter, "shape": [int(dimension) for dimension in initializer.shape], "tensor_sha256": hashlib.sha256(_tensor_bytes(initializer.const_value)).hexdigest(), } diff --git a/src/mobius/adapters_test.py b/src/mobius/adapters_test.py index d5a646a90..61281342d 100644 --- a/src/mobius/adapters_test.py +++ b/src/mobius/adapters_test.py @@ -609,7 +609,8 @@ def test_exact_onnx_genai_catalog_and_portable_bundle_serialization() -> None: assert target == { "id": "layers.0.self_attn.q_proj", "component": "decoder", - "parameter": "projection.weight", + "initializer": "projection.weight", + "layer_index": 0, "node_name": "projection", "output_name": "projection.output", "activation_dtype": "float32", @@ -625,6 +626,8 @@ def test_exact_onnx_genai_catalog_and_portable_bundle_serialization() -> None: "role": "q", "offset": 0, "width": 3, + "rank": 2, + "alpha": 4.0, } assert service["cache"] == {"max_entries": 2, "eviction": "lru"} assert service["planning"] == { @@ -746,6 +749,35 @@ def test_wire_contract_emits_per_target_rank_override() -> None: 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) diff --git a/tests/fixtures/onnx_genai_workflows/README.md b/tests/fixtures/onnx_genai_workflows/README.md index 66f4b998f..669a0ac0c 100644 --- a/tests/fixtures/onnx_genai_workflows/README.md +++ b/tests/fixtures/onnx_genai_workflows/README.md @@ -1,7 +1,7 @@ # ONNX GenAI workflow conformance fixtures Generated by `tests/generate_onnx_genai_validation_packages.py` for semantic -validation and runtime conformance against `justinchuby/onnx-genai@d9482bca`. +validation and runtime conformance against `justinchuby/onnx-genai@793bfe9b`. The decoder, VLM, diffusion, masked diffusion, speculative, and codec packages contain executable synthetic models. The TTS fixture uses the real tiny diff --git a/tests/fixtures/onnx_genai_workflows/adapter/inference_metadata.yaml b/tests/fixtures/onnx_genai_workflows/adapter/inference_metadata.yaml index 9f3f070c1..bf601d63a 100644 --- a/tests/fixtures/onnx_genai_workflows/adapter/inference_metadata.yaml +++ b/tests/fixtures/onnx_genai_workflows/adapter/inference_metadata.yaml @@ -157,17 +157,20 @@ pipeline: output: result mode: replace adapters: - base_model_fingerprint: onnx-genai-targeted-base-v1:sha256:e7001822e692c6c2e51b0665bc105db3a4ae9a517d90c116e6f3dca1ee992316 + base_model_fingerprint: onnx-genai-targeted-base-v1:sha256:702704827be590661fd77676445ddaf1557e4f4d446eaa7bcfb78754466cccdb target_manifest: targets: - id: projection component: decoder - parameter: projection + 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 @@ -195,7 +198,7 @@ adapters: index: 0 identity: blue version: '1' - base_model_fingerprint: onnx-genai-targeted-base-v1:sha256:e7001822e692c6c2e51b0665bc105db3a4ae9a517d90c116e6f3dca1ee992316 + base_model_fingerprint: onnx-genai-targeted-base-v1:sha256:702704827be590661fd77676445ddaf1557e4f4d446eaa7bcfb78754466cccdb rank: 1 alpha: 1.0 dtype: float32 @@ -214,7 +217,7 @@ adapters: index: 1 identity: green version: '1' - base_model_fingerprint: onnx-genai-targeted-base-v1:sha256:e7001822e692c6c2e51b0665bc105db3a4ae9a517d90c116e6f3dca1ee992316 + base_model_fingerprint: onnx-genai-targeted-base-v1:sha256:702704827be590661fd77676445ddaf1557e4f4d446eaa7bcfb78754466cccdb rank: 1 alpha: 1.0 dtype: float32 @@ -233,7 +236,7 @@ adapters: index: 2 identity: peft version: '1' - base_model_fingerprint: onnx-genai-targeted-base-v1:sha256:e7001822e692c6c2e51b0665bc105db3a4ae9a517d90c116e6f3dca1ee992316 + base_model_fingerprint: onnx-genai-targeted-base-v1:sha256:702704827be590661fd77676445ddaf1557e4f4d446eaa7bcfb78754466cccdb rank: 1 alpha: 1.0 dtype: float32 @@ -261,7 +264,7 @@ adapters: index: 3 identity: red version: '1' - base_model_fingerprint: onnx-genai-targeted-base-v1:sha256:e7001822e692c6c2e51b0665bc105db3a4ae9a517d90c116e6f3dca1ee992316 + base_model_fingerprint: onnx-genai-targeted-base-v1:sha256:702704827be590661fd77676445ddaf1557e4f4d446eaa7bcfb78754466cccdb rank: 1 alpha: 1.0 dtype: float32 diff --git a/tests/generate_onnx_genai_validation_packages.py b/tests/generate_onnx_genai_validation_packages.py index b32d695a3..908e237a0 100644 --- a/tests/generate_onnx_genai_validation_packages.py +++ b/tests/generate_onnx_genai_validation_packages.py @@ -504,6 +504,9 @@ def _adapter_package(source_root: Path) -> ModelPackage: 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", @@ -777,7 +780,7 @@ def main() -> None: """# ONNX GenAI workflow conformance fixtures Generated by `tests/generate_onnx_genai_validation_packages.py` for semantic -validation and runtime conformance against `justinchuby/onnx-genai@d9482bca`. +validation and runtime conformance against `justinchuby/onnx-genai@793bfe9b`. The decoder, VLM, diffusion, masked diffusion, speculative, and codec packages contain executable synthetic models. The TTS fixture uses the real tiny From fa868cecf4c7cd1f2f79148d0463c97c14edc53e Mon Sep 17 00:00:00 2001 From: justinchuby Date: Sat, 15 Aug 2026 09:56:48 +0000 Subject: [PATCH 106/151] Test heterogeneous PEFT binding overrides Verify one adapter alias preserves PEFT rank_pattern and alpha_pattern as per-target binding overrides without splitting selection or artifact identity. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: justinchuby --- src/mobius/adapters_test.py | 84 +++++++++++++++++++++++++++++++++++++ 1 file changed, 84 insertions(+) diff --git a/src/mobius/adapters_test.py b/src/mobius/adapters_test.py index 61281342d..7f1b18648 100644 --- a/src/mobius/adapters_test.py +++ b/src/mobius/adapters_test.py @@ -456,6 +456,90 @@ def test_peft_migration_source_preserves_rank_alpha_and_provenance() -> None: 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) From 219f77e596df2b60971fe21ee2bf02aafabb8d25 Mon Sep 17 00:00:00 2001 From: justinchuby Date: Sat, 15 Aug 2026 10:02:19 +0000 Subject: [PATCH 107/151] Pin executable heterogeneous adapter regression Validate Mobius fixtures and PEFT binding overrides against ONNX GenAI 2af34dca, which executes distinct effective rank and alpha values under one adapter alias. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: justinchuby --- .github/workflows/main.yml | 2 +- tests/fixtures/onnx_genai_workflows/README.md | 2 +- tests/generate_onnx_genai_validation_packages.py | 2 +- 3 files changed, 3 insertions(+), 3 deletions(-) diff --git a/.github/workflows/main.yml b/.github/workflows/main.yml index f52c039cf..bc4184aa2 100644 --- a/.github/workflows/main.yml +++ b/.github/workflows/main.yml @@ -23,7 +23,7 @@ jobs: - uses: actions/checkout@v7 with: repository: justinchuby/onnx-genai - ref: 793bfe9b5489257a8ae4126ba66b760e3b2cfe19 + ref: 2af34dcad0e429604ca66d7ba1388ee2e688756e path: validation/onnx-genai - uses: actions/setup-python@v7 with: diff --git a/tests/fixtures/onnx_genai_workflows/README.md b/tests/fixtures/onnx_genai_workflows/README.md index 669a0ac0c..bd380598a 100644 --- a/tests/fixtures/onnx_genai_workflows/README.md +++ b/tests/fixtures/onnx_genai_workflows/README.md @@ -1,7 +1,7 @@ # ONNX GenAI workflow conformance fixtures Generated by `tests/generate_onnx_genai_validation_packages.py` for semantic -validation and runtime conformance against `justinchuby/onnx-genai@793bfe9b`. +validation and runtime conformance against `justinchuby/onnx-genai@2af34dca`. The decoder, VLM, diffusion, masked diffusion, speculative, and codec packages contain executable synthetic models. The TTS fixture uses the real tiny diff --git a/tests/generate_onnx_genai_validation_packages.py b/tests/generate_onnx_genai_validation_packages.py index 908e237a0..33791845c 100644 --- a/tests/generate_onnx_genai_validation_packages.py +++ b/tests/generate_onnx_genai_validation_packages.py @@ -780,7 +780,7 @@ def main() -> None: """# ONNX GenAI workflow conformance fixtures Generated by `tests/generate_onnx_genai_validation_packages.py` for semantic -validation and runtime conformance against `justinchuby/onnx-genai@793bfe9b`. +validation and runtime conformance against `justinchuby/onnx-genai@2af34dca`. The decoder, VLM, diffusion, masked diffusion, speculative, and codec packages contain executable synthetic models. The TTS fixture uses the real tiny From 2a46e924919c1c408cf2183cee212a3c99e7f6a7 Mon Sep 17 00:00:00 2001 From: Justin Chu Date: Thu, 20 Aug 2026 05:26:25 +0000 Subject: [PATCH 108/151] Express encoder-conditioned decode loops and align serving with the runtime schema Whisper-style speech-to-text is an autoregressive decoder conditioned on an encoder result that never changes during decoding. That shape was not expressible: the workflow builder assumed a single decoder component, a mandatory attention mask, and a KV cache that is the only state the loop carries. Three generic changes make it expressible: * `_build_autoregressive_workflow_metadata` now accepts an optional encoder. Its inputs become externally suppliable workflow inputs, it is invoked in the loop setup, and every decoder input that structurally matches an encoder output becomes a shape-invariant, request-aligned state cell with an identity carry. Cross state is therefore not a special case: it is ordinary loop-invariant state, which is exactly what keeps encoder rows and decoder rows aligned when the runtime compacts a batch. * The attention mask is now optional. A shared full-capacity KV buffer needs a mask to convey each row's logical length, so fixed-capacity KV is gated on the mask being present and mask-free decoders get growable caches instead. `build_decoder_state_initializer` and `build_decoder_step_update` follow. * `serving` now emits `state_service` groups with the semantic `kind`, `aliasing`, `reuse`, and `capabilities` the runtime schema requires, instead of the removed `kv_service`/`slot_ids` shape. Aliasing is derived from the admitted graph: a shared buffer may alias past onto present, a growable cache must not. `declare_request_alignment` stamps the request-aligned row axis onto every batch-leading contract in state, ports, inputs, and outputs after policy components are attached. Without it the runtime cannot permute a tensor when it compacts finished rows, so state and results would silently drift apart. Runtime-managed cells also gain the release boundary that externally owned state needs, since it has no SSA liveness to free it. `genai_config_import` performs the one-way import of an existing `genai_config.json` package into this workflow IR. The import is structural: an encoder-conditioned package is recognised because the config declares an encoder whose outputs the decoder consumes, not because of a model name. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: Justin Chu (cherry picked from commit 24373a38337d68ce43b0dc0761586480afa6fd70) --- src/mobius/generation/_policy_components.py | 159 +++--- .../onnx_genai/auto_export_test.py | 15 +- .../codec_workflow_metadata_test.py | 11 +- .../onnx_genai/genai_config_import.py | 205 ++++++++ .../onnx_genai/inference_metadata.py | 41 ++ .../onnx_genai/inference_metadata_test.py | 2 + .../onnx_genai/workflow_metadata.py | 473 ++++++++++++++---- .../onnx_genai/workflow_metadata_test.py | 18 +- 8 files changed, 729 insertions(+), 195 deletions(-) create mode 100644 src/mobius/integrations/onnx_genai/genai_config_import.py diff --git a/src/mobius/generation/_policy_components.py b/src/mobius/generation/_policy_components.py index 90587a639..04fd84bfa 100644 --- a/src/mobius/generation/_policy_components.py +++ b/src/mobius/generation/_policy_components.py @@ -621,13 +621,18 @@ def build_decoder_state_initializer( *, token_input: str | None, prompt_dtype: ir.DataType | None = None, - attention_mask_input: str, + attention_mask_input: str | None, position_ids_input: str | None, cache_inputs: list[str], fixed_capacity: bool = False, ragged: bool = False, ) -> PolicyComponent: """Build prompt-derived decoder state, optionally with capture-stable storage.""" + 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} @@ -687,41 +692,44 @@ def build_decoder_state_initializer( attention_shape, ) - 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]), - ), + 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, - ), - axis=1, - ) - body_attention.shape = ir.Shape(["batch", "prompt_sequence + 1"]) + ) + 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), @@ -747,11 +755,15 @@ def build_decoder_state_initializer( [-1], ) body_position.shape = ir.Shape(["batch", 1]) - builder.add_output(attention, attention_mask_input) + 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) - builder.add_output(body_attention, "body_attention_mask") + 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") @@ -813,46 +825,53 @@ def build_decoder_state_initializer( def build_decoder_step_update( *, - attention_dtype: ir.DataType, + attention_dtype: ir.DataType | None, position_dtype: ir.DataType | None, fixed_capacity: bool = False, ) -> 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 - 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"], + if attention_dtype is not None: + attention = builder.input( + "attention_mask", + dtype=attention_dtype, + shape=["batch", "context"], ) - 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 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: position = builder.input( "position_ids", diff --git a/src/mobius/integrations/onnx_genai/auto_export_test.py b/src/mobius/integrations/onnx_genai/auto_export_test.py index 1c2bef966..1ea0c05f2 100644 --- a/src/mobius/integrations/onnx_genai/auto_export_test.py +++ b/src/mobius/integrations/onnx_genai/auto_export_test.py @@ -208,7 +208,12 @@ def test_dispatch_decoder(tmp_path): 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]}, + "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"}, @@ -245,11 +250,13 @@ def test_seeded_decoder_sampler_uses_request_controls_and_direct_kv_carry(): "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 @@ -277,11 +284,8 @@ def test_seeded_decoder_sampler_uses_request_controls_and_direct_kv_carry(): ] 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 emit["row_ids"] == "slot_ids" - assert "emit_row_identity" in workflow["manifest"]["capabilities"] assert set( - workflow["serving"]["kv_service"]["groups"]["decoder_cache"]["ports"]["model"] + 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"] == { @@ -530,6 +534,7 @@ def test_dispatch_vision_multimodal_pipeline(tmp_path): "dtype": "float32", "rank": 2, "shape": ["batch", 128], + "batch_layout": {"kind": "request_aligned", "axis": 0}, } assert workflow["state"]["logits"]["initializer"] == "decoder.setup.last_logits" assert (tmp_path / "policies" / "token_sampler.onnx").is_file() diff --git a/src/mobius/integrations/onnx_genai/codec_workflow_metadata_test.py b/src/mobius/integrations/onnx_genai/codec_workflow_metadata_test.py index 436ca93be..abc41ad0a 100644 --- a/src/mobius/integrations/onnx_genai/codec_workflow_metadata_test.py +++ b/src/mobius/integrations/onnx_genai/codec_workflow_metadata_test.py @@ -51,6 +51,7 @@ def test_codec_workflow_has_typed_ssa_and_audio_emit(): "dtype": "float32", "rank": 3, "shape": ["batch", 1, "audio_samples"], + "batch_layout": {"kind": "request_aligned", "axis": 0}, } assert "ports" not in workflow["components"]["encoder"] assert "effects" not in workflow["components"]["encoder"] @@ -177,14 +178,18 @@ def test_real_qwen3_tts_workflow_carries_trained_transitions_and_kv_state(): 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"]["kv_service"]["compaction"] is True assert workflow["inputs"]["package.slot_ids"] == { - "contract": {"dtype": "int64", "rank": 1, "shape": ["batch"]}, + "contract": { + "dtype": "int64", + "rank": 1, + "shape": ["batch"], + "batch_layout": {"kind": "request_aligned", "axis": 0}, + }, "role": {"kind": "opaque"}, "source": {"kind": "application", "name": "serving.slot_ids"}, "required": True, } - assert workflow["serving"]["kv_service"]["groups"]["talker_cache"]["ports"]["talker"][ + 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" 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..f89d20a06 --- /dev/null +++ b/src/mobius/integrations/onnx_genai/genai_config_import.py @@ -0,0 +1,205 @@ +# 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 9ed0b8e2d..4aceb8fd4 100644 --- a/src/mobius/integrations/onnx_genai/inference_metadata.py +++ b/src/mobius/integrations/onnx_genai/inference_metadata.py @@ -1479,9 +1479,50 @@ def tensor_contract(value: Any) -> dict[str, Any]: if component.contract.get("role") == "token_sampler": declaration["application_overridable"] = True components[name] = declaration + declare_request_alignment(workflow) return metadata +_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, diff --git a/src/mobius/integrations/onnx_genai/inference_metadata_test.py b/src/mobius/integrations/onnx_genai/inference_metadata_test.py index 2c7f7312f..f7ae2e3b8 100644 --- a/src/mobius/integrations/onnx_genai/inference_metadata_test.py +++ b/src/mobius/integrations/onnx_genai/inference_metadata_test.py @@ -156,6 +156,7 @@ def test_workflow_policy_components_reference_saved_onnx_artifacts(tmp_path): "dtype": "float32", "rank": 2, "shape": ["batch", "vocabulary"], + "batch_layout": {"kind": "request_aligned", "axis": 0}, } }, "outputs": { @@ -163,6 +164,7 @@ def test_workflow_policy_components_reference_saved_onnx_artifacts(tmp_path): "dtype": "int64", "rank": 1, "shape": ["batch"], + "batch_layout": {"kind": "request_aligned", "axis": 0}, } }, } diff --git a/src/mobius/integrations/onnx_genai/workflow_metadata.py b/src/mobius/integrations/onnx_genai/workflow_metadata.py index 6169bbe95..539acfb29 100644 --- a/src/mobius/integrations/onnx_genai/workflow_metadata.py +++ b/src/mobius/integrations/onnx_genai/workflow_metadata.py @@ -49,6 +49,7 @@ add_adapter_service_to_metadata, add_policy_components_to_workflow, build_native_vlm_package_metadata, + declare_request_alignment, ) @@ -236,8 +237,6 @@ def convert(node: dict[str, Any]) -> dict[str, Any]: result["valid_length"] = rewrite(node["valid_length"]) if "when" in node: result["when"] = rewrite(node["when"]) - if "row_ids" in node: - result["row_ids"] = rewrite(node["row_ids"]) return result if kind == "branch": result = { @@ -316,9 +315,12 @@ def convert(node: dict[str, Any]) -> dict[str, Any]: collect_carried(graph) published = convert(graph) workflow["steps"] = published["steps"] if published["kind"] == "sequence" else [published] + declare_request_alignment(workflow) return workflow + + def _name_image_preprocessing_program(image: dict[str, Any]) -> None: """Convert structural preprocessing transforms into explicit typed SSA values.""" transforms = image["transforms"] @@ -502,6 +504,7 @@ def _model_cache_pairs(model: ir.Model) -> list[tuple[ir.Value, ir.Value]]: for name in ( past.name.replace("past_key_values", "present"), past.name.replace("past.", "present."), + past.name.replace("past_", "present_"), ) if name in outputs ), @@ -534,6 +537,41 @@ def _kv_storage_contract(model: ir.Model) -> dict[str, Any]: } +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 ("permitted" if storage in ("shared_buffer", "paged") else "forbidden") + ), + "reuse": {"prefix_reusable": True, "evictable_prefix": False}, + "capabilities": {"snapshot": True, "fork": True}, + "ports": ports, + } + if logical_lengths is not None: + group["logical_lengths"] = logical_lengths + return 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"] @@ -1625,25 +1663,15 @@ def frame_generation_nodes(prefix: str, hidden: str, logits: str) -> list[dict[s "active": "active", "done": "done", "accepted_len": "accepted_len", - "slot_ids": "slot_ids", - "kv_service": { - "paging": ( - "paged" - if talker_kv["paging"] == "paged" - or predictor_kv["paging"] == "paged" - else "none" - ), - "allocation": "runtime", - "compaction": (talker_kv["compaction"] or predictor_kv["compaction"]), + "state_service": { "groups": { **( { - "talker_cache": { - "sequence_axis": 2, - "layout": "bnsh", - "logical_lengths": "talker_cache_lengths", - "storage": talker_kv["storage"], - "ports": { + "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, @@ -1654,19 +1682,18 @@ def frame_generation_nodes(prefix: str, hidden: str, logits: str) -> list[dict[s ) } }, - } + ) } if talker_caches else {} ), **( { - "predictor_cache": { - "sequence_axis": 2, - "layout": "bnsh", - "logical_lengths": "predictor_cache_lengths", - "storage": predictor_kv["storage"], - "ports": { + "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, @@ -1677,7 +1704,7 @@ def frame_generation_nodes(prefix: str, hidden: str, logits: str) -> list[dict[s ) } }, - } + ) } if predictor_caches else {} @@ -3422,7 +3449,6 @@ def build_vlm_workflow_metadata( "mode": "append", "when": "state.active.body", "valid_length": "token.emitted_length", - "row_ids": "state.slot_ids.body", "effect_name": "emit", "effect": _effect("emit.0", "emit.1"), }, @@ -3473,7 +3499,6 @@ def build_vlm_workflow_metadata( "loop_induction_values", "typed_emit", "emit_valid_length", - "emit_row_identity", *(["input_presence"] if text_only_vision is not None else []), *( ["serving_service_contract", "bounded_state_recurrence"] @@ -3502,14 +3527,10 @@ def build_vlm_workflow_metadata( "active": "active", "done": "done", "accepted_len": "accepted_len", - "slot_ids": "slot_ids", - "kv_service": { - "paging": decoder_kv["paging"], - "allocation": "runtime", - "compaction": decoder_kv["compaction"], + "state_service": { "groups": { - "decoder_cache": { - "sequence_axis": next( + "decoder_cache": _state_group( + sequence_axis=next( ( axis for axis, dimension in enumerate( @@ -3519,10 +3540,9 @@ def build_vlm_workflow_metadata( ), 2, ), - "layout": "bnsh", - "logical_lengths": "cache_lengths", - "storage": decoder_kv["storage"], - "ports": { + logical_lengths="cache_lengths", + storage=decoder_kv["storage"], + ports={ "decoder": { f"cache_{index}": { "input": past.name, @@ -3531,7 +3551,7 @@ def build_vlm_workflow_metadata( for index, (past, present) in enumerate(cache_pairs) } }, - } + ) }, }, } @@ -4048,7 +4068,6 @@ def build_speculative_workflow_metadata( "valid_length": emit_length, "output": "tokens", "mode": "append", - "row_ids": "state.slot_ids.body", "effect_name": "emit", "effect": _effect("emit.0", "emit.1"), }, @@ -4061,7 +4080,6 @@ def build_speculative_workflow_metadata( "valid_length": "grammar.forced_length", "output": "tokens", "mode": "append", - "row_ids": "state.slot_ids.body", "effect_name": "emit", "effect": _effect("emit.1", "emit.2"), } @@ -4291,7 +4309,6 @@ def build_speculative_workflow_metadata( "loop_induction_values", "typed_emit", "emit_valid_length", - "emit_row_identity", "bounded_state_recurrence", "serving_service_contract", ], @@ -4318,19 +4335,14 @@ def build_speculative_workflow_metadata( "active": "active", "done": "done", "accepted_len": "accepted_len", - "slot_ids": "slot_ids", - "kv_service": { - "paging": verifier_kv["paging"], - "allocation": "runtime", - "compaction": verifier_kv["compaction"], + "state_service": { "groups": { - "verifier_cache": { - "sequence_axis": kv_sequence_axis, - "layout": "bnsh", - "logical_lengths": "cache_lengths", - "storage": verifier_kv["storage"], - "ports": {"verifier": kv_ports}, - } + "verifier_cache": _state_group( + sequence_axis=kv_sequence_axis, + logical_lengths="cache_lengths", + storage=verifier_kv["storage"], + ports={"verifier": kv_ports}, + ) }, }, }, @@ -4389,6 +4401,70 @@ def write_speculative_workflow_metadata( 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, @@ -4398,7 +4474,38 @@ def build_decoder_workflow_metadata( """Build the exact workflow-policy contract for an autoregressive decoder.""" if len(pkg) != 1: raise ValueError("decoder workflow requires exactly one neural component") - decoder_name, decoder = next(iter(pkg.items())) + 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) @@ -4455,9 +4562,12 @@ def build_decoder_workflow_metadata( 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 candidates = [ value.name.replace("past_key_values", "present"), value.name.replace("past.", "present."), + value.name.replace("past_", "present_"), ] present = next( (output_by_suffix.get(name) for name in candidates if name in output_by_suffix), @@ -4479,8 +4589,11 @@ def build_decoder_workflow_metadata( ( value for value in integer_rank2 - if "mask" in value.name - or "past" in str(getattr(list(value.shape)[1], "value", list(value.shape)[1])) + 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, ) @@ -4488,15 +4601,22 @@ def build_decoder_workflow_metadata( ( value for value in integer_rank2 - if value is not attention_input and "position" in value.name + 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, ), - next((value for value in integer_rank2 if value is not attention_input), None), ) - if attention_input is None: - raise ValueError( - "standard decoder workflow requires a derived rank-2 attention-mask input" - ) - derived_names = cache_names | {attention_input.name} + derived_names = cache_names | set(cross_bindings) + 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 = [ @@ -4506,27 +4626,43 @@ def build_decoder_workflow_metadata( ] if unsupported: raise ValueError(f"decoder workflow has unsupported non-request inputs: {unsupported}") - fixed_capacity = bool(cache_pairs) and decoder_kv_contract["storage"] == "shared_buffer" + # 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 + ) pkg.add_policy_component( "decoder_state_initializer", build_decoder_state_initializer( decoder, token_input=token_input.name, - attention_mask_input=attention_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), ), ) - 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, - ), - ) + 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, + ), + ) + 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, + ), + ) 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)) @@ -4557,6 +4693,35 @@ def build_decoder_workflow_metadata( 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] = {} + if encoder is not None: + assert encoder_name is not None + preprocessing_outputs = { + binding["name"]: binding + for binding in (audio_preprocessing or {}).get("outputs", []) + } + for value in encoder.graph.inputs: + ssa = f"encoder.input.{value.name}" + if value.name in preprocessing_outputs: + encoder_invoke_inputs[value.name] = preprocessing_outputs[value.name][ + "source_value" + ] + 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}" batch_dimension = _shape_metadata(_port(token_input))[0] batch_int = {"dtype": "int64", "rank": 1, "shape": [batch_dimension]} batch_bool = {"dtype": "bool", "rank": 1, "shape": [batch_dimension]} @@ -4789,13 +4954,22 @@ def build_decoder_workflow_metadata( 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}" - 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 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" + # 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" ) @@ -4980,8 +5154,9 @@ def build_decoder_workflow_metadata( "write_effect": _effect("state:rng_counter.read", "state:rng_counter.1"), } ) - decoder_state_specs = { - "attention_mask": ( + 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, @@ -5003,8 +5178,7 @@ def build_decoder_workflow_metadata( "max": "package.max_context", } ), - ), - } + ) if position_input is not None: decoder_state_specs["position_ids"] = ( { @@ -5080,9 +5254,45 @@ def build_decoder_workflow_metadata( } ) + # 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"), + } + ) + setup = { "kind": "sequence", "nodes": [ + *( + [_invoke(encoder_name, encoder_invoke_inputs, encoder_invoke_outputs)] + if encoder is not None + else [] + ), _invoke( "decoder_state_initializer", { @@ -5091,8 +5301,14 @@ def build_decoder_workflow_metadata( **({"max_iterations": "request.max_iterations"} if fixed_capacity else {}), }, { - attention_input.name: f"initializer.{attention_input.name}", - "body_attention_mask": "initializer.body_attention_mask", + **( + { + 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"} @@ -5144,10 +5360,15 @@ def build_decoder_workflow_metadata( ), ], } + 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", + **( + {"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"} @@ -5156,7 +5377,11 @@ def build_decoder_workflow_metadata( ), }, { - "next_attention_mask": "decoder_step.body_attention_mask", + **( + {"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 @@ -5313,27 +5538,29 @@ def build_decoder_workflow_metadata( { "when": "state.active.body", "valid_length": "token.emitted_length", - "row_ids": "state.slot_ids.body", - } + } if cache_pairs else {} ), "effect_name": "emit", "effect": _effect("emit.0", "emit.1"), }, - *([decoder_step_invoke] if fixed_capacity else []), + *([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"}, ), - *([] if fixed_capacity else [decoder_step_invoke]), + *([decoder_step_invoke] if has_step_update and not fixed_capacity else []), ], } + artifacts = artifacts or {} use_subfolders = len(pkg) > 1 - artifact = f"{decoder_name}/model.onnx" if use_subfolders else "model.onnx" + artifact = artifacts.get( + decoder_name, f"{decoder_name}/model.onnx" if use_subfolders else "model.onnx" + ) workflow = { "manifest": { "ir_version": "1.0", @@ -5344,7 +5571,6 @@ def build_decoder_workflow_metadata( "nested_control_flow", "typed_emit", "emit_valid_length", - *(["emit_row_identity"] if cache_pairs else []), "loop_induction_values", *(["serving_service_contract"] if cache_pairs else []), *(["bounded_state_recurrence"] if cache_pairs else []), @@ -5362,7 +5588,19 @@ def build_decoder_workflow_metadata( "stage": "pre_adapter", } }, - "components": {decoder_name: _component(decoder, artifact)}, + "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 {} + ), + }, "state": state, **( { @@ -5370,19 +5608,18 @@ def build_decoder_workflow_metadata( "active": "active", "done": "done", "accepted_len": "accepted_len", - "slot_ids": "slot_ids", - "kv_service": { - "paging": decoder_kv_contract["paging"], - "allocation": "runtime", - "compaction": decoder_kv_contract["compaction"], + "state_service": { "groups": { - "decoder_cache": { - "sequence_axis": decoder_kv_axis, - "layout": "bnsh", - "logical_lengths": "cache_lengths", - "storage": decoder_kv_contract["storage"], - "ports": {decoder_name: decoder_kv_ports}, - } + "decoder_cache": _state_group( + sequence_axis=decoder_kv_axis, + logical_lengths="cache_lengths" if fixed_capacity else None, + storage=( + decoder_kv_contract["storage"] + if fixed_capacity + else "growable" + ), + ports={decoder_name: decoder_kv_ports}, + ) }, }, } @@ -5694,6 +5931,30 @@ def write_decoder_workflow_metadata( 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, diff --git a/src/mobius/integrations/onnx_genai/workflow_metadata_test.py b/src/mobius/integrations/onnx_genai/workflow_metadata_test.py index 1777ddfa4..3e6d94471 100644 --- a/src/mobius/integrations/onnx_genai/workflow_metadata_test.py +++ b/src/mobius/integrations/onnx_genai/workflow_metadata_test.py @@ -48,9 +48,7 @@ def test_speculative_emit_uses_accepted_prefix_length(): ] emit = next(node for node in workflow["steps"][0]["steps"] if node["kind"] == "emit") assert emit["valid_length"] == "acceptance.length" - assert emit["row_ids"] == "slot_ids" assert "emit_valid_length" in workflow["manifest"]["capabilities"] - assert "emit_row_identity" in workflow["manifest"]["capabilities"] assert workflow["inputs"]["request.slot_ids"]["source"]["name"] == "serving.slot_ids" assert workflow["outputs"]["tokens"]["contract"]["shape"][-1] == "accepted_sequence" assert workflow["state"]["cache_0"]["recurrence"] == { @@ -236,11 +234,13 @@ def collect_decoder_invokes(node): "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}, } assert workflow["state"]["attention_mask"]["recurrence"] == {"kind": "invariant"} assert workflow["state"]["cache_103"]["recurrence"] == { @@ -292,8 +292,6 @@ def collect_emits(node): ) assert emit["when"] == "active" assert emit["valid_length"] == "token.emitted_length" - assert emit["row_ids"] == "slot_ids" - assert "emit_row_identity" 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( @@ -310,9 +308,7 @@ def collect_emits(node): assert media_branch["predicate"] == "request.image_present" assert set(media_branch["cases"]) == {"true", "false"} assert media_branch["cases"]["false"]["component"] == "empty_image_features" - kv_service = workflow["serving"]["kv_service"] - assert kv_service["paging"] == "none" - assert kv_service["compaction"] is True + state_service = workflow["serving"]["state_service"] carried = {item["cell"] for item in workflow["steps"][0]["carried"]} assert { "slot_ids", @@ -327,11 +323,12 @@ def collect_emits(node): "attention_mask", "cache_0", } <= carried - decoder_cache = kv_service["groups"]["decoder_cache"] + 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) - assert decoder_cache["storage"] == "shared_buffer" + assert decoder_cache["aliasing"] == "permitted" + assert decoder_cache["kind"] == "full_attention" kv_ports = decoder_cache["ports"]["decoder"] assert len(kv_ports) == 104 assert kv_ports["cache_103"] == { @@ -533,11 +530,10 @@ def test_speculative_workflow_uses_per_row_ragged_state_and_rng(): 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 emit["row_ids"] == "slot_ids" assert not any(node["kind"] == "branch" for node in body) assert workflow["serving"]["active"] == "active" assert workflow["serving"]["done"] == "done" - assert workflow["serving"]["kv_service"]["groups"]["verifier_cache"]["ports"]["verifier"][ + assert workflow["serving"]["state_service"]["groups"]["verifier_cache"]["ports"]["verifier"][ "cache_0" ] == { "input": "past_key_values.0.key", From 16c93edf44841a12cc9e463893d5a5023c43b182 Mon Sep 17 00:00:00 2001 From: Justin Chu Date: Thu, 20 Aug 2026 05:20:22 +0000 Subject: [PATCH 109/151] Fix Gemma 4 text-only export and prefill-prefix pruning on the Attention path MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Two independent defects surfaced while running the shipped `google/gemma-4-E2B-it` checkpoint through the real `mobius build` CLI. 1. `--features text-only` was unusable for every shipped Gemma 4 checkpoint. Those configs declare `model_type="gemma4"` with a nested `text_config` of `model_type="gemma4_text"`, but neither spelling had an entry in `_TEXT_ONLY_MODEL_TYPE`, so `build(..., text_only=True)` raised before it could resolve the registered `Gemma4CausalLMModel` backbone. Add both the multimodal key and the idempotent self-mapping, matching how `gemma3n` and `muse_glimmer` are wired. 2. `--features prune-prefill-prefix` produced a graph that could not run at all on the default (opset-24 `Attention` + `RotaryEmbedding`) path. Gemma 4 prunes mid-stack: at the first KV-shared layer the hidden states narrow to a single query position, because the shared tail borrows K/V from earlier layers and only the final row reaches `lm_head`. The RoPE `(cos, sin)` caches and the additive attention bias were left at the full prompt length, so `RotaryEmbedding` rejected the mismatch (`Inputs 'cos_cache' and 'sin_cache' are expected to have the same shape as input 'x', got 1 and 21` at `model/layers.15/self_attn`). Narrow both alongside the hidden states. The fused `GroupQueryAttention` path is unaffected — it derives its rotary offset from `total_seq_len - q_len` and takes a full-length cache — which is why the existing WebGPU graph-shape tests never caught this. Both fixes are verified numerically, not structurally: * new `test_pruned_prefill_matches_unpruned_final_row` builds the same tiny hybrid config twice (pruned/unpruned), copies the weights across by name and runs both under ORT, asserting the pruned logits equal the unpruned final row and that every cache-owning layer's present K/V is unchanged. It covers the `default` and `cpu` execution providers and a sequence length past the sliding window, and fails on the pre-fix code. * on the real 4.65 B checkpoint, the pruned f32 export now matches HuggingFace with `max_abs=7.06e-05` on the prefill logits and produces identical greedy tokens over four cached decode steps. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: Justin Chu (cherry picked from commit 0c8bf8f10fc5f2ab1ba87c1cba9a9eb262afefa2) --- src/mobius/_registry.py | 6 ++ src/mobius/_registry_test.py | 49 +++++++++ src/mobius/models/gemma4.py | 59 ++++++++++- tests/gemma4_prefill_prefix_test.py | 152 ++++++++++++++++++++++++++++ 4 files changed, 265 insertions(+), 1 deletion(-) diff --git a/src/mobius/_registry.py b/src/mobius/_registry.py index b49cc0788..eaa6d4be8 100644 --- a/src/mobius/_registry.py +++ b/src/mobius/_registry.py @@ -872,6 +872,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/models/gemma4.py b/src/mobius/models/gemma4.py index 54d7de2c2..c74ed467b 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) @@ -2210,6 +2243,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/tests/gemma4_prefill_prefix_test.py b/tests/gemma4_prefill_prefix_test.py index dbccb285d..c3f91cf67 100644 --- a/tests/gemma4_prefill_prefix_test.py +++ b/tests/gemma4_prefill_prefix_test.py @@ -3,11 +3,15 @@ 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 @@ -133,3 +137,151 @@ def test_splits_per_layer_projection_weight() -> None: 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 From e83a1df8deb0ee426203886640dd47011b39ffad Mon Sep 17 00:00:00 2001 From: justinchuby Date: Thu, 20 Aug 2026 07:09:04 +0000 Subject: [PATCH 110/151] Declare Gemma 4's heterogeneous KV cache geometry in workflow metadata MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Gemma 4 owns two structurally different caches in one decoder. Local (sliding) layers store `head_dim`-wide entries and may drop their prefix; global (full-attention) layers store `global_head_dim`-wide entries and must retain the entire history. The last `num_kv_shared_layers` borrow K/V from an earlier source layer and own no cache at all. The emitter previously declared every cache cell in one undifferentiated `kv_service` group carrying a `storage` allocator flag, and the model-level `sliding_window` was the only eviction signal. A runtime reading that document has no way to tell the two geometries apart, so it would apply sliding-window eviction to the global layers and corrupt them. State groups are now split by attention kind. Each group declares `kind`, `aliasing`, and per-group `reuse.evictable_prefix`, so eviction is a property of the group that actually evicts. On `google/gemma-4-E2B-it` this yields `decoder_cache_sliding_attention` (24 cells at head_dim 256, evictable) and `decoder_cache_full_attention` (6 cells at head_dim 512, not evictable) — exactly the 15 cache-owning layers, with no slots generated for the 20 KV-shared layers. Alongside that, the document is realigned with the runtime's workflow IR: * Runtime-private scheduler identity is no longer serialized. `slot_ids`, `emit.row_ids`, and `emit_row_identity` are removed in favour of `TensorContract.batch_layout`, which states the structural fact a runtime needs — that a value carries one entry per in-flight request on a given axis — without pinning the scheduler's private row table into the package. Policy-component ports derive this from the graph's leading `batch` dimension rather than being hand-annotated. * Deployment policy the runtime owns (`paging`, `allocation`, `compaction`, `storage`) is dropped; the graph ABI constraint it stood for is expressed as `aliasing`. * Cache cells declare `management: runtime` and `release_boundary: invocation`, so cache lifetime is stated rather than inferred. * The package-level `kv_cache.native_dtype` key is removed. Cache storage representation is derived from the exported graph's port dtypes, so declaring it in the package could only contradict the graph. * Native VLM packages emit the workflow IR through the same path as every other package instead of a parallel composite `pipeline.models` document, and image preprocessing outputs bind the processor-local value that produces them via an explicit `source`. Verified end to end against real exports of `google/gemma-4-E2B-it` (snapshot 9dbdf8a839e4e9e0eb56ed80cc8886661d3817cf): the text f32, text f32 + prune-prefill-prefix, text f16/CUDA, and full multimodal (vision + audio + embedding + decoder) packages all validate against the runtime's metadata validator, where previously none did. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: justinchuby (cherry picked from commit b23c7ebaf6d9a6ba837a963851fb24b1cfbf6ffb) --- src/mobius/__main__.py | 30 +- .../integrations/onnx_genai/auto_export.py | 29 +- .../onnx_genai/auto_export_test.py | 8 +- .../codec_workflow_metadata_test.py | 3 + .../onnx_genai/decoder_metadata.py | 30 +- .../onnx_genai/decoder_metadata_test.py | 15 +- .../onnx_genai/inference_metadata.py | 80 +++- .../onnx_genai/inference_metadata_test.py | 8 + .../onnx_genai/workflow_metadata.py | 360 +++++++++++------- .../onnx_genai/workflow_metadata_test.py | 29 +- tests/cli_test.py | 35 +- tests/gemma4_prefill_prefix_test.py | 104 +++++ 12 files changed, 484 insertions(+), 247 deletions(-) diff --git a/src/mobius/__main__.py b/src/mobius/__main__.py index 19159163a..e9251f84c 100644 --- a/src/mobius/__main__.py +++ b/src/mobius/__main__.py @@ -435,35 +435,13 @@ 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: - artifacts = write_onnx_genai_config( - pkg, - output_dir, - config=config, - source=source, - **revision_kwargs, - ) + try: + artifacts = write_onnx_genai_config(pkg, output_dir, config=config, source=source) + 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/integrations/onnx_genai/auto_export.py b/src/mobius/integrations/onnx_genai/auto_export.py index 5d10d9e0b..b7dfd0b90 100644 --- a/src/mobius/integrations/onnx_genai/auto_export.py +++ b/src/mobius/integrations/onnx_genai/auto_export.py @@ -25,6 +25,7 @@ ) from mobius.integrations.onnx_genai.inference_metadata import ( SchedulerConfig, + _copy_runtime_assets, add_adapter_service_to_metadata, add_explicit_package_io, add_policy_components_to_workflow, @@ -591,17 +592,14 @@ def write_onnx_genai_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 - if "audio_encoder" in pkg: - audio_processor_path = _write_hf_audio_processor( - output_dir, - source, - revision=revision, - ) - if audio_processor_path is not None: - artifacts["audio_processor"] = audio_processor_path + # 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)) + if "tokenizer" not in artifacts: + tokenizer_path = _write_hf_tokenizer(output_dir, source) + if tokenizer_path is not None: + artifacts["tokenizer"] = tokenizer_path return artifacts if _looks_like_speech_to_text(pkg): @@ -612,9 +610,12 @@ def write_onnx_genai_config( "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 - ) + if kv_native_dtype is not None: + raise ValueError( + "speech-to-text export derives KV state dtype from ONNX ports; " + "kv_native_dtype overrides are unsupported" + ) + decoder_metadata = decoder_metadata_from_config(resolved_config) path = write_speech_to_text_pipeline_metadata( output_dir, decoder_metadata=decoder_metadata, diff --git a/src/mobius/integrations/onnx_genai/auto_export_test.py b/src/mobius/integrations/onnx_genai/auto_export_test.py index 1ea0c05f2..64de3ecf0 100644 --- a/src/mobius/integrations/onnx_genai/auto_export_test.py +++ b/src/mobius/integrations/onnx_genai/auto_export_test.py @@ -284,6 +284,9 @@ def test_seeded_decoder_sampler_uses_request_controls_and_direct_kv_carry(): ] 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"} @@ -623,12 +626,13 @@ def test_dispatch_speech_to_text_pipeline(tmp_path): "decoder": _FakeModel(["decoder_input_ids", "encoder_hidden_states"], ["logits"]), } ) - artifacts = write_onnx_genai_config(pkg, str(tmp_path), kv_native_dtype="bf16") + artifacts = write_onnx_genai_config(pkg, str(tmp_path)) with open(artifacts["inference_metadata"]) as handle: metadata = yaml.safe_load(handle) - assert metadata["kv_cache"] == {"native_dtype": "bfloat16"} + # Cache storage representation is derived from the graph, never declared. + assert "kv_cache" not in metadata pipeline = metadata["pipeline"] assert pipeline["models"]["encoder"]["filename"] == "encoder/model.onnx" assert pipeline["models"]["encoder"]["type"] == "encoder" diff --git a/src/mobius/integrations/onnx_genai/codec_workflow_metadata_test.py b/src/mobius/integrations/onnx_genai/codec_workflow_metadata_test.py index abc41ad0a..feb8c287f 100644 --- a/src/mobius/integrations/onnx_genai/codec_workflow_metadata_test.py +++ b/src/mobius/integrations/onnx_genai/codec_workflow_metadata_test.py @@ -178,6 +178,9 @@ def test_real_qwen3_tts_workflow_carries_trained_transitions_and_kv_state(): 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" + ) assert workflow["inputs"]["package.slot_ids"] == { "contract": { "dtype": "int64", 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/inference_metadata.py b/src/mobius/integrations/onnx_genai/inference_metadata.py index 4aceb8fd4..74b99f98b 100644 --- a/src/mobius/integrations/onnx_genai/inference_metadata.py +++ b/src/mobius/integrations/onnx_genai/inference_metadata.py @@ -113,6 +113,10 @@ def _port(value: Any) -> _Port: ) +_BATCH_DIMENSION = "batch" +"""Symbolic leading dimension mobius uses for per-request batching.""" + + def _shape_metadata(port: _Port) -> list[int | str | None]: """Return a YAML-safe graph shape without losing symbolic dimensions.""" shape: list[int | str | None] = [] @@ -125,6 +129,62 @@ def _shape_metadata(port: _Port) -> list[int | str | None]: 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 { @@ -1451,11 +1511,19 @@ def tensor_contract(value: Any) -> dict[str, Any]: "fp16": "float16", "bf16": "bfloat16", }.get(port.dtype, port.dtype) - return { + shape = _shape_metadata(port) + contract: dict[str, Any] = { "dtype": dtype, "rank": port.rank, - "shape": _shape_metadata(port), + "shape": shape, } + # Policy components are per-row operators: when the graph's leading + # dimension is the batch symbol, axis 0 carries exactly one entry per + # in-flight request. Declaring it lets the runtime permute/compact the + # batch without the producer serializing any row identity. + if shape and shape[0] == _BATCH_DIMENSION: + contract["batch_layout"] = {"kind": "request_aligned", "axis": 0} + return contract for name, component in policy_components.items(): declaration = { @@ -1932,6 +2000,9 @@ def build_native_vlm_package_metadata( "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, @@ -2005,8 +2076,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.""" @@ -2018,8 +2087,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) diff --git a/src/mobius/integrations/onnx_genai/inference_metadata_test.py b/src/mobius/integrations/onnx_genai/inference_metadata_test.py index f7ae2e3b8..c39bbfb4e 100644 --- a/src/mobius/integrations/onnx_genai/inference_metadata_test.py +++ b/src/mobius/integrations/onnx_genai/inference_metadata_test.py @@ -906,6 +906,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")[ @@ -932,17 +934,21 @@ 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) @@ -1478,6 +1484,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" diff --git a/src/mobius/integrations/onnx_genai/workflow_metadata.py b/src/mobius/integrations/onnx_genai/workflow_metadata.py index 539acfb29..0def2c65f 100644 --- a/src/mobius/integrations/onnx_genai/workflow_metadata.py +++ b/src/mobius/integrations/onnx_genai/workflow_metadata.py @@ -6,6 +6,7 @@ from __future__ import annotations import os +import re from typing import Any import onnx_ir as ir @@ -44,6 +45,7 @@ build_tts_state_initializer, ) from mobius.integrations.onnx_genai.inference_metadata import ( + _name_image_preprocessing_program, _port, _shape_metadata, add_adapter_service_to_metadata, @@ -105,6 +107,16 @@ def _contract(value: ir.Value) -> dict[str, Any]: } +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}} + + def _component( model: ir.Model, artifact: str, @@ -319,59 +331,6 @@ def convert(node: dict[str, Any]) -> dict[str, Any]: return workflow - - -def _name_image_preprocessing_program(image: dict[str, Any]) -> None: - """Convert structural preprocessing transforms into explicit typed SSA values.""" - transforms = image["transforms"] - 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 _invoke( component: str, inputs: dict[str, str], @@ -537,6 +496,17 @@ def _kv_storage_contract(model: ir.Model) -> dict[str, Any]: } +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 _state_group( *, ports: dict[str, Any], @@ -558,11 +528,7 @@ def _state_group( "kind": kind, "sequence_axis": sequence_axis, "layout": layout, - "aliasing": ( - aliasing - if aliasing is not None - else ("permitted" if storage in ("shared_buffer", "paged") else "forbidden") - ), + "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": ports, @@ -572,6 +538,102 @@ def _state_group( 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 ``past_key_values.N.key`` port name.""" + match = re.search(r"\.(\d+)\.(?:key|value)$", port_name) + return int(match.group(1)) if match else fallback + + +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, +) -> tuple[dict[str, Any], dict[str, str]]: + """Build ``serving.state_service.groups`` plus each cell's owning group. + + One group per semantic kind: 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). 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. + """ + kinds = _state_group_kinds(config, cache_pairs) + distinct = sorted(set(kinds)) + names = { + kind: (base_name if len(distinct) == 1 else f"{base_name}_{kind}") for kind in distinct + } + cell_group = {} + grouped_ports: dict[str, dict[str, dict[str, dict[str, str]]]] = { + name: {} for name in names.values() + } + for index, kind in enumerate(kinds): + cell = f"cache_{index}" + cell_group[cell] = names[kind] + for component, aliases in ports.items(): + for cell, alias in aliases.items(): + grouped_ports[cell_group[cell]].setdefault(component, {})[cell] = alias + groups = { + names[kind]: { + "kind": kind, + "sequence_axis": sequence_axis, + "layout": "bnsh", + **({"logical_lengths": logical_lengths} if logical_lengths else {}), + "aliasing": aliasing, + "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", + }, + "ports": grouped_ports[names[kind]], + } + for kind in distinct + } + 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"] @@ -712,8 +774,8 @@ def _build_real_tts_workflow_metadata(pkg: Any, config: Any) -> dict[str, Any]: pkg.add_policy_component("codec_layout", build_codec_layout_transpose(num_groups)) batch = _contract(prompt)["shape"][0] - batch_int = {"dtype": "int64", "rank": 1, "shape": [batch]} - batch_bool = {"dtype": "bool", "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 = { "request.prompt_tokens": { @@ -1473,7 +1535,7 @@ def frame_generation_nodes(prefix: str, hidden: str, logits: str) -> list[dict[s } for index, (past, present) in enumerate(talker_caches): state[f"talker_cache_{index}"] = { - "contract": _contract(past), + "contract": _request_aligned(_contract(past)), "class": "semantic", "scope": "invocation", "initializer": f"talker.setup.{present.name}", @@ -1483,10 +1545,12 @@ def frame_generation_nodes(prefix: str, hidden: str, logits: str) -> list[dict[s "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": _contract(past), + "contract": _request_aligned(_contract(past)), "class": "semantic", "scope": "invocation", "initializer": f"frame.predictor.{present.name}", @@ -1496,6 +1560,8 @@ def frame_generation_nodes(prefix: str, hidden: str, logits: str) -> list[dict[s "max": "package.predictor_context_limit", }, "service_group": "predictor_cache", + "management": "runtime", + "release_boundary": "invocation", } outer_carried = [ @@ -1843,8 +1909,8 @@ def build_tts_workflow_metadata(pkg: Any, config: Any) -> dict[str, Any]: pkg.add_policy_component("codec_layout", build_codec_layout_transpose(num_groups)) batch = _contract(prompt_input)["shape"][0] - batch_int = {"dtype": "int64", "rank": 1, "shape": [batch]} - batch_bool = {"dtype": "bool", "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.prompt_tokens": { @@ -2333,8 +2399,8 @@ def build_diffusion_workflow_metadata( pkg.add_policy_component("schedule_lookup", build_schedule_lookup(timestep_input.dtype)) batch = _contract(sample_input)["shape"][0] - batch_int = {"dtype": "int64", "rank": 1, "shape": [batch]} - batch_bool = {"dtype": "bool", "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.latent": { @@ -2740,8 +2806,8 @@ def build_vlm_workflow_metadata( pkg.add_policy_component("cache_length_update", build_selective_integer_add()) batch = _contract(token_input)["shape"][0] - batch_int = {"dtype": "int64", "rank": 1, "shape": [batch]} - batch_bool = {"dtype": "bool", "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]} eos = _source_token_id(source, "eos_token_id", getattr(config, "eos_token_id", 0)) inputs: dict[str, Any] = { @@ -3166,7 +3232,7 @@ def build_vlm_workflow_metadata( for index, (past, present) in enumerate(cache_pairs): cell = f"cache_{index}" state[cell] = { - "contract": _contract(past), + "contract": _request_aligned(_contract(past)), "scope": "invocation", "initializer": f"decoder.setup.{present.name}", "recurrence": { @@ -3181,7 +3247,8 @@ def build_vlm_workflow_metadata( ), "max": "package.max_context", }, - "service_group": "decoder_cache", + "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}" @@ -3194,6 +3261,33 @@ def build_vlm_workflow_metadata( 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", + ) + for cell, group_name in vlm_cell_groups.items(): + state[cell]["service_group"] = group_name carried = [] initial_effects = { "sample": "sample.0", @@ -3510,11 +3604,13 @@ def build_vlm_workflow_metadata( "inputs": inputs, "outputs": { "tokens": { - "contract": { - "dtype": "int64", - "rank": 2, - "shape": [batch, "generated_sequence"], - }, + "contract": _request_aligned( + { + "dtype": "int64", + "rank": 2, + "shape": [batch, "generated_sequence"], + } + ), "role": "tokens", "stage": "pre_adapter", } @@ -3527,33 +3623,7 @@ def build_vlm_workflow_metadata( "active": "active", "done": "done", "accepted_len": "accepted_len", - "state_service": { - "groups": { - "decoder_cache": _state_group( - sequence_axis=next( - ( - axis - for axis, dimension in enumerate( - _contract(cache_pairs[0][0])["shape"] - ) - if "sequence" in str(dimension) - ), - 2, - ), - logical_lengths="cache_lengths", - storage=decoder_kv["storage"], - ports={ - "decoder": { - f"cache_{index}": { - "input": past.name, - "output": present.name, - } - for index, (past, present) in enumerate(cache_pairs) - } - }, - ) - }, - }, + "state_service": {"groups": vlm_state_groups}, } } if cache_pairs @@ -3679,8 +3749,8 @@ def build_speculative_workflow_metadata( 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 = {"dtype": "int64", "rank": 1, "shape": [batch]} - batch_bool = {"dtype": "bool", "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.tokens": { @@ -4250,7 +4320,7 @@ def build_speculative_workflow_metadata( 2, ) state[cell] = { - "contract": _contract(past), + "contract": _request_aligned(_contract(past)), "class": "semantic", "scope": "invocation", "initializer": initializer, @@ -4260,6 +4330,8 @@ def build_speculative_workflow_metadata( "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( @@ -4316,13 +4388,15 @@ def build_speculative_workflow_metadata( "inputs": inputs, "outputs": { "tokens": { - "contract": { - **_contract(proposed_tokens), - "shape": [ - *_contract(proposed_tokens)["shape"][:-1], - "accepted_sequence", - ], - }, + "contract": _request_aligned( + { + **_contract(proposed_tokens), + "shape": [ + *_contract(proposed_tokens)["shape"][:-1], + "accepted_sequence", + ], + } + ), "role": "tokens", "stage": "pre_adapter", }, @@ -4723,8 +4797,8 @@ def _build_autoregressive_workflow_metadata( for decoder_input, (encoder_output, _) in sorted(cross_bindings.items()): encoder_invoke_outputs[encoder_output] = f"encoder.{encoder_output}" batch_dimension = _shape_metadata(_port(token_input))[0] - batch_int = {"dtype": "int64", "rank": 1, "shape": [batch_dimension]} - batch_bool = {"dtype": "bool", "rank": 1, "shape": [batch_dimension]} + 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): @@ -5211,6 +5285,7 @@ def _build_autoregressive_workflow_metadata( } ) 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 @@ -5229,7 +5304,7 @@ def _build_autoregressive_workflow_metadata( 2, ) state[cell] = { - "contract": _contract(past), + "contract": _request_aligned(_contract(past)), "scope": "invocation", "initializer": setup_value, "recurrence": { @@ -5237,9 +5312,13 @@ def _build_autoregressive_workflow_metadata( "axis": decoder_kv_axis, "max": "package.max_context", }, - "service_group": "decoder_cache", + # 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( @@ -5284,6 +5363,26 @@ def _build_autoregressive_workflow_metadata( "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", + ) + for cell in decoder_cache_cells: + state[cell]["service_group"] = decoder_cell_groups[cell] setup = { "kind": "sequence", @@ -5579,11 +5678,13 @@ def _build_autoregressive_workflow_metadata( "inputs": workflow_inputs, "outputs": { "tokens": { - "contract": { - "dtype": "int64", - "rank": 2, - "shape": [batch_dimension, "generated_sequence"], - }, + "contract": _request_aligned( + { + "dtype": "int64", + "rank": 2, + "shape": [batch_dimension, "generated_sequence"], + } + ), "role": "tokens", "stage": "pre_adapter", } @@ -5608,20 +5709,7 @@ def _build_autoregressive_workflow_metadata( "active": "active", "done": "done", "accepted_len": "accepted_len", - "state_service": { - "groups": { - "decoder_cache": _state_group( - sequence_axis=decoder_kv_axis, - logical_lengths="cache_lengths" if fixed_capacity else None, - storage=( - decoder_kv_contract["storage"] - if fixed_capacity - else "growable" - ), - ports={decoder_name: decoder_kv_ports}, - ) - }, - }, + "state_service": {"groups": decoder_state_groups}, } } if cache_pairs @@ -5699,7 +5787,7 @@ def build_language_diffusion_pipeline_metadata( "shape": token_contract["shape"], } batch_dimension = token_contract["shape"][0] - batch_int = {"dtype": "int64", "rank": 1, "shape": [batch_dimension]} + batch_int = _request_aligned({"dtype": "int64", "rank": 1, "shape": [batch_dimension]}) control_int = {"dtype": "int64", "rank": 1, "shape": [1]} inputs = { "request.input_ids": { diff --git a/src/mobius/integrations/onnx_genai/workflow_metadata_test.py b/src/mobius/integrations/onnx_genai/workflow_metadata_test.py index 3e6d94471..b5aae2fe3 100644 --- a/src/mobius/integrations/onnx_genai/workflow_metadata_test.py +++ b/src/mobius/integrations/onnx_genai/workflow_metadata_test.py @@ -48,7 +48,10 @@ def test_speculative_emit_uses_accepted_prefix_length(): ] 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 workflow["inputs"]["request.slot_ids"]["source"]["name"] == "serving.slot_ids" assert workflow["outputs"]["tokens"]["contract"]["shape"][-1] == "accepted_sequence" assert workflow["state"]["cache_0"]["recurrence"] == { @@ -292,6 +295,8 @@ def collect_emits(node): ) 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( @@ -309,6 +314,17 @@ def collect_emits(node): 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 { "slot_ids", @@ -327,8 +343,6 @@ def collect_emits(node): # 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) - assert decoder_cache["aliasing"] == "permitted" - assert decoder_cache["kind"] == "full_attention" kv_ports = decoder_cache["ports"]["decoder"] assert len(kv_ports) == 104 assert kv_ports["cache_103"] == { @@ -530,12 +544,17 @@ def test_speculative_workflow_uses_per_row_ragged_state_and_rng(): 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"]["ports"]["verifier"][ - "cache_0" - ] == { + 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", } diff --git a/tests/cli_test.py b/tests/cli_test.py index 17ea88b6f..d16abc818 100644 --- a/tests/cli_test.py +++ b/tests/cli_test.py @@ -612,10 +612,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 +631,18 @@ 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", ) - generic_writer.assert_not_called() def test_runtime_onnx_genai_does_not_fallback_for_unsupported_vlm(self): pkg = mock.MagicMock() @@ -670,25 +662,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/gemma4_prefill_prefix_test.py b/tests/gemma4_prefill_prefix_test.py index c3f91cf67..3560184bc 100644 --- a/tests/gemma4_prefill_prefix_test.py +++ b/tests/gemma4_prefill_prefix_test.py @@ -285,3 +285,107 @@ def test_pruned_prefill_matches_unpruned_final_row( # 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 + assert group["aliasing"] == "permitted" + 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"] From 9e1e5b35ee66f43d10e9a16558fd63178386ed8c Mon Sep 17 00:00:00 2001 From: Justin Chu Date: Thu, 20 Aug 2026 09:17:33 +0000 Subject: [PATCH 111/151] Derive KV storage, ship chat templates, and drop vestigial row identity Four independent producer defects surfaced while driving real Gemma 4 E2B packages through the onnx-genai workflow runtime end to end. Derive KV storage from the attention operator, not from port names. `_kv_storage_contract` declared `shared_buffer` for every non-paged model that had cache ports. Binding past and present to one capacity-sized buffer is only sound when the attention operator takes the logical cache length as a separate input: `GroupQueryAttention` does (`seqlens_k` / `total_sequence_length`), and a paged layout carries lengths in its block tables. The opset-24 `Attention` operator instead derives the total length from the past tensor's own shape, so a preallocated buffer both attends over unwritten slots and contradicts an exactly sized attention mask. ORT rejects that outright: the f32 Gemma 4 package failed at prefill with "inconsistent total_sequence_length (between attn_mask and past_key and past_value)". Storage is now derived structurally from the operators that consume the cache. This flips every default-EP decoder package to the dynamic branch, which previously existed but was never exercised, and exposed a latent bug there: the growing mask recurrence used `package.one` / `package.one_token`, which are `[batch]`, where the runtime requires a static `[1]` control. Both growing recurrences now use a new `package.one_step` literal. Ship chat-template assets with text packages. Text and decoder packages copied only `tokenizer.json`, so the runtime had no chat template and fed the model raw user text with no BOS and no turn markers. Gemma 4 answers such a prompt with unbounded repetition. `_write_text_runtime_assets` now copies the tokenizer, template, and special-token maps, while excluding the image/audio processor configs a text package cannot consume. Materialize narrow-float caches through a cast. ORT's `ConstantOfShape` kernel has no fp8 output implementation, so an fp8 KV package failed session initialization with "Unsupported value attribute datatype: 17". The decoder state initializer now fills a supported dtype and casts. Resolve the static-cache task after the text-only substitution. `--features text-only,static-cache` resolved the task from the checkpoint's raw `model_type`, pairing a text-only module with the multimodal task, which then failed looking for a decoder sub-module that a text-only module does not have. Static-cache graphs additionally expose in-place ring buffers and write indices instead of past/present pairs and a rank-2 mask, so they cannot be described by the onnx-genai workflow contract at all; the CLI now rejects that combination up front instead of exporting the weights first. Also drop the vestigial `slot_ids` surface. Row identity is runtime-private and the runtime's contracts no longer carry it, so the workflow inputs, state cell, emit keys, and the adapter selection fields that published it are removed in favour of the declared `batch_layout`. The artifact-backed VLM writer test now wires its synthetic decoder through `GroupQueryAttention`, because a decoder built from `Identity` nodes describes a dynamic cache and can no longer stand in for a shared-buffer one. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: Justin Chu (cherry picked from commit dcb2913f4e8ae98afebf9c16f9823ecfe671ea20) --- src/mobius/__main__.py | 31 +- src/mobius/adapters.py | 2 - src/mobius/adapters_test.py | 15 +- src/mobius/generation/_policy_components.py | 33 +- .../generation/_policy_components_test.py | 52 ++ .../integrations/onnx_genai/auto_export.py | 50 +- .../onnx_genai/auto_export_test.py | 37 +- .../codec_workflow_metadata_test.py | 13 +- .../onnx_genai/inference_metadata.py | 45 +- .../onnx_genai/workflow_metadata.py | 160 ++-- .../onnx_genai/workflow_metadata_test.py | 114 ++- tests/cli_test.py | 61 +- tests/gemma4_prefill_prefix_test.py | 785 +++++++++--------- 13 files changed, 830 insertions(+), 568 deletions(-) diff --git a/src/mobius/__main__.py b/src/mobius/__main__.py index e9251f84c..82e17f940 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 @@ -201,6 +213,23 @@ def _resolve_static_cache_task(model_type: str) -> ModelTask: "Remove --task to use --features static-cache." ) + # Validate static-cache + onnx-genai compatibility. + # + # A static-cache decoder exposes in-place ring buffers (``key_cache.N`` / + # ``updated_key_cache.N``) plus ``write_indices`` and ``nonpad_kv_seqlen``, + # and drops the rank-2 attention mask. The onnx-genai workflow decoder + # contract is built on ``past_key_values.N`` -> ``present.N`` pairs and that + # mask, so the metadata emitter cannot describe a static-cache graph. + # Reject the combination up front rather than after exporting the weights. + if args.static_cache and args.runtime == "onnx-genai": + raise SystemExit( + "Error: --features static-cache cannot be combined with " + "--runtime onnx-genai. The onnx-genai workflow decoder contract " + "requires dynamic past/present KV ports and a rank-2 attention " + "mask, which a static-cache graph does not expose. Build without " + "--features static-cache, or omit --runtime onnx-genai." + ) + # text-only resolution lives in build() (model_type remap + config # stripping), which is only reached on the HuggingFace model-ID path. if args.text_only and args.config: diff --git a/src/mobius/adapters.py b/src/mobius/adapters.py index 3fa8aa7ff..f3ca6b449 100644 --- a/src/mobius/adapters.py +++ b/src/mobius/adapters.py @@ -613,8 +613,6 @@ def __post_init__(self) -> None: class AdapterServiceOptions: """Producer-neutral runtime lifecycle, planning, and artifact format options.""" - slot_ids: str | None = None - request_epochs: str | None = None segments: str = "request.adapter_segments" adapter_counts: str = "request.adapter_counts" scales: str = "request.adapter_scales" diff --git a/src/mobius/adapters_test.py b/src/mobius/adapters_test.py index 7f1b18648..edba87d6b 100644 --- a/src/mobius/adapters_test.py +++ b/src/mobius/adapters_test.py @@ -621,7 +621,6 @@ def test_exact_onnx_genai_catalog_and_portable_bundle_serialization() -> None: {"decoder": model}, adapter_target_manifest=manifest, adapter_service_options=AdapterServiceOptions( - slot_ids="request.slot_ids", active="request.active", max_adapters=2, cache_max_entries=2, @@ -679,8 +678,6 @@ def test_exact_onnx_genai_catalog_and_portable_bundle_serialization() -> None: "onnx-genai-targeted-base-v1:sha256:" ) assert service["selection"] == { - "slot_ids": "request.slot_ids", - "request_epochs": "request.request_epochs", "segments": "request.adapter_segments", "adapter_counts": "request.adapter_counts", "scales": "request.adapter_scales", @@ -719,15 +716,9 @@ def test_exact_onnx_genai_catalog_and_portable_bundle_serialization() -> None: "stable_buffers": True, "invalidate_capture_on_eviction": True, } - assert metadata["pipeline"]["workflow"]["inputs"]["request.request_epochs"] == { - "contract": {"dtype": "int64", "rank": 1, "shape": ["batch"]}, - "role": { - "kind": "runtime", - "version": "1.0", - "role": "request_epochs", - }, - "source": {"kind": "request"}, - } + 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" diff --git a/src/mobius/generation/_policy_components.py b/src/mobius/generation/_policy_components.py index 04fd84bfa..3fe03af97 100644 --- a/src/mobius/generation/_policy_components.py +++ b/src/mobius/generation/_policy_components.py @@ -112,6 +112,28 @@ def attach_policy_components( 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, @@ -803,11 +825,18 @@ def build_decoder_state_initializer( f"dimension {dimension_text!r}" ) cache_shape = op.Concat(*shape_parts, axis=0) - zero = 0.0 if value.dtype.is_floating_point else 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=value.dtype), + value=ir.tensor([zero], dtype=fill_dtype), ) + if fill_dtype != value.dtype: + empty = op.Cast(empty, to=value.dtype) empty.shape = ( ir.Shape( [ diff --git a/src/mobius/generation/_policy_components_test.py b/src/mobius/generation/_policy_components_test.py index 363a21e57..62fd36125 100644 --- a/src/mobius/generation/_policy_components_test.py +++ b/src/mobius/generation/_policy_components_test.py @@ -954,3 +954,55 @@ def test_capability_driven_attachment_is_model_agnostic(): "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) diff --git a/src/mobius/integrations/onnx_genai/auto_export.py b/src/mobius/integrations/onnx_genai/auto_export.py index b7dfd0b90..d077b6c48 100644 --- a/src/mobius/integrations/onnx_genai/auto_export.py +++ b/src/mobius/integrations/onnx_genai/auto_export.py @@ -24,6 +24,7 @@ decoder_metadata_from_config, ) from mobius.integrations.onnx_genai.inference_metadata import ( + _TEXT_RUNTIME_ASSET_NAMES, SchedulerConfig, _copy_runtime_assets, add_adapter_service_to_metadata, @@ -178,12 +179,31 @@ def _write_clip_tokenizer( return path -def _write_hf_tokenizer( - output_dir: str, - source: str | None, - *, - revision: str | None = None, -) -> str | None: +def _write_text_runtime_assets(output_dir: str, source: str | 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) + if "tokenizer" not in artifacts: + fallback = _write_hf_tokenizer(output_dir, source) + if fallback is not None: + artifacts["tokenizer"] = fallback + return artifacts + + +def _write_hf_tokenizer(output_dir: str, source: str | None) -> str | None: """Emit ``tokenizer.json`` for a text-producing package from its HF source. Decoder-LM, multimodal (VLM / speech-language ASR), and Whisper-style ASR @@ -500,9 +520,7 @@ def write_onnx_genai_config( num_inference_steps=num_inference_steps, ) artifacts = {"inference_metadata": path} - tokenizer_path = _write_hf_tokenizer(output_dir, source) - if tokenizer_path is not None: - artifacts["tokenizer"] = tokenizer_path + artifacts.update(_write_text_runtime_assets(output_dir, source)) return artifacts if _looks_like_diffusion(pkg): @@ -624,14 +642,8 @@ def write_onnx_genai_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 - audio_processor_path = _write_hf_audio_processor( - output_dir, - source, - revision=revision, - ) + artifacts.update(_write_text_runtime_assets(output_dir, source)) + audio_processor_path = _write_hf_audio_processor(output_dir, source) if audio_processor_path is not None: artifacts["audio_processor"] = audio_processor_path return artifacts @@ -684,7 +696,5 @@ def write_onnx_genai_config( sampler=str(getattr(resolved_config, "workflow_sampler", "greedy")), ) 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)) 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 64de3ecf0..bf9feb454 100644 --- a/src/mobius/integrations/onnx_genai/auto_export_test.py +++ b/src/mobius/integrations/onnx_genai/auto_export_test.py @@ -178,7 +178,6 @@ def test_dispatch_decoder(tmp_path): "request.eos_ids", "request.eos_lengths", "request.row_max_iterations", - "package.slot_ids", "request.rng_counter", } assert workflow["inputs"]["request.prompt_lengths"]["default"] == -1 @@ -830,3 +829,39 @@ def test_decoder_without_source_skips_tokenizer(tmp_path): 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 index feb8c287f..16a10cdb4 100644 --- a/src/mobius/integrations/onnx_genai/codec_workflow_metadata_test.py +++ b/src/mobius/integrations/onnx_genai/codec_workflow_metadata_test.py @@ -181,17 +181,8 @@ def test_real_qwen3_tts_workflow_carries_trained_transitions_and_kv_state(): assert workflow["serving"]["state_service"]["groups"]["talker_cache"]["kind"] == ( "full_attention" ) - assert workflow["inputs"]["package.slot_ids"] == { - "contract": { - "dtype": "int64", - "rank": 1, - "shape": ["batch"], - "batch_layout": {"kind": "request_aligned", "axis": 0}, - }, - "role": {"kind": "opaque"}, - "source": {"kind": "application", "name": "serving.slot_ids"}, - "required": True, - } + # 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.") diff --git a/src/mobius/integrations/onnx_genai/inference_metadata.py b/src/mobius/integrations/onnx_genai/inference_metadata.py index 74b99f98b..8c576ab2a 100644 --- a/src/mobius/integrations/onnx_genai/inference_metadata.py +++ b/src/mobius/integrations/onnx_genai/inference_metadata.py @@ -35,7 +35,7 @@ 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 @@ -61,6 +61,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: @@ -1652,30 +1665,13 @@ def ensure_input( } if not compatible_input(name, dtype=dtype, shape=shape, role=role): raise ValueError( - f"adapter {role or 'slot_ids'} must reference a required " + f"adapter {role} must reference a required " "request/application-sourced " f"{dtype}{shape} workflow input" ) - serving = workflow.get("serving") if workflow is not None else None - serving_slot_ids = serving.get("slot_ids") if isinstance(serving, dict) else None - slot_ids = options.slot_ids or serving_slot_ids or "request.slot_ids" - request_epochs = options.request_epochs or "request.request_epochs" active = options.active if workflow is not None: - ensure_input( - slot_ids, - dtype="int64", - shape=["batch"], - role=None, - source={"kind": "application", "name": "serving.slot_ids"}, - ) - ensure_input( - request_epochs, - dtype="int64", - shape=["batch"], - role="request_epochs", - ) ensure_input( options.segments, dtype="int64", @@ -1710,8 +1706,6 @@ def ensure_input( "target_manifest": pkg.adapter_target_manifest_metadata(), "discovery_fallback": options.discovery_fallback, "selection": { - "slot_ids": slot_ids, - "request_epochs": request_epochs, "segments": options.segments, "adapter_counts": options.adapter_counts, "scales": options.scales, @@ -2018,14 +2012,13 @@ def build_native_vlm_package_metadata( def _copy_runtime_assets( output_dir: str, source: str | None, - *, - revision: str | None = None, + names: Sequence[str] = _RUNTIME_ASSET_NAMES, ) -> dict[str, str]: if not source: return {} os.makedirs(output_dir, exist_ok=True) - for filename in _RUNTIME_ASSET_NAMES: - source_path = _source_asset_path(source, filename, revision=revision) + for filename in names: + source_path = _source_asset_path(source, filename) if source_path is not None: shutil.copy2(source_path, os.path.join(output_dir, filename)) @@ -2065,7 +2058,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)) } diff --git a/src/mobius/integrations/onnx_genai/workflow_metadata.py b/src/mobius/integrations/onnx_genai/workflow_metadata.py index 0def2c65f..c73b74256 100644 --- a/src/mobius/integrations/onnx_genai/workflow_metadata.py +++ b/src/mobius/integrations/onnx_genai/workflow_metadata.py @@ -478,7 +478,19 @@ 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. It does not require an attention-node attribute. + 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( @@ -487,13 +499,51 @@ def _kv_storage_contract(model: ir.Model) -> dict[str, Any]: 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": "paged" if paged else "shared_buffer", + "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 + } + 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: @@ -867,12 +917,6 @@ def _build_real_tts_workflow_metadata(pkg: Any, config: Any) -> dict[str, Any]: "required": False, "default": True, }, - "package.slot_ids": { - "contract": batch_int, - "role": {"kind": "opaque"}, - "source": {"kind": "application", "name": "serving.slot_ids"}, - "required": True, - }, } for iteration in range(num_groups - 2): inputs[f"package.setup_predictor_iteration_{iteration}"] = { @@ -1511,13 +1555,6 @@ def frame_generation_nodes(prefix: str, hidden: str, logits: str) -> list[dict[s "initializer": "package.zero_batch", "recurrence": {"kind": "invariant"}, }, - "slot_ids": { - "contract": batch_int, - "class": "semantic", - "scope": "invocation", - "initializer": "package.slot_ids", - "recurrence": {"kind": "invariant"}, - }, "talker_cache_lengths": { "contract": batch_int, "class": "semantic", @@ -1590,13 +1627,6 @@ def frame_generation_nodes(prefix: str, hidden: str, logits: str) -> list[dict[s ), "next": "state.accepted_len.final", }, - { - "cell": "slot_ids", - "current": "package.slot_ids", - "body_input": "state.slot_ids.body", - "body_output": "state.slot_ids.body", - "next": "state.slot_ids.final", - }, { "cell": "talker_cache_lengths", "current": "package.zero_batch", @@ -2891,6 +2921,16 @@ def build_vlm_workflow_metadata( "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"}, @@ -2912,12 +2952,6 @@ def build_vlm_workflow_metadata( "required": False, "default": 0, }, - "package.slot_ids": { - "contract": batch_int, - "role": {"kind": "opaque"}, - "source": {"kind": "application", "name": "serving.slot_ids"}, - "required": True, - }, } inputs.update( { @@ -3089,7 +3123,7 @@ def build_vlm_workflow_metadata( else { "kind": "growing", "axis": 1, - "increment": "package.one", + "increment": "package.one_step", "max": "package.max_context", } ), @@ -3115,13 +3149,6 @@ def build_vlm_workflow_metadata( "initializer": "package.zero_batch", "recurrence": {"kind": "invariant"}, }, - "slot_ids": { - "contract": batch_int, - "class": "semantic", - "scope": "invocation", - "initializer": "package.slot_ids", - "recurrence": {"kind": "invariant"}, - }, "cache_lengths": { "contract": batch_int, "class": "semantic", @@ -3194,13 +3221,6 @@ def build_vlm_workflow_metadata( "accepted_len.next", "state.accepted_len.final", ), - ( - "slot_ids", - "package.slot_ids", - "state.slot_ids.body", - "state.slot_ids.body", - "state.slot_ids.final", - ), ( "cache_lengths", "initializer.cache_lengths" if fixed_capacity else "package.zero_batch", @@ -3811,12 +3831,6 @@ def build_speculative_workflow_metadata( "required": False, "default": True, }, - "request.slot_ids": { - "contract": batch_int, - "role": {"kind": "opaque"}, - "source": {"kind": "application", "name": "serving.slot_ids"}, - "required": True, - }, "request.cache_lengths": { "contract": batch_int, "role": {"kind": "opaque"}, @@ -4190,13 +4204,6 @@ def build_speculative_workflow_metadata( "initializer": "package.zero", "recurrence": {"kind": "invariant"}, }, - "slot_ids": { - "contract": batch_int, - "class": "semantic", - "scope": "invocation", - "initializer": "request.slot_ids", - "recurrence": {"kind": "invariant"}, - }, "cache_lengths": { "contract": batch_int, "class": "semantic", @@ -4241,13 +4248,6 @@ def build_speculative_workflow_metadata( emit_length, "state.accepted_len.final", ), - ( - "slot_ids", - "request.slot_ids", - "state.slot_ids.body", - "state.slot_ids.body", - "state.slot_ids.final", - ), ( "cache_lengths", "request.cache_lengths", @@ -4830,6 +4830,16 @@ def _build_autoregressive_workflow_metadata( "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"}, @@ -4999,12 +5009,6 @@ def _build_autoregressive_workflow_metadata( "required": False, "default": False, }, - "package.slot_ids": { - "contract": batch_int, - "role": {"kind": "opaque"}, - "source": {"kind": "application", "name": "serving.slot_ids"}, - "required": True, - }, "package.cache_lengths": { "contract": batch_int, "role": {"kind": "opaque"}, @@ -5105,13 +5109,6 @@ def _build_autoregressive_workflow_metadata( "initializer": "package.zero_batch", "recurrence": {"kind": "invariant"}, }, - "slot_ids": { - "contract": batch_int, - "class": "semantic", - "scope": "invocation", - "initializer": "package.slot_ids", - "recurrence": {"kind": "invariant"}, - }, "cache_lengths": { "contract": batch_int, "class": "semantic", @@ -5199,13 +5196,6 @@ def _build_autoregressive_workflow_metadata( "body_output": "accepted_len.next", "next": "state.accepted_len.final", }, - { - "cell": "slot_ids", - "current": "package.slot_ids", - "body_input": "state.slot_ids.body", - "body_output": "state.slot_ids.body", - "next": "state.slot_ids.final", - }, ] ) if sampler_with_rng: @@ -5248,7 +5238,7 @@ def _build_autoregressive_workflow_metadata( else { "kind": "growing", "axis": 1, - "increment": "package.one_token", + "increment": "package.one_step", "max": "package.max_context", } ), @@ -5637,7 +5627,7 @@ def _build_autoregressive_workflow_metadata( { "when": "state.active.body", "valid_length": "token.emitted_length", - } + } if cache_pairs else {} ), diff --git a/src/mobius/integrations/onnx_genai/workflow_metadata_test.py b/src/mobius/integrations/onnx_genai/workflow_metadata_test.py index b5aae2fe3..407ee4634 100644 --- a/src/mobius/integrations/onnx_genai/workflow_metadata_test.py +++ b/src/mobius/integrations/onnx_genai/workflow_metadata_test.py @@ -17,6 +17,7 @@ _VlmConfig, ) from mobius.integrations.onnx_genai.workflow_metadata import ( + _kv_storage_contract, build_language_diffusion_pipeline_metadata, build_speculative_workflow_metadata, build_vlm_workflow_metadata, @@ -52,7 +53,7 @@ def test_speculative_emit_uses_accepted_prefix_length(): assert "row_ids" not in emit assert "emit_valid_length" in workflow["manifest"]["capabilities"] assert "emit_row_identity" not in workflow["manifest"]["capabilities"] - assert workflow["inputs"]["request.slot_ids"]["source"]["name"] == "serving.slot_ids" + 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", @@ -108,6 +109,53 @@ def test_vlm_preprocessing_is_explicit_typed_ssa(tmp_path): 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() @@ -163,7 +211,7 @@ def test_vlm_writer_derives_real_decoder_contract_from_artifact(tmp_path): present_shape, ) ) - decoder = _model("decoder", decoder_inputs, decoder_outputs) + decoder = _decoder_with_capacity_addressable_attention(decoder_inputs, decoder_outputs) vision = _model( "vision_encoder", [ @@ -245,6 +293,7 @@ def collect_decoder_invokes(node): "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", @@ -326,8 +375,8 @@ def collect_emits(node): } 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 { - "slot_ids", "token", "logits", "generated_lengths", @@ -559,3 +608,62 @@ def test_speculative_workflow_uses_per_row_ragged_state_and_rng(): "output": "present.0.key", } 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" diff --git a/tests/cli_test.py b/tests/cli_test.py index d16abc818..676f7326c 100644 --- a/tests/cli_test.py +++ b/tests/cli_test.py @@ -193,14 +193,48 @@ def test_text_only_skips_diffusers_autodetect(self): mock_build.assert_called_once() assert mock_build.call_args.kwargs.get("text_only") is True - def test_revision_is_forwarded_to_detection_and_build(self): - revision = "61ba4e0b3309b6656edea3e93e419f7bd5c61957" + def test_static_cache_with_onnx_genai_runtime_errors(self): + """static-cache graphs cannot be described by the onnx-genai contract. + + A static-cache decoder exposes in-place ring buffers plus write indices + instead of past/present KV pairs and a rank-2 attention mask, so the + workflow metadata emitter cannot describe it. The CLI must say so + before exporting the weights, not after. + """ with ( tempfile.TemporaryDirectory() as tmpdir, - mock.patch( - "mobius.integrations.diffusers._builder._load_diffusers_pipeline_index", - return_value=None, - ) as mock_diffusers, + pytest.raises(SystemExit, match=r"static-cache.*--runtime onnx-genai"), + ): + main( + [ + "build", + "--model", + "Qwen/Qwen2.5-0.5B", + tmpdir, + "--no-weights", + "--features", + "static-cache", + "--runtime", + "onnx-genai", + ] + ) + + 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"), ): @@ -208,19 +242,18 @@ def test_revision_is_forwarded_to_detection_and_build(self): [ "build", "--model", - "zai-org/GLM-ASR-Nano-2512", + "google/gemma-4-E2B-it", tmpdir, - "--revision", - revision, "--no-weights", + "--features", + "text-only,static-cache", + "--max-seq-len", + "128", ] ) - mock_diffusers.assert_called_once_with( - "zai-org/GLM-ASR-Nano-2512", - revision=revision, - ) - assert mock_build.call_args.kwargs["revision"] == revision + task = mock_build.call_args.kwargs["task"] + assert isinstance(task, Gemma4TextCausalLMTask) def test_build_static_cache(self): with tempfile.TemporaryDirectory() as tmpdir: diff --git a/tests/gemma4_prefill_prefix_test.py b/tests/gemma4_prefill_prefix_test.py index 3560184bc..2fb5cb724 100644 --- a/tests/gemma4_prefill_prefix_test.py +++ b/tests/gemma4_prefill_prefix_test.py @@ -1,391 +1,394 @@ -# 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 - assert group["aliasing"] == "permitted" - 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"] +# 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"] From 87accf70e00277ecb55112dd047daf0a9176c7af Mon Sep 17 00:00:00 2001 From: justinchuby Date: Thu, 20 Aug 2026 07:56:22 +0000 Subject: [PATCH 112/151] Declare Whisper audio preprocessing in speech-to-text workflow metadata Speech-to-text packages were still emitted through the legacy `pipeline.models`/`dataflow`/`strategy` writer, which the redesigned onnx-genai runtime rejects outright, and they carried no description of how encoded audio becomes an encoder feature tensor. That left the whole encoder-decoder audio family unrunnable from metadata: a consumer had to know, out of band, that the rank-3 encoder input is a Whisper log-mel spectrogram over a fixed 30 s window. Route the speech-to-text branch of the auto exporter at the workflow writer and derive a declarative audio preprocessing program from the HuggingFace feature-extractor config that is already written next to the package. The program is data, not a model-family branch: decode -> resample -> pad/trim to the declared window -> log-mel -> normalize, with a single output bound to the encoder's feature input. Every parameter (sampling rate, mel bins, FFT size, hop, window length) comes from the checkpoint, so a non-log-mel extractor simply produces no program and the feature tensor stays an ordinary request input. The workflow builder wires that program as an `onnx-genai.audio-preprocess` adapter component invoked once in the loop prologue, ahead of the encoder, with the encoded bytes entering as a request-level media input. The program itself is emitted at the metadata document level, a sibling of `pipeline`, because `pipeline.workflow` denies unknown fields. Two `WhisperConfig` defects surfaced while running the exported package and are fixed at the source rather than patched in the writer: `from_transformers` dropped `bos_token_id`/`eos_token_id`, so the package advertised EOS 0 and never terminated, and `max_position_embeddings` was never derived from `max_target_positions`, so the sentinel -42 leaked into `package.max_context` and the runtime refused to load the package. Verified end to end against openai/whisper-tiny: the exported package validates, and the log-mel features, encoder hidden states, token sequence, and transcript all match HuggingFace (mel max abs diff 1.2e-05, encoder relative L2 6.4e-06, identical tokens and text) as well as the Foundry ORT GenAI baseline. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: justinchuby (cherry picked from commit 8a34fcf69373cf026432db1de37e649b617df2ba) --- src/mobius/_configs/_base.py | 7 + .../integrations/onnx_genai/auto_export.py | 112 ++++++-- .../onnx_genai/auto_export_test.py | 142 ++++++---- .../speech_to_text_workflow_metadata_test.py | 249 ++++++++++++++++++ .../onnx_genai/workflow_metadata.py | 91 +++++-- 5 files changed, 512 insertions(+), 89 deletions(-) create mode 100644 src/mobius/integrations/onnx_genai/speech_to_text_workflow_metadata_test.py diff --git a/src/mobius/_configs/_base.py b/src/mobius/_configs/_base.py index 0dd76977c..8117f4162 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 diff --git a/src/mobius/integrations/onnx_genai/auto_export.py b/src/mobius/integrations/onnx_genai/auto_export.py index d077b6c48..bde636355 100644 --- a/src/mobius/integrations/onnx_genai/auto_export.py +++ b/src/mobius/integrations/onnx_genai/auto_export.py @@ -20,9 +20,6 @@ import onnx_ir as ir import yaml -from mobius.integrations.onnx_genai.decoder_metadata import ( - decoder_metadata_from_config, -) from mobius.integrations.onnx_genai.inference_metadata import ( _TEXT_RUNTIME_ASSET_NAMES, SchedulerConfig, @@ -31,7 +28,6 @@ add_explicit_package_io, add_policy_components_to_workflow, load_diffusers_scheduler_config, - write_speech_to_text_pipeline_metadata, ) from mobius.integrations.onnx_genai.workflow_metadata import ( write_audio_codec_workflow_metadata, @@ -39,6 +35,7 @@ write_diffusion_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, ) @@ -286,6 +283,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()) @@ -621,29 +700,24 @@ def write_onnx_genai_config( 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, - ) if kv_native_dtype is not None: raise ValueError( - "speech-to-text export derives KV state dtype from ONNX ports; " + "workflow speech-to-text export derives KV state dtype from ONNX ports; " "kv_native_dtype overrides are unsupported" ) - decoder_metadata = decoder_metadata_from_config(resolved_config) - path = write_speech_to_text_pipeline_metadata( + audio_processor_path = _write_hf_audio_processor(output_dir, source) + 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} + # 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)) - audio_processor_path = _write_hf_audio_processor(output_dir, source) if audio_processor_path is not None: artifacts["audio_processor"] = audio_processor_path 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 bf9feb454..3a5a77d9d 100644 --- a/src/mobius/integrations/onnx_genai/auto_export_test.py +++ b/src/mobius/integrations/onnx_genai/auto_export_test.py @@ -617,75 +617,107 @@ 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)) + 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) - # Cache storage representation is derived from the graph, never declared. + # The redesigned schema has no legacy pipeline description at all. assert "kv_cache" not in metadata - 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", - ] + 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", + } -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): 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..1a6b176f2 --- /dev/null +++ b/src/mobius/integrations/onnx_genai/speech_to_text_workflow_metadata_test.py @@ -0,0 +1,249 @@ +# 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", + } + 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 index c73b74256..e250c3c9a 100644 --- a/src/mobius/integrations/onnx_genai/workflow_metadata.py +++ b/src/mobius/integrations/onnx_genai/workflow_metadata.py @@ -5,6 +5,7 @@ from __future__ import annotations +import copy import os import re from typing import Any @@ -4566,15 +4567,11 @@ def _build_autoregressive_workflow_metadata( 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 - ] + 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 {} - ) + 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 " @@ -4771,18 +4768,32 @@ def _build_autoregressive_workflow_metadata( # 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 - preprocessing_outputs = { - binding["name"]: binding - for binding in (audio_preprocessing or {}).get("outputs", []) - } + 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 preprocessing_outputs: - encoder_invoke_inputs[value.name] = preprocessing_outputs[value.name][ - "source_value" - ] + 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), @@ -4794,8 +4805,17 @@ def _build_autoregressive_workflow_metadata( "externally_suppliable": True, } encoder_invoke_inputs[value.name] = ssa - for decoder_input, (encoder_output, _) in sorted(cross_bindings.items()): + 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]}) @@ -5377,6 +5397,17 @@ def _build_autoregressive_workflow_metadata( 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 @@ -5654,6 +5685,11 @@ def _build_autoregressive_workflow_metadata( "manifest": { "ir_version": "1.0", "onnx_opsets": {"ai.onnx": OPSET_VERSION}, + **( + {"adapter_abis": {"onnx-genai.audio-preprocess": "1"}} + if audio_program is not None + else {} + ), "capabilities": [ "workflow_ssa", "linear_effects", @@ -5661,6 +5697,7 @@ def _build_autoregressive_workflow_metadata( "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 []), ], @@ -5691,6 +5728,29 @@ def _build_autoregressive_workflow_metadata( 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, **( @@ -5723,6 +5783,7 @@ def _build_autoregressive_workflow_metadata( } 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) From 7616c08e024bee62f33731c88fcb6ad88fb4aa1e Mon Sep 17 00:00:00 2001 From: justinchuby Date: Thu, 20 Aug 2026 06:32:45 +0000 Subject: [PATCH 113/151] Fix Wav2Vec2 CTC correctness and represent it end-to-end in inference metadata MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Proves a real CTC ASR checkpoint (facebook/wav2vec2-base-960h) is fully described by the one-file inference metadata and executes to a correct transcript driven by that document alone. Wav2Vec2 was substantially broken and could not export a loadable model: * `conv_bias` was ignored — the bias initializer was declared unconditionally but never set, so ORT refused to load any checkpoint with `conv_bias=False`. * Convolution strides were ignored, producing `time - 19` frames instead of `time / 320`. The feature extractor now honors conv_dim/conv_kernel/ conv_stride/conv_bias/feat_extract_norm from the config. * `pos_conv_embed` was dropped during weight loading. Its weight-normalized grouped convolution is now materialized from `weight_g`/`weight_v` (and the newer `parametrizations.weight.original0/1` spelling). * The encoder implemented only the pre-norm layout, which is wrong whenever `do_stable_layer_norm=False`. Post-norm and pre-norm are now separate classes selected by the config. * `attention_mask` was accepted but never reached `op.Attention`, so padded frames leaked into every row. * Module-level `config_class` was missing, so config resolution — which keys off the un-rerouted `model_type` — silently fell back to `ArchitectureConfig` and hard-coded wav2vec2-base geometry for every wav2vec2 variant. `MMSConfig` gains the convolution and encoder geometry it needs plus `feature_extract_output_length`, and `CTCAsrTask` emits a `frame_lengths` output so a padded batch can be segmented without a caller re-deriving the downsampling ratio. Metadata now carries the whole contract: an audio preprocessing program, a `transcription` profile with a CTC decoding spec (blank id, repeat collapsing, time/class axes, lengths binding, inline vocabulary), and a workflow that is a plain sequence — the encoder runs once, with no loop and no carried state. A row's outputs depend on the padded width of its batch whenever the feature extractor group-normalizes over the time axis, so the module answers `batch_padding_sensitive`, the task records it on the built graph, and the producer publishes it as `batch_invariance`. When nobody states the fact it is omitted rather than guessed. `ctc_runtime` executes a package using only the emitted document, which is how the contract is proven complete: preprocessing, encoder, frame argmax, CTC collapse, blank removal, transcript. Verified against HuggingFace on testdata/652-129742-0006.flac: identical logits (max abs 2.99e-3, corr 0.99999999), 455/455 identical argmax frames, identical 111-token collapsed sequence, and an identical transcript. A B=2 batch with unequal lengths (145840 and 64000 samples) segments to frame_lengths [455, 199] matching the analytic formula, and both rows match HuggingFace fed the same padded batch exactly. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: justinchuby (cherry picked from commit de2d46fabf5b88b349f09405b39e0fc6d53b00c1) --- src/mobius/_configs/_base.py | 63 ++- .../integrations/onnx_genai/auto_export.py | 45 ++ .../integrations/onnx_genai/ctc_runtime.py | 366 ++++++++++++ .../onnx_genai/ctc_runtime_test.py | 311 ++++++++++ .../onnx_genai/workflow_metadata.py | 331 +++++++++++ src/mobius/models/wav2vec2.py | 535 ++++++++++++++---- src/mobius/models/wav2vec2_ctc.py | 108 ++-- src/mobius/tasks/_ctc_asr.py | 36 +- .../wav2vec2_ctc_metadata_integration_test.py | 132 +++++ 9 files changed, 1741 insertions(+), 186 deletions(-) create mode 100644 src/mobius/integrations/onnx_genai/ctc_runtime.py create mode 100644 src/mobius/integrations/onnx_genai/ctc_runtime_test.py create mode 100644 tests/wav2vec2_ctc_metadata_integration_test.py diff --git a/src/mobius/_configs/_base.py b/src/mobius/_configs/_base.py index 8117f4162..4ee199de1 100644 --- a/src/mobius/_configs/_base.py +++ b/src/mobius/_configs/_base.py @@ -2983,6 +2983,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 @@ -2991,24 +2998,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=tuple(getattr(config, "conv_dim", None) or defaults.conv_dim), + 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/integrations/onnx_genai/auto_export.py b/src/mobius/integrations/onnx_genai/auto_export.py index bde636355..98675da93 100644 --- a/src/mobius/integrations/onnx_genai/auto_export.py +++ b/src/mobius/integrations/onnx_genai/auto_export.py @@ -31,6 +31,7 @@ ) 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_language_diffusion_workflow_metadata, @@ -431,6 +432,33 @@ 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_audio_codec(pkg: Any) -> bool: """Detect an audio-to-audio neural codec package. @@ -656,6 +684,23 @@ def write_onnx_genai_config( 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) + 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( 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/workflow_metadata.py b/src/mobius/integrations/onnx_genai/workflow_metadata.py index e250c3c9a..360408e39 100644 --- a/src/mobius/integrations/onnx_genai/workflow_metadata.py +++ b/src/mobius/integrations/onnx_genai/workflow_metadata.py @@ -6,6 +6,7 @@ from __future__ import annotations import copy +import json import os import re from typing import Any @@ -49,11 +50,13 @@ _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, ) +from mobius.tasks._ctc_asr import BATCH_PADDING_SENSITIVE_KEY class _NoAliasSafeDumper(yaml.SafeDumper): @@ -6112,3 +6115,331 @@ def write_language_diffusion_workflow_metadata( 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": { + "ir_version": "1.0", + "onnx_opsets": {"ai.onnx": OPSET_VERSION}, + "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 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/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/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] + ) From d93abcd7bed8f58be06ca2601f86123a8766295f Mon Sep 17 00:00:00 2001 From: justinchuby Date: Thu, 20 Aug 2026 07:32:56 +0000 Subject: [PATCH 114/151] Make diffusion component exports faithful to their schedulers Sampler schedules are not integral. EulerDiscreteScheduler with five steps over the SD1.x training schedule produces timesteps 999, 749.25, 499.5, 249.75, 0, and the same is true of every scheduler that interpolates the training grid. Declaring the denoiser's timestep port as INT64 silently truncated those values, so an exported UNet could not reproduce its own pipeline no matter how accurate its weights were. Declare the port as FLOAT in DenoisingTask and ControlNetTask, which is also what diffusers feeds its UNet. Shape inference gives the denoiser estimate and the VAE decoder sample anonymous symbols (_d0, _d1, _d2) because the spatial relationship is a property of the configured block stack, not of the graph. A consumer wiring these components into a loop cannot tell that the estimate has the latent's shape or that the decoded image scales with it. Republish named dimensions on both. Denoisers that learn the variance emit twice the latent channels, so derive the count from out_channels when the config has one and from in_channels * (2 if learn_sigma else 1) otherwise, which keeps HunyuanDiT and other learn_sigma configs correct without naming any model. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: justinchuby (cherry picked from commit df045d6a7cddbfd09e11541c059dbcf3a8b0b4e4) --- src/mobius/models/unet_parity_test.py | 6 +++--- src/mobius/tasks/_controlnet.py | 2 +- src/mobius/tasks/_denoising.py | 12 +++++++++++- src/mobius/tasks/_task_test.py | 2 +- src/mobius/tasks/_vae.py | 6 ++++++ 5 files changed, 22 insertions(+), 6 deletions(-) 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/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/_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/_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/_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) From 904495fe52d6206aa0de087e2096ff638e1f347f Mon Sep 17 00:00:00 2001 From: justinchuby Date: Thu, 20 Aug 2026 07:33:45 +0000 Subject: [PATCH 115/151] Produce a diffusion workflow that a runtime can actually execute The diffusion workflow the producer emitted described a denoiser loop but left every part that makes an image out of it implicit: the runtime had to know how to seed a latent, how to apply classifier-free guidance, what a solver's carried state was, and how to normalize a latent before decoding. None of that is inferable from the graphs, so no runtime could execute the package end to end. Rebuild the workflow structurally so the metadata alone is sufficient, and prove it by running it. The non-network parts become real ONNX policy graphs. Sigma schedules, timestep tables, latent row shapes, and scalar scales are constant components, so the solver reads its constants from the graph rather than from a runtime-side scheduler registry. A counter-based RNG draws latent noise from a per-row seed and returns the advanced counter, which makes a row's noise depend only on its own seed and not on its position in the batch. A zeros-like component initializes solver history from the shape of the noise, so nothing needs a build-time resolution. A guidance component computes uncond + scale * (cond - uncond) with a per-row scale. Classifier-free guidance is two denoiser invocations rather than a doubled batch. BatchLayout can say a tensor has one row per request but has no kind meaning "k rows per request", so a doubled batch cannot be declared truthfully and the runtime would be unable to attribute a row to a request. Two invocations of the same component keep every tensor request-aligned and leave batching to the runtime. Multistep solvers carry their previous x0 estimate as ordinary loop state rather than as hidden scheduler memory, which is what lets a runtime batch rows that are at different steps. The DPM++ 2M component masks itself down to first order on the first and last steps the way lower_order_final does. Scale steps are conditional. A sampler whose state already lives in the denoiser's space and a VAE whose latents are unnormalized need no rescaling, and emitting an identity multiply there would be a lie about the pipeline. Append-mode outputs accumulate along the last axis, so the trajectory and noise-estimate outputs declare their own symbol for that axis. Reusing the per-step symbol binds it to both the chunk width and the accumulated width and fails validation on the final output. auto_export now derives the schedule from the diffusers scheduler config, reads the VAE scaling factor from vae/config.json instead of assuming one, and keeps the guard that a text-conditioned pipeline must declare how it is guided, since a silently unguided export produces plausible garbage. Adds a diffusion_guided validation fixture covering guidance, a multistep solver with history, trajectory emits, seeded noise, and a scaled decoder input, plus unit tests for each new policy component against closed-form references. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: justinchuby (cherry picked from commit 4a35a2fc08c732fb08415995648b9401c1ef23fd) --- docs/onnx-genai-workflows.md | 133 ++++- src/mobius/generation/__init__.py | 16 + src/mobius/generation/_policy_components.py | 274 ++++++++- .../generation/_policy_components_test.py | 164 ++++++ .../integrations/onnx_genai/auto_export.py | 97 +++- .../onnx_genai/auto_export_test.py | 92 ++- .../onnx_genai/inference_metadata.py | 66 ++- .../onnx_genai/workflow_metadata.py | 531 ++++++++++++++---- ...generate_onnx_genai_validation_packages.py | 22 + 9 files changed, 1233 insertions(+), 162 deletions(-) diff --git a/docs/onnx-genai-workflows.md b/docs/onnx-genai-workflows.md index af5a80e78..c5bfc7a62 100644 --- a/docs/onnx-genai-workflows.md +++ b/docs/onnx-genai-workflows.md @@ -184,35 +184,114 @@ Only the decoder body runs per generated token. ```yaml steps: - - kind: invoke - component: initialize_latent - inputs: { noise: noise } - outputs: { latent: latent.initial } - kind: loop - setup: [] + 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: denoiser - inputs: { sample: latent, step: diffusion_step } - outputs: { estimate: estimate } - - kind: invoke - component: solver - inputs: { state: latent, estimate: estimate, step: diffusion_step } - outputs: { next_state: latent.next } - condition: continue - max_iterations: num_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: diffusion_step - contract: { dtype: int64, rank: 0, shape: [] } - carried: [{ cell: latent, initial: latent.initial, next: latent.next }] - - kind: invoke - component: vae_decoder - inputs: { latent: latent } - outputs: { image: image } - - kind: emit - value: image - output: image - mode: replace + 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 } ``` -Latent initialization and VAE decoding run once; denoiser and solver run per iteration. +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/src/mobius/generation/__init__.py b/src/mobius/generation/__init__.py index f34a6c97b..e99576b68 100644 --- a/src/mobius/generation/__init__.py +++ b/src/mobius/generation/__init__.py @@ -6,6 +6,7 @@ from __future__ import annotations from mobius.generation._policy_components import ( + SOLVER_BUILDERS, PolicyCapabilities, PolicyComponent, attach_policy_components, @@ -15,6 +16,7 @@ build_code_frame_update, build_code_history_append, build_codec_layout_transpose, + build_counter_rng_normal, build_decoder_state_initializer, build_decoder_step_update, build_effectful_identity, @@ -24,6 +26,7 @@ build_euler_solver_step, build_grammar_logits_processor, build_greedy_sampler, + build_guidance_combine, build_integer_add, build_integer_minimum, build_integer_row_broadcast, @@ -31,14 +34,18 @@ build_last_token_logits, build_masked_token_update, build_model_token_cast, + build_multistep_solver_step, build_proposal_metrics, + build_scalar_constant, build_schedule_constant, build_schedule_lookup, build_seeded_categorical_sampler, build_selective_integer_add, 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, @@ -46,15 +53,18 @@ build_tts_decoder_state_initializer, build_tts_decoder_step_update, build_tts_state_initializer, + build_zeros_like, ) __all__ = [ + "SOLVER_BUILDERS", "PolicyCapabilities", "PolicyComponent", "attach_policy_components", "build_adaptive_k_policy", "build_batch_minimum", "build_boolean_not", + "build_counter_rng_normal", "build_code_frame_update", "build_code_history_append", "build_codec_layout_transpose", @@ -65,6 +75,7 @@ "build_eos_termination", "build_euler_model_input", "build_euler_solver_step", + "build_guidance_combine", "build_grammar_logits_processor", "build_greedy_sampler", "build_integer_add", @@ -76,6 +87,9 @@ "build_masked_token_update", "build_model_token_cast", "build_proposal_metrics", + "build_multistep_solver_step", + "build_scalar_constant", + "build_shape_constant", "build_schedule_constant", "build_schedule_lookup", "build_seeded_categorical_sampler", @@ -83,10 +97,12 @@ "build_speculative_acceptance", "build_speculative_state_rollback", "build_token_block_identity", + "build_tensor_scale", "build_token_state_update", "build_token_to_slot", "build_termination_batch_initializer", "build_tts_decoder_state_initializer", "build_tts_decoder_step_update", "build_tts_state_initializer", + "build_zeros_like", ] diff --git a/src/mobius/generation/_policy_components.py b/src/mobius/generation/_policy_components.py index 3fe03af97..a3148374a 100644 --- a/src/mobius/generation/_policy_components.py +++ b/src/mobius/generation/_policy_components.py @@ -78,7 +78,7 @@ def attach_policy_components( "greedy": build_greedy_sampler, "seeded_categorical": build_seeded_categorical_sampler, } - solvers = {"euler": build_euler_solver_step} + 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}: @@ -1524,6 +1524,171 @@ def build_eos_termination(*, row_selective: bool = False) -> PolicyComponent: ) +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", + }, + ) + + def build_euler_model_input(dtype: ir.DataType = ir.DataType.FLOAT) -> PolicyComponent: """Scale a latent for the Euler denoiser input at the current sigma.""" graph, builder = _make_graph("euler_model_input") @@ -1584,6 +1749,113 @@ def build_euler_solver_step( ) +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", + }, + ) + + +SOLVER_BUILDERS = { + "euler": build_euler_solver_step, + "multistep": build_multistep_solver_step, +} + + def build_masked_token_update() -> PolicyComponent: """Build replacement of masked positions with explicit RNG-counter threading.""" graph, builder = _make_graph("masked_token_update") diff --git a/src/mobius/generation/_policy_components_test.py b/src/mobius/generation/_policy_components_test.py index 62fd36125..f1f500d1b 100644 --- a/src/mobius/generation/_policy_components_test.py +++ b/src/mobius/generation/_policy_components_test.py @@ -15,6 +15,7 @@ build_batch_minimum, build_boolean_not, build_code_frame_update, + build_counter_rng_normal, build_decoder_state_initializer, build_decoder_step_update, build_empty_features, @@ -23,17 +24,23 @@ build_euler_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_proposal_metrics, + build_scalar_constant, build_seeded_categorical_sampler, + build_shape_constant, build_speculative_acceptance, build_speculative_state_rollback, + build_tensor_scale, build_termination_batch_initializer, build_token_state_update, + build_zeros_like, ) from mobius.generation._policy_components import _make_graph @@ -899,6 +906,163 @@ def test_euler_model_input_scales_by_sigma(tmp_path): 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), diff --git a/src/mobius/integrations/onnx_genai/auto_export.py b/src/mobius/integrations/onnx_genai/auto_export.py index 98675da93..21ecc94f7 100644 --- a/src/mobius/integrations/onnx_genai/auto_export.py +++ b/src/mobius/integrations/onnx_genai/auto_export.py @@ -28,6 +28,7 @@ add_explicit_package_io, add_policy_components_to_workflow, load_diffusers_scheduler_config, + load_diffusers_vae_scaling_factor, ) from mobius.integrations.onnx_genai.workflow_metadata import ( write_audio_codec_workflow_metadata, @@ -44,16 +45,43 @@ _LOGGER = logging.getLogger(__name__) -def _euler_schedule( +#: 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 Euler timesteps and sigma values.""" - if scheduler.kind != "euler" or scheduler.prediction_type != "epsilon": + """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 Euler epsilon " - f"schedulers, got kind={scheduler.kind!r}, " + "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" @@ -80,12 +108,21 @@ def _euler_schedule( 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)) - timesteps = np.linspace( - scheduler.num_train_timesteps - 1, - 0, - num_inference_steps, - dtype=np.float64, - ) + 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), @@ -649,25 +686,45 @@ def write_onnx_genai_config( derived = _diffusion_component_kwargs(pkg) for name, value in derived.items(): kwargs.setdefault(name, value) - if "text_encoder_filename" in kwargs and guidance_scale is None: + 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( - "text-conditioned workflow diffusion does not implement " - "classifier-free guidance; pass guidance_scale=1.0 explicitly " - "to request unguided generation" + "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 np.isclose(guidance_scale, 1.0): + if guidance_scale is not None and not conditioned: raise ValueError( - "workflow diffusion requires an explicit classifier-free guidance " - "component before guidance_scale can differ from 1.0" + "classifier-free guidance requires a text-conditioned diffusion package" ) - resolved_scheduler = scheduler or SchedulerConfig(kind="euler") - timesteps, sigma_schedule = _euler_schedule(resolved_scheduler, num_inference_steps) + 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, 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 diff --git a/src/mobius/integrations/onnx_genai/auto_export_test.py b/src/mobius/integrations/onnx_genai/auto_export_test.py index 3a5a77d9d..8aaa5c4a9 100644 --- a/src/mobius/integrations/onnx_genai/auto_export_test.py +++ b/src/mobius/integrations/onnx_genai/auto_export_test.py @@ -551,14 +551,102 @@ def test_workflow_vlm_rejects_kv_dtype_override(tmp_path): ) -def test_text_diffusion_requires_explicit_unguided_mode(tmp_path): - with pytest.raises(ValueError, match=r"pass guidance_scale=1\.0 explicitly"): +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. diff --git a/src/mobius/integrations/onnx_genai/inference_metadata.py b/src/mobius/integrations/onnx_genai/inference_metadata.py index 8c576ab2a..309235a5f 100644 --- a/src/mobius/integrations/onnx_genai/inference_metadata.py +++ b/src/mobius/integrations/onnx_genai/inference_metadata.py @@ -1462,6 +1462,19 @@ def annotate_strategy(strategy: dict[str, Any]) -> None: return metadata +#: 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, @@ -1530,12 +1543,9 @@ def tensor_contract(value: Any) -> dict[str, Any]: "rank": port.rank, "shape": shape, } - # Policy components are per-row operators: when the graph's leading - # dimension is the batch symbol, axis 0 carries exactly one entry per - # in-flight request. Declaring it lets the runtime permute/compact the - # batch without the producer serializing any row identity. - if shape and shape[0] == _BATCH_DIMENSION: - contract["batch_layout"] = {"kind": "request_aligned", "axis": 0} + layout = request_batch_layout(shape) + if layout is not None: + contract["batch_layout"] = layout return contract for name, component in policy_components.items(): @@ -2113,6 +2123,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] = { @@ -2233,6 +2248,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)), ) @@ -2284,6 +2304,40 @@ def load_diffusers_scheduler_config( return None +def load_diffusers_vae_scaling_factor(source: str | None) -> float | None: + """Best-effort load of a diffusers ``vae/config.json`` ``scaling_factor``. + + 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 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 + + 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( *, num_inference_steps: int, diff --git a/src/mobius/integrations/onnx_genai/workflow_metadata.py b/src/mobius/integrations/onnx_genai/workflow_metadata.py index 360408e39..493cbb973 100644 --- a/src/mobius/integrations/onnx_genai/workflow_metadata.py +++ b/src/mobius/integrations/onnx_genai/workflow_metadata.py @@ -7,6 +7,7 @@ import copy import json +import math import os import re from typing import Any @@ -16,35 +17,41 @@ from mobius._constants import OPSET_VERSION 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_decoder_state_initializer, build_decoder_step_update, build_empty_features, build_eos_termination, build_euler_model_input, - build_euler_solver_step, build_greedy_sampler, + build_guidance_combine, build_integer_add, build_integer_minimum, build_last_token_logits, build_model_token_cast, build_proposal_metrics, + build_scalar_constant, build_schedule_constant, build_schedule_lookup, build_seeded_categorical_sampler, build_selective_integer_add, build_sequence_length, + build_shape_constant, + build_tensor_scale, build_termination_batch_initializer, build_token_state_update, build_token_to_slot, build_tts_decoder_state_initializer, build_tts_decoder_step_update, build_tts_state_initializer, + build_zeros_like, ) from mobius.integrations.onnx_genai.inference_metadata import ( _name_image_preprocessing_program, @@ -55,6 +62,7 @@ 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 @@ -104,11 +112,16 @@ def _contract(value: ir.Value) -> dict[str, Any]: dtype = {"fp16": "float16", "bf16": "bfloat16", "fp32": "float32"}.get( port.dtype, port.dtype ) - return { + shape = _shape_metadata(port) + contract: dict[str, Any] = { "dtype": dtype, "rank": port.rank, - "shape": _shape_metadata(port), + "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]: @@ -190,6 +203,28 @@ def _publish_workflow_v1(workflow: dict[str, Any]) -> dict[str, Any]: 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 ports in component.get("ports", {}).values(): + for port, contract in ports.items(): + ports[port] = _declare_row_alignment(contract) substitutions: dict[str, str] = {} loop_index = 0 cell_aliases = { @@ -845,7 +880,7 @@ def _build_real_tts_workflow_metadata(pkg: Any, config: Any) -> dict[str, Any]: "required": True, }, "package.false": { - "contract": batch_bool, + "contract": _request_aligned(batch_bool), "role": {"kind": "opaque"}, "source": {"kind": "literal"}, "required": False, @@ -2344,16 +2379,50 @@ def _contracts_compatible(left: ir.Value, right: ir.Value) -> bool: ) +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 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.""" + """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), @@ -2384,7 +2453,7 @@ def build_diffusion_workflow_metadata( if len(sample_input.shape or []) != 4 or _contract(sample_input) != _contract( estimate_output ): - raise ValueError("Euler diffusion workflow requires matching rank-4 latent/estimate") + 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") @@ -2413,11 +2482,23 @@ def build_diffusion_workflow_metadata( ), 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("euler_model_input", build_euler_model_input(sample_input.dtype)) - pkg.add_policy_component("solver_step", build_euler_solver_step(sample_input.dtype)) + 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) ] @@ -2431,18 +2512,41 @@ def build_diffusion_workflow_metadata( 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.latent": { - "contract": _contract(sample_input), - "role": {"kind": "opaque"}, - "source": {"kind": "application", "name": "latent"}, - "required": True, - }, "request.max_iterations": { "contract": control_int, "role": {"kind": "runtime", "version": "1.0", "role": "max_iterations"}, @@ -2462,8 +2566,100 @@ def build_diffusion_workflow_metadata( _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 - if text_encoder is not None and conditioning_output is not 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}" @@ -2474,18 +2670,15 @@ def build_diffusion_workflow_metadata( if index == 0 else {"kind": "opaque"} ), - "source": { - "kind": "request" if index == 0 else "application", - "field": "prompt_tokens" if index == 0 else None, - "name": value.name if index else None, - }, + "source": ( + {"kind": "request", "field": "prompt_tokens"} + if index == 0 + else {"kind": "application", "name": value.name} + ), "required": True, - } - inputs[name]["source"] = { - key: item for key, item in inputs[name]["source"].items() if item is not None + "externally_suppliable": True, } text_inputs[value.name] = name - conditioning_value = "conditioning.hidden_states" setup_nodes.append( _invoke( text_name, @@ -2493,6 +2686,78 @@ def build_diffusion_workflow_metadata( {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", @@ -2501,52 +2766,84 @@ def build_diffusion_workflow_metadata( ) ) - 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]] = [] - body_nodes.append( + body_nodes: list[dict[str, Any]] = [ _invoke( "schedule_lookup", - { - "schedule": "diffusion.timesteps", - "step": "loop.iteration", - }, + {"schedule": "diffusion.timesteps", "step": "loop.iteration"}, {"timestep": "diffusion.timestep"}, ) - ) - body_nodes.append( - _invoke( - "euler_model_input", - { - "sample": "state.latent.body", - "step": "loop.iteration", - "schedule": "diffusion.schedule", - }, - {"model_input": "diffusion.model_input"}, - ) - ) - body_nodes.extend( - [ - _invoke( - denoiser_name, - denoiser_inputs, - {estimate_output.name: "denoiser.estimate"}, - ), + ] + if scale_model_input: + model_input_value = "diffusion.model_input" + body_nodes.append( _invoke( - "solver_step", + "model_input_scale", { "sample": "state.latent.body", - "derivative": "denoiser.estimate", "step": "loop.iteration", "schedule": "diffusion.schedule", }, - {"next_state": "latent.body"}, - {"solver": _effect("solver.0", "solver.1")}, - ), + {"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"}, @@ -2554,7 +2851,55 @@ def build_diffusion_workflow_metadata( ), ] ) - latent_effect = "state:latent" + + 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": { "ir_version": "1.0", @@ -2568,28 +2913,14 @@ def build_diffusion_workflow_metadata( ], }, "inputs": inputs, - "outputs": { - "image": { - "contract": _contract(vae_output), - "role": "image", - "stage": "pre_adapter", - } - }, + "outputs": outputs, "components": { name: _component(model, _artifact(name, len(pkg))) for name, model in pkg.items() }, - "state": { - "latent": { - "contract": _contract(sample_input), - "scope": "invocation", - "initializer": "request.latent", - "recurrence": {"kind": "invariant"}, - } - }, + "state": state, "initial_effects": { - "solver": "solver.0", - latent_effect: f"{latent_effect}.0", "emit": "emit.0", + **{f"state:{cell}": f"state:{cell}.0" for cell in state}, }, "graph": { "kind": "sequence", @@ -2601,35 +2932,9 @@ def build_diffusion_workflow_metadata( "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( - vae_name, - {vae_input.name: "latent.final"}, - {vae_output.name: "vae.image"}, - ), - { - "kind": "emit", - "value": "vae.image", - "output": "image", - "mode": "replace", - "effect_name": "emit", - "effect": _effect("emit.0", "emit.1"), + "carried": carried, }, + *tail_nodes, ], }, } @@ -2648,6 +2953,13 @@ def write_diffusion_workflow_metadata( 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( @@ -2655,6 +2967,13 @@ def write_diffusion_workflow_metadata( 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) diff --git a/tests/generate_onnx_genai_validation_packages.py b/tests/generate_onnx_genai_validation_packages.py index 33791845c..2a3949502 100644 --- a/tests/generate_onnx_genai_validation_packages.py +++ b/tests/generate_onnx_genai_validation_packages.py @@ -32,6 +32,7 @@ ) from mobius.integrations.onnx_genai.workflow_metadata import ( write_audio_codec_workflow_metadata, + write_diffusion_workflow_metadata, write_language_diffusion_workflow_metadata, write_speculative_workflow_metadata, ) @@ -745,6 +746,27 @@ def main() -> None: 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) From 6743366bf1d36977977d36dd0597ea2c936bf68c Mon Sep 17 00:00:00 2001 From: justinchuby Date: Thu, 20 Aug 2026 09:38:37 +0000 Subject: [PATCH 116/151] Make the Qwen Image VAE exportable and loadable in bfloat16 Exporting the Qwen Image Edit VAE at the pinned 2509 checkpoint failed at three independent points before onnxruntime could even create a session. Each cause is generic, so each fix is generic. 1. `_RMSNorm3d` reduced in the model dtype. onnxruntime ships no bfloat16 `ReduceL2` kernel on any execution provider, and a single unassignable node aborts session creation outright; the same reduction also overflows in float16. The L2 reduction now runs in float32 and casts back. 2. `_Resample` applied `time_conv` and the temporal pixel shuffle unconditionally. Diffusers' `QwenImageResample.forward` only runs that branch from the *second* cached chunk onward, so a single-frame image never takes it. Applying it to a T=1 input drove a Conv output dimension to zero and produced a graph onnxruntime rejected outright. The forward pass now matches the single-chunk image path, and `Resize` -- which also has no bfloat16 kernel -- is sandwiched between casts. 3. The decoder's output clamp lowered to `Clip`, which has no bfloat16 kernel either; onnxruntime expands its ONNX function body into `Less`/`Where`, which are equally unavailable. A new `clip_to_min_max_rules()` rewrite rule lowers `Clip` to the exactly equivalent, kernel-backed `Min(Max(...))` form, and `_optimizations` applies it as a lowering stage whenever the export dtype is bfloat16. Tests cover all three: a real diffusers comparison with `temperal_downsample=(True,)` (the only configuration that instantiates `time_conv`) pinning both the surviving frame axis and 1e-4 numerical agreement, a graph-level assertion that no bfloat16 `ReduceL2`, `Resize` or `Clip` node survives optimization, and four unit tests for the rewrite rule itself. With these applied, the bfloat16 CUDA export of the official checkpoint reaches cos 0.99997 on the encoder and 53.31 dB PSNR on the decoder against diffusers 0.39.0. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: justinchuby (cherry picked from commit 0df5e91718c880d1305ff0364480bcc355263481) --- src/mobius/_optimizations.py | 8 ++ src/mobius/models/qwen_image_test.py | 112 ++++++++++++++++++ src/mobius/models/qwen_image_vae.py | 73 +++++------- src/mobius/rewrite_rules/__init__.py | 2 + src/mobius/rewrite_rules/_clip_to_min_max.py | 81 +++++++++++++ .../rewrite_rules/_clip_to_min_max_test.py | 90 ++++++++++++++ 6 files changed, 324 insertions(+), 42 deletions(-) create mode 100644 src/mobius/rewrite_rules/_clip_to_min_max.py create mode 100644 src/mobius/rewrite_rules/_clip_to_min_max_test.py diff --git a/src/mobius/_optimizations.py b/src/mobius/_optimizations.py index af5af8f93..efc83d70f 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 diff --git a/src/mobius/models/qwen_image_test.py b/src/mobius/models/qwen_image_test.py index fe185adbd..1c9922862 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._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/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)) From c6acf9f2ba72f18deb65a19e078d8b7fd6fe4d02 Mon Sep 17 00:00:00 2001 From: justinchuby Date: Thu, 20 Aug 2026 09:39:10 +0000 Subject: [PATCH 117/151] Add flow-matching image-edit policy components Five model-agnostic ONNX policy components that a source-image-conditioned flow-matching editing pipeline needs and the existing latent-diffusion builders cannot supply: - `build_flow_match_solver_step`: the Euler update on rank-3 packed tokens. `build_euler_solver_step` assumes rank-4 spatial latents, which patchified transformers never carry. - `build_pack_latents_2x2` / `build_unpack_latents_2x2`: the `(B,C,T,H,W) <-> (B,T*H/2*W/2,C*4)` patchify pair, with shapes derived from the input at runtime so one component serves every resolution. - `build_sequence_concat`: joins target and source token blocks so the denoiser sees both while the loop carries only the target. - `build_true_cfg`: true classifier-free guidance, i.e. `neg + s*(cond-neg)` renormalized to the conditional's channel norm. This is distinct from the plain-scaling CFG the diffusion workflow rejects, and it collapses to the identity at `guidance_scale=1.0`. Each is verified against a numpy transcription of the corresponding diffusers reference, plus an exact pack/unpack round trip. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: justinchuby (cherry picked from commit a465908cf0e36e8681c0286d7fd79757ba7f93b5) --- src/mobius/generation/__init__.py | 10 ++ src/mobius/generation/_policy_components.py | 167 +++++++++++++++++- .../generation/_policy_components_test.py | 100 +++++++++++ 3 files changed, 274 insertions(+), 3 deletions(-) diff --git a/src/mobius/generation/__init__.py b/src/mobius/generation/__init__.py index e99576b68..ca468524e 100644 --- a/src/mobius/generation/__init__.py +++ b/src/mobius/generation/__init__.py @@ -24,6 +24,7 @@ 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, @@ -35,12 +36,14 @@ build_masked_token_update, build_model_token_cast, build_multistep_solver_step, + build_pack_latents_2x2, build_proposal_metrics, build_scalar_constant, build_schedule_constant, build_schedule_lookup, build_seeded_categorical_sampler, build_selective_integer_add, + build_sequence_concat, build_sequence_length, build_shape_constant, build_speculative_acceptance, @@ -50,9 +53,11 @@ 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_zeros_like, ) @@ -76,6 +81,7 @@ "build_euler_model_input", "build_euler_solver_step", "build_guidance_combine", + "build_flow_match_solver_step", "build_grammar_logits_processor", "build_greedy_sampler", "build_integer_add", @@ -91,7 +97,11 @@ "build_scalar_constant", "build_shape_constant", "build_schedule_constant", + "build_pack_latents_2x2", "build_schedule_lookup", + "build_sequence_concat", + "build_true_cfg", + "build_unpack_latents_2x2", "build_seeded_categorical_sampler", "build_sequence_length", "build_speculative_acceptance", diff --git a/src/mobius/generation/_policy_components.py b/src/mobius/generation/_policy_components.py index a3148374a..9ad484dd6 100644 --- a/src/mobius/generation/_policy_components.py +++ b/src/mobius/generation/_policy_components.py @@ -895,9 +895,7 @@ def build_decoder_step_update( 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 - ) + 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") @@ -1850,6 +1848,169 @@ def row_scalar(value): ) +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_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, diff --git a/src/mobius/generation/_policy_components_test.py b/src/mobius/generation/_policy_components_test.py index f1f500d1b..17a374417 100644 --- a/src/mobius/generation/_policy_components_test.py +++ b/src/mobius/generation/_policy_components_test.py @@ -22,6 +22,7 @@ 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, @@ -31,15 +32,19 @@ build_masked_token_update, build_model_token_cast, build_multistep_solver_step, + build_pack_latents_2x2, build_proposal_metrics, build_scalar_constant, 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_zeros_like, ) from mobius.generation._policy_components import _make_graph @@ -1170,3 +1175,98 @@ def test_state_initializer_allocates_fp8_cache_through_a_cast(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) From 26ba902e76d3d4e88684294acf6235c9c39ec212 Mon Sep 17 00:00:00 2001 From: justinchuby Date: Thu, 20 Aug 2026 09:39:30 +0000 Subject: [PATCH 118/151] Emit executable onnx-genai metadata for image-edit pipelines `write_onnx_genai_config` previously refused any Qwen Image Edit package outright, so the exported components had no runnable deployment path. It now recognizes the shape structurally and emits a complete workflow. Detection is structural, not name-based: `_looks_like_image_edit` keys off a VAE encoder/decoder pair plus a denoiser that takes rank-3 packed latents and exposes a `target_sequence_length` port, so any pipeline built the same way dispatches here without a registry entry. `build_image_edit_workflow_metadata` emits the full path -- encode and pack the source image once in loop setup, then per step concatenate target and source tokens, run the denoiser twice for positive and negative conditioning, combine with true CFG, and integrate; finally unpack and decode. Positive and negative conditioning are separate application inputs because the two prompts tokenize to different lengths and cannot share a contract dimension. `_flow_match_euler_schedule` materializes the diffusers `FlowMatchEulerDiscreteScheduler` timesteps, including resolution-dependent dynamic shifting and terminal stretching. The schedule is resolution dependent, so `image_seq_len` must be supplied rather than guessed; both it and the scheduler config are hard requirements with explicit errors. The emitted 8-step schedule matches the one the upstream pipeline computed for the 1216x864 reference edit to 1.2e-7. The test that pinned the old rejection is replaced by tests asserting the emitted loop structure, the separated conditioning contracts, the policy artifacts on disk, the structural detector, and the schedule values. The `testdata` case records the official-weight validation this enables: the emitted metadata executes in the onnx-genai Rust runtime and reproduces the upstream edit at 37.10 dB PSNR (cos 0.99971). Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: justinchuby (cherry picked from commit 0f74f066de15d537a4a73849043ec7a237029d9d) --- .../integrations/onnx_genai/__init__.py | 4 + .../integrations/onnx_genai/auto_export.py | 106 +++++- .../onnx_genai/auto_export_test.py | 236 ++++++++++-- .../onnx_genai/workflow_metadata.py | 340 ++++++++++++++++++ .../cases/diffusion/qwen-image-edit-2509.yaml | 19 +- 5 files changed, 655 insertions(+), 50 deletions(-) diff --git a/src/mobius/integrations/onnx_genai/__init__.py b/src/mobius/integrations/onnx_genai/__init__.py index 7abcae6b1..8a26706a3 100644 --- a/src/mobius/integrations/onnx_genai/__init__.py +++ b/src/mobius/integrations/onnx_genai/__init__.py @@ -64,6 +64,7 @@ build_audio_codec_workflow_metadata, build_decoder_workflow_metadata, build_diffusion_workflow_metadata, + build_image_edit_workflow_metadata, build_language_diffusion_pipeline_metadata, build_speculative_workflow_metadata, build_tts_workflow_metadata, @@ -71,6 +72,7 @@ write_audio_codec_workflow_metadata, write_decoder_workflow_metadata, write_diffusion_workflow_metadata, + write_image_edit_workflow_metadata, write_language_diffusion_workflow_metadata, write_speculative_workflow_metadata, write_tts_workflow_metadata, @@ -85,6 +87,7 @@ "build_decoder_metadata", "build_decoder_workflow_metadata", "build_diffusion_workflow_metadata", + "build_image_edit_workflow_metadata", "build_diffusion_pipeline_metadata", "build_language_diffusion_pipeline_metadata", "build_speculative_workflow_metadata", @@ -105,6 +108,7 @@ "write_decoder_metadata", "write_decoder_workflow_metadata", "write_diffusion_workflow_metadata", + "write_image_edit_workflow_metadata", "write_language_diffusion_workflow_metadata", "write_speculative_workflow_metadata", "write_diffusion_pipeline_metadata", diff --git a/src/mobius/integrations/onnx_genai/auto_export.py b/src/mobius/integrations/onnx_genai/auto_export.py index 21ecc94f7..c31c2f734 100644 --- a/src/mobius/integrations/onnx_genai/auto_export.py +++ b/src/mobius/integrations/onnx_genai/auto_export.py @@ -35,6 +35,7 @@ write_ctc_asr_workflow_metadata, write_decoder_workflow_metadata, write_diffusion_workflow_metadata, + write_image_edit_workflow_metadata, write_language_diffusion_workflow_metadata, write_speculative_workflow_metadata, write_speech_to_text_workflow_metadata, @@ -134,6 +135,70 @@ def _diffusion_schedule( _DENOISER_KEYS = ("denoiser", "transformer", "unet") +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: + names = set(pkg.keys()) + except AttributeError: + 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 _add_explicit_io_to_file(path: str, pkg: Any, config: Any) -> None: """Augment an emitted sidecar with roles derived from the actual ONNX ports.""" try: @@ -637,6 +702,7 @@ 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 typed SSA workflow Multimodal VLM decoder + vision/audio encoder + fusion ``composite`` (encoders→fuse→AR) @@ -668,17 +734,37 @@ def write_onnx_genai_config( return artifacts if _looks_like_diffusion(pkg): - is_qwen_image_edit = getattr(getattr(pkg, "config", None), "model_type", None) == ( - "qwen_image_edit" - ) - 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." + 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) + 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 diff --git a/src/mobius/integrations/onnx_genai/auto_export_test.py b/src/mobius/integrations/onnx_genai/auto_export_test.py index 8aaa5c4a9..7cd24d6eb 100644 --- a/src/mobius/integrations/onnx_genai/auto_export_test.py +++ b/src/mobius/integrations/onnx_genai/auto_export_test.py @@ -16,6 +16,11 @@ 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, @@ -382,53 +387,206 @@ def test_single_diffusion_component_requires_explicit_vae(tmp_path): write_onnx_genai_config(pkg, str(tmp_path), num_inference_steps=2) -def test_rejects_unsupported_qwen_image_edit_runtime_export(tmp_path): +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_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"] - pkg = _DiffusionPkg( - { - "transformer": object(), - "text_encoder": object(), - "text_encoder_vision_encoder": object(), - "text_encoder_embedding": object(), - "vae_encoder": object(), - "vae_decoder": object(), - } + 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, + ] + + 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, ) diff --git a/src/mobius/integrations/onnx_genai/workflow_metadata.py b/src/mobius/integrations/onnx_genai/workflow_metadata.py index 493cbb973..3e0101f42 100644 --- a/src/mobius/integrations/onnx_genai/workflow_metadata.py +++ b/src/mobius/integrations/onnx_genai/workflow_metadata.py @@ -30,27 +30,32 @@ build_empty_features, build_eos_termination, build_euler_model_input, + build_flow_match_solver_step, build_greedy_sampler, build_guidance_combine, 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_schedule_constant, 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_zeros_like, ) from mobius.integrations.onnx_genai.inference_metadata import ( @@ -2983,6 +2988,341 @@ def write_diffusion_workflow_metadata( 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": { + "ir_version": "1.0", + "onnx_opsets": {"ai.onnx": OPSET_VERSION}, + "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_vlm_workflow_metadata( pkg: Any, config: Any, 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)" From 228a2d7b2b081ec55cde40fc41939687824b359d Mon Sep 17 00:00:00 2001 From: Copilot <223556219+Copilot@users.noreply.github.com> Date: Thu, 20 Aug 2026 06:49:07 +0000 Subject: [PATCH 119/151] Free the CogVideoX denoiser from its baked frame count and resolution The patch embedder built its 3D sincos positional table from the checkpoint config (sample_frames/sample_height/sample_width) and baked it into an initializer, so an exported graph silently produced wrong positions for any request that asked for a different frame count or resolution. diffusers recomputes that table per call, so the export was not a faithful replacement. Build the table inside the graph instead. `_sincos_in_graph` evaluates a 1-D sincos embedding for a dynamic extent, and `_sincos_3d_in_graph` composes the separable temporal (D/4) and spatial (3D/4) halves the way diffusers does, including the `meshgrid(..., indexing="xy")` ordering that embeds width first and the zero positions assigned to the text prefix. Checkpoints that really do learn the table (`use_learned_positional_embeddings`) keep the lookup path. Also reshape the stride-p, p x p patch-embed convolution weight into the equivalent linear weight, and declare `noise_pred.shape` on the task output. Unpatchify recovers the spatial extents through Shape ops, which mints fresh symbolic dimensions; a solver that carries the latent across steps needs the latent and the estimate contracts to unify, and they only do if the task says the prediction is elementwise with the latent. Verified against diffusers CogVideoXTransformer3DModel (finetrainers/dummy-cogvideox @ bed7eacd) at 9 frames, 17 frames, 48x32, and batch 2: max abs difference 3.0e-5 to 4.6e-5, correlation 1.0. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: Copilot <223556219+Copilot@users.noreply.github.com> (cherry picked from commit d4d5af4f32199542155f4456fcfe2a4394fc2ce4) --- src/mobius/integrations/diffusers/_configs.py | 4 + src/mobius/models/cogvideox.py | 181 +++++++++++++++--- src/mobius/tasks/_video_denoising.py | 7 + 3 files changed, 164 insertions(+), 28 deletions(-) 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/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/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") From d23245e987f8adbc7127c44ccaf95e1a64535425 Mon Sep 17 00:00:00 2001 From: Copilot <223556219+Copilot@users.noreply.github.com> Date: Thu, 20 Aug 2026 06:50:27 +0000 Subject: [PATCH 120/151] Add a faithful CogVideoX causal 3D VAE decoder with conv-cache chunking `AutoencoderKLCogVideoX` was mapped onto the generic `VideoAutoencoderModel`, which is not the diffusers architecture: it has neither the causal temporal padding nor the spatial-norm conditioning, so a decode of real CogVideoX weights produced the wrong video. Replace the mapping with a decoder that reproduces `AutoencoderKLCogVideoX` exactly. The interesting part is chunking. CogVideoX decodes long clips a few latent frames at a time, and this is not a free optimization: `CogVideoXSpatialNorm3D` runs a GroupNorm that reduces over the whole time axis, so the result depends on where the boundaries are. diffusers resolves this by threading a `conv_cache` of the last k-1 temporally-padded frames through every causal convolution, and the decode is only reproducible if the export does the same. `CogVideoXCausalConv3d` therefore exposes its cache as graph I/O: the task declares 22 paired `conv_cache.* / conv_cache_out.*` ports at two spatial resolutions and walks the same chunk schedule diffusers does (`num_batches = max(T_lat // 2, 1)`, with the remainder folded into the first chunk). The first chunk passes zero-element caches and the graph replicates frame 0 instead, branch-free, via `Concat` followed by an edge `Pad` sized from the actual cache extent. Each cache port records its spatial scale in `metadata_props`, because ONNX shape inference is free to replace a declared symbolic dimension with an anonymous one and a consumer must not have to parse dimension name strings to size an empty cache. Verified against diffusers with real weights (finetrainers/dummy-cogvideox @ bed7eacd): single chunk max abs 4.8e-7 (149.1 dB PSNR), two chunks 5.4e-7 (148.3 dB), batch 2 9.5e-7 (145.9 dB). Causality is checked directly by perturbing the final latent frame: the already-published frames stay bit-identical (delta 0.0) while the current chunk moves by 16.2. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: Copilot <223556219+Copilot@users.noreply.github.com> (cherry picked from commit d9b653b083a2f8cf955a7acecc73dbd48fa35c38) --- src/mobius/integrations/diffusers/_builder.py | 11 +- .../integrations/diffusers/_builder_test.py | 1 + src/mobius/models/__init__.py | 6 + src/mobius/models/cogvideox_vae.py | 665 ++++++++++++++++++ src/mobius/tasks/__init__.py | 3 + src/mobius/tasks/_video_vae.py | 103 +++ tests/build_graph_test.py | 71 ++ 7 files changed, 856 insertions(+), 4 deletions(-) create mode 100644 src/mobius/models/cogvideox_vae.py create mode 100644 src/mobius/tasks/_video_vae.py 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/models/__init__.py b/src/mobius/models/__init__.py index 51b829c2f..1f4863fc8 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", @@ -188,6 +190,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 diff --git a/src/mobius/models/cogvideox_vae.py b/src/mobius/models/cogvideox_vae.py new file mode 100644 index 000000000..e718feab1 --- /dev/null +++ b/src/mobius/models/cogvideox_vae.py @@ -0,0 +1,665 @@ +# 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)) + temporal_levels = int(math.log2(config.temporal_compression_ratio)) + 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/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/_video_vae.py b/src/mobius/tasks/_video_vae.py new file mode 100644 index 000000000..3519854fc --- /dev/null +++ b/src/mobius/tasks/_video_vae.py @@ -0,0 +1,103 @@ +# 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/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.""" From e58d35b42b5e2134f30436ebd3382c407c7ad3d8 Mon Sep 17 00:00:00 2001 From: Copilot <223556219+Copilot@users.noreply.github.com> Date: Thu, 20 Aug 2026 06:52:01 +0000 Subject: [PATCH 121/151] Produce video diffusion workflow metadata with temporal state The workflow metadata producer could only describe image diffusion. Its latent state was rank-4, its solver components indexed the trailing axes directly, and its published output grew along the last axis, so nothing in the emitted document could express "the latent has a time dimension" or "frames are published a chunk at a time". Generalize it instead of adding a second, parallel image-shaped path. The Euler components now take an explicit `latent_dims`, so the same graphs serve rank-4 and rank-5 latents. Alongside them: - `build_ddim_solver_step` and `build_identity_model_input`, because the CogVideoX pipeline uses DDIM with `clip_sample`, not Euler, and reporting Euler numbers for a DDIM pipeline would not be parity. - `build_video_latent_initializer` seeds the rank-5 latent and an int64 scheduler history; `build_schedule_history_append` grows it. - `build_video_decode_chunk_count` and `build_video_decode_chunk` walk the causal decoder's chunk schedule inside the workflow. - `build_video_conv_cache_initializer` sizes each zero-element cache from the spatial scale the decoder recorded in `metadata_props`. - `build_video_latent_permute` and `build_video_latent_unscale` bridge the denoiser's [B, T, C, H, W] to the VAE's [B, C, T, H, W]. `build_video_diffusion_workflow_metadata` assembles these into a document whose latent state is rank-5, whose decode loop is bounded by a setup-computed chunk count, and whose `decode.frames` emit carries `axis: 2` so the runtime concatenates along time rather than assuming the sequence axis. Packages without a text encoder declare `request.encoder_hidden_states` as application-supplied conditioning rather than silently invoking the denoiser without it. `tests/generate_onnx_genai_validation_packages.py` gains an executable `video` package (rank-5 denoiser plus a causal two-resolution chunked decoder) and the conformance suite asserts the runtime publishes causal temporal chunks from it. End-to-end against diffusers with real weights (finetrainers/dummy-cogvideox @ bed7eacd), driving the Rust onnx-genai engine from the generated document alone: 9 frames at 16x16 max abs 7.6e-5 / 81.6 dB, 17 frames 7.5e-5 / 84.7 dB, 16 frames at 48x32 8.6e-5 / 87.1 dB, batch 2 7.6e-5 / 83.7 dB, correlation 1.0 everywhere. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: Copilot <223556219+Copilot@users.noreply.github.com> (cherry picked from commit b65f7f60b256a254b3920e1a973a2a046a4637e1) --- src/mobius/generation/__init__.py | 42 +- src/mobius/generation/_policy_components.py | 306 +++++++++- .../generation/_policy_components_test.py | 104 ++++ .../integrations/onnx_genai/__init__.py | 20 +- .../onnx_genai/workflow_metadata.py | 577 ++++++++++++++++++ .../onnx_genai/workflow_metadata_test.py | 120 ++++ ...generate_onnx_genai_validation_packages.py | 205 +++++++ tests/onnx_genai_workflow_conformance.rs | 67 ++ 8 files changed, 1405 insertions(+), 36 deletions(-) diff --git a/src/mobius/generation/__init__.py b/src/mobius/generation/__init__.py index ca468524e..69ca0840c 100644 --- a/src/mobius/generation/__init__.py +++ b/src/mobius/generation/__init__.py @@ -17,6 +17,7 @@ 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_effectful_identity, @@ -28,6 +29,7 @@ build_grammar_logits_processor, build_greedy_sampler, build_guidance_combine, + build_identity_model_input, build_integer_add, build_integer_minimum, build_integer_row_broadcast, @@ -40,6 +42,7 @@ build_proposal_metrics, build_scalar_constant, build_schedule_constant, + build_schedule_history_append, build_schedule_lookup, build_seeded_categorical_sampler, build_selective_integer_add, @@ -58,6 +61,12 @@ 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, ) @@ -69,10 +78,11 @@ "build_adaptive_k_policy", "build_batch_minimum", "build_boolean_not", - "build_counter_rng_normal", "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_effectful_identity", @@ -80,39 +90,47 @@ "build_eos_termination", "build_euler_model_input", "build_euler_solver_step", - "build_guidance_combine", "build_flow_match_solver_step", "build_grammar_logits_processor", "build_greedy_sampler", + "build_guidance_combine", + "build_identity_model_input", "build_integer_add", - "build_integer_row_broadcast", - "build_selective_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_proposal_metrics", "build_multistep_solver_step", + "build_pack_latents_2x2", + "build_proposal_metrics", "build_scalar_constant", - "build_shape_constant", "build_schedule_constant", - "build_pack_latents_2x2", + "build_schedule_history_append", "build_schedule_lookup", - "build_sequence_concat", - "build_true_cfg", - "build_unpack_latents_2x2", "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_token_block_identity", "build_tensor_scale", + "build_termination_batch_initializer", + "build_token_block_identity", "build_token_state_update", "build_token_to_slot", - "build_termination_batch_initializer", + "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", ] diff --git a/src/mobius/generation/_policy_components.py b/src/mobius/generation/_policy_components.py index 9ad484dd6..6fff6c73b 100644 --- a/src/mobius/generation/_policy_components.py +++ b/src/mobius/generation/_policy_components.py @@ -11,6 +11,7 @@ from __future__ import annotations import json +from collections.abc import Sequence from dataclasses import dataclass from typing import Protocol @@ -1687,17 +1688,28 @@ def build_counter_rng_normal(dtype: ir.DataType = ir.DataType.FLOAT) -> PolicyCo ) -def build_euler_model_input(dtype: ir.DataType = ir.DataType.FLOAT) -> PolicyComponent: - """Scale a latent for the Euler denoiser input at the current sigma.""" +_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, ["batch", "channels", "height", "width"]) + 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=[1, 2, 3])) + 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") @@ -1706,20 +1718,17 @@ def build_euler_model_input(dtype: ir.DataType = ir.DataType.FLOAT) -> PolicyCom 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)``.""" + """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, - ["batch", "channels", "height", "width"], - ) - derivative = builder.input( - "derivative", - dtype, - ["batch", "channels", "height", "width"], - ) + 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])) @@ -1728,7 +1737,7 @@ def build_euler_solver_step( 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=[1, 2, 3])) + 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( @@ -1889,6 +1898,69 @@ def build_flow_match_solver_step( ) +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. @@ -2017,6 +2089,30 @@ def build_true_cfg( } +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") @@ -2336,3 +2432,181 @@ def build_token_to_slot() -> PolicyComponent: 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, {}) diff --git a/src/mobius/generation/_policy_components_test.py b/src/mobius/generation/_policy_components_test.py index 17a374417..731ffcd6e 100644 --- a/src/mobius/generation/_policy_components_test.py +++ b/src/mobius/generation/_policy_components_test.py @@ -16,6 +16,7 @@ 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, @@ -35,6 +36,7 @@ build_pack_latents_2x2, build_proposal_metrics, build_scalar_constant, + build_schedule_history_append, build_seeded_categorical_sampler, build_sequence_concat, build_shape_constant, @@ -45,6 +47,10 @@ 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, ) from mobius.generation._policy_components import _make_graph @@ -757,6 +763,104 @@ def test_euler_solver_runtime_parity(tmp_path): 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 diff --git a/src/mobius/integrations/onnx_genai/__init__.py b/src/mobius/integrations/onnx_genai/__init__.py index 8a26706a3..f389b92c1 100644 --- a/src/mobius/integrations/onnx_genai/__init__.py +++ b/src/mobius/integrations/onnx_genai/__init__.py @@ -68,6 +68,7 @@ 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, @@ -76,6 +77,7 @@ write_language_diffusion_workflow_metadata, write_speculative_workflow_metadata, write_tts_workflow_metadata, + write_video_diffusion_workflow_metadata, write_vlm_workflow_metadata, ) @@ -84,38 +86,40 @@ "ConversionResult", "SchedulerConfig", "add_policy_components_to_workflow", + "build_audio_codec_workflow_metadata", "build_decoder_metadata", "build_decoder_workflow_metadata", + "build_diffusion_pipeline_metadata", "build_diffusion_workflow_metadata", "build_image_edit_workflow_metadata", - "build_diffusion_pipeline_metadata", "build_language_diffusion_pipeline_metadata", - "build_speculative_workflow_metadata", - "build_audio_codec_workflow_metadata", "build_multimodal_pipeline_metadata", "build_pipeline_metadata_for_workflow", + "build_speculative_workflow_metadata", "build_speech_to_text_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_decoder_metadata", "write_decoder_workflow_metadata", + "write_diffusion_pipeline_metadata", "write_diffusion_workflow_metadata", "write_image_edit_workflow_metadata", "write_language_diffusion_workflow_metadata", - "write_speculative_workflow_metadata", - "write_diffusion_pipeline_metadata", - "write_audio_codec_workflow_metadata", "write_multimodal_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", - "write_onnx_genai_config", ] diff --git a/src/mobius/integrations/onnx_genai/workflow_metadata.py b/src/mobius/integrations/onnx_genai/workflow_metadata.py index 3e0101f42..c5a530ff3 100644 --- a/src/mobius/integrations/onnx_genai/workflow_metadata.py +++ b/src/mobius/integrations/onnx_genai/workflow_metadata.py @@ -25,6 +25,7 @@ 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_empty_features, @@ -33,6 +34,7 @@ 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, @@ -41,6 +43,7 @@ build_proposal_metrics, build_scalar_constant, build_schedule_constant, + build_schedule_history_append, build_schedule_lookup, build_seeded_categorical_sampler, build_selective_integer_add, @@ -56,6 +59,12 @@ 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, ) from mobius.integrations.onnx_genai.inference_metadata import ( @@ -289,6 +298,8 @@ def convert(node: dict[str, Any]) -> dict[str, Any]: "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: @@ -2396,6 +2407,29 @@ def _accumulated_contract(contract: dict[str, Any], symbol: str) -> dict[str, An 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( @@ -3323,6 +3357,549 @@ def write_image_edit_workflow_metadata( 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": { + "ir_version": "1.0", + "onnx_opsets": {"ai.onnx": OPSET_VERSION}, + "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, diff --git a/src/mobius/integrations/onnx_genai/workflow_metadata_test.py b/src/mobius/integrations/onnx_genai/workflow_metadata_test.py index 407ee4634..732285406 100644 --- a/src/mobius/integrations/onnx_genai/workflow_metadata_test.py +++ b/src/mobius/integrations/onnx_genai/workflow_metadata_test.py @@ -20,6 +20,7 @@ _kv_storage_contract, 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, @@ -463,6 +464,125 @@ def test_language_diffusion_uses_exclusive_ssa_workflow(): 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( diff --git a/tests/generate_onnx_genai_validation_packages.py b/tests/generate_onnx_genai_validation_packages.py index 2a3949502..4b9c28e8d 100644 --- a/tests/generate_onnx_genai_validation_packages.py +++ b/tests/generate_onnx_genai_validation_packages.py @@ -35,6 +35,7 @@ write_diffusion_workflow_metadata, write_language_diffusion_workflow_metadata, write_speculative_workflow_metadata, + write_video_diffusion_workflow_metadata, ) from mobius.models.qwen3_tts import Qwen3TTSForConditionalGeneration from mobius.models.qwen3_tts_test import _TINY_CONFIG @@ -327,6 +328,196 @@ def _executable_diffusion_package() -> ModelPackage: ) +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]) @@ -786,6 +977,20 @@ def main() -> None: 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) diff --git a/tests/onnx_genai_workflow_conformance.rs b/tests/onnx_genai_workflow_conformance.rs index 0b5db4e06..4f40e17d5 100644 --- a/tests/onnx_genai_workflow_conformance.rs +++ b/tests/onnx_genai_workflow_conformance.rs @@ -660,3 +660,70 @@ fn mobius_speculative_workflow_executes_rejection_and_correction() -> anyhow::Re 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(()) +} From 2dc969cbd193df4d1d7164d4fd2b37856fe11c86 Mon Sep 17 00:00:00 2001 From: justinchuby Date: Thu, 20 Aug 2026 08:12:25 +0000 Subject: [PATCH 122/151] Shape generation positions for decoders with multi-axis rotary embeddings The policy builders assumed every decoder reads (batch, sequence) positions. A decoder with multi-axis rotary embeddings reads (sections, batch, sequence), where the leading axis is a fixed count of rotary axes, so the prefill range and the per-step increment were built one rank too small. ONNX Runtime accepts the rank-2 tensor by broadcasting it into the first rotary axis alone, which means the decoder is fed positions for one axis and zeros for the rest: the model still runs, still emits plausible text, and is quietly wrong. `rotary_axis_count` reads the count from the decoder's own declared `position_ids` input rather than from a flag, so a package states it once by being exported. A symbolic leading dimension is refused: the number of rotary axes is fixed by the export, and guessing it would mis-shape every position the decoder reads. `build_decoder_state_initializer` now emits (sections, batch, prompt) prefill positions and (sections, batch, 1) body positions when the decoder declares them, and `build_decoder_step_update` takes the same `position_sections` so the +1 per step keeps the rank. Both keep their rank-2 behaviour untouched when the decoder declares rank-2 positions. Exercised end to end against the cached Foundry Local qwen3.5-0.8b package, whose text decoder declares position_ids as i64[3, batch, sequence]: three prompts reproduce Foundry Local's own output exactly, including a 16-token sentence byte for byte. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: justinchuby (cherry picked from commit c0e8873a698c9169b6434bd0432a7d0f3f525cfe) --- src/mobius/generation/__init__.py | 2 + src/mobius/generation/_policy_components.py | 62 +++++++++++- .../generation/_policy_components_test.py | 98 +++++++++++++++++++ .../onnx_genai/workflow_metadata.py | 10 ++ 4 files changed, 169 insertions(+), 3 deletions(-) diff --git a/src/mobius/generation/__init__.py b/src/mobius/generation/__init__.py index 69ca0840c..09edc410a 100644 --- a/src/mobius/generation/__init__.py +++ b/src/mobius/generation/__init__.py @@ -68,6 +68,7 @@ build_video_latent_permute, build_video_latent_unscale, build_zeros_like, + rotary_axis_count, ) __all__ = [ @@ -133,4 +134,5 @@ "build_video_latent_permute", "build_video_latent_unscale", "build_zeros_like", + "rotary_axis_count", ] diff --git a/src/mobius/generation/_policy_components.py b/src/mobius/generation/_policy_components.py index 6fff6c73b..c704bda3e 100644 --- a/src/mobius/generation/_policy_components.py +++ b/src/mobius/generation/_policy_components.py @@ -639,6 +639,30 @@ def build_empty_features(dtype: ir.DataType, feature_size: int) -> PolicyCompone 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, *, @@ -763,6 +787,7 @@ def build_decoder_state_initializer( 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])), @@ -772,12 +797,36 @@ def build_decoder_state_initializer( op.Unsqueeze(positions, op.Constant(value_ints=[0])), prompt_shape ) positions = op.Cast(positions, to=position_value.dtype) - positions.shape = position_value.shape + # (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) @@ -858,6 +907,7 @@ 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: @@ -901,13 +951,19 @@ def build_decoder_step_update( 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=["batch", 1], + shape=position_shape, ) next_position = op.Add(position, op.CastLike(op.Constant(value_int=1), position)) - next_position.shape = ir.Shape(["batch", 1]) + next_position.shape = ir.Shape(position_shape) builder.add_output(next_position, "next_position_ids") 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 index 731ffcd6e..3dae63518 100644 --- a/src/mobius/generation/_policy_components_test.py +++ b/src/mobius/generation/_policy_components_test.py @@ -6,6 +6,7 @@ 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 ( @@ -46,6 +47,7 @@ build_termination_batch_initializer, build_token_state_update, build_true_cfg, + rotary_axis_count, build_unpack_latents_2x2, build_video_conv_cache_initializer, build_video_decode_chunk, @@ -1374,3 +1376,99 @@ def test_true_cfg_is_identity_at_unit_guidance(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/onnx_genai/workflow_metadata.py b/src/mobius/integrations/onnx_genai/workflow_metadata.py index c5a530ff3..79d5d7ea8 100644 --- a/src/mobius/integrations/onnx_genai/workflow_metadata.py +++ b/src/mobius/integrations/onnx_genai/workflow_metadata.py @@ -66,6 +66,7 @@ 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, @@ -2407,6 +2408,8 @@ def _accumulated_contract(contract: dict[str, Any], symbol: str) -> dict[str, An 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(".", "_") @@ -4058,6 +4061,9 @@ def build_vlm_workflow_metadata( 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()) @@ -5962,6 +5968,9 @@ def _build_autoregressive_workflow_metadata( 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: @@ -5971,6 +5980,7 @@ def _build_autoregressive_workflow_metadata( 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 From 6b4511e90bb947beb21a02a6e3540864a86e8099 Mon Sep 17 00:00:00 2001 From: Justin Chu Date: Thu, 20 Aug 2026 06:20:27 +0000 Subject: [PATCH 123/151] Add full-duplex speech-to-speech workflow metadata producer PersonaPlex (nvidia/personaplex-7b-v1) and other Moshi-family models consume and produce audio simultaneously at a fixed 12.5 Hz frame rate. One invocation is exactly one codec frame: packed audio in, packed audio out, with the whole conversation carried in session-scoped state between calls. This adds the structural description of that workflow so a runtime can execute it without a Python driver: * Nine ONNX policy components for the delay ring buffer that Moshi-family models use to stagger text and acoustic streams -- frame assemble/commit, teacher selection, stream append/tail, user-stream merge, and the frame selectors. These are a direct transcription of upstream `LMGen.prepare_step_input` and its commit path, and they are validated against a NumPy reference in `_duplex_policy_test.py`. * `build_scalar_integer_add`, a rank-0 companion to `build_integer_add` for loop induction values and substep indices. * `build_full_duplex_workflow_metadata` / `write_full_duplex_workflow_metadata`, which emit a workflow whose frame body is a loop carrying every session cell: the temporal KV cache (bound to a runtime state-service group), the delay ring, the frame offset, the attention mask, and the codec prefixes. Two details are worth calling out because they were found by execution, not by reading: * The codec prefix cells use `release_boundary: invocation` while every language-model cell uses `session`. The stateless Mimi graphs replay an accumulated prefix instead of carrying convolution state, so upstream's `reset_streaming()` at a phase boundary has to be expressed as dropping that prefix. Keeping the prompt frames in the prefix drops output waveform correlation from 1.0 to 0.988. * `build_duplex_frame_assemble` uses `Or(existing, mask)` rather than `Where(mask, true, existing)` because onnxruntime has no bool `Where` kernel. Verified against the real 16.7 GB checkpoint at revision fdaf4090a61cb315c138a1faee287ffd6c716309: the emitted metadata for the exported 33 GB fp32 package passes the authoritative `validate_metadata` checker, and replaying the recorded 122-step upstream reference trace through the emitted policy graphs reproduces the ring state, the assembled input frames, and the emit schedule exactly. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: Justin Chu (cherry picked from commit b9afd13320ac7e2cec98cb3ece60636da6d58450) --- src/mobius/generation/__init__.py | 20 + src/mobius/generation/_duplex_policy_test.py | 239 +++++ src/mobius/generation/_policy_components.py | 337 ++++++- .../integrations/onnx_genai/__init__.py | 7 + .../duplex_workflow_metadata_test.py | 310 ++++++ .../onnx_genai/genai_config_import.py | 8 +- .../onnx_genai/workflow_metadata.py | 933 ++++++++++++++++++ 7 files changed, 1847 insertions(+), 7 deletions(-) create mode 100644 src/mobius/generation/_duplex_policy_test.py create mode 100644 src/mobius/integrations/onnx_genai/duplex_workflow_metadata_test.py diff --git a/src/mobius/generation/__init__.py b/src/mobius/generation/__init__.py index 09edc410a..a9c7951a3 100644 --- a/src/mobius/generation/__init__.py +++ b/src/mobius/generation/__init__.py @@ -20,6 +20,15 @@ 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, @@ -41,6 +50,7 @@ build_pack_latents_2x2, build_proposal_metrics, build_scalar_constant, + build_scalar_integer_add, build_schedule_constant, build_schedule_history_append, build_schedule_lookup, @@ -86,6 +96,15 @@ "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", @@ -107,6 +126,7 @@ "build_pack_latents_2x2", "build_proposal_metrics", "build_scalar_constant", + "build_scalar_integer_add", "build_schedule_constant", "build_schedule_history_append", "build_schedule_lookup", 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 index c704bda3e..cfd9ff27b 100644 --- a/src/mobius/generation/_policy_components.py +++ b/src/mobius/generation/_policy_components.py @@ -13,7 +13,7 @@ import json from collections.abc import Sequence from dataclasses import dataclass -from typing import Protocol +from typing import Any, Protocol import onnx_ir as ir from onnxscript import GraphBuilder @@ -2666,3 +2666,338 @@ def build_video_latent_unscale(scaling_factor: float) -> PolicyComponent: _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/integrations/onnx_genai/__init__.py b/src/mobius/integrations/onnx_genai/__init__.py index f389b92c1..589429600 100644 --- a/src/mobius/integrations/onnx_genai/__init__.py +++ b/src/mobius/integrations/onnx_genai/__init__.py @@ -15,6 +15,9 @@ 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. +* **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 @@ -64,6 +67,7 @@ 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, @@ -73,6 +77,7 @@ 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, @@ -87,6 +92,7 @@ "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", @@ -109,6 +115,7 @@ "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", 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/genai_config_import.py b/src/mobius/integrations/onnx_genai/genai_config_import.py index f89d20a06..8800ca0fb 100644 --- a/src/mobius/integrations/onnx_genai/genai_config_import.py +++ b/src/mobius/integrations/onnx_genai/genai_config_import.py @@ -132,14 +132,10 @@ def unbound_decoder_ports(package: ModelPackage) -> dict[str, list[str]]: decoder = package["decoder"] return { "inputs": sorted( - value.name - for value in decoder.graph.inputs - if value.name not in declared_inputs + 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 + value.name for value in decoder.graph.outputs if value.name not in declared_outputs ), } diff --git a/src/mobius/integrations/onnx_genai/workflow_metadata.py b/src/mobius/integrations/onnx_genai/workflow_metadata.py index 79d5d7ea8..38f1811f2 100644 --- a/src/mobius/integrations/onnx_genai/workflow_metadata.py +++ b/src/mobius/integrations/onnx_genai/workflow_metadata.py @@ -28,6 +28,15 @@ 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, @@ -42,6 +51,7 @@ build_pack_latents_2x2, build_proposal_metrics, build_scalar_constant, + build_scalar_integer_add, build_schedule_constant, build_schedule_history_append, build_schedule_lookup, @@ -7689,3 +7699,926 @@ def write_ctc_asr_workflow_metadata( 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": { + "ir_version": "1.0", + "onnx_opsets": {"ai.onnx": OPSET_VERSION}, + "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}": { + "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 From df4b6f58277ce045f2588ab469f78be9371307fa Mon Sep 17 00:00:00 2001 From: justinchuby Date: Thu, 20 Aug 2026 12:23:44 +0000 Subject: [PATCH 124/151] Reconcile the E2E metadata producers after the rebase onto main Integrating the isolated end-to-end branches surfaced four defects that only appear once the commits sit on the current main and next to each other. `qwen_image_test` imported `QwenImageVAEConfig` from `mobius._diffusers_configs`, a module main relocated to `mobius.integrations.diffusers._configs`, so every temporal-VAE test errored at collection. The video producer calls `build_euler_solver_step` directly rather than through `SOLVER_BUILDERS`, but the symbol never reached the import list while the overlapping import hunks from the video and image-edit branches were merged. `_CogVideoXDecoder3D` and `AutoencoderKLCogVideoXModel.conv_cache_entries` both computed `temporal_levels`; only the former consumes it. The dead copy in the cache-entry walk is removed, since the entry list is scaled by `scale`, not by the temporal level. The remainder is formatting the merged regions to the repository's ruff profile. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: justinchuby --- src/mobius/generation/_policy_components_test.py | 2 +- src/mobius/integrations/onnx_genai/workflow_metadata.py | 1 + src/mobius/models/cogvideox_vae.py | 9 ++------- src/mobius/models/qwen_image_test.py | 2 +- src/mobius/tasks/_video_vae.py | 6 ++---- 5 files changed, 7 insertions(+), 13 deletions(-) diff --git a/src/mobius/generation/_policy_components_test.py b/src/mobius/generation/_policy_components_test.py index 3dae63518..feb29153b 100644 --- a/src/mobius/generation/_policy_components_test.py +++ b/src/mobius/generation/_policy_components_test.py @@ -47,13 +47,13 @@ build_termination_batch_initializer, build_token_state_update, build_true_cfg, - rotary_axis_count, 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 diff --git a/src/mobius/integrations/onnx_genai/workflow_metadata.py b/src/mobius/integrations/onnx_genai/workflow_metadata.py index 38f1811f2..19558cb16 100644 --- a/src/mobius/integrations/onnx_genai/workflow_metadata.py +++ b/src/mobius/integrations/onnx_genai/workflow_metadata.py @@ -40,6 +40,7 @@ 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, diff --git a/src/mobius/models/cogvideox_vae.py b/src/mobius/models/cogvideox_vae.py index e718feab1..d659dadfe 100644 --- a/src/mobius/models/cogvideox_vae.py +++ b/src/mobius/models/cogvideox_vae.py @@ -115,9 +115,7 @@ def _nearest_indices(op: OpBuilder, out_len: ir.Value, in_len: ir.Value) -> ir.V """ 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]) - ) + 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) @@ -140,9 +138,7 @@ def _causal_nearest_indices(op: OpBuilder, out_len: ir.Value, in_len: ir.Value) 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]) - ) + 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 @@ -593,7 +589,6 @@ def conv_cache_spec(self) -> list[ConvCacheEntry]: """ config = self.config reversed_channels = list(reversed(config.block_out_channels)) - temporal_levels = int(math.log2(config.temporal_compression_ratio)) frames = 2 # kernel_t - 1 for every cached convolution entries = [ConvCacheEntry("conv_in", config.latent_channels, frames, 1)] diff --git a/src/mobius/models/qwen_image_test.py b/src/mobius/models/qwen_image_test.py index 1c9922862..642c4b147 100644 --- a/src/mobius/models/qwen_image_test.py +++ b/src/mobius/models/qwen_image_test.py @@ -372,7 +372,7 @@ def _temporal_vae(dtype: ir.DataType = ir.DataType.FLOAT): 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._diffusers_configs import QwenImageVAEConfig + from mobius.integrations.diffusers._configs import QwenImageVAEConfig from mobius.models.qwen_image_vae import AutoencoderKLQwenImageModel from mobius.tasks import QwenImageEditVAETask diff --git a/src/mobius/tasks/_video_vae.py b/src/mobius/tasks/_video_vae.py index 3519854fc..ccf440595 100644 --- a/src/mobius/tasks/_video_vae.py +++ b/src/mobius/tasks/_video_vae.py @@ -71,14 +71,12 @@ def build( shape=["batch", entry.channels, "cache_frames", height, width], ) - sample, updated_cache = module( - op, latent_sample=latent_sample, conv_cache=conv_cache - ) + 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) + spatial = 2 ** (len(config.block_out_channels) - 1) sample.shape = ir.Shape( [ "batch", From d5353d3c123402c05a70ae3673b155cead0539cf Mon Sep 17 00:00:00 2001 From: justinchuby Date: Thu, 20 Aug 2026 13:26:58 +0000 Subject: [PATCH 125/151] Refresh the artifacts the integrated producers now emit CI caught three things the per-branch test runs could not, because each one only appears once the commits are combined. `qwen-image-edit-2509.yaml` records an official-weight validation run that the case schema never learned about, so every schema check rejected the file. The block is real evidence for a case whose golden is a reduced-config reference, so the schema now describes it rather than the case dropping it. The validation package generator still selected adapters by row identity, a surface removed with the vestigial `slot_ids` fields. Regenerating the checked-in packages also picks up the diffusion policy renames and the new guided-diffusion and video workflows, which the fixtures predate. `conv_dim`, `conv_kernel` and `conv_stride` describe a single convolutional feature encoder, and the CTC work made their disagreement an error. M-CTC-T declares a one-layer subsampler through kernel and stride alone, so inheriting wav2vec2's seven-layer `conv_dim` default described a stack the checkpoint does not have. The widths are now sized to the declared depth. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: justinchuby --- src/mobius/_configs/_base.py | 26 +- testdata/cases/schema.json | 49 + .../adapter/inference_metadata.yaml | 2 - .../adapter/model.onnx.data | 0 .../codec/inference_metadata.yaml | 6 + .../decoder/inference_metadata.yaml | 364 ++++++-- .../policies/decoder_state_initializer.onnx | Bin 9244 -> 8581 bytes .../decoder/policies/decoder_step_update.onnx | Bin 3028 -> 2037 bytes .../diffusion/inference_metadata.yaml | 274 ++++-- .../policies/initial_state_scale.onnx | Bin 0 -> 363 bytes ...odel_input.onnx => model_input_scale.onnx} | Bin .../diffusion/policies/tensor_scale.onnx | Bin 0 -> 615 bytes .../diffusion_guided/denoiser/model.onnx | Bin 0 -> 1885 bytes .../diffusion_guided/denoiser/model.onnx.data | 0 .../diffusion_guided/inference_metadata.yaml | 777 ++++++++++++++++ .../policies/continue_predicate.onnx | Bin 0 -> 910 bytes .../policies/decoder_input_scale.onnx | Bin 0 -> 363 bytes .../policies/diffusion_schedule.onnx | Bin 0 -> 387 bytes .../policies/diffusion_timesteps.onnx | Bin 0 -> 383 bytes .../policies/guidance_combine.onnx | Bin 0 -> 1635 bytes .../policies/history_initializer.onnx | Bin 0 -> 777 bytes .../policies/latent_noise.onnx | Bin 0 -> 18187 bytes .../policies/latent_row_shape.onnx | Bin 0 -> 382 bytes .../policies/schedule_lookup.onnx | Bin 0 -> 607 bytes .../policies/solver_step.onnx | Bin 0 -> 13995 bytes .../policies/tensor_scale.onnx | Bin 0 -> 615 bytes .../diffusion_guided/text_encoder/model.onnx | Bin 0 -> 1331 bytes .../text_encoder/model.onnx.data | 0 .../diffusion_guided/vae_decoder/model.onnx | Bin 0 -> 957 bytes .../vae_decoder/model.onnx.data | 0 .../masked/inference_metadata.yaml | 69 ++ .../speculative/inference_metadata.yaml | 305 ++++++- .../tts/inference_metadata.yaml | 338 ++++++- .../video/inference_metadata.yaml | 850 ++++++++++++++++++ .../video/policies/continue_predicate.onnx | Bin 0 -> 910 bytes .../video/policies/diffusion_schedule.onnx | Bin 0 -> 387 bytes .../video/policies/diffusion_timesteps.onnx | Bin 0 -> 383 bytes .../video/policies/model_input.onnx | Bin 0 -> 720 bytes .../policies/schedule_history_append.onnx | Bin 0 -> 829 bytes .../video/policies/schedule_lookup.onnx | Bin 0 -> 607 bytes .../video/policies/solver_step.onnx | Bin 0 -> 6366 bytes .../video/policies/video_conv_cache_init.onnx | Bin 0 -> 2801 bytes .../video/policies/video_decode_chunk.onnx | Bin 0 -> 3115 bytes .../video/policies/video_decode_chunks.onnx | Bin 0 -> 1130 bytes .../video/policies/video_latent_init.onnx | Bin 0 -> 1580 bytes .../video/policies/video_latent_permute.onnx | Bin 0 -> 488 bytes .../video/policies/video_latent_unscale.onnx | Bin 0 -> 824 bytes .../video/transformer/model.onnx | Bin 0 -> 5197 bytes .../video/transformer/model.onnx.data | 0 .../video/vae_decoder/model.onnx | Bin 0 -> 14988 bytes .../video/vae_decoder/model.onnx.data | 0 .../vlm/inference_metadata.yaml | 370 ++++++-- .../policies/decoder_state_initializer.onnx | Bin 10460 -> 9620 bytes .../vlm/policies/decoder_step_update.onnx | Bin 3028 -> 2037 bytes ...generate_onnx_genai_validation_packages.py | 2 - 55 files changed, 3164 insertions(+), 268 deletions(-) create mode 100644 tests/fixtures/onnx_genai_workflows/adapter/model.onnx.data create mode 100644 tests/fixtures/onnx_genai_workflows/diffusion/policies/initial_state_scale.onnx rename tests/fixtures/onnx_genai_workflows/diffusion/policies/{euler_model_input.onnx => model_input_scale.onnx} (100%) create mode 100644 tests/fixtures/onnx_genai_workflows/diffusion/policies/tensor_scale.onnx create mode 100644 tests/fixtures/onnx_genai_workflows/diffusion_guided/denoiser/model.onnx create mode 100644 tests/fixtures/onnx_genai_workflows/diffusion_guided/denoiser/model.onnx.data create mode 100644 tests/fixtures/onnx_genai_workflows/diffusion_guided/inference_metadata.yaml create mode 100644 tests/fixtures/onnx_genai_workflows/diffusion_guided/policies/continue_predicate.onnx create mode 100644 tests/fixtures/onnx_genai_workflows/diffusion_guided/policies/decoder_input_scale.onnx create mode 100644 tests/fixtures/onnx_genai_workflows/diffusion_guided/policies/diffusion_schedule.onnx create mode 100644 tests/fixtures/onnx_genai_workflows/diffusion_guided/policies/diffusion_timesteps.onnx create mode 100644 tests/fixtures/onnx_genai_workflows/diffusion_guided/policies/guidance_combine.onnx create mode 100644 tests/fixtures/onnx_genai_workflows/diffusion_guided/policies/history_initializer.onnx create mode 100644 tests/fixtures/onnx_genai_workflows/diffusion_guided/policies/latent_noise.onnx create mode 100644 tests/fixtures/onnx_genai_workflows/diffusion_guided/policies/latent_row_shape.onnx create mode 100644 tests/fixtures/onnx_genai_workflows/diffusion_guided/policies/schedule_lookup.onnx create mode 100644 tests/fixtures/onnx_genai_workflows/diffusion_guided/policies/solver_step.onnx create mode 100644 tests/fixtures/onnx_genai_workflows/diffusion_guided/policies/tensor_scale.onnx create mode 100644 tests/fixtures/onnx_genai_workflows/diffusion_guided/text_encoder/model.onnx create mode 100644 tests/fixtures/onnx_genai_workflows/diffusion_guided/text_encoder/model.onnx.data create mode 100644 tests/fixtures/onnx_genai_workflows/diffusion_guided/vae_decoder/model.onnx create mode 100644 tests/fixtures/onnx_genai_workflows/diffusion_guided/vae_decoder/model.onnx.data create mode 100644 tests/fixtures/onnx_genai_workflows/video/inference_metadata.yaml create mode 100644 tests/fixtures/onnx_genai_workflows/video/policies/continue_predicate.onnx create mode 100644 tests/fixtures/onnx_genai_workflows/video/policies/diffusion_schedule.onnx create mode 100644 tests/fixtures/onnx_genai_workflows/video/policies/diffusion_timesteps.onnx create mode 100644 tests/fixtures/onnx_genai_workflows/video/policies/model_input.onnx create mode 100644 tests/fixtures/onnx_genai_workflows/video/policies/schedule_history_append.onnx create mode 100644 tests/fixtures/onnx_genai_workflows/video/policies/schedule_lookup.onnx create mode 100644 tests/fixtures/onnx_genai_workflows/video/policies/solver_step.onnx create mode 100644 tests/fixtures/onnx_genai_workflows/video/policies/video_conv_cache_init.onnx create mode 100644 tests/fixtures/onnx_genai_workflows/video/policies/video_decode_chunk.onnx create mode 100644 tests/fixtures/onnx_genai_workflows/video/policies/video_decode_chunks.onnx create mode 100644 tests/fixtures/onnx_genai_workflows/video/policies/video_latent_init.onnx create mode 100644 tests/fixtures/onnx_genai_workflows/video/policies/video_latent_permute.onnx create mode 100644 tests/fixtures/onnx_genai_workflows/video/policies/video_latent_unscale.onnx create mode 100644 tests/fixtures/onnx_genai_workflows/video/transformer/model.onnx create mode 100644 tests/fixtures/onnx_genai_workflows/video/transformer/model.onnx.data create mode 100644 tests/fixtures/onnx_genai_workflows/video/vae_decoder/model.onnx create mode 100644 tests/fixtures/onnx_genai_workflows/video/vae_decoder/model.onnx.data diff --git a/src/mobius/_configs/_base.py b/src/mobius/_configs/_base.py index 4ee199de1..91b96448f 100644 --- a/src/mobius/_configs/_base.py +++ b/src/mobius/_configs/_base.py @@ -2972,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. @@ -3059,7 +3083,7 @@ def from_transformers(cls, config, parent_config=None) -> MMSConfig: 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=tuple(getattr(config, "conv_dim", None) or defaults.conv_dim), + 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)), 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/fixtures/onnx_genai_workflows/adapter/inference_metadata.yaml b/tests/fixtures/onnx_genai_workflows/adapter/inference_metadata.yaml index bf601d63a..8916e4c6e 100644 --- a/tests/fixtures/onnx_genai_workflows/adapter/inference_metadata.yaml +++ b/tests/fixtures/onnx_genai_workflows/adapter/inference_metadata.yaml @@ -177,8 +177,6 @@ adapters: scale: lora.projection.scale discovery_fallback: disabled selection: - slot_ids: request.slot_ids - request_epochs: request.request_epochs segments: request.adapter_segments adapter_counts: request.adapter_counts scales: request.adapter_scales diff --git a/tests/fixtures/onnx_genai_workflows/adapter/model.onnx.data b/tests/fixtures/onnx_genai_workflows/adapter/model.onnx.data new file mode 100644 index 000000000..e69de29bb diff --git a/tests/fixtures/onnx_genai_workflows/codec/inference_metadata.yaml b/tests/fixtures/onnx_genai_workflows/codec/inference_metadata.yaml index 032e47071..aa6c5e570 100644 --- a/tests/fixtures/onnx_genai_workflows/codec/inference_metadata.yaml +++ b/tests/fixtures/onnx_genai_workflows/codec/inference_metadata.yaml @@ -18,6 +18,9 @@ pipeline: - batch - 1 - audio_samples + batch_layout: + kind: request_aligned + axis: 0 role: kind: runtime version: '1.0' @@ -34,6 +37,9 @@ pipeline: - batch - 1 - audio_samples + batch_layout: + kind: request_aligned + axis: 0 role: audio stage: post_adapter components: diff --git a/tests/fixtures/onnx_genai_workflows/decoder/inference_metadata.yaml b/tests/fixtures/onnx_genai_workflows/decoder/inference_metadata.yaml index 7712409f3..b62107279 100644 --- a/tests/fixtures/onnx_genai_workflows/decoder/inference_metadata.yaml +++ b/tests/fixtures/onnx_genai_workflows/decoder/inference_metadata.yaml @@ -11,7 +11,6 @@ pipeline: - nested_control_flow - typed_emit - emit_valid_length - - emit_row_identity - loop_induction_values - serving_service_contract - bounded_state_recurrence @@ -23,6 +22,9 @@ pipeline: shape: - batch - sequence + batch_layout: + kind: request_aligned + axis: 0 role: kind: runtime version: '1.0' @@ -61,6 +63,21 @@ pipeline: 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: @@ -85,6 +102,9 @@ pipeline: rank: 1 shape: - batch + batch_layout: + kind: request_aligned + axis: 0 role: kind: opaque source: @@ -99,6 +119,9 @@ pipeline: shape: - batch - num_eos + batch_layout: + kind: request_aligned + axis: 0 role: kind: opaque source: @@ -112,6 +135,9 @@ pipeline: rank: 1 shape: - batch + batch_layout: + kind: request_aligned + axis: 0 role: kind: opaque source: @@ -125,6 +151,9 @@ pipeline: rank: 1 shape: - batch + batch_layout: + kind: request_aligned + axis: 0 role: kind: opaque source: @@ -138,6 +167,9 @@ pipeline: rank: 1 shape: - batch + batch_layout: + kind: request_aligned + axis: 0 role: kind: runtime version: '1.0' @@ -152,6 +184,9 @@ pipeline: rank: 1 shape: - batch + batch_layout: + kind: request_aligned + axis: 0 role: kind: runtime version: '1.0' @@ -166,6 +201,9 @@ pipeline: rank: 1 shape: - batch + batch_layout: + kind: request_aligned + axis: 0 role: kind: runtime version: '1.0' @@ -180,6 +218,9 @@ pipeline: rank: 1 shape: - batch + batch_layout: + kind: request_aligned + axis: 0 role: kind: runtime version: '1.0' @@ -194,6 +235,9 @@ pipeline: rank: 1 shape: - batch + batch_layout: + kind: request_aligned + axis: 0 role: kind: runtime version: '1.0' @@ -208,6 +252,9 @@ pipeline: rank: 1 shape: - batch + batch_layout: + kind: request_aligned + axis: 0 role: kind: opaque source: @@ -221,6 +268,9 @@ pipeline: rank: 1 shape: - batch + batch_layout: + kind: request_aligned + axis: 0 role: kind: opaque source: @@ -233,30 +283,24 @@ pipeline: rank: 1 shape: - batch + batch_layout: + kind: request_aligned + axis: 0 role: kind: opaque source: kind: literal required: false default: false - package.slot_ids: - contract: - dtype: int64 - rank: 1 - shape: - - batch - role: - kind: opaque - source: - kind: application - name: serving.slot_ids - required: true package.cache_lengths: contract: dtype: int64 rank: 1 shape: - batch + batch_layout: + kind: request_aligned + axis: 0 role: kind: opaque source: @@ -269,6 +313,9 @@ pipeline: rank: 1 shape: - batch + batch_layout: + kind: request_aligned + axis: 0 role: kind: opaque source: @@ -283,6 +330,9 @@ pipeline: shape: - batch - generated_sequence + batch_layout: + kind: request_aligned + axis: 0 role: tokens stage: pre_adapter components: @@ -302,57 +352,90 @@ pipeline: 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' @@ -384,17 +467,26 @@ pipeline: 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 @@ -405,22 +497,34 @@ pipeline: 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 @@ -454,22 +558,34 @@ pipeline: 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 @@ -477,6 +593,9 @@ pipeline: shape: - batch - 1 + batch_layout: + kind: request_aligned + axis: 0 contract: id: onnx-genai.state-update version: '2' @@ -502,6 +621,9 @@ pipeline: - batch - sequence - vocabulary + batch_layout: + kind: request_aligned + axis: 0 outputs: last_logits: dtype: float32 @@ -509,6 +631,9 @@ pipeline: shape: - batch - vocabulary + batch_layout: + kind: request_aligned + axis: 0 decoder_state_initializer: implementation: kind: onnx @@ -521,65 +646,82 @@ pipeline: shape: - batch - prompt_sequence + batch_layout: + kind: request_aligned + axis: 0 prompt_lengths: dtype: int64 rank: 1 shape: - batch - max_iterations: - dtype: int64 - rank: 1 - shape: - - 1 + batch_layout: + kind: request_aligned + axis: 0 outputs: attention_mask: dtype: int64 rank: 2 shape: - batch - - capacity + - total_sequence + batch_layout: + kind: request_aligned + axis: 0 position_ids: dtype: int64 rank: 2 shape: - batch - - sequence + - prompt_sequence + batch_layout: + kind: request_aligned + axis: 0 body_attention_mask: dtype: int64 rank: 2 shape: - batch - - capacity + - 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 - cache_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 - - capacity + - past_sequence - 8 + batch_layout: + kind: request_aligned + axis: 0 decoder_step_update: implementation: kind: onnx @@ -592,30 +734,37 @@ pipeline: shape: - batch - context - logical_length: - dtype: int64 - rank: 1 - shape: - - batch + 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 + - 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 @@ -627,27 +776,42 @@ pipeline: 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 @@ -660,16 +824,25 @@ pipeline: 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 @@ -680,6 +853,9 @@ pipeline: rank: 1 shape: - batch + batch_layout: + kind: request_aligned + axis: 0 outputs: row_eos_ids: dtype: int64 @@ -687,16 +863,25 @@ pipeline: 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 @@ -708,6 +893,9 @@ pipeline: rank: 1 shape: - batch + batch_layout: + kind: request_aligned + axis: 0 outputs: slot: dtype: int64 @@ -715,6 +903,9 @@ pipeline: shape: - batch - 1 + batch_layout: + kind: request_aligned + axis: 0 generated_length_update: implementation: kind: onnx @@ -726,27 +917,42 @@ pipeline: 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: @@ -755,6 +961,9 @@ pipeline: shape: - batch - 1 + batch_layout: + kind: request_aligned + axis: 0 scope: invocation initializer: initializer.token_slot recurrence: @@ -766,6 +975,9 @@ pipeline: shape: - batch - 128 + batch_layout: + kind: request_aligned + axis: 0 scope: invocation initializer: decoder.setup.last_logits recurrence: @@ -776,6 +988,9 @@ pipeline: rank: 1 shape: - batch + batch_layout: + kind: request_aligned + axis: 0 class: semantic scope: invocation initializer: initializer.generated_lengths @@ -787,6 +1002,9 @@ pipeline: rank: 1 shape: - batch + batch_layout: + kind: request_aligned + axis: 0 class: semantic scope: invocation initializer: package.active @@ -798,6 +1016,9 @@ pipeline: rank: 1 shape: - batch + batch_layout: + kind: request_aligned + axis: 0 class: semantic scope: invocation initializer: package.not_done @@ -809,31 +1030,26 @@ pipeline: rank: 1 shape: - batch + batch_layout: + kind: request_aligned + axis: 0 class: semantic scope: invocation initializer: package.zero_batch recurrence: kind: invariant - slot_ids: - contract: - dtype: int64 - rank: 1 - shape: - - batch - class: semantic - scope: invocation - initializer: package.slot_ids - 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 + initializer: package.cache_lengths recurrence: kind: invariant rng_counter: @@ -842,6 +1058,9 @@ pipeline: rank: 1 shape: - batch + batch_layout: + kind: request_aligned + axis: 0 scope: invocation class: semantic initializer: request.rng_counter @@ -854,10 +1073,16 @@ pipeline: shape: - batch - context + batch_layout: + kind: request_aligned + axis: 0 scope: invocation - initializer: initializer.attention_mask + initializer: initializer.body_attention_mask recurrence: - kind: invariant + kind: growing + axis: 1 + increment: package.one_step + max: package.max_context position_ids: contract: dtype: int64 @@ -865,6 +1090,9 @@ pipeline: shape: - batch - 1 + batch_layout: + kind: request_aligned + axis: 0 scope: invocation initializer: initializer.body_position_ids recurrence: @@ -878,28 +1106,32 @@ pipeline: - 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 - slot_ids: slot_ids - kv_service: - paging: none - allocation: runtime - compaction: true + state_service: groups: decoder_cache: + kind: full_attention sequence_axis: 2 layout: bnsh - logical_lengths: cache_lengths - storage: shared_buffer + aliasing: forbidden + reuse: + prefix_reusable: true + evictable_prefix: false ports: model: cache_0: @@ -913,13 +1145,11 @@ pipeline: inputs: prompt_tokens: request.input_ids prompt_lengths: request.prompt_lengths - max_iterations: request.max_iterations outputs: attention_mask: initializer.attention_mask body_attention_mask: initializer.body_attention_mask token_slot: initializer.token_slot generated_lengths: initializer.generated_lengths - cache_lengths: initializer.cache_lengths position_ids: initializer.position_ids body_position_ids: initializer.body_position_ids past_key_values.0.key: initializer.past_key_values.0.key @@ -1037,22 +1267,12 @@ pipeline: mode: append valid_length: token.emitted_length when: active - row_ids: slot_ids - - kind: invoke - component: decoder_step_update - inputs: - attention_mask: attention_mask - logical_length: cache_lengths - position_ids: position_ids - outputs: - next_attention_mask: decoder_step.body_attention_mask - next_position_ids: decoder_step.body_position_ids - kind: invoke component: model inputs: input_ids: token.body past_key_values.0.key: cache_0 - attention_mask: decoder_step.body_attention_mask + attention_mask: attention_mask position_ids: position_ids outputs: logits: decoder.body.logits @@ -1063,6 +1283,14 @@ pipeline: 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: @@ -1080,8 +1308,6 @@ pipeline: next: cache_lengths.next - cell: accepted_len next: accepted_len.next - - cell: slot_ids - next: slot_ids - cell: rng_counter next: sample.next_counter - cell: attention_mask diff --git a/tests/fixtures/onnx_genai_workflows/decoder/policies/decoder_state_initializer.onnx b/tests/fixtures/onnx_genai_workflows/decoder/policies/decoder_state_initializer.onnx index 6e93156a3863f12fdf03a7fe95b4529558b3278e..9d9e76ce3f7ad77bd2e6743247d6c2e56e4aec91 100644 GIT binary patch literal 8581 zcmd5?O>Y}T7|t3e-b`&LnUXf$qLNrqsf-kPcl?Qj)D|Hm3keB?kc!Y~ZEq7N&aUxp zqNIpR51hC_3kU(C5^&?p1tGy-;t#+(vtKi_vC3r~a@xu4%shENo_XGPqhG-*$KGIk z;=lIyee{`zD$|*FJPl~z9XYO#S5N7C2lmvV4Wk@BxZVgxwPil)SJ0|!A3Ofk9y)lb zzl&~7kM{RG*FE!xv+*?88%}KBrw3zaX3vHP=eY3UNBwPNC~vVns6X_kj?dme&*lR$ zB_L)5q>j)Fjyu}Yo@s@887=#PJqxy$U~~x(EeWEfKrBcQ8Q{DiKGAUbl)mA)U=25* z#8^w7xfe-IkAG3$KDh zo=t6cMB7Fsdg931wUuy~7pyPkAnT|>))9hi9U#X+mTPi;r=e;@WE02T4-R}O&^j#8 z-kLdf;LK>(sK$@nV%=IP8s}v)U^nsn3{=6HjDSBlaApqe8RhVSgY4Co`6w@lPc;_u zF*{JRVXP+u4!~?&hF?`!W(Jly4_2N;`!i3?#6V_4CKAb#JXhbf-F=6emQe|xLX*Nx z!H0P{^#=_|p2lZO{8km2#U#c?LT5C@Mi#(u-$2&NI8ushu#}C6&4jx_C}pBR9{7fz}qV2D2cvOc&(cukgh-=vM2%u$xRgfB#0uD z-Sa7#bXnxXaA<=T=^54NF|=H-wi1r(S5VPD8~eFltAFx1OR(S;LB4m;UcZirt7Th5V z%%tDMOtP97@p_y36EB#TKdtBIPn14Tz2%laiUPCR8Mn6u0$rHkSEI3$fhE=EM-GSU64H`tp$)E4 zB#WtaBxx@oynYu**Tft9jtkxAz!}k~7c}RFC7u7p4Xg4>rbHq0OBCoW=T@SqxH8uT z5wDBI@tuOJnQ2uu#a5M!OV80>e~D2j)3AJrk)AZ$I~ZrZ&xm%G5tZnc6h* z=9CrfN6tCrb(_D}*n`2Gd}>JPH5b3Ki?ZyTHO)PAaopiM z=(euGADx)Kpcl{@uJ9<0jf4uiez+xER1J)f4seSwqe(N+n-ylRi=rf*C>#nCb5~f> zR)+NXfk<~0Yqz?|< z!mC2?9c~KkO<}w15;2*LWItuXOeQ$6xWg>$bW7<;g0cdOmS7$D0D#{RWS0r$aO8$1 zm{OOvlBG-Bf!rH$mys=81<$gD+=6j0ikjVISLSwpU_`PeR1t~a3x~Bi9isRhwl_aC z?#+l3@}bajh$q!USXDqGZQCl*yU>m?I{e@H{{k9-Q$2`M&B z^#HLF4`iJOnmJy|?8E~r(u(tslU<`=5uKwcPKCQNK25bR+*avS@v-&_+6n*BVK35y z(23rgdXw?+oDMxVnAyVs7e4v)v33vbKKD&J9^o6&Tj3^q_Q~0JGPY;ucg@GzZL~di T59i@SXBhB*bhrM;*==9yS$oExzd4uhoJoBF&mE6;hbO@+ zpS^)T)lg$R86A&9I~*N)et_ps?RWQ`v1hmR#puEx4Pa!<^3hHM&HK)=7mOX(!!tWK z(Dm`*-uB4%&jNQc9EaQPkrM>={?MB^6L0 z5HkkSMCcjMAFOMSwMugl%?6<}3D;*}^aLOh2_jJ-RwRfFa8VE+YIyO~zBBTH4L`K0 zzMNk3U8F`m`JuUr(5Z8D;@Lw#3~mB(wW3gctD)4Y2$BHe1u(adPCW0vXIuJGa>>D3 zMlBu}1?(pciSI|qWz5;;z;5Bjsr{ZG#L(LMN;cr&+QxkPOHuG1flo*%G8#>cPVLvu z#*RO*J9;C&;>bG2Tr?~S*4MI;byXwl3P!d9kP{;-)H#3D&{7=8jy!)a+z%u}>oP-o zed0NxH?e#AQgX=|>lw9pToh#x-Ng6jpbDGG0`U9$-o&%}`eJm!LH3PVJ}L?#BkO`Owg>C4GAhi0}4_tY;HPnF={`l)?uFtANiUts@7piJW#7!*a< zHyVo&NieldqOYWb=>0@+gYshS&l*bEV9_H5WfHt31{0I&t7!y_DWfF>wcmwgU?K$) z^HMYNROLgUFJWq%<9q*O+D+QR;p7Uon~ogA#rG+0II zj$e7EN(nOC(~Pjg^42?!zvtOzS8qgD5uaUSmJiGF*j|phop5n9p zywCPa@Yx?AH%PIg&(c}!ouKoXIqHa60(B30 z;!4?(c_b#w_HTJmFA*pq_yuX|RK%QK(-@a>m>Dgk}Tj z8foDg-)QP+mRAzZrMk9K+o{jQ{~l}AqPAU0wj`G^h?p@-(*UY9_SHm2`cEX0L6#88 ze90R7tSyn$3{ZU(r-9&nPo}dZY9~r$x~MDwsFf1{)G9#$e5Yt9YGq#15|k5d%gTwii&jn`0y5!7 z+jwz23Wi(-3deG%B zOz0^o-^z8Htf#EoWGUSyE8lIh%IG$!#93oOG^lbUQZH~r{}J6*%3GGw6tk#& zn^)Vhn6@Bi#FxBUvmj?g<5GHLe~ZH2YU8yrs~Zo!bDJymV7s*qgDDZn>f|HPR;l}Y^K^gZ4fcrKi6PHZ?W3)wjsWIyL5}hycem8@0KAMEFl#=Yb@!UuoW}^&;gKq; z)V1B0bCvT|v;lY5(ghI0b=QG2Td13NUqVu<#}_@EfZHYLqAL8tIq3^tMEUfHFQHX{ z)W{$)&bZ<>M#Kd!`3OvsoTxD@1>+rXwIhR?_G44ur za2?ITlnRJepp%@NSB)nk% z;>VaW0~>NXQl?rgM9p^kx+_vxH9P4IPEAy1nO)&q6ns4i-*p8a6uf|<28)7~1Fsxm z+67Ff&cH2YD)a=CsWWVHWGqRmj(yrZcs-ja1A>lE_T5_a!Z8MON%Tc*-DyirMwyxsfqmwJW?{1M=I~6LI$-`GV(k^$&!2)ql24ZEY+vJ&XWC2 z{$da<$%2e{#pWy-sWR1)rium$i7{KX`apXDZAE{7vmJ|d=(}!@M@K{V+;&HPIC0z% zS3dsafp!bs_}?e(;Q(*MPeqGtJ11wuqoFf7ziB?uo=2Ng*YNJ$^W2dC{mpGKsg?Er E00-+mbN~PV diff --git a/tests/fixtures/onnx_genai_workflows/decoder/policies/decoder_step_update.onnx b/tests/fixtures/onnx_genai_workflows/decoder/policies/decoder_step_update.onnx index d61b374c0aba21bd49e09839e485c098162ccaa4..b2a4c7288de43c04c5ae83a2119d81bd95d87b9b 100644 GIT binary patch literal 2037 zcmcIlL2uJA6pp)g&EDE=bNE1fDntYBIciqR=~*3Mjd7=fZcwIPfSiytkQGDCx<%OFbj*ZLZHA5R z0qpq~=M9hZsjvgrmkoPFg&+f$1vIb+S2*+GW48*{ovg415w_eLL33nl&>H)Jp)DbTnbau_5_X$hYnVBQ2$lc&v}oQtFrm(!%*`$wXQYF^m>2b2gPE>9eFT zj+cZH`YSEKUy_b$-*r@b!BIO(=CPx+>C*!B9VJ}UhJx=2SNA2?5xND$jCSr`awq*u z{-)q90!mzQ2bc5}FD6)U>Nt_@U)I*tI3?n6*z)^W2PyddILz$9NAOs@%#vw}>}qb*s;+&E8X3 z=4Ryv!}|*z6$6;6Y!=kcLo1<-X5wEZ(c|W;GQ<6q-Q#w#0miIwEUQ-H-)>^lHeQFU z0%Q!Mc)X?)$FI>7*pD98#yR8ai3_FM@V$|1Uy+cIfZ7sgK7GDMr*QD!OyWBD!7M8( h(x8*6J924ob<(^>kD>a{8m!-Ew$!iZ>ns>MmA^m2fFb|@ literal 3028 zcmc(hK~EDw6vw-T((R*#GKkboB(M@S#5CO&3iTuyFO4USF~OM0ba!aCl-;E}C7|)b zM=*HMgW}mw;Wu-3y1R5*(9=p!^Jd<@`St(i%^P(V%^bNcd*DC)y$qiflt01GUB_p{ zVYs5r$79E}D6Z&gG_g^H+zIIoD7GC|N?j`HW&v^zIikKtOo~#?by)EZ+hx~rPJOd) zd#r5si0|W$P5Y#8cFs`x-TUS$Xk%VM4D6e(NBzJ9ekw3W7^4oeT?Zc#{}AO)@aqoo zC^oeGa3JsuJtM3QfGO%&r8KA7$U}zzn?5Tk2{K+NT#bkczzh_`hJ`C=mQNoKDE&Zj zRhx^30=KGXLpv$y#`0Mk&u8tXd?u;Yr2^6+gHeg3rCN~aSHx*kys71f1EJieo)K13 zx_uMr40o2&SvbLO96wCxmR7tn1lX3Ii)WI8`@!p>=eBK=^l*}ACpCGLN z2f{B3NEBPLSr5|s;&{OMx3&CmAk^E|Gr~$rxv_HC883&O+bDs68-`YljUjD`P?4bA4O1G~4n8b0n)#<;;;WWUeNk=CAx5jv@)(h{C9% z1(DiO4{dhWvck01^eJv6lvK26QIo$v`q*dG!voJEjBaSEQq1m+^xCGH3V#!8MpIS7 zwf!d(r&fcKnniaYi!v?3%nsyGmbY-m&+`7-_+^Xcp}@U!vG@K1Sva%^3*4a)J3v}h zS_WKK8|W`vw;sc!=GO-AmPR^21p%mh|LI^`X~mFM?n(~F^l^}L_BcTPDS956g zO72uwX4;9FhV#*(Zp8D;W5qeP5PiU_AXMpGc?4_2cXPQ-9sZyVD!J^rJ=;9P!P)5( zlcDs-Pv^=mtpC#!+ZMVPdJTh=$>7xP*`$B=%s5vb!s?AN=)nOsnRqMjO+%_!`U4Jr BuulL0 diff --git a/tests/fixtures/onnx_genai_workflows/diffusion/inference_metadata.yaml b/tests/fixtures/onnx_genai_workflows/diffusion/inference_metadata.yaml index 8a369caf6..9bdb0a5a0 100644 --- a/tests/fixtures/onnx_genai_workflows/diffusion/inference_metadata.yaml +++ b/tests/fixtures/onnx_genai_workflows/diffusion/inference_metadata.yaml @@ -12,21 +12,6 @@ pipeline: - loop_induction_values - typed_emit inputs: - request.latent: - contract: - dtype: float32 - rank: 4 - shape: - - batch - - 4 - - height - - width - role: - kind: opaque - source: - kind: application - name: latent - required: true request.max_iterations: contract: dtype: int64 @@ -47,12 +32,34 @@ pipeline: 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 @@ -60,6 +67,9 @@ pipeline: shape: - batch - prompt_sequence + batch_layout: + kind: request_aligned + axis: 0 role: kind: runtime version: '1.0' @@ -67,6 +77,7 @@ pipeline: source: kind: request required: true + externally_suppliable: true package.loop_0_active: contract: dtype: bool @@ -89,8 +100,53 @@ pipeline: - 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: @@ -104,39 +160,6 @@ pipeline: implementation: kind: onnx artifact: vae_decoder/model.onnx - euler_model_input: - implementation: - kind: onnx - artifact: policies/euler_model_input.onnx - ports: - inputs: - sample: - dtype: float32 - rank: 4 - shape: - - batch - - channels - - height - - width - step: - dtype: int64 - rank: 1 - shape: - - batch - schedule: - dtype: float32 - rank: 1 - shape: - - schedule_length - outputs: - model_input: - dtype: float32 - rank: 4 - shape: - - batch - - channels - - height - - width solver_step: implementation: kind: onnx @@ -151,6 +174,9 @@ pipeline: - channels - height - width + batch_layout: + kind: request_aligned + axis: 0 derivative: dtype: float32 rank: 4 @@ -159,11 +185,17 @@ pipeline: - 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 @@ -178,6 +210,9 @@ pipeline: - channels - height - width + batch_layout: + kind: request_aligned + axis: 0 contract: id: onnx-genai.solver-step version: '1' @@ -198,12 +233,57 @@ pipeline: 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 @@ -244,14 +324,66 @@ pipeline: 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: + latent_state: contract: dtype: float32 rank: 4 @@ -260,8 +392,11 @@ pipeline: - 4 - height - width + batch_layout: + kind: request_aligned + axis: 0 scope: invocation - initializer: request.latent + initializer: diffusion.initial_state recurrence: kind: invariant loop_0_active: @@ -287,6 +422,18 @@ pipeline: 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: @@ -308,9 +455,9 @@ pipeline: outputs: timestep: diffusion.timestep - kind: invoke - component: euler_model_input + component: model_input_scale inputs: - sample: latent + sample: latent_state step: loop.iteration schedule: diffusion.schedule outputs: @@ -326,12 +473,20 @@ pipeline: - kind: invoke component: solver_step inputs: - sample: latent - derivative: denoiser.estimate + 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: @@ -341,7 +496,7 @@ pipeline: continue_when: loop_0_active max_iterations: request.max_iterations carried: - - cell: latent + - cell: latent_state next: latent.body - cell: loop_0_active next: loop.continue @@ -352,12 +507,19 @@ pipeline: rank: 1 shape: - batch + batch_layout: + kind: request_aligned + axis: 0 - kind: invoke component: vae_decoder inputs: - latent: latent + latent: latent_state outputs: image: vae.image + - kind: emit + value: latent_state + output: latent + mode: replace - kind: emit value: vae.image output: image diff --git a/tests/fixtures/onnx_genai_workflows/diffusion/policies/initial_state_scale.onnx b/tests/fixtures/onnx_genai_workflows/diffusion/policies/initial_state_scale.onnx new file mode 100644 index 0000000000000000000000000000000000000000..9bfb5425335bfa9b86c25844a44f8b3f0534e31e GIT binary patch literal 363 zcmaiv!Ab)$5Qdv&m2^Z}xQ7+D#%zQ1+C?sN3_+~5^62HdWXauqHiF?V zs_qa$TBkxyuwIauqZhLEnjf|r41iL?Qb;JN$OW7>@9>o}Ul*dbhC7>US=Q0$QAU~fa~S>Ep_WS}{p`< zr$Xm_J7sj3m{{>8sU=FI9Y6H_W^d= zaffL)uKPOsL~}wRAuGEjQIz*=zxRIry+642@=pVPA=p)H8cqb+o70aT3k5x~>q|if zuF;&GMC$;Qg>+kJfvip)c9dq3Zk3ES+=o5Mrb2~`3sM>$!}IXwx*JG&r+DOty2~d_ zDeCzmViEVg66^Ejunl(ZRP2K)9)v>0C*UU{^*MbLNE1qG+Or?6_ zPe>z%=+3x!FPNlf_CZEVfSx&!E)|86l+$kHmD5k!xhX|}<6Oe(g1PvqMW^lYH5(mMip5mxc$H7*!dXcHAVH1^>k9Gea z`RmC?;SwYA32f5h!88n{kh;6_)xBbxdtl)TJRwI9?0dq$_VgCxTi?~*sGIg8aiDF2 zJ>7CX2{sa2x`e!gp)r2@KVvv1&I2gJ$lOd*CBDo|0}6AW_SF~!a literal 0 HcmV?d00001 diff --git a/tests/fixtures/onnx_genai_workflows/diffusion_guided/denoiser/model.onnx.data b/tests/fixtures/onnx_genai_workflows/diffusion_guided/denoiser/model.onnx.data new file mode 100644 index 000000000..e69de29bb 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..fdbc5a51b --- /dev/null +++ b/tests/fixtures/onnx_genai_workflows/diffusion_guided/inference_metadata.yaml @@ -0,0 +1,777 @@ +schema_version: v1 +pipeline: + workflow: + manifest: + ir_version: '1.0' + onnx_opsets: + ai.onnx: 24 + 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 + denoiser: + implementation: + kind: onnx + artifact: denoiser/model.onnx + 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/diffusion_guided/policies/continue_predicate.onnx b/tests/fixtures/onnx_genai_workflows/diffusion_guided/policies/continue_predicate.onnx new file mode 100644 index 0000000000000000000000000000000000000000..3177d8905c97dc120048da84e39b00bcae00c3f8 GIT binary patch literal 910 zcmcIiJx{|h5Ur!oCReIf9-ssXL4`o-5Vf!|@gbH>m zU9gO*mp+iysemhKpX-P+jtiY#*oq$Rd?n>jvsgsNXCc*^^aPJ-%z7hiUEg%J!TW{E zdJxT2#C3*%cLXy9xH2G@JV+S7phMF3YCl1)s@-sBnPwH72hO8F^mQFAQQ#}WtLciA z)P_n!4!r8@o{KtYINyn84cq|@6HbIQ^}-TGy&+iK6)b4P4pY(ynQ(9Q{4t>+$)z4P zR`O^`e~HUor9w_+`kJ{~pZco`701Ocl<-QI8rJKs!AunNk3AdU?Y-kD5lI{>*;xnPGI+6BoP7L kn4%_4e3}eJC}=!7ZjaF(Z2#GV-`;X&a?79EP^i_v08>02xBvhE literal 0 HcmV?d00001 diff --git a/tests/fixtures/onnx_genai_workflows/diffusion_guided/policies/decoder_input_scale.onnx b/tests/fixtures/onnx_genai_workflows/diffusion_guided/policies/decoder_input_scale.onnx new file mode 100644 index 0000000000000000000000000000000000000000..b716c92f07b3983eba24bd837d04da54cc88c380 GIT binary patch literal 363 zcmaiv!Ab)$5Qdv&l{g};@uCH#N>S)x(W|FgZ~F=&VVb6GFik>|l`hiDK7~hLKztn^ z!iUgx*NY$?XZZeqzM;Epdu6Ap2-oikJQKF{h`Nyd?yWT;qKVwtndFye^k*?ThK3pK zQ_>@oX2{cz*SGNtN%K9}F}RY!p^&VT58=c;EHZ1%ItZ^^%!I}e_)MYX4jzPcGSmd?27wwq7p>Rypw(ailoFOgLP@pGy(Hn*P!vTQe#W-;9j*^RVFLGmP?eL5e( z)MyWaco-OdzW<;3**@Q0*qLtJ?MDFblryIa%$gDhhxgVvFO3(*W+<-D*^e-~f}Lks zHaLn|7}4q0lSYJ)2_diJZtl9BywNZ{fITA@=qjm@lXM8D)zdt&#w?wxb>$OPO6NqO zv6i(eRy=rkOk)`QqBa@ADO;i25S(W^&l{($>D2z`*%54i;Q)+L9zwvvmR`bHUo)Ab z>F)T>mQ=RXY9*94zLv`K;AP#?2@L-wh0gfNH`RwEvRUd<%X&4wY3T*Te`a8wqw-%B JrXi5%_!B6Vd({8{ literal 0 HcmV?d00001 diff --git a/tests/fixtures/onnx_genai_workflows/diffusion_guided/policies/diffusion_timesteps.onnx b/tests/fixtures/onnx_genai_workflows/diffusion_guided/policies/diffusion_timesteps.onnx new file mode 100644 index 0000000000000000000000000000000000000000..f6344fdf2fae202e0e65edd53b18c2471d558a47 GIT binary patch literal 383 zcmaivu};G<6h-YcAoBoGd8nX5sGtf^hp4Pf&@$y05VBn7rLjouD0WaI#FCF-X5#Dk z2vQ0ikPyRLuFk#hUbf4(7j~vAcl{B-JLSx&43jFy{=w85=cVz&*aXG+jQt3sOW1mn zd4;26*6Gs8#uM)mBJT+4wf8!GqhY)cJ4P|2hjgTtulmDwm`QcI8St%R!&>fs_VY99&CDW0LCZ}Az+<`Ucl*kVl+q7 z{_(9XDr~M*EtECBl*;qq`K6(^F#MMkI^jp()H)=R)l%nLmi5h5L(gIGX9lJzD*rWM J5(4QSegeS9d)EK} literal 0 HcmV?d00001 diff --git a/tests/fixtures/onnx_genai_workflows/diffusion_guided/policies/guidance_combine.onnx b/tests/fixtures/onnx_genai_workflows/diffusion_guided/policies/guidance_combine.onnx new file mode 100644 index 0000000000000000000000000000000000000000..4a71c4253dad3241fd648dd706c5611779bfa343 GIT binary patch literal 1635 zcmcJPPiqrF6u>)4+H77L8U_k+ixQV2QqrYq>%o8+6)yo#;-wv(YME zawsTx_TZ%-zz^d`a%N*Te?TOM9F}?S&6_vBee>RP%6MlGcSWLFzjxsV!M0*F;@BCJ z7gT9-;+E4@#sMcMUco9yCCKTxhHg==Q-PAC1FnXYalF+zgnPrQo*zs3QL&L2YM(_^ zDG~}kq9YbwUnzyM`c^N>iyx)7{JDl%(l7GsDq7z?en} zCqin~M(tcXe~XHnBKwIT%MC0s;>wu3k?MWI`6o_J-RfLPnNPh^X4w?ocLZe?%Rp!m zOBz8Xkqas=j>+q!OU_)dOG^3AX; zrNgqM2-x9DiviW#g=ax9g?Z%Vjl3ZkT5Pk?nyxpam zg#g%LA(fIx3LLx>a?uO5-Toi~9bPANo2Oq#4gSDSw880;Z`j6*zwPxC;N2g`5=FtiIh2Z{ldip&s z&4{yY=MUpZuxseR__~j*MIvwNB3$RtK F{Q(ll`9J^w literal 0 HcmV?d00001 diff --git a/tests/fixtures/onnx_genai_workflows/diffusion_guided/policies/history_initializer.onnx b/tests/fixtures/onnx_genai_workflows/diffusion_guided/policies/history_initializer.onnx new file mode 100644 index 0000000000000000000000000000000000000000..962692d28c0b298d171192c6e22a2612e45548d5 GIT binary patch literal 777 zcmcJNO-}+b5QZrNE`vtX1cQ4pNQ@d{HsHlXFMixe`~w=(lrpfDrQK|cipC4^>e(y5 z(zCk@5;QT!gS~V*ednDg?KmY`80fYfn&bC8yb-!`$1b!omMhB|?%F(7x`V7yb3POG zGL-IkFhtf3G`HnUye5RSR-vf)08PjR(oAayw!(hT(@H%U5y{YcBH+d_Ut+`~;g4wc z`leL{_nRv5U`FT=O`?ET0$<3e6LWH7S<@}0d&<|W<&s0a2wAHuWJ0q4#q&a-7-1J9 zR03fRL-UUqY2(7>FkmUtVQndn{|SRa+KvUXk5?xPMAU-=htPdkrujCv!Uv#^@VQbL z7;tF8N7?gja(yp5)^8udl>D!Sw;SL7r2x6OapsmTJjyILCVQ|w9X79rDjonUwemt8 zNHJnUD;sfPY4+)POwM5ES0$4jx;ax#O*}q)kb&gU=(IT|yHNen1G`-mHl3~W9Awrj EA6w1yN&o-= literal 0 HcmV?d00001 diff --git a/tests/fixtures/onnx_genai_workflows/diffusion_guided/policies/latent_noise.onnx b/tests/fixtures/onnx_genai_workflows/diffusion_guided/policies/latent_noise.onnx new file mode 100644 index 0000000000000000000000000000000000000000..c9c6bbb9dc86870b09df5299d11a3472a2ee800d GIT binary patch literal 18187 zcmc&+O>Z1Y8TL4iJ#NQqyV-m&u-PyQ2t~V=ySu)6PhrC;Vg(99LW@vqGH%C{*fZXl zv6BTwT17%20Ychk5zB`ZaplB6;L2}6`y&ts#0lQFdb+;4XEvAW@nxpEUENPT_4ZS5 zJ=Hbst#JKx^l)%KdFMaBZUxVSwMm{IhMO1J_m6vLd3MmbQhXVX4x!Y&EKB#dgUw;@ zG@qRH`gyo=|8{Wm?8&43(Qx=|(jO1bru+Sq-ei&;5Atzu+&}&>T>a6H@81eK&VS)G zvPpk*mQV0Mf7sWeU!CLRxU~*&S zWoy+JzP|uI4eGUv>_L8b-p{``9vxCOKPWZ+&%p=hdH$0;qn+2vuaZXE-I$eqp>sNtxtuYBGpi$gM%T-hMNtx% z6j`Sjs>p-h@KK)e&XwYeWR7<)%aSkNoLCbsZ^GcMxmf#?)V3KP7EvPHzQ}$soXi3( z?d)_nQ*8C!p773q#0S;UDfT8%8)KxF;ic_3MZG*kay zHwW_n!QY0*4C8kWG+u$m4Xp@GXutXIsz+eIb%2QqFjM;*8#6z4MS|K9C)7iknh#~_ z^`Ug)E3OEBYzn>=Uw`}e=l|OD*o}pJ6)I&h6(5vp8b^kj#*tSw{kc?A2`4ReHIB+Y zi$T|wkj9|v9)WX@Bm&^5`#Tbm(5gfv zbSV;%*t#0&xMzk$B&;EcNZ7o%2btG#4>I3W+#3?XZETtB#Gv)MvIaLn_7$F^!q4 zu&5wcrH)QUz3J7Jm6h+lY_)vl@v!W<3C?%Hj}`1h8Z%NbK;*SxfyjFlEGNEcMc1jR z>sEaG^wtU!ORr~UwoXxH%}_Iir+!m-!l>7_qTB`J%ntD$C>bVKP!>BiEq3bHVlT2N9Pz202~#;)#_WliUJP4`7sG5XhLX~*Wu}lbze27P zPG);C234jP!L6xcN%&)E+5@CBW234jP!R$%0yFzE?geJ{O~k!$l6e*P z0+abN?gb|E4aL1HeDC)= zqeqzlEyqu}0w6%kq2S9|hwQ+7XIpTDQ-CNv27wv7Qc29%{i)-MjtJlq=Q&!$a=HqNU#6e(u+BD@Nnf!B*2w1J>)2jhH;Te91$$6=TXzA{ z^(LwsYFHH4|HB*j54d^}CE;GM16SD6SEyjH8FDeu`deJL+o36QPH@Ggb3&*SvN;lM z`YQK=tpsIH%XT6W@KS$GI~5IaJohxttm0ah^u@{I+O+Pp-CB4J+ro?~Zp`cI-LAZv zT=!oZTRQ(Vv@AK8e;Q~@=F^6}B(=2JErYj5D=kZE33FSP)bgfUmazHI5nC;Ly0%zc zyOuV)W$@BlrDdd+Ft=r-mN$isEC@Gf8L4H@!A5Ot`1HWrLN(98P2&|AdA?<9fpgvB zsceCYO1K!OFayEsF02Rdkd$0qE5!-k<}lS-UJr57T6Sz4V33yAb54+XCp05UPbhm8 zJ)v5nR2R|t;Uk++8$LyC=s~vlsbjT4q-(zdKd??%@+0eEKf#J!P5-;G$re9#tTu== zw#82!n@=0+39F^e?qFEnsI&|g8eXAgyq32DI%pkV{M7MU_H^w9v_Z>wEp1NAoH)Zi z1Ap{YT2{-O!X_4k8(sopOlTy6y;(7&F zs@}|csHvw}sFf)RSGojg!P3IprjZR7Qwx~H&b#ewVma@$tX)ge>(2D6$b#)IfqBNP zOJ218+6ul_{MO6?t({V-{6?~c@#XX3vfnG4)JL51&`jzYUX{yU~l*H-cMpFOeS|<^8Gr-I;gSf|cDH{|8gd$c_L2 literal 0 HcmV?d00001 diff --git a/tests/fixtures/onnx_genai_workflows/diffusion_guided/policies/latent_row_shape.onnx b/tests/fixtures/onnx_genai_workflows/diffusion_guided/policies/latent_row_shape.onnx new file mode 100644 index 0000000000000000000000000000000000000000..254b8624d88e6e5ce78ba8ce729117ce4d6813d4 GIT binary patch literal 382 zcmaiv!A=4(5QewQB6W;P>IFF%G^h!OMy{R|Z~F?yG^G@_acP@v*RUa8_DOsI-;9^Q zrFbDR&SjFxH~)NtPsq{IOjYR~H!*x8(smiQvVZ%L*!Shm% zZgH#InS_P3mqeis`xi>mQ#jUqDV^m)k|-U*m3?1OqxH&(LRn8m&YfeKk_9hBR+0Gi zEgeArfNC(6@6s9t?&9p-x#tmB^z^xDD zLpcMeNMc;M>&^LcdhWSi3FoJ3AQE%=?!qeq&#;gWlZb>#EJ%kV8U>oCdY{PzaQ)&guUI7=OfcFN}vT$$&UmM{}L|fwH@sHJaEtwt`n^v1;P4*lUd7|fBpRjpBja;p^Z5bk zrtw9RP2)jgZajaFygU8z+0i6Tud~5yG|i6&<2cKr;V7BKv%&C{x%tB%J%2*1+OO~( zQ8t)NlMH`R2W8hm**AgO&IJz%u5_e*xM;?5BJD+d_BsvdiOtNWF=M1Jps<{TW`te?EZS;;jqc@ z>>50KrSSf4mh6rfw;>VJUZQmihzuKxq-8jHipHLh-hlwI@W})~n*5 zayUh=JRJCBmPfV+;Q$tj!SZ_Bc&BT${%ELnpnS3|aG*)PGrEdw-`c}#0SKfVZgk

FpTy~PouWvY1GylFIApK zZQ*GY18?DJl&;p(0^w%K00wcjMgY=igym>3 zp}0&|T(%0uK@!3dMy)jk2O3yYr<*DL4$&7fDwT3)EMm;M5oQjK}uaJM$0^rf?WYu}B;- zkYbIAqtDtlhiyC?@zE@Sgu^Uifpxbm7ZwV7yYmOS+z%wFZOa<$gBuSF)LCqIFi>Z> z(S-yuj-uCU*SZUHRssh+yuP*13*0bXzqgHdy0E|kRT&pd#?E0Nbv{T>BM;{~v(Y81H zCv|_#t1`+Slh+&`TiKce##70OdVEDzh$XrK&*yp!&n^4_PAw#n=UE3usc;3acZc8V z@=XXB6o5`xCQVCK_heP~^s0KDU5y-xr{+lN?9#dOUuRb!R*m}i4@EEup6l_mNOq*T#j?dQYX_Wn~bj@R>Ja&KzRNU zdB00&W1Tw|73KU4_`c!HhPwA8&+ZMvP~b( z4}IMw2c2EBE8xEPs&w-69r6&q=@^k@*+p6dr-xGYLoIrqggi|yiOcg zYp=tg-@mCHKS|bT_a!BHgw=*LUgz(snVrS;Ht$+#|4<2)Y|3qBfGr=Fu4X&8wsO^0 zC>1tqw)&ztN??1*pgpRd6RACThwh${(Vm*Z4x%FI@PAVu+j;Lsyk89F#D&kg$d^ zScoxPj8P>t>~<*BL=(HMR+^6;Yzj2^Re^U~sVP&)br*qm*<$eAhM22#TWnZWwwe$= zfPY=k{D7)ui!X~x$Wt~xomb*(A$8bZS!5HNm6g%htww|F)&f?tTZ|@r2eLj(PMGJy zXPLir+b^pb1a5+^JcV%Mp%}8l>~cJC&h1lX9MH5aP5F|$?Y%PNU=_Z=cwXVs-J3*9 zbr)EkmhSSlyc8@SLUbYp^lxlzZ+Tniy=FI z&QBQOqw=g1J{4tnKF`|Xz?d@tsd`MR07tHoiul-1v&*<;!G{^S^+7dz){p`< zr$Xm_J7sj3m{{>8QIvS)OCjbgN@xw<;X> z2kgQPEIR1jAvr{h4>#Yq9M?oRrwjr zh*V3HuYhlebeHU9s*L5zvI)Jn85i{@A^tNlx(~4A*9*xKWzCD>8nv6ET03brdcYPS zfdW@>%Z_NzvE`!z^(fpg!%+uXe3h7iw?=+bh~Mko0}#1_Y!b6W(p$3k%6wi(`Bk!G z+OL8F^EeDPGflzWBG{%N*r~=H;R)>&-vZ?%Y~{fy0NrMZgB#y$`C%djE2NY~1hWxq*A6r+Cotk%&8`(F8cxTs&* z*Km&8B!UC-NThR}Dyi(ae$@DqTQP-a$t{v|}*VV${xi_yzD|?VU MhpOCJ3%vf|7vDgf`~Uy| literal 0 HcmV?d00001 diff --git a/tests/fixtures/onnx_genai_workflows/diffusion_guided/text_encoder/model.onnx.data b/tests/fixtures/onnx_genai_workflows/diffusion_guided/text_encoder/model.onnx.data new file mode 100644 index 000000000..e69de29bb diff --git a/tests/fixtures/onnx_genai_workflows/diffusion_guided/vae_decoder/model.onnx b/tests/fixtures/onnx_genai_workflows/diffusion_guided/vae_decoder/model.onnx new file mode 100644 index 0000000000000000000000000000000000000000..773c8b4b2b4432b5f206ae035966db2fd0edc2f7 GIT binary patch literal 957 zcmchVu};G<5QdvJCE;2GctQaxR539`N!XYWY+2br$Z{ML8xkkPCcty>0xY}$Pe@u+ z)I8dvx{mr4_`m>(#nCoEMbZa`h1BY_ z_i9+-2CPexF;x)CP%k`$NpY9Bxs*?emZH!uO^H&N3s#bn@@Le)x(*M(X{ly2u%dav z)J)N)^4nB?C)MEpsDd_C&`D*ikm}8VwUlTkH8hfJRDPoh&`@L~Vb1s>6_nw5(#O-` zyjoGFXjBn~V@B&=RbG}0X2jvfHc$lnXqV<$k;~eMyCiiZqA7>IZNdcYeJ|lmBwTlo scoebb(LH>+{{70qA6MQwjc4}|fNfjItlb!`p8rEHSni#f(A(Pi06iTsFaQ7m literal 0 HcmV?d00001 diff --git a/tests/fixtures/onnx_genai_workflows/diffusion_guided/vae_decoder/model.onnx.data b/tests/fixtures/onnx_genai_workflows/diffusion_guided/vae_decoder/model.onnx.data new file mode 100644 index 000000000..e69de29bb diff --git a/tests/fixtures/onnx_genai_workflows/masked/inference_metadata.yaml b/tests/fixtures/onnx_genai_workflows/masked/inference_metadata.yaml index 8c1a28821..7aaf1bc73 100644 --- a/tests/fixtures/onnx_genai_workflows/masked/inference_metadata.yaml +++ b/tests/fixtures/onnx_genai_workflows/masked/inference_metadata.yaml @@ -20,6 +20,9 @@ pipeline: shape: - batch - sequence + batch_layout: + kind: request_aligned + axis: 0 role: kind: runtime version: '1.0' @@ -34,6 +37,9 @@ pipeline: shape: - batch - sequence + batch_layout: + kind: request_aligned + axis: 0 role: kind: opaque source: @@ -46,6 +52,9 @@ pipeline: rank: 1 shape: - batch + batch_layout: + kind: request_aligned + axis: 0 role: kind: runtime version: '1.0' @@ -60,6 +69,9 @@ pipeline: rank: 1 shape: - batch + batch_layout: + kind: request_aligned + axis: 0 role: kind: opaque source: @@ -87,6 +99,9 @@ pipeline: rank: 1 shape: - batch + batch_layout: + kind: request_aligned + axis: 0 role: kind: opaque source: @@ -113,6 +128,9 @@ pipeline: shape: - batch - sequence + batch_layout: + kind: request_aligned + axis: 0 role: tokens stage: pre_adapter components: @@ -132,12 +150,18 @@ pipeline: 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 @@ -145,32 +169,50 @@ pipeline: - 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 @@ -178,22 +220,34 @@ pipeline: 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 @@ -221,6 +275,9 @@ pipeline: shape: - batch - sequence + batch_layout: + kind: request_aligned + axis: 0 scope: invocation initializer: request.input_ids recurrence: @@ -232,6 +289,9 @@ pipeline: shape: - batch - sequence + batch_layout: + kind: request_aligned + axis: 0 scope: invocation initializer: request.mask recurrence: @@ -242,6 +302,9 @@ pipeline: rank: 1 shape: - batch + batch_layout: + kind: request_aligned + axis: 0 scope: invocation initializer: request.rng_offset recurrence: @@ -254,6 +317,9 @@ pipeline: - batch - sequence - 128 + batch_layout: + kind: request_aligned + axis: 0 scope: invocation initializer: denoiser.setup.logits recurrence: @@ -265,6 +331,9 @@ pipeline: shape: - batch - sequence + batch_layout: + kind: request_aligned + axis: 0 scope: invocation initializer: denoiser.setup.proposal recurrence: diff --git a/tests/fixtures/onnx_genai_workflows/speculative/inference_metadata.yaml b/tests/fixtures/onnx_genai_workflows/speculative/inference_metadata.yaml index d8db9b97f..a9e57845a 100644 --- a/tests/fixtures/onnx_genai_workflows/speculative/inference_metadata.yaml +++ b/tests/fixtures/onnx_genai_workflows/speculative/inference_metadata.yaml @@ -12,7 +12,6 @@ pipeline: - loop_induction_values - typed_emit - emit_valid_length - - emit_row_identity - bounded_state_recurrence - serving_service_contract - grammar_guidance_adapter @@ -28,6 +27,9 @@ pipeline: shape: - batch - 4 + batch_layout: + kind: request_aligned + axis: 0 role: kind: runtime version: '1.0' @@ -41,6 +43,9 @@ pipeline: rank: 1 shape: - batch + batch_layout: + kind: request_aligned + axis: 0 role: kind: runtime version: '1.0' @@ -68,6 +73,9 @@ pipeline: rank: 1 shape: - batch + batch_layout: + kind: request_aligned + axis: 0 role: kind: opaque source: @@ -80,6 +88,9 @@ pipeline: rank: 1 shape: - batch + batch_layout: + kind: request_aligned + axis: 0 role: kind: opaque source: @@ -104,6 +115,9 @@ pipeline: rank: 1 shape: - batch + batch_layout: + kind: request_aligned + axis: 0 role: kind: opaque source: @@ -116,30 +130,24 @@ pipeline: rank: 1 shape: - batch + batch_layout: + kind: request_aligned + axis: 0 role: kind: opaque source: kind: literal required: false default: true - request.slot_ids: - contract: - dtype: int64 - rank: 1 - shape: - - batch - role: - kind: opaque - source: - kind: application - name: serving.slot_ids - required: true request.cache_lengths: contract: dtype: int64 rank: 1 shape: - batch + batch_layout: + kind: request_aligned + axis: 0 role: kind: opaque source: @@ -153,6 +161,9 @@ pipeline: rank: 1 shape: - batch + batch_layout: + kind: request_aligned + axis: 0 role: kind: opaque source: @@ -178,6 +189,9 @@ pipeline: rank: 1 shape: - batch + batch_layout: + kind: request_aligned + axis: 0 role: kind: opaque source: @@ -192,6 +206,9 @@ pipeline: shape: - batch - 24 + batch_layout: + kind: request_aligned + axis: 0 role: kind: opaque source: @@ -205,6 +222,9 @@ pipeline: rank: 1 shape: - batch + batch_layout: + kind: request_aligned + axis: 0 role: kind: opaque source: @@ -217,6 +237,9 @@ pipeline: rank: 1 shape: - batch + batch_layout: + kind: request_aligned + axis: 0 role: kind: opaque source: @@ -232,6 +255,9 @@ pipeline: - 2 - past_sequence - 8 + batch_layout: + kind: request_aligned + axis: 0 role: kind: opaque source: @@ -246,6 +272,9 @@ pipeline: shape: - batch - accepted_sequence + batch_layout: + kind: request_aligned + axis: 0 role: tokens stage: pre_adapter components: @@ -269,17 +298,26 @@ pipeline: 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 @@ -292,28 +330,43 @@ pipeline: 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' @@ -343,17 +396,26 @@ pipeline: 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 @@ -366,28 +428,43 @@ pipeline: 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' @@ -417,17 +494,26 @@ pipeline: 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 @@ -440,28 +526,43 @@ pipeline: 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' @@ -492,22 +593,34 @@ pipeline: - 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 @@ -515,31 +628,49 @@ pipeline: 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' @@ -565,23 +696,35 @@ pipeline: 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 @@ -589,6 +732,9 @@ pipeline: shape: - batch - 1 + batch_layout: + kind: request_aligned + axis: 0 adaptive_k: implementation: kind: onnx @@ -600,54 +746,84 @@ pipeline: 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' @@ -673,17 +849,26 @@ pipeline: 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 @@ -697,6 +882,9 @@ pipeline: - batch - sequence - vocabulary + batch_layout: + kind: request_aligned + axis: 0 outputs: last_logits: dtype: float32 @@ -704,6 +892,9 @@ pipeline: shape: - batch - vocabulary + batch_layout: + kind: request_aligned + axis: 0 proposal_metrics: implementation: kind: onnx @@ -716,22 +907,34 @@ pipeline: 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 @@ -743,17 +946,26 @@ pipeline: 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: @@ -762,6 +974,9 @@ pipeline: shape: - batch - 4 + batch_layout: + kind: request_aligned + axis: 0 class: semantic scope: invocation initializer: request.tokens @@ -773,6 +988,9 @@ pipeline: rank: 1 shape: - batch + batch_layout: + kind: request_aligned + axis: 0 class: semantic scope: invocation initializer: package.zero @@ -784,6 +1002,9 @@ pipeline: rank: 1 shape: - batch + batch_layout: + kind: request_aligned + axis: 0 class: semantic scope: invocation initializer: package.active @@ -795,6 +1016,9 @@ pipeline: rank: 1 shape: - batch + batch_layout: + kind: request_aligned + axis: 0 class: semantic scope: invocation initializer: package.false @@ -806,28 +1030,23 @@ pipeline: rank: 1 shape: - batch + batch_layout: + kind: request_aligned + axis: 0 class: semantic scope: invocation initializer: package.zero recurrence: kind: invariant - slot_ids: - contract: - dtype: int64 - rank: 1 - shape: - - batch - class: semantic - scope: invocation - initializer: request.slot_ids - 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 @@ -839,6 +1058,9 @@ pipeline: rank: 1 shape: - batch + batch_layout: + kind: request_aligned + axis: 0 class: semantic scope: invocation initializer: request.grammar_state @@ -850,6 +1072,9 @@ pipeline: rank: 1 shape: - batch + batch_layout: + kind: request_aligned + axis: 0 class: advisory scope: invocation initializer: request.adaptive_k @@ -862,6 +1087,9 @@ pipeline: shape: - batch - 24 + batch_layout: + kind: request_aligned + axis: 0 class: advisory scope: invocation initializer: request.adaptive_estimates @@ -876,6 +1104,9 @@ pipeline: - 2 - past_sequence - 8 + batch_layout: + kind: request_aligned + axis: 0 class: semantic scope: invocation initializer: request.verifier.past_key_values.0.key @@ -884,26 +1115,31 @@ pipeline: axis: 2 max: package.max_context service_group: verifier_cache + management: runtime + release_boundary: invocation serving: active: active done: done accepted_len: accepted_len - slot_ids: slot_ids - kv_service: - paging: none - allocation: runtime - compaction: true + state_service: groups: verifier_cache: + kind: full_attention sequence_axis: 2 layout: bnsh - logical_lengths: cache_lengths - storage: shared_buffer + 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 + logical_lengths: cache_lengths steps: - kind: loop setup: [] @@ -1033,13 +1269,11 @@ pipeline: output: tokens mode: append valid_length: grammar.committed_length - row_ids: slot_ids - kind: emit value: grammar.token output: tokens mode: append valid_length: grammar.forced_length - row_ids: slot_ids continue_when: active max_iterations: request.max_iterations carried: @@ -1053,8 +1287,6 @@ pipeline: next: acceptance.done - cell: accepted_len next: grammar.committed_length - - cell: slot_ids - next: slot_ids - cell: cache_lengths next: cache_lengths.next - cell: grammar @@ -1072,3 +1304,6 @@ pipeline: rank: 1 shape: - batch + batch_layout: + kind: request_aligned + axis: 0 diff --git a/tests/fixtures/onnx_genai_workflows/tts/inference_metadata.yaml b/tests/fixtures/onnx_genai_workflows/tts/inference_metadata.yaml index 37a9faad7..d382059f6 100644 --- a/tests/fixtures/onnx_genai_workflows/tts/inference_metadata.yaml +++ b/tests/fixtures/onnx_genai_workflows/tts/inference_metadata.yaml @@ -21,6 +21,9 @@ pipeline: shape: - batch - text_sequence_len + batch_layout: + kind: request_aligned + axis: 0 role: kind: runtime version: '1.0' @@ -47,6 +50,9 @@ pipeline: rank: 1 shape: - batch + batch_layout: + kind: request_aligned + axis: 0 role: kind: opaque source: @@ -141,6 +147,9 @@ pipeline: rank: 1 shape: - batch + batch_layout: + kind: request_aligned + axis: 0 role: kind: opaque source: @@ -153,6 +162,9 @@ pipeline: rank: 1 shape: - batch + batch_layout: + kind: request_aligned + axis: 0 role: kind: opaque source: @@ -165,24 +177,15 @@ pipeline: rank: 1 shape: - batch + batch_layout: + kind: request_aligned + axis: 0 role: kind: opaque source: kind: literal required: false default: true - package.slot_ids: - contract: - dtype: int64 - rank: 1 - shape: - - batch - role: - kind: opaque - source: - kind: application - name: serving.slot_ids - required: true package.setup_predictor_iteration_0: contract: dtype: int64 @@ -238,6 +241,9 @@ pipeline: - batch - 1 - frames + batch_layout: + kind: request_aligned + axis: 0 role: audio stage: post_adapter components: @@ -294,6 +300,9 @@ pipeline: - batch - sequence - vocabulary + batch_layout: + kind: request_aligned + axis: 0 outputs: last_logits: dtype: float32 @@ -301,6 +310,9 @@ pipeline: shape: - batch - vocabulary + batch_layout: + kind: request_aligned + axis: 0 setup_talker_sampler: implementation: kind: onnx @@ -313,12 +325,18 @@ pipeline: 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' @@ -340,12 +358,18 @@ pipeline: 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' @@ -367,12 +391,18 @@ pipeline: 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' @@ -394,12 +424,18 @@ pipeline: 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' @@ -421,12 +457,18 @@ pipeline: 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' @@ -447,6 +489,9 @@ pipeline: rank: 1 shape: - batch + batch_layout: + kind: request_aligned + axis: 0 outputs: continue: dtype: bool @@ -465,6 +510,9 @@ pipeline: shape: - batch - sequence + batch_layout: + kind: request_aligned + axis: 0 outputs: frame_codes: dtype: int64 @@ -472,12 +520,18 @@ pipeline: 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 @@ -485,6 +539,9 @@ pipeline: - batch - 0 - 4 + batch_layout: + kind: request_aligned + axis: 0 token_to_slot: implementation: kind: onnx @@ -496,6 +553,9 @@ pipeline: rank: 1 shape: - batch + batch_layout: + kind: request_aligned + axis: 0 outputs: slot: dtype: int64 @@ -503,6 +563,9 @@ pipeline: shape: - batch - 1 + batch_layout: + kind: request_aligned + axis: 0 code_frame_update: implementation: kind: onnx @@ -515,11 +578,17 @@ pipeline: 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 @@ -531,6 +600,9 @@ pipeline: shape: - batch - 4 + batch_layout: + kind: request_aligned + axis: 0 code_history_append: implementation: kind: onnx @@ -544,12 +616,18 @@ pipeline: - 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 @@ -558,6 +636,9 @@ pipeline: - batch - frames + 1 - 4 + batch_layout: + kind: request_aligned + axis: 0 cache_length_update: implementation: kind: onnx @@ -569,17 +650,26 @@ pipeline: 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 @@ -593,6 +683,9 @@ pipeline: - batch - prefill_sequence - 8 + batch_layout: + kind: request_aligned + axis: 0 outputs: attention_mask: dtype: int64 @@ -600,6 +693,9 @@ pipeline: shape: - batch - past_seq_len + seq_len + batch_layout: + kind: request_aligned + axis: 0 position_ids: dtype: int64 rank: 3 @@ -613,6 +709,9 @@ pipeline: shape: - batch - prefill_sequence + 1 + batch_layout: + kind: request_aligned + axis: 0 body_position_ids: dtype: int64 rank: 3 @@ -628,6 +727,9 @@ pipeline: - 1 - past_sequence_len - 4 + batch_layout: + kind: request_aligned + axis: 0 past_key_values.0.value: dtype: float32 rank: 4 @@ -636,6 +738,9 @@ pipeline: - 1 - past_sequence_len - 4 + batch_layout: + kind: request_aligned + axis: 0 predictor_state_initializer: implementation: kind: onnx @@ -649,6 +754,9 @@ pipeline: - batch - prefill_sequence - 8 + batch_layout: + kind: request_aligned + axis: 0 outputs: attention_mask: dtype: int64 @@ -656,24 +764,36 @@ pipeline: 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 @@ -682,6 +802,9 @@ pipeline: - 8 - past_sequence_len - 128 + batch_layout: + kind: request_aligned + axis: 0 past_key_values.0.value: dtype: float32 rank: 4 @@ -690,6 +813,9 @@ pipeline: - 8 - past_sequence_len - 128 + batch_layout: + kind: request_aligned + axis: 0 past_key_values.1.key: dtype: float32 rank: 4 @@ -698,6 +824,9 @@ pipeline: - 8 - past_sequence_len - 128 + batch_layout: + kind: request_aligned + axis: 0 past_key_values.1.value: dtype: float32 rank: 4 @@ -706,6 +835,9 @@ pipeline: - 8 - past_sequence_len - 128 + batch_layout: + kind: request_aligned + axis: 0 past_key_values.2.key: dtype: float32 rank: 4 @@ -714,6 +846,9 @@ pipeline: - 8 - past_sequence_len - 128 + batch_layout: + kind: request_aligned + axis: 0 past_key_values.2.value: dtype: float32 rank: 4 @@ -722,6 +857,9 @@ pipeline: - 8 - past_sequence_len - 128 + batch_layout: + kind: request_aligned + axis: 0 past_key_values.3.key: dtype: float32 rank: 4 @@ -730,6 +868,9 @@ pipeline: - 8 - past_sequence_len - 128 + batch_layout: + kind: request_aligned + axis: 0 past_key_values.3.value: dtype: float32 rank: 4 @@ -738,6 +879,9 @@ pipeline: - 8 - past_sequence_len - 128 + batch_layout: + kind: request_aligned + axis: 0 past_key_values.4.key: dtype: float32 rank: 4 @@ -746,6 +890,9 @@ pipeline: - 8 - past_sequence_len - 128 + batch_layout: + kind: request_aligned + axis: 0 past_key_values.4.value: dtype: float32 rank: 4 @@ -754,6 +901,9 @@ pipeline: - 8 - past_sequence_len - 128 + batch_layout: + kind: request_aligned + axis: 0 talker_step_update: implementation: kind: onnx @@ -766,6 +916,9 @@ pipeline: shape: - batch - context + batch_layout: + kind: request_aligned + axis: 0 position_ids: dtype: int64 rank: 3 @@ -780,6 +933,9 @@ pipeline: shape: - batch - context + 1 + batch_layout: + kind: request_aligned + axis: 0 next_position_ids: dtype: int64 rank: 3 @@ -799,12 +955,18 @@ pipeline: 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 @@ -812,12 +974,18 @@ pipeline: 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 @@ -831,6 +999,9 @@ pipeline: - batch - frames - 4 + batch_layout: + kind: request_aligned + axis: 0 outputs: codes: dtype: int64 @@ -839,6 +1010,9 @@ pipeline: - batch - 4 - frames + batch_layout: + kind: request_aligned + axis: 0 state: last_frame: contract: @@ -847,6 +1021,9 @@ pipeline: shape: - batch - 4 + batch_layout: + kind: request_aligned + axis: 0 scope: invocation initializer: setup.predictor.remaining_1.frame recurrence: @@ -859,6 +1036,9 @@ pipeline: - batch - frames - 4 + batch_layout: + kind: request_aligned + axis: 0 scope: invocation initializer: history.setup recurrence: @@ -873,6 +1053,9 @@ pipeline: shape: - batch - talker_context + batch_layout: + kind: request_aligned + axis: 0 scope: invocation initializer: talker.initializer.body_attention_mask recurrence: @@ -899,6 +1082,9 @@ pipeline: shape: - batch - 4 + batch_layout: + kind: request_aligned + axis: 0 scope: invocation initializer: frame.frame_prefill recurrence: @@ -909,6 +1095,9 @@ pipeline: rank: 1 shape: - batch + batch_layout: + kind: request_aligned + axis: 0 scope: invocation initializer: frame.group1 recurrence: @@ -920,6 +1109,9 @@ pipeline: shape: - batch - predictor_context + batch_layout: + kind: request_aligned + axis: 0 scope: invocation initializer: frame.predictor.initializer.body_attention_mask recurrence: @@ -934,6 +1126,9 @@ pipeline: shape: - batch - 1 + batch_layout: + kind: request_aligned + axis: 0 scope: invocation initializer: frame.predictor.initializer.body_position_ids recurrence: @@ -944,6 +1139,9 @@ pipeline: rank: 1 shape: - batch + batch_layout: + kind: request_aligned + axis: 0 class: semantic scope: invocation initializer: package.true @@ -955,6 +1153,9 @@ pipeline: rank: 1 shape: - batch + batch_layout: + kind: request_aligned + axis: 0 class: semantic scope: invocation initializer: package.false @@ -966,28 +1167,23 @@ pipeline: rank: 1 shape: - batch + batch_layout: + kind: request_aligned + axis: 0 class: semantic scope: invocation initializer: package.zero_batch recurrence: kind: invariant - slot_ids: - contract: - dtype: int64 - rank: 1 - shape: - - batch - class: semantic - scope: invocation - initializer: package.slot_ids - 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 @@ -999,6 +1195,9 @@ pipeline: rank: 1 shape: - batch + batch_layout: + kind: request_aligned + axis: 0 class: semantic scope: invocation initializer: package.zero_batch @@ -1013,6 +1212,9 @@ pipeline: - 1 - past_sequence_len - 4 + batch_layout: + kind: request_aligned + axis: 0 class: semantic scope: invocation initializer: talker.setup.present.0.key @@ -1021,6 +1223,8 @@ pipeline: axis: 2 max: package.talker_context_limit service_group: talker_cache + management: runtime + release_boundary: invocation talker_cache_1: contract: dtype: float32 @@ -1030,6 +1234,9 @@ pipeline: - 1 - past_sequence_len - 4 + batch_layout: + kind: request_aligned + axis: 0 class: semantic scope: invocation initializer: talker.setup.present.0.value @@ -1038,6 +1245,8 @@ pipeline: axis: 2 max: package.talker_context_limit service_group: talker_cache + management: runtime + release_boundary: invocation predictor_cache_0: contract: dtype: float32 @@ -1047,6 +1256,9 @@ pipeline: - 8 - past_sequence_len - 128 + batch_layout: + kind: request_aligned + axis: 0 class: semantic scope: invocation initializer: frame.predictor.present.0.key @@ -1055,6 +1267,8 @@ pipeline: axis: 2 max: package.predictor_context_limit service_group: predictor_cache + management: runtime + release_boundary: invocation predictor_cache_1: contract: dtype: float32 @@ -1064,6 +1278,9 @@ pipeline: - 8 - past_sequence_len - 128 + batch_layout: + kind: request_aligned + axis: 0 class: semantic scope: invocation initializer: frame.predictor.present.0.value @@ -1072,6 +1289,8 @@ pipeline: axis: 2 max: package.predictor_context_limit service_group: predictor_cache + management: runtime + release_boundary: invocation predictor_cache_2: contract: dtype: float32 @@ -1081,6 +1300,9 @@ pipeline: - 8 - past_sequence_len - 128 + batch_layout: + kind: request_aligned + axis: 0 class: semantic scope: invocation initializer: frame.predictor.present.1.key @@ -1089,6 +1311,8 @@ pipeline: axis: 2 max: package.predictor_context_limit service_group: predictor_cache + management: runtime + release_boundary: invocation predictor_cache_3: contract: dtype: float32 @@ -1098,6 +1322,9 @@ pipeline: - 8 - past_sequence_len - 128 + batch_layout: + kind: request_aligned + axis: 0 class: semantic scope: invocation initializer: frame.predictor.present.1.value @@ -1106,6 +1333,8 @@ pipeline: axis: 2 max: package.predictor_context_limit service_group: predictor_cache + management: runtime + release_boundary: invocation predictor_cache_4: contract: dtype: float32 @@ -1115,6 +1344,9 @@ pipeline: - 8 - past_sequence_len - 128 + batch_layout: + kind: request_aligned + axis: 0 class: semantic scope: invocation initializer: frame.predictor.present.2.key @@ -1123,6 +1355,8 @@ pipeline: axis: 2 max: package.predictor_context_limit service_group: predictor_cache + management: runtime + release_boundary: invocation predictor_cache_5: contract: dtype: float32 @@ -1132,6 +1366,9 @@ pipeline: - 8 - past_sequence_len - 128 + batch_layout: + kind: request_aligned + axis: 0 class: semantic scope: invocation initializer: frame.predictor.present.2.value @@ -1140,6 +1377,8 @@ pipeline: axis: 2 max: package.predictor_context_limit service_group: predictor_cache + management: runtime + release_boundary: invocation predictor_cache_6: contract: dtype: float32 @@ -1149,6 +1388,9 @@ pipeline: - 8 - past_sequence_len - 128 + batch_layout: + kind: request_aligned + axis: 0 class: semantic scope: invocation initializer: frame.predictor.present.3.key @@ -1157,6 +1399,8 @@ pipeline: axis: 2 max: package.predictor_context_limit service_group: predictor_cache + management: runtime + release_boundary: invocation predictor_cache_7: contract: dtype: float32 @@ -1166,6 +1410,9 @@ pipeline: - 8 - past_sequence_len - 128 + batch_layout: + kind: request_aligned + axis: 0 class: semantic scope: invocation initializer: frame.predictor.present.3.value @@ -1174,6 +1421,8 @@ pipeline: axis: 2 max: package.predictor_context_limit service_group: predictor_cache + management: runtime + release_boundary: invocation predictor_cache_8: contract: dtype: float32 @@ -1183,6 +1432,9 @@ pipeline: - 8 - past_sequence_len - 128 + batch_layout: + kind: request_aligned + axis: 0 class: semantic scope: invocation initializer: frame.predictor.present.4.key @@ -1191,6 +1443,8 @@ pipeline: axis: 2 max: package.predictor_context_limit service_group: predictor_cache + management: runtime + release_boundary: invocation predictor_cache_9: contract: dtype: float32 @@ -1200,6 +1454,9 @@ pipeline: - 8 - past_sequence_len - 128 + batch_layout: + kind: request_aligned + axis: 0 class: semantic scope: invocation initializer: frame.predictor.present.4.value @@ -1208,6 +1465,8 @@ pipeline: axis: 2 max: package.predictor_context_limit service_group: predictor_cache + management: runtime + release_boundary: invocation loop_1_active: contract: dtype: bool @@ -1232,17 +1491,19 @@ pipeline: active: active done: done accepted_len: accepted_len - slot_ids: slot_ids - kv_service: - paging: none - allocation: runtime - compaction: true + state_service: groups: talker_cache: + kind: full_attention sequence_axis: 2 layout: bnsh - logical_lengths: talker_cache_lengths - storage: shared_buffer + aliasing: forbidden + reuse: + prefix_reusable: true + evictable_prefix: false + capabilities: + snapshot: true + fork: true ports: talker: talker_cache_0: @@ -1251,11 +1512,18 @@ pipeline: talker_cache_1: input: past_key_values.0.value output: present.0.value + logical_lengths: talker_cache_lengths predictor_cache: + kind: full_attention sequence_axis: 2 layout: bnsh - logical_lengths: predictor_cache_lengths - storage: shared_buffer + aliasing: forbidden + reuse: + prefix_reusable: true + evictable_prefix: false + capabilities: + snapshot: true + fork: true ports: code_predictor: predictor_cache_0: @@ -1288,6 +1556,7 @@ pipeline: predictor_cache_9: input: past_key_values.4.value output: present.4.value + logical_lengths: predictor_cache_lengths steps: - kind: loop setup: @@ -1920,8 +2189,6 @@ pipeline: next: done - cell: accepted_len next: accepted_len.next - - cell: slot_ids - next: slot_ids - cell: talker_cache_lengths next: talker_cache_lengths.next - cell: last_frame @@ -1946,6 +2213,9 @@ pipeline: rank: 1 shape: - batch + batch_layout: + kind: request_aligned + axis: 0 - kind: invoke component: codec_layout inputs: 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..1146780ea --- /dev/null +++ b/tests/fixtures/onnx_genai_workflows/video/inference_metadata.yaml @@ -0,0 +1,850 @@ +schema_version: v1 +pipeline: + workflow: + manifest: + ir_version: '1.0' + onnx_opsets: + ai.onnx: 24 + 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 + 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/video/policies/continue_predicate.onnx b/tests/fixtures/onnx_genai_workflows/video/policies/continue_predicate.onnx new file mode 100644 index 0000000000000000000000000000000000000000..3177d8905c97dc120048da84e39b00bcae00c3f8 GIT binary patch literal 910 zcmcIiJx{|h5Ur!oCReIf9-ssXL4`o-5Vf!|@gbH>m zU9gO*mp+iysemhKpX-P+jtiY#*oq$Rd?n>jvsgsNXCc*^^aPJ-%z7hiUEg%J!TW{E zdJxT2#C3*%cLXy9xH2G@JV+S7phMF3YCl1)s@-sBnPwH72hO8F^mQFAQQ#}WtLciA z)P_n!4!r8@o{KtYINyn84cq|@6HbIQ^}-TGy&+iK6)b4P4pY(ynQ(9Q{4t>+$)z4P zR`O^`e~HUor9w_+`kJ{~pZco`701Ocl<-QI8rJKs!AunNk3AdU?Y-kD5lI{>*;xnPGI+6BoP7L kn4%_4e3}eJC}=!7ZjaF(Z2#GV-`;X&a?79EP^i_v08>02xBvhE literal 0 HcmV?d00001 diff --git a/tests/fixtures/onnx_genai_workflows/video/policies/diffusion_schedule.onnx b/tests/fixtures/onnx_genai_workflows/video/policies/diffusion_schedule.onnx new file mode 100644 index 0000000000000000000000000000000000000000..be69ec803c55d622285a3798f4b93bc2d11de2cf GIT binary patch literal 387 zcmaiv!AiqG7=*JS%KD2K_8_(giBc4b1igB&wchdyBFlEO+bpKLA-j==B1qn%J$d%& zd<0XYJqY6AzYO2ZKZEzl?%YgOWp6%wct-?Vo{6-|MgQ>LXzQ4ERGU=Lt26u~3@>5l zndOxjMmPx3Wa~*nDdlLr>b1AJ^}N^78w!#G*wbt-Y{9r7UNV4_;%OEetrwP;s&FyS znYA=iqGTn{YT`dUCK2?1QJV~ExhaI*5QL;k%E~IETea&vJ%Vi)0)R0h0r)s*&;^`! zGvk@i?H||L;=<&L*OVLWO2!@WmoE*vgTcQfRVnFxQ(Z{Rs)foGE9>!fgU%uPGebxz Nxcizg3BU`FKLOz1f7$>5 literal 0 HcmV?d00001 diff --git a/tests/fixtures/onnx_genai_workflows/video/policies/diffusion_timesteps.onnx b/tests/fixtures/onnx_genai_workflows/video/policies/diffusion_timesteps.onnx new file mode 100644 index 0000000000000000000000000000000000000000..328fdaf6198fa76d1d1ac7110c1233ede3c3b9a6 GIT binary patch literal 383 zcmaiv!AiqG6h&v!DC1MH%%W%!sZvzB2zBd9TUYr7kzqP{X$I4okeNt>xXDMjcIDUk z5vEqV5X85`Yhg`XhjM%9&FcCRL98gQ+#nOXG#H35xL<`w>Q$u=OPK z3P;DR)1{M*C*C8ZcSp$Ul(hFo!+0Ndj9j2Aq(V;OA)FMC^T-;rbgI;ak5n$56Pdzk{Fp2OhJ3`|p0{%gV{ I1kyeH1OwK4G5`Po literal 0 HcmV?d00001 diff --git a/tests/fixtures/onnx_genai_workflows/video/policies/model_input.onnx b/tests/fixtures/onnx_genai_workflows/video/policies/model_input.onnx new file mode 100644 index 0000000000000000000000000000000000000000..93a460da6cbb943b3fa24a7d1d6b092f84443eb8 GIT binary patch literal 720 zcmcJN%}T>S6os48#&}DqUIlAGQSk>{1nWY?m3AZc5h^lFC%2u^bP_UCZ4o!!_y}%& z5#PVRu_FAfPyy(plNDINE1t#T8Pw#p8f z9;@tR&zlR)GAwJ}BNKB$m}~CCcKp!wBdrHUBr3MP2)Qv#DoS`Fcg-EBeRHK5 z%!nu^lRDtd0nZ3YK|drY7)vq2BfI8SQo`!B^6WLKQ{Qu^<8cWJe4tD<_wMBXsTr=T zIg1RTR$~F~9L_ukPhkm*ZEl4G58Y0ZA_729NUk-7CRHV=uC%G~Nd;E6HfA5KK-pWu z J3Xof=egZ=1-+urA literal 0 HcmV?d00001 diff --git a/tests/fixtures/onnx_genai_workflows/video/policies/schedule_history_append.onnx b/tests/fixtures/onnx_genai_workflows/video/policies/schedule_history_append.onnx new file mode 100644 index 0000000000000000000000000000000000000000..de2caace0d541bb215ae412598925e82ec63fd97 GIT binary patch literal 829 zcmcJNKTpFj5XEs-+Vn~UHAs{op;8r)7$VBZ#GkD@3kX@Rdk{mgL+n&7h^1eF39&Hq zDM?cZY9(Om<~u(<|L*ymDlOgXu1w6)M+sgCUFx$lt&HW$vIB29KP%lscF=Ub8~ZhY zJ`WRQQdx6M-buORLPha=G!Yj_7dkE2itcWFt<=DXSVqSDglU zGpex%GeSpb#tL{N;8>aZT8vpT7-p_`X^6=25$fW literal 0 HcmV?d00001 diff --git a/tests/fixtures/onnx_genai_workflows/video/policies/schedule_lookup.onnx b/tests/fixtures/onnx_genai_workflows/video/policies/schedule_lookup.onnx new file mode 100644 index 0000000000000000000000000000000000000000..fbd67690d437efe41d1a898b5802585ab8260161 GIT binary patch literal 607 zcmcJM%}xR_6ou(zgqceSS{M8n4HDFZ#RRsl{8^ar2pH3pQf4Yc+e|wIV%)Ix0o?jP zK9mfgB8hS3t~d9)r{|vQ6>)x|1|l_=?=HL|@C=LiFpW95#&V5eO75tQctSe$(&i=A zkdsa`zb@NFaOq5#YVHj=y%PA6PPoyOah&V5;6#tZKuI|>ED_oUET+bgNbrOvESh8I z{-M`^`VN(85W|$_CL`bl!Il&3S+V?M>91bArN)w#qi)m9inb4qRn@}(xxwYnkVvf5 zIMw~L@7!<=DjvdzaIlSW8QQ2{*@Rk#l2^b!6!6}FS{4nC;A3M;!q4g(8#1Pju)LJB<=}g3eCiAQA0-ZzS SR}DTIab{N+yKy16U;O|hCcUfx literal 0 HcmV?d00001 diff --git a/tests/fixtures/onnx_genai_workflows/video/policies/solver_step.onnx b/tests/fixtures/onnx_genai_workflows/video/policies/solver_step.onnx new file mode 100644 index 0000000000000000000000000000000000000000..f53af2ab7ccfd1231e307da8cc93ec1b3e4684ea GIT binary patch literal 6366 zcmc&(TW=dh6!ykW;*66f85EK(w1idRA+(WaeD4dP6p$*)3*sdRjW+fq-qcyU-nB_V zNL3Jsi-bgZKou40Z{d|c#W^$Hn^TmMXu^ZFb7psEzCGWW@0=MIOLXq2zvqpD+kZ3i zhee7(yDK}Ro}>$7@oCo`NKv*Ib;Iv>z;b4lwM9r?`kf8yxiz!3NM?i39flh@uZfA!N6@xn!9&{Tpv6**zWuOanK%mgK)dubAv#1JvnrT?d~Hw^Z6I|HiL-gNo+255 zbRJwd+7sL+cu^p^lZUn`Fi)~Bs!Y46l5)`s;Er6RSTqIdJBt*8P!2NAsi1S->vTlb zCU{ZKsXBRRo5Jy}MWQ*ep_oAw_dQ?KY=ResqUPkGZ3@NHEJgKcis~tf*0WK(GqaQ5 zDHx39-&5;fb4JNneB_2*ITQ_hY0^?&YB+P++$xcRd+Y^9*Z$8IxFSYEi=YE;5Q?T< zQdfnj=@ih|S|BrF-)KamPDThaQC+7?WAT|EJRC{+l@u*|B{3Anma`D|n2NAxX%?gS zrlr_@tZZLNGs()AO64yWNn^o0qc~*DUQEbj?9$l$Qz`fa!xWuY(sVLzug2_Ytzi5T z2nET2u!?e7HKOdU!rY-rciodQ~>+%l2TCD1qO

e*vnK;9Exnnrm698g$wmz%t7#tri>QVi)Hzl~4(besDW;Na#TyAn%?w9S zc3)Dq;>{OjE8a?z)G|sE1Pri%S&G-=R*Mo}9)|*sB|>SJ;@xZo z?(T;f7oUHe^~wHmnJg9O=|YTT_pQ_s(^aw#Fkd07bU8Y95U4$@QUovsGE=k;?M$W5GWXV!^MUj$uV!AU&^c0o;c_rkbgn?BKUd#VCk!vTmp(44 zryvv#uVwbTIINZO;^Au;eQL|nRO4oDmZwXZH^?=3{c21CIw1W!yq?{Xoto}-C616d z1Y@V?W2dC*%n8?3CW?>pv@b32rM?lRC{|TY#vAY87o+a68a;FPn@QJ62moNrF79dhH$j}TsmZs30fx8OrvZ(D!IElhAb?8JJTZ2s#I PdP?+1tv7r)^z8XjBtS5vrPQbS)C1fH+VIap1;;A}@6pmn^9gw`?FT zJ90$`aoaxuLgIhnZ^22|)~-5n>efr2{rp}&fByXZ^hKOI4h~sAdh~YGbWs{K!kfwgSh|ZB_E`E7_|zp;FDt-^@u}psQZhY*__cjC5YSeM5L@U_TFhaYR11ca7z-ot)@yNgqv^s{jMy^{LH#9yK#f zah~~k;rfMOqA^e)P8`<6tYS#E;Sg@`83lnkwvsLEm2$(Oxd<6@!lJnf@i!7I#e^nE zFm?ieU^~QdX*n^OZ>$S%Lig*Abax(>#?)WTo9mEG2sc4TE2x!8iQhTO%!~7{supU~ zTBu1{xT$DCo1+rGNJ_92CDd?!V87~n_Pzn}Uvk>}rL?fGh@$#5MfGc^sH#v@A5qkh zC~C}=qVEV7q>5ryjFpkzhbLrJB$CvclAdc#`6(%>0xC!tNoiy>@Z$IyTb6M{>|tC{ zEbB^3${z$W7QC7i5~-#HeQs=wNLGr9RlV-2f^QF4n+Em-vv@ujXT)P4Y54Z#xTSak zk9C9_kjL3W!W|a?>yAr&pL!82usdPgJ%EmLE@ek9kJo=-ZSEd6VMejVChUl` z+mp1Ko&ily&C2isW>JgF@Ia*4kyE6T6arz@ynwNQ^Hnne(%qOmOBD>D>vlT%VW)hX{s!=s53U${;0^e7r)sVX48;3NM>047*NOEf3O;xfI zx-)bi?hc#E@)7k#o5xw;mwSQ7oYPn)JR}as+NaNF=qYUfccsnRc=KY_FiDy8PngF@ ac-pki&^_3?GzaZ;sKb*E@>vE_E5(0kSXTT1 literal 0 HcmV?d00001 diff --git a/tests/fixtures/onnx_genai_workflows/video/policies/video_decode_chunk.onnx b/tests/fixtures/onnx_genai_workflows/video/policies/video_decode_chunk.onnx new file mode 100644 index 0000000000000000000000000000000000000000..b493c53ff4910834340002212a11597128b83700 GIT binary patch literal 3115 zcmcguOK;Oa5RQ{3ai$3^3sr-WP>Cv#dWf<~`Uv9Df)Jt-Ck{v;v~ulj?bclfKSC?S zr4k%KNJt0?E~q#DCcEpzv0Wlg=H%Jg^~~pQz8#M`k1EIRo*nvkel+2mg6bpk$Yp+@ zvw(E8h1h1UNlB-rP8@5O0Y9xA*jnv7AJAcn9h_SAQ!6aX?AQh+>%7d;knFAV zGK*DS{C$mz8D+i*>G7LTcgTiTj|>TEqg9rxydZJybA637ry;>aUxKa|7WF9E)asET zv2C`>a+MdhvG)H$S4V0;bjTLhK7U9!TU`60lNXo|isVD&6pW4e5IA;BlMFNv4~b*Y zDsek69N!hxNHlib(H19-JW7vexAr`5;g^)Q(-4swv4pWUAJc?!t~`BKAkLc7bw{8P zqQPBlJ{Aw|=Bl2L5>1qKdaot`dr4Y&3;&Pu>0gk~>?WS)WbiBg2McOl!oN&PmMrM;1B1Zu-Np?M14XGTnQA-+TW0qxbM5(z-TdmD<;TG5ki*ikK7Hk*$n=lq?qX z(C<`%HesIinq}4pXb2s48-I}3??BT!mbk8$XQOS{((GE;m~lZiMhEabzM2L`>zU<= zie11XW-XOUBrM@_PQ1@wM*HBesEQucauW+%DBu?&os6CvonNIL9r!!TxO6t?;oF1$ zV}OiBsh~1b3vMW{|WqdZ99UDT&8fL?hRmVe{Sl{yx}us^T6*x8`|)3l!$oJs-#D!1|>S>vQXiEBdta}jR^^^dR(d?n%TSg489(&Jn(7j zc388#V?#Mn)46A%t!m#gvTuP(1<8t=6Fy)8N!?Gl;EGe9e`L|y^L%8V!l{Mu5ma&c zk}B5&3s;6wyfXl><$6>I=F2$iF+TECT=;vZ}O6n~ta4=QNs~b?4|1_W#<04TsEC M$@;x0fw9&619_9^MgRZ+ literal 0 HcmV?d00001 diff --git a/tests/fixtures/onnx_genai_workflows/video/policies/video_latent_permute.onnx b/tests/fixtures/onnx_genai_workflows/video/policies/video_latent_permute.onnx new file mode 100644 index 0000000000000000000000000000000000000000..a21b3400fe68f2e59fda13048d9a89693ee0e66d GIT binary patch literal 488 zcmaiw!Ab)$6h)hMY^P6z5?yE)Ql%)k7}SM`tFgH1+La<9&Ac|jX%doDJBXY9h8us% z|Iu+u7cJuUa^B(I7jM(%)Qr^3-+VRTgFrLq0d=6UMK_%V%(|U>$F;Y{Bhw~Bc7EJm zC1<-3m0(&+>nTYaY4k~2!yRaAK1FZ2Kw2C2;m|%$5~KCp3#V*IM9#fuQlaBc$b!}% zpN2i?u2H3fd0{O2lAwDpDnrAnf^6CEVX&%v4OcOtm(ZeHBOZi=m^R0*ECN7dA-UF= z`!bcNCNh-ZtIC2L-N9-u-LS0PE{&U(1v!PI>X*p`_0P7Tj83e{l_;1nIyf!@t-rl5 r$Sw5$C7H_T!IG*Z37^eXuDC1Giv>A>-k%N}k5Po*w+$kwb@smj0F{~> literal 0 HcmV?d00001 diff --git a/tests/fixtures/onnx_genai_workflows/video/policies/video_latent_unscale.onnx b/tests/fixtures/onnx_genai_workflows/video/policies/video_latent_unscale.onnx new file mode 100644 index 0000000000000000000000000000000000000000..d72a843f8e13836a35e3b9ec1ebd3c7b50112f10 GIT binary patch literal 824 zcmcJNK}*9h7>1WQYyDJ|;6!(jIYhxj*>(`|qRyR$zd(_YrP&%+lhR~!L%i&wcR{aS z{A2zS+pbk6h#>6t^78WL$@hKs61HwtTlDquy9KWZFAvDMlG;#dNVT;3rKNH?saEXF zI4Hu>fckw-I=-T&`0@CB@`g}r1r{Z}y3q@cY=2NYSCwNFh z=8bUv>bg~i(vMW!K{QhV*Kq<~5KNJQI+}B`Zs}shHihurs+$ zr^{t@FGeWo1j-j|M3|B$q|D&_{lgfY!S0_)B3!(crY0p0?GJ@7XgE5pjnN*If9>F% O4reAg#YO>gYnz`*K?h3! literal 0 HcmV?d00001 diff --git a/tests/fixtures/onnx_genai_workflows/video/transformer/model.onnx b/tests/fixtures/onnx_genai_workflows/video/transformer/model.onnx new file mode 100644 index 0000000000000000000000000000000000000000..b28bcff1be8b640711a07365c2d5aa90eddfd970 GIT binary patch literal 5197 zcmd5=U2hvj6!kbxycs)*JCF};OB!ukwemycT|2SO3nla^q7q8qg3xHYp2T+Twb|X6 z(jNeV3WPuiA%qkO35mbK6aNbC%zAfc)-muh2>W5@dgtTJxp(e8W25r+?~m}02(N|i zcsgQaeMX=4-6^9Dvl=)2Ne?V*RakY5SD4@1L)WNSuj17(a)W5E0M_XSUiaNG3#YEf zNa6H8-k-j>Je>Ic%g_r3)9BC}xnW5A0~WY}*Z-6hKm6$Q0Gqj9d=46VlPL>%55PvM zGy2}dho$_8+U8c;7RYSNNPejAVmxz4SBwt)D140K?=>px!gh^t#urv#5-dP_ZbTh( z0}%3-0PI+6v01O+lKXNHYNhWBgpFtbZi(2G7=XzQBib}8;wkVntr9ot>$n(AgkyEl z^@HllR^F8@$(7r<61(zZG;yQtLgCeuYlL(m<18iID1K~_@{B$i%;=Gc`9m;vWUWA3 z7nW}kPS%8Q=8-Rgh*W0ujD_NC9-EsoK3gd5lB0)xEC*y6kDVNk9ql}RryhZB{s?rX zBXCQ}W{_jG<*~1j67*f3ai4vB#n>xG?S{FPw1p4thE<;bpv!|lW%US7KwO4w+uV=} zl5OWIQcC|WIeD%e?Oc72$cD%g^ZM8{YmhQxBlcHwX-ol2+G6T! zgjHM$QIPx9plWF+>K75@$*OjPhkR-|iq+!P`^@z(A^eZbjkpz)d1S2$Q&(`oQIyu^ za!gce``D~zCEGs!&uaFsx$tT)gjZV%FH1>T=<%P)N52+>te=skhF+V8UQ#gaj=7bz zW0&gB`k-Nzv(Q~2QO^4d(&vv8jnd~EaV!!>?2bxZmup5Ye4K_^Nd=J8(5irPbUAjO zE=Lyq3v?CU(MZ>qX=zf3a%@i2)2k769J9ioBDNgIS^-02vXE6FNai%mI~>ouy&iQy z?DJ=gr)jN#p$pIFLPp1{azPyjPKu-@BWYeb928w8(j9I<;#3B=tcc=Te?m@!1s(Yxx0Lb zKz9CO$<=|J;SXSI@BhPi{E$4j8Syj#J|Y@2kgC4NV0#D8e}?**Ihp+5_#5si7bc7i zV6X(pi^E;%zXX!KKtodz=eG2&#HW+4h}o#s;hWiwmKb;qh%Vy7a7US|Mb>Z$-V}y! zq<2oK^IbWi9FrQ}guWOTlJC&UU$mq?MFYXM&`_DABXyWhliOZb+8&DGrXsB&;FbZd z*mkPGvu)=`IO&MnO1ie2f3r-?3SVYN(PEs%vS4vrk{S|Uxi#9-8mq%cocXNGl0Ccw zY!c29@HuBoFt)2g?WRt;FXhyKI~)pA;^TcjJWh$QZVgti3QzK^0L~cRNfTxB1~+AI MTZ?Z`V20Y>U!>#9Nt}4Gr^wcNq<4W|4A;+|^ZOn+vcjE)gps_O=MMI<|LWuVY8C z6NWPv1ZXdOX<%VTNQiGCggC%~{h#;?cvb!NUbQPbAj)254r#mG?O*-gtKX~lUb(sd zh4+^1_W=tB^VwB0NC&5*BXc>Pv0Kx0Ihrk#MS6BV9fj;Fd1F2sq|2n!x|Lte=EG6a zY2V0R9`3W9^yzqUfBSi`eSC|(G)vD$i}Q3a3b&3Qu?OexJ~^7tW=|J`i}CsLXfRC| zi{x}Xx=1eur{4~@zwyoEhpc6NNgYc$d4896TW(DTO4WOC_RPlL!FlO(HD6 z>e;mQoV0;(d0Tp*vPeSFb(d`s9(&r+zLpM?ev9cp#?kui23#M$%y!N|bK5(k;giwx zpyA8&-ZwTjAhcC;X$9fO+S=Kexc)c@Zy|AwqSiqnrWsvPZQ z(z4o&GHhmI9jVrS8Mh+kHE`KPb!$X*okaOu2cW@)wo+#otYnD}bkQp4(N)m%yb8Xb zSv6mQW}#-+J#Er24~I$AZ|&>9*{XyP0WW;5$!An@dB-*>?;w-jVf)#R_MPcGUEbN+ z+WO-2An=9l7XdLE?;x^b6DiWg@~h)_M+w)FqIh9w<2q90@A$&^gFw4iW-)`bK)@qy z$~P~k3FobS{WrtJVc&f@^##XTojk6YGDfCcUvhGdjPbaTF)omdaj{Mr<2AkKZoT&@ zV}C_5mg^Q5&PdK>cy~~y_sRScQg8ioJR^F>m1jNsj#&knvsar+4ygj#RPt`Vsr2Ri zzeBN-cNqgJ7Rw|WCLkK+b%d@vRRrKP0&w3Ez=;*(;|{ESYTdqP>pHPweDP|mu}&1; zPx#dFf1w>7?^T-?-mk9X4%yqp;G(dfH;W~X+(X!XFK;>J+xDy}n-s=b%L;VUxyH`d6a?n-OmA^?--$Tc(-MGwTNXfBE#h>^eUqD@rau5yU}Myec848kRNktw1fr>JpA2qI;v zRzeW**kiP4lC@TfbhlD`*4lq%wuz?8NQ1w`;X+#BLasMl*yPm`9_dbad?pruHeoE^ zrH#QrBuWleYbBvjVQexroS+BzGKP{M_cK?b9eT`7s|O8Ie-U8!s}k?v-^ zPuaCSuoe!aI~@4*z&{~-Qh!rs3t(cYlynF+TO=*DN#CV5o5xzTknU*VQwA+us@8T< zxb3o;TD@gP)hXhbg50BwV+wMQ^~Et&YrU(@wBBLYUJ4uab`$tDJ&Goi3U%HU?Yt}B zop;=Gl&B)5Oq9?aECo@b>^Ulk66@9{1V_Khzb!&3pb<4faaQ-)}V&GR9FV0Bb?S%KLKr%WV1oHq8;g#cfoqH(H@975@o}@^8WYFO9oLvVZtmt|Z~p}{ zM0)o6gMhssG=mVTwVw%J8O+bl=QB8(eS~IrbdoLyr);~q!ybp9Tm6yzq3>MibJ-vH zvD4A`$?1|k3ftE{?OWsF^7L(RkQHahLpUq?a}6Ad{x1OB`Bmm%dg253KOB`GNIyT& zz<>4vz#DFk2=LE`k5@Cv4|B`pl)n}}`iKu2(ocR9uw2SRII#IlM8HA6$i=6y3U@kEW{~z3Ba-VjGS+ltp?toGRWCUu(juF_8I!0pWh;WWb zGV)Nzki1aQ2}5=OK712O6k5l&^#92{^os>@Ym1zkrXcTBr}i;+dJDy9Zqad#0;IRV z8FxxW2_ACSa%6i#@9LqkVDd_N1~EKfw_(04o2i*C*b1Vn$&Mp9bx=N1--gRJ5jB@J zNXd9mZB3(i#dlgaQ(CM+O3q!q$JCrflo;zK%uO)CaC$4?yjpBrqyCD~dxQqw32(3+ z;86RC=BL=@!pXW~z&tNMc2p;vLWOf?nDcUe{5n$LfEB(8nULd|@9O1N5*#AQutWtsLl_RM>(^Hl3_hfZwNv6Dznex<n6+&5jUur)Zvuk zrdGKQ6*nbI4Wr_wQ}|r8oQrPguL+I7#d2gqR7Cq<4^(W52#dAAb z=-T99cYEx4r~Y6%ngrW}L&x{+y^%Y0rh~mRT>9{%olRsYU$Hf8e=wf7KKlYa$p=Ct zAS43PKBRAdcH5{%Esm^RUkJNJ!TL%LvW^;L9U;h80CF5;g(m0s8mdKw?9laggFRmgva)>xKq1p)Z!+$Shrq@`bC)xs++if2CCppR)9a)bEmG|GnT>z2idF7^IlO9A8Rb+ zQAH~`YT2*0YZNCpdq{e`DyVj$9#3yEY&o~!RV-mYsC)2N26(4=ru z@NQ8~{ZWG=Z^wH}^{pmW7DA2Hq@39yHadd}_bp_tj3X^^O_s7zVIy!hh&7H_MHe|Z z8q!1*3`ejGxRo{BERKdpyr3l3QLNZvMdWQ4Xc%c5OS}>Kv|TT=UQyxpB{y2MF|Ij5 z;CjJm?Ab?-e_$-#b>MIffpqF+)@v-HQZU}A6iA7_R_aDVbzaHIvMId2unBWKtoMOgKZtOfumN znPDdVE@qNxVZ`;e?H`VVy!>fBJ%6I~0rgf`{wNI0Y-ilw77FOh2)`PQoeV5dhsbrk zNU{5!P}fbVuA3bu&SQydc8%q@9hOUwAqDGQ)s?-B&3bvREV8itDwAom&7@mw8?!4L zvB#8e$Y^O(V$&4vs0wA%O32ZS7NMdRnO(?%x)f|8nTslXm6{p8O3e!(`qVt1=u;{a zeM*Za`oLLJok1yHnvDGsC#|EQFA@W4&5Rg)A&QTLr_h!Jr!_M;KXEvmm(Z4^g;qI7 zku0X_NZOu3c>N(1U4mD4T@P-b19xagcR@KPEbaU+PFS_BWF!jBmnd+zT$n^rab=DR zDjpZJt9J@6D|1z4iC0xJF1Sb91 z^^#+@NV=-p3tl>6F&|#lSgl-K7n4gCm6OMkyDdTTS#2ek;B91vlJ74mm)dwen@c>Y z^0Mii(`mNZgs%?W0lfZBZMfYD*h?WWBj<4Rz@0V>c7gM{@M5;3$6_=yMt1anqn~@= zJ^8`8jyClQei2pheBTKMduSCeM_>3ZyuWz^_uh5&#RM;-D$H6;X7yiF?^5tB^a_0P zQt}Bh@CNrlyzJ(!&A8czcS(`P;AcPuTb_4EgKwZM*mpDCcMV@w{vg^i}J$w z*%#U05LjdrBc#LMT!!T}X#^flgpqYoF{A^fePJMXhAAy&NN;a2bU>^7YfL|3f5rMm^$PIHaq)u%mQ|Gt?xi{h`05)+E zJj*6>3npg^=_ac(xAP4n21G#OCw?#O*K@9K<8Rp7eAD>iR2+~uh5NvGP~DW>$O)Uu z7B$<6MU@3R!YnQHcY1M=Oj1_sB?E$gXjAZyybFl<6q}}cfSQR1vd#lR8&73+;(-+j zx%};9*C<#-=V*#kfylWHU|D#8ukFEmgWyfP#p2?D1X@LcItfA|wO|Ar&FZ@w&~XiFa8$ zQBuT9A9w;JXt@alMZjO+1tBEf_*HObe48`Q3SR6qFY)Yn#-G1Amv7FvaaNz%A8rp0 z!WaL3lYFd^h0%DpKZ=}ac+U?)eeTeCd&eF5PRp2&FM?qoTFq(E+FBrUfxGXABiHlw zsjX||%IMzR&0!E6h2D5DiZ;DHHw>Mffj@S~-p;XJd++_Nbz&%Y@j0B(8;<;t-$5Qv z20{@KngH1#Llu^JRb?m#5KX$stVtOeU>o#YTc3G4`bd#T7fGRB}55Vv4_+#Jc8T0Xl0NHC!i&j|> zIf47`kuv~>+-LxffZPGRHg!(gGR~#f0vc_V?x~+Co{He;jYH=~zwc0LoaO&xP_%gl z8fDS-rN$#f8cZE(8cSIte!eNWL3OeA2aRNG@aU0&qNctm2NP-;%UJ|VD6=XEYQIRy zfD#1~_fkOYxe$aQ9VC}ApR73(w>Qp`x_dMTON}Rwl=X!SkWm@qv8DB`EX=8}+67qc zNwCV)9Y2Xol@UbkNk-V_dFw4Vxa&K#V=Tm1F`u30v}ji4vE2f7U5UEtJQi|;#0mD< zuHv)3qR;j!@Y!!6H=x+jgN=?8gsj;}e<_PDCE1&W zWMIn40EZ3;yQOWcenz=e&i1A)rBbtvQ>h7CgbGvXZ^8wVXkkJI5S0ce+#5#5#WZKU zaYrgPn4M;}aZNfl3^FZt5*v#RZE1F^F_rv1*6L*~yOM0pLdG!Vj8T~eP_412CNi^k zDv1oTgtW~^*4X1^iFD6ksrv`3{QZNnlonMvv#7)wTFNqY)2~=JnO(kV&3Qs%z!H1@ zSTR{9(oPmsoT}7LmQ|qdbSg`-cCt)NmxTp@#R>wzVigF0uNCcNvD`~ql5(K|hp5fYq){1biZ^IWCTb6tOx8PNyb8UcT}BPErM%kYTOf6Omy*Ma#Y5x!#VSAysm@SZ!ieRGU=dn3}j4Y#&k$svLg8pR%{xRKHmCcH#CD{M|NB+u!xI zzVE@e3}XjQ$0L634Qkij8{GHD8wM{b+?39L^8!Pmk8U;oN&ap_aK@icG|75nR=-SU z_33Rl@^;9welEEq^x-r)@ci3X)$>3799e)-bJ?idm({~Kosb6HQYW>Bw*5k(Qg(@~ zLGMa70U>(BbKx5esFQ6!Pmr|NPr5k;w{4P>s%|Ha$sX_=Dav8}JXr=v3ppf?{3?;E zKPTUEjTBXEfuQYmg4JU2j+e;E`AC17Ty1FjQ!=oc%okAH*nXKzvhV*jes@*cEwtpt zLp`vi&zE)gqP#QAJ8E*r#1ST%%0Q>z7V29;A3f+Zqz>yewR?$2p8Q16Mq@7UGm4~& z-PRjq28Ptg)b0iJoK7~`6?EP$7H3Yf@7)!tXs=F})<0+~n@-Qqld|6g-ZYrQ0W87` z{v*E(#SECh=YI5w8Fw9j8{4x$Dhhyg)`IDc;~eqfSlc#SqQOKyEKC1y8! zb7jJwG)jvtBV9@>UrIb16tzNE1fDntYBIciqR=~*3Mjd7=fZcwIPfSiytkQGDCx<%OFbj*ZLZHA5R z0qpq~=M9hZsjvgrmkoPFg&+f$1vIb+S2*+GW48*{ovg415w_eLL33nl&>H)Jp)DbTnbau_5_X$hYnVBQ2$lc&v}oQtFrm(!%*`$wXQYF^m>2b2gPE>9eFT zj+cZH`YSEKUy_b$-*r@b!BIO(=CPx+>C*!B9VJ}UhJx=2SNA2?5xND$jCSr`awq*u z{-)q90!mzQ2bc5}FD6)U>Nt_@U)I*tI3?n6*z)^W2PyddILz$9NAOs@%#vw}>}qb*s;+&E8X3 z=4Ryv!}|*z6$6;6Y!=kcLo1<-X5wEZ(c|W;GQ<6q-Q#w#0miIwEUQ-H-)>^lHeQFU z0%Q!Mc)X?)$FI>7*pD98#yR8ai3_FM@V$|1Uy+cIfZ7sgK7GDMr*QD!OyWBD!7M8( h(x8*6J924ob<(^>kD>a{8m!-Ew$!iZ>ns>MmA^m2fFb|@ literal 3028 zcmc(hK~EDw6vw-T((R*#GKkboB(M@S#5CO&3iTuyFO4USF~OM0ba!aCl-;E}C7|)b zM=*HMgW}mw;Wu-3y1R5*(9=p!^Jd<@`St(i%^P(V%^bNcd*DC)y$qiflt01GUB_p{ zVYs5r$79E}D6Z&gG_g^H+zIIoD7GC|N?j`HW&v^zIikKtOo~#?by)EZ+hx~rPJOd) zd#r5si0|W$P5Y#8cFs`x-TUS$Xk%VM4D6e(NBzJ9ekw3W7^4oeT?Zc#{}AO)@aqoo zC^oeGa3JsuJtM3QfGO%&r8KA7$U}zzn?5Tk2{K+NT#bkczzh_`hJ`C=mQNoKDE&Zj zRhx^30=KGXLpv$y#`0Mk&u8tXd?u;Yr2^6+gHeg3rCN~aSHx*kys71f1EJieo)K13 zx_uMr40o2&SvbLO96wCxmR7tn1lX3Ii)WI8`@!p>=eBK=^l*}ACpCGLN z2f{B3NEBPLSr5|s;&{OMx3&CmAk^E|Gr~$rxv_HC883&O+bDs68-`YljUjD`P?4bA4O1G~4n8b0n)#<;;;WWUeNk=CAx5jv@)(h{C9% z1(DiO4{dhWvck01^eJv6lvK26QIo$v`q*dG!voJEjBaSEQq1m+^xCGH3V#!8MpIS7 zwf!d(r&fcKnniaYi!v?3%nsyGmbY-m&+`7-_+^Xcp}@U!vG@K1Sva%^3*4a)J3v}h zS_WKK8|W`vw;sc!=GO-AmPR^21p%mh|LI^`X~mFM?n(~F^l^}L_BcTPDS956g zO72uwX4;9FhV#*(Zp8D;W5qeP5PiU_AXMpGc?4_2cXPQ-9sZyVD!J^rJ=;9P!P)5( zlcDs-Pv^=mtpC#!+ZMVPdJTh=$>7xP*`$B=%s5vb!s?AN=)nOsnRqMjO+%_!`U4Jr BuulL0 diff --git a/tests/generate_onnx_genai_validation_packages.py b/tests/generate_onnx_genai_validation_packages.py index 4b9c28e8d..1935a638d 100644 --- a/tests/generate_onnx_genai_validation_packages.py +++ b/tests/generate_onnx_genai_validation_packages.py @@ -710,8 +710,6 @@ def _adapter_package(source_root: Path) -> ModelPackage: {"decoder": model}, adapter_target_manifest=manifest, adapter_service_options=AdapterServiceOptions( - slot_ids="request.slot_ids", - request_epochs="request.request_epochs", active="request.active", max_adapters=2, cache_max_entries=2, From 8c38757e71261309c811086f11b461fa4ad68cfc Mon Sep 17 00:00:00 2001 From: justinchuby Date: Thu, 20 Aug 2026 14:18:31 +0000 Subject: [PATCH 126/151] Execute the integrated workflows against their coordinated runtime MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The E2E producer work merged into this branch targets a newer onnx-genai than the ref this branch pinned. `batch_layout` — the structural batching declaration that replaced the vestigial `slot_ids` surface — does not exist anywhere in `2af34dca`, so its parser rejected nine of the ten checked-in packages outright. Only one failure was visible in CI because the validation step's `for` loop runs under `set -e` and aborts on the first package, alphabetically the adapter one. Pin the coordinated runtime instead. `90e43c62` is the head of `justinchuby/onnx-genai@justinchuby/simplify-composite-metadata`, the runtime side of this same effort: it understands `batch_layout`, keys adapter composition by batch row instead of a serialized scheduler slot, and requires a declared row axis for any component holding per-request state. Validating against it turns the metadata back into a contract that something actually enforces, and it exposed four real gaps. Declare the grammar adapter's row scope. Its three actions keep one grammar FSM per in-flight request and every port is request-aligned on axis 0, so without a declared row axis the runtime cannot drive the mandatory `compact(selection)`/`release(row)` ABI and the FSM rows would drift out of correspondence with the sequences they guide. Declare the grammar effect domain. `clone`/`lookahead`/`commit` are a speculation protocol: abandoning a rejected proposal is a transaction abort, and the explicit `clone` is exactly what makes the domain safe to enter speculatively. Both facts are now stated rather than defaulted. Stamp request alignment onto the adapter selection inputs. They carry one entry per in-flight request by construction; reusing the existing `declare_request_alignment` pass also covers the declarations a producer wrote by hand before attaching the adapter service. Move the conformance test onto the row-positional adapter ABI, drop the retired `package.slot_ids`/`serving.slot_ids` inputs, and feed the diffusion workflow the `request.noise` it now declares. Finally, state what a dynamic KV cache costs. This decoder concatenates `present` onto `past`, so its attention mask is a dense carry that grows a column per step for every row at once, and preserving an inactive row would need that row to keep a width the rest of the batch has outgrown. The runtime refuses instead of corrupting the held row; the test asserts that contract rather than leaving it to chance. All ten packages now validate and execute against the pinned runtime. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: justinchuby --- .github/workflows/main.yml | 2 +- .../onnx_genai/inference_metadata.py | 5 + .../onnx_genai/workflow_metadata.py | 18 ++ .../adapter/inference_metadata.yaml | 47 ++-- .../speculative/inference_metadata.yaml | 14 ++ ...generate_onnx_genai_validation_packages.py | 25 -- tests/onnx_genai_workflow_conformance.rs | 233 +++++++----------- 7 files changed, 151 insertions(+), 193 deletions(-) diff --git a/.github/workflows/main.yml b/.github/workflows/main.yml index bc4184aa2..4a1dbfe6a 100644 --- a/.github/workflows/main.yml +++ b/.github/workflows/main.yml @@ -23,7 +23,7 @@ jobs: - uses: actions/checkout@v7 with: repository: justinchuby/onnx-genai - ref: 2af34dcad0e429604ca66d7ba1388ee2e688756e + ref: 90e43c62e6051a9a7074b8b8e96593d43d0d389c path: validation/onnx-genai - uses: actions/setup-python@v7 with: diff --git a/src/mobius/integrations/onnx_genai/inference_metadata.py b/src/mobius/integrations/onnx_genai/inference_metadata.py index 309235a5f..4f946e440 100644 --- a/src/mobius/integrations/onnx_genai/inference_metadata.py +++ b/src/mobius/integrations/onnx_genai/inference_metadata.py @@ -1740,6 +1740,11 @@ def ensure_input( 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 diff --git a/src/mobius/integrations/onnx_genai/workflow_metadata.py b/src/mobius/integrations/onnx_genai/workflow_metadata.py index 19558cb16..1d01db535 100644 --- a/src/mobius/integrations/onnx_genai/workflow_metadata.py +++ b/src/mobius/integrations/onnx_genai/workflow_metadata.py @@ -214,6 +214,12 @@ def port(dtype: str, shape: list[int | str]) -> dict[str, Any]: "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}, } @@ -5699,6 +5705,18 @@ def build_speculative_workflow_metadata( "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"] diff --git a/tests/fixtures/onnx_genai_workflows/adapter/inference_metadata.yaml b/tests/fixtures/onnx_genai_workflows/adapter/inference_metadata.yaml index 8916e4c6e..88b1c819c 100644 --- a/tests/fixtures/onnx_genai_workflows/adapter/inference_metadata.yaml +++ b/tests/fixtures/onnx_genai_workflows/adapter/inference_metadata.yaml @@ -13,41 +13,21 @@ pipeline: - parameter_adapters - heterogeneous_adapter_batching 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 + batch_layout: + kind: request_aligned + axis: 0 role: kind: runtime version: '1.0' role: adapter_active source: kind: request - request.request_epochs: - contract: - dtype: int64 - rank: 1 - shape: - - batch - role: - kind: runtime - version: '1.0' - role: request_epochs - source: - kind: request activations: contract: dtype: float32 @@ -55,6 +35,9 @@ pipeline: shape: - batch - 2 + batch_layout: + kind: request_aligned + axis: 0 role: kind: opaque source: @@ -67,6 +50,9 @@ pipeline: shape: - batch - 2 + batch_layout: + kind: request_aligned + axis: 0 role: kind: runtime version: '1.0' @@ -79,6 +65,9 @@ pipeline: rank: 1 shape: - batch + batch_layout: + kind: request_aligned + axis: 0 role: kind: runtime version: '1.0' @@ -92,6 +81,9 @@ pipeline: shape: - batch - 2 + batch_layout: + kind: request_aligned + axis: 0 role: kind: runtime version: '1.0' @@ -106,6 +98,9 @@ pipeline: shape: - batch - 2 + batch_layout: + kind: request_aligned + axis: 0 role: tensor stage: pre_adapter components: @@ -127,6 +122,9 @@ pipeline: shape: - batch - 2 + batch_layout: + kind: request_aligned + axis: 0 outputs: output: dtype: float32 @@ -134,6 +132,9 @@ pipeline: shape: - batch - 2 + batch_layout: + kind: request_aligned + axis: 0 contract: id: onnx-genai.parameter-overlay version: '1' diff --git a/tests/fixtures/onnx_genai_workflows/speculative/inference_metadata.yaml b/tests/fixtures/onnx_genai_workflows/speculative/inference_metadata.yaml index a9e57845a..186d134c8 100644 --- a/tests/fixtures/onnx_genai_workflows/speculative/inference_metadata.yaml +++ b/tests/fixtures/onnx_genai_workflows/speculative/inference_metadata.yaml @@ -384,6 +384,9 @@ pipeline: action: clone effects: - grammar + row_scope: + axis: 0 + stateful: true grammar_lookahead: implementation: kind: adapter @@ -482,6 +485,9 @@ pipeline: action: lookahead effects: - grammar + row_scope: + axis: 0 + stateful: true grammar_commit: implementation: kind: adapter @@ -580,6 +586,9 @@ pipeline: action: commit effects: - grammar + row_scope: + axis: 0 + stateful: true speculative_acceptance: implementation: kind: onnx @@ -1140,6 +1149,11 @@ pipeline: input: past_key_values.0.key output: present.0.key logical_lengths: cache_lengths + effects: + grammar: + retry: transactional + speculation_safety: + kind: clonable steps: - kind: loop setup: [] diff --git a/tests/generate_onnx_genai_validation_packages.py b/tests/generate_onnx_genai_validation_packages.py index 1935a638d..816c44894 100644 --- a/tests/generate_onnx_genai_validation_packages.py +++ b/tests/generate_onnx_genai_validation_packages.py @@ -792,18 +792,6 @@ def _write_adapter_metadata(package: ModelPackage, directory: Path) -> None: ], }, "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", @@ -817,19 +805,6 @@ def _write_adapter_metadata(package: ModelPackage, directory: Path) -> None: }, "source": {"kind": "request"}, }, - "request.request_epochs": { - "contract": { - "dtype": "int64", - "rank": 1, - "shape": ["batch"], - }, - "role": { - "kind": "runtime", - "version": "1.0", - "role": "request_epochs", - }, - "source": {"kind": "request"}, - }, "activations": { "contract": { "dtype": "float32", diff --git a/tests/onnx_genai_workflow_conformance.rs b/tests/onnx_genai_workflow_conformance.rs index 4f40e17d5..16aa90547 100644 --- a/tests/onnx_genai_workflow_conformance.rs +++ b/tests/onnx_genai_workflow_conformance.rs @@ -31,43 +31,30 @@ fn options(max_new_tokens: usize) -> GenerateOptions { } fn adapter_request( - slot_ids: &[i64], - request_epochs: &[i64], active: &[bool], values: &[f32], - selection: AdapterSelection, + selection: &AdapterSelection, ) -> anyhow::Result { - let batch = i64::try_from(slot_ids.len())?; - let mut segments = vec![-1i64; slot_ids.len() * 2]; - let mut adapter_counts = vec![0i64; slot_ids.len()]; - let mut adapter_scales = vec![0.0f32; slot_ids.len() * 2]; - for (row, (&slot_id, &request_epoch)) in slot_ids.iter().zip(request_epochs).enumerate() { - let identity = onnx_genai_engine::AdapterSlotIdentity { - slot_id, - request_epoch, - }; - if let Some(activations) = selection.rows.get(&identity) { - 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; - } + 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.slot_ids", Value::from_slice_i64(slot_ids, &[batch])?) - .with_input( - "request.request_epochs", - Value::from_slice_i64(request_epochs, &[batch])?, - ) .with_input( "request.adapter_segments", Value::from_slice_i64(&segments, &[batch, 2])?, @@ -94,90 +81,62 @@ fn adapter_request( )) } +/// 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_compaction_and_epochs() -> anyhow::Result<()> { +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_slot(10, 0, [AdapterActivation::new("red", 1.0)]) - .with_slot( - 20, - 0, - [AdapterActivation::new("blue", 1.0)], - ) - .with_slot( - 30, - 0, - [ - AdapterActivation::new("red", 0.5), - AdapterActivation::new("blue", 1.0), - ], - ); + .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( - &[10, 20, 30], - &[0, 0, 0], &[true, false, true], &[1.0, 2.0, 3.0, 4.0, 5.0, 6.0], - selection.clone(), + &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( - &[30, 10], - &[0, 0], &[true, true], &[5.0, 6.0, 1.0, 2.0], - selection, + &compacted_selection, )?)?; assert_eq!( compacted["result"].to_vec_f32()?, vec![25.5, 35.0, 2.0, 4.0] ); - let reused = - AdapterSelection::default().with_slot(10, 1, [AdapterActivation::new("blue", 1.0)]); - let stale = engine.run_pipeline(adapter_request( - &[10], - &[1], - &[true], - &[1.0, 2.0], - AdapterSelection::default().with_slot( - 10, - 0, - [AdapterActivation::new("red", 1.0)], - ), - )?)?; - assert_eq!(stale["result"].to_vec_f32()?, vec![1.0, 2.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( - &[10], - &[1], - &[true], - &[1.0, 2.0], - reused.clone(), - )?)?; + 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_slot(40, 0, [AdapterActivation::new("green", 1.0)]); - let output = engine.run_pipeline(adapter_request( - &[40], - &[0], - &[true], - &[1.0, 2.0], - green, - )?)?; + 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_slot(50, 0, [AdapterActivation::new("red", 1.0)]); + let red = AdapterSelection::default().with_row([AdapterActivation::new("red", 1.0)]); for _ in 0..2 { - let output = engine.run_pipeline(adapter_request( - &[50], - &[0], - &[true], - &[1.0, 2.0], - red.clone(), - )?)?; + 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(); @@ -229,25 +188,25 @@ fn decoder_batch_request( active: &[bool], max_new_tokens: usize, ) -> anyhow::Result { - let slot_ids = (0..batch).collect::>(); - decoder_batch_request_with_slots( + let seeds = (0..batch).collect::>(); + decoder_batch_request_with_seeds( input_ids, batch, sequence, prompt_lengths, active, - &slot_ids, + &seeds, max_new_tokens, ) } -fn decoder_batch_request_with_slots( +fn decoder_batch_request_with_seeds( input_ids: &[i64], batch: i64, sequence: i64, prompt_lengths: &[i64], active: &[bool], - slot_ids: &[i64], + seeds: &[i64], max_new_tokens: usize, ) -> anyhow::Result { let bool_bytes = active.iter().map(|value| u8::from(*value)).collect(); @@ -256,7 +215,7 @@ fn decoder_batch_request_with_slots( 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!(slot_ids.len(), 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), @@ -278,10 +237,6 @@ fn decoder_batch_request_with_slots( 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( - "package.slot_ids", - Value::from_slice_i64(slot_ids, &[batch])?, - ) .with_input( "request.eos_ids", Value::from_slice_i64(&vec![2_i64; usize::try_from(batch)?], &[batch, 1])?, @@ -307,7 +262,7 @@ fn decoder_batch_request_with_slots( "request.min_p", Value::from_slice_f32(&floats_zero, &[batch])?, ) - .with_input("request.seed", Value::from_slice_i64(slot_ids, &[batch])?) + .with_input("request.seed", Value::from_slice_i64(seeds, &[batch])?) .with_input( "request.rng_counter", Value::from_slice_i64(&zeros, &[batch])?, @@ -329,8 +284,7 @@ fn mobius_decoder_workflow_executes() -> anyhow::Result<()> { PipelineGenerateRequest::new(GenerateRequest { prompt: GeneratePrompt::TokenIds(vec![4, 5]), options: options(3), - }) - .with_input("package.slot_ids", Value::from_slice_i64(&[0], &[1])?), + }), )?; assert_eq!( engine @@ -403,7 +357,7 @@ fn mobius_decoder_rows_match_independent_runs_and_dynamic_batch_replay() -> anyh .iter() .map(|island| island.stable_binding_runs) .sum::(); - let compacted = decoder_batch_request_with_slots( + let compacted = decoder_batch_request_with_seeds( &[6, 0, 4, 5], 2, 2, @@ -423,8 +377,11 @@ fn mobius_decoder_rows_match_independent_runs_and_dynamic_batch_replay() -> anyh .1 .to_vec_i64() }; - assert_eq!(compacted_row(0)?, first_tokens); - assert_eq!(compacted_row(1)?, second_tokens); + // 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() @@ -435,27 +392,23 @@ fn mobius_decoder_rows_match_independent_runs_and_dynamic_batch_replay() -> anyh "same-shape row compaction must reuse stable island bindings" ); - let inactive = decoder_batch_request(&[4, 5, 6, 0], 2, 2, &[2, 1], &[true, false], 3)?; - let inactive_output = engine.run_pipeline_outputs(inactive)?; - let inactive_rows = engine.output_rows_for_role(&inactive_output, WorkflowOutputRole::Tokens); - assert_eq!(inactive_rows.len(), 1); - assert_eq!(inactive_rows[0].0, 0); - assert_eq!(inactive_rows[0].1.to_vec_i64()?, first_tokens); - - let first_inactive = decoder_batch_request(&[4, 5, 6, 0], 2, 2, &[2, 1], &[false, true], 3)?; - let first_inactive_output = engine.run_pipeline_outputs(first_inactive)?; - let first_inactive_rows = - engine.output_rows_for_role(&first_inactive_output, WorkflowOutputRole::Tokens); - assert_eq!(first_inactive_rows.len(), 1); - assert_eq!(first_inactive_rows[0].0, 1); - assert_eq!(first_inactive_rows[0].1.to_vec_i64()?, second_tokens); - assert_eq!( - engine - .structured_output_for_role(&first_inactive_output, WorkflowOutputRole::Tokens) - .expect("semantic lookup must return the first emitted row") - .to_vec_i64()?, - second_tokens - ); + // 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)?; @@ -485,8 +438,7 @@ fn mobius_vlm_workflow_executes_complete_image_path() -> anyhow::Result<()> { .with_input( "request.image", Value::from_raw_bytes(png, &[png_len], DataType::Uint8)?, - ) - .with_input("package.slot_ids", Value::from_slice_i64(&[0], &[1])?); + ); let output = engine.run_pipeline_outputs(request)?; assert_eq!( engine @@ -506,7 +458,10 @@ fn mobius_euler_diffusion_workflow_executes_complete_path() -> anyhow::Result<() prompt: GeneratePrompt::TokenIds(vec![1, 2]), options: options(2), }) - .with_input("latent", Value::from_slice_f32(&[1.0; 64], &[1, 4, 4, 4])?); + .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!( @@ -549,14 +504,9 @@ fn mobius_codec_workflow_executes() -> anyhow::Result<()> { Ok(()) } -fn tts_request( - prompt_tokens: &[i64], - batch: i64, - slot_ids: &[i64], -) -> anyhow::Result { +fn tts_request(prompt_tokens: &[i64], batch: i64) -> anyhow::Result { let rows = usize::try_from(batch)?; assert_eq!(prompt_tokens.len(), rows * 2); - assert_eq!(slot_ids.len(), rows); Ok( PipelineGenerateRequest::new(GenerateRequest { prompt: GeneratePrompt::TokenIds(vec![0]), @@ -581,10 +531,6 @@ fn tts_request( .with_input( "package.true", Value::from_raw_bytes(vec![1; rows], &[batch], DataType::Bool)?, - ) - .with_input( - "package.slot_ids", - Value::from_slice_i64(slot_ids, &[batch])?, ), ) } @@ -592,16 +538,16 @@ fn tts_request( #[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, &[0])?)?; + 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, &[1])?)?; + 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, &[0, 1])?)?; + 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()?; @@ -613,7 +559,7 @@ fn mobius_tts_workflow_executes_real_producer_graphs() -> anyhow::Result<()> { .iter() .map(|island| island.stable_binding_runs) .sum::(); - let compacted = engine.run_pipeline_outputs(tts_request(&[3, 4, 1, 2], 2, &[1, 0])?)?; + 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); @@ -627,7 +573,7 @@ fn mobius_tts_workflow_executes_real_producer_graphs() -> anyhow::Result<()> { "same-shape nested TTS compaction must preserve stable bindings" ); - let reused = engine.run_pipeline_outputs(tts_request(&[3, 4], 1, &[0])?)?; + let reused = engine.run_pipeline_outputs(tts_request(&[3, 4], 1)?)?; assert_eq!(reused["waveform"].to_vec_f32()?, second); Ok(()) } @@ -639,7 +585,6 @@ fn mobius_speculative_workflow_executes_rejection_and_correction() -> anyhow::Re prompt: GeneratePrompt::TokenIds(vec![1, 2, 3, 4]), options: options(1), }) - .with_input("serving.slot_ids", Value::from_slice_i64(&[0], &[1])?) .with_input( "verifier.past_key_values.0.key", Value::from_slice_f32(&[], &[1, 2, 0, 8])?, From 5a37fe47d89612fe1b929aa8dcf9e50c558f1b35 Mon Sep 17 00:00:00 2001 From: justinchuby Date: Thu, 20 Aug 2026 16:04:47 +0000 Subject: [PATCH 127/151] Describe fixed-capacity and FP8 KV caches instead of refusing to MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit A static-cache export was rejected outright for onnx-genai on the grounds that the metadata could not describe it. That was never true of the schema: ModelIoSpec carries a StaticCacheIoSpec, and the onnx-genai ORT decode backend actually *requires* it — a static-cache graph without model.io.static_cache is rejected there rather than guessed. What was missing was a producer that emits it. The producer now reads the ABI out of the graph rather than assuming it. The two control ports are per-row integer vectors and so are shape-indistinguishable from each other and from every other per-row integer input, which is why they are a declared ABI minted once in _constants.py and read back by the producers. Everything else — the buffers, their constant capacity, the scatter axis — is derived from the graph and disagreements are hard errors, so a graph that scatters on an axis other than the one the metadata publishes cannot ship. Nothing dispatches on a model name. The result is published twice, for two kinds of consumer. model.io.static_cache names the ports directly. The workflow says the same thing operationally: the buffers are invariant loop cells rather than growing ones, the capacity is a package.cache_capacity literal input, and the state service publishes an indexed_scatter discipline naming the write cursor, the capacity and the port that carries it. The write cursor and the logical length are the same quantity, so both name the single carried cache_lengths cell instead of 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 what makes a fixed buffer safe to keep serving a batch whose rows finish at different times. Ragged prefill is not claimed, because ONNX leaves the region between nonpad_kv_seqlen and the query length undefined. Heterogeneous caches keep their own disciplines. Gemma 4 mixes a growing rank-4 BNSH sliding cache with rank-3 fixed-capacity full-attention buffers, and its KV-shared suffix owns no buffer at all, so these surface as separate groups and model.io.static_cache lists only the layers that own something. Two real bugs surfaced while getting that right. The shared-KV fallback pinned a 4-D BNSH shape onto any borrowed KV whose rank was not 4, which overwrote the fully-known rank-3 buffer a static source hands over — corrupting the declared shape of updated_key_cache.N and defeating the rank-3 test sixty lines below, which would then have transposed a rank-3 tensor as BNSH. And the 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, so a hybrid decoder lost all padding information. Both builders now share one rule — a mask exists exactly when some layer still has a dynamic cache. FP8 was never coerced to fp16 in the metadata; the graph dtype is published as-is and a package whose cache is float8_e4m3fn validates even where no kernel can run it. What was wrong is that requesting it could silently do nothing: the gate tested whether GQA fusion was *expected*, not whether any cache was converted. The pass now reports its count and the build fails when it is zero, naming the reason. FP8 KV storage needs an attention operator with k_scale/v_scale inputs to dequantize the cache on read; a TensorScatter plus ai.onnx Attention graph has none, so --features static-cache,fp8-kv-cache is refused rather than quietly producing float16. Verified on an H200 against a real export: batch-of-two prefill matches per-row batch-of-one exactly, the scatter stays inside the declared prefix, decode with divergent per-row cursors matches to 3.1e-6 over three steps, and a finished row reclaims its own slot without leaking past it. FP8 does not execute on the shipped CUDA kernel, and it is not an illegal memory access: the node is never assigned a provider. A three-way comparison of the same graph isolates the cause to the KV type constraint — the identical 14-input GQA with scale inputs and quantization attributes loads and runs with FLOAT KV. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: justinchuby --- .github/workflows/main.yml | 13 +- CHANGELOG.md | 44 + docs/onnx-genai-performance-conformance.md | 42 + docs/onnx-genai-workflows.md | 67 + src/mobius/__main__.py | 17 - src/mobius/_constants.py | 22 + src/mobius/_optimizations.py | 23 +- src/mobius/_passes/_fp8_kv_cache.py | 5 + src/mobius/generation/_policy_components.py | 20 +- .../onnx_genai/inference_metadata.py | 10 +- .../onnx_genai/workflow_metadata.py | 440 +++++- .../onnx_genai/workflow_metadata_test.py | 81 + src/mobius/models/gemma4.py | 9 +- src/mobius/tasks/_causal_lm.py | 8 +- src/mobius/tasks/_gemma4.py | 40 +- tests/cli_test.py | 49 +- .../static_cache/inference_metadata.yaml | 1354 +++++++++++++++++ .../static_cache/model.onnx | Bin 0 -> 2311 bytes .../static_cache/model.onnx.data | 0 .../policies/cache_length_update.onnx | Bin 0 -> 930 bytes .../policies/decoder_state_initializer.onnx | Bin 0 -> 7037 bytes .../policies/decoder_step_update.onnx | Bin 0 -> 780 bytes .../policies/generated_length_update.onnx | Bin 0 -> 930 bytes .../policies/last_token_logits.onnx | Bin 0 -> 631 bytes .../static_cache/policies/termination.onnx | Bin 0 -> 5880 bytes .../termination_batch_initializer.onnx | Bin 0 -> 1677 bytes .../static_cache/policies/token_sampler.onnx | Bin 0 -> 58720 bytes .../policies/token_state_update.onnx | Bin 0 -> 1229 bytes .../static_cache/policies/token_to_slot.onnx | Bin 0 -> 438 bytes ...generate_onnx_genai_validation_packages.py | 50 + tests/onnx_genai_workflow_conformance.rs | 25 + tests/static_cache_metadata_test.py | 408 +++++ 32 files changed, 2611 insertions(+), 116 deletions(-) create mode 100644 tests/fixtures/onnx_genai_workflows/static_cache/inference_metadata.yaml create mode 100644 tests/fixtures/onnx_genai_workflows/static_cache/model.onnx create mode 100644 tests/fixtures/onnx_genai_workflows/static_cache/model.onnx.data create mode 100644 tests/fixtures/onnx_genai_workflows/static_cache/policies/cache_length_update.onnx create mode 100644 tests/fixtures/onnx_genai_workflows/static_cache/policies/decoder_state_initializer.onnx create mode 100644 tests/fixtures/onnx_genai_workflows/static_cache/policies/decoder_step_update.onnx create mode 100644 tests/fixtures/onnx_genai_workflows/static_cache/policies/generated_length_update.onnx create mode 100644 tests/fixtures/onnx_genai_workflows/static_cache/policies/last_token_logits.onnx create mode 100644 tests/fixtures/onnx_genai_workflows/static_cache/policies/termination.onnx create mode 100644 tests/fixtures/onnx_genai_workflows/static_cache/policies/termination_batch_initializer.onnx create mode 100644 tests/fixtures/onnx_genai_workflows/static_cache/policies/token_sampler.onnx create mode 100644 tests/fixtures/onnx_genai_workflows/static_cache/policies/token_state_update.onnx create mode 100644 tests/fixtures/onnx_genai_workflows/static_cache/policies/token_to_slot.onnx create mode 100644 tests/static_cache_metadata_test.py diff --git a/.github/workflows/main.yml b/.github/workflows/main.yml index 4a1dbfe6a..41f848cb7 100644 --- a/.github/workflows/main.yml +++ b/.github/workflows/main.yml @@ -23,7 +23,7 @@ jobs: - uses: actions/checkout@v7 with: repository: justinchuby/onnx-genai - ref: 90e43c62e6051a9a7074b8b8e96593d43d0d389c + ref: 2f02f6ec5990d80d3907ace8b9f1f2ec0dbcecab path: validation/onnx-genai - uses: actions/setup-python@v7 with: @@ -42,12 +42,21 @@ jobs: 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="" for package in tests/fixtures/onnx_genai_workflows/*; 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" + -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 \ diff --git a/CHANGELOG.md b/CHANGELOG.md index dc6f8a782..eb221ba86 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -7,6 +7,50 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 ## [Unreleased] +### 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. + The same ABI is also published authoritatively as `model.io.static_cache`. + 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 `model.io.static_cache` lists only the layers that + own a buffer — 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/docs/onnx-genai-performance-conformance.md b/docs/onnx-genai-performance-conformance.md index 5dff1217c..fcb14b38d 100644 --- a/docs/onnx-genai-performance-conformance.md +++ b/docs/onnx-genai-performance-conformance.md @@ -79,6 +79,48 @@ demonstration even when the selected implementation is pure same-device ONNX. Pe 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. + ## Current measured baseline ONNX GenAI `8bacf8c` reports paired five-sample synthetic native/composite measurements over diff --git a/docs/onnx-genai-workflows.md b/docs/onnx-genai-workflows.md index c5bfc7a62..ef9acfd11 100644 --- a/docs/onnx-genai-workflows.md +++ b/docs/onnx-genai-workflows.md @@ -112,6 +112,73 @@ This preserves fixed `[B,V]` shapes. The generated component is marked 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. All of this is +published twice, for two different kinds of consumer: + +* `model.io.static_cache` names the ports directly, for a consumer that binds + the graph without interpreting a workflow. +* the workflow declares the same thing operationally — 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, and the per-component port that carries it. + +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 `model.io.static_cache` lists only the layers +that actually own a buffer. 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 ### Decoder diff --git a/src/mobius/__main__.py b/src/mobius/__main__.py index 82e17f940..f6f609e5e 100644 --- a/src/mobius/__main__.py +++ b/src/mobius/__main__.py @@ -213,23 +213,6 @@ def _resolve_static_cache_task(model_type: str) -> ModelTask: "Remove --task to use --features static-cache." ) - # Validate static-cache + onnx-genai compatibility. - # - # A static-cache decoder exposes in-place ring buffers (``key_cache.N`` / - # ``updated_key_cache.N``) plus ``write_indices`` and ``nonpad_kv_seqlen``, - # and drops the rank-2 attention mask. The onnx-genai workflow decoder - # contract is built on ``past_key_values.N`` -> ``present.N`` pairs and that - # mask, so the metadata emitter cannot describe a static-cache graph. - # Reject the combination up front rather than after exporting the weights. - if args.static_cache and args.runtime == "onnx-genai": - raise SystemExit( - "Error: --features static-cache cannot be combined with " - "--runtime onnx-genai. The onnx-genai workflow decoder contract " - "requires dynamic past/present KV ports and a rank-2 attention " - "mask, which a static-cache graph does not expose. Build without " - "--features static-cache, or omit --runtime onnx-genai." - ) - # text-only resolution lives in build() (model_type remap + config # stripping), which is only reached on the HuggingFace model-ID path. if args.text_only and args.config: 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/_optimizations.py b/src/mobius/_optimizations.py index efc83d70f..defbdc8d8 100644 --- a/src/mobius/_optimizations.py +++ b/src/mobius/_optimizations.py @@ -601,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/generation/_policy_components.py b/src/mobius/generation/_policy_components.py index cfd9ff27b..a261710e5 100644 --- a/src/mobius/generation/_policy_components.py +++ b/src/mobius/generation/_policy_components.py @@ -673,8 +673,15 @@ def build_decoder_state_initializer( 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.""" + """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 " @@ -846,10 +853,19 @@ def build_decoder_state_initializer( ) generated_lengths.shape = ir.Shape(["batch"]) builder.add_output(generated_lengths, "generated_lengths") - if fixed_capacity: + 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] diff --git a/src/mobius/integrations/onnx_genai/inference_metadata.py b/src/mobius/integrations/onnx_genai/inference_metadata.py index 4f946e440..914fd6bcc 100644 --- a/src/mobius/integrations/onnx_genai/inference_metadata.py +++ b/src/mobius/integrations/onnx_genai/inference_metadata.py @@ -42,6 +42,10 @@ 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 @@ -867,7 +871,7 @@ 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: @@ -876,8 +880,8 @@ def _static_cache_io( f"ABI is 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], diff --git a/src/mobius/integrations/onnx_genai/workflow_metadata.py b/src/mobius/integrations/onnx_genai/workflow_metadata.py index 1d01db535..3931bc412 100644 --- a/src/mobius/integrations/onnx_genai/workflow_metadata.py +++ b/src/mobius/integrations/onnx_genai/workflow_metadata.py @@ -15,7 +15,13 @@ import onnx_ir as ir import yaml -from mobius._constants import OPSET_VERSION +from mobius._constants import ( + OPSET_VERSION, + STATIC_CACHE_KV_SEQUENCE_LENGTH, + STATIC_CACHE_LAYOUT, + STATIC_CACHE_SEQUENCE_AXIS, + STATIC_CACHE_WRITE_INDICES, +) from mobius.generation import ( SOLVER_BUILDERS, PolicyCapabilities, @@ -84,6 +90,7 @@ _port, _shape_metadata, _source_asset_path, + _static_cache_io, add_adapter_service_to_metadata, add_policy_components_to_workflow, build_native_vlm_package_metadata, @@ -533,11 +540,7 @@ def _model_cache_pairs(model: ir.Model) -> list[tuple[ir.Value, ir.Value]]: present = next( ( outputs.get(name) - for name in ( - past.name.replace("past_key_values", "present"), - past.name.replace("past.", "present."), - past.name.replace("past_", "present_"), - ) + for name in _cache_output_candidates(past.name or "") if name in outputs ), None, @@ -547,6 +550,109 @@ def _model_cache_pairs(model: ir.Model) -> list[tuple[ir.Value, ir.Value]]: 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 _static_cache_model_io(model: ir.Model) -> dict[str, Any]: + """Return ``model.io`` declaring the static-cache port ABI of *model*.""" + static_cache = _static_cache_io( + [_port(value) for value in model.graph.inputs], + [_port(value) for value in model.graph.outputs], + ) + if static_cache is None: + raise ValueError( + "decoder was classified as a static-cache graph but exposes no " + "updated__cache. ports to declare" + ) + return {"kv_ownership": "owned", "static_cache": static_cache} + + def _kv_storage_contract(model: ir.Model) -> dict[str, Any]: """Derive physical KV storage from the admitted model interface. @@ -608,6 +714,12 @@ def _consumes_explicit_cache_length(model: ir.Model) -> bool: 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 = { @@ -674,9 +786,17 @@ def _state_aliasing(kv_contract: dict[str, Any]) -> str: def _cache_layer_index(port_name: str, fallback: int) -> int: - """Recover the decoder layer index from a ``past_key_values.N.key`` port name.""" - match = re.search(r"\.(\d+)\.(?:key|value)$", port_name) - return int(match.group(1)) if match else fallback + """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 _state_group_kinds(config: Any, cache_pairs: list[tuple[ir.Value, ir.Value]]) -> list[str]: @@ -711,49 +831,86 @@ def _state_service_groups( 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: 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). 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. + 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) - distinct = sorted(set(kinds)) + 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 = { - kind: (base_name if len(distinct) == 1 else f"{base_name}_{kind}") for kind in distinct + 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, str]]]] = { name: {} for name in names.values() } - for index, kind in enumerate(kinds): + for index, identity in enumerate(identities): cell = f"cache_{index}" - cell_group[cell] = names[kind] + cell_group[cell] = names[identity] for component, aliases in ports.items(): for cell, alias in aliases.items(): grouped_ports[cell_group[cell]].setdefault(component, {})[cell] = alias - groups = { - names[kind]: { + groups = {} + for kind, update in distinct: + is_scattered = update == "indexed_scatter" + name = names[(kind, update)] + group: dict[str, Any] = { "kind": kind, - "sequence_axis": sequence_axis, - "layout": "bnsh", - **({"logical_lengths": logical_lengths} if logical_lengths else {}), - "aliasing": aliasing, - "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", - }, - "ports": grouped_ports[names[kind]], + "sequence_axis": (STATIC_CACHE_SEQUENCE_AXIS if is_scattered else sequence_axis), + "layout": STATIC_CACHE_LAYOUT if is_scattered else "bnsh", } - for kind in distinct - } + 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"] + ), + } + # 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 @@ -3985,10 +4142,7 @@ def build_vlm_workflow_metadata( present = next( ( decoder_outputs.get(name) - for name in ( - value.name.replace("past_key_values", "present"), - value.name.replace("past.", "present."), - ) + for name in _cache_output_candidates(value.name or "") if name in decoder_outputs ), None, @@ -3998,6 +4152,10 @@ def build_vlm_workflow_metadata( 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 @@ -4012,6 +4170,10 @@ def build_vlm_workflow_metadata( 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") @@ -4070,6 +4232,9 @@ def build_vlm_workflow_metadata( 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( @@ -4177,6 +4342,21 @@ def build_vlm_workflow_metadata( ) ), }, + **( + { + # 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"}, @@ -4336,6 +4516,15 @@ def build_vlm_workflow_metadata( 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" @@ -4417,7 +4606,7 @@ def build_vlm_workflow_metadata( "class": "semantic", "scope": "invocation", "initializer": ( - "initializer.cache_lengths" if fixed_capacity else "package.zero_batch" + "initializer.cache_lengths" if tracks_cache_lengths else "package.zero_batch" ), "recurrence": {"kind": "invariant"}, }, @@ -4486,7 +4675,7 @@ def build_vlm_workflow_metadata( ), ( "cache_lengths", - "initializer.cache_lengths" if fixed_capacity else "package.zero_batch", + "initializer.cache_lengths" if tracks_cache_lengths else "package.zero_batch", "state.cache_lengths.body", "cache_lengths.next", "state.cache_lengths.final", @@ -4514,22 +4703,29 @@ def build_vlm_workflow_metadata( ) 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": "bounded", - "axis": next( - ( - axis - for axis, dimension in enumerate(_contract(past)["shape"]) - if "sequence" in str(dimension) + "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, ), - 2, - ), - "max": "package.max_context", - }, + "max": "package.max_context", + } + ), "management": "runtime", "release_boundary": "invocation", } @@ -4568,6 +4764,19 @@ def build_vlm_workflow_metadata( 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"], + } + if static_cache is not None + else None + ), ) for cell, group_name in vlm_cell_groups.items(): state[cell]["service_group"] = group_name @@ -4663,7 +4872,16 @@ def build_vlm_workflow_metadata( "generated_lengths": "initializer.generated_lengths", **( {"cache_lengths": "initializer.cache_lengths"} - if fixed_capacity + if tracks_cache_lengths + else {} + ), + **( + { + static_cache["write_indices"]: ( + f"initializer.{static_cache['write_indices']}" + ) + } + if static_cache is not None else {} ), **( @@ -4930,6 +5148,14 @@ def build_vlm_workflow_metadata( metadata = { "schema_version": "v1", "preprocessing": preprocessing, + # The scatter ABI's control ports are rank-1 integer vectors and are + # indistinguishable by shape, so which is the write cursor and which is + # the non-pad length is declared, never inferred. + **( + {"model": {"io": _static_cache_model_io(decoder)}} + if static_cache is not None + else {} + ), "pipeline": {"workflow": _publish_workflow_v1(workflow)}, } add_policy_components_to_workflow(metadata, pkg) @@ -5909,18 +6135,27 @@ def _build_autoregressive_workflow_metadata( for value in inputs: if value.name in cross_bindings: continue - candidates = [ - value.name.replace("past_key_values", "present"), - value.name.replace("past.", "present."), - value.name.replace("past_", "present_"), - ] present = next( - (output_by_suffix.get(name) for name in candidates if name in output_by_suffix), + ( + 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 @@ -5959,7 +6194,7 @@ def _build_autoregressive_workflow_metadata( None, ), ) - derived_names = cache_names | set(cross_bindings) + 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: @@ -5978,6 +6213,10 @@ def _build_autoregressive_workflow_metadata( 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( @@ -5988,6 +6227,9 @@ def _build_autoregressive_workflow_metadata( 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: @@ -6147,6 +6389,16 @@ def _build_autoregressive_workflow_metadata( }, } ) + 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( { @@ -6340,6 +6592,18 @@ def _build_autoregressive_workflow_metadata( 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. @@ -6413,7 +6677,7 @@ def _build_autoregressive_workflow_metadata( "scope": "invocation", "initializer": ( "initializer.cache_lengths" - if fixed_capacity + if tracks_cache_lengths else "package.cache_lengths" ), "recurrence": {"kind": "invariant"}, @@ -6480,7 +6744,7 @@ def _build_autoregressive_workflow_metadata( "cell": "cache_lengths", "current": ( "initializer.cache_lengths" - if fixed_capacity + if tracks_cache_lengths else "package.cache_lengths" ), "body_input": "state.cache_lengths.body", @@ -6591,15 +6855,23 @@ def _build_autoregressive_workflow_metadata( ), 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": "bounded", - "axis": decoder_kv_axis, - "max": "package.max_context", - }, + "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", @@ -6668,6 +6940,19 @@ def _build_autoregressive_workflow_metadata( 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"], + } + if static_cache is not None + else None + ), ) for cell in decoder_cache_cells: state[cell]["service_group"] = decoder_cell_groups[cell] @@ -6715,7 +7000,16 @@ def _build_autoregressive_workflow_metadata( ), **( {"cache_lengths": "initializer.cache_lengths"} - if fixed_capacity + if tracks_cache_lengths + else {} + ), + **( + { + static_cache["write_indices"]: ( + f"initializer.{static_cache['write_indices']}" + ) + } + if static_cache is not None else {} ), **( @@ -7062,6 +7356,16 @@ def _build_autoregressive_workflow_metadata( metadata = { "schema_version": "1.0", **({"preprocessing": {"audio": audio_program}} if audio_program is not None else {}), + # The scatter ABI's two control ports are integer vectors and so are + # indistinguishable by shape. ``model.io.static_cache`` is the + # authoritative declaration of which port is which, and a runtime that + # drives the graph directly rather than through the workflow reads it + # from here; it is deliberately redundant with the workflow bindings. + **( + {"model": {"io": _static_cache_model_io(decoder)}} + if static_cache is not None + else {} + ), "pipeline": {"workflow": _publish_workflow_v1(workflow)}, } add_policy_components_to_workflow(metadata, pkg) diff --git a/src/mobius/integrations/onnx_genai/workflow_metadata_test.py b/src/mobius/integrations/onnx_genai/workflow_metadata_test.py index 732285406..a7d60593b 100644 --- a/src/mobius/integrations/onnx_genai/workflow_metadata_test.py +++ b/src/mobius/integrations/onnx_genai/workflow_metadata_test.py @@ -18,6 +18,7 @@ ) 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, @@ -787,3 +788,83 @@ def test_paged_cache_inputs_take_precedence_over_operator_derivation(): 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/models/gemma4.py b/src/mobius/models/gemma4.py index c74ed467b..d69ce7a4e 100644 --- a/src/mobius/models/gemma4.py +++ b/src/mobius/models/gemma4.py @@ -986,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", 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/_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/tests/cli_test.py b/tests/cli_test.py index 676f7326c..0aaa8e9f2 100644 --- a/tests/cli_test.py +++ b/tests/cli_test.py @@ -193,18 +193,18 @@ def test_text_only_skips_diffusers_autodetect(self): mock_build.assert_called_once() assert mock_build.call_args.kwargs.get("text_only") is True - def test_static_cache_with_onnx_genai_runtime_errors(self): - """static-cache graphs cannot be described by the onnx-genai contract. - - A static-cache decoder exposes in-place ring buffers plus write indices - instead of past/present KV pairs and a rank-2 attention mask, so the - workflow metadata emitter cannot describe it. The CLI must say so - before exporting the weights, not after. + 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: ``model.io.static_cache`` names + which port is the write cursor and which is the non-pad length, and the + workflow binds those same names. """ - with ( - tempfile.TemporaryDirectory() as tmpdir, - pytest.raises(SystemExit, match=r"static-cache.*--runtime onnx-genai"), - ): + import yaml + + with tempfile.TemporaryDirectory() as tmpdir: main( [ "build", @@ -214,10 +214,37 @@ def test_static_cache_with_onnx_genai_runtime_errors(self): "--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) + + static_cache = metadata["model"]["io"]["static_cache"] + assert static_cache["write_indices_input"] == "write_indices" + assert static_cache["kv_sequence_length_input"] == "nonpad_kv_seqlen" + assert static_cache["key_cache_inputs"][0] == "key_cache.0" + assert static_cache["key_cache_outputs"][0] == "updated_key_cache.0" + assert ( + len(static_cache["key_cache_inputs"]) + == len(static_cache["value_cache_inputs"]) + == len(static_cache["key_cache_outputs"]) + == len(static_cache["value_cache_outputs"]) + ) + + 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" def test_static_cache_task_follows_text_only_substitution(self): """``text-only`` + ``static-cache`` must resolve the *text* task. 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..d77ee1492 --- /dev/null +++ b/tests/fixtures/onnx_genai_workflows/static_cache/inference_metadata.yaml @@ -0,0 +1,1354 @@ +schema_version: '1.0' +model: + io: + kv_ownership: owned + static_cache: + write_indices_input: write_indices + kv_sequence_length_input: nonpad_kv_seqlen + key_cache_inputs: + - key_cache.0 + value_cache_inputs: + - value_cache.0 + key_cache_outputs: + - updated_key_cache.0 + value_cache_outputs: + - updated_value_cache.0 +pipeline: + workflow: + manifest: + ir_version: '1.0' + onnx_opsets: + ai.onnx: 24 + 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 + 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 + aliasing: permitted + reuse: + prefix_reusable: true + evictable_prefix: false + ports: + model: + cache_0: + input: key_cache.0 + output: updated_key_cache.0 + cache_1: + input: value_cache.0 + output: updated_value_cache.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/static_cache/model.onnx b/tests/fixtures/onnx_genai_workflows/static_cache/model.onnx new file mode 100644 index 0000000000000000000000000000000000000000..6098ce3c1f740dd5f0a1f6e9849eb9a3b1d86fe0 GIT binary patch literal 2311 zcmc&#-EI;=6lQ?Ja)3yuw#7~SLyZX+gtmaSCJoUWgBM=ZD~#C;vtwBbv+K@)?Oo$5 z=mYQozJO2Rg|Fj|IJ zYuy7qa>~2I5Y%><@J_%NgmWX~h0TQ|TpD}&joH|Lf3KuX7s?o5MEXN!SX_$cB^smj zs$h*3H#0&Ygb;&|mlA30>v;vv8~IrGOR34UoP{Jgw?r{EM)6fa{E8GE9&;pXR{0%Dxg*}roJ&st_tqwu+X)r(6xZjn<1gKwdh<6 z=v)u$T>q~+e+%kd$6MFgGF7;-K=j7iMgNWGq_Q53SC!P8EigW!(=1`AHG_8=tnPK7~x@dUd!4 zuW;IAR8@Po{W91Rp21Tsjq9+Bx2ZjFY|gl>4Ai-Pr9&j`f{u&P4Z%JP8KTO}e^svs}0st0xwy$5rv49+QQgi{dn{^v=E6^g$A!mP|O literal 0 HcmV?d00001 diff --git a/tests/fixtures/onnx_genai_workflows/static_cache/model.onnx.data b/tests/fixtures/onnx_genai_workflows/static_cache/model.onnx.data new file mode 100644 index 000000000..e69de29bb diff --git a/tests/fixtures/onnx_genai_workflows/static_cache/policies/cache_length_update.onnx b/tests/fixtures/onnx_genai_workflows/static_cache/policies/cache_length_update.onnx new file mode 100644 index 0000000000000000000000000000000000000000..af3bc3fca3b73acdb986806a33ddcf35d9032ffc GIT binary patch literal 930 zcmcIjO-}+b5OsxRVU%c^7;z6qiP75z^yW!Ca`0-RF->X9ZgpulyG0>J4_rO?muG^?f4nPmxkXM$A4+~*~a+4JlnkM**TWH zl!Sb(G)cqV`-JDb)N&3b?u5Fmd*G)*2_ zi@B{ujuyK`Hq=-yGHCALRs#cr=2~_o%RJBF?}e;Zl`%KKHL7{c&U6iP=3T&P`3s_$ zOLMwfC>doc6?BTJk~%lXiMHQAW?m1@{>{XKp`&?L2@;XO>g5w7>?7Vai>X|4CIwWyU}uSc5|BXjAtG{?=$bmL|%uh=g!Ew@LvC? z2R@fTZRR@XGaviT8L>UMc8TAc7&C&~N;PcQ&Y1SpGV2}Hz?yBG6K`gi1Qw1Cz~1ca zvd_l+)f;&n*ZZfbPwG#F70``*x z`27Je1Lm+c#%)-;#P8c)1g)zy<{bvEtF9$a@`Cr2o(TcPoJI|TOZ?{5%&^C}r_`br zL)KGQ!){)%z7|E+Pm8S2iL5~($0E!1aQ-NPdgREa#6IySo*=Y-Ahfq!V)(?x1En4} z8Dj&r67}<ln1KWSU_xAihe|bUFpxvF%zAl2{8kDiAA1L^wUtKFVF0u? zXIIgAP?wgQ^Z12x9$Ii75>#3{Fe>eu20gQyQjeM;T1{OI`>l0QGOjEySNZ+P44hzr zfJIa6O&wp^rjMCFe&?7As2Wle=Nt{R6i{ZTrB#=uM&(E@M88|9l$Y6`rp&Hq%+AM* zr5)$A8ao+SorP-5GnD;$1V4q#eW*@)Yr9=o|mv$QFD8P0kfXM!p7%t?M_ zQ7k?KSknW?9CM>}MG|{KTzl8BPYBlfN-b=K1oqW3>*jS5(vllw=o@hKi4HJN$doKi&DG!bMX%CI~(nIvnjT}?d zG_FkKK#L=XoFlH^cuNa`1NFu3fXk%OP|T#!aH*N}?|deWhH(9iTw``<4%-XlskgK8 z+=^3aKvQYys1#(fEs;kEqQB1uLt>j z1OAU?N3Vfwq5@w6O}cU|Ss4^%X><>8b@|;>4nQ`)$imzQVjRD7h7eh^GwuW4wZUz$ zDNC@)18z{zGW|SJXA5VZ?r_;#8v-Cx21>z)3nTREkQ>?NGxhmE@st}__y!V|QiK~f zT^QRGE?VI#C{b*M)0cxVG_P2?!aG|*-4mtmJt%`!`brTLPPc@8`pIH@!W|nUe2vrA z>1Ob$aoRm!j$xRWry_B?!J6E)PQqM3kZ>8LCpc0Ic2XZDRKLkfsnvSP0AnGph&dKg zy9|U~_R{qjtWw)n=%Q2<T7d^QtJedw~uymZEB4;dZDmY$swtA{L%P8qedsit_&P696WO-EfnEyG%H->(73p=rk3GeKhSQZ`(S_J8szwxm_GZ1>tP8L8aw|1nJ$6l literal 0 HcmV?d00001 diff --git a/tests/fixtures/onnx_genai_workflows/static_cache/policies/decoder_step_update.onnx b/tests/fixtures/onnx_genai_workflows/static_cache/policies/decoder_step_update.onnx new file mode 100644 index 0000000000000000000000000000000000000000..f07a3a6e4e316142535edc6f60901042c484c18b GIT binary patch literal 780 zcmcJNy-or_6opv?S*{RaNDNXKM54xG!&+L2cEl&pn8`3Zz{oN)*_kB}W5XlZ*jV_Q zb}q}J;7_dRZca{4zT9){9ke}?T^{O-_bfbHs4&LYQfNbj!L8cv=PjhiaI0Z2g|iam z$HWgA=E791H!IyKLSB#&(~3}rQk^CosA1ocLQFIbxH1m)iPm_)SwI3hn4JB{=MH$aU=0ciN@{Mn6qtKD1}56v@FBx)1iZXsU^fCzgiYbSM0l$}G3*k; zMP$Io>Wjlw8kVDxnF?0ZF z+v=XePwZcv4}0s~pl!%Raj8WAdPAbmtP?m~?9J&j5e;G9qobtH=@e5bOhBkX>Bpy; r)rRK3&S(K>KXENW91>2r&q*-7bZ6Ev)c?$2y&j_`-q=+JQl;t#ajN`^ literal 0 HcmV?d00001 diff --git a/tests/fixtures/onnx_genai_workflows/static_cache/policies/generated_length_update.onnx b/tests/fixtures/onnx_genai_workflows/static_cache/policies/generated_length_update.onnx new file mode 100644 index 0000000000000000000000000000000000000000..af3bc3fca3b73acdb986806a33ddcf35d9032ffc GIT binary patch literal 930 zcmcIjO-}+b5OsxRVU%c^7;z6qiP75z^yW!Ca`0-RF->X9ZgpulyG0>J4_rO?muG^?f4nPmxkXM$A4+~*~a+4JlnkM**TWH zl!Sb(G)cqV`-JDb)N&3b?u5Fmd*G)*2_ zi@B{ujuyK`Hq=-yGHCALRs#cr=2~_o%RJBF?}e;Zl`%KKHL7{c&U6iP=3T&P`3s_$ zOLMwfC>doc6?BTJk~%lXiMHQAW?m1@{>{XKp`&?L2@;XVSA`N5$W@H3tkY;kI1!>+E8gov%LAWm#V`_v+67_8YRe% zs2_16q$z$Z7kWeKb^)@I4!I5}<2co>!9g(Sxk}2hW}yg-%Y3Re=?fmxko6}xeRtok zKzWUtWf09&!1auPR|K}NdctT}vBbGJ&XDA(T)3g8&qFd7ZdL8%*vNrJ$3hqXQ4nv> zJ&YP~xv;rCTPouds5=NBK?ZMl)UZAPb~JyCxMUpi7)FXwFY;+P@lN6AkMKV1#g)6? zl`Hha&RREYu$&A|O;H^VlhL|8F5^}kp`;rqU$6;bN}7-|gVRsXQ*;TnKS?4wxU`^> g5SK<{;fvV#qBTXwQ28~3ce|XK`4q1#NNp8A0X>A#xc~qF literal 0 HcmV?d00001 diff --git a/tests/fixtures/onnx_genai_workflows/static_cache/policies/termination.onnx b/tests/fixtures/onnx_genai_workflows/static_cache/policies/termination.onnx new file mode 100644 index 0000000000000000000000000000000000000000..3183438a88f3ecce329e4afc84272274af35cdc8 GIT binary patch literal 5880 zcmc&&-EQMV6ppiLlT6a4on^bhETwCeMT%$@J4xH@)v_yu6akf3aZ`jwt~1*@P3)}g zgf6f*6&D}`NU%4c6)Fz^kHTB<2HeA$@$cBFxN^yj=bRJIeE!aypP5PxFHh`abLQUv zYX`lQaD7f5*_P|+mPguhBN|)wkdn4i2^B3Fp}9VuQDR!2`{1AOCtXVA!3wHc`h>bu z-JrNI*hZVvvy)xhvd&$@F{j?HG1gs|oSM|p9pm&FE`IXqU<=8qf3Z1;YuHojvVWkT zCA93>XVh|$028mw$;X!abVljtl&JDrEO6MWQjIiuVf#TsA_`QdwoA;Ri>q_;?$en* zCYoIL2OO2AEO99>L_dno?4+I9NjP&2RU>CM3zt$cFHpZo0xyqG^(iG?x$X}*FI{Dc zOD%+MP;1yM0-RPIEy132ye90g+`zq%fCds^LdaloK-3{=;~J=g?bVZ4!y$*;Q_6vy z7iTF)?xh{Mcau2EguK^x}f=H4m_L)!Lr;x+dwp`}}?ozd_ zl$ey4OTXbR4YDcoZh!$-=Hwk~NK{p>v8NA4Rc?Tg7nW}&6s^kWc%-vz0Xd!zd<2?O zVp6M$ik{sp=a{>G=gb5cgx$0X(MN|fgFc>3h}x0YW61~EQL2%u)j*XqN~c3};^tZi zUPx@4hus8GyYgBj_^`THrL{ck)2|Xr@x?ZNgiP(S2I)PA>K=87+LIgM*vAaTAW-DB z>(adKWqA8v%E5QA{Y>CQ_e|RoDcj@0IfUTchv0;#Pt`sI=P4lXJR6?c8vr|rGY%P!Q?9p6SgwzCGnlMF#Nr|{wiiu}gmRDC$3W}+!)^sNyanLFochTA;*Z=+o3ukoiJ@$U*faIJB(`1736#NYBHopcVede5+EQKNrD6ykaQvaT|{^NThQGTY6V^F?$jQe z#xugQpaWrw|15r?IQEz-_qUWpZX#2M4ow421j^fB;P;dKFa8L0!9^fbhC|Rs4Jp&d z0$3(OyUNTdO)zJ0cM#4SieUA#3Nr*F(DAHiS|?09h0}{;f5reqhUE<&lU;6es7swW z{~Fd_FdcuwpA(sfsgojD!+OC?(HQjf$rEaLOq-xhd43@c(Dr}Q4p*pHVOPgH;2PcB dO-!*9G5fNe&7;#>9{6+eFJ z9f6*STIRrk<$5$&M!=LGsZ@<0$Sa93}x?dbfjp{^Dl?Sd8(XG$^SqE+~%6!U`m zt-v1R-c;ZTnPDd6l?0QBa03<9Lj2aJM1IDn`p(LXuuqL@;^)QqQ-S><7JKiJf#IKq8ulp`-M+8+uu z;+9i4z-$npn(uzbiK92)4Ja2(%ww8ohPqnF3+PgTJJ&!C-G-#;Gr;}w_`|aEA*Mx@LpshFb zPi}uVXe03c)^R`ps!UpfqN?-{)`GXpprwv7A40n2=pXFnA*36z0b?FR<}Gi8jQ6FO ztrRy+IKG=vnF^F7ePJxmOzBOOFpV$wrDsWniriI%!a_NL<5=)JBgzvg>v*oitU0z^ zE`0e1K*g`$7RpO#|2I<320BQh;v^k1oiT?Ie}3LwC=cLhYYsXbQcFa_|EdIq-PT_U CB^YA> literal 0 HcmV?d00001 diff --git a/tests/fixtures/onnx_genai_workflows/static_cache/policies/token_sampler.onnx b/tests/fixtures/onnx_genai_workflows/static_cache/policies/token_sampler.onnx new file mode 100644 index 0000000000000000000000000000000000000000..1a5d5e311eae66f17403b7c4307e7609913be913 GIT binary patch literal 58720 zcmd5_ON?a6RqgCwb!UGkJ>xFd_|r~X1`2zW@Bf~Kggs_347Nv_HUf=hrLwDXx@)F8 zvs{_g9)l2QkPYH%jAg)l+_L#vd<3#3SRf%5ED#b)#EKOXOZF^=6Y*aB?v2;wO+;32 zI^U<>yYJq(C+>|KC*snT(Zz?$y9dXs5C7Z0o*RCCcz%Dmn2)v|O+U3 zPcM(A)y`Aak!)$zcxX1XZXh8Ml!O#_J zbh(4`B;C1c@H~kJ=x9F5{c=E2+qophXqj!aEU3}i=F>J6A(Z@;>E;74H_ayxDhX{q z1#3PwggedWWb7a#=eE1?+&0a%+opx>wrLS92K&x{8BF7HTDaa@GVd)d@7*LWr|2@=fGtIrq0YFO z<=Z0~xH1N=9Rt_V4E)_E8Mt;0+%N+-mVq}4&uQtxbK}BuHev9bO*R0}*+ju}nh|); z5)RK<5*MD+jCgyND0t2=LZsv6ER7j2XBqY0tan`KYCP}FGVQ%tZhLPQZ)k|YbCzg$ z&T_|lGh`h)e36}hhO9%K*=X3Dq32Lw+{PstdJc8Qh37n?flFiH$}w=YIe4}WT)75D z@}bPxX%WeXI8lLma@!lLci1O9edvoL< z>dZ#rIcMWldGC1Ddm|dyj92B(cvT*6yoyE7wt?MYtDFs6<@T`EM&UWn9eB|OBSY87dmAvOsg)e zXrm}%MjHjo6^hR*l}bzi55qi;)!M1a^eJ#C5qB2Tx``>qA0D-Y&QKXYA3!Z z6F0snQgO(!ZCwYbfomp$ywKkm50g=Y^sjWYgOX(jrJWsY5|>L>)Gm9A+A$4mirS@9)Gp%{ zwORUS8`!OFSFE;OotoNq>DIO@QroV0ZM)h;ZM%~Bl?abZTHCHXwe2c$ZM&qk?MhF3 zs_fc!6}7fqvD$XkTif1_2hWwMZC6fhyV_uFyGk`YS8i>)Vzuq+wA8jMR@<(6Yun!+ z(ZHs*T{*SwYJ;_H+rVyZyJEHN>eSS>E4Q{?k=l06YuoiEYTLCQD5z*{yY|$!>&Uh3 ziq^JkUE8kh+IAhawq3K@cHLXsekPvx)~2>yJGJe4gSGA242somZM$Z*?fSITwre(C zRriip{a{1`oAIjJ8Lz72jaRYg**35{Y*n*itJ)s6+Gz2lc5B-;scqN1wq0+cwq5Jm zc1>&BjiPFvXL`gz=3oSu}RZq6u5FFvXKbr>u=+#);&$0W(hgFXv!2 zWz6D9!*YeDH&=*QJZVg>&^Wn56E9aV#gnGcSzzNP2@OjU8aqjdI!O1mffz$q)e*4a zkOk-?*|czK6XlcNmW%`~3Zxhf6%TBZ>|C>GX_F*c1?6vbM+5U(@^poWj#)Hm3lZ&6 zr$iK|H&jIVQtQU1j4HCstkx89Van+^*lndCcf zgF+P*x6^VSqm=SnrV*4wCYC7k30zJl7?9EcjiaCl&HT9Je%zZdge0RdKQ1jlu3SHk zmoxnRfE&*8XEkw2xkx*4$?<19Z`dqiBQel0*||Y)!8;@?g;}~WmIlk(M(iYB&HEpo zXz4`NC^}1nk}W9xwxyHJmB0T?vGQ~~?2;QTErD+EwG}r{LpNC4;rWB;V@$#$pgWPl zbBBWzyOYd=-AR-(_*VuZD(mE;7?8;PLtadfe;Dvi!j{4*cqavO>7=(99+Ij;T{Xa0GQl}`r>OyN zaKT+CK^jBiv2_swK|yp;y^aZ8#is5GoC z@jN5z%`-xlY*<^e@{G*QGqQMj28Q(&W@e>{MlzCMfV@qk5pcK7xc)}R0YTlSZ~+2$ z4_t&KzEBf`WU1i-BySro!0?U*7k}qBG3eb+#}1A+GciizDh7?rf1zOrXu(8U5{7{O z&B726z&9`e2rjU8sfvanpn~%-1Z?mKYo}T{WLD82AUBAB=B5zR8$?zLE=t^ji(5=Wj zFbkc6H)!b8rTcm4A(=GPr9nj3APFYAs+}6TsKCoQXK_)O@uG^Buoyi4K zI_*xtOz%nW&qu7`fS9gf8?1C2WI#!emEK!5oXSc9DpNH9E#0aa$Qj;_T9CF%I zfkJM^h&DPJ2nIPdF?i%Xa1pbH1Fokrapsyh-ol1f4aYHY=7yc1nllsIi{Nf{Dhb%h zPA~!vx*Lpuh0cNz@X$9f31$a=?nWaxp0UkVl2fV>pcflLHbq3yGuWtan8&iO< zSpgn2b`IX4v&RMRnib$dZ&y7I`0W;W2PU}F!67i*nM?N;;A1jyz;V~$4U)U#(wi*6 z+jMBq-CZZ)1^81#cuzUu9jE%hJ0b%IQfN!3F5w-A8w795(h$V4bjpeEIAwLyn~l5) z^BtnO;GKGy?>M&s@WxwpsgL@Olaax@p#^w|pujtIxbHZM8F{A@R)EK;K9`Qj zz=0H&OFISlbh1J2I}E4COCShwX;y$wC#R(VpHA427*6$BIwAuHQW#6S-1l^o+;_*) zg!`Uy?mJGM7Vt)+b0TkuM!?&{eNQ*SeaBlBcmw&!uu1CUzT;#J@Ww{nU2qM#`syXA z!+gh4Q0Qhn9H>H9mozsh1tH$#0yqjP6UhbAgm))NZ+CJb-arlnvp~k}PLg=tiK*eh zFs2_#bk!tH?5YWl;#!Y)5lL?saV4hptS$n~bt{ug!CH?sIAey?;p)bEdSCK(W^LN- zOK!0r#q^2$GJip`1jqV)Kaj}$ru7oS;5IO!z+(w502;4 zS05~A>u1iKIrGZn!64KQ?_`zhM36iQ22xD#Du_g+guPK7)L1>N?~klTx~M#ToEQipR+8r55G{rp=28`U4U zi}XR1;l&5bdk5>4NFu}%eV;gZeGVInT8KL83@K_M-KkL4WXj`xf#8@c4nN3Ujz|7u z^`#=9@-QFr%6k6raDFsfA0N#{K&B*t#^iQUpPaohO+hWcDt@$j1F8HRjzX#IkE8>1 z?;Q^JmW$PTn(R;IjH#botUh-Dg8Pjkr;iD zt}CeI1~LSdJRoGyMh4C}g?(qAani^x{r3;OJ76PY(lJtb>X6EbbZ5@|!Vd?O?i+*Q zg@*@=>7m>B!RL0Ip(C6>I5>ov1>s!$>e`$joZ~Q5bAyDUgX&)WB!?(H4pByMh`%vi z7B+S!9{ciq|9EfSx^uNNGKZb^s_|ugZuiFU((Cj2;r_wH)pp1}?Mrg8INNn)$^H|| zdoV9@1@WVk54i%4LKzp?Con`S0B^5E_umPk2Q$G-x3W*r(nC%opbjfI=ek>HwIVp@ z>J;s>5JjE|=&*n6cUWv7{uJ~!U#(s`df{`&vj@`@bo8g~6P@VE_=-mwaDjjy{GLh1GQ zsORSC&?~9fAQ;O)rr@JrQ7cj6V0`waqfoLvmaFo%FV#Uj5Rn9;qbrZ5KXre8gct@a z^y}@%*-Fv!i?SRSPS-PY=(} z-Z)rozoq_#G(T&XJTn2Qd;?(fb*XY~cPng21l-(FDKJRpH`PI>=VH9ooe9+*nAUi4 zqJJavkmme0843c`*HxAa-Qc3bxn7y8BZ)1f}ioyUN)rSQ^*wY-K^THd4vjw775Y2CNzgshc zYJRaj5lXynw5mmA%ZhZg3fyvLMUczGTMF)JQ zgPU6~UJ?nh`1g2lK9nkdQ;rcSE1IeSxa4*z3OsVVYOobBUPjd5*Q98TI%G&p?FH;| z{K1{Q*%}$=3lHWG=Zm#yljY7fJ=v+g9N!>UcCU2dL&$E1_J7#$oDPo&in(ePyU|jV zKbI;eDiIoV`@PalS;VI3Zps1H`sB$iB1T(-@rmF?!}` z1P!|#vl5lJ*W(wR?RfU+llm zf@o0S!w1nAb#@%U?=)h3=Fah5=n~}`@uPDOlxuJl$~V3$$BVl6lIXsYM9VfV6ufvV zx);3o=+Ql>?@n|N-h0I89{Z&GVp34vd!lA z_q0C*37+ZD&X8Dt7_@-=Eiz093aA(+G(o*PaZHlv-HBsj$M(<@$Ns|2kib`Vg>-Wy z5b8~qf_>VZ90`>AusIUeJ+a?TJMeaPgA5(T{_N(k?H<9R?`DADmS{6TXdC4W`qW|2 zlQtT7&eYeeX(65BTsH@=UKNQztG9|%Vi)(Y#VNzYOLa>F;LNl90QK@jY=J8f|MgncX@rja6*X;RQj-? z1Q}e|2g90i@FiGD>J5*>_~wTD+~AxBpi z^EcMhPGn-mz_ZLoJFEG8f4)COF=u{nd33Oc>6ok8!@~#jqum`*6nowp<9YFt>PXnN z#1zk7`93x2e)O*4vzG>=Tf?oZKz}0eZcJ~-lNxtyPyB$u$=*I9G?!K(= zer3Cx9wo1Ad9ECK|3*#{=T~kG&x?=c58E-O5F$7)ddKkMY;S$=XfFTkvMmFix4&HQ z3wNd1a7gy1yYCq~wq`fHw(UgZ!{H@JbAEXC(!jZbLa=oO;y;>xYPlA)3y5#a_?o;X zA%tIh*5b~JKfGpr+-q-_qiAYaIU8LZo`)!BUK{B@%H_FhIOpD5c7jKE=OvfeQg{u#Hn~42HsyIzjEv@`o`KVR{ z;%m+0u@)bFH+j6#h4px&40*R^CV~Vjgc0u`kEwl9HtyZ}F<@h=$81~57IR(OuN5yC zi?GKs1NOsZ-=aSh>_2wVl@JU(u19Rossy<$v~DCsF!*G}rL8JNgENXtTQvyr$@4so zz-~nN^vu%~s1CphVq$3>WF7^%I>?NZ7@0vv@#_FP53-k$=e6y`)j?+3h^vFl<0VuF z9yL)NU~WDWOX~o8v$2b=%-A4}w}agHjw^G1q61-zEzk3e1Ci&Q048V$*zSpa0@Hy< zPg)0`*l}Arn2PpQkcepZFK48Xj-aRs(2+UaCn(Gv$tDfh1Jc;N4MH2Q2R0k*qv3s-* zu-GuW=;{Ck4K6bl()G)XQTdae=kqPBn%MI^k8e>Oc+^C7(0KHub$}s8CbPMCqJpq4 zC09^m6l5Od(T%xyV&lPbgr4{GHJ-+!6&cW&=nbphsRur-M<5a;#L5&((cPABVhuOy zleVJKK_nB8PrZ*W8;r%v*n{3F0?0P(^r(qSZrYG7xoJdHa!)Iwk2DZ%JrXf+k21uX zNGB-+p)k9LDI@V{%qS!AsESCQ7@)Q#H%LvCfk2r_MIULf+J2+~Yw}10&Fqo1WgvWJ z_b_Gn0%@ZRUnET=w?b)4ax0dml6%~aKGOJ|{Yc|@^pPIVqss6F)Jz$^h}tN_7gAeS zJ*|p-R~KB12-*5BMAoXsTTk<7$+~)SgF<}mff{fx)w_reeStI)AK|p;kzEKaf-3q* zk4C6jd?7V63!-Y)kc~1h2a#Q+l6xBzV(UI-cp4OWB;sn1GO%Z7&m+Ztm?xnZ;%ase zQwBDIWH%XQ6dqMk$-Pis^g_8W*rt#4Lb(-gla}F&w`t4p1>8&-zKGi>!xwTB$*q{% za>=q6%6(BceWcg(thk#f!xwl{W%wd*rVL-`ZIt1Qy{)UBRz<$6i^2(dVQ!yl@z&Ek z_+(u@IX1wiD<1Dc?5*BKbm)t?iTH@SeUCJ$EPbR$Bh)Ou*qfOJTQ)qwVG)P3t5kAt zgTf{PJ<9Nf-sF)8z&*;qYW01O^dJ*e1_E&QB}^Gu0-fDtT85~)QjnyQd!Za<#XiYT zgmSSKHCKjM{n`IWxw0~qk|+a99P=-s$`DHmv(MGa5UXpltFAJ{lDj_1`=A`ljGp{R zpXXs&*FI%@4^cUm4CSw7%6MW_j^+FKBaEsLke{R~EGOwvl{JKIJ?&&IC0;FABiSyP zh&B9}fbY^7AQp#{m#^3b4Pt6U;zIs`6FZ8P%@~gTKmV?srw<_t@KKdTSEd652&v91S@!xk_e2!Xt$NRzZ z4v~WiLrPlOdN_m)Rr`vFmF6LlLz z^(Wl=L2Ftc!uu1<`wo|KWET^uN=)}4$J*5dRV<&u)q!%Zp@y~ngk4Let;!1dO$$fF zFPzQ!yY_0!+vQhx%tU6nvVr)T0g9$)hG$Uz#KZ|hVbo<)agBT_@&E;Ba`;e@^;4t( z>H_@}gMqiir)z@x=+f{aKH)6h@*d)yi2koa!Yw=#RbAX*+XAEP#ujH==?!68UlY>P zriTDPzmOvnL1J7p;~ciFN(gOmSU&%eaR>V*fq) z6l4bKQ{-C>8Osru{#Dz|JQ&-h1mcQ?u^5RLj`GB=00Z-0+`A>Pv!kI9nRIIr>;N>DT96!wKU1xR z{#o0|jE;3bj9J;@f#Qnvi>`$sDG-lSElfJInG~HOGZ4-^z~7$fCTTeRewP-sPld{5 zAC9cR{Ud3OZa%K_IwE~HA}bZOK zg;*f9`2a9%MkT_n+Ye)0d{OHkKTznwcwH0T( zUpObX|sDe&Vb>D@|N-j&Xhd>Z|j;wYYKw<^0^6gXf3uKhX};gZhL$KBC;5n(34|OxnK9asU_`vXX??c)*u104+^V#Ro!Qb5ee|UGf>Hq)$ literal 0 HcmV?d00001 diff --git a/tests/fixtures/onnx_genai_workflows/static_cache/policies/token_state_update.onnx b/tests/fixtures/onnx_genai_workflows/static_cache/policies/token_state_update.onnx new file mode 100644 index 0000000000000000000000000000000000000000..7dfd0e7f3b7409692a79ae731043950283511932 GIT binary patch literal 1229 zcmcIk&2G~`5U$%au?Hf;a-cP$3Rc9SqE-@aJ<%R>;D(Tb(8^g)6Q}8}+1;cqh)W)U z6VHOg)9@g$yLOt^6mg0#GxLpi=G*yM{{g9Ah*36I$3FtNgjy;%lE#9*6`FQKu!4&e zX*b$2`k)B`PcBg92}80q=)ggKK8^+FSBl9j*D;$VO3_J%GLdX@O}0OL9JC><{AFWM z#YB$E{sBu5>Iu`?0%4Wgk~rZh?S){4Bhia?jJ_!n-#we$G0wE2}&%g^*Y*h)`mzC zqQyCKsi6 zggk<#ACTbY^zdmpw^hgJ8Mw0C8*sXJr^^7op;WU>rr~-Dt04=Y8xx7K; zo=fkS2QP}Hi^s^#rsz!YI2W^wT~jl7TAE!TKfN!J6f=yD+tKYBN3YsZ*>vk#sRk{2 zV8K!|m@99_#}>^|QYkKNfQ%RGW^Lts5*5k?T4&+1XsH>@z6%y{(%5#|tnH++@$3wl yLmb!W>(V=c&VO%7vy|-H=i|s!BFSQR_^6=&roZ%FLi?XFij?jFIjzA~Ywr(oA9$t! literal 0 HcmV?d00001 diff --git a/tests/fixtures/onnx_genai_workflows/static_cache/policies/token_to_slot.onnx b/tests/fixtures/onnx_genai_workflows/static_cache/policies/token_to_slot.onnx new file mode 100644 index 0000000000000000000000000000000000000000..eb8e4b6599f5d6095617f9fe0d294eb1d098d216 GIT binary patch literal 438 zcmaiwK}*9h7>3un+4PfQjh9(LrVIrSWxIOP9eekpNXW8mjq8%yq_9rB^zN_lhicoE zp@NuW-iMdxd(i>jU#Y3c_07i~yj##SYRM%4Q&Jl;o{O`5qznA(Y=X}G7R5MOS z7w+yS^1zjJ#dS&<$DL>hXX$bgDk;~RWg;~pOQ_Z)7CfUFiwkT&K1Blveo=J?(M+XW z*97(&8DbQv6|Mth`Of5HNf(^F^;z{hY@1>Pkfwhl0?k#NlR@@LYl8d bB4}3Jj!Wwj27h+&`J6M;Ofu<0$M1gu>uQY1 literal 0 HcmV?d00001 diff --git a/tests/generate_onnx_genai_validation_packages.py b/tests/generate_onnx_genai_validation_packages.py index 816c44894..1b23feb9a 100644 --- a/tests/generate_onnx_genai_validation_packages.py +++ b/tests/generate_onnx_genai_validation_packages.py @@ -127,6 +127,54 @@ def _executable_decoder_package() -> ModelPackage: 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]) @@ -896,8 +944,10 @@ def main() -> None: parser.add_argument("output", type=Path) args = parser.parse_args() 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(), diff --git a/tests/onnx_genai_workflow_conformance.rs b/tests/onnx_genai_workflow_conformance.rs index 16aa90547..ce5413c16 100644 --- a/tests/onnx_genai_workflow_conformance.rs +++ b/tests/onnx_genai_workflow_conformance.rs @@ -298,6 +298,31 @@ fn mobius_decoder_workflow_executes() -> anyhow::Result<()> { 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())?; diff --git a/tests/static_cache_metadata_test.py b/tests/static_cache_metadata_test.py new file mode 100644 index 000000000..26f4b3d43 --- /dev/null +++ b/tests/static_cache_metadata_test.py @@ -0,0 +1,408 @@ +# 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 + +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"] + + +@pytest.fixture(scope="module") +def static_workflow(): + pkg, config = _static_package() + return build_decoder_workflow_metadata(pkg, config) + + +@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 TestStaticCacheModelIo: + """``model.io.static_cache`` is the authoritative port ABI.""" + + def test_declares_control_and_buffer_ports(self, static_workflow): + static_cache = static_workflow["model"]["io"]["static_cache"] + assert static_cache["write_indices_input"] == STATIC_CACHE_WRITE_INDICES + assert static_cache["kv_sequence_length_input"] == STATIC_CACHE_KV_SEQUENCE_LENGTH + assert static_cache["key_cache_inputs"] == ["key_cache.0", "key_cache.1"] + assert static_cache["value_cache_inputs"] == ["value_cache.0", "value_cache.1"] + assert static_cache["key_cache_outputs"] == [ + "updated_key_cache.0", + "updated_key_cache.1", + ] + assert static_cache["value_cache_outputs"] == [ + "updated_value_cache.0", + "updated_value_cache.1", + ] + + def test_per_layer_lists_are_paired(self, static_workflow): + # A runtime binds these four lists positionally; unequal lengths would + # silently pair layer i's key buffer with layer j's output. + static_cache = static_workflow["model"]["io"]["static_cache"] + lengths = { + len(static_cache[key]) + for key in ( + "key_cache_inputs", + "value_cache_inputs", + "key_cache_outputs", + "value_cache_outputs", + ) + } + assert lengths == {2} + + def test_owns_its_cache_and_declares_no_appending_ports(self, static_workflow): + io = static_workflow["model"]["io"] + assert io["kv_ownership"] == "owned" + # A static cache has no past/present pair to advertise; declaring one + # would invite a runtime to concatenate into a fixed buffer. + assert "kv_inputs" not in io + assert "kv_outputs" not in io + + +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}, + } + 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", + } + + 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) + io = build_decoder_workflow_metadata(pkg, config)["model"]["io"]["static_cache"] + assert io["key_cache_inputs"] == ["key_cache.0", "key_cache.1", "key_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. + static_cache = mixed["model"]["io"]["static_cache"] + assert static_cache["key_cache_inputs"] == ["key_cache.1"] + assert static_cache["value_cache_inputs"] == ["value_cache.1"] + assert static_cache["key_cache_outputs"] == ["updated_key_cache.1"] + assert static_cache["value_cache_outputs"] == ["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 = next( + iter(metadata["pipeline"]["workflow"]["serving"]["state_service"]["groups"].values()) + ) + assert group["update"]["kind"] == "indexed_scatter" + assert metadata["model"]["io"]["static_cache"]["key_cache_inputs"] == [ + "key_cache.0", + "key_cache.1", + ] From ac540b5084f0bf87ba0cb18ad8f17bb153f9482b Mon Sep 17 00:00:00 2001 From: justinchuby Date: Thu, 20 Aug 2026 16:09:02 +0000 Subject: [PATCH 128/151] Format the static cache metadata tests The formatter disagrees with the hand-wrapping in the new test module. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: justinchuby --- tests/static_cache_metadata_test.py | 21 +++++++++------------ 1 file changed, 9 insertions(+), 12 deletions(-) diff --git a/tests/static_cache_metadata_test.py b/tests/static_cache_metadata_test.py index 26f4b3d43..2c3676b20 100644 --- a/tests/static_cache_metadata_test.py +++ b/tests/static_cache_metadata_test.py @@ -59,7 +59,6 @@ def _static_package(**overrides): 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())) @@ -88,9 +87,7 @@ def mixed(): # 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 - ) + pkg = Gemma4TextCausalLMTask(static_cache=True, max_seq_len=CAPACITY).build(module, config) return build_decoder_workflow_metadata(pkg, config) @@ -177,9 +174,7 @@ def test_buffers_are_carried_by_the_loop_not_regrown(self, static_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" - ] + groups = static_workflow["pipeline"]["workflow"]["serving"]["state_service"]["groups"] assert len(groups) == 1 group = next(iter(groups.values())) assert group["update"] == { @@ -282,9 +277,7 @@ def test_dual_head_dims_survive_the_split(self, mixed): 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"]) - ) + 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. @@ -354,7 +347,9 @@ 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()) + iter( + metadata["pipeline"]["workflow"]["serving"]["state_service"]["groups"].values() + ) ) assert group["sequence_axis"] == 2 assert group["layout"] == "bnsh" @@ -399,7 +394,9 @@ def test_static_cache_survives_cuda_optimization(self): ) metadata = build_decoder_workflow_metadata(pkg, config) group = next( - iter(metadata["pipeline"]["workflow"]["serving"]["state_service"]["groups"].values()) + iter( + metadata["pipeline"]["workflow"]["serving"]["state_service"]["groups"].values() + ) ) assert group["update"]["kind"] == "indexed_scatter" assert metadata["model"]["io"]["static_cache"]["key_cache_inputs"] == [ From ed73edbffebf322d90d153b31bd69eee05ec661e Mon Sep 17 00:00:00 2001 From: justinchuby Date: Thu, 20 Aug 2026 17:36:21 +0000 Subject: [PATCH 129/151] Pin the runtime that accepts a port ABI beside a workflow The static-cache package declares both a `model.io.static_cache` port ABI and a workflow that binds it, because the two answer different questions and two different consumers read them. The pinned runtime rejected that pair outright and told authors to move the ABI to `pipeline.models..io`, a key the workflow IR had already removed, while its own decode backend refused to load a static-cache graph that did not declare one. No package could satisfy both. The new pin permits the overlap and cross-checks it instead: the ABI's write-destination port must be bound by a state group whose update is an indexed scatter, exactly one component may claim it, and every per-layer key and value pair the ABI names must be a pair that group actually advances. All eleven fixture packages validate against it, which is what this pin buys. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: justinchuby --- .github/workflows/main.yml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/.github/workflows/main.yml b/.github/workflows/main.yml index 41f848cb7..111d63c85 100644 --- a/.github/workflows/main.yml +++ b/.github/workflows/main.yml @@ -23,7 +23,7 @@ jobs: - uses: actions/checkout@v7 with: repository: justinchuby/onnx-genai - ref: 2f02f6ec5990d80d3907ace8b9f1f2ec0dbcecab + ref: 72c21e07bf66465383acb1e03cd2d518ee19e430 path: validation/onnx-genai - uses: actions/setup-python@v7 with: From e2f736fc2b8592269ffd7c77067396fc737638e3 Mon Sep 17 00:00:00 2001 From: justinchuby Date: Thu, 20 Aug 2026 19:42:25 +0000 Subject: [PATCH 130/151] Make the workflow the only place a package describes itself MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit A package could state its decode-step ABI twice: once as `model.io`, and once as the component ports, invoke bindings and state groups a workflow engine executes. Two writable answers to one question is a defect whatever their contents, because nothing forces them to agree and a reader of one never learns the other said something else. The previous commit made that worse by adding `model.io` to static-cache exports specifically, which also meant a bare single-file decoder and a composite package were suddenly different *kinds* of document rather than one shape with a different number of components. `pipeline.workflow` is now the only description, for every export. `model` keeps package-wide geometry and capabilities and never a port ABI. A runtime that wants an optimized single-graph path derives it by lowering the one-component workflow, and a derivation cannot disagree with its source. Three facts had to move into the workflow for that to be true rather than merely asserted. Every ONNX component now declares `ports.inputs`/`ports.outputs`: a full contract for exactly the graph's inputs and outputs, no more and no fewer. A subset would let a consumer quietly fall back to opening the artifact; a superset would be a promise the graph does not keep. 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, and recovering the difference from a port's spelling is the name-guessing this format refuses everywhere else. Mobius mints these names in its own task builders, so it states the mapping between its vocabulary and the runtime's rather than inferring one; a port outside that vocabulary gets no role, because a workflow that guesses is worse than one that stays silent. State port aliases declare `role` and `layer`. A layer's key buffer and its value buffer are the same dtype and the same shape, and a cell's label sorts lexicographically so `cache_10` precedes `cache_2` — a consumer pairing per-layer buffers positionally would silently transpose two layers' caches with nothing failing. Both fields are emitted together or not at all, so a recurrent or convolution cache is never handed a fabricated index that would corrupt the very ordering the index exists to fix. `IndexedScatter` gains `kv_length_ports` beside `write_indices_ports`, because a valid length and a write cursor are both rank-1 integer vectors and are indistinguishable from each other by shape. `_add_explicit_io_to_file` is deleted. It was the only route by which an export could gain a `model.io`, it had no caller, and dead code that can mint a forbidden key is exactly the kind of thing that gets resurrected. The `adapter` fixture's hand-written decoder declared `ports: {}`; it now declares its real ports, since a component that describes nothing is not a description. `tests/canonical_workflow_contract_test.py` is what keeps this from lapsing. It asks one set of shape-agnostic questions — workflow present, no `model.io`, declared ports exactly equal to the graph's, every invoke binding and state alias resolving to a declared port, scatter control ports declared, decode step reconstructible from the workflow alone — of dynamic, static-cache, FP8, heterogeneous and composite packages, and of all 11 checked-in fixtures. None of the assertions names a feature, so a future feature cannot grow a private top-level block while every feature-specific test keeps passing. Verified against ONNX GenAI 02e22dd6, which lands the matching lowering: 11/11 packages validate and 11/11 execute under the runtime conformance suite, including the fixed-capacity decode path, with no `model.io` in any package. That last result is what makes removing the second copy safe rather than merely tidy. The CI pin moves to that commit. Tests: 4654 passed, 58 skipped in the fast suite; lintrunner clean; fixtures regenerate byte-identically. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: justinchuby --- .github/workflows/main.yml | 2 +- CHANGELOG.md | 50 +- docs/onnx-genai-performance-conformance.md | 15 + docs/onnx-genai-workflows.md | 77 ++- .../integrations/onnx_genai/auto_export.py | 28 +- .../onnx_genai/auto_export_test.py | 2 + .../codec_workflow_metadata_test.py | 10 +- .../speech_to_text_workflow_metadata_test.py | 2 + .../onnx_genai/workflow_metadata.py | 149 ++-- .../onnx_genai/workflow_metadata_test.py | 4 + tests/canonical_workflow_contract_test.py | 390 +++++++++++ tests/cli_test.py | 38 +- .../adapter/inference_metadata.yaml | 22 +- .../codec/inference_metadata.yaml | 46 ++ .../decoder/inference_metadata.yaml | 69 ++ .../diffusion/inference_metadata.yaml | 95 +++ .../diffusion_guided/inference_metadata.yaml | 95 +++ .../masked/inference_metadata.yaml | 34 + .../speculative/inference_metadata.yaml | 85 +++ .../static_cache/inference_metadata.yaml | 111 ++- .../tts/inference_metadata.yaml | 649 ++++++++++++++++++ .../video/inference_metadata.yaml | 122 ++++ .../vlm/inference_metadata.yaml | 154 +++++ ...generate_onnx_genai_validation_packages.py | 23 +- tests/static_cache_metadata_test.py | 148 ++-- 25 files changed, 2251 insertions(+), 169 deletions(-) create mode 100644 tests/canonical_workflow_contract_test.py diff --git a/.github/workflows/main.yml b/.github/workflows/main.yml index 111d63c85..773933ea0 100644 --- a/.github/workflows/main.yml +++ b/.github/workflows/main.yml @@ -23,7 +23,7 @@ jobs: - uses: actions/checkout@v7 with: repository: justinchuby/onnx-genai - ref: 72c21e07bf66465383acb1e03cd2d518ee19e430 + ref: 02e22dd6516f3851261ebfc5de8b6a93c48694f9 path: validation/onnx-genai - uses: actions/setup-python@v7 with: diff --git a/CHANGELOG.md b/CHANGELOG.md index eb221ba86..965e36757 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -7,6 +7,51 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 ## [Unreleased] +### 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 + +- Every ONNX component declares `ports.inputs` / `ports.outputs` — a contract + (dtype, rank, shape, batch layout) for exactly the graph's inputs and outputs. + A subset would let a consumer fall back to opening the artifact; a superset + would be a promise the graph does not keep. +- 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 @@ -18,12 +63,11 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 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. - The same ABI is also published authoritatively as `model.io.static_cache`. 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 `model.io.static_cache` lists only the layers that - own a buffer — its KV-shared suffix owns none. + 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. diff --git a/docs/onnx-genai-performance-conformance.md b/docs/onnx-genai-performance-conformance.md index fcb14b38d..8664d3088 100644 --- a/docs/onnx-genai-performance-conformance.md +++ b/docs/onnx-genai-performance-conformance.md @@ -121,6 +121,21 @@ 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 `02e22dd6`, with no `model.io` in any package: + +| Check | Result | +| --- | --- | +| `validate_metadata` over the checked-in fixtures | 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. + ## Current measured baseline ONNX GenAI `8bacf8c` reports paired five-sample synthetic native/composite measurements over diff --git a/docs/onnx-genai-workflows.md b/docs/onnx-genai-workflows.md index ef9acfd11..34c53e5e0 100644 --- a/docs/onnx-genai-workflows.md +++ b/docs/onnx-genai-workflows.md @@ -3,6 +3,55 @@ 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 has to carry everything such a lowering needs, so +every ONNX component declares: + +* `ports.inputs` / `ports.outputs` — a contract (dtype, rank, shape, batch + layout) for every graph input and output, no more and no fewer. +* `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. + +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 the same dtype and the same shape. +* `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. + +Both are emitted together or not at all. A recurrent or convolution cache has no +halves and no layer index to state, and inventing one would corrupt the very +ordering the index exists to fix. + +`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. @@ -125,16 +174,16 @@ ports drive that write: | `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. All of this is -published twice, for two different kinds of consumer: - -* `model.io.static_cache` names the ports directly, for a consumer that binds - the graph without interpreting a workflow. -* the workflow declares the same thing operationally — 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, and the per-component port that carries it. +`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 @@ -155,10 +204,10 @@ 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 `model.io.static_cache` lists only the layers -that actually own a buffer. 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. +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 diff --git a/src/mobius/integrations/onnx_genai/auto_export.py b/src/mobius/integrations/onnx_genai/auto_export.py index c31c2f734..a4491f015 100644 --- a/src/mobius/integrations/onnx_genai/auto_export.py +++ b/src/mobius/integrations/onnx_genai/auto_export.py @@ -18,15 +18,11 @@ import numpy as np import onnx_ir as ir -import yaml from mobius.integrations.onnx_genai.inference_metadata import ( _TEXT_RUNTIME_ASSET_NAMES, SchedulerConfig, _copy_runtime_assets, - add_adapter_service_to_metadata, - add_explicit_package_io, - add_policy_components_to_workflow, load_diffusers_scheduler_config, load_diffusers_vae_scaling_factor, ) @@ -199,29 +195,7 @@ def _looks_like_image_edit(pkg: Any) -> bool: ) -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.""" - try: - models = list(pkg.values()) - 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) - add_policy_components_to_workflow(metadata, pkg) - add_adapter_service_to_metadata(metadata, pkg, os.path.dirname(path)) - with open(path, "w", encoding="utf-8") as handle: - yaml.safe_dump(metadata, handle, sort_keys=False) - - -def _write_clip_tokenizer( - output_dir: str, - source: str | None, - *, - revision: str | None = None, -) -> str | None: +def _write_clip_tokenizer(output_dir: str, source: str | None) -> str | None: """Emit ``tokenizer.json`` for a text-conditioned diffusion package. Classic Stable Diffusion conditions on a CLIP text encoder, and the diff --git a/src/mobius/integrations/onnx_genai/auto_export_test.py b/src/mobius/integrations/onnx_genai/auto_export_test.py index 7cd24d6eb..593fe2803 100644 --- a/src/mobius/integrations/onnx_genai/auto_export_test.py +++ b/src/mobius/integrations/onnx_genai/auto_export_test.py @@ -944,6 +944,8 @@ def test_dispatch_speech_to_text_workflow(tmp_path): assert groups["decoder_cache"]["ports"]["decoder"]["cache_0"] == { "input": "past_key_values.0.key", "output": "present.0.key", + "role": "key", + "layer": 0, } diff --git a/src/mobius/integrations/onnx_genai/codec_workflow_metadata_test.py b/src/mobius/integrations/onnx_genai/codec_workflow_metadata_test.py index 16a10cdb4..f55d2f4de 100644 --- a/src/mobius/integrations/onnx_genai/codec_workflow_metadata_test.py +++ b/src/mobius/integrations/onnx_genai/codec_workflow_metadata_test.py @@ -53,7 +53,15 @@ def test_codec_workflow_has_typed_ssa_and_audio_emit(): "shape": ["batch", 1, "audio_samples"], "batch_layout": {"kind": "request_aligned", "axis": 0}, } - assert "ports" not in workflow["components"]["encoder"] + # An ONNX component declares the contract of every port it exposes, so the + # workflow describes the package without opening the artifact. + assert workflow["components"]["encoder"]["ports"]["inputs"]["waveform"] == { + "dtype": "float32", + "rank": 3, + "shape": ["batch", 1, "audio_samples"], + "batch_layout": {"kind": "request_aligned", "axis": 0}, + } + assert set(workflow["components"]["decoder"]["ports"]["outputs"]) == {"waveform"} assert "effects" not in workflow["components"]["encoder"] assert "effects" not in workflow["components"]["decoder"] 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 index 1a6b176f2..76fd6d13a 100644 --- 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 @@ -142,6 +142,8 @@ def test_self_attention_cache_is_the_only_served_group(): assert ports["cache_0"] == { "input": "past_key_values.0.key", "output": "present.0.key", + "role": "key", + "layer": 0, } assert len(ports) == 4 diff --git a/src/mobius/integrations/onnx_genai/workflow_metadata.py b/src/mobius/integrations/onnx_genai/workflow_metadata.py index 3931bc412..3aeaa17f8 100644 --- a/src/mobius/integrations/onnx_genai/workflow_metadata.py +++ b/src/mobius/integrations/onnx_genai/workflow_metadata.py @@ -90,7 +90,6 @@ _port, _shape_metadata, _source_asset_path, - _static_cache_io, add_adapter_service_to_metadata, add_policy_components_to_workflow, build_native_vlm_package_metadata, @@ -167,14 +166,56 @@ def _request_aligned(contract: dict[str, Any], axis: int = 0) -> dict[str, Any]: 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]: - del model, effects - return {"implementation": {"kind": "onnx", "artifact": artifact}} + """Declare one ONNX-backed workflow component: ports, contracts and roles. + + The workflow is the package's only description of itself, so a component + states the contract of every port it exposes rather than leaving a consumer + to open the artifact and infer one. That is what lets the scatter ABI, the + token and logits roles, and the per-layer cache pairs all be resolved from + the workflow alone: an integer control vector is indistinguishable from its + neighbours by shape, so the binding that names it has to sit next to a + declared port for the name to mean anything. + + A contract says what a value *is*; ``roles`` says what the component *does* + with it. 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. + Only ports in this producer's own vocabulary get a role; state ports never + need one, because the group that carries them already names its pairs. + """ + del effects + inputs = {str(value.name): _contract(value) for value in model.graph.inputs} + outputs = {str(value.name): _contract(value) for value in model.graph.outputs} + roles = {name: _PORT_ROLES[name] for name in (*inputs, *outputs) if name in _PORT_ROLES} + ports: dict[str, Any] = {"inputs": inputs, "outputs": outputs} + if roles: + ports["roles"] = roles + return {"implementation": {"kind": "onnx", "artifact": artifact}, "ports": ports} def _grammar_adapter_component(action: str) -> dict[str, Any]: @@ -639,20 +680,6 @@ def _static_cache_ports(model: ir.Model) -> dict[str, Any] | None: } -def _static_cache_model_io(model: ir.Model) -> dict[str, Any]: - """Return ``model.io`` declaring the static-cache port ABI of *model*.""" - static_cache = _static_cache_io( - [_port(value) for value in model.graph.inputs], - [_port(value) for value in model.graph.outputs], - ) - if static_cache is None: - raise ValueError( - "decoder was classified as a static-cache graph but exposes no " - "updated__cache. ports to declare" - ) - return {"kv_ownership": "owned", "static_cache": static_cache} - - def _kv_storage_contract(model: ir.Model) -> dict[str, Any]: """Derive physical KV storage from the admitted model interface. @@ -742,6 +769,28 @@ def _aliasing_for_storage(storage: str) -> str: 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], @@ -766,7 +815,10 @@ def _state_group( "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": ports, + "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 @@ -799,6 +851,22 @@ def _cache_layer_index(port_name: str, fallback: int) -> int: 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. @@ -868,15 +936,20 @@ def _state_service_groups( 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, str]]]] = { + 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] = alias + grouped_ports[cell_group[cell]].setdefault(component, {})[cell] = _annotated_alias( + alias + ) groups = {} for kind, update in distinct: is_scattered = update == "indexed_scatter" @@ -897,6 +970,15 @@ def _state_service_groups( "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 @@ -4773,6 +4855,7 @@ def build_vlm_workflow_metadata( "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 @@ -5148,14 +5231,6 @@ def build_vlm_workflow_metadata( metadata = { "schema_version": "v1", "preprocessing": preprocessing, - # The scatter ABI's control ports are rank-1 integer vectors and are - # indistinguishable by shape, so which is the write cursor and which is - # the non-pad length is declared, never inferred. - **( - {"model": {"io": _static_cache_model_io(decoder)}} - if static_cache is not None - else {} - ), "pipeline": {"workflow": _publish_workflow_v1(workflow)}, } add_policy_components_to_workflow(metadata, pkg) @@ -6949,6 +7024,7 @@ def _build_autoregressive_workflow_metadata( "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 @@ -7356,16 +7432,6 @@ def _build_autoregressive_workflow_metadata( metadata = { "schema_version": "1.0", **({"preprocessing": {"audio": audio_program}} if audio_program is not None else {}), - # The scatter ABI's two control ports are integer vectors and so are - # indistinguishable by shape. ``model.io.static_cache`` is the - # authoritative declaration of which port is which, and a runtime that - # drives the graph directly rather than through the workflow reads it - # from here; it is deliberately redundant with the workflow bindings. - **( - {"model": {"io": _static_cache_model_io(decoder)}} - if static_cache is not None - else {} - ), "pipeline": {"workflow": _publish_workflow_v1(workflow)}, } add_policy_components_to_workflow(metadata, pkg) @@ -8922,10 +8988,9 @@ def _annotate_duplex_state_service( "capabilities": {"snapshot": True, "fork": False}, "ports": { "temporal": { - f"temporal_cache_{index}": { - "input": past.name, - "output": present.name, - } + f"temporal_cache_{index}": _annotated_alias( + {"input": past.name, "output": present.name} + ) for index, (past, present) in enumerate(temporal_caches) } }, diff --git a/src/mobius/integrations/onnx_genai/workflow_metadata_test.py b/src/mobius/integrations/onnx_genai/workflow_metadata_test.py index a7d60593b..525a2e7aa 100644 --- a/src/mobius/integrations/onnx_genai/workflow_metadata_test.py +++ b/src/mobius/integrations/onnx_genai/workflow_metadata_test.py @@ -399,6 +399,8 @@ def collect_emits(node): 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() @@ -727,6 +729,8 @@ def test_speculative_workflow_uses_per_row_ragged_state_and_rng(): ]["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"]) diff --git a/tests/canonical_workflow_contract_test.py b/tests/canonical_workflow_contract_test.py new file mode 100644 index 000000000..bfaaadc12 --- /dev/null +++ b/tests/canonical_workflow_contract_test.py @@ -0,0 +1,390 @@ +# 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 glob +import os +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 + +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 _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 {} + + +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_every_graph_declares_exactly_the_ports_it_has(self, package): + """Declared ports are the graph's ports — not a subset, not a superset. + + A subset lets a runtime silently fall back to opening the artifact, and + a superset is a promise the graph does not keep. Either way the + declaration stops being usable as the single source of truth. + """ + pkg, metadata = package + components = _onnx_components(metadata["pipeline"]["workflow"]) + graphs = {name: pkg[name] for name in components if name in pkg} + assert graphs, "a package with no ONNX component describes nothing" + for name, model in graphs.items(): + ports = components[name]["ports"] + assert set(ports["inputs"]) == {str(v.name) for v in model.graph.inputs} + assert set(ports["outputs"]) == {str(v.name) for v in model.graph.outputs} + + def test_declared_contracts_match_the_graph_dtype_and_rank(self, package): + """A contract that disagrees with its graph would mis-size every buffer.""" + pkg, metadata = package + dtypes = { + ir.DataType.FLOAT: "float32", + ir.DataType.FLOAT16: "float16", + ir.DataType.BFLOAT16: "bfloat16", + ir.DataType.INT64: "int64", + ir.DataType.INT32: "int32", + ir.DataType.BOOL: "bool", + ir.DataType.FLOAT8E4M3FN: "float8_e4m3fn", + } + components = _onnx_components(metadata["pipeline"]["workflow"]) + for name in (name for name in components if name in pkg): + model = pkg[name] + ports = components[name]["ports"] + declared = {**ports["inputs"], **ports["outputs"]} + for value in (*model.graph.inputs, *model.graph.outputs): + contract = declared[str(value.name)] + assert contract["rank"] == len(value.shape) + if value.dtype in dtypes: + assert contract["dtype"] == dtypes[value.dtype] + + def test_every_invocation_binds_a_declared_port(self, package): + """A binding to an undeclared port names nothing a consumer can resolve.""" + _, metadata = package + workflow = metadata["pipeline"]["workflow"] + components = workflow["components"] + for step in _walk_steps(workflow["steps"]): + if step.get("kind") != "invoke": + continue + ports = components[step["component"]]["ports"] + assert set(step.get("inputs", {})) <= set(ports["inputs"]) + assert set(step.get("outputs", {})) <= set(ports["outputs"]) + + def test_every_state_pair_names_declared_ports(self, package): + """State is carried through ports, so both halves have to be declared.""" + _, metadata = package + workflow = metadata["pipeline"]["workflow"] + components = workflow["components"] + for group in _groups(workflow).values(): + for component, aliases in (group.get("ports") or {}).items(): + ports = components[component]["ports"] + for alias in aliases.values(): + assert alias["input"] in ports["inputs"] + assert alias["output"] in ports["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_are_declared_ports(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 that + declares 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. + """ + _, metadata = package + workflow = metadata["pipeline"]["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: + declared = workflow["components"][component]["ports"]["inputs"] + assert update["write_indices_ports"][component] in declared + assert update["kv_length_ports"][component] in declared + + 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. + """ + _, metadata = package + workflow = metadata["pipeline"]["workflow"] + decoders = [] + for name, component in _onnx_components(workflow).items(): + roles = component["ports"].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, component, roles)) + assert decoders, "no component declares what it does with the sequence" + for name, component, roles in decoders: + inputs = component["ports"]["inputs"] + outputs = component["ports"]["outputs"] + sequence = [ + port + for port, role in roles.items() + if role in {"token_ids", "inputs_embeds"} and port in inputs + ] + assert len(sequence) == 1 + assert [port for port, role in roles.items() if role == "logits"] == [ + port for port in outputs if roles.get(port) == "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 _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) + + def test_describes_itself_only_through_the_workflow(self, directory): + metadata = self._metadata(directory) + assert "workflow" in metadata["pipeline"] + assert "io" not in (metadata.get("model") or {}) + + def test_every_onnx_component_declares_its_ports(self, directory): + workflow = self._metadata(directory)["pipeline"]["workflow"] + for component in _onnx_components(workflow).values(): + 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_binding_and_state_pair_resolves(self, directory): + workflow = self._metadata(directory)["pipeline"]["workflow"] + components = workflow["components"] + for step in _walk_steps(workflow["steps"]): + if step.get("kind") != "invoke": + continue + ports = components[step["component"]]["ports"] + assert set(step.get("inputs", {})) <= set(ports["inputs"]) + assert set(step.get("outputs", {})) <= set(ports["outputs"]) + for group in _groups(workflow).values(): + for component, aliases in (group.get("ports") or {}).items(): + ports = components[component]["ports"] + for alias in aliases.values(): + assert alias["input"] in ports["inputs"] + assert alias["output"] in ports["outputs"] diff --git a/tests/cli_test.py b/tests/cli_test.py index 0aaa8e9f2..b8e08bd86 100644 --- a/tests/cli_test.py +++ b/tests/cli_test.py @@ -198,9 +198,10 @@ def test_static_cache_with_onnx_genai_runtime_emits_scatter_abi(self): 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: ``model.io.static_cache`` names - which port is the write cursor and which is the non-pad length, and the - workflow binds those same names. + 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 yaml @@ -225,17 +226,8 @@ def test_static_cache_with_onnx_genai_runtime_emits_scatter_abi(self): ) as handle: metadata = yaml.safe_load(handle) - static_cache = metadata["model"]["io"]["static_cache"] - assert static_cache["write_indices_input"] == "write_indices" - assert static_cache["kv_sequence_length_input"] == "nonpad_kv_seqlen" - assert static_cache["key_cache_inputs"][0] == "key_cache.0" - assert static_cache["key_cache_outputs"][0] == "updated_key_cache.0" - assert ( - len(static_cache["key_cache_inputs"]) - == len(static_cache["value_cache_inputs"]) - == len(static_cache["key_cache_outputs"]) - == len(static_cache["value_cache_outputs"]) - ) + # 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 @@ -245,6 +237,24 @@ def test_static_cache_with_onnx_genai_runtime_emits_scatter_abi(self): 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"} + + declared = workflow["components"]["model"]["ports"]["inputs"] + assert declared["write_indices"]["rank"] == 1 + assert declared["nonpad_kv_seqlen"]["rank"] == 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. diff --git a/tests/fixtures/onnx_genai_workflows/adapter/inference_metadata.yaml b/tests/fixtures/onnx_genai_workflows/adapter/inference_metadata.yaml index 88b1c819c..86e3b3a75 100644 --- a/tests/fixtures/onnx_genai_workflows/adapter/inference_metadata.yaml +++ b/tests/fixtures/onnx_genai_workflows/adapter/inference_metadata.yaml @@ -108,7 +108,27 @@ pipeline: implementation: kind: onnx artifact: model.onnx - ports: {} + ports: + inputs: + activations: + dtype: float32 + rank: 2 + shape: + - batch + - 2 + batch_layout: + kind: request_aligned + axis: 0 + outputs: + projection.output: + dtype: float32 + rank: 2 + shape: + - batch + - 2 + batch_layout: + kind: request_aligned + axis: 0 overlay: implementation: kind: adapter diff --git a/tests/fixtures/onnx_genai_workflows/codec/inference_metadata.yaml b/tests/fixtures/onnx_genai_workflows/codec/inference_metadata.yaml index aa6c5e570..394c7be99 100644 --- a/tests/fixtures/onnx_genai_workflows/codec/inference_metadata.yaml +++ b/tests/fixtures/onnx_genai_workflows/codec/inference_metadata.yaml @@ -47,10 +47,56 @@ pipeline: implementation: kind: onnx artifact: encoder/model.onnx + ports: + inputs: + waveform: + dtype: float32 + rank: 3 + shape: + - batch + - 1 + - audio_samples + batch_layout: + kind: request_aligned + axis: 0 + outputs: + codes: + dtype: float32 + rank: 3 + shape: + - batch + - 1 + - audio_samples + batch_layout: + kind: request_aligned + axis: 0 decoder: implementation: kind: onnx artifact: decoder/model.onnx + ports: + inputs: + codes: + dtype: float32 + rank: 3 + shape: + - batch + - 1 + - audio_samples + batch_layout: + kind: request_aligned + axis: 0 + outputs: + waveform: + dtype: float32 + rank: 3 + shape: + - batch + - 1 + - audio_samples + batch_layout: + kind: request_aligned + axis: 0 steps: - kind: invoke component: encoder diff --git a/tests/fixtures/onnx_genai_workflows/decoder/inference_metadata.yaml b/tests/fixtures/onnx_genai_workflows/decoder/inference_metadata.yaml index b62107279..7632647cd 100644 --- a/tests/fixtures/onnx_genai_workflows/decoder/inference_metadata.yaml +++ b/tests/fixtures/onnx_genai_workflows/decoder/inference_metadata.yaml @@ -340,6 +340,73 @@ pipeline: implementation: kind: onnx artifact: model.onnx + ports: + inputs: + input_ids: + dtype: int64 + rank: 2 + shape: + - batch + - sequence + batch_layout: + kind: request_aligned + axis: 0 + 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 + - sequence + 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 + outputs: + logits: + dtype: float32 + rank: 3 + shape: + - batch + - sequence + - 128 + batch_layout: + kind: request_aligned + axis: 0 + present.0.key: + dtype: float32 + rank: 4 + shape: + - batch + - 2 + - present_sequence + - 8 + batch_layout: + kind: request_aligned + axis: 0 + roles: + input_ids: token_ids + attention_mask: attention_mask + position_ids: position_ids + logits: logits token_sampler: implementation: kind: onnx @@ -1137,6 +1204,8 @@ pipeline: cache_0: input: past_key_values.0.key output: present.0.key + role: key + layer: 0 steps: - kind: loop setup: diff --git a/tests/fixtures/onnx_genai_workflows/diffusion/inference_metadata.yaml b/tests/fixtures/onnx_genai_workflows/diffusion/inference_metadata.yaml index 9bdb0a5a0..e1bcc0dbf 100644 --- a/tests/fixtures/onnx_genai_workflows/diffusion/inference_metadata.yaml +++ b/tests/fixtures/onnx_genai_workflows/diffusion/inference_metadata.yaml @@ -152,14 +152,109 @@ pipeline: implementation: kind: onnx artifact: text_encoder/model.onnx + ports: + inputs: + input_ids: + dtype: int64 + rank: 2 + shape: + - batch + - prompt_sequence + batch_layout: + kind: request_aligned + axis: 0 + outputs: + encoder_hidden_states: + dtype: float32 + rank: 3 + shape: + - batch + - prompt_sequence + - 32 + batch_layout: + kind: request_aligned + axis: 0 + roles: + input_ids: token_ids + encoder_hidden_states: encoder_hidden_states denoiser: implementation: kind: onnx artifact: denoiser/model.onnx + ports: + inputs: + sample: + dtype: float32 + rank: 4 + shape: + - batch + - 4 + - height + - width + batch_layout: + kind: request_aligned + axis: 0 + timestep: + dtype: float32 + rank: 1 + shape: + - batch + batch_layout: + kind: request_aligned + axis: 0 + encoder_hidden_states: + dtype: float32 + rank: 3 + shape: + - batch + - prompt_sequence + - 32 + batch_layout: + kind: request_aligned + axis: 0 + outputs: + noise_pred: + dtype: float32 + rank: 4 + shape: + - batch + - 4 + - height + - width + batch_layout: + kind: request_aligned + axis: 0 + roles: + encoder_hidden_states: encoder_hidden_states vae_decoder: implementation: kind: onnx artifact: vae_decoder/model.onnx + ports: + inputs: + latent: + dtype: float32 + rank: 4 + shape: + - batch + - 4 + - height + - width + batch_layout: + kind: request_aligned + axis: 0 + outputs: + image: + dtype: float32 + rank: 4 + shape: + - batch + - 3 + - height + - width + batch_layout: + kind: request_aligned + axis: 0 solver_step: implementation: kind: onnx diff --git a/tests/fixtures/onnx_genai_workflows/diffusion_guided/inference_metadata.yaml b/tests/fixtures/onnx_genai_workflows/diffusion_guided/inference_metadata.yaml index fdbc5a51b..bd47b38b3 100644 --- a/tests/fixtures/onnx_genai_workflows/diffusion_guided/inference_metadata.yaml +++ b/tests/fixtures/onnx_genai_workflows/diffusion_guided/inference_metadata.yaml @@ -210,14 +210,109 @@ pipeline: implementation: kind: onnx artifact: text_encoder/model.onnx + ports: + inputs: + input_ids: + dtype: int64 + rank: 2 + shape: + - batch + - prompt_sequence + batch_layout: + kind: request_aligned + axis: 0 + outputs: + encoder_hidden_states: + dtype: float32 + rank: 3 + shape: + - batch + - prompt_sequence + - 32 + batch_layout: + kind: request_aligned + axis: 0 + roles: + input_ids: token_ids + encoder_hidden_states: encoder_hidden_states denoiser: implementation: kind: onnx artifact: denoiser/model.onnx + ports: + inputs: + sample: + dtype: float32 + rank: 4 + shape: + - batch + - 4 + - height + - width + batch_layout: + kind: request_aligned + axis: 0 + timestep: + dtype: float32 + rank: 1 + shape: + - batch + batch_layout: + kind: request_aligned + axis: 0 + encoder_hidden_states: + dtype: float32 + rank: 3 + shape: + - batch + - prompt_sequence + - 32 + batch_layout: + kind: request_aligned + axis: 0 + outputs: + noise_pred: + dtype: float32 + rank: 4 + shape: + - batch + - 4 + - height + - width + batch_layout: + kind: request_aligned + axis: 0 + roles: + encoder_hidden_states: encoder_hidden_states vae_decoder: implementation: kind: onnx artifact: vae_decoder/model.onnx + ports: + inputs: + latent: + dtype: float32 + rank: 4 + shape: + - batch + - 4 + - height + - width + batch_layout: + kind: request_aligned + axis: 0 + outputs: + image: + dtype: float32 + rank: 4 + shape: + - batch + - 3 + - height + - width + batch_layout: + kind: request_aligned + axis: 0 solver_step: implementation: kind: onnx diff --git a/tests/fixtures/onnx_genai_workflows/masked/inference_metadata.yaml b/tests/fixtures/onnx_genai_workflows/masked/inference_metadata.yaml index 7aaf1bc73..fe8e5ab3b 100644 --- a/tests/fixtures/onnx_genai_workflows/masked/inference_metadata.yaml +++ b/tests/fixtures/onnx_genai_workflows/masked/inference_metadata.yaml @@ -138,6 +138,40 @@ pipeline: implementation: kind: onnx artifact: model.onnx + ports: + inputs: + input_ids: + dtype: int64 + rank: 2 + shape: + - batch + - sequence + batch_layout: + kind: request_aligned + axis: 0 + outputs: + logits: + dtype: float32 + rank: 3 + shape: + - batch + - sequence + - 128 + batch_layout: + kind: request_aligned + axis: 0 + proposed_tokens: + dtype: int64 + rank: 2 + shape: + - batch + - sequence + batch_layout: + kind: request_aligned + axis: 0 + roles: + input_ids: token_ids + logits: logits masked_update: implementation: kind: onnx diff --git a/tests/fixtures/onnx_genai_workflows/speculative/inference_metadata.yaml b/tests/fixtures/onnx_genai_workflows/speculative/inference_metadata.yaml index 186d134c8..1dbf87fd5 100644 --- a/tests/fixtures/onnx_genai_workflows/speculative/inference_metadata.yaml +++ b/tests/fixtures/onnx_genai_workflows/speculative/inference_metadata.yaml @@ -282,10 +282,93 @@ pipeline: implementation: kind: onnx artifact: proposer/model.onnx + ports: + inputs: + tokens: + dtype: int64 + rank: 2 + shape: + - batch + - 4 + batch_layout: + kind: request_aligned + axis: 0 + proposal_budget: + dtype: int64 + rank: 1 + shape: + - batch + batch_layout: + kind: request_aligned + axis: 0 + outputs: + proposed_tokens: + dtype: int64 + rank: 2 + shape: + - batch + - 4 + batch_layout: + kind: request_aligned + axis: 0 + proposal_scores: + dtype: float32 + rank: 3 + shape: + - batch + - 4 + - 32 + batch_layout: + kind: request_aligned + axis: 0 verifier: implementation: kind: onnx artifact: verifier/model.onnx + ports: + inputs: + proposed_tokens: + dtype: int64 + rank: 2 + shape: + - batch + - 4 + 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 + outputs: + target_scores: + dtype: float32 + rank: 3 + shape: + - batch + - 4 + - 32 + batch_layout: + kind: request_aligned + axis: 0 + present.0.key: + dtype: float32 + rank: 4 + shape: + - batch + - 2 + - past_sequence + 4 + - 8 + batch_layout: + kind: request_aligned + axis: 0 grammar_clone: implementation: kind: adapter @@ -1148,6 +1231,8 @@ pipeline: cache_0: input: past_key_values.0.key output: present.0.key + role: key + layer: 0 logical_lengths: cache_lengths effects: grammar: diff --git a/tests/fixtures/onnx_genai_workflows/static_cache/inference_metadata.yaml b/tests/fixtures/onnx_genai_workflows/static_cache/inference_metadata.yaml index d77ee1492..06163a59f 100644 --- a/tests/fixtures/onnx_genai_workflows/static_cache/inference_metadata.yaml +++ b/tests/fixtures/onnx_genai_workflows/static_cache/inference_metadata.yaml @@ -1,18 +1,4 @@ schema_version: '1.0' -model: - io: - kv_ownership: owned - static_cache: - write_indices_input: write_indices - kv_sequence_length_input: nonpad_kv_seqlen - key_cache_inputs: - - key_cache.0 - value_cache_inputs: - - value_cache.0 - key_cache_outputs: - - updated_key_cache.0 - value_cache_outputs: - - updated_value_cache.0 pipeline: workflow: manifest: @@ -366,6 +352,97 @@ pipeline: implementation: kind: onnx artifact: model.onnx + ports: + inputs: + input_ids: + dtype: int64 + rank: 2 + shape: + - batch + - sequence + batch_layout: + kind: request_aligned + axis: 0 + position_ids: + dtype: int64 + rank: 2 + shape: + - batch + - sequence + 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 + write_indices: + dtype: int64 + rank: 1 + shape: + - batch + batch_layout: + kind: request_aligned + axis: 0 + nonpad_kv_seqlen: + dtype: int64 + rank: 1 + shape: + - batch + batch_layout: + kind: request_aligned + axis: 0 + outputs: + logits: + dtype: float32 + rank: 3 + shape: + - batch + - sequence + - 128 + batch_layout: + kind: request_aligned + axis: 0 + updated_key_cache.0: + dtype: float32 + rank: 3 + shape: + - batch + - 16 + - 8 + batch_layout: + kind: request_aligned + axis: 0 + updated_value_cache.0: + dtype: float32 + rank: 3 + shape: + - batch + - 16 + - 8 + batch_layout: + kind: request_aligned + axis: 0 + roles: + input_ids: token_ids + position_ids: position_ids + logits: logits token_sampler: implementation: kind: onnx @@ -1148,6 +1225,8 @@ pipeline: capacity: package.cache_capacity write_indices_ports: model: write_indices + kv_length_ports: + model: nonpad_kv_seqlen aliasing: permitted reuse: prefix_reusable: true @@ -1157,9 +1236,13 @@ pipeline: 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: diff --git a/tests/fixtures/onnx_genai_workflows/tts/inference_metadata.yaml b/tests/fixtures/onnx_genai_workflows/tts/inference_metadata.yaml index d382059f6..eda94d097 100644 --- a/tests/fixtures/onnx_genai_workflows/tts/inference_metadata.yaml +++ b/tests/fixtures/onnx_genai_workflows/tts/inference_metadata.yaml @@ -251,42 +251,667 @@ pipeline: implementation: kind: onnx artifact: talker/model.onnx + ports: + inputs: + inputs_embeds: + dtype: float32 + rank: 3 + shape: + - batch + - sequence_len + - 8 + batch_layout: + kind: request_aligned + axis: 0 + 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 + 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 + outputs: + logits: + dtype: float32 + rank: 3 + shape: + - batch + - sequence_len + - 2160 + batch_layout: + kind: request_aligned + axis: 0 + last_hidden_state: + dtype: float32 + rank: 3 + shape: + - batch + - 1 + - 8 + batch_layout: + kind: request_aligned + axis: 0 + present.0.key: + dtype: float32 + rank: 4 + shape: + - batch + - 1 + - total_sequence_len + - 4 + batch_layout: + kind: request_aligned + axis: 0 + present.0.value: + dtype: float32 + rank: 4 + shape: + - batch + - 1 + - total_sequence_len + - 4 + batch_layout: + kind: request_aligned + axis: 0 + 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: + inputs: + inputs_embeds: + dtype: float32 + rank: 3 + shape: + - batch + - sequence_len + - 8 + batch_layout: + kind: request_aligned + axis: 0 + step_index: + dtype: int64 + rank: 0 + shape: [] + 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 + 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 + outputs: + logits: + dtype: float32 + rank: 3 + shape: + - batch + - sequence_len + - 6 + batch_layout: + kind: request_aligned + axis: 0 + codec_embeddings: + dtype: float32 + rank: 3 + shape: + - 3 + - 6 + - 8 + present.0.key: + dtype: float32 + rank: 4 + shape: + - batch + - 8 + - total_sequence_len + - 128 + batch_layout: + kind: request_aligned + axis: 0 + present.0.value: + dtype: float32 + rank: 4 + shape: + - batch + - 8 + - total_sequence_len + - 128 + batch_layout: + kind: request_aligned + axis: 0 + present.1.key: + dtype: float32 + rank: 4 + shape: + - batch + - 8 + - total_sequence_len + - 128 + batch_layout: + kind: request_aligned + axis: 0 + present.1.value: + dtype: float32 + rank: 4 + shape: + - batch + - 8 + - total_sequence_len + - 128 + batch_layout: + kind: request_aligned + axis: 0 + present.2.key: + dtype: float32 + rank: 4 + shape: + - batch + - 8 + - total_sequence_len + - 128 + batch_layout: + kind: request_aligned + axis: 0 + present.2.value: + dtype: float32 + rank: 4 + shape: + - batch + - 8 + - total_sequence_len + - 128 + batch_layout: + kind: request_aligned + axis: 0 + present.3.key: + dtype: float32 + rank: 4 + shape: + - batch + - 8 + - total_sequence_len + - 128 + batch_layout: + kind: request_aligned + axis: 0 + present.3.value: + dtype: float32 + rank: 4 + shape: + - batch + - 8 + - total_sequence_len + - 128 + batch_layout: + kind: request_aligned + axis: 0 + present.4.key: + dtype: float32 + rank: 4 + shape: + - batch + - 8 + - total_sequence_len + - 128 + batch_layout: + kind: request_aligned + axis: 0 + present.4.value: + dtype: float32 + rank: 4 + shape: + - batch + - 8 + - total_sequence_len + - 128 + batch_layout: + kind: request_aligned + axis: 0 + roles: + inputs_embeds: inputs_embeds + attention_mask: attention_mask + position_ids: position_ids + logits: logits embedding: implementation: kind: onnx artifact: embedding/model.onnx + ports: + inputs: + text_ids: + dtype: int64 + rank: 2 + shape: + - batch + - text_sequence_len + batch_layout: + kind: request_aligned + axis: 0 + codec_ids: + dtype: int64 + rank: 2 + shape: + - batch + - codec_sequence_len + batch_layout: + kind: request_aligned + axis: 0 + outputs: + text_embeds: + dtype: float32 + rank: 3 + shape: + - batch + - text_sequence_len + - 8 + batch_layout: + kind: request_aligned + axis: 0 + codec_embeds: + dtype: float32 + rank: 3 + shape: + - batch + - codec_sequence_len + - 8 + batch_layout: + kind: request_aligned + axis: 0 talker_step_embedder: implementation: kind: onnx artifact: talker_step_embedder/model.onnx + ports: + inputs: + frame_codes: + dtype: int64 + rank: 2 + shape: + - batch + - 4 + batch_layout: + kind: request_aligned + axis: 0 + text_embed: + dtype: float32 + rank: 3 + shape: + - batch + - 1 + - 8 + batch_layout: + kind: request_aligned + axis: 0 + outputs: + inputs_embeds: + dtype: float32 + rank: 3 + shape: + - batch + - 1 + - 8 + batch_layout: + kind: request_aligned + axis: 0 + roles: + inputs_embeds: inputs_embeds talker_prefill_embedder: implementation: kind: onnx artifact: talker_prefill_embedder/model.onnx + ports: + inputs: + text_ids: + dtype: int64 + rank: 2 + shape: + - batch + - text_sequence_len + batch_layout: + kind: request_aligned + axis: 0 + outputs: + prefill_embeds: + dtype: float32 + rank: 3 + shape: + - batch + - prefill_sequence_len + - 8 + batch_layout: + kind: request_aligned + axis: 0 + trailing_text_embeds: + dtype: float32 + rank: 3 + shape: + - batch + - trailing_sequence_len + - 8 + batch_layout: + kind: request_aligned + axis: 0 code_predictor_prefill: implementation: kind: onnx artifact: code_predictor_prefill/model.onnx + ports: + inputs: + talker_hidden: + dtype: float32 + rank: 3 + shape: + - batch + - 1 + - 8 + batch_layout: + kind: request_aligned + axis: 0 + group_0_embed: + dtype: float32 + rank: 3 + shape: + - batch + - 1 + - 8 + batch_layout: + kind: request_aligned + axis: 0 + outputs: + inputs_embeds: + dtype: float32 + rank: 3 + shape: + - batch + - 2 + - 8 + batch_layout: + kind: request_aligned + axis: 0 + roles: + inputs_embeds: inputs_embeds code_predictor_step_embedder: implementation: kind: onnx artifact: code_predictor_step_embedder/model.onnx + ports: + inputs: + codec_embeddings: + dtype: float32 + rank: 3 + shape: + - 3 + - 6 + - 8 + token: + dtype: int64 + rank: 1 + shape: + - batch + batch_layout: + kind: request_aligned + axis: 0 + embedding_index: + dtype: int64 + rank: 0 + shape: [] + outputs: + inputs_embeds: + dtype: float32 + rank: 3 + shape: + - batch + - 1 + - 8 + batch_layout: + kind: request_aligned + axis: 0 + roles: + inputs_embeds: inputs_embeds code_predictor_indices: implementation: kind: onnx artifact: code_predictor_indices/model.onnx + ports: + inputs: + iteration: + dtype: int64 + rank: 0 + shape: [] + outputs: + embedding_index: + dtype: int64 + rank: 0 + shape: [] + step_index: + dtype: int64 + rank: 0 + shape: [] + frame_index: + dtype: int64 + rank: 0 + shape: [] talker_text_step: implementation: kind: onnx artifact: talker_text_step/model.onnx + ports: + inputs: + trailing_text_embeds: + dtype: float32 + rank: 3 + shape: + - batch + - trailing_sequence + - 8 + batch_layout: + kind: request_aligned + axis: 0 + iteration: + dtype: int64 + rank: 1 + shape: + - batch + batch_layout: + kind: request_aligned + axis: 0 + outputs: + text_embed: + dtype: float32 + rank: 3 + shape: + - batch + - 1 + - 8 + batch_layout: + kind: request_aligned + axis: 0 codec: implementation: kind: onnx artifact: codec/model.onnx + ports: + inputs: + codes: + dtype: int64 + rank: 3 + shape: + - batch + - 4 + - frames + batch_layout: + kind: request_aligned + axis: 0 + outputs: + waveform: + dtype: float32 + rank: 3 + shape: + - batch + - 1 + - frames + batch_layout: + kind: request_aligned + axis: 0 last_token_logits: implementation: kind: onnx @@ -1509,9 +2134,13 @@ pipeline: 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 @@ -1529,33 +2158,53 @@ pipeline: 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 diff --git a/tests/fixtures/onnx_genai_workflows/video/inference_metadata.yaml b/tests/fixtures/onnx_genai_workflows/video/inference_metadata.yaml index 1146780ea..95a382664 100644 --- a/tests/fixtures/onnx_genai_workflows/video/inference_metadata.yaml +++ b/tests/fixtures/onnx_genai_workflows/video/inference_metadata.yaml @@ -159,10 +159,132 @@ pipeline: implementation: kind: onnx artifact: transformer/model.onnx + ports: + inputs: + sample: + dtype: float32 + rank: 5 + shape: + - batch + - num_frames + - 4 + - height + - width + batch_layout: + kind: request_aligned + axis: 0 + timestep: + dtype: int64 + rank: 1 + shape: + - batch + batch_layout: + kind: request_aligned + axis: 0 + encoder_hidden_states: + dtype: float32 + rank: 3 + shape: + - batch + - prompt_sequence + - 32 + batch_layout: + kind: request_aligned + axis: 0 + outputs: + noise_pred: + dtype: float32 + rank: 5 + shape: + - batch + - num_frames + - 4 + - height + - width + batch_layout: + kind: request_aligned + axis: 0 + roles: + encoder_hidden_states: encoder_hidden_states vae_decoder: implementation: kind: onnx artifact: vae_decoder/model.onnx + ports: + inputs: + latent_sample: + dtype: float32 + rank: 5 + shape: + - batch + - 4 + - latent_frames + - latent_height + - latent_width + batch_layout: + kind: request_aligned + axis: 0 + conv_cache.conv_in: + dtype: float32 + rank: 5 + shape: + - batch + - 4 + - cache_frames + - latent_height + - latent_width + batch_layout: + kind: request_aligned + axis: 0 + conv_cache.conv_out: + dtype: float32 + rank: 5 + shape: + - batch + - 3 + - cache_frames + - 2*latent_height + - 2*latent_width + batch_layout: + kind: request_aligned + axis: 0 + outputs: + sample: + dtype: float32 + rank: 5 + shape: + - batch + - 3 + - frames + - 2*latent_height + - 2*latent_width + batch_layout: + kind: request_aligned + axis: 0 + conv_cache_out.conv_in: + dtype: float32 + rank: 5 + shape: + - batch + - 4 + - cache_frames + - latent_height + - latent_width + batch_layout: + kind: request_aligned + axis: 0 + conv_cache_out.conv_out: + dtype: float32 + rank: 5 + shape: + - batch + - 3 + - cache_frames + - 2*latent_height + - 2*latent_width + batch_layout: + kind: request_aligned + axis: 0 model_input: implementation: kind: onnx diff --git a/tests/fixtures/onnx_genai_workflows/vlm/inference_metadata.yaml b/tests/fixtures/onnx_genai_workflows/vlm/inference_metadata.yaml index f56d4c0b5..2e3582ccb 100644 --- a/tests/fixtures/onnx_genai_workflows/vlm/inference_metadata.yaml +++ b/tests/fixtures/onnx_genai_workflows/vlm/inference_metadata.yaml @@ -423,14 +423,164 @@ pipeline: implementation: kind: onnx artifact: vision_encoder/model.onnx + ports: + inputs: + pixel_values: + dtype: float32 + rank: 2 + shape: + - 4 + - 1176 + grid_thw: + dtype: int64 + rank: 2 + shape: + - 1 + - 3 + outputs: + image_features: + dtype: float32 + rank: 3 + shape: + - batch + - 4 + - 32 + batch_layout: + kind: request_aligned + axis: 0 embedding: implementation: kind: onnx artifact: embedding/model.onnx + ports: + inputs: + input_ids: + dtype: int64 + rank: 2 + shape: + - batch + - sequence + batch_layout: + kind: request_aligned + axis: 0 + image_features: + dtype: float32 + rank: 3 + shape: + - batch + - 4 + - 32 + batch_layout: + kind: request_aligned + axis: 0 + outputs: + inputs_embeds: + dtype: float32 + rank: 3 + shape: + - batch + - sequence + - 32 + batch_layout: + kind: request_aligned + axis: 0 + roles: + input_ids: token_ids + inputs_embeds: inputs_embeds decoder: implementation: kind: onnx artifact: decoder/model.onnx + ports: + inputs: + inputs_embeds: + dtype: float32 + rank: 3 + shape: + - batch + - sequence + - 32 + batch_layout: + kind: request_aligned + axis: 0 + 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 + - sequence + 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 + outputs: + logits: + dtype: float32 + rank: 3 + shape: + - batch + - sequence + - 128 + batch_layout: + kind: request_aligned + axis: 0 + present.0.key: + dtype: float32 + rank: 4 + shape: + - batch + - 2 + - total_sequence + - 8 + batch_layout: + kind: request_aligned + axis: 0 + present.0.value: + dtype: float32 + rank: 4 + shape: + - batch + - 2 + - total_sequence + - 8 + batch_layout: + kind: request_aligned + axis: 0 + roles: + inputs_embeds: inputs_embeds + attention_mask: attention_mask + position_ids: position_ids + logits: logits image_preprocess: implementation: kind: adapter @@ -1296,9 +1446,13 @@ pipeline: 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: diff --git a/tests/generate_onnx_genai_validation_packages.py b/tests/generate_onnx_genai_validation_packages.py index 1b23feb9a..2515c4de4 100644 --- a/tests/generate_onnx_genai_validation_packages.py +++ b/tests/generate_onnx_genai_validation_packages.py @@ -880,7 +880,28 @@ def _write_adapter_metadata(package: ModelPackage, directory: Path) -> None: "kind": "onnx", "artifact": "model.onnx", }, - "ports": {}, + # Every ONNX component declares its ports, including the + # ones a hand-written fixture builds: the workflow is the + # only description a package has of itself, and a + # component that declares nothing describes nothing. + "ports": { + "inputs": { + "activations": { + "dtype": "float32", + "rank": 2, + "shape": ["batch", 2], + "batch_layout": {"kind": "request_aligned", "axis": 0}, + } + }, + "outputs": { + "projection.output": { + "dtype": "float32", + "rank": 2, + "shape": ["batch", 2], + "batch_layout": {"kind": "request_aligned", "axis": 0}, + } + }, + }, }, "overlay": { "implementation": { diff --git a/tests/static_cache_metadata_test.py b/tests/static_cache_metadata_test.py index 2c3676b20..05351a09e 100644 --- a/tests/static_cache_metadata_test.py +++ b/tests/static_cache_metadata_test.py @@ -70,6 +70,46 @@ def _model_invoke(steps) -> dict[str, str]: 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. + """ + workflow = metadata["pipeline"]["workflow"] + _, group = _scatter_group(metadata) + update = group["update"] + component = next(iter(update["write_indices_ports"])) + declared = workflow["components"][component]["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], + "declared_inputs": declared["inputs"], + "declared_outputs": declared["outputs"], + } + + @pytest.fixture(scope="module") def static_workflow(): pkg, config = _static_package() @@ -91,46 +131,46 @@ def mixed(): return build_decoder_workflow_metadata(pkg, config) -class TestStaticCacheModelIo: - """``model.io.static_cache`` is the authoritative port ABI.""" - - def test_declares_control_and_buffer_ports(self, static_workflow): - static_cache = static_workflow["model"]["io"]["static_cache"] - assert static_cache["write_indices_input"] == STATIC_CACHE_WRITE_INDICES - assert static_cache["kv_sequence_length_input"] == STATIC_CACHE_KV_SEQUENCE_LENGTH - assert static_cache["key_cache_inputs"] == ["key_cache.0", "key_cache.1"] - assert static_cache["value_cache_inputs"] == ["value_cache.0", "value_cache.1"] - assert static_cache["key_cache_outputs"] == [ - "updated_key_cache.0", - "updated_key_cache.1", - ] - assert static_cache["value_cache_outputs"] == [ - "updated_value_cache.0", - "updated_value_cache.1", +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): + # Naming a port the component does not expose would bind nothing. + abi = _static_cache_abi(static_workflow) + for role in ("write_indices_input", "kv_sequence_length_input"): + contract = abi["declared_inputs"][abi[role]] + # Both are per-row integer vectors, which is exactly why they are + # declared rather than recognized by shape. + assert contract["dtype"] == "int64" + assert contract["rank"] == 1 + + def test_every_buffer_pair_is_declared(self, static_workflow): + 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"]): + declared = abi["declared_inputs"] if name in abi["cache_inputs"] else None + declared = declared or abi["declared_outputs"] + assert declared[name]["shape"][STATIC_CACHE_SEQUENCE_AXIS] == CAPACITY - def test_per_layer_lists_are_paired(self, static_workflow): - # A runtime binds these four lists positionally; unequal lengths would - # silently pair layer i's key buffer with layer j's output. - static_cache = static_workflow["model"]["io"]["static_cache"] - lengths = { - len(static_cache[key]) - for key in ( - "key_cache_inputs", - "value_cache_inputs", - "key_cache_outputs", - "value_cache_outputs", - ) - } - assert lengths == {2} - - def test_owns_its_cache_and_declares_no_appending_ports(self, static_workflow): - io = static_workflow["model"]["io"] - assert io["kv_ownership"] == "owned" - # A static cache has no past/present pair to advertise; declaring one - # would invite a runtime to concatenate into a fixed buffer. - assert "kv_inputs" not in io - assert "kv_outputs" not in io + def test_capacity_is_recoverable(self, static_workflow): + assert _static_cache_abi(static_workflow)["capacity"] == CAPACITY class TestStaticCacheWorkflow: @@ -182,6 +222,7 @@ def test_state_service_publishes_an_indexed_scatter_discipline(self, static_work "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 @@ -194,6 +235,8 @@ def test_state_service_publishes_an_indexed_scatter_discipline(self, static_work 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): @@ -227,8 +270,15 @@ def test_grouped_query_layouts_are_declared_flat(self, static_workflow): def test_layer_count_follows_the_config(self): pkg, config = _static_package(num_hidden_layers=3) - io = build_decoder_workflow_metadata(pkg, config)["model"]["io"]["static_cache"] - assert io["key_cache_inputs"] == ["key_cache.0", "key_cache.1", "key_cache.2"] + 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: @@ -244,11 +294,9 @@ class TestHeterogeneousStaticCache: 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. - static_cache = mixed["model"]["io"]["static_cache"] - assert static_cache["key_cache_inputs"] == ["key_cache.1"] - assert static_cache["value_cache_inputs"] == ["value_cache.1"] - assert static_cache["key_cache_outputs"] == ["updated_key_cache.1"] - assert static_cache["value_cache_outputs"] == ["updated_value_cache.1"] + 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"] @@ -393,13 +441,11 @@ def test_static_cache_survives_cuda_optimization(self): pkg["model"], ep="cuda", dtype=ir.DataType.FLOAT16, model_role="decoder" ) metadata = build_decoder_workflow_metadata(pkg, config) - group = next( - iter( - metadata["pipeline"]["workflow"]["serving"]["state_service"]["groups"].values() - ) - ) + _, group = _scatter_group(metadata) assert group["update"]["kind"] == "indexed_scatter" - assert metadata["model"]["io"]["static_cache"]["key_cache_inputs"] == [ + assert _static_cache_abi(metadata)["cache_inputs"] == [ "key_cache.0", + "value_cache.0", "key_cache.1", + "value_cache.1", ] From d8da2459ff107f7a53fb7877b4fb6ca04440271c Mon Sep 17 00:00:00 2001 From: justinchuby Date: Thu, 20 Aug 2026 20:08:15 +0000 Subject: [PATCH 131/151] Declare a role for every component a task builds MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Rebasing onto main brought in inspect_components, which reports a package's components from the task's model_roles. Qwen3-TTS builds four loop-wiring graphs — code_predictor_prefill, code_predictor_step_embedder, code_predictor_indices and talker_text_step — that were never declared there. Two things followed from that omission. build_from_module looks each component up in model_roles and falls back to the "decoder" role when it is absent, so these parameter-free graphs were offered the GQA and QKV-packing passes meant for attention stacks. And inspect_components reported six components for a package that contains ten, so a caller planning per-component work would silently skip four artifacts that ModelPackage.save writes to disk. Declare them with a "glue" role: a graph that reads every tensor it uses from its own graph inputs, carries no weights, and has no fusion to gate. Give them empty HF_COMPONENT_SOURCES entries, because mobius authors them and no HuggingFace sub-module backs them. Guard both directions. arch_validation_test now fails any task that builds a component it does not declare, and exempts glue components from the "must have initializers" check — a wiring graph holding only hoisted constants is what correctness looks like, not a lost-weights bug. A network-free unit test pins the same invariant for Qwen3-TTS in the fast suite. The rebase also resolved two conflicts semantically: __all__ keeps main's inspect_components alongside the branch's generation and fingerprint_model_weights, and unet_parity_test keeps main's _run_onnx helper while taking the branch's FLOAT timestep, which is the dtype the denoiser port now declares. Repin the onnx-genai validation checkout to f19f2e71, the current head of the schema branch after its rebase; the previous pin is no longer on that lineage. 11/11 fixtures validate and 11/11 conformance tests execute against it. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: justinchuby --- .github/workflows/main.yml | 2 +- CHANGELOG.md | 13 +++++++++++ docs/onnx-genai-performance-conformance.md | 2 +- src/mobius/_inspect.py | 8 ++++--- src/mobius/_inspect_test.py | 16 ++++++++++--- src/mobius/models/qwen3_tts.py | 6 +++++ src/mobius/models/qwen3_tts_test.py | 24 ++++++++++++++++++++ src/mobius/tasks/_tts.py | 11 +++++++++ tests/arch_validation_test.py | 26 +++++++++++++++++----- 9 files changed, 95 insertions(+), 13 deletions(-) diff --git a/.github/workflows/main.yml b/.github/workflows/main.yml index 773933ea0..a3004ac35 100644 --- a/.github/workflows/main.yml +++ b/.github/workflows/main.yml @@ -23,7 +23,7 @@ jobs: - uses: actions/checkout@v7 with: repository: justinchuby/onnx-genai - ref: 02e22dd6516f3851261ebfc5de8b6a93c48694f9 + ref: f19f2e7145698bfe4fa25b0df909f334f2304757 path: validation/onnx-genai - uses: actions/setup-python@v7 with: diff --git a/CHANGELOG.md b/CHANGELOG.md index 965e36757..68e1a85c5 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -7,6 +7,19 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 ## [Unreleased] +### Fixed + +- 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 diff --git a/docs/onnx-genai-performance-conformance.md b/docs/onnx-genai-performance-conformance.md index 8664d3088..88a967e94 100644 --- a/docs/onnx-genai-performance-conformance.md +++ b/docs/onnx-genai-performance-conformance.md @@ -123,7 +123,7 @@ well-formed and validate; only the local kernel is missing. ### Canonical representation — lowering verified -Against ONNX GenAI `02e22dd6`, with no `model.io` in any package: +Against ONNX GenAI `f19f2e71`, with no `model.io` in any package: | Check | Result | | --- | --- | 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 a7f4dade2..2da97bde1 100644 --- a/src/mobius/_inspect_test.py +++ b/src/mobius/_inspect_test.py @@ -239,7 +239,8 @@ def test_explicit_task_without_config_returns_roles_without_source_paths(monkeyp staticmethod(lambda *a, **k: (_ for _ in ()).throw(OSError("no config"))), ) monkeypatch.setattr( - "mobius.integrations.transformers._config_resolver._try_load_config_json", lambda *_a, **_k: None + "mobius.integrations.transformers._config_resolver._try_load_config_json", + lambda *_a, **_k: None, ) components = inspect_components("anything", task="vision-language") @@ -280,7 +281,8 @@ def test_unresolvable_config_raises(monkeypatch): staticmethod(lambda *a, **k: (_ for _ in ()).throw(OSError("no config"))), ) monkeypatch.setattr( - "mobius.integrations.transformers._config_resolver._try_load_config_json", lambda *_a, **_k: None + "mobius.integrations.transformers._config_resolver._try_load_config_json", + lambda *_a, **_k: None, ) with pytest.raises(ValueError, match="Could not load a HuggingFace config"): inspect_components("fake/diffusers-pipeline") @@ -320,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/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 7cc3eb301..cec831214 100644 --- a/src/mobius/models/qwen3_tts_test.py +++ b/src/mobius/models/qwen3_tts_test.py @@ -169,6 +169,30 @@ def test_talker_text_step_clamps_to_last_trailing_embedding(): 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/tasks/_tts.py b/src/mobius/tasks/_tts.py index fb05403b1..73ec89600 100644 --- a/src/mobius/tasks/_tts.py +++ b/src/mobius/tasks/_tts.py @@ -57,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", @@ -64,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", 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" From c7e453a53b1e3025688816c9f629f418fc374f8f Mon Sep 17 00:00:00 2001 From: justinchuby Date: Thu, 20 Aug 2026 20:54:27 +0000 Subject: [PATCH 132/151] Pin GenAI validation to the branch head, not an ancestor of it The validation checkout pinned f19f2e71, which is a reachable ancestor of justinchuby/onnx-genai@justinchuby/simplify-composite-metadata but no longer its head. An ancestor pin is only as durable as the branch's history: that branch has been force-pushed repeatedly, and each rewrite risks orphaning the commit CI resolves, which would fail the checkout rather than the assertion. Pinning the head keeps the reference on the branch for as long as possible and makes the SHA verifiable with a single git ls-remote. c344c2c7 adds only documentation and a test doc-comment on top of f19f2e71 -- no schema, validator or runtime change -- so it cannot move the contract Mobius emits against. Re-verified against the new pin rather than assumed: all 11 emitted packages validate, and all 11 workflow conformance tests pass, including mobius_static_cache_workflow_executes. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: justinchuby --- .github/workflows/main.yml | 2 +- docs/onnx-genai-performance-conformance.md | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/.github/workflows/main.yml b/.github/workflows/main.yml index a3004ac35..c9a3a5752 100644 --- a/.github/workflows/main.yml +++ b/.github/workflows/main.yml @@ -23,7 +23,7 @@ jobs: - uses: actions/checkout@v7 with: repository: justinchuby/onnx-genai - ref: f19f2e7145698bfe4fa25b0df909f334f2304757 + ref: c344c2c7af19aed7a18b719ed82270a818a2c84b path: validation/onnx-genai - uses: actions/setup-python@v7 with: diff --git a/docs/onnx-genai-performance-conformance.md b/docs/onnx-genai-performance-conformance.md index 88a967e94..147cff3ef 100644 --- a/docs/onnx-genai-performance-conformance.md +++ b/docs/onnx-genai-performance-conformance.md @@ -123,7 +123,7 @@ well-formed and validate; only the local kernel is missing. ### Canonical representation — lowering verified -Against ONNX GenAI `f19f2e71`, with no `model.io` in any package: +Against ONNX GenAI `c344c2c7`, with no `model.io` in any package: | Check | Result | | --- | --- | From 16fd0822a8fbd07d386bdcd1e6353466df2d673f Mon Sep 17 00:00:00 2001 From: justinchuby Date: Thu, 20 Aug 2026 21:50:26 +0000 Subject: [PATCH 133/151] Cover the cache layer annotation where it can actually be wrong Every package these tests build has two layers, and two layers cannot distinguish a correct layer annotation from a broken one. Below ten cells a label sorts the same way under every rule -- lexicographic, numeric, insertion order -- so a `layer` taken from a cell's position instead of parsed from its port name satisfies each existing assertion exactly as the correct value does. Real decoders have twenty to eighty layers, which puts production permanently in the region the suite could not reach. The failure that region hides is silent. A layer's key and value buffers share a shape and a dtype, so two transposed caches raise nothing: no validation error, no shape mismatch, only subtly wrong tokens. So build twelve-layer dynamic, static-cache and hybrid decoders and assert what only holds there: that the labels really do sort out of layer order at this depth (otherwise the rest proves nothing), that each declared layer restates the index its port name carries, that ordering by layer and half recovers the buffer lists the runtime collects, and that a hybrid's two alternating groups own layers a cell's position within its group never equals. Verified by mutation rather than assumed: dropping the parse from _cache_layer_index leaves all 101 pre-existing metadata assertions green and fails seven of these, across all three shapes. Also pin GenAI validation to 0d4738e7, which makes both facts these tests rely on enforceable rather than merely documented -- a group binding key/value roles must now declare kv_length_ports, and every such alias must carry a layer. A real twelve-layer static-cache export validates against it, declaring layers 0 through 11; all 11 fixtures validate and all 11 conformance tests pass. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: justinchuby --- .github/workflows/main.yml | 2 +- CHANGELOG.md | 12 ++ docs/onnx-genai-performance-conformance.md | 2 +- tests/canonical_workflow_contract_test.py | 144 +++++++++++++++++++++ 4 files changed, 158 insertions(+), 2 deletions(-) diff --git a/.github/workflows/main.yml b/.github/workflows/main.yml index c9a3a5752..ad2e03b52 100644 --- a/.github/workflows/main.yml +++ b/.github/workflows/main.yml @@ -23,7 +23,7 @@ jobs: - uses: actions/checkout@v7 with: repository: justinchuby/onnx-genai - ref: c344c2c7af19aed7a18b719ed82270a818a2c84b + ref: 0d4738e70f3b25fdd2623d198557dc9db30f3262 path: validation/onnx-genai - uses: actions/setup-python@v7 with: diff --git a/CHANGELOG.md b/CHANGELOG.md index 68e1a85c5..760089aa9 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -9,6 +9,18 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 ### Fixed +- 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`) diff --git a/docs/onnx-genai-performance-conformance.md b/docs/onnx-genai-performance-conformance.md index 147cff3ef..85fffd5ea 100644 --- a/docs/onnx-genai-performance-conformance.md +++ b/docs/onnx-genai-performance-conformance.md @@ -123,7 +123,7 @@ well-formed and validate; only the local kernel is missing. ### Canonical representation — lowering verified -Against ONNX GenAI `c344c2c7`, with no `model.io` in any package: +Against ONNX GenAI `0d4738e7`, with no `model.io` in any package: | Check | Result | | --- | --- | diff --git a/tests/canonical_workflow_contract_test.py b/tests/canonical_workflow_contract_test.py index bfaaadc12..481a8aec5 100644 --- a/tests/canonical_workflow_contract_test.py +++ b/tests/canonical_workflow_contract_test.py @@ -26,6 +26,7 @@ import glob import os +import re from typing import Any import onnx_ir as ir @@ -43,6 +44,11 @@ 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") @@ -333,6 +339,144 @@ def test_the_decode_step_is_recoverable_from_the_workflow_alone(self, package): ) +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) From c3dfeca64ca9ea3fcde5ddd5fadeb15e603eeabf Mon Sep 17 00:00:00 2001 From: justinchuby Date: Thu, 20 Aug 2026 23:28:49 +0000 Subject: [PATCH 134/151] Pin GenAI validation to the workflow-derived scatter proof 1d8cfefe is additive over the previously pinned 0d4738e7: it adds an end-to-end test that a package carrying no `model:` block still resolves the static-cache ABI and executes through the scatter driver, plus the canonical fixture that test runs against. No validator or schema behaviour changes. That test is worth pinning because it closes the last gap between what this producer emits and what is proven to run. Until now the evidence that a workflow-only package was sufficient came from validation and from Mobius's own conformance run; the ABI resolution itself was argued from call sites. It is now asserted against a real graph on the consumer side. Re-verified against the new pin rather than inherited from the previous run: all 11 emitted packages validate and all 11 conformance tests pass. Compared the emitted static-cache package against the new canonical fixture directly -- `update.write_indices_ports`, `update.kv_length_ports`, per-alias `role` and `layer`, `components..ports.roles` and the absence of a top-level `model:` block all agree, so the two describe the same contract. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: justinchuby --- .github/workflows/main.yml | 2 +- docs/onnx-genai-performance-conformance.md | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/.github/workflows/main.yml b/.github/workflows/main.yml index ad2e03b52..84bd60b19 100644 --- a/.github/workflows/main.yml +++ b/.github/workflows/main.yml @@ -23,7 +23,7 @@ jobs: - uses: actions/checkout@v7 with: repository: justinchuby/onnx-genai - ref: 0d4738e70f3b25fdd2623d198557dc9db30f3262 + ref: 1d8cfefe98f45a5cc791a09e45a4b164466e9346 path: validation/onnx-genai - uses: actions/setup-python@v7 with: diff --git a/docs/onnx-genai-performance-conformance.md b/docs/onnx-genai-performance-conformance.md index 85fffd5ea..6551ff648 100644 --- a/docs/onnx-genai-performance-conformance.md +++ b/docs/onnx-genai-performance-conformance.md @@ -123,7 +123,7 @@ well-formed and validate; only the local kernel is missing. ### Canonical representation — lowering verified -Against ONNX GenAI `0d4738e7`, with no `model.io` in any package: +Against ONNX GenAI `1d8cfefe`, with no `model.io` in any package: | Check | Result | | --- | --- | From 564af035e193e537a92ebbb10cfb6cd891c8c9bc Mon Sep 17 00:00:00 2001 From: justinchuby Date: Thu, 20 Aug 2026 23:56:10 +0000 Subject: [PATCH 135/151] Re-pin GenAI validation after the upstream branch was rebased The previously pinned 1d8cfefe is no longer reachable from justinchuby/onnx-genai@justinchuby/simplify-composite-metadata: that branch was rebased onto current main, so every SHA on it changed and the old head became an orphan. `git merge-base --is-ancestor 1d8cfefe 60f41354` is false. CI still resolved the old SHA because GitHub serves unreferenced objects, which is the failure mode worth avoiding -- the pin kept working while pointing at a commit that is on no branch and whose content nobody is maintaining. 60f41354 is the current head of that branch and of PR #828. The rebase carried no content change, so the two facts this pin exists to check are unchanged: the workflow-derived scatter proof and the kv_length_ports / role / layer enforcement. Re-verified against the new head rather than carried over: all 11 emitted packages validate and all 11 workflow conformance tests pass. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: justinchuby --- .github/workflows/main.yml | 2 +- docs/onnx-genai-performance-conformance.md | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/.github/workflows/main.yml b/.github/workflows/main.yml index 84bd60b19..207342fcc 100644 --- a/.github/workflows/main.yml +++ b/.github/workflows/main.yml @@ -23,7 +23,7 @@ jobs: - uses: actions/checkout@v7 with: repository: justinchuby/onnx-genai - ref: 1d8cfefe98f45a5cc791a09e45a4b164466e9346 + ref: 60f413548e23695975f47be6fe91cef263d5559e path: validation/onnx-genai - uses: actions/setup-python@v7 with: diff --git a/docs/onnx-genai-performance-conformance.md b/docs/onnx-genai-performance-conformance.md index 6551ff648..c0dc3f1a4 100644 --- a/docs/onnx-genai-performance-conformance.md +++ b/docs/onnx-genai-performance-conformance.md @@ -123,7 +123,7 @@ well-formed and validate; only the local kernel is missing. ### Canonical representation — lowering verified -Against ONNX GenAI `1d8cfefe`, with no `model.io` in any package: +Against ONNX GenAI `60f41354`, with no `model.io` in any package: | Check | Result | | --- | --- | From a02bf940ded393e05504f675f2db5ae2f0b0ae4c Mon Sep 17 00:00:00 2001 From: justinchuby Date: Fri, 21 Aug 2026 00:51:21 +0000 Subject: [PATCH 136/151] Re-pin GenAI validation after a second upstream rebase The branch was rebased onto current main again, so 60f41354 is now unreachable from it (`git merge-base --is-ancestor 60f41354 2498e0bc` is false), the same way 1d8cfefe was before it. As noted last time, an orphaned pin keeps reporting green because GitHub serves unreferenced objects, so this has to be caught by checking lineage rather than by waiting for a red job. 2498e0bc also changes how a component's port roles resolve. A declared role now names its port on its own; the port contracts are consulted only to break a tie between two ports claiming the same role. Previously a role was honoured only if the port also appeared in `ports.inputs`, so a producer that declared roles without transcribing its graph had the declaration silently dropped and was matched by port spelling instead. This producer declares both, so it was never affected and nothing here has to change. Keeping the contracts is deliberate: they are optional now, not forbidden, and they live inside `pipeline.workflow`, which is the one canonical representation -- they are not a competing statement of the ABI the way a second top-level block would be. They also carry the dtype, rank and shape a consumer needs to allocate without opening the ONNX file, and the tests that guard against contract drift assert on them. Dropping them would remove information and weaken those guards to satisfy a preference, not a rule. Re-verified against the new head: 11 of 11 packages validate, 11 of 11 conformance tests pass, 111 metadata tests pass. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: justinchuby --- .github/workflows/main.yml | 2 +- docs/onnx-genai-performance-conformance.md | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/.github/workflows/main.yml b/.github/workflows/main.yml index 207342fcc..422c17e77 100644 --- a/.github/workflows/main.yml +++ b/.github/workflows/main.yml @@ -23,7 +23,7 @@ jobs: - uses: actions/checkout@v7 with: repository: justinchuby/onnx-genai - ref: 60f413548e23695975f47be6fe91cef263d5559e + ref: 2498e0bcc524565b80af3c13f28dc56c2dc44e8b path: validation/onnx-genai - uses: actions/setup-python@v7 with: diff --git a/docs/onnx-genai-performance-conformance.md b/docs/onnx-genai-performance-conformance.md index c0dc3f1a4..437c08c3c 100644 --- a/docs/onnx-genai-performance-conformance.md +++ b/docs/onnx-genai-performance-conformance.md @@ -123,7 +123,7 @@ well-formed and validate; only the local kernel is missing. ### Canonical representation — lowering verified -Against ONNX GenAI `60f41354`, with no `model.io` in any package: +Against ONNX GenAI `2498e0bc`, with no `model.io` in any package: | Check | Result | | --- | --- | From 969495393fab4378294871d4d0426306862a7916 Mon Sep 17 00:00:00 2001 From: justinchuby Date: Fri, 21 Aug 2026 01:33:07 +0000 Subject: [PATCH 137/151] Stop transcribing an exported graph into its own workflow component MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit A component backed by a shipped `.onnx` now declares only `ports.roles`. The artifact travels inside the package and is authoritative for which ports exist and what dtype, rank and shape each one has, so restating that in YAML created a second writable statement of one fact with nothing keeping the two in agreement — the same defect as `model.io`, one level further down. The runtime resolves ports against the live session, which catches a name the graph does not expose rather than agreeing with a stale echo of it. What no graph can state is what a port *means*: `input_ids` and `position_ids` are both rank-2 int64, and nothing in the file says which is the autoregressive sequence. That is the whole remaining declaration, and it is what lets the decode ABI resolve without recognizing a spelling. The static-cache ABI is unaffected: it was already derived from the state-service group's aliases and the scatter's `write_indices_ports` / `kv_length_ports`. Policy graphs keep their contracts, and that boundary was measured rather than assumed. A workflow SSA value inherits its dtype, rank and request axis from the port that produced it, so those contracts are the type annotations of the workflow's own dataflow, not a description of an external interface. Dropping them made 4 of 11 packages invalid — `.when is row-wise but .value declares no request_aligned batch_layout` — because a validator reads metadata without the artifacts and has no other source for the axis. The contract tests move with the truth they check. Roles, invocation bindings, state pairs and the two scatter control ports now resolve against the graph itself instead of against the metadata's agreement with its own copy, which is strictly stronger: a role naming a port the artifact does not expose now fails where before it only had to match a line the same producer wrote. `TestRolesAloneCarryTheDecodeAbi` is the regression guard for the omission. Every port of every roled component is renamed to an opaque label and the ABI must resolve identically, so nothing recognizable is left to have matched on; deleting the role table must then break exactly the sequence and logits binding and leave the cache half untouched, which is what proves the roles are carrying the fact rather than decorating it. Verified against ONNX GenAI 2498e0bc (confirmed as the remote branch and PR #828 head): 11/11 packages valid, 11/11 runtime conformance including `mobius_static_cache_workflow_executes`, 4532 fast-suite tests green, lintrunner clean. 1376 lines of transcription removed from the fixtures. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: justinchuby --- CHANGELOG.md | 33 +- docs/onnx-genai-performance-conformance.md | 18 +- .../codec_workflow_metadata_test.py | 14 +- .../onnx_genai/inference_metadata.py | 7 + .../onnx_genai/workflow_metadata.py | 53 +- tests/canonical_workflow_contract_test.py | 476 ++++++++++--- tests/cli_test.py | 21 +- .../adapter/inference_metadata.yaml | 21 - .../codec/inference_metadata.yaml | 46 -- .../decoder/inference_metadata.yaml | 61 -- .../diffusion/inference_metadata.yaml | 88 --- .../diffusion_guided/inference_metadata.yaml | 88 --- .../masked/inference_metadata.yaml | 30 - .../speculative/inference_metadata.yaml | 83 --- .../static_cache/inference_metadata.yaml | 86 --- .../tts/inference_metadata.yaml | 625 +----------------- .../video/inference_metadata.yaml | 119 ---- .../vlm/inference_metadata.yaml | 140 ---- ...generate_onnx_genai_validation_packages.py | 28 +- tests/static_cache_metadata_test.py | 56 +- 20 files changed, 558 insertions(+), 1535 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 760089aa9..f89ee93ff 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -9,6 +9,24 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 ### 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, @@ -49,10 +67,17 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 #### Added -- Every ONNX component declares `ports.inputs` / `ports.outputs` — a contract - (dtype, rank, shape, batch layout) for exactly the graph's inputs and outputs. - A subset would let a consumer fall back to opening the artifact; a superset - would be a promise the graph does not keep. +- 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 diff --git a/docs/onnx-genai-performance-conformance.md b/docs/onnx-genai-performance-conformance.md index 437c08c3c..3165dd09e 100644 --- a/docs/onnx-genai-performance-conformance.md +++ b/docs/onnx-genai-performance-conformance.md @@ -123,7 +123,8 @@ well-formed and validate; only the local kernel is missing. ### Canonical representation — lowering verified -Against ONNX GenAI `2498e0bc`, with no `model.io` in any package: +Against ONNX GenAI `2498e0bc`, with no `model.io` in any package and no port +contracts on any component that ships an artifact: | Check | Result | | --- | --- | @@ -136,6 +137,21 @@ 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. +### 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 diff --git a/src/mobius/integrations/onnx_genai/codec_workflow_metadata_test.py b/src/mobius/integrations/onnx_genai/codec_workflow_metadata_test.py index f55d2f4de..ac52e262c 100644 --- a/src/mobius/integrations/onnx_genai/codec_workflow_metadata_test.py +++ b/src/mobius/integrations/onnx_genai/codec_workflow_metadata_test.py @@ -53,15 +53,11 @@ def test_codec_workflow_has_typed_ssa_and_audio_emit(): "shape": ["batch", 1, "audio_samples"], "batch_layout": {"kind": "request_aligned", "axis": 0}, } - # An ONNX component declares the contract of every port it exposes, so the - # workflow describes the package without opening the artifact. - assert workflow["components"]["encoder"]["ports"]["inputs"]["waveform"] == { - "dtype": "float32", - "rank": 3, - "shape": ["batch", 1, "audio_samples"], - "batch_layout": {"kind": "request_aligned", "axis": 0}, - } - assert set(workflow["components"]["decoder"]["ports"]["outputs"]) == {"waveform"} + # 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"] diff --git a/src/mobius/integrations/onnx_genai/inference_metadata.py b/src/mobius/integrations/onnx_genai/inference_metadata.py index 914fd6bcc..be613152c 100644 --- a/src/mobius/integrations/onnx_genai/inference_metadata.py +++ b/src/mobius/integrations/onnx_genai/inference_metadata.py @@ -1553,6 +1553,13 @@ def tensor_contract(value: Any) -> dict[str, Any]: 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", diff --git a/src/mobius/integrations/onnx_genai/workflow_metadata.py b/src/mobius/integrations/onnx_genai/workflow_metadata.py index 3aeaa17f8..7628183a1 100644 --- a/src/mobius/integrations/onnx_genai/workflow_metadata.py +++ b/src/mobius/integrations/onnx_genai/workflow_metadata.py @@ -191,31 +191,35 @@ def _component( *, effects: tuple[str, ...] = (), ) -> dict[str, Any]: - """Declare one ONNX-backed workflow component: ports, contracts and roles. - - The workflow is the package's only description of itself, so a component - states the contract of every port it exposes rather than leaving a consumer - to open the artifact and infer one. That is what lets the scatter ABI, the - token and logits roles, and the per-layer cache pairs all be resolved from - the workflow alone: an integer control vector is indistinguishable from its - neighbours by shape, so the binding that names it has to sit next to a - declared port for the name to mean anything. - - A contract says what a value *is*; ``roles`` says what the component *does* - with it. 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. - Only ports in this producer's own vocabulary get a role; state ports never - need one, because the group that carries them already names its pairs. + """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 - inputs = {str(value.name): _contract(value) for value in model.graph.inputs} - outputs = {str(value.name): _contract(value) for value in model.graph.outputs} - roles = {name: _PORT_ROLES[name] for name in (*inputs, *outputs) if name in _PORT_ROLES} - ports: dict[str, Any] = {"inputs": inputs, "outputs": outputs} + 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: - ports["roles"] = roles - return {"implementation": {"kind": "onnx", "artifact": artifact}, "ports": ports} + declaration["ports"] = {"roles": roles} + return declaration def _grammar_adapter_component(action: str) -> dict[str, Any]: @@ -302,7 +306,10 @@ def _declare_row_alignment(contract: Any) -> Any: for declaration in workflow.get(section, {}).values(): declaration["contract"] = _declare_row_alignment(declaration.get("contract")) for component in workflow.get("components", {}).values(): - for ports in component.get("ports", {}).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] = {} diff --git a/tests/canonical_workflow_contract_test.py b/tests/canonical_workflow_contract_test.py index 481a8aec5..52914ad7f 100644 --- a/tests/canonical_workflow_contract_test.py +++ b/tests/canonical_workflow_contract_test.py @@ -24,6 +24,7 @@ from __future__ import annotations +import copy import glob import os import re @@ -152,6 +153,30 @@ def _onnx_components(workflow: dict[str, Any]) -> dict[str, dict[str, Any]]: } +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: @@ -187,68 +212,94 @@ def test_no_package_states_its_graph_abi_a_second_time(self, package): _, metadata = package assert "io" not in (metadata.get("model") or {}) - def test_every_graph_declares_exactly_the_ports_it_has(self, package): - """Declared ports are the graph's ports — not a subset, not a superset. + def test_an_exported_graph_is_not_transcribed_into_the_workflow(self, package): + """A model component declares its roles and nothing the artifact says. - A subset lets a runtime silently fall back to opening the artifact, and - a superset is a promise the graph does not keep. Either way the - declaration stops being usable as the single source of truth. + 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"]) - graphs = {name: pkg[name] for name in components if name in pkg} - assert graphs, "a package with no ONNX component describes nothing" - for name, model in graphs.items(): - ports = components[name]["ports"] - assert set(ports["inputs"]) == {str(v.name) for v in model.graph.inputs} - assert set(ports["outputs"]) == {str(v.name) for v in model.graph.outputs} - - def test_declared_contracts_match_the_graph_dtype_and_rank(self, package): - """A contract that disagrees with its graph would mis-size every buffer.""" + 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 - dtypes = { - ir.DataType.FLOAT: "float32", - ir.DataType.FLOAT16: "float16", - ir.DataType.BFLOAT16: "bfloat16", - ir.DataType.INT64: "int64", - ir.DataType.INT32: "int32", - ir.DataType.BOOL: "bool", - ir.DataType.FLOAT8E4M3FN: "float8_e4m3fn", - } - components = _onnx_components(metadata["pipeline"]["workflow"]) - for name in (name for name in components if name in pkg): - model = pkg[name] - ports = components[name]["ports"] - declared = {**ports["inputs"], **ports["outputs"]} - for value in (*model.graph.inputs, *model.graph.outputs): - contract = declared[str(value.name)] - assert contract["rank"] == len(value.shape) - if value.dtype in dtypes: - assert contract["dtype"] == dtypes[value.dtype] - - def test_every_invocation_binds_a_declared_port(self, package): - """A binding to an undeclared port names nothing a consumer can resolve.""" - _, metadata = package workflow = metadata["pipeline"]["workflow"] - components = workflow["components"] + 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": + if step.get("kind") != "invoke" or step["component"] not in graphs: continue - ports = components[step["component"]]["ports"] - assert set(step.get("inputs", {})) <= set(ports["inputs"]) - assert set(step.get("outputs", {})) <= set(ports["outputs"]) + 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_declared_ports(self, package): - """State is carried through ports, so both halves have to be declared.""" - _, metadata = package + 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"] - components = workflow["components"] + graphs = _artifacts(pkg, workflow) for group in _groups(workflow).values(): for component, aliases in (group.get("ports") or {}).items(): - ports = components[component]["ports"] + if component not in graphs: + continue + inputs, outputs = _graph_ports(graphs[component]) for alias in aliases.values(): - assert alias["input"] in ports["inputs"] - assert alias["output"] in ports["outputs"] + 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. @@ -274,17 +325,18 @@ def test_split_cache_halves_and_layers_are_stated_not_positional(self, package): alias["layer"] for alias in values } - def test_control_ports_of_a_fixed_capacity_cache_are_declared_ports(self, package): + 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 that - declares 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. + 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. """ - _, metadata = package + 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": @@ -293,9 +345,9 @@ def test_control_ports_of_a_fixed_capacity_cache_are_declared_ports(self, packag assert set(update["write_indices_ports"]) == bound assert set(update["kv_length_ports"]) == bound for component in bound: - declared = workflow["components"][component]["ports"]["inputs"] - assert update["write_indices_ports"][component] in declared - assert update["kv_length_ports"][component] in declared + 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. @@ -303,32 +355,32 @@ def test_the_decode_step_is_recoverable_from_the_workflow_alone(self, package): 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. + 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. """ - _, metadata = package + pkg, metadata = package workflow = metadata["pipeline"]["workflow"] + graphs = _artifacts(pkg, workflow) decoders = [] for name, component in _onnx_components(workflow).items(): - roles = component["ports"].get("roles", {}) + 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, component, roles)) + decoders.append((name, roles)) assert decoders, "no component declares what it does with the sequence" - for name, component, roles in decoders: - inputs = component["ports"]["inputs"] - outputs = component["ports"]["outputs"] + 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 [port for port, role in roles.items() if role == "logits"] == [ - port for port in outputs if roles.get(port) == "logits" - ] + assert all(port in outputs for port, role in roles.items() if role == "logits") pairs = [ alias for group in _groups(workflow).values() @@ -339,6 +391,229 @@ def test_the_decode_step_is_recoverable_from_the_workflow_alone(self, package): ) +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 _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) @@ -504,31 +779,66 @@ def _metadata(directory: str) -> dict[str, Any]: ) as handle: return yaml.safe_load(handle) + @staticmethod + def _ports(directory: str, workflow: dict[str, Any]) -> dict[str, tuple[set, set]]: + """Load each shipped artifact and read the ports it really exposes.""" + ports: dict[str, tuple[set, set]] = {} + for name, component in _onnx_components(workflow).items(): + artifact = os.path.join(directory, component["implementation"]["artifact"]) + if not os.path.exists(artifact): + continue + ports[name] = _graph_ports(ir.load(artifact)) + return ports + def test_describes_itself_only_through_the_workflow(self, directory): metadata = self._metadata(directory) assert "workflow" in metadata["pipeline"] assert "io" not in (metadata.get("model") or {}) - def test_every_onnx_component_declares_its_ports(self, directory): + def test_no_component_transcribes_the_ports_of_its_own_artifact(self, directory): + """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"] - for component in _onnx_components(workflow).values(): - ports = component["ports"] - assert ports["inputs"] or ports["outputs"] - for contract in (*ports["inputs"].values(), *ports["outputs"].values()): - assert contract["rank"] == len(contract["shape"]) + graphs = self._ports(directory, workflow) + 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(self, directory): + def test_every_binding_and_state_pair_resolves_in_the_artifact(self, directory): workflow = self._metadata(directory)["pipeline"]["workflow"] - components = workflow["components"] + graphs = self._ports(directory, workflow) for step in _walk_steps(workflow["steps"]): - if step.get("kind") != "invoke": + if step.get("kind") != "invoke" or step["component"] not in graphs: continue - ports = components[step["component"]]["ports"] - assert set(step.get("inputs", {})) <= set(ports["inputs"]) - assert set(step.get("outputs", {})) <= set(ports["outputs"]) + 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(): - ports = components[component]["ports"] + if component not in graphs: + continue + inputs, outputs = graphs[component] for alias in aliases.values(): - assert alias["input"] in ports["inputs"] - assert alias["output"] in ports["outputs"] + assert alias["input"] in inputs + assert alias["output"] in outputs + + def test_every_declared_role_resolves_in_the_artifact(self, directory): + workflow = self._metadata(directory)["pipeline"]["workflow"] + graphs = self._ports(directory, workflow) + 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 diff --git a/tests/cli_test.py b/tests/cli_test.py index b8e08bd86..155a35785 100644 --- a/tests/cli_test.py +++ b/tests/cli_test.py @@ -203,6 +203,7 @@ def test_static_cache_with_onnx_genai_runtime_emits_scatter_abi(self): 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: @@ -225,6 +226,16 @@ def test_static_cache_with_onnx_genai_runtime_emits_scatter_abi(self): 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", {}) @@ -240,9 +251,13 @@ def test_static_cache_with_onnx_genai_runtime_emits_scatter_abi(self): assert update["write_indices_ports"] == {"model": "write_indices"} assert update["kv_length_ports"] == {"model": "nonpad_kv_seqlen"} - declared = workflow["components"]["model"]["ports"]["inputs"] - assert declared["write_indices"]["rank"] == 1 - assert declared["nonpad_kv_seqlen"]["rank"] == 1 + # 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"] diff --git a/tests/fixtures/onnx_genai_workflows/adapter/inference_metadata.yaml b/tests/fixtures/onnx_genai_workflows/adapter/inference_metadata.yaml index 86e3b3a75..c2e2ff17b 100644 --- a/tests/fixtures/onnx_genai_workflows/adapter/inference_metadata.yaml +++ b/tests/fixtures/onnx_genai_workflows/adapter/inference_metadata.yaml @@ -108,27 +108,6 @@ pipeline: implementation: kind: onnx artifact: model.onnx - ports: - inputs: - activations: - dtype: float32 - rank: 2 - shape: - - batch - - 2 - batch_layout: - kind: request_aligned - axis: 0 - outputs: - projection.output: - dtype: float32 - rank: 2 - shape: - - batch - - 2 - batch_layout: - kind: request_aligned - axis: 0 overlay: implementation: kind: adapter diff --git a/tests/fixtures/onnx_genai_workflows/codec/inference_metadata.yaml b/tests/fixtures/onnx_genai_workflows/codec/inference_metadata.yaml index 394c7be99..aa6c5e570 100644 --- a/tests/fixtures/onnx_genai_workflows/codec/inference_metadata.yaml +++ b/tests/fixtures/onnx_genai_workflows/codec/inference_metadata.yaml @@ -47,56 +47,10 @@ pipeline: implementation: kind: onnx artifact: encoder/model.onnx - ports: - inputs: - waveform: - dtype: float32 - rank: 3 - shape: - - batch - - 1 - - audio_samples - batch_layout: - kind: request_aligned - axis: 0 - outputs: - codes: - dtype: float32 - rank: 3 - shape: - - batch - - 1 - - audio_samples - batch_layout: - kind: request_aligned - axis: 0 decoder: implementation: kind: onnx artifact: decoder/model.onnx - ports: - inputs: - codes: - dtype: float32 - rank: 3 - shape: - - batch - - 1 - - audio_samples - batch_layout: - kind: request_aligned - axis: 0 - outputs: - waveform: - dtype: float32 - rank: 3 - shape: - - batch - - 1 - - audio_samples - batch_layout: - kind: request_aligned - axis: 0 steps: - kind: invoke component: encoder diff --git a/tests/fixtures/onnx_genai_workflows/decoder/inference_metadata.yaml b/tests/fixtures/onnx_genai_workflows/decoder/inference_metadata.yaml index 7632647cd..46066a212 100644 --- a/tests/fixtures/onnx_genai_workflows/decoder/inference_metadata.yaml +++ b/tests/fixtures/onnx_genai_workflows/decoder/inference_metadata.yaml @@ -341,67 +341,6 @@ pipeline: kind: onnx artifact: model.onnx ports: - inputs: - input_ids: - dtype: int64 - rank: 2 - shape: - - batch - - sequence - batch_layout: - kind: request_aligned - axis: 0 - 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 - - sequence - 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 - outputs: - logits: - dtype: float32 - rank: 3 - shape: - - batch - - sequence - - 128 - batch_layout: - kind: request_aligned - axis: 0 - present.0.key: - dtype: float32 - rank: 4 - shape: - - batch - - 2 - - present_sequence - - 8 - batch_layout: - kind: request_aligned - axis: 0 roles: input_ids: token_ids attention_mask: attention_mask diff --git a/tests/fixtures/onnx_genai_workflows/diffusion/inference_metadata.yaml b/tests/fixtures/onnx_genai_workflows/diffusion/inference_metadata.yaml index e1bcc0dbf..64d063dfa 100644 --- a/tests/fixtures/onnx_genai_workflows/diffusion/inference_metadata.yaml +++ b/tests/fixtures/onnx_genai_workflows/diffusion/inference_metadata.yaml @@ -153,27 +153,6 @@ pipeline: kind: onnx artifact: text_encoder/model.onnx ports: - inputs: - input_ids: - dtype: int64 - rank: 2 - shape: - - batch - - prompt_sequence - batch_layout: - kind: request_aligned - axis: 0 - outputs: - encoder_hidden_states: - dtype: float32 - rank: 3 - shape: - - batch - - prompt_sequence - - 32 - batch_layout: - kind: request_aligned - axis: 0 roles: input_ids: token_ids encoder_hidden_states: encoder_hidden_states @@ -182,79 +161,12 @@ pipeline: kind: onnx artifact: denoiser/model.onnx ports: - inputs: - sample: - dtype: float32 - rank: 4 - shape: - - batch - - 4 - - height - - width - batch_layout: - kind: request_aligned - axis: 0 - timestep: - dtype: float32 - rank: 1 - shape: - - batch - batch_layout: - kind: request_aligned - axis: 0 - encoder_hidden_states: - dtype: float32 - rank: 3 - shape: - - batch - - prompt_sequence - - 32 - batch_layout: - kind: request_aligned - axis: 0 - outputs: - noise_pred: - dtype: float32 - rank: 4 - shape: - - batch - - 4 - - height - - width - batch_layout: - kind: request_aligned - axis: 0 roles: encoder_hidden_states: encoder_hidden_states vae_decoder: implementation: kind: onnx artifact: vae_decoder/model.onnx - ports: - inputs: - latent: - dtype: float32 - rank: 4 - shape: - - batch - - 4 - - height - - width - batch_layout: - kind: request_aligned - axis: 0 - outputs: - image: - dtype: float32 - rank: 4 - shape: - - batch - - 3 - - height - - width - batch_layout: - kind: request_aligned - axis: 0 solver_step: implementation: kind: onnx diff --git a/tests/fixtures/onnx_genai_workflows/diffusion_guided/inference_metadata.yaml b/tests/fixtures/onnx_genai_workflows/diffusion_guided/inference_metadata.yaml index bd47b38b3..c2fd5a1c8 100644 --- a/tests/fixtures/onnx_genai_workflows/diffusion_guided/inference_metadata.yaml +++ b/tests/fixtures/onnx_genai_workflows/diffusion_guided/inference_metadata.yaml @@ -211,27 +211,6 @@ pipeline: kind: onnx artifact: text_encoder/model.onnx ports: - inputs: - input_ids: - dtype: int64 - rank: 2 - shape: - - batch - - prompt_sequence - batch_layout: - kind: request_aligned - axis: 0 - outputs: - encoder_hidden_states: - dtype: float32 - rank: 3 - shape: - - batch - - prompt_sequence - - 32 - batch_layout: - kind: request_aligned - axis: 0 roles: input_ids: token_ids encoder_hidden_states: encoder_hidden_states @@ -240,79 +219,12 @@ pipeline: kind: onnx artifact: denoiser/model.onnx ports: - inputs: - sample: - dtype: float32 - rank: 4 - shape: - - batch - - 4 - - height - - width - batch_layout: - kind: request_aligned - axis: 0 - timestep: - dtype: float32 - rank: 1 - shape: - - batch - batch_layout: - kind: request_aligned - axis: 0 - encoder_hidden_states: - dtype: float32 - rank: 3 - shape: - - batch - - prompt_sequence - - 32 - batch_layout: - kind: request_aligned - axis: 0 - outputs: - noise_pred: - dtype: float32 - rank: 4 - shape: - - batch - - 4 - - height - - width - batch_layout: - kind: request_aligned - axis: 0 roles: encoder_hidden_states: encoder_hidden_states vae_decoder: implementation: kind: onnx artifact: vae_decoder/model.onnx - ports: - inputs: - latent: - dtype: float32 - rank: 4 - shape: - - batch - - 4 - - height - - width - batch_layout: - kind: request_aligned - axis: 0 - outputs: - image: - dtype: float32 - rank: 4 - shape: - - batch - - 3 - - height - - width - batch_layout: - kind: request_aligned - axis: 0 solver_step: implementation: kind: onnx diff --git a/tests/fixtures/onnx_genai_workflows/masked/inference_metadata.yaml b/tests/fixtures/onnx_genai_workflows/masked/inference_metadata.yaml index fe8e5ab3b..b898a308c 100644 --- a/tests/fixtures/onnx_genai_workflows/masked/inference_metadata.yaml +++ b/tests/fixtures/onnx_genai_workflows/masked/inference_metadata.yaml @@ -139,36 +139,6 @@ pipeline: kind: onnx artifact: model.onnx ports: - inputs: - input_ids: - dtype: int64 - rank: 2 - shape: - - batch - - sequence - batch_layout: - kind: request_aligned - axis: 0 - outputs: - logits: - dtype: float32 - rank: 3 - shape: - - batch - - sequence - - 128 - batch_layout: - kind: request_aligned - axis: 0 - proposed_tokens: - dtype: int64 - rank: 2 - shape: - - batch - - sequence - batch_layout: - kind: request_aligned - axis: 0 roles: input_ids: token_ids logits: logits diff --git a/tests/fixtures/onnx_genai_workflows/speculative/inference_metadata.yaml b/tests/fixtures/onnx_genai_workflows/speculative/inference_metadata.yaml index 1dbf87fd5..e818b8b2c 100644 --- a/tests/fixtures/onnx_genai_workflows/speculative/inference_metadata.yaml +++ b/tests/fixtures/onnx_genai_workflows/speculative/inference_metadata.yaml @@ -282,93 +282,10 @@ pipeline: implementation: kind: onnx artifact: proposer/model.onnx - ports: - inputs: - tokens: - dtype: int64 - rank: 2 - shape: - - batch - - 4 - batch_layout: - kind: request_aligned - axis: 0 - proposal_budget: - dtype: int64 - rank: 1 - shape: - - batch - batch_layout: - kind: request_aligned - axis: 0 - outputs: - proposed_tokens: - dtype: int64 - rank: 2 - shape: - - batch - - 4 - batch_layout: - kind: request_aligned - axis: 0 - proposal_scores: - dtype: float32 - rank: 3 - shape: - - batch - - 4 - - 32 - batch_layout: - kind: request_aligned - axis: 0 verifier: implementation: kind: onnx artifact: verifier/model.onnx - ports: - inputs: - proposed_tokens: - dtype: int64 - rank: 2 - shape: - - batch - - 4 - 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 - outputs: - target_scores: - dtype: float32 - rank: 3 - shape: - - batch - - 4 - - 32 - batch_layout: - kind: request_aligned - axis: 0 - present.0.key: - dtype: float32 - rank: 4 - shape: - - batch - - 2 - - past_sequence + 4 - - 8 - batch_layout: - kind: request_aligned - axis: 0 grammar_clone: implementation: kind: adapter diff --git a/tests/fixtures/onnx_genai_workflows/static_cache/inference_metadata.yaml b/tests/fixtures/onnx_genai_workflows/static_cache/inference_metadata.yaml index 06163a59f..73912050c 100644 --- a/tests/fixtures/onnx_genai_workflows/static_cache/inference_metadata.yaml +++ b/tests/fixtures/onnx_genai_workflows/static_cache/inference_metadata.yaml @@ -353,92 +353,6 @@ pipeline: kind: onnx artifact: model.onnx ports: - inputs: - input_ids: - dtype: int64 - rank: 2 - shape: - - batch - - sequence - batch_layout: - kind: request_aligned - axis: 0 - position_ids: - dtype: int64 - rank: 2 - shape: - - batch - - sequence - 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 - write_indices: - dtype: int64 - rank: 1 - shape: - - batch - batch_layout: - kind: request_aligned - axis: 0 - nonpad_kv_seqlen: - dtype: int64 - rank: 1 - shape: - - batch - batch_layout: - kind: request_aligned - axis: 0 - outputs: - logits: - dtype: float32 - rank: 3 - shape: - - batch - - sequence - - 128 - batch_layout: - kind: request_aligned - axis: 0 - updated_key_cache.0: - dtype: float32 - rank: 3 - shape: - - batch - - 16 - - 8 - batch_layout: - kind: request_aligned - axis: 0 - updated_value_cache.0: - dtype: float32 - rank: 3 - shape: - - batch - - 16 - - 8 - batch_layout: - kind: request_aligned - axis: 0 roles: input_ids: token_ids position_ids: position_ids diff --git a/tests/fixtures/onnx_genai_workflows/tts/inference_metadata.yaml b/tests/fixtures/onnx_genai_workflows/tts/inference_metadata.yaml index eda94d097..8176c4723 100644 --- a/tests/fixtures/onnx_genai_workflows/tts/inference_metadata.yaml +++ b/tests/fixtures/onnx_genai_workflows/tts/inference_metadata.yaml @@ -252,380 +252,17 @@ pipeline: kind: onnx artifact: talker/model.onnx ports: - inputs: - inputs_embeds: - dtype: float32 - rank: 3 - shape: - - batch - - sequence_len - - 8 - batch_layout: - kind: request_aligned - axis: 0 - 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 - 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 - outputs: - logits: - dtype: float32 - rank: 3 - shape: - - batch - - sequence_len - - 2160 - batch_layout: - kind: request_aligned - axis: 0 - last_hidden_state: - dtype: float32 - rank: 3 - shape: - - batch - - 1 - - 8 - batch_layout: - kind: request_aligned - axis: 0 - present.0.key: - dtype: float32 - rank: 4 - shape: - - batch - - 1 - - total_sequence_len - - 4 - batch_layout: - kind: request_aligned - axis: 0 - present.0.value: - dtype: float32 - rank: 4 - shape: - - batch - - 1 - - total_sequence_len - - 4 - batch_layout: - kind: request_aligned - axis: 0 - 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: - inputs: - inputs_embeds: - dtype: float32 - rank: 3 - shape: - - batch - - sequence_len - - 8 - batch_layout: - kind: request_aligned - axis: 0 - step_index: - dtype: int64 - rank: 0 - shape: [] - 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 - 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 - outputs: - logits: - dtype: float32 - rank: 3 - shape: - - batch - - sequence_len - - 6 - batch_layout: - kind: request_aligned - axis: 0 - codec_embeddings: - dtype: float32 - rank: 3 - shape: - - 3 - - 6 - - 8 - present.0.key: - dtype: float32 - rank: 4 - shape: - - batch - - 8 - - total_sequence_len - - 128 - batch_layout: - kind: request_aligned - axis: 0 - present.0.value: - dtype: float32 - rank: 4 - shape: - - batch - - 8 - - total_sequence_len - - 128 - batch_layout: - kind: request_aligned - axis: 0 - present.1.key: - dtype: float32 - rank: 4 - shape: - - batch - - 8 - - total_sequence_len - - 128 - batch_layout: - kind: request_aligned - axis: 0 - present.1.value: - dtype: float32 - rank: 4 - shape: - - batch - - 8 - - total_sequence_len - - 128 - batch_layout: - kind: request_aligned - axis: 0 - present.2.key: - dtype: float32 - rank: 4 - shape: - - batch - - 8 - - total_sequence_len - - 128 - batch_layout: - kind: request_aligned - axis: 0 - present.2.value: - dtype: float32 - rank: 4 - shape: - - batch - - 8 - - total_sequence_len - - 128 - batch_layout: - kind: request_aligned - axis: 0 - present.3.key: - dtype: float32 - rank: 4 - shape: - - batch - - 8 - - total_sequence_len - - 128 - batch_layout: - kind: request_aligned - axis: 0 - present.3.value: - dtype: float32 - rank: 4 - shape: - - batch - - 8 - - total_sequence_len - - 128 - batch_layout: - kind: request_aligned - axis: 0 - present.4.key: - dtype: float32 - rank: 4 - shape: - - batch - - 8 - - total_sequence_len - - 128 - batch_layout: - kind: request_aligned - axis: 0 - present.4.value: - dtype: float32 - rank: 4 - shape: - - batch - - 8 - - total_sequence_len - - 128 - batch_layout: - kind: request_aligned - axis: 0 + 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 @@ -635,158 +272,22 @@ pipeline: implementation: kind: onnx artifact: embedding/model.onnx - ports: - inputs: - text_ids: - dtype: int64 - rank: 2 - shape: - - batch - - text_sequence_len - batch_layout: - kind: request_aligned - axis: 0 - codec_ids: - dtype: int64 - rank: 2 - shape: - - batch - - codec_sequence_len - batch_layout: - kind: request_aligned - axis: 0 - outputs: - text_embeds: - dtype: float32 - rank: 3 - shape: - - batch - - text_sequence_len - - 8 - batch_layout: - kind: request_aligned - axis: 0 - codec_embeds: - dtype: float32 - rank: 3 - shape: - - batch - - codec_sequence_len - - 8 - batch_layout: - kind: request_aligned - axis: 0 talker_step_embedder: implementation: kind: onnx artifact: talker_step_embedder/model.onnx ports: - inputs: - frame_codes: - dtype: int64 - rank: 2 - shape: - - batch - - 4 - batch_layout: - kind: request_aligned - axis: 0 - text_embed: - dtype: float32 - rank: 3 - shape: - - batch - - 1 - - 8 - batch_layout: - kind: request_aligned - axis: 0 - outputs: - inputs_embeds: - dtype: float32 - rank: 3 - shape: - - batch - - 1 - - 8 - batch_layout: - kind: request_aligned - axis: 0 roles: inputs_embeds: inputs_embeds talker_prefill_embedder: implementation: kind: onnx artifact: talker_prefill_embedder/model.onnx - ports: - inputs: - text_ids: - dtype: int64 - rank: 2 - shape: - - batch - - text_sequence_len - batch_layout: - kind: request_aligned - axis: 0 - outputs: - prefill_embeds: - dtype: float32 - rank: 3 - shape: - - batch - - prefill_sequence_len - - 8 - batch_layout: - kind: request_aligned - axis: 0 - trailing_text_embeds: - dtype: float32 - rank: 3 - shape: - - batch - - trailing_sequence_len - - 8 - batch_layout: - kind: request_aligned - axis: 0 code_predictor_prefill: implementation: kind: onnx artifact: code_predictor_prefill/model.onnx ports: - inputs: - talker_hidden: - dtype: float32 - rank: 3 - shape: - - batch - - 1 - - 8 - batch_layout: - kind: request_aligned - axis: 0 - group_0_embed: - dtype: float32 - rank: 3 - shape: - - batch - - 1 - - 8 - batch_layout: - kind: request_aligned - axis: 0 - outputs: - inputs_embeds: - dtype: float32 - rank: 3 - shape: - - batch - - 2 - - 8 - batch_layout: - kind: request_aligned - axis: 0 roles: inputs_embeds: inputs_embeds code_predictor_step_embedder: @@ -794,124 +295,20 @@ pipeline: kind: onnx artifact: code_predictor_step_embedder/model.onnx ports: - inputs: - codec_embeddings: - dtype: float32 - rank: 3 - shape: - - 3 - - 6 - - 8 - token: - dtype: int64 - rank: 1 - shape: - - batch - batch_layout: - kind: request_aligned - axis: 0 - embedding_index: - dtype: int64 - rank: 0 - shape: [] - outputs: - inputs_embeds: - dtype: float32 - rank: 3 - shape: - - batch - - 1 - - 8 - batch_layout: - kind: request_aligned - axis: 0 roles: inputs_embeds: inputs_embeds code_predictor_indices: implementation: kind: onnx artifact: code_predictor_indices/model.onnx - ports: - inputs: - iteration: - dtype: int64 - rank: 0 - shape: [] - outputs: - embedding_index: - dtype: int64 - rank: 0 - shape: [] - step_index: - dtype: int64 - rank: 0 - shape: [] - frame_index: - dtype: int64 - rank: 0 - shape: [] talker_text_step: implementation: kind: onnx artifact: talker_text_step/model.onnx - ports: - inputs: - trailing_text_embeds: - dtype: float32 - rank: 3 - shape: - - batch - - trailing_sequence - - 8 - batch_layout: - kind: request_aligned - axis: 0 - iteration: - dtype: int64 - rank: 1 - shape: - - batch - batch_layout: - kind: request_aligned - axis: 0 - outputs: - text_embed: - dtype: float32 - rank: 3 - shape: - - batch - - 1 - - 8 - batch_layout: - kind: request_aligned - axis: 0 codec: implementation: kind: onnx artifact: codec/model.onnx - ports: - inputs: - codes: - dtype: int64 - rank: 3 - shape: - - batch - - 4 - - frames - batch_layout: - kind: request_aligned - axis: 0 - outputs: - waveform: - dtype: float32 - rank: 3 - shape: - - batch - - 1 - - frames - batch_layout: - kind: request_aligned - axis: 0 last_token_logits: implementation: kind: onnx diff --git a/tests/fixtures/onnx_genai_workflows/video/inference_metadata.yaml b/tests/fixtures/onnx_genai_workflows/video/inference_metadata.yaml index 95a382664..effc1c829 100644 --- a/tests/fixtures/onnx_genai_workflows/video/inference_metadata.yaml +++ b/tests/fixtures/onnx_genai_workflows/video/inference_metadata.yaml @@ -160,131 +160,12 @@ pipeline: kind: onnx artifact: transformer/model.onnx ports: - inputs: - sample: - dtype: float32 - rank: 5 - shape: - - batch - - num_frames - - 4 - - height - - width - batch_layout: - kind: request_aligned - axis: 0 - timestep: - dtype: int64 - rank: 1 - shape: - - batch - batch_layout: - kind: request_aligned - axis: 0 - encoder_hidden_states: - dtype: float32 - rank: 3 - shape: - - batch - - prompt_sequence - - 32 - batch_layout: - kind: request_aligned - axis: 0 - outputs: - noise_pred: - dtype: float32 - rank: 5 - shape: - - batch - - num_frames - - 4 - - height - - width - batch_layout: - kind: request_aligned - axis: 0 roles: encoder_hidden_states: encoder_hidden_states vae_decoder: implementation: kind: onnx artifact: vae_decoder/model.onnx - ports: - inputs: - latent_sample: - dtype: float32 - rank: 5 - shape: - - batch - - 4 - - latent_frames - - latent_height - - latent_width - batch_layout: - kind: request_aligned - axis: 0 - conv_cache.conv_in: - dtype: float32 - rank: 5 - shape: - - batch - - 4 - - cache_frames - - latent_height - - latent_width - batch_layout: - kind: request_aligned - axis: 0 - conv_cache.conv_out: - dtype: float32 - rank: 5 - shape: - - batch - - 3 - - cache_frames - - 2*latent_height - - 2*latent_width - batch_layout: - kind: request_aligned - axis: 0 - outputs: - sample: - dtype: float32 - rank: 5 - shape: - - batch - - 3 - - frames - - 2*latent_height - - 2*latent_width - batch_layout: - kind: request_aligned - axis: 0 - conv_cache_out.conv_in: - dtype: float32 - rank: 5 - shape: - - batch - - 4 - - cache_frames - - latent_height - - latent_width - batch_layout: - kind: request_aligned - axis: 0 - conv_cache_out.conv_out: - dtype: float32 - rank: 5 - shape: - - batch - - 3 - - cache_frames - - 2*latent_height - - 2*latent_width - batch_layout: - kind: request_aligned - axis: 0 model_input: implementation: kind: onnx diff --git a/tests/fixtures/onnx_genai_workflows/vlm/inference_metadata.yaml b/tests/fixtures/onnx_genai_workflows/vlm/inference_metadata.yaml index 2e3582ccb..c4306fbe4 100644 --- a/tests/fixtures/onnx_genai_workflows/vlm/inference_metadata.yaml +++ b/tests/fixtures/onnx_genai_workflows/vlm/inference_metadata.yaml @@ -423,67 +423,11 @@ pipeline: implementation: kind: onnx artifact: vision_encoder/model.onnx - ports: - inputs: - pixel_values: - dtype: float32 - rank: 2 - shape: - - 4 - - 1176 - grid_thw: - dtype: int64 - rank: 2 - shape: - - 1 - - 3 - outputs: - image_features: - dtype: float32 - rank: 3 - shape: - - batch - - 4 - - 32 - batch_layout: - kind: request_aligned - axis: 0 embedding: implementation: kind: onnx artifact: embedding/model.onnx ports: - inputs: - input_ids: - dtype: int64 - rank: 2 - shape: - - batch - - sequence - batch_layout: - kind: request_aligned - axis: 0 - image_features: - dtype: float32 - rank: 3 - shape: - - batch - - 4 - - 32 - batch_layout: - kind: request_aligned - axis: 0 - outputs: - inputs_embeds: - dtype: float32 - rank: 3 - shape: - - batch - - sequence - - 32 - batch_layout: - kind: request_aligned - axis: 0 roles: input_ids: token_ids inputs_embeds: inputs_embeds @@ -492,90 +436,6 @@ pipeline: kind: onnx artifact: decoder/model.onnx ports: - inputs: - inputs_embeds: - dtype: float32 - rank: 3 - shape: - - batch - - sequence - - 32 - batch_layout: - kind: request_aligned - axis: 0 - 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 - - sequence - 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 - outputs: - logits: - dtype: float32 - rank: 3 - shape: - - batch - - sequence - - 128 - batch_layout: - kind: request_aligned - axis: 0 - present.0.key: - dtype: float32 - rank: 4 - shape: - - batch - - 2 - - total_sequence - - 8 - batch_layout: - kind: request_aligned - axis: 0 - present.0.value: - dtype: float32 - rank: 4 - shape: - - batch - - 2 - - total_sequence - - 8 - batch_layout: - kind: request_aligned - axis: 0 roles: inputs_embeds: inputs_embeds attention_mask: attention_mask diff --git a/tests/generate_onnx_genai_validation_packages.py b/tests/generate_onnx_genai_validation_packages.py index 2515c4de4..bbaff149d 100644 --- a/tests/generate_onnx_genai_validation_packages.py +++ b/tests/generate_onnx_genai_validation_packages.py @@ -880,28 +880,12 @@ def _write_adapter_metadata(package: ModelPackage, directory: Path) -> None: "kind": "onnx", "artifact": "model.onnx", }, - # Every ONNX component declares its ports, including the - # ones a hand-written fixture builds: the workflow is the - # only description a package has of itself, and a - # component that declares nothing describes nothing. - "ports": { - "inputs": { - "activations": { - "dtype": "float32", - "rank": 2, - "shape": ["batch", 2], - "batch_layout": {"kind": "request_aligned", "axis": 0}, - } - }, - "outputs": { - "projection.output": { - "dtype": "float32", - "rank": 2, - "shape": ["batch", 2], - "batch_layout": {"kind": "request_aligned", "axis": 0}, - } - }, - }, + # 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": { diff --git a/tests/static_cache_metadata_test.py b/tests/static_cache_metadata_test.py index 05351a09e..eef86097f 100644 --- a/tests/static_cache_metadata_test.py +++ b/tests/static_cache_metadata_test.py @@ -18,6 +18,9 @@ from __future__ import annotations +from typing import Any + +import onnx_ir as ir import pytest from mobius import registry @@ -88,12 +91,16 @@ def _static_cache_abi(metadata) -> dict: 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"])) - declared = workflow["components"][component]["ports"] aliases = group["ports"][component] buffers = [ aliases[cell] for cell in sorted(aliases, key=lambda cell: int(cell.rsplit("_")[-1])) @@ -105,15 +112,35 @@ def _static_cache_abi(metadata) -> dict: "capacity": workflow["inputs"][update["capacity"]]["default"], "cache_inputs": [alias["input"] for alias in buffers], "cache_outputs": [alias["output"] for alias in buffers], - "declared_inputs": declared["inputs"], - "declared_outputs": declared["outputs"], } +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_workflow(): +def static_built(): pkg, config = _static_package() - return build_decoder_workflow_metadata(pkg, config) + 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") @@ -145,17 +172,20 @@ def test_control_ports_are_declared_by_the_scatter_discipline(self, static_workf 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): - # Naming a port the component does not expose would bind nothing. + 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"): - contract = abi["declared_inputs"][abi[role]] + value = static_ports[abi[role]] # Both are per-row integer vectors, which is exactly why they are # declared rather than recognized by shape. - assert contract["dtype"] == "int64" - assert contract["rank"] == 1 + assert value.dtype == ir.DataType.INT64 + assert len(value.shape) == 1 - def test_every_buffer_pair_is_declared(self, static_workflow): + 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", @@ -165,9 +195,7 @@ def test_every_buffer_pair_is_declared(self, static_workflow): ] assert abi["cache_outputs"] == [f"updated_{name}" for name in abi["cache_inputs"]] for name in (*abi["cache_inputs"], *abi["cache_outputs"]): - declared = abi["declared_inputs"] if name in abi["cache_inputs"] else None - declared = declared or abi["declared_outputs"] - assert declared[name]["shape"][STATIC_CACHE_SEQUENCE_AXIS] == CAPACITY + 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 From e9bd70f7752205a5e63e6c2735dac3df23d92a14 Mon Sep 17 00:00:00 2001 From: justinchuby Date: Fri, 21 Aug 2026 01:54:21 +0000 Subject: [PATCH 138/151] Pin the ONNX GenAI head rather than an ancestor of it MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit `2498e0bc` was the branch head when it was pinned and is still reachable — `git merge-base --is-ancestor 2498e0bc 0d1f1702` succeeds — so this is not the orphaned-pin failure. It is the weaker version of it: the pin names a commit that is no longer what the branch says, so CI was answering a question about a state that had moved on, and any new check added upstream would not have run here until something happened to notice. Nothing about the pinned commit was wrong. The rule is simply that a pin should name the ref's head, because an ancestor is a state nobody is maintaining and the gap only ever widens. `0d1f1702` adds the artifact to the canonical `tiny-llm-scatter-workflow` fixture and a test that validates that package as a directory rather than as a YAML string, which is the same entry point our fixtures go through. Re-ran both gates against it: 11/11 packages valid and 11/11 runtime conformance, with TensorScatter firing per layer per step on the fixed-capacity path. Unchanged from the previous head, which is the expected result — that head added a test and a fixture and no validation logic — but it is checked rather than assumed. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: justinchuby --- .github/workflows/main.yml | 2 +- docs/onnx-genai-performance-conformance.md | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/.github/workflows/main.yml b/.github/workflows/main.yml index 422c17e77..f8bb071ed 100644 --- a/.github/workflows/main.yml +++ b/.github/workflows/main.yml @@ -23,7 +23,7 @@ jobs: - uses: actions/checkout@v7 with: repository: justinchuby/onnx-genai - ref: 2498e0bcc524565b80af3c13f28dc56c2dc44e8b + ref: 0d1f17024d76918edb388e6dd556e0bc3c22ed69 path: validation/onnx-genai - uses: actions/setup-python@v7 with: diff --git a/docs/onnx-genai-performance-conformance.md b/docs/onnx-genai-performance-conformance.md index 3165dd09e..f5c2932be 100644 --- a/docs/onnx-genai-performance-conformance.md +++ b/docs/onnx-genai-performance-conformance.md @@ -123,7 +123,7 @@ well-formed and validate; only the local kernel is missing. ### Canonical representation — lowering verified -Against ONNX GenAI `2498e0bc`, with no `model.io` in any package and no port +Against ONNX GenAI `0d1f1702`, with no `model.io` in any package and no port contracts on any component that ships an artifact: | Check | Result | From 10cd48f1309e813d2016e4567f421f2ba1a74a49 Mon Sep 17 00:00:00 2001 From: justinchuby Date: Fri, 21 Aug 2026 02:57:00 +0000 Subject: [PATCH 139/151] Re-verify the canonical packages against a stricter package loader MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit `7324351a` moves the document-level invariants onto `load_metadata_package`. Until now they ran only for callers who already held a parsed document, which excluded everyone loading a package from disk — including the `validate_metadata` binary our CI job invokes. So the entry point our fixtures actually go through was enforcing a strictly weaker rule set than the one we believed we were being checked against. All 11 packages were re-run against it rather than assumed to be unaffected, because a strictness increase on the exact path we validate through is the case most likely to surface something: 11/11 valid, 11/11 runtime conformance, with TensorScatter still firing per layer per step on the fixed-capacity path. They pass for a structural reason rather than a lucky one. The rule that gained teeth forbids a package from carrying both `model.io` and a workflow, and ours carry no `model:` block at all — there is nothing for a coexistence check to find. This is the first external confirmation that emitting one serialized ABI was worth doing for a reason beyond tidiness: the packages that were relying on the loader's silence are the ones this commit would have broken. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: justinchuby --- .github/workflows/main.yml | 2 +- docs/onnx-genai-performance-conformance.md | 11 ++++++++++- 2 files changed, 11 insertions(+), 2 deletions(-) diff --git a/.github/workflows/main.yml b/.github/workflows/main.yml index f8bb071ed..3152c4888 100644 --- a/.github/workflows/main.yml +++ b/.github/workflows/main.yml @@ -23,7 +23,7 @@ jobs: - uses: actions/checkout@v7 with: repository: justinchuby/onnx-genai - ref: 0d1f17024d76918edb388e6dd556e0bc3c22ed69 + ref: 7324351ad9bac08bfc8a5f942ea4948b44bbff7d path: validation/onnx-genai - uses: actions/setup-python@v7 with: diff --git a/docs/onnx-genai-performance-conformance.md b/docs/onnx-genai-performance-conformance.md index f5c2932be..68bda530e 100644 --- a/docs/onnx-genai-performance-conformance.md +++ b/docs/onnx-genai-performance-conformance.md @@ -123,7 +123,7 @@ well-formed and validate; only the local kernel is missing. ### Canonical representation — lowering verified -Against ONNX GenAI `0d1f1702`, with no `model.io` in any package and no port +Against ONNX GenAI `7324351a`, with no `model.io` in any package and no port contracts on any component that ships an artifact: | Check | Result | @@ -137,6 +137,15 @@ 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. +`7324351a` 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. + ### Where the transcription boundary actually falls A component backed by a shipped `.onnx` declares `ports.roles` and nothing else: From 2a7bace0a765a84d4e02e20b23f16e831fe2d49f Mon Sep 17 00:00:00 2001 From: justinchuby Date: Fri, 21 Aug 2026 04:29:39 +0000 Subject: [PATCH 140/151] Require a sequence role wherever a fixed-capacity cache is bound MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit A component named in an `indexed_scatter` group's `write_indices_ports` or `kv_length_ports` is being handed the write cursor and valid length of a preallocated cache. Exactly one consumer reads those: the driver that scatters into the buffer at an index, and it binds its ports from the resolved decode ABI, every field of which is found by role. So a component handed the cursor while declaring no sequence role cannot be resolved as a decoder at all, and the package silently degrades to inferring ports from shapes — the behaviour the canonical form exists to remove. Nothing upstream catches it, and the reason is structural rather than an oversight. Identifying the decoder requires a sequence role, so a component that omits one is invisible to the check that would have caught it. The upstream sole-decoder guard works around that by firing only on workflows with a single ONNX component — but policy graphs are ONNX components, and every package we emit ships ten, so the guard is disabled by construction on precisely the packages that matter. Measured against `6e2ddc78`: dropping the role from the shipped `static_cache` fixture still reports `valid`. That leaves the producer as the only place the contradiction is visible, so it is asserted here, on both the built packages and the shipped fixtures. The second copy is not redundant — a fixture is hand-edited and regenerated far more often than the producer is changed, and the mutation above is exactly the edit that would slip through. The rule is scoped by what the ports are for rather than by counting anything. `speculative` owns attention state and declares no sequence role, which is correct: it is driven through explicit invoke bindings and never asks for a single-decoder ABI. It binds no scatter cursor, so the obligation does not reach it, and no exemption list is needed to say so. Pin moves to `6e2ddc78`, which adds that upstream guard. Re-ran both gates against it: 11/11 valid, 11/11 conformance. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: justinchuby --- .github/workflows/main.yml | 2 +- docs/onnx-genai-performance-conformance.md | 22 +++++++- tests/canonical_workflow_contract_test.py | 63 ++++++++++++++++++++++ 3 files changed, 85 insertions(+), 2 deletions(-) diff --git a/.github/workflows/main.yml b/.github/workflows/main.yml index 3152c4888..9945c2e8a 100644 --- a/.github/workflows/main.yml +++ b/.github/workflows/main.yml @@ -23,7 +23,7 @@ jobs: - uses: actions/checkout@v7 with: repository: justinchuby/onnx-genai - ref: 7324351ad9bac08bfc8a5f942ea4948b44bbff7d + ref: 6e2ddc78e956f84401d0fc65adbcf4538fb76ecc path: validation/onnx-genai - uses: actions/setup-python@v7 with: diff --git a/docs/onnx-genai-performance-conformance.md b/docs/onnx-genai-performance-conformance.md index 68bda530e..3a4e736dc 100644 --- a/docs/onnx-genai-performance-conformance.md +++ b/docs/onnx-genai-performance-conformance.md @@ -123,7 +123,7 @@ well-formed and validate; only the local kernel is missing. ### Canonical representation — lowering verified -Against ONNX GenAI `7324351a`, with no `model.io` in any package and no port +Against ONNX GenAI `6e2ddc78`, with no `model.io` in any package and no port contracts on any component that ships an artifact: | Check | Result | @@ -146,6 +146,26 @@ 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. +### 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: diff --git a/tests/canonical_workflow_contract_test.py b/tests/canonical_workflow_contract_test.py index 52914ad7f..70e394b58 100644 --- a/tests/canonical_workflow_contract_test.py +++ b/tests/canonical_workflow_contract_test.py @@ -194,6 +194,29 @@ 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.""" @@ -613,6 +636,32 @@ def test_state_pairs_still_bind_without_a_role_table(self, package): ): 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.""" @@ -842,3 +891,17 @@ def test_every_declared_role_resolves_in_the_artifact(self, directory): 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" + ) From 76f158b75aaedfa8949a0e2ba101c2e89a2e1029 Mon Sep 17 00:00:00 2001 From: justinchuby Date: Fri, 21 Aug 2026 05:33:26 +0000 Subject: [PATCH 141/151] Stop committing generated fixture graphs; drop the last model.io producer MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Two changes that both remove a second copy of something the tree already states once. `add_explicit_package_io` wrote `metadata["model"]["io"]`. It had no production callers, and ONNX GenAI is removing the serialized field with no compatibility shim, so the only thing it could still do was tempt a future caller into emitting an ABI the runtime discards at load. Removed with its six tests; the surviving cross-attention test keeps its coverage under a name that says what it checks. The static-cache error message no longer names `model.io.static_cache`, a key nothing emits. The eleven conformance packages committed 146 binary files and about 14 MB of graphs and weights, including a 5 MB blob. They are a deterministic function of the generator, CI already regenerated and compared the whole tree, and no reviewer can read a diff of them. Only the metadata is committed now. Textproto is the right answer for ONNX GenAI's 26 KB synthetic fixtures and the wrong one here: these carry real weight blobs that a text encoding would grow, and this repository does not use the protobuf APIs a textproto writer needs. Tests that need a graph build them once per session through the new `materialized_workflow_packages` fixture; generation takes about four seconds. This decides where CI must point, and the failure would otherwise have been silent in the wrong direction. A checkout of the committed tree alone does not validate — `component 'cache_length_update' artifact ... cannot be opened` — which is correct, because a workflow claims to describe something executable. Verified by archiving the staged tree and running the validator against it. Validation and conformance now run against the regenerated tree, and the conformance harness no longer guesses a default path that resolved inside the ONNX GenAI checkout it is copied into. Two guards keep the removal from hollowing out the assertions that remain. The comparison step checks that every artifact the committed metadata names was really produced, so a generator that stopped emitting one is caught rather than leaving a test with nothing to compare. `_ports` asserts the same per package, because two of its three callers skip components they cannot resolve and would have passed vacuously against an empty directory; a mutation pointing the fixture at a missing package fails all eleven. `test_describes_itself_only_through_the_workflow` also tightened from "no `model.io`" to "no second ABI at all", including the legacy `pipeline.models`, which is the rule the runtime actually enforces. Verified: 11/11 validate_metadata and 11/11 runtime conformance against the regenerated tree at ONNX GenAI 6e2ddc78, canonical contract suite 145 passed, fast suite 4541 passed, lintrunner clean. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: justinchuby --- .github/workflows/main.yml | 8 +- .gitignore | 10 +- docs/onnx-genai-performance-conformance.md | 31 +++- .../onnx_genai/inference_metadata.py | 79 +--------- .../onnx_genai/inference_metadata_test.py | 144 +----------------- tests/canonical_workflow_contract_test.py | 66 ++++++-- .../compare_onnx_genai_validation_packages.py | 41 ++++- tests/conftest.py | 24 +++ tests/fixtures/onnx_genai_workflows/README.md | 11 +- .../adapters/peft/adapter_model.safetensors | Bin 208 -> 0 bytes .../onnx_genai_workflows/adapter/model.onnx | Bin 356 -> 0 bytes .../adapter/model.onnx.data | 0 .../codec/decoder/model.onnx | Bin 323 -> 0 bytes .../codec/decoder/model.onnx.data | 0 .../codec/encoder/model.onnx | Bin 323 -> 0 bytes .../codec/encoder/model.onnx.data | 0 .../onnx_genai_workflows/decoder/model.onnx | Bin 2577 -> 0 bytes .../decoder/model.onnx.data | 0 .../decoder/policies/cache_length_update.onnx | Bin 930 -> 0 bytes .../policies/decoder_state_initializer.onnx | Bin 8581 -> 0 bytes .../decoder/policies/decoder_step_update.onnx | Bin 2037 -> 0 bytes .../policies/generated_length_update.onnx | Bin 930 -> 0 bytes .../decoder/policies/last_token_logits.onnx | Bin 631 -> 0 bytes .../decoder/policies/termination.onnx | Bin 5880 -> 0 bytes .../termination_batch_initializer.onnx | Bin 1677 -> 0 bytes .../decoder/policies/token_sampler.onnx | Bin 58720 -> 0 bytes .../decoder/policies/token_state_update.onnx | Bin 1229 -> 0 bytes .../decoder/policies/token_to_slot.onnx | Bin 438 -> 0 bytes .../diffusion/denoiser/model.onnx | Bin 1885 -> 0 bytes .../diffusion/denoiser/model.onnx.data | 0 .../policies/continue_predicate.onnx | Bin 910 -> 0 bytes .../policies/diffusion_schedule.onnx | Bin 497 -> 0 bytes .../policies/diffusion_timesteps.onnx | Bin 493 -> 0 bytes .../policies/initial_state_scale.onnx | Bin 363 -> 0 bytes .../diffusion/policies/model_input_scale.onnx | Bin 1926 -> 0 bytes .../diffusion/policies/schedule_lookup.onnx | Bin 607 -> 0 bytes .../diffusion/policies/solver_step.onnx | Bin 3060 -> 0 bytes .../diffusion/policies/tensor_scale.onnx | Bin 615 -> 0 bytes .../diffusion/text_encoder/model.onnx | Bin 1331 -> 0 bytes .../diffusion/text_encoder/model.onnx.data | 0 .../diffusion/vae_decoder/model.onnx | Bin 957 -> 0 bytes .../diffusion/vae_decoder/model.onnx.data | 0 .../diffusion_guided/denoiser/model.onnx | Bin 1885 -> 0 bytes .../diffusion_guided/denoiser/model.onnx.data | 0 .../policies/continue_predicate.onnx | Bin 910 -> 0 bytes .../policies/decoder_input_scale.onnx | Bin 363 -> 0 bytes .../policies/diffusion_schedule.onnx | Bin 387 -> 0 bytes .../policies/diffusion_timesteps.onnx | Bin 383 -> 0 bytes .../policies/guidance_combine.onnx | Bin 1635 -> 0 bytes .../policies/history_initializer.onnx | Bin 777 -> 0 bytes .../policies/latent_noise.onnx | Bin 18187 -> 0 bytes .../policies/latent_row_shape.onnx | Bin 382 -> 0 bytes .../policies/schedule_lookup.onnx | Bin 607 -> 0 bytes .../policies/solver_step.onnx | Bin 13995 -> 0 bytes .../policies/tensor_scale.onnx | Bin 615 -> 0 bytes .../diffusion_guided/text_encoder/model.onnx | Bin 1331 -> 0 bytes .../text_encoder/model.onnx.data | 0 .../diffusion_guided/vae_decoder/model.onnx | Bin 957 -> 0 bytes .../vae_decoder/model.onnx.data | 0 .../onnx_genai_workflows/masked/model.onnx | Bin 1104 -> 0 bytes .../masked/model.onnx.data | 0 .../masked/policies/masked_update.onnx | Bin 13182 -> 0 bytes .../speculative/policies/adaptive_k.onnx | Bin 34702 -> 0 bytes .../policies/cache_length_update.onnx | Bin 386 -> 0 bytes .../policies/grammar_guidance.onnx | Bin 1926 -> 0 bytes .../speculative/policies/grammar_length.onnx | Bin 394 -> 0 bytes .../policies/grammar_sampler_logits.onnx | Bin 631 -> 0 bytes .../policies/proposal_metrics.onnx | Bin 1026 -> 0 bytes .../policies/speculative_acceptance.onnx | Bin 8736 -> 0 bytes .../speculative/proposer/model.onnx | Bin 1095 -> 0 bytes .../speculative/proposer/model.onnx.data | 0 .../speculative/verifier/model.onnx | Bin 4371 -> 0 bytes .../speculative/verifier/model.onnx.data | 0 .../static_cache/model.onnx | Bin 2311 -> 0 bytes .../static_cache/model.onnx.data | 0 .../policies/cache_length_update.onnx | Bin 930 -> 0 bytes .../policies/decoder_state_initializer.onnx | Bin 7037 -> 0 bytes .../policies/decoder_step_update.onnx | Bin 780 -> 0 bytes .../policies/generated_length_update.onnx | Bin 930 -> 0 bytes .../policies/last_token_logits.onnx | Bin 631 -> 0 bytes .../static_cache/policies/termination.onnx | Bin 5880 -> 0 bytes .../termination_batch_initializer.onnx | Bin 1677 -> 0 bytes .../static_cache/policies/token_sampler.onnx | Bin 58720 -> 0 bytes .../policies/token_state_update.onnx | Bin 1229 -> 0 bytes .../static_cache/policies/token_to_slot.onnx | Bin 438 -> 0 bytes .../tts/code_predictor/model.onnx | Bin 106338 -> 0 bytes .../tts/code_predictor/model.onnx.data | Bin 2529408 -> 0 bytes .../tts/code_predictor_indices/model.onnx | Bin 942 -> 0 bytes .../code_predictor_indices/model.onnx.data | 0 .../tts/code_predictor_prefill/model.onnx | Bin 465 -> 0 bytes .../code_predictor_prefill/model.onnx.data | 0 .../code_predictor_step_embedder/model.onnx | Bin 888 -> 0 bytes .../model.onnx.data | 0 .../onnx_genai_workflows/tts/codec/model.onnx | Bin 1040 -> 0 bytes .../tts/codec/model.onnx.data | 0 .../tts/embedding/model.onnx | Bin 5313 -> 0 bytes .../tts/embedding/model.onnx.data | Bin 4984640 -> 0 bytes .../tts/policies/cache_length_update.onnx | Bin 386 -> 0 bytes .../tts/policies/code_frame_update.onnx | Bin 1580 -> 0 bytes .../tts/policies/code_history_append.onnx | Bin 853 -> 0 bytes .../tts/policies/codec_layout.onnx | Bin 428 -> 0 bytes .../tts/policies/continue_predicate.onnx | Bin 910 -> 0 bytes .../tts/policies/last_token_logits.onnx | Bin 631 -> 0 bytes .../tts/policies/predictor_body_sampler.onnx | Bin 536 -> 0 bytes .../policies/predictor_prefill_sampler.onnx | Bin 539 -> 0 bytes .../policies/predictor_state_initializer.onnx | Bin 14690 -> 0 bytes .../tts/policies/predictor_step_update.onnx | Bin 2039 -> 0 bytes .../tts/policies/setup_predictor_sampler.onnx | Bin 537 -> 0 bytes .../tts/policies/setup_talker_sampler.onnx | Bin 534 -> 0 bytes .../tts/policies/talker_sampler.onnx | Bin 528 -> 0 bytes .../policies/talker_state_initializer.onnx | Bin 7457 -> 0 bytes .../tts/policies/talker_step_update.onnx | Bin 2044 -> 0 bytes .../tts/policies/token_to_slot.onnx | Bin 438 -> 0 bytes .../tts/policies/tts_state_initializer.onnx | Bin 2194 -> 0 bytes .../tts/talker/model.onnx | Bin 31685 -> 0 bytes .../tts/talker/model.onnx.data | Bin 72704 -> 0 bytes .../tts/talker_prefill_embedder/model.onnx | Bin 21848 -> 0 bytes .../talker_prefill_embedder/model.onnx.data | Bin 4984640 -> 0 bytes .../tts/talker_step_embedder/model.onnx | Bin 7537 -> 0 bytes .../tts/talker_step_embedder/model.onnx.data | Bin 69696 -> 0 bytes .../tts/talker_text_step/model.onnx | Bin 1860 -> 0 bytes .../tts/talker_text_step/model.onnx.data | 0 .../video/policies/continue_predicate.onnx | Bin 910 -> 0 bytes .../video/policies/diffusion_schedule.onnx | Bin 387 -> 0 bytes .../video/policies/diffusion_timesteps.onnx | Bin 383 -> 0 bytes .../video/policies/model_input.onnx | Bin 720 -> 0 bytes .../policies/schedule_history_append.onnx | Bin 829 -> 0 bytes .../video/policies/schedule_lookup.onnx | Bin 607 -> 0 bytes .../video/policies/solver_step.onnx | Bin 6366 -> 0 bytes .../video/policies/video_conv_cache_init.onnx | Bin 2801 -> 0 bytes .../video/policies/video_decode_chunk.onnx | Bin 3115 -> 0 bytes .../video/policies/video_decode_chunks.onnx | Bin 1130 -> 0 bytes .../video/policies/video_latent_init.onnx | Bin 1580 -> 0 bytes .../video/policies/video_latent_permute.onnx | Bin 488 -> 0 bytes .../video/policies/video_latent_unscale.onnx | Bin 824 -> 0 bytes .../video/transformer/model.onnx | Bin 5197 -> 0 bytes .../video/transformer/model.onnx.data | 0 .../video/vae_decoder/model.onnx | Bin 14988 -> 0 bytes .../video/vae_decoder/model.onnx.data | 0 .../vlm/decoder/model.onnx | Bin 3469 -> 0 bytes .../vlm/decoder/model.onnx.data | 0 .../vlm/embedding/model.onnx | Bin 1715 -> 0 bytes .../vlm/embedding/model.onnx.data | 0 .../vlm/policies/cache_length_update.onnx | Bin 930 -> 0 bytes .../policies/decoder_state_initializer.onnx | Bin 9620 -> 0 bytes .../vlm/policies/decoder_step_update.onnx | Bin 2037 -> 0 bytes .../vlm/policies/generated_length_update.onnx | Bin 930 -> 0 bytes .../vlm/policies/last_token_logits.onnx | Bin 631 -> 0 bytes .../vlm/policies/termination.onnx | Bin 5880 -> 0 bytes .../termination_batch_initializer.onnx | Bin 1677 -> 0 bytes .../vlm/policies/token_sampler.onnx | Bin 58720 -> 0 bytes .../vlm/policies/token_state_update.onnx | Bin 1229 -> 0 bytes .../vlm/policies/token_to_slot.onnx | Bin 438 -> 0 bytes .../vlm/vision_encoder/model.onnx | Bin 652 -> 0 bytes .../vlm/vision_encoder/model.onnx.data | 0 ...generate_onnx_genai_validation_packages.py | 31 +++- tests/onnx_genai_workflow_conformance.rs | 15 +- 157 files changed, 208 insertions(+), 252 deletions(-) delete mode 100644 tests/fixtures/onnx_genai_workflows/adapter/adapters/peft/adapter_model.safetensors delete mode 100644 tests/fixtures/onnx_genai_workflows/adapter/model.onnx delete mode 100644 tests/fixtures/onnx_genai_workflows/adapter/model.onnx.data delete mode 100644 tests/fixtures/onnx_genai_workflows/codec/decoder/model.onnx delete mode 100644 tests/fixtures/onnx_genai_workflows/codec/decoder/model.onnx.data delete mode 100644 tests/fixtures/onnx_genai_workflows/codec/encoder/model.onnx delete mode 100644 tests/fixtures/onnx_genai_workflows/codec/encoder/model.onnx.data delete mode 100644 tests/fixtures/onnx_genai_workflows/decoder/model.onnx delete mode 100644 tests/fixtures/onnx_genai_workflows/decoder/model.onnx.data delete mode 100644 tests/fixtures/onnx_genai_workflows/decoder/policies/cache_length_update.onnx delete mode 100644 tests/fixtures/onnx_genai_workflows/decoder/policies/decoder_state_initializer.onnx delete mode 100644 tests/fixtures/onnx_genai_workflows/decoder/policies/decoder_step_update.onnx delete mode 100644 tests/fixtures/onnx_genai_workflows/decoder/policies/generated_length_update.onnx delete mode 100644 tests/fixtures/onnx_genai_workflows/decoder/policies/last_token_logits.onnx delete mode 100644 tests/fixtures/onnx_genai_workflows/decoder/policies/termination.onnx delete mode 100644 tests/fixtures/onnx_genai_workflows/decoder/policies/termination_batch_initializer.onnx delete mode 100644 tests/fixtures/onnx_genai_workflows/decoder/policies/token_sampler.onnx delete mode 100644 tests/fixtures/onnx_genai_workflows/decoder/policies/token_state_update.onnx delete mode 100644 tests/fixtures/onnx_genai_workflows/decoder/policies/token_to_slot.onnx delete mode 100644 tests/fixtures/onnx_genai_workflows/diffusion/denoiser/model.onnx delete mode 100644 tests/fixtures/onnx_genai_workflows/diffusion/denoiser/model.onnx.data delete mode 100644 tests/fixtures/onnx_genai_workflows/diffusion/policies/continue_predicate.onnx delete mode 100644 tests/fixtures/onnx_genai_workflows/diffusion/policies/diffusion_schedule.onnx delete mode 100644 tests/fixtures/onnx_genai_workflows/diffusion/policies/diffusion_timesteps.onnx delete mode 100644 tests/fixtures/onnx_genai_workflows/diffusion/policies/initial_state_scale.onnx delete mode 100644 tests/fixtures/onnx_genai_workflows/diffusion/policies/model_input_scale.onnx delete mode 100644 tests/fixtures/onnx_genai_workflows/diffusion/policies/schedule_lookup.onnx delete mode 100644 tests/fixtures/onnx_genai_workflows/diffusion/policies/solver_step.onnx delete mode 100644 tests/fixtures/onnx_genai_workflows/diffusion/policies/tensor_scale.onnx delete mode 100644 tests/fixtures/onnx_genai_workflows/diffusion/text_encoder/model.onnx delete mode 100644 tests/fixtures/onnx_genai_workflows/diffusion/text_encoder/model.onnx.data delete mode 100644 tests/fixtures/onnx_genai_workflows/diffusion/vae_decoder/model.onnx delete mode 100644 tests/fixtures/onnx_genai_workflows/diffusion/vae_decoder/model.onnx.data delete mode 100644 tests/fixtures/onnx_genai_workflows/diffusion_guided/denoiser/model.onnx delete mode 100644 tests/fixtures/onnx_genai_workflows/diffusion_guided/denoiser/model.onnx.data delete mode 100644 tests/fixtures/onnx_genai_workflows/diffusion_guided/policies/continue_predicate.onnx delete mode 100644 tests/fixtures/onnx_genai_workflows/diffusion_guided/policies/decoder_input_scale.onnx delete mode 100644 tests/fixtures/onnx_genai_workflows/diffusion_guided/policies/diffusion_schedule.onnx delete mode 100644 tests/fixtures/onnx_genai_workflows/diffusion_guided/policies/diffusion_timesteps.onnx delete mode 100644 tests/fixtures/onnx_genai_workflows/diffusion_guided/policies/guidance_combine.onnx delete mode 100644 tests/fixtures/onnx_genai_workflows/diffusion_guided/policies/history_initializer.onnx delete mode 100644 tests/fixtures/onnx_genai_workflows/diffusion_guided/policies/latent_noise.onnx delete mode 100644 tests/fixtures/onnx_genai_workflows/diffusion_guided/policies/latent_row_shape.onnx delete mode 100644 tests/fixtures/onnx_genai_workflows/diffusion_guided/policies/schedule_lookup.onnx delete mode 100644 tests/fixtures/onnx_genai_workflows/diffusion_guided/policies/solver_step.onnx delete mode 100644 tests/fixtures/onnx_genai_workflows/diffusion_guided/policies/tensor_scale.onnx delete mode 100644 tests/fixtures/onnx_genai_workflows/diffusion_guided/text_encoder/model.onnx delete mode 100644 tests/fixtures/onnx_genai_workflows/diffusion_guided/text_encoder/model.onnx.data delete mode 100644 tests/fixtures/onnx_genai_workflows/diffusion_guided/vae_decoder/model.onnx delete mode 100644 tests/fixtures/onnx_genai_workflows/diffusion_guided/vae_decoder/model.onnx.data delete mode 100644 tests/fixtures/onnx_genai_workflows/masked/model.onnx delete mode 100644 tests/fixtures/onnx_genai_workflows/masked/model.onnx.data delete mode 100644 tests/fixtures/onnx_genai_workflows/masked/policies/masked_update.onnx delete mode 100644 tests/fixtures/onnx_genai_workflows/speculative/policies/adaptive_k.onnx delete mode 100644 tests/fixtures/onnx_genai_workflows/speculative/policies/cache_length_update.onnx delete mode 100644 tests/fixtures/onnx_genai_workflows/speculative/policies/grammar_guidance.onnx delete mode 100644 tests/fixtures/onnx_genai_workflows/speculative/policies/grammar_length.onnx delete mode 100644 tests/fixtures/onnx_genai_workflows/speculative/policies/grammar_sampler_logits.onnx delete mode 100644 tests/fixtures/onnx_genai_workflows/speculative/policies/proposal_metrics.onnx delete mode 100644 tests/fixtures/onnx_genai_workflows/speculative/policies/speculative_acceptance.onnx delete mode 100644 tests/fixtures/onnx_genai_workflows/speculative/proposer/model.onnx delete mode 100644 tests/fixtures/onnx_genai_workflows/speculative/proposer/model.onnx.data delete mode 100644 tests/fixtures/onnx_genai_workflows/speculative/verifier/model.onnx delete mode 100644 tests/fixtures/onnx_genai_workflows/speculative/verifier/model.onnx.data delete mode 100644 tests/fixtures/onnx_genai_workflows/static_cache/model.onnx delete mode 100644 tests/fixtures/onnx_genai_workflows/static_cache/model.onnx.data delete mode 100644 tests/fixtures/onnx_genai_workflows/static_cache/policies/cache_length_update.onnx delete mode 100644 tests/fixtures/onnx_genai_workflows/static_cache/policies/decoder_state_initializer.onnx delete mode 100644 tests/fixtures/onnx_genai_workflows/static_cache/policies/decoder_step_update.onnx delete mode 100644 tests/fixtures/onnx_genai_workflows/static_cache/policies/generated_length_update.onnx delete mode 100644 tests/fixtures/onnx_genai_workflows/static_cache/policies/last_token_logits.onnx delete mode 100644 tests/fixtures/onnx_genai_workflows/static_cache/policies/termination.onnx delete mode 100644 tests/fixtures/onnx_genai_workflows/static_cache/policies/termination_batch_initializer.onnx delete mode 100644 tests/fixtures/onnx_genai_workflows/static_cache/policies/token_sampler.onnx delete mode 100644 tests/fixtures/onnx_genai_workflows/static_cache/policies/token_state_update.onnx delete mode 100644 tests/fixtures/onnx_genai_workflows/static_cache/policies/token_to_slot.onnx delete mode 100644 tests/fixtures/onnx_genai_workflows/tts/code_predictor/model.onnx delete mode 100644 tests/fixtures/onnx_genai_workflows/tts/code_predictor/model.onnx.data delete mode 100644 tests/fixtures/onnx_genai_workflows/tts/code_predictor_indices/model.onnx delete mode 100644 tests/fixtures/onnx_genai_workflows/tts/code_predictor_indices/model.onnx.data delete mode 100644 tests/fixtures/onnx_genai_workflows/tts/code_predictor_prefill/model.onnx delete mode 100644 tests/fixtures/onnx_genai_workflows/tts/code_predictor_prefill/model.onnx.data delete mode 100644 tests/fixtures/onnx_genai_workflows/tts/code_predictor_step_embedder/model.onnx delete mode 100644 tests/fixtures/onnx_genai_workflows/tts/code_predictor_step_embedder/model.onnx.data delete mode 100644 tests/fixtures/onnx_genai_workflows/tts/codec/model.onnx delete mode 100644 tests/fixtures/onnx_genai_workflows/tts/codec/model.onnx.data delete mode 100644 tests/fixtures/onnx_genai_workflows/tts/embedding/model.onnx delete mode 100644 tests/fixtures/onnx_genai_workflows/tts/embedding/model.onnx.data delete mode 100644 tests/fixtures/onnx_genai_workflows/tts/policies/cache_length_update.onnx delete mode 100644 tests/fixtures/onnx_genai_workflows/tts/policies/code_frame_update.onnx delete mode 100644 tests/fixtures/onnx_genai_workflows/tts/policies/code_history_append.onnx delete mode 100644 tests/fixtures/onnx_genai_workflows/tts/policies/codec_layout.onnx delete mode 100644 tests/fixtures/onnx_genai_workflows/tts/policies/continue_predicate.onnx delete mode 100644 tests/fixtures/onnx_genai_workflows/tts/policies/last_token_logits.onnx delete mode 100644 tests/fixtures/onnx_genai_workflows/tts/policies/predictor_body_sampler.onnx delete mode 100644 tests/fixtures/onnx_genai_workflows/tts/policies/predictor_prefill_sampler.onnx delete mode 100644 tests/fixtures/onnx_genai_workflows/tts/policies/predictor_state_initializer.onnx delete mode 100644 tests/fixtures/onnx_genai_workflows/tts/policies/predictor_step_update.onnx delete mode 100644 tests/fixtures/onnx_genai_workflows/tts/policies/setup_predictor_sampler.onnx delete mode 100644 tests/fixtures/onnx_genai_workflows/tts/policies/setup_talker_sampler.onnx delete mode 100644 tests/fixtures/onnx_genai_workflows/tts/policies/talker_sampler.onnx delete mode 100644 tests/fixtures/onnx_genai_workflows/tts/policies/talker_state_initializer.onnx delete mode 100644 tests/fixtures/onnx_genai_workflows/tts/policies/talker_step_update.onnx delete mode 100644 tests/fixtures/onnx_genai_workflows/tts/policies/token_to_slot.onnx delete mode 100644 tests/fixtures/onnx_genai_workflows/tts/policies/tts_state_initializer.onnx delete mode 100644 tests/fixtures/onnx_genai_workflows/tts/talker/model.onnx delete mode 100644 tests/fixtures/onnx_genai_workflows/tts/talker/model.onnx.data delete mode 100644 tests/fixtures/onnx_genai_workflows/tts/talker_prefill_embedder/model.onnx delete mode 100644 tests/fixtures/onnx_genai_workflows/tts/talker_prefill_embedder/model.onnx.data delete mode 100644 tests/fixtures/onnx_genai_workflows/tts/talker_step_embedder/model.onnx delete mode 100644 tests/fixtures/onnx_genai_workflows/tts/talker_step_embedder/model.onnx.data delete mode 100644 tests/fixtures/onnx_genai_workflows/tts/talker_text_step/model.onnx delete mode 100644 tests/fixtures/onnx_genai_workflows/tts/talker_text_step/model.onnx.data delete mode 100644 tests/fixtures/onnx_genai_workflows/video/policies/continue_predicate.onnx delete mode 100644 tests/fixtures/onnx_genai_workflows/video/policies/diffusion_schedule.onnx delete mode 100644 tests/fixtures/onnx_genai_workflows/video/policies/diffusion_timesteps.onnx delete mode 100644 tests/fixtures/onnx_genai_workflows/video/policies/model_input.onnx delete mode 100644 tests/fixtures/onnx_genai_workflows/video/policies/schedule_history_append.onnx delete mode 100644 tests/fixtures/onnx_genai_workflows/video/policies/schedule_lookup.onnx delete mode 100644 tests/fixtures/onnx_genai_workflows/video/policies/solver_step.onnx delete mode 100644 tests/fixtures/onnx_genai_workflows/video/policies/video_conv_cache_init.onnx delete mode 100644 tests/fixtures/onnx_genai_workflows/video/policies/video_decode_chunk.onnx delete mode 100644 tests/fixtures/onnx_genai_workflows/video/policies/video_decode_chunks.onnx delete mode 100644 tests/fixtures/onnx_genai_workflows/video/policies/video_latent_init.onnx delete mode 100644 tests/fixtures/onnx_genai_workflows/video/policies/video_latent_permute.onnx delete mode 100644 tests/fixtures/onnx_genai_workflows/video/policies/video_latent_unscale.onnx delete mode 100644 tests/fixtures/onnx_genai_workflows/video/transformer/model.onnx delete mode 100644 tests/fixtures/onnx_genai_workflows/video/transformer/model.onnx.data delete mode 100644 tests/fixtures/onnx_genai_workflows/video/vae_decoder/model.onnx delete mode 100644 tests/fixtures/onnx_genai_workflows/video/vae_decoder/model.onnx.data delete mode 100644 tests/fixtures/onnx_genai_workflows/vlm/decoder/model.onnx delete mode 100644 tests/fixtures/onnx_genai_workflows/vlm/decoder/model.onnx.data delete mode 100644 tests/fixtures/onnx_genai_workflows/vlm/embedding/model.onnx delete mode 100644 tests/fixtures/onnx_genai_workflows/vlm/embedding/model.onnx.data delete mode 100644 tests/fixtures/onnx_genai_workflows/vlm/policies/cache_length_update.onnx delete mode 100644 tests/fixtures/onnx_genai_workflows/vlm/policies/decoder_state_initializer.onnx delete mode 100644 tests/fixtures/onnx_genai_workflows/vlm/policies/decoder_step_update.onnx delete mode 100644 tests/fixtures/onnx_genai_workflows/vlm/policies/generated_length_update.onnx delete mode 100644 tests/fixtures/onnx_genai_workflows/vlm/policies/last_token_logits.onnx delete mode 100644 tests/fixtures/onnx_genai_workflows/vlm/policies/termination.onnx delete mode 100644 tests/fixtures/onnx_genai_workflows/vlm/policies/termination_batch_initializer.onnx delete mode 100644 tests/fixtures/onnx_genai_workflows/vlm/policies/token_sampler.onnx delete mode 100644 tests/fixtures/onnx_genai_workflows/vlm/policies/token_state_update.onnx delete mode 100644 tests/fixtures/onnx_genai_workflows/vlm/policies/token_to_slot.onnx delete mode 100644 tests/fixtures/onnx_genai_workflows/vlm/vision_encoder/model.onnx delete mode 100644 tests/fixtures/onnx_genai_workflows/vlm/vision_encoder/model.onnx.data diff --git a/.github/workflows/main.yml b/.github/workflows/main.yml index 9945c2e8a..01a3be18e 100644 --- a/.github/workflows/main.yml +++ b/.github/workflows/main.yml @@ -46,7 +46,11 @@ jobs: # `bash -e`, so an unguarded failure inside the loop would abort it and # hide the remaining results. failed="" - for package in tests/fixtures/onnx_genai_workflows/*; do + # 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 \ @@ -63,7 +67,7 @@ jobs: 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/tests/fixtures/onnx_genai_workflows" \ + 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 diff --git a/.gitignore b/.gitignore index 38ca53b2b..93550c1a6 100644 --- a/.gitignore +++ b/.gitignore @@ -212,8 +212,14 @@ __marimo__/ *.onnx *.onnx.data *.gguf -!tests/fixtures/onnx_genai_workflows/**/*.onnx -!tests/fixtures/onnx_genai_workflows/**/*.onnx.data + +# 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 # Common test dirs output/** diff --git a/docs/onnx-genai-performance-conformance.md b/docs/onnx-genai-performance-conformance.md index 3a4e736dc..31d73201a 100644 --- a/docs/onnx-genai-performance-conformance.md +++ b/docs/onnx-genai-performance-conformance.md @@ -128,7 +128,7 @@ contracts on any component that ships an artifact: | Check | Result | | --- | --- | -| `validate_metadata` over the checked-in fixtures | 11/11 valid | +| `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 | @@ -146,6 +146,35 @@ 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. +### 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 diff --git a/src/mobius/integrations/onnx_genai/inference_metadata.py b/src/mobius/integrations/onnx_genai/inference_metadata.py index be613152c..d91fd9b4a 100644 --- a/src/mobius/integrations/onnx_genai/inference_metadata.py +++ b/src/mobius/integrations/onnx_genai/inference_metadata.py @@ -876,8 +876,8 @@ def _static_cache_io( 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": STATIC_CACHE_WRITE_INDICES, @@ -1391,81 +1391,6 @@ def validate_executable_closure(pkg: Any, metadata: dict[str, Any]) -> None: ) -def add_explicit_package_io( - 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 - 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], - } - 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) - return metadata - - #: 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. diff --git a/src/mobius/integrations/onnx_genai/inference_metadata_test.py b/src/mobius/integrations/onnx_genai/inference_metadata_test.py index c39bbfb4e..8e529fff6 100644 --- a/src/mobius/integrations/onnx_genai/inference_metadata_test.py +++ b/src/mobius/integrations/onnx_genai/inference_metadata_test.py @@ -29,7 +29,6 @@ _match_max_token_grid, _port, _processor_values, - add_explicit_package_io, add_policy_components_to_workflow, build_diffusion_pipeline_metadata, build_multimodal_pipeline_metadata, @@ -393,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"]), @@ -612,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, diff --git a/tests/canonical_workflow_contract_test.py b/tests/canonical_workflow_contract_test.py index 70e394b58..2e9bea22d 100644 --- a/tests/canonical_workflow_contract_test.py +++ b/tests/canonical_workflow_contract_test.py @@ -829,22 +829,58 @@ def _metadata(directory: str) -> dict[str, Any]: return yaml.safe_load(handle) @staticmethod - def _ports(directory: str, workflow: dict[str, Any]) -> dict[str, tuple[set, set]]: - """Load each shipped artifact and read the ports it really exposes.""" + 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(directory, component["implementation"]["artifact"]) - if not os.path.exists(artifact): - continue + 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 "io" not in (metadata.get("model") or {}) - - def test_no_component_transcribes_the_ports_of_its_own_artifact(self, directory): + 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 @@ -854,7 +890,7 @@ def test_no_component_transcribes_the_ports_of_its_own_artifact(self, directory) declaration into a partial second truth that drifts silently. """ workflow = self._metadata(directory)["pipeline"]["workflow"] - graphs = self._ports(directory, 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 {} @@ -864,9 +900,11 @@ def test_no_component_transcribes_the_ports_of_its_own_artifact(self, directory) 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): + 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) + 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 @@ -882,9 +920,11 @@ def test_every_binding_and_state_pair_resolves_in_the_artifact(self, directory): assert alias["input"] in inputs assert alias["output"] in outputs - def test_every_declared_role_resolves_in_the_artifact(self, directory): + 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) + graphs = self._ports(directory, workflow, materialized_workflow_packages) for name, component in _onnx_components(workflow).items(): if name not in graphs: continue diff --git a/tests/compare_onnx_genai_validation_packages.py b/tests/compare_onnx_genai_validation_packages.py index 1c0e35794..4ff252196 100644 --- a/tests/compare_onnx_genai_validation_packages.py +++ b/tests/compare_onnx_genai_validation_packages.py @@ -53,14 +53,44 @@ def _model(path: Path) -> dict[str, Any]: return _graph(model.graph) -def _relative_files(directory: Path) -> set[Path]: +# 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.name != "model.onnx.data" + 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) @@ -77,13 +107,16 @@ def main() -> None: parser.add_argument("actual", type=Path) args = parser.parse_args() - expected_files = _relative_files(args.expected) - actual_files = _relative_files(args.actual) + 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) 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 index bd380598a..70ea478d6 100644 --- a/tests/fixtures/onnx_genai_workflows/README.md +++ b/tests/fixtures/onnx_genai_workflows/README.md @@ -1,7 +1,16 @@ # ONNX GenAI workflow conformance fixtures Generated by `tests/generate_onnx_genai_validation_packages.py` for semantic -validation and runtime conformance against `justinchuby/onnx-genai@2af34dca`. +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 diff --git a/tests/fixtures/onnx_genai_workflows/adapter/adapters/peft/adapter_model.safetensors b/tests/fixtures/onnx_genai_workflows/adapter/adapters/peft/adapter_model.safetensors deleted file mode 100644 index ee02117f4d5e1d232a04c0481fd4b240f961c302..0000000000000000000000000000000000000000 GIT binary patch literal 0 HcmV?d00001 literal 208 zcmdnN00GrXNr}a&@wxdasX2NDMfq8&$t9Wjd3rhdMTzl_dgZB^=@}(TR@F)=C6xuK zN>)m4#zsmyO2rw8AYQbgj!~?RQc7Y;VtjsDT5)PgF;LJz$0D{?2cN-C1PwOQF~l_3 ZLdVc7wzgIQ2pAX|?16Yc&=(-u4ggvtK*#_9 diff --git a/tests/fixtures/onnx_genai_workflows/adapter/model.onnx b/tests/fixtures/onnx_genai_workflows/adapter/model.onnx deleted file mode 100644 index caf3d7ffbe1b82ff0fdad6034b7d4248cc7bcbf7..0000000000000000000000000000000000000000 GIT binary patch literal 0 HcmV?d00001 literal 356 zcmdsh0KPwd|ke??ch+ymGmzETimPjE$8j+Lz0whs#$Wy?q(Yf!26<}9mFuHirs?1)tEme^m_TWnZpb@IDjs0Gnp$X p3M2nT`4P1&0T($i#Qoox9^+IJLvxjJ6RqI$8U9;qP=->g^8k`SXA1xT diff --git a/tests/fixtures/onnx_genai_workflows/codec/decoder/model.onnx.data b/tests/fixtures/onnx_genai_workflows/codec/decoder/model.onnx.data deleted file mode 100644 index e69de29bb..000000000 diff --git a/tests/fixtures/onnx_genai_workflows/codec/encoder/model.onnx b/tests/fixtures/onnx_genai_workflows/codec/encoder/model.onnx deleted file mode 100644 index 6ddede8d66d2bc407625a55688c764f6fa2b9be4..0000000000000000000000000000000000000000 GIT binary patch literal 0 HcmV?d00001 literal 323 zcmd8I!9o6cvKx8V@cv4n`p-E=eKrm$k`haj rGq{*I7`YUM3vyD6lLEOQJ`?u|Dm)?MrOE z|JYmyC$=igL4@lEOoS`oJHokh^2+BzQZ7lwxi=n1@2?c>(TySoIH%n+Mm#RX?iCuK z*|e|*N}I7DU=U)6cwH*VhO?mH!coM!QJ5J{i}R4B7oNyx2FOYa@h4Jpgvv!Ec&6*0 zcv3hkBe~vv(zfwgVXW)*{o)EtM`{+qMoXxf&8TaDvT5P|xV_q{MzvM1tE~ty#ds^P zw%97>U|JPNU#6;N0$$C)tHtmv122S4OTBNVP3m!*)HR#jH*BI9Q*%j-YBNr?nMh@) zNcCIuZq@K^ePZgG5!lpB-7rnv_@|~;>=RQr@PkQf9NBp*!Qor~-Qm9wEC*Cba>_1A z)LKNTTv8Lb?KAK~&}V{iIb7o`>YcPwzM*>PwUpe2XIsowhh6_@Uq2CdV9TDxHCV;V zuHOrM9>Vt2x^C&1O1BLT&c_3SeLiE{Wk-jwNTpOycs?gRDo*iBc%jnmB&S=)Yw>i_ zm$Xa9Ne>_w_`-{FhG+I)5z>U+!u#;XMtCn7)Uqvjj332eqrD%of@RysO}IXl@Gd-8 zIk%H@-k~K-*;nZJ1DFMa58$mz`8qjeh*`kvH!z1~H&!(y?W|5ZLs|XJu+>gwRlw>v zMxUr?G1SRxUfs|lQ+m3p7C?HSW(3Ul)bEI?UNz=385&w7+j)As?G_=YWV0bN{Ox5S Hvy}e>E4m%3 diff --git a/tests/fixtures/onnx_genai_workflows/decoder/model.onnx.data b/tests/fixtures/onnx_genai_workflows/decoder/model.onnx.data deleted file mode 100644 index e69de29bb..000000000 diff --git a/tests/fixtures/onnx_genai_workflows/decoder/policies/cache_length_update.onnx b/tests/fixtures/onnx_genai_workflows/decoder/policies/cache_length_update.onnx deleted file mode 100644 index af3bc3fca3b73acdb986806a33ddcf35d9032ffc..0000000000000000000000000000000000000000 GIT binary patch literal 0 HcmV?d00001 literal 930 zcmcIjO-}+b5OsxRVU%c^7;z6qiP75z^yW!Ca`0-RF->X9ZgpulyG0>J4_rO?muG^?f4nPmxkXM$A4+~*~a+4JlnkM**TWH zl!Sb(G)cqV`-JDb)N&3b?u5Fmd*G)*2_ zi@B{ujuyK`Hq=-yGHCALRs#cr=2~_o%RJBF?}e;Zl`%KKHL7{c&U6iP=3T&P`3s_$ zOLMwfC>doc6?BTJk~%lXiMHQAW?m1@{>{XKp`&?L2@;XY}T7|t3e-b`&LnUXf$qLNrqsf-kPcl?Qj)D|Hm3keB?kc!Y~ZEq7N&aUxp zqNIpR51hC_3kU(C5^&?p1tGy-;t#+(vtKi_vC3r~a@xu4%shENo_XGPqhG-*$KGIk z;=lIyee{`zD$|*FJPl~z9XYO#S5N7C2lmvV4Wk@BxZVgxwPil)SJ0|!A3Ofk9y)lb zzl&~7kM{RG*FE!xv+*?88%}KBrw3zaX3vHP=eY3UNBwPNC~vVns6X_kj?dme&*lR$ zB_L)5q>j)Fjyu}Yo@s@887=#PJqxy$U~~x(EeWEfKrBcQ8Q{DiKGAUbl)mA)U=25* z#8^w7xfe-IkAG3$KDh zo=t6cMB7Fsdg931wUuy~7pyPkAnT|>))9hi9U#X+mTPi;r=e;@WE02T4-R}O&^j#8 z-kLdf;LK>(sK$@nV%=IP8s}v)U^nsn3{=6HjDSBlaApqe8RhVSgY4Co`6w@lPc;_u zF*{JRVXP+u4!~?&hF?`!W(Jly4_2N;`!i3?#6V_4CKAb#JXhbf-F=6emQe|xLX*Nx z!H0P{^#=_|p2lZO{8km2#U#c?LT5C@Mi#(u-$2&NI8ushu#}C6&4jx_C}pBR9{7fz}qV2D2cvOc&(cukgh-=vM2%u$xRgfB#0uD z-Sa7#bXnxXaA<=T=^54NF|=H-wi1r(S5VPD8~eFltAFx1OR(S;LB4m;UcZirt7Th5V z%%tDMOtP97@p_y36EB#TKdtBIPn14Tz2%laiUPCR8Mn6u0$rHkSEI3$fhE=EM-GSU64H`tp$)E4 zB#WtaBxx@oynYu**Tft9jtkxAz!}k~7c}RFC7u7p4Xg4>rbHq0OBCoW=T@SqxH8uT z5wDBI@tuOJnQ2uu#a5M!OV80>e~D2j)3AJrk)AZ$I~ZrZ&xm%G5tZnc6h* z=9CrfN6tCrb(_D}*n`2Gd}>JPH5b3Ki?ZyTHO)PAaopiM z=(euGADx)Kpcl{@uJ9<0jf4uiez+xER1J)f4seSwqe(N+n-ylRi=rf*C>#nCb5~f> zR)+NXfk<~0Yqz?|< z!mC2?9c~KkO<}w15;2*LWItuXOeQ$6xWg>$bW7<;g0cdOmS7$D0D#{RWS0r$aO8$1 zm{OOvlBG-Bf!rH$mys=81<$gD+=6j0ikjVISLSwpU_`PeR1t~a3x~Bi9isRhwl_aC z?#+l3@}bajh$q!USXDqGZQCl*yU>m?I{e@H{{k9-Q$2`M&B z^#HLF4`iJOnmJy|?8E~r(u(tslU<`=5uKwcPKCQNK25bR+*avS@v-&_+6n*BVK35y z(23rgdXw?+oDMxVnAyVs7e4v)v33vbKKD&J9^o6&Tj3^q_Q~0JGPY;ucg@GzZL~di T59i@SXBhB*bhrNE1fDntYBIciqR=~*3Mjd7=fZcwIPfSiytkQGDCx<%OFbj*ZLZHA5R z0qpq~=M9hZsjvgrmkoPFg&+f$1vIb+S2*+GW48*{ovg415w_eLL33nl&>H)Jp)DbTnbau_5_X$hYnVBQ2$lc&v}oQtFrm(!%*`$wXQYF^m>2b2gPE>9eFT zj+cZH`YSEKUy_b$-*r@b!BIO(=CPx+>C*!B9VJ}UhJx=2SNA2?5xND$jCSr`awq*u z{-)q90!mzQ2bc5}FD6)U>Nt_@U)I*tI3?n6*z)^W2PyddILz$9NAOs@%#vw}>}qb*s;+&E8X3 z=4Ryv!}|*z6$6;6Y!=kcLo1<-X5wEZ(c|W;GQ<6q-Q#w#0miIwEUQ-H-)>^lHeQFU z0%Q!Mc)X?)$FI>7*pD98#yR8ai3_FM@V$|1Uy+cIfZ7sgK7GDMr*QD!OyWBD!7M8( h(x8*6J924ob<(^>kD>a{8m!-Ew$!iZ>ns>MmA^m2fFb|@ diff --git a/tests/fixtures/onnx_genai_workflows/decoder/policies/generated_length_update.onnx b/tests/fixtures/onnx_genai_workflows/decoder/policies/generated_length_update.onnx deleted file mode 100644 index af3bc3fca3b73acdb986806a33ddcf35d9032ffc..0000000000000000000000000000000000000000 GIT binary patch literal 0 HcmV?d00001 literal 930 zcmcIjO-}+b5OsxRVU%c^7;z6qiP75z^yW!Ca`0-RF->X9ZgpulyG0>J4_rO?muG^?f4nPmxkXM$A4+~*~a+4JlnkM**TWH zl!Sb(G)cqV`-JDb)N&3b?u5Fmd*G)*2_ zi@B{ujuyK`Hq=-yGHCALRs#cr=2~_o%RJBF?}e;Zl`%KKHL7{c&U6iP=3T&P`3s_$ zOLMwfC>doc6?BTJk~%lXiMHQAW?m1@{>{XKp`&?L2@;XVSA`N5$W@H3tkY;kI1!>+E8gov%LAWm#V`_v+67_8YRe% zs2_16q$z$Z7kWeKb^)@I4!I5}<2co>!9g(Sxk}2hW}yg-%Y3Re=?fmxko6}xeRtok zKzWUtWf09&!1auPR|K}NdctT}vBbGJ&XDA(T)3g8&qFd7ZdL8%*vNrJ$3hqXQ4nv> zJ&YP~xv;rCTPouds5=NBK?ZMl)UZAPb~JyCxMUpi7)FXwFY;+P@lN6AkMKV1#g)6? zl`Hha&RREYu$&A|O;H^VlhL|8F5^}kp`;rqU$6;bN}7-|gVRsXQ*;TnKS?4wxU`^> g5SK<{;fvV#qBTXwQ28~3ce|XK`4q1#NNp8A0X>A#xc~qF diff --git a/tests/fixtures/onnx_genai_workflows/decoder/policies/termination.onnx b/tests/fixtures/onnx_genai_workflows/decoder/policies/termination.onnx deleted file mode 100644 index 3183438a88f3ecce329e4afc84272274af35cdc8..0000000000000000000000000000000000000000 GIT binary patch literal 0 HcmV?d00001 literal 5880 zcmc&&-EQMV6ppiLlT6a4on^bhETwCeMT%$@J4xH@)v_yu6akf3aZ`jwt~1*@P3)}g zgf6f*6&D}`NU%4c6)Fz^kHTB<2HeA$@$cBFxN^yj=bRJIeE!aypP5PxFHh`abLQUv zYX`lQaD7f5*_P|+mPguhBN|)wkdn4i2^B3Fp}9VuQDR!2`{1AOCtXVA!3wHc`h>bu z-JrNI*hZVvvy)xhvd&$@F{j?HG1gs|oSM|p9pm&FE`IXqU<=8qf3Z1;YuHojvVWkT zCA93>XVh|$028mw$;X!abVljtl&JDrEO6MWQjIiuVf#TsA_`QdwoA;Ri>q_;?$en* zCYoIL2OO2AEO99>L_dno?4+I9NjP&2RU>CM3zt$cFHpZo0xyqG^(iG?x$X}*FI{Dc zOD%+MP;1yM0-RPIEy132ye90g+`zq%fCds^LdaloK-3{=;~J=g?bVZ4!y$*;Q_6vy z7iTF)?xh{Mcau2EguK^x}f=H4m_L)!Lr;x+dwp`}}?ozd_ zl$ey4OTXbR4YDcoZh!$-=Hwk~NK{p>v8NA4Rc?Tg7nW}&6s^kWc%-vz0Xd!zd<2?O zVp6M$ik{sp=a{>G=gb5cgx$0X(MN|fgFc>3h}x0YW61~EQL2%u)j*XqN~c3};^tZi zUPx@4hus8GyYgBj_^`THrL{ck)2|Xr@x?ZNgiP(S2I)PA>K=87+LIgM*vAaTAW-DB z>(adKWqA8v%E5QA{Y>CQ_e|RoDcj@0IfUTchv0;#Pt`sI=P4lXJR6?c8vr|rGY%P!Q?9p6SgwzCGnlMF#Nr|{wiiu}gmRDC$3W}+!)^sNyanLFochTA;*Z=+o3ukoiJ@$U*faIJB(`1736#NYBHopcVede5+EQKNrD6ykaQvaT|{^NThQGTY6V^F?$jQe z#xugQpaWrw|15r?IQEz-_qUWpZX#2M4ow421j^fB;P;dKFa8L0!9^fbhC|Rs4Jp&d z0$3(OyUNTdO)zJ0cM#4SieUA#3Nr*F(DAHiS|?09h0}{;f5reqhUE<&lU;6es7swW z{~Fd_FdcuwpA(sfsgojD!+OC?(HQjf$rEaLOq-xhd43@c(Dr}Q4p*pHVOPgH;2PcB dO-!*9G5fNe&7;#>9{6+eFJ z9f6*STIRrk<$5$&M!=LGsZ@<0$Sa93}x?dbfjp{^Dl?Sd8(XG$^SqE+~%6!U`m zt-v1R-c;ZTnPDd6l?0QBa03<9Lj2aJM1IDn`p(LXuuqL@;^)QqQ-S><7JKiJf#IKq8ulp`-M+8+uu z;+9i4z-$npn(uzbiK92)4Ja2(%ww8ohPqnF3+PgTJJ&!C-G-#;Gr;}w_`|aEA*Mx@LpshFb zPi}uVXe03c)^R`ps!UpfqN?-{)`GXpprwv7A40n2=pXFnA*36z0b?FR<}Gi8jQ6FO ztrRy+IKG=vnF^F7ePJxmOzBOOFpV$wrDsWniriI%!a_NL<5=)JBgzvg>v*oitU0z^ zE`0e1K*g`$7RpO#|2I<320BQh;v^k1oiT?Ie}3LwC=cLhYYsXbQcFa_|EdIq-PT_U CB^YA> diff --git a/tests/fixtures/onnx_genai_workflows/decoder/policies/token_sampler.onnx b/tests/fixtures/onnx_genai_workflows/decoder/policies/token_sampler.onnx deleted file mode 100644 index 1a5d5e311eae66f17403b7c4307e7609913be913..0000000000000000000000000000000000000000 GIT binary patch literal 0 HcmV?d00001 literal 58720 zcmd5_ON?a6RqgCwb!UGkJ>xFd_|r~X1`2zW@Bf~Kggs_347Nv_HUf=hrLwDXx@)F8 zvs{_g9)l2QkPYH%jAg)l+_L#vd<3#3SRf%5ED#b)#EKOXOZF^=6Y*aB?v2;wO+;32 zI^U<>yYJq(C+>|KC*snT(Zz?$y9dXs5C7Z0o*RCCcz%Dmn2)v|O+U3 zPcM(A)y`Aak!)$zcxX1XZXh8Ml!O#_J zbh(4`B;C1c@H~kJ=x9F5{c=E2+qophXqj!aEU3}i=F>J6A(Z@;>E;74H_ayxDhX{q z1#3PwggedWWb7a#=eE1?+&0a%+opx>wrLS92K&x{8BF7HTDaa@GVd)d@7*LWr|2@=fGtIrq0YFO z<=Z0~xH1N=9Rt_V4E)_E8Mt;0+%N+-mVq}4&uQtxbK}BuHev9bO*R0}*+ju}nh|); z5)RK<5*MD+jCgyND0t2=LZsv6ER7j2XBqY0tan`KYCP}FGVQ%tZhLPQZ)k|YbCzg$ z&T_|lGh`h)e36}hhO9%K*=X3Dq32Lw+{PstdJc8Qh37n?flFiH$}w=YIe4}WT)75D z@}bPxX%WeXI8lLma@!lLci1O9edvoL< z>dZ#rIcMWldGC1Ddm|dyj92B(cvT*6yoyE7wt?MYtDFs6<@T`EM&UWn9eB|OBSY87dmAvOsg)e zXrm}%MjHjo6^hR*l}bzi55qi;)!M1a^eJ#C5qB2Tx``>qA0D-Y&QKXYA3!Z z6F0snQgO(!ZCwYbfomp$ywKkm50g=Y^sjWYgOX(jrJWsY5|>L>)Gm9A+A$4mirS@9)Gp%{ zwORUS8`!OFSFE;OotoNq>DIO@QroV0ZM)h;ZM%~Bl?abZTHCHXwe2c$ZM&qk?MhF3 zs_fc!6}7fqvD$XkTif1_2hWwMZC6fhyV_uFyGk`YS8i>)Vzuq+wA8jMR@<(6Yun!+ z(ZHs*T{*SwYJ;_H+rVyZyJEHN>eSS>E4Q{?k=l06YuoiEYTLCQD5z*{yY|$!>&Uh3 ziq^JkUE8kh+IAhawq3K@cHLXsekPvx)~2>yJGJe4gSGA242somZM$Z*?fSITwre(C zRriip{a{1`oAIjJ8Lz72jaRYg**35{Y*n*itJ)s6+Gz2lc5B-;scqN1wq0+cwq5Jm zc1>&BjiPFvXL`gz=3oSu}RZq6u5FFvXKbr>u=+#);&$0W(hgFXv!2 zWz6D9!*YeDH&=*QJZVg>&^Wn56E9aV#gnGcSzzNP2@OjU8aqjdI!O1mffz$q)e*4a zkOk-?*|czK6XlcNmW%`~3Zxhf6%TBZ>|C>GX_F*c1?6vbM+5U(@^poWj#)Hm3lZ&6 zr$iK|H&jIVQtQU1j4HCstkx89Van+^*lndCcf zgF+P*x6^VSqm=SnrV*4wCYC7k30zJl7?9EcjiaCl&HT9Je%zZdge0RdKQ1jlu3SHk zmoxnRfE&*8XEkw2xkx*4$?<19Z`dqiBQel0*||Y)!8;@?g;}~WmIlk(M(iYB&HEpo zXz4`NC^}1nk}W9xwxyHJmB0T?vGQ~~?2;QTErD+EwG}r{LpNC4;rWB;V@$#$pgWPl zbBBWzyOYd=-AR-(_*VuZD(mE;7?8;PLtadfe;Dvi!j{4*cqavO>7=(99+Ij;T{Xa0GQl}`r>OyN zaKT+CK^jBiv2_swK|yp;y^aZ8#is5GoC z@jN5z%`-xlY*<^e@{G*QGqQMj28Q(&W@e>{MlzCMfV@qk5pcK7xc)}R0YTlSZ~+2$ z4_t&KzEBf`WU1i-BySro!0?U*7k}qBG3eb+#}1A+GciizDh7?rf1zOrXu(8U5{7{O z&B726z&9`e2rjU8sfvanpn~%-1Z?mKYo}T{WLD82AUBAB=B5zR8$?zLE=t^ji(5=Wj zFbkc6H)!b8rTcm4A(=GPr9nj3APFYAs+}6TsKCoQXK_)O@uG^Buoyi4K zI_*xtOz%nW&qu7`fS9gf8?1C2WI#!emEK!5oXSc9DpNH9E#0aa$Qj;_T9CF%I zfkJM^h&DPJ2nIPdF?i%Xa1pbH1Fokrapsyh-ol1f4aYHY=7yc1nllsIi{Nf{Dhb%h zPA~!vx*Lpuh0cNz@X$9f31$a=?nWaxp0UkVl2fV>pcflLHbq3yGuWtan8&iO< zSpgn2b`IX4v&RMRnib$dZ&y7I`0W;W2PU}F!67i*nM?N;;A1jyz;V~$4U)U#(wi*6 z+jMBq-CZZ)1^81#cuzUu9jE%hJ0b%IQfN!3F5w-A8w795(h$V4bjpeEIAwLyn~l5) z^BtnO;GKGy?>M&s@WxwpsgL@Olaax@p#^w|pujtIxbHZM8F{A@R)EK;K9`Qj zz=0H&OFISlbh1J2I}E4COCShwX;y$wC#R(VpHA427*6$BIwAuHQW#6S-1l^o+;_*) zg!`Uy?mJGM7Vt)+b0TkuM!?&{eNQ*SeaBlBcmw&!uu1CUzT;#J@Ww{nU2qM#`syXA z!+gh4Q0Qhn9H>H9mozsh1tH$#0yqjP6UhbAgm))NZ+CJb-arlnvp~k}PLg=tiK*eh zFs2_#bk!tH?5YWl;#!Y)5lL?saV4hptS$n~bt{ug!CH?sIAey?;p)bEdSCK(W^LN- zOK!0r#q^2$GJip`1jqV)Kaj}$ru7oS;5IO!z+(w502;4 zS05~A>u1iKIrGZn!64KQ?_`zhM36iQ22xD#Du_g+guPK7)L1>N?~klTx~M#ToEQipR+8r55G{rp=28`U4U zi}XR1;l&5bdk5>4NFu}%eV;gZeGVInT8KL83@K_M-KkL4WXj`xf#8@c4nN3Ujz|7u z^`#=9@-QFr%6k6raDFsfA0N#{K&B*t#^iQUpPaohO+hWcDt@$j1F8HRjzX#IkE8>1 z?;Q^JmW$PTn(R;IjH#botUh-Dg8Pjkr;iD zt}CeI1~LSdJRoGyMh4C}g?(qAani^x{r3;OJ76PY(lJtb>X6EbbZ5@|!Vd?O?i+*Q zg@*@=>7m>B!RL0Ip(C6>I5>ov1>s!$>e`$joZ~Q5bAyDUgX&)WB!?(H4pByMh`%vi z7B+S!9{ciq|9EfSx^uNNGKZb^s_|ugZuiFU((Cj2;r_wH)pp1}?Mrg8INNn)$^H|| zdoV9@1@WVk54i%4LKzp?Con`S0B^5E_umPk2Q$G-x3W*r(nC%opbjfI=ek>HwIVp@ z>J;s>5JjE|=&*n6cUWv7{uJ~!U#(s`df{`&vj@`@bo8g~6P@VE_=-mwaDjjy{GLh1GQ zsORSC&?~9fAQ;O)rr@JrQ7cj6V0`waqfoLvmaFo%FV#Uj5Rn9;qbrZ5KXre8gct@a z^y}@%*-Fv!i?SRSPS-PY=(} z-Z)rozoq_#G(T&XJTn2Qd;?(fb*XY~cPng21l-(FDKJRpH`PI>=VH9ooe9+*nAUi4 zqJJavkmme0843c`*HxAa-Qc3bxn7y8BZ)1f}ioyUN)rSQ^*wY-K^THd4vjw775Y2CNzgshc zYJRaj5lXynw5mmA%ZhZg3fyvLMUczGTMF)JQ zgPU6~UJ?nh`1g2lK9nkdQ;rcSE1IeSxa4*z3OsVVYOobBUPjd5*Q98TI%G&p?FH;| z{K1{Q*%}$=3lHWG=Zm#yljY7fJ=v+g9N!>UcCU2dL&$E1_J7#$oDPo&in(ePyU|jV zKbI;eDiIoV`@PalS;VI3Zps1H`sB$iB1T(-@rmF?!}` z1P!|#vl5lJ*W(wR?RfU+llm zf@o0S!w1nAb#@%U?=)h3=Fah5=n~}`@uPDOlxuJl$~V3$$BVl6lIXsYM9VfV6ufvV zx);3o=+Ql>?@n|N-h0I89{Z&GVp34vd!lA z_q0C*37+ZD&X8Dt7_@-=Eiz093aA(+G(o*PaZHlv-HBsj$M(<@$Ns|2kib`Vg>-Wy z5b8~qf_>VZ90`>AusIUeJ+a?TJMeaPgA5(T{_N(k?H<9R?`DADmS{6TXdC4W`qW|2 zlQtT7&eYeeX(65BTsH@=UKNQztG9|%Vi)(Y#VNzYOLa>F;LNl90QK@jY=J8f|MgncX@rja6*X;RQj-? z1Q}e|2g90i@FiGD>J5*>_~wTD+~AxBpi z^EcMhPGn-mz_ZLoJFEG8f4)COF=u{nd33Oc>6ok8!@~#jqum`*6nowp<9YFt>PXnN z#1zk7`93x2e)O*4vzG>=Tf?oZKz}0eZcJ~-lNxtyPyB$u$=*I9G?!K(= zer3Cx9wo1Ad9ECK|3*#{=T~kG&x?=c58E-O5F$7)ddKkMY;S$=XfFTkvMmFix4&HQ z3wNd1a7gy1yYCq~wq`fHw(UgZ!{H@JbAEXC(!jZbLa=oO;y;>xYPlA)3y5#a_?o;X zA%tIh*5b~JKfGpr+-q-_qiAYaIU8LZo`)!BUK{B@%H_FhIOpD5c7jKE=OvfeQg{u#Hn~42HsyIzjEv@`o`KVR{ z;%m+0u@)bFH+j6#h4px&40*R^CV~Vjgc0u`kEwl9HtyZ}F<@h=$81~57IR(OuN5yC zi?GKs1NOsZ-=aSh>_2wVl@JU(u19Rossy<$v~DCsF!*G}rL8JNgENXtTQvyr$@4so zz-~nN^vu%~s1CphVq$3>WF7^%I>?NZ7@0vv@#_FP53-k$=e6y`)j?+3h^vFl<0VuF z9yL)NU~WDWOX~o8v$2b=%-A4}w}agHjw^G1q61-zEzk3e1Ci&Q048V$*zSpa0@Hy< zPg)0`*l}Arn2PpQkcepZFK48Xj-aRs(2+UaCn(Gv$tDfh1Jc;N4MH2Q2R0k*qv3s-* zu-GuW=;{Ck4K6bl()G)XQTdae=kqPBn%MI^k8e>Oc+^C7(0KHub$}s8CbPMCqJpq4 zC09^m6l5Od(T%xyV&lPbgr4{GHJ-+!6&cW&=nbphsRur-M<5a;#L5&((cPABVhuOy zleVJKK_nB8PrZ*W8;r%v*n{3F0?0P(^r(qSZrYG7xoJdHa!)Iwk2DZ%JrXf+k21uX zNGB-+p)k9LDI@V{%qS!AsESCQ7@)Q#H%LvCfk2r_MIULf+J2+~Yw}10&Fqo1WgvWJ z_b_Gn0%@ZRUnET=w?b)4ax0dml6%~aKGOJ|{Yc|@^pPIVqss6F)Jz$^h}tN_7gAeS zJ*|p-R~KB12-*5BMAoXsTTk<7$+~)SgF<}mff{fx)w_reeStI)AK|p;kzEKaf-3q* zk4C6jd?7V63!-Y)kc~1h2a#Q+l6xBzV(UI-cp4OWB;sn1GO%Z7&m+Ztm?xnZ;%ase zQwBDIWH%XQ6dqMk$-Pis^g_8W*rt#4Lb(-gla}F&w`t4p1>8&-zKGi>!xwTB$*q{% za>=q6%6(BceWcg(thk#f!xwl{W%wd*rVL-`ZIt1Qy{)UBRz<$6i^2(dVQ!yl@z&Ek z_+(u@IX1wiD<1Dc?5*BKbm)t?iTH@SeUCJ$EPbR$Bh)Ou*qfOJTQ)qwVG)P3t5kAt zgTf{PJ<9Nf-sF)8z&*;qYW01O^dJ*e1_E&QB}^Gu0-fDtT85~)QjnyQd!Za<#XiYT zgmSSKHCKjM{n`IWxw0~qk|+a99P=-s$`DHmv(MGa5UXpltFAJ{lDj_1`=A`ljGp{R zpXXs&*FI%@4^cUm4CSw7%6MW_j^+FKBaEsLke{R~EGOwvl{JKIJ?&&IC0;FABiSyP zh&B9}fbY^7AQp#{m#^3b4Pt6U;zIs`6FZ8P%@~gTKmV?srw<_t@KKdTSEd652&v91S@!xk_e2!Xt$NRzZ z4v~WiLrPlOdN_m)Rr`vFmF6LlLz z^(Wl=L2Ftc!uu1<`wo|KWET^uN=)}4$J*5dRV<&u)q!%Zp@y~ngk4Let;!1dO$$fF zFPzQ!yY_0!+vQhx%tU6nvVr)T0g9$)hG$Uz#KZ|hVbo<)agBT_@&E;Ba`;e@^;4t( z>H_@}gMqiir)z@x=+f{aKH)6h@*d)yi2koa!Yw=#RbAX*+XAEP#ujH==?!68UlY>P zriTDPzmOvnL1J7p;~ciFN(gOmSU&%eaR>V*fq) z6l4bKQ{-C>8Osru{#Dz|JQ&-h1mcQ?u^5RLj`GB=00Z-0+`A>Pv!kI9nRIIr>;N>DT96!wKU1xR z{#o0|jE;3bj9J;@f#Qnvi>`$sDG-lSElfJInG~HOGZ4-^z~7$fCTTeRewP-sPld{5 zAC9cR{Ud3OZa%K_IwE~HA}bZOK zg;*f9`2a9%MkT_n+Ye)0d{OHkKTznwcwH0T( zUpObX|sDe&Vb>D@|N-j&Xhd>Z|j;wYYKw<^0^6gXf3uKhX};gZhL$KBC;5n(34|OxnK9asU_`vXX??c)*u104+^V#Ro!Qb5ee|UGf>Hq)$ diff --git a/tests/fixtures/onnx_genai_workflows/decoder/policies/token_state_update.onnx b/tests/fixtures/onnx_genai_workflows/decoder/policies/token_state_update.onnx deleted file mode 100644 index 7dfd0e7f3b7409692a79ae731043950283511932..0000000000000000000000000000000000000000 GIT binary patch literal 0 HcmV?d00001 literal 1229 zcmcIk&2G~`5U$%au?Hf;a-cP$3Rc9SqE-@aJ<%R>;D(Tb(8^g)6Q}8}+1;cqh)W)U z6VHOg)9@g$yLOt^6mg0#GxLpi=G*yM{{g9Ah*36I$3FtNgjy;%lE#9*6`FQKu!4&e zX*b$2`k)B`PcBg92}80q=)ggKK8^+FSBl9j*D;$VO3_J%GLdX@O}0OL9JC><{AFWM z#YB$E{sBu5>Iu`?0%4Wgk~rZh?S){4Bhia?jJ_!n-#we$G0wE2}&%g^*Y*h)`mzC zqQyCKsi6 zggk<#ACTbY^zdmpw^hgJ8Mw0C8*sXJr^^7op;WU>rr~-Dt04=Y8xx7K; zo=fkS2QP}Hi^s^#rsz!YI2W^wT~jl7TAE!TKfN!J6f=yD+tKYBN3YsZ*>vk#sRk{2 zV8K!|m@99_#}>^|QYkKNfQ%RGW^Lts5*5k?T4&+1XsH>@z6%y{(%5#|tnH++@$3wl yLmb!W>(V=c&VO%7vy|-H=i|s!BFSQR_^6=&roZ%FLi?XFij?jFIjzA~Ywr(oA9$t! diff --git a/tests/fixtures/onnx_genai_workflows/decoder/policies/token_to_slot.onnx b/tests/fixtures/onnx_genai_workflows/decoder/policies/token_to_slot.onnx deleted file mode 100644 index eb8e4b6599f5d6095617f9fe0d294eb1d098d216..0000000000000000000000000000000000000000 GIT binary patch literal 0 HcmV?d00001 literal 438 zcmaiwK}*9h7>3un+4PfQjh9(LrVIrSWxIOP9eekpNXW8mjq8%yq_9rB^zN_lhicoE zp@NuW-iMdxd(i>jU#Y3c_07i~yj##SYRM%4Q&Jl;o{O`5qznA(Y=X}G7R5MOS z7w+yS^1zjJ#dS&<$DL>hXX$bgDk;~RWg;~pOQ_Z)7CfUFiwkT&K1Blveo=J?(M+XW z*97(&8DbQv6|Mth`Of5HNf(^F^;z{hY@1>Pkfwhl0?k#NlR@@LYl8d bB4}3Jj!Wwj27h+&`J6M;Ofu<0$M1gu>uQY1 diff --git a/tests/fixtures/onnx_genai_workflows/diffusion/denoiser/model.onnx b/tests/fixtures/onnx_genai_workflows/diffusion/denoiser/model.onnx deleted file mode 100644 index b2f0012c61e80b875e79bade7be93ad228a581a4..0000000000000000000000000000000000000000 GIT binary patch literal 0 HcmV?d00001 literal 1885 zcmcIkPjA{V6vu%U_dyx8RE=hTgrc4}pi)XF?X-2$4rsesU=FI9Y6H_W^d= zaffL)uKPOsL~}wRAuGEjQIz*=zxRIry+642@=pVPA=p)H8cqb+o70aT3k5x~>q|if zuF;&GMC$;Qg>+kJfvip)c9dq3Zk3ES+=o5Mrb2~`3sM>$!}IXwx*JG&r+DOty2~d_ zDeCzmViEVg66^Ejunl(ZRP2K)9)v>0C*UU{^*MbLNE1qG+Or?6_ zPe>z%=+3x!FPNlf_CZEVfSx&!E)|86l+$kHmD5k!xhX|}<6Oe(g1PvqMW^lYH5(mMip5mxc$H7*!dXcHAVH1^>k9Gea z`RmC?;SwYA32f5h!88n{kh;6_)xBbxdtl)TJRwI9?0dq$_VgCxTi?~*sGIg8aiDF2 zJ>7CX2{sa2x`e!gp)r2@KVvv1&I2gJ$lOd*CBDo|0}6AW_SF~!a diff --git a/tests/fixtures/onnx_genai_workflows/diffusion/denoiser/model.onnx.data b/tests/fixtures/onnx_genai_workflows/diffusion/denoiser/model.onnx.data deleted file mode 100644 index e69de29bb..000000000 diff --git a/tests/fixtures/onnx_genai_workflows/diffusion/policies/continue_predicate.onnx b/tests/fixtures/onnx_genai_workflows/diffusion/policies/continue_predicate.onnx deleted file mode 100644 index 3177d8905c97dc120048da84e39b00bcae00c3f8..0000000000000000000000000000000000000000 GIT binary patch literal 0 HcmV?d00001 literal 910 zcmcIiJx{|h5Ur!oCReIf9-ssXL4`o-5Vf!|@gbH>m zU9gO*mp+iysemhKpX-P+jtiY#*oq$Rd?n>jvsgsNXCc*^^aPJ-%z7hiUEg%J!TW{E zdJxT2#C3*%cLXy9xH2G@JV+S7phMF3YCl1)s@-sBnPwH72hO8F^mQFAQQ#}WtLciA z)P_n!4!r8@o{KtYINyn84cq|@6HbIQ^}-TGy&+iK6)b4P4pY(ynQ(9Q{4t>+$)z4P zR`O^`e~HUor9w_+`kJ{~pZco`701Ocl<-QI8rJKs!AunNk3AdU?Y-kD5lI{>*;xnPGI+6BoP7L kn4%_4e3}eJC}=!7ZjaF(Z2#GV-`;X&a?79EP^i_v08>02xBvhE diff --git a/tests/fixtures/onnx_genai_workflows/diffusion/policies/diffusion_schedule.onnx b/tests/fixtures/onnx_genai_workflows/diffusion/policies/diffusion_schedule.onnx deleted file mode 100644 index 74c32b348f8e298f13299c0755c5726b615162b7..0000000000000000000000000000000000000000 GIT binary patch literal 0 HcmV?d00001 literal 497 zcmdvta;&1&l0Syj+}l ziMgr81&PV2LJVFiTyh23>3aEjc@@RUMVSR9ddWG7#l`U%nW;sIMadbJLQK)IUJ6`N zXsSSF#1|*$7o--0R0s*BWTvH+7H8(?0YxDplqALl_qPxi7Y7HU5DOO*hkO&GI+r3O zQ1#MN^MLl00JZ5Asj=am#CCYJ~?Ro6B#I&!HHR2iR{A|we_1u;Y~v9uyH hCo{3A(!sEaQH@Ihrx~efX{pI2U@thaa4`r-001ORvu6MR diff --git a/tests/fixtures/onnx_genai_workflows/diffusion/policies/diffusion_timesteps.onnx b/tests/fixtures/onnx_genai_workflows/diffusion/policies/diffusion_timesteps.onnx deleted file mode 100644 index 96f6134f2fc6bc5066af61d0c0d2cd3eb16d8568..0000000000000000000000000000000000000000 GIT binary patch literal 0 HcmV?d00001 literal 493 zcmdSF zJ9D1ybOxFU6!!y)=>YkhAT>@vaUkY&VgP~#j4WQfT%37{xv9kkiOH!#3|=Z+as}Dx zdii;I6~)O#nFS?!$vKI|#qk-LsYQuJ$r+VGOwqAk3S3fXsz7GM7boWzq!xoz2nnTR zrlpk@XXfVtMIoV+B*q2zvk(^-2M41N3l|fIToa=@mm(xU_0m)GfcBICwdob)=VT^V z#wX|Jl@ui=mk2Rc*ETUaa;Xqh8K0RVBnedoF+?x1v?4PnGqI@B!LW%@jY|Qi8L4S$ QsmUc^FF3JqF$hQi0BOmnd;kCd diff --git a/tests/fixtures/onnx_genai_workflows/diffusion/policies/initial_state_scale.onnx b/tests/fixtures/onnx_genai_workflows/diffusion/policies/initial_state_scale.onnx deleted file mode 100644 index 9bfb5425335bfa9b86c25844a44f8b3f0534e31e..0000000000000000000000000000000000000000 GIT binary patch literal 0 HcmV?d00001 literal 363 zcmaiv!Ab)$5Qdv&m2^Z}xQ7+D#%zQ1+C?sN3_+~5^62HdWXauqHiF?V zs_qa$TBkxyuwIauqZhLEnjf|r41iL?Qb;JN$OW7>@9>o}Ul*dbhC7>US=Q0$QAU~fa~S>Ep_WS} zp*?}b`yP1Nf8-C~Bx~1nAjCr+qI+%M>vME;a&0znweKJBNIv^@2QCnp(s5Zg@)#IW zu>eC%UQ^|=kgQnEt4#P^Mpo=@9#bR+ZdAq~ixnCKd1Tq0Byvu_9n zYCw)+BGOK`OV%w&LP~SpuIRidDrX6k6cU>;d3zixvSBsSn4)jkMzSahUi)PTn-JKg zQjtxok;at3rfnpPr8<1V1^rcKT; zZQP8>J0XuF#y&H$Wldf$QuVg%`f#Txc0UkQCG7`#7UkLCMqf`Lk8lx0%7Rxs&K$jM zS9D$w9X!SKR1C=o^#1mMAa233iSSdX;adk(IW7P;9hVBhJP9VQx{UW+l^lQKUFGi2 zT#eTWXqq)#ANK7(x}Lyz4rcVFcVP-o45l*inCL0@;5m$H@W1d!_j408A%R9n99hFL zf^VGN%Nb=2SHaNDrK~6W;cxWJkNk{(KFmCL6V1w@wK0IH^)jGrR*sg(Z7~mXn#%N0 zOkITr-#K~$53`r9)nh`xK8pLI75E-^PKo0S6;emx@|Ux7v;&L(tR%dP=dP-(6@6s9t?&9p-x#tmB^z^xDD zLpcMeNMc;M>&^LcdhWSi3FoJ3AQE%=?!qeq&#;gWlZb>#EJ%kV8U>oCdY{PzaQ)&guUI7=OfcFN}vT$$&UmM{}L|fwH@sHJa|DkqA(gGirXX&cJS6ifdDxbyJQ#$MP?$oGNp>N zo3`jBMbWL70_|bgS`7PJ_MhxODM>{!?|Y=x#PvaZ6b|{`KTGf# zfyMoR_J$F~jS+bgcnKw&_H1UbxFmuyth}SO2r%-ZAtfObyel!zl!fX&f-zH20ElNKg~KDVWrJ!o z$StQT?XtjpEK<~&q^M(1v>^1uewXcP~esH$majDw+9pO2%CT`jbc@4~x+1)fGf4x{0yOnOW#LD=m%%_|kD5AvQ zjL46SKO0i|k`mXo7ssAv&~=^0MORr;{2+6xN7HcRiEu<=0+YJhgJDG6ExW1O4416m z(kTnkI|Oy^4U&i|>ayq<7(Eq4K6--M&5K&;cFHW4c#+CKKuUNVjZ`Dc_u+1i4Y%Q* zh45XN!P7^c@B;vB`GLn6jW}4i9?9&5#w8uC!ivNtDao@%vZol-EbbYJe#NvpLw_UqHv7G{PbibWNfE}Zoy1Y-zaW0h zMC$nhe?hODG>)jVx8@l6?A(*C5|!@PqElUzi*Rt>J1!2{WSZ!?QViNZ|0(Of+?G3x zlTD2np4$ERnEFDm(8_XNpP??S|K}k{*u%@}2fZN=>V@s>92-iPUH1&#fwh0#A`7X% L&_C2*YGL^+<@Bd( diff --git a/tests/fixtures/onnx_genai_workflows/diffusion/policies/tensor_scale.onnx b/tests/fixtures/onnx_genai_workflows/diffusion/policies/tensor_scale.onnx deleted file mode 100644 index 10f87f5240227f36273f427df29c21481b4b48af..0000000000000000000000000000000000000000 GIT binary patch literal 0 HcmV?d00001 literal 615 zcmcJM%}N6?7)6_LI;LM0B!bi~q?V#0GPFB!<{p`< zr$Xm_J7sj3m{{>8QIvS)OCjbgN@xw<;X> z2kgQPEIR1jAvr{h4>#Yq9M?oRrwjr zh*V3HuYhlebeHU9s*L5zvI)Jn85i{@A^tNlx(~4A*9*xKWzCD>8nv6ET03brdcYPS zfdW@>%Z_NzvE`!z^(fpg!%+uXe3h7iw?=+bh~Mko0}#1_Y!b6W(p$3k%6wi(`Bk!G z+OL8F^EeDPGflzWBG{%N*r~=H;R)>&-vZ?%Y~{fy0NrMZgB#y$`C%djE2NY~1hWxq*A6r+Cotk%&8`(F8cxTs&* z*Km&8B!UC-NThR}Dyi(ae$@DqTQP-a$t{v|}*VV${xi_yzD|?VU MhpOCJ3%vf|7vDgf`~Uy| diff --git a/tests/fixtures/onnx_genai_workflows/diffusion/text_encoder/model.onnx.data b/tests/fixtures/onnx_genai_workflows/diffusion/text_encoder/model.onnx.data deleted file mode 100644 index e69de29bb..000000000 diff --git a/tests/fixtures/onnx_genai_workflows/diffusion/vae_decoder/model.onnx b/tests/fixtures/onnx_genai_workflows/diffusion/vae_decoder/model.onnx deleted file mode 100644 index 773c8b4b2b4432b5f206ae035966db2fd0edc2f7..0000000000000000000000000000000000000000 GIT binary patch literal 0 HcmV?d00001 literal 957 zcmchVu};G<5QdvJCE;2GctQaxR539`N!XYWY+2br$Z{ML8xkkPCcty>0xY}$Pe@u+ z)I8dvx{mr4_`m>(#nCoEMbZa`h1BY_ z_i9+-2CPexF;x)CP%k`$NpY9Bxs*?emZH!uO^H&N3s#bn@@Le)x(*M(X{ly2u%dav z)J)N)^4nB?C)MEpsDd_C&`D*ikm}8VwUlTkH8hfJRDPoh&`@L~Vb1s>6_nw5(#O-` zyjoGFXjBn~V@B&=RbG}0X2jvfHc$lnXqV<$k;~eMyCiiZqA7>IZNdcYeJ|lmBwTlo scoebb(LH>+{{70qA6MQwjc4}|fNfjItlb!`p8rEHSni#f(A(Pi06iTsFaQ7m diff --git a/tests/fixtures/onnx_genai_workflows/diffusion/vae_decoder/model.onnx.data b/tests/fixtures/onnx_genai_workflows/diffusion/vae_decoder/model.onnx.data deleted file mode 100644 index e69de29bb..000000000 diff --git a/tests/fixtures/onnx_genai_workflows/diffusion_guided/denoiser/model.onnx b/tests/fixtures/onnx_genai_workflows/diffusion_guided/denoiser/model.onnx deleted file mode 100644 index b2f0012c61e80b875e79bade7be93ad228a581a4..0000000000000000000000000000000000000000 GIT binary patch literal 0 HcmV?d00001 literal 1885 zcmcIkPjA{V6vu%U_dyx8RE=hTgrc4}pi)XF?X-2$4rsesU=FI9Y6H_W^d= zaffL)uKPOsL~}wRAuGEjQIz*=zxRIry+642@=pVPA=p)H8cqb+o70aT3k5x~>q|if zuF;&GMC$;Qg>+kJfvip)c9dq3Zk3ES+=o5Mrb2~`3sM>$!}IXwx*JG&r+DOty2~d_ zDeCzmViEVg66^Ejunl(ZRP2K)9)v>0C*UU{^*MbLNE1qG+Or?6_ zPe>z%=+3x!FPNlf_CZEVfSx&!E)|86l+$kHmD5k!xhX|}<6Oe(g1PvqMW^lYH5(mMip5mxc$H7*!dXcHAVH1^>k9Gea z`RmC?;SwYA32f5h!88n{kh;6_)xBbxdtl)TJRwI9?0dq$_VgCxTi?~*sGIg8aiDF2 zJ>7CX2{sa2x`e!gp)r2@KVvv1&I2gJ$lOd*CBDo|0}6AW_SF~!a diff --git a/tests/fixtures/onnx_genai_workflows/diffusion_guided/denoiser/model.onnx.data b/tests/fixtures/onnx_genai_workflows/diffusion_guided/denoiser/model.onnx.data deleted file mode 100644 index e69de29bb..000000000 diff --git a/tests/fixtures/onnx_genai_workflows/diffusion_guided/policies/continue_predicate.onnx b/tests/fixtures/onnx_genai_workflows/diffusion_guided/policies/continue_predicate.onnx deleted file mode 100644 index 3177d8905c97dc120048da84e39b00bcae00c3f8..0000000000000000000000000000000000000000 GIT binary patch literal 0 HcmV?d00001 literal 910 zcmcIiJx{|h5Ur!oCReIf9-ssXL4`o-5Vf!|@gbH>m zU9gO*mp+iysemhKpX-P+jtiY#*oq$Rd?n>jvsgsNXCc*^^aPJ-%z7hiUEg%J!TW{E zdJxT2#C3*%cLXy9xH2G@JV+S7phMF3YCl1)s@-sBnPwH72hO8F^mQFAQQ#}WtLciA z)P_n!4!r8@o{KtYINyn84cq|@6HbIQ^}-TGy&+iK6)b4P4pY(ynQ(9Q{4t>+$)z4P zR`O^`e~HUor9w_+`kJ{~pZco`701Ocl<-QI8rJKs!AunNk3AdU?Y-kD5lI{>*;xnPGI+6BoP7L kn4%_4e3}eJC}=!7ZjaF(Z2#GV-`;X&a?79EP^i_v08>02xBvhE diff --git a/tests/fixtures/onnx_genai_workflows/diffusion_guided/policies/decoder_input_scale.onnx b/tests/fixtures/onnx_genai_workflows/diffusion_guided/policies/decoder_input_scale.onnx deleted file mode 100644 index b716c92f07b3983eba24bd837d04da54cc88c380..0000000000000000000000000000000000000000 GIT binary patch literal 0 HcmV?d00001 literal 363 zcmaiv!Ab)$5Qdv&l{g};@uCH#N>S)x(W|FgZ~F=&VVb6GFik>|l`hiDK7~hLKztn^ z!iUgx*NY$?XZZeqzM;Epdu6Ap2-oikJQKF{h`Nyd?yWT;qKVwtndFye^k*?ThK3pK zQ_>@oX2{cz*SGNtN%K9}F}RY!p^&VT58=c;EHZ1%ItZ^^%!I}e_)MYX4jzPcGSmd?27wwq7p>Rypw(ailoFOgLP@pGy(Hn*P!vTQe#W-;9j*^RVFLGmP?eL5e( z)MyWaco-OdzW<;3**@Q0*qLtJ?MDFblryIa%$gDhhxgVvFO3(*W+<-D*^e-~f}Lks zHaLn|7}4q0lSYJ)2_diJZtl9BywNZ{fITA@=qjm@lXM8D)zdt&#w?wxb>$OPO6NqO zv6i(eRy=rkOk)`QqBa@ADO;i25S(W^&l{($>D2z`*%54i;Q)+L9zwvvmR`bHUo)Ab z>F)T>mQ=RXY9*94zLv`K;AP#?2@L-wh0gfNH`RwEvRUd<%X&4wY3T*Te`a8wqw-%B JrXi5%_!B6Vd({8{ diff --git a/tests/fixtures/onnx_genai_workflows/diffusion_guided/policies/diffusion_timesteps.onnx b/tests/fixtures/onnx_genai_workflows/diffusion_guided/policies/diffusion_timesteps.onnx deleted file mode 100644 index f6344fdf2fae202e0e65edd53b18c2471d558a47..0000000000000000000000000000000000000000 GIT binary patch literal 0 HcmV?d00001 literal 383 zcmaivu};G<6h-YcAoBoGd8nX5sGtf^hp4Pf&@$y05VBn7rLjouD0WaI#FCF-X5#Dk z2vQ0ikPyRLuFk#hUbf4(7j~vAcl{B-JLSx&43jFy{=w85=cVz&*aXG+jQt3sOW1mn zd4;26*6Gs8#uM)mBJT+4wf8!GqhY)cJ4P|2hjgTtulmDwm`QcI8St%R!&>fs_VY99&CDW0LCZ}Az+<`Ucl*kVl+q7 z{_(9XDr~M*EtECBl*;qq`K6(^F#MMkI^jp()H)=R)l%nLmi5h5L(gIGX9lJzD*rWM J5(4QSegeS9d)EK} diff --git a/tests/fixtures/onnx_genai_workflows/diffusion_guided/policies/guidance_combine.onnx b/tests/fixtures/onnx_genai_workflows/diffusion_guided/policies/guidance_combine.onnx deleted file mode 100644 index 4a71c4253dad3241fd648dd706c5611779bfa343..0000000000000000000000000000000000000000 GIT binary patch literal 0 HcmV?d00001 literal 1635 zcmcJPPiqrF6u>)4+H77L8U_k+ixQV2QqrYq>%o8+6)yo#;-wv(YME zawsTx_TZ%-zz^d`a%N*Te?TOM9F}?S&6_vBee>RP%6MlGcSWLFzjxsV!M0*F;@BCJ z7gT9-;+E4@#sMcMUco9yCCKTxhHg==Q-PAC1FnXYalF+zgnPrQo*zs3QL&L2YM(_^ zDG~}kq9YbwUnzyM`c^N>iyx)7{JDl%(l7GsDq7z?en} zCqin~M(tcXe~XHnBKwIT%MC0s;>wu3k?MWI`6o_J-RfLPnNPh^X4w?ocLZe?%Rp!m zOBz8Xkqas=j>+q!OU_)dOG^3AX; zrNgqM2-x9DiviW#g=ax9g?Z%Vjl3ZkT5Pk?nyxpam zg#g%LA(fIx3LLx>a?uO5-Toi~9bPANo2Oq#4gSDSw880;Z`j6*zwPxC;N2g`5=FtiIh2Z{ldip&s z&4{yY=MUpZuxseR__~j*MIvwNB3$RtK F{Q(ll`9J^w diff --git a/tests/fixtures/onnx_genai_workflows/diffusion_guided/policies/history_initializer.onnx b/tests/fixtures/onnx_genai_workflows/diffusion_guided/policies/history_initializer.onnx deleted file mode 100644 index 962692d28c0b298d171192c6e22a2612e45548d5..0000000000000000000000000000000000000000 GIT binary patch literal 0 HcmV?d00001 literal 777 zcmcJNO-}+b5QZrNE`vtX1cQ4pNQ@d{HsHlXFMixe`~w=(lrpfDrQK|cipC4^>e(y5 z(zCk@5;QT!gS~V*ednDg?KmY`80fYfn&bC8yb-!`$1b!omMhB|?%F(7x`V7yb3POG zGL-IkFhtf3G`HnUye5RSR-vf)08PjR(oAayw!(hT(@H%U5y{YcBH+d_Ut+`~;g4wc z`leL{_nRv5U`FT=O`?ET0$<3e6LWH7S<@}0d&<|W<&s0a2wAHuWJ0q4#q&a-7-1J9 zR03fRL-UUqY2(7>FkmUtVQndn{|SRa+KvUXk5?xPMAU-=htPdkrujCv!Uv#^@VQbL z7;tF8N7?gja(yp5)^8udl>D!Sw;SL7r2x6OapsmTJjyILCVQ|w9X79rDjonUwemt8 zNHJnUD;sfPY4+)POwM5ES0$4jx;ax#O*}q)kb&gU=(IT|yHNen1G`-mHl3~W9Awrj EA6w1yN&o-= diff --git a/tests/fixtures/onnx_genai_workflows/diffusion_guided/policies/latent_noise.onnx b/tests/fixtures/onnx_genai_workflows/diffusion_guided/policies/latent_noise.onnx deleted file mode 100644 index c9c6bbb9dc86870b09df5299d11a3472a2ee800d..0000000000000000000000000000000000000000 GIT binary patch literal 0 HcmV?d00001 literal 18187 zcmc&+O>Z1Y8TL4iJ#NQqyV-m&u-PyQ2t~V=ySu)6PhrC;Vg(99LW@vqGH%C{*fZXl zv6BTwT17%20Ychk5zB`ZaplB6;L2}6`y&ts#0lQFdb+;4XEvAW@nxpEUENPT_4ZS5 zJ=Hbst#JKx^l)%KdFMaBZUxVSwMm{IhMO1J_m6vLd3MmbQhXVX4x!Y&EKB#dgUw;@ zG@qRH`gyo=|8{Wm?8&43(Qx=|(jO1bru+Sq-ei&;5Atzu+&}&>T>a6H@81eK&VS)G zvPpk*mQV0Mf7sWeU!CLRxU~*&S zWoy+JzP|uI4eGUv>_L8b-p{``9vxCOKPWZ+&%p=hdH$0;qn+2vuaZXE-I$eqp>sNtxtuYBGpi$gM%T-hMNtx% z6j`Sjs>p-h@KK)e&XwYeWR7<)%aSkNoLCbsZ^GcMxmf#?)V3KP7EvPHzQ}$soXi3( z?d)_nQ*8C!p773q#0S;UDfT8%8)KxF;ic_3MZG*kay zHwW_n!QY0*4C8kWG+u$m4Xp@GXutXIsz+eIb%2QqFjM;*8#6z4MS|K9C)7iknh#~_ z^`Ug)E3OEBYzn>=Uw`}e=l|OD*o}pJ6)I&h6(5vp8b^kj#*tSw{kc?A2`4ReHIB+Y zi$T|wkj9|v9)WX@Bm&^5`#Tbm(5gfv zbSV;%*t#0&xMzk$B&;EcNZ7o%2btG#4>I3W+#3?XZETtB#Gv)MvIaLn_7$F^!q4 zu&5wcrH)QUz3J7Jm6h+lY_)vl@v!W<3C?%Hj}`1h8Z%NbK;*SxfyjFlEGNEcMc1jR z>sEaG^wtU!ORr~UwoXxH%}_Iir+!m-!l>7_qTB`J%ntD$C>bVKP!>BiEq3bHVlT2N9Pz202~#;)#_WliUJP4`7sG5XhLX~*Wu}lbze27P zPG);C234jP!L6xcN%&)E+5@CBW234jP!R$%0yFzE?geJ{O~k!$l6e*P z0+abN?gb|E4aL1HeDC)= zqeqzlEyqu}0w6%kq2S9|hwQ+7XIpTDQ-CNv27wv7Qc29%{i)-MjtJlq=Q&!$a=HqNU#6e(u+BD@Nnf!B*2w1J>)2jhH;Te91$$6=TXzA{ z^(LwsYFHH4|HB*j54d^}CE;GM16SD6SEyjH8FDeu`deJL+o36QPH@Ggb3&*SvN;lM z`YQK=tpsIH%XT6W@KS$GI~5IaJohxttm0ah^u@{I+O+Pp-CB4J+ro?~Zp`cI-LAZv zT=!oZTRQ(Vv@AK8e;Q~@=F^6}B(=2JErYj5D=kZE33FSP)bgfUmazHI5nC;Ly0%zc zyOuV)W$@BlrDdd+Ft=r-mN$isEC@Gf8L4H@!A5Ot`1HWrLN(98P2&|AdA?<9fpgvB zsceCYO1K!OFayEsF02Rdkd$0qE5!-k<}lS-UJr57T6Sz4V33yAb54+XCp05UPbhm8 zJ)v5nR2R|t;Uk++8$LyC=s~vlsbjT4q-(zdKd??%@+0eEKf#J!P5-;G$re9#tTu== zw#82!n@=0+39F^e?qFEnsI&|g8eXAgyq32DI%pkV{M7MU_H^w9v_Z>wEp1NAoH)Zi z1Ap{YT2{-O!X_4k8(sopOlTy6y;(7&F zs@}|csHvw}sFf)RSGojg!P3IprjZR7Qwx~H&b#ewVma@$tX)ge>(2D6$b#)IfqBNP zOJ218+6ul_{MO6?t({V-{6?~c@#XX3vfnG4)JL51&`jzYUX{yU~l*H-cMpFOeS|<^8Gr-I;gSf|cDH{|8gd$c_L2 diff --git a/tests/fixtures/onnx_genai_workflows/diffusion_guided/policies/latent_row_shape.onnx b/tests/fixtures/onnx_genai_workflows/diffusion_guided/policies/latent_row_shape.onnx deleted file mode 100644 index 254b8624d88e6e5ce78ba8ce729117ce4d6813d4..0000000000000000000000000000000000000000 GIT binary patch literal 0 HcmV?d00001 literal 382 zcmaiv!A=4(5QewQB6W;P>IFF%G^h!OMy{R|Z~F?yG^G@_acP@v*RUa8_DOsI-;9^Q zrFbDR&SjFxH~)NtPsq{IOjYR~H!*x8(smiQvVZ%L*!Shm% zZgH#InS_P3mqeis`xi>mQ#jUqDV^m)k|-U*m3?1OqxH&(LRn8m&YfeKk_9hBR+0Gi zEgeArfNC(6@6s9t?&9p-x#tmB^z^xDD zLpcMeNMc;M>&^LcdhWSi3FoJ3AQE%=?!qeq&#;gWlZb>#EJ%kV8U>oCdY{PzaQ)&guUI7=OfcFN}vT$$&UmM{}L|fwH@sHJaEtwt`n^v1;P4*lUd7|fBpRjpBja;p^Z5bk zrtw9RP2)jgZajaFygU8z+0i6Tud~5yG|i6&<2cKr;V7BKv%&C{x%tB%J%2*1+OO~( zQ8t)NlMH`R2W8hm**AgO&IJz%u5_e*xM;?5BJD+d_BsvdiOtNWF=M1Jps<{TW`te?EZS;;jqc@ z>>50KrSSf4mh6rfw;>VJUZQmihzuKxq-8jHipHLh-hlwI@W})~n*5 zayUh=JRJCBmPfV+;Q$tj!SZ_Bc&BT${%ELnpnS3|aG*)PGrEdw-`c}#0SKfVZgk

FpTy~PouWvY1GylFIApK zZQ*GY18?DJl&;p(0^w%K00wcjMgY=igym>3 zp}0&|T(%0uK@!3dMy)jk2O3yYr<*DL4$&7fDwT3)EMm;M5oQjK}uaJM$0^rf?WYu}B;- zkYbIAqtDtlhiyC?@zE@Sgu^Uifpxbm7ZwV7yYmOS+z%wFZOa<$gBuSF)LCqIFi>Z> z(S-yuj-uCU*SZUHRssh+yuP*13*0bXzqgHdy0E|kRT&pd#?E0Nbv{T>BM;{~v(Y81H zCv|_#t1`+Slh+&`TiKce##70OdVEDzh$XrK&*yp!&n^4_PAw#n=UE3usc;3acZc8V z@=XXB6o5`xCQVCK_heP~^s0KDU5y-xr{+lN?9#dOUuRb!R*m}i4@EEup6l_mNOq*T#j?dQYX_Wn~bj@R>Ja&KzRNU zdB00&W1Tw|73KU4_`c!HhPwA8&+ZMvP~b( z4}IMw2c2EBE8xEPs&w-69r6&q=@^k@*+p6dr-xGYLoIrqggi|yiOcg zYp=tg-@mCHKS|bT_a!BHgw=*LUgz(snVrS;Ht$+#|4<2)Y|3qBfGr=Fu4X&8wsO^0 zC>1tqw)&ztN??1*pgpRd6RACThwh${(Vm*Z4x%FI@PAVu+j;Lsyk89F#D&kg$d^ zScoxPj8P>t>~<*BL=(HMR+^6;Yzj2^Re^U~sVP&)br*qm*<$eAhM22#TWnZWwwe$= zfPY=k{D7)ui!X~x$Wt~xomb*(A$8bZS!5HNm6g%htww|F)&f?tTZ|@r2eLj(PMGJy zXPLir+b^pb1a5+^JcV%Mp%}8l>~cJC&h1lX9MH5aP5F|$?Y%PNU=_Z=cwXVs-J3*9 zbr)EkmhSSlyc8@SLUbYp^lxlzZ+Tniy=FI z&QBQOqw=g1J{4tnKF`|Xz?d@tsd`MR07tHoiul-1v&*<;!G{^S^+7dz){p`< zr$Xm_J7sj3m{{>8QIvS)OCjbgN@xw<;X> z2kgQPEIR1jAvr{h4>#Yq9M?oRrwjr zh*V3HuYhlebeHU9s*L5zvI)Jn85i{@A^tNlx(~4A*9*xKWzCD>8nv6ET03brdcYPS zfdW@>%Z_NzvE`!z^(fpg!%+uXe3h7iw?=+bh~Mko0}#1_Y!b6W(p$3k%6wi(`Bk!G z+OL8F^EeDPGflzWBG{%N*r~=H;R)>&-vZ?%Y~{fy0NrMZgB#y$`C%djE2NY~1hWxq*A6r+Cotk%&8`(F8cxTs&* z*Km&8B!UC-NThR}Dyi(ae$@DqTQP-a$t{v|}*VV${xi_yzD|?VU MhpOCJ3%vf|7vDgf`~Uy| diff --git a/tests/fixtures/onnx_genai_workflows/diffusion_guided/text_encoder/model.onnx.data b/tests/fixtures/onnx_genai_workflows/diffusion_guided/text_encoder/model.onnx.data deleted file mode 100644 index e69de29bb..000000000 diff --git a/tests/fixtures/onnx_genai_workflows/diffusion_guided/vae_decoder/model.onnx b/tests/fixtures/onnx_genai_workflows/diffusion_guided/vae_decoder/model.onnx deleted file mode 100644 index 773c8b4b2b4432b5f206ae035966db2fd0edc2f7..0000000000000000000000000000000000000000 GIT binary patch literal 0 HcmV?d00001 literal 957 zcmchVu};G<5QdvJCE;2GctQaxR539`N!XYWY+2br$Z{ML8xkkPCcty>0xY}$Pe@u+ z)I8dvx{mr4_`m>(#nCoEMbZa`h1BY_ z_i9+-2CPexF;x)CP%k`$NpY9Bxs*?emZH!uO^H&N3s#bn@@Le)x(*M(X{ly2u%dav z)J)N)^4nB?C)MEpsDd_C&`D*ikm}8VwUlTkH8hfJRDPoh&`@L~Vb1s>6_nw5(#O-` zyjoGFXjBn~V@B&=RbG}0X2jvfHc$lnXqV<$k;~eMyCiiZqA7>IZNdcYeJ|lmBwTlo scoebb(LH>+{{70qA6MQwjc4}|fNfjItlb!`p8rEHSni#f(A(Pi06iTsFaQ7m diff --git a/tests/fixtures/onnx_genai_workflows/diffusion_guided/vae_decoder/model.onnx.data b/tests/fixtures/onnx_genai_workflows/diffusion_guided/vae_decoder/model.onnx.data deleted file mode 100644 index e69de29bb..000000000 diff --git a/tests/fixtures/onnx_genai_workflows/masked/model.onnx b/tests/fixtures/onnx_genai_workflows/masked/model.onnx deleted file mode 100644 index 608298e124b93bfcaded679bebf1224a171cd06e..0000000000000000000000000000000000000000 GIT binary patch literal 0 HcmV?d00001 literal 1104 zcmcJO%T59@6ox5`T86P3)un29SE#Op#$jM0tGLKp5drYWULxwLg!i0+6F;K~Q! z0elVL%Q6T8OATF}OD|vla}HiU{=)EvATO2aI4WzF|h6c zAhOb>7g|}zl%u<0VQi=G??!kOC=-A#Yjp&b%Gt9X>Z7H3VfDQ>V?khwkU+}InWKB* zs)rj`DeGReICRU)5U>|%OPM~(&I|Dqc{v8kr6gprrJtm;;rd9fo5zVV?nh%?f8ynJ zSV+|@14cy@2N}QmD3}-S`&73TOrnjpqDGx=3T4B>RZS=--Rn@wg9wiTa{?lTHnU9+ z;NR8KisbHcavVOu}@IrI2?AM<-} zM)ei*@@cZ)8>g@T^MHI}kjtl0dJ=cd###9AC>qA$p0&zv21yrM_A+bj5OOIVbhj3t z8;d)uWI4^EQMP43>%kgn45HIG9Y&qlTzK#@xidUD+)W09$7yHO8)mzmew3!+Q7;}v zqt4M2bMaf>ey~lf(ywrgFzqD6IK^L(j}3G6EPN{&z*K`QbgXr8&m?o~dhuXqgEY>f z{x}YMgKTR-qfoz4g>uV6xsp(ugb1P1H$bYIE~tKO!1ovGBQk2PxdI&@j^p?{aY(Hz z#VtdS+O@o`OR$m-y)s;$1ec;iord+?P#|)M43X01p;vqm4G!bbw^sQLL(aFCSxXmg zB{K`kxCIh!^E0#GF*D2YP=E@|wX^UWgETj^rnNC?FuQDdVap=j(gqj>jtDZhJ&|HhfBXnGALo_%Lau~Oi zF#bjH2N6nJJpa@{fH*8f71g{VJDNkSP2pF$66qXV?PCMQn#|lJC9#BBJz1@uR<-`B zq*lM8R$p?!O_?(Mqo&aVfn*ih^i*y7%ncK;{s7z%uH5sj?Ch!*!WdCB2_2DQPCTZLa zv*aWmq>5TyM~$_vTZy%f86Mwb^7;PS%5 zPqFN}u$~Au<-&Rb4PAk#6trC+|0oewzbCZxO{DUnKg9@SBl?s zudG(OC#m%Ud==7o2^_^m>7E+8sZ$Bv)R{x*E@N0pH0bO^<6p%ptuxU$4@0+vVKvpD z^GG#5Gt>o-(o);OaQx0Fjxvx;*9QEF2fJQdTr#Jdr}th~02W ze1%8r%sBo!;i}Xpq=Plo(CLQEEk205<4*i=d>Ycg+MKkwkb&J0k9O8c{UnZu-QH=c zH4FKX2!tXBp-K)y)f{AA*fr0=u2sdZbrIP8Qq7rYyAr2omYsCce<=2>WJwnFLo7a} zyat8zS*mCVEdfqy9(iVs`IS~HZV=!uAl$F?)xPqNU5 zNRC&Ak`Iv_8oE%D)CF?z+06z61Gm&Ab@{GVb-@ZZnN|=epUULZrx!s!eL9;!`Cf%U z`QHBlf$}Tc2EM#)psdI`Pa2e0WMADo@O^Wsn+(9^-u`e5EQ?!wU@x<#E}O^AM4WYE zWeQFA=SS23FiCOzCN3&KJ}i9w}*!PEu9kdVYnP`!lJzPRo)PK#rHP zPWuo|A!E#XgRyE|0R(id)>YcT3o2U%fxKlfzjF_;fd>n6emY9QDqzKjq7oaxaiHR! z`D5T}uLWc3%+EW>HQ|8qb>yo7@O4oM4uG#~tOS?rJ*brT1G&6Ew@CU!JOP0RL*$+G zz+>4T9DN9Yy^A9d0kC(CM<9MIwt@qm(Xjc=6~-F48fA59E!>+|`N==d|VPj&@i68A&xTAUQ@2S|( ze2u(XH_W?Jhb5~}6HhwP{(vL5K z?NmE%7JgUh+JB8mocP>t-v5eB-4LBwU$75?N>);x6a-}}Gk!rs`ojI|M1*sbkUDfO zk;S@kyrpn6?5X8FH>7aKUMV`RD~E_2^DI|wbSM>|b6u5s?TEjF$J>Sg2R8;m$^>Af~vdBgz(^S}4r<6Fw7z>d3YhK`CB z>~z0KZUY!MCtyh7h9MVB-M(e-8RxxMlEYKv_^Q0Jx;}GgV@9CiSfO#i4&xN8RNc8v zfidrpHvZ7J5d0VAU9M?Esw>k|wY7;>diRJ$#YLaU>2U%QJ? zNjo8E-VNEj7r-I6g&k7A>HzqQyC}6(jO0+3+yFvnZqBrMZje2GsWe=x>b9z0Xekop z<>K<<0OmWXN>)16A)$g=L`(J-8+O$qN@_1^DEJYruaZQapEJPHQ zCp6!N8X|aP;w6_^-FYQsO7}!#KK-9G=1;*wg%{G{0^S#dT|i=3(aH-GTj05m;x45@ zWiyb$1==q=8CyjP;t$ljkVNP*n4)NPDxjmvTS!O)zasuKGwKMLQLC6Q71G1xcHGhO8T+=`nH|3JkIazuWj3`REIrkArHIwd8nw5y`{5J z*o{L?hU8<#LwtgpJlLMp0`aNTD*B2cSQl$RPmyloJK`Rmj|PYE=p!2$7@GyJ@SHPc ze3IW4qfa^{u4F%j?)cy!?qv8%4#d-E#(nbge}ppZbxjNZm*pP79rboOn)mo1Z#d74 VFO%(=gYZe&zq8z5A`2T^{|1aL7~%i` diff --git a/tests/fixtures/onnx_genai_workflows/speculative/policies/adaptive_k.onnx b/tests/fixtures/onnx_genai_workflows/speculative/policies/adaptive_k.onnx deleted file mode 100644 index 58de18d70f6992b9663deaece49f999960fe8032..0000000000000000000000000000000000000000 GIT binary patch literal 0 HcmV?d00001 literal 34702 zcmd5_TaP5kRi5hUneNWH>`Zy>ZOVYhn2&LR=*TFT5D>16ckBPe{D;0$)Vr<(#-w^-E;CU%DbA zD)Y+|=Mv|fI8nQ|hBu!s9v@#UKmWhq+zP%m7~XzXeQ_~guBP)_1cGnor62U-hBFWwmh32&W6_xUJicp?A^B?E#~v*%fs{I zv(=-+)9G?qy>mP}pPnDS^Nr!wSHE`fAQ)+*#67Cz;o@wz6eDz0p^hqaqPo!^sz^r_ zIZ+KfqWZ@{us2(-j-O6fvt__0leI$DiiUg7s<%#$4`)?8x~)FQ36IA&WUHsCO>`nA zP9pX=wW*FOb)pJ9POT=KnT{%RqPiM}OBjo#td(^PmotXT`(?O7Ct~3wVvkc>>ZnR5 zsuw&m+;;}UJH#4D_ljZ)&spAU9WM#nN4}(U@h2iK9(Q;wvnd@{) zMg}F_)}qfDjtyL~lk;Bek@FG*m)kJ-;HFW8oPS}pc>J{)ZMS{i?w zTi?jA84<|0AxMqgYPu>G4?=C74?A zK}~IfsYT1zd^%-;@_*!z1#UAU-!_BELfXkf#$+MuO%|NAqx6a$w9C03@$F$S+T|VX z3P!u)>S;&ml{jd}?M39|H{lba?k!^(*t` z*Dq$XUqj3!9Nq6)DsU2x_u5OoRElOIV#7qlKME!yHcUk9G!a-p^<}wY!>Nd=Q^73P zmPc5>JqUJ>&ZkdS)zf9Lx0;^6Jrf^?yU(g$I6A80baYGnS6r2jucP6MpMDEanx~3| z*lWu-rf*bnHoBueDGb^8hHUnS;g-}3!Ryc+`&#E`i ztGE~i_@~emRkDlSL2#DDeha(o>PUr+y@@Oj1B5)tJiwz7&N zkuCV8^BH>1s|4AC=Cj0=OvXF)1z*@c90d0c7f+uauYj#uE#95Ym+K%Rb-*Ep-XVUE zFXIzKM<-NAe>}LQmhn$c7t__pui@{*!N3>!{)EUoBU{2}HrB}&dneoP4n!!gnXDuh zib-+4RE+bZDoLOi@n7*;0>wD@rWn672tM)T`1BM^c6PovTP&xi)#HmJXlki2C7KqO zP}nUlp|Jho(x?=+Xwnu6R2dEZu%6 z{na|<7HF%<;m5%G)73k(^Dmvwo?`L2tder{3GI?%{BnGkUw?Qj*q%N=Uhex`OHugw zCE%wMJN$HVCGh`^#CpQ;GWbIW{_yJYrK zT36$}pY!W#utiRqSZJA)Vxjof#p5cAN4La(MYwo;9SvV(gC?Vy+$<97Rs}-9dHn)2 z3136}S8(8Kpy3P7AKADpDFrA6MJ}cxz8Q@SWqu<=NUspK<_iu?C9sJYRhA0o|5dD8 zXH|x%k@}>FhNzKj`Xbxk4>+o%5QrNrc4q4aAh<-qBDmxgEF5nx4W$KOI%%b$pwduW zku(&oG$8uaNdux!eM>_sT}vqqQez0-%O`Lx@n3RVK7niLFHX8v8lem721ZPu?;z~4LP?B`BmD6JhL-fCk@;-05;37qpIEB7yX6j7VVjHhE-yo5=Mx(bc@o zM{5Iv5-hr1PYmrd42l4Dv*s4|CPid8-s=~jfjX?Ys0c!x{nfpSR#nnR6!$Rh2MzS zzcQ64f~;!)oh8vx6!u7UQCvx+nj+6-T+xl(k)kQCCPkwOOKC*X$~BQR&%$C|M-;qX zqm=Hikd`teQd5RrrN$b4Ddb{8L21ONj7UJajJyVv#TJZKW>~O_Z^0<_M1AQ$QI9Qq z0-M|_&G(_l(nS&rz*3~@n}uL0QuY28g5TE`f^kS_NB66gDC=wX7f5|vrrvio=rU>4 zVwAKNLo;B#Yz7gQq2%VP%oy2noospUWHVf%a*a#$pC5hMr%QzBacxj_vbFKDlE%wQ zzwxpTPN8f*!UAO}oSiSvPyT?5|f1oKWc&;rZS(mNQAt_F1gs7KCNbgCV zUO!h*MeMu^Cn%R_Hyxtob=M(kJZt~GOt{5MgsW03?l;d>el47&qnEU6HG`*EHmV(p z^yDiJj@v;hAO*5*n2_U~av(|CEESOM`P%uFtA?tyb zugUwtK8){=ToWuJ{qqd@eGVK?IN^dSG^eEN7MXvTm zM-*D?bVRXrcZi5a$N&E2=RR#UtWjsH86Cq=$fk&)har@|_)yag${#IX^!kSu5zKXm zz^39t1edrbYaqsbsW@yMF}cmA>O#5n$L{zr#L~anC69*BZStsIpq5>htx0tpVvkg3 zxI$z@(DF5-x+V!TQv%j5@?Z5S0slgRLb(jrwpuO>Gb`KWF1E}2Vms^>Y}O!QYGt=v zc;%AakSXuj&1+Q#;%m)A5D?3mzD_)$A1sQD}Xaq=V4*oDcfT zmXV?m^tOVb{kVh3l_~^P#J=&>lR_*6b=HGQcW`>D8=?m zH7w2>s>KsdwMd-fRr?O09O|)~h7uRlWBabu{I*F!j;UQZp)Q*&RZx@dbE!gDO(&XU zHL(piL<#oob~7?sUiV+BA{4o3i>2MFA&|+XT&9_a*tdB=C8V3npiAm_9ICbNlCAGy z45Q$6@4N1Hp*mYP3NbYbTHpb>l~+A(4xnyVA=jauEr<)7+l2~mySE>UY;rFFO>E1g z6KZq3;?BcB`8N?;)PDOE+cgoyL>}EF7fM0cRtfY`ce)x-3HN|hjREVC({i+oJF>POA=;M6YNCdCN=S&OokAU zc@q}Dfe(&JpCEXoKneXNE{5QRZKo?jP)3m~pq8LMHwK{&D;!HoE zN)8jgGZ?rJW72BkGg1eni0L0mOMfdhr~-o6*r3@p@!RUECjLOI>H5zEOO(D z`lD8xL-bvt*TsSXk%>b@$cp`)o$bi9kP^zd{gUyoct$lQz#`f}(K|t@r6w}psR``( zq)%sO8i25G)rT~ODB`}?G?b;vlJQP`#g~#Cs4lZFI@sNE^~DFf*Ij+d z!S1%JFFn}(zxwYhKThZQh2YNa&Tuc-8Qyq2T^+vjl~0mWwTw=DwLSZEX!L%@Jblz0 z^fRubBjwON4N7wouQ~Wh6uEK2zs<4ozIA9bhL$`vzsytd^Bzt*pzwqAeUkYZA)mre3;nHL7Mo!$ z^MB(+`#vX@*2!m7>*R!@PfjoOWKJ;c=J~pJ=ER~;E^SZVJ|RQ!qaYTt3{K>Psoh12 za64S|1}85u!bj2YX0Q#}xu$o-ZHKyNQ)&uvHttD*S*gZ=WfJ<@n^>8R-K5xj6Dt+z zO|Y{>5Ya#a6~s@}w-4i#gm-MGBuWn*9ZA~(y|1=##HpcdPkus3Lv#GOp{Seg`aE!< zjHw=nT4)dKHu~ga#cpKePL0k_(H+dlWNsJ`ngck2klGV;7=YFItz+|>iXLK_l&tGi z#A2`8eKTA4&D`wNW<60Hwx`R_L(35@I4jO7jkd7}{$GBqRrw80toM!V_l2ph*TRTJ zi#epG`b2wa1c;$C zHT={O%0Wr;e4J-uh4?wE>uROyjT8G|D}7yqt+aLxw$lBEwno5NB8(v7P=dliZAme8 z)Q6oMHl0>D7}+}3F%x_aagc)nxwO)J04oK#RyF8$$e*DWvQ+dn^h(PrrJ|{^+TeFl z{A}DsdyA}n3;u1z$D6*b@v#szlx16)RPwRRPoq3r4l!;nmxnUhIdP~~GD30ex_dat zDDNRFs3wYz#yzBn@Pst#gdOxGkQBojw<9AQR3#%EYlVHWkVF7QIf5wUY^Uv~EyZhg z@@n^hKh!X=e7=}GlnL5r&q@;1!-=UjURW}!EF5wWSUyiqLFk~Wjltw8;b2%G|FN#Z6m5y_J`B8;g#TZZ5rWVUJVNIy@*Cn95mIx-1(_|PO!6Y6TFSe zHF4+w9m57%)*hH4<=SoFd^OBY=#V;yXq`1|HQMmf@Uu3JSjVc;l%@{SwyH5*N7;Ip zhJ;~Aa3`JEydi)ZBd9#2m(T-nv^L(ScDvR-D|2W9il=tMjQc#$(kar_q_K6T-D*e7 zxC?vII!&OQwJ>w&SazRFP&so@S+4Hd42+e`by%>;<;j`Q{O9cxeji0kQH_80 z!MOdcN-?l&pr`ra5P66CM?}Oo+U$nt=o;c)+V85#Tk6*^&i=btVzi?N-)XWOVpRJX z*V_b#=;#{gY13Y!r)!``gc> zJ{o>@F`qww1yeRXe$;&bm1ag>QNIoLYWV)(r-BF0`zZv=;nr8b_J!@>+P(e%2VDnB Apa1{> diff --git a/tests/fixtures/onnx_genai_workflows/speculative/policies/cache_length_update.onnx b/tests/fixtures/onnx_genai_workflows/speculative/policies/cache_length_update.onnx deleted file mode 100644 index 72ad12c9f4e344e870cd12220f36d4f2593cc2b3..0000000000000000000000000000000000000000 GIT binary patch literal 0 HcmV?d00001 literal 386 zcmaiv!A`C9_8vCK3xkAz0u^oRWg3)KzA zxO>Sx@8;a$n4NBHs<-~}GlX~OmzWVmPUi)&h%AY61}mkwu?qS9%}IMF6Brq>L0<`p ztd}g{s(Q_1Ys}6|rz?tODZJ-}MkkysY8F1fBy*S@QTq<=rLEBK2{zV-Fh|FQQtOLO zg`L4_~e(FQ-eEJ2Sf_QEK diff --git a/tests/fixtures/onnx_genai_workflows/speculative/policies/grammar_guidance.onnx b/tests/fixtures/onnx_genai_workflows/speculative/policies/grammar_guidance.onnx deleted file mode 100644 index cb4ad9693ecdab7f45483069d859c941cafc93d6..0000000000000000000000000000000000000000 GIT binary patch literal 0 HcmV?d00001 literal 1926 zcmcIlUuzRV5YL7*xtla?Tx|`fEo~?*1RIhzt>BYZL`qCYR(wt50`Nv%}oXZ+`om9W9TuuHUi;{K=m+_>S=EkUaN2E~qC+ z%~+q=p5JDqR@Nd%Wdrg<>I@j^I6f5{qtTbqHA2k-=pJ<$52(p-s<{W-LEozSo_E2` z&<;e^bSUSfYqOAsX7>`O-@R|{fUy)c&Ox~82aJy+;2VM%GRQiEz@3)$nI|cDza)1t6f-J?-)A0&1xex+ z<%`Y7<7DRpZ(1RBT^bT=V7KK3bLs{T3oC|(in+4|GMieY|MP=r#PJq9&=9^4IlR)M z!j!_0Bj%g5HE?KndHyK&FsNw#_UOGOYVcjSqvdd6o;|Ih^PO4LM2mPG3NlKcMYSHn zgvT5I@0DT-0np^n8IYNKn=x6}a27JMxzyA@f5{?$W$|VCHxs?x*gfCW*i!k1DxZ_( zw`2NN?dl8d_9SgFfvL9jCEBjg0qjn;T-9Qp+&seey=ve)ws}e9$cT`d0;fNHzCtfx x@4u16ZsV=-N;{OY40m{R-{F&aHV^r^8ISyX0a!%Rs7F`U}-(V^IJA diff --git a/tests/fixtures/onnx_genai_workflows/speculative/policies/grammar_length.onnx b/tests/fixtures/onnx_genai_workflows/speculative/policies/grammar_length.onnx deleted file mode 100644 index 36bafd06a4dc6b8b1db3c42a428a6ab6ab0b2ac2..0000000000000000000000000000000000000000 GIT binary patch literal 0 HcmV?d00001 literal 394 zcmaiv!AiqG5QdxBns%gA_R^YzR4IB3diA6p1HM2-md$pO3?{o_cLNQghrF~`+oru# z#M{7p^UwDW&)8|H6RPdQw+~;?FEDj5FqGwv1tlpe>oSBFQe5d3^1JJk%^Z(mq(q6f z5)xT2UchDbo<&-zwUve{7s*0c%X2~_jLaM6zrMwD2#=_32DegIXtxBLkaC!z;rqwS zX(z-cFdhbM1OXc)!pVF&J2)P1k`NBszph&E7OuLsQHE;c=18lk(gn$eORby{(lP(z sv-O@}@h{0~#V+<#7ZQnjO$7~0MAQ&I{*Lx diff --git a/tests/fixtures/onnx_genai_workflows/speculative/policies/grammar_sampler_logits.onnx b/tests/fixtures/onnx_genai_workflows/speculative/policies/grammar_sampler_logits.onnx deleted file mode 100644 index f674183330680e7f8e3468568778f80f8327ea0b..0000000000000000000000000000000000000000 GIT binary patch literal 0 HcmV?d00001 literal 631 zcmcJM%T59@6ox4ohv5KGyHLPr5Q#)$CJ=Y7yev$31dVArr9fq9$7zQ^j2m9V#Q1Q& zhYk#=kQi5X)ysGOd_Bj>VSA`N5$W@H3tkY;kI1!>+E8gov%LAWm#V`_v+67_8YRe% zs2_16q$z$Z7kWeKb^)@I4!I5}<2co>!9g(Sxk}2hW}yg-%Y3Re=?fmxko6}xeRtok zKzWUtWf09&!1auPR|K}NdctT}vBbGJ&XDA(T)3g8&qFd7ZdL8%*vNrJ$3hqXQ4nv> zJ&YP~xv;rCTPouds5=NBK?ZMl)UZAPb~JyCxMUpi7)FXwFY;+P@lN6AkMKV1#g)6? zl`Hha&RREYu$&A|O;H^VlhL|8F5^}kp`;rqU$6;bN}7-|gVRsXQ*;TnKS?4wxU`^> g5SK<{;fvV#qBTXwQ28~3ce|XK`4q1#NNp8A0X>A#xc~qF diff --git a/tests/fixtures/onnx_genai_workflows/speculative/policies/proposal_metrics.onnx b/tests/fixtures/onnx_genai_workflows/speculative/policies/proposal_metrics.onnx deleted file mode 100644 index 0084a107f34b0d3c63c7a853108d6ea7a925af8d..0000000000000000000000000000000000000000 GIT binary patch literal 0 HcmV?d00001 literal 1026 zcmcJO!EO^V5Qei$)9g+YwX#%51FA|?QPIPy3pY-H)Jv}1ijd`aCyB|%c5QE!2=$T| z-~~7!UXWM7*@duaRXF7~zvlPMH)A{6BfB#-6oq~Bvjsl_*w;pBWig{pO;K7B-_Q>e zrjcIs+EpM`=2gF0s&jxxD6?+xJqXX+(6o*j*KK$;=s+yl46SAyNn>yd$9g(Wl$3MJ zjnFRPIkT2d1R7@eWI@7@p9Uw;yW?8bpq49*cIEgtWV!*F-hm7^AXmuSg8Ul5dU5IB z@Pfp6!}5YT%zB;6xn?q>UUX_qyLK)t2zc!Mgbjoe|t#CHLRrQ#guRT;zh8#oKflynvGjYv5>v+?7A+{SGt^ Gy1xMaltV@U diff --git a/tests/fixtures/onnx_genai_workflows/speculative/policies/speculative_acceptance.onnx b/tests/fixtures/onnx_genai_workflows/speculative/policies/speculative_acceptance.onnx deleted file mode 100644 index 912d023e99436d4787610497e7782c1fee02fd73..0000000000000000000000000000000000000000 GIT binary patch literal 0 HcmV?d00001 literal 8736 zcmd5?%WvF77>_rb?CxZ@NlewIE=p6YAhBxF^&|V}sSh|1QIJ|8q#$HD-bu2_#@^cA zD53Py6R#eq3I~u-L2y7^xNznV;g2cbj9)W$n=Xe5?8P(V@yzGn_xm36d8>2s?18s6 z7zfuMzm5KoP$RVc9UNMLj|Rcg^W-PN=tcw8T>AhABiq69)W#LGJlfw`^IZ2RaQwk2TyutY5Lmkd z?AyMxdn}he{$yhXDY;o>4ITVNq0)+b7+Z>XImh;Uf7Gb8J8 zM%EE*WM4{|m3300cUhu$PmSomWmM`SABfw`^zM8e&&jxQqK^2NPaV zHucn6I#O_@qLlenCaXb|M;(LKQtFw9&{`aFHG$N6D#;DTMQEC`khp{hBTai5`TNT25{5uTN-$Ja;5ar40RehqKE_*ScteF(!E`K zge}cf8c{0}#Z+f#SClY+Q5Ti$EuYOe^}%*Lx~`m0sFjRGvaPFWT#9cr6*NPmQS(ey zTBvMPr{do;sVwRZ7`yCdcG>3avVCfH`78Aj5v8~!>xIY?i}+|4`=IxZ(ui6Swh#tr zR}^YL(J&BMM7zL8B6%VgBV(7Bu`5`{UrscX9&aeU*if1XO=GveR08MQtfLu7q5N=3 zq|bXSAue1I@M6>H$>qN1V$Rv=Ey2$27|~6MYoi&orp#q2v?er%j%+d7HMK=uXKhi} zIa?H&Mt_}GOcJ$qo%;QL-_nWShbuKYo%sE}FuxBxA<>5J%5kJ_I}RR&aQ-66&fJ}W z8_Q#=)6fy+0#cqp-W3kqF=lzPE|e#^>_u5vY#C6qiCPpu1H3iuih?i4Iirxs8*CzfdYpSkIhSC< z!DVRjnM3PCH;7qh=*nW&i(u2$S~?^Ox*vISHH=}5oM52 zxQG*2qeoiN{_>oQk3uU_Bcm=8%Jt&1Y580*!p?YThl4|GMV4>74sI%>@NkQ%C2mv` zU%z=H|E1;{TCGa*Tc{$>Y}ui+ix%WFec#@OPdONZ8?1-wkQ#W7y#-tO$D3y1s4QPb zm#P){;t6YPo{v;0&!YxRt)){pFC_FY+}eVja{g?BowhFV0jIbuh(b=J^hS*lWosF+ zdMrK&E3-A(;)IK+nj9nwtM4sw)m?NYW8ho5@F79MSWAAMP+!ispzeeGU&fwsGTmvylrQax6S15PO>dsKCX*S zviy0@4QW~`CskuwRh<;*s7^M?#z3I7s$$9+fl7PoOVWX4&J>!8I1?AUCf+S8fW>qr zE7A8%5|K#|W(9|12E%+Xn505eG?>>3%c;artjPO~J^E-=Nhsdrb);upy2vRHPL8ZJ zl*!roC{=83wO)GT9MMQ}K9H+u7KmR0g!ohe|!lGJ!5m8|6hb4<nkc7^A52;*oF`P>eUsM z`~Lx%$Pd7b+$VlACqD$UauJ8j$`7%#pj^dGA$_<{0t^|aSLi2UfsCbX828;B7S5XZYUt?RUBifJa}T zFX7FHbCWdb$3eU_5SW>r$^7!)KdxN-VEBNtmM>f=vN5LjeHLQRmgl!3SJBL!k1oK3h}1=C5UXhGok5`UC-lW|B-{Ug1ZGUX2Vxb1EJMbbAr}aMKwV%(*qU&2_(2a)AZP` z8MZh|P>-E*nz|LR*ju0)g$XLI3h^5mHDcis5>gK%8A`>5vA7Xe_I4)Q-Sgaao;bO_ z)4ID2r9@R1Fxo_(RWQ#4S*zmyM7iar%)Coe58CzHE*Cn`@2Q}t4$cZIq9HhyjYkiM zIo3Y3@H#QpB@JJWLMaSWJL_wAnD+Zn#MZ+Z%tY_Jc6(eC2byMd_pi~;c{Y~+H~Htr z&Sf^5Uo&yx03*`+Yhw&+-sWIu&eQ||8`SuqIIx#3o}ZyXW5H&#skLhh3Y(2zezsJo diff --git a/tests/fixtures/onnx_genai_workflows/speculative/proposer/model.onnx.data b/tests/fixtures/onnx_genai_workflows/speculative/proposer/model.onnx.data deleted file mode 100644 index e69de29bb..000000000 diff --git a/tests/fixtures/onnx_genai_workflows/speculative/verifier/model.onnx b/tests/fixtures/onnx_genai_workflows/speculative/verifier/model.onnx deleted file mode 100644 index d994492e184284f711782abf7a235bdcd1a045e8..0000000000000000000000000000000000000000 GIT binary patch literal 0 HcmV?d00001 literal 4371 zcmc&%-EQMV6wWwq;-02y7`B_>6-#7A)yhawoS&uvAYTpD$~P~EbZn=PofhkUW#2S$n-eTR~ zzM}LyO4@pP*%V0IMsZ=^5?7rxSDkfS@o;7ZYK}k(AS5Z0hB`d^p4;vTC0)G~H3V~A zV?$V3K}_XXFU?!;|KRNhkz)%wB%XS>I3u5pMr2+i(Ti;~Q0-$FqDS-rM7l~+&(kiv=Yuv^Ko!;nrNozn~Ih5@lU zdNth_G2AhBQUm4|sP^SJ5Z8WIyr=WDyOO8fb$MFPFi(F=9iQF`tmN%*j};=JK#tEn zI}AR_v4?$q|LGhRvLfzp1lPt_K3`qg?BQ)D-A-u8wbqmddikkKAA8{;lAcFE~^V)kAHLidoV8hot$W+STU;ixpLN#E5mEmlw_G zI2{;;$c|4@mU8(oQN3a;{Nyu3Xy>QLk)XE*dNth_K^+*exL!;)@4|-Q*aQt#&`{+G zJx6L*B>#zE*SDFNKBKRRa2zySO=j>EPlL7wc-v>W%j#J)NoKy?T47$vu25iYMyKOr z_I6|HjH=?p#9P5oqMW7_7dI?Lk8;V)mAVYmRr}doKU4OGnKkJ9JJ# z(bN|wYiPq?!V?YQ$2V-CfPw4RI1r%-DAO;lLzyG(=wp!U1~%aVbI`cKfyx%wj~tNW zg`dMSw)n{nixt?zjhoQbKdg5-!bvT0W4ADX#t7u@ku+jSBh{q%NG`c62}=tXKxNzI zCcDyGOx!I`@5KGPeE)@%#~Kva9tz|p^7LLj{oPpCuASVMCky47nEZZ95#SR0S4Fla zCzzXdorM&hz=5QQzogwej;@F$$K2} cRlU0ko;ulEQ#8s@Vp*p#F2v`d3b{({KgN7amjD0& diff --git a/tests/fixtures/onnx_genai_workflows/speculative/verifier/model.onnx.data b/tests/fixtures/onnx_genai_workflows/speculative/verifier/model.onnx.data deleted file mode 100644 index e69de29bb..000000000 diff --git a/tests/fixtures/onnx_genai_workflows/static_cache/model.onnx b/tests/fixtures/onnx_genai_workflows/static_cache/model.onnx deleted file mode 100644 index 6098ce3c1f740dd5f0a1f6e9849eb9a3b1d86fe0..0000000000000000000000000000000000000000 GIT binary patch literal 0 HcmV?d00001 literal 2311 zcmc&#-EI;=6lQ?Ja)3yuw#7~SLyZX+gtmaSCJoUWgBM=ZD~#C;vtwBbv+K@)?Oo$5 z=mYQozJO2Rg|Fj|IJ zYuy7qa>~2I5Y%><@J_%NgmWX~h0TQ|TpD}&joH|Lf3KuX7s?o5MEXN!SX_$cB^smj zs$h*3H#0&Ygb;&|mlA30>v;vv8~IrGOR34UoP{Jgw?r{EM)6fa{E8GE9&;pXR{0%Dxg*}roJ&st_tqwu+X)r(6xZjn<1gKwdh<6 z=v)u$T>q~+e+%kd$6MFgGF7;-K=j7iMgNWGq_Q53SC!P8EigW!(=1`AHG_8=tnPK7~x@dUd!4 zuW;IAR8@Po{W91Rp21Tsjq9+Bx2ZjFY|gl>4Ai-Pr9&j`f{u&P4Z%JP8KTO}e^svs}0st0xwy$5rv49+QQgi{dn{^v=E6^g$A!mP|O diff --git a/tests/fixtures/onnx_genai_workflows/static_cache/model.onnx.data b/tests/fixtures/onnx_genai_workflows/static_cache/model.onnx.data deleted file mode 100644 index e69de29bb..000000000 diff --git a/tests/fixtures/onnx_genai_workflows/static_cache/policies/cache_length_update.onnx b/tests/fixtures/onnx_genai_workflows/static_cache/policies/cache_length_update.onnx deleted file mode 100644 index af3bc3fca3b73acdb986806a33ddcf35d9032ffc..0000000000000000000000000000000000000000 GIT binary patch literal 0 HcmV?d00001 literal 930 zcmcIjO-}+b5OsxRVU%c^7;z6qiP75z^yW!Ca`0-RF->X9ZgpulyG0>J4_rO?muG^?f4nPmxkXM$A4+~*~a+4JlnkM**TWH zl!Sb(G)cqV`-JDb)N&3b?u5Fmd*G)*2_ zi@B{ujuyK`Hq=-yGHCALRs#cr=2~_o%RJBF?}e;Zl`%KKHL7{c&U6iP=3T&P`3s_$ zOLMwfC>doc6?BTJk~%lXiMHQAW?m1@{>{XKp`&?L2@;XO>g5w7>?7Vai>X|4CIwWyU}uSc5|BXjAtG{?=$bmL|%uh=g!Ew@LvC? z2R@fTZRR@XGaviT8L>UMc8TAc7&C&~N;PcQ&Y1SpGV2}Hz?yBG6K`gi1Qw1Cz~1ca zvd_l+)f;&n*ZZfbPwG#F70``*x z`27Je1Lm+c#%)-;#P8c)1g)zy<{bvEtF9$a@`Cr2o(TcPoJI|TOZ?{5%&^C}r_`br zL)KGQ!){)%z7|E+Pm8S2iL5~($0E!1aQ-NPdgREa#6IySo*=Y-Ahfq!V)(?x1En4} z8Dj&r67}<ln1KWSU_xAihe|bUFpxvF%zAl2{8kDiAA1L^wUtKFVF0u? zXIIgAP?wgQ^Z12x9$Ii75>#3{Fe>eu20gQyQjeM;T1{OI`>l0QGOjEySNZ+P44hzr zfJIa6O&wp^rjMCFe&?7As2Wle=Nt{R6i{ZTrB#=uM&(E@M88|9l$Y6`rp&Hq%+AM* zr5)$A8ao+SorP-5GnD;$1V4q#eW*@)Yr9=o|mv$QFD8P0kfXM!p7%t?M_ zQ7k?KSknW?9CM>}MG|{KTzl8BPYBlfN-b=K1oqW3>*jS5(vllw=o@hKi4HJN$doKi&DG!bMX%CI~(nIvnjT}?d zG_FkKK#L=XoFlH^cuNa`1NFu3fXk%OP|T#!aH*N}?|deWhH(9iTw``<4%-XlskgK8 z+=^3aKvQYys1#(fEs;kEqQB1uLt>j z1OAU?N3Vfwq5@w6O}cU|Ss4^%X><>8b@|;>4nQ`)$imzQVjRD7h7eh^GwuW4wZUz$ zDNC@)18z{zGW|SJXA5VZ?r_;#8v-Cx21>z)3nTREkQ>?NGxhmE@st}__y!V|QiK~f zT^QRGE?VI#C{b*M)0cxVG_P2?!aG|*-4mtmJt%`!`brTLPPc@8`pIH@!W|nUe2vrA z>1Ob$aoRm!j$xRWry_B?!J6E)PQqM3kZ>8LCpc0Ic2XZDRKLkfsnvSP0AnGph&dKg zy9|U~_R{qjtWw)n=%Q2<T7d^QtJedw~uymZEB4;dZDmY$swtA{L%P8qedsit_&P696WO-EfnEyG%H->(73p=rk3GeKhSQZ`(S_J8szwxm_GZ1>tP8L8aw|1nJ$6l diff --git a/tests/fixtures/onnx_genai_workflows/static_cache/policies/decoder_step_update.onnx b/tests/fixtures/onnx_genai_workflows/static_cache/policies/decoder_step_update.onnx deleted file mode 100644 index f07a3a6e4e316142535edc6f60901042c484c18b..0000000000000000000000000000000000000000 GIT binary patch literal 0 HcmV?d00001 literal 780 zcmcJNy-or_6opv?S*{RaNDNXKM54xG!&+L2cEl&pn8`3Zz{oN)*_kB}W5XlZ*jV_Q zb}q}J;7_dRZca{4zT9){9ke}?T^{O-_bfbHs4&LYQfNbj!L8cv=PjhiaI0Z2g|iam z$HWgA=E791H!IyKLSB#&(~3}rQk^CosA1ocLQFIbxH1m)iPm_)SwI3hn4JB{=MH$aU=0ciN@{Mn6qtKD1}56v@FBx)1iZXsU^fCzgiYbSM0l$}G3*k; zMP$Io>Wjlw8kVDxnF?0ZF z+v=XePwZcv4}0s~pl!%Raj8WAdPAbmtP?m~?9J&j5e;G9qobtH=@e5bOhBkX>Bpy; r)rRK3&S(K>KXENW91>2r&q*-7bZ6Ev)c?$2y&j_`-q=+JQl;t#ajN`^ diff --git a/tests/fixtures/onnx_genai_workflows/static_cache/policies/generated_length_update.onnx b/tests/fixtures/onnx_genai_workflows/static_cache/policies/generated_length_update.onnx deleted file mode 100644 index af3bc3fca3b73acdb986806a33ddcf35d9032ffc..0000000000000000000000000000000000000000 GIT binary patch literal 0 HcmV?d00001 literal 930 zcmcIjO-}+b5OsxRVU%c^7;z6qiP75z^yW!Ca`0-RF->X9ZgpulyG0>J4_rO?muG^?f4nPmxkXM$A4+~*~a+4JlnkM**TWH zl!Sb(G)cqV`-JDb)N&3b?u5Fmd*G)*2_ zi@B{ujuyK`Hq=-yGHCALRs#cr=2~_o%RJBF?}e;Zl`%KKHL7{c&U6iP=3T&P`3s_$ zOLMwfC>doc6?BTJk~%lXiMHQAW?m1@{>{XKp`&?L2@;XVSA`N5$W@H3tkY;kI1!>+E8gov%LAWm#V`_v+67_8YRe% zs2_16q$z$Z7kWeKb^)@I4!I5}<2co>!9g(Sxk}2hW}yg-%Y3Re=?fmxko6}xeRtok zKzWUtWf09&!1auPR|K}NdctT}vBbGJ&XDA(T)3g8&qFd7ZdL8%*vNrJ$3hqXQ4nv> zJ&YP~xv;rCTPouds5=NBK?ZMl)UZAPb~JyCxMUpi7)FXwFY;+P@lN6AkMKV1#g)6? zl`Hha&RREYu$&A|O;H^VlhL|8F5^}kp`;rqU$6;bN}7-|gVRsXQ*;TnKS?4wxU`^> g5SK<{;fvV#qBTXwQ28~3ce|XK`4q1#NNp8A0X>A#xc~qF diff --git a/tests/fixtures/onnx_genai_workflows/static_cache/policies/termination.onnx b/tests/fixtures/onnx_genai_workflows/static_cache/policies/termination.onnx deleted file mode 100644 index 3183438a88f3ecce329e4afc84272274af35cdc8..0000000000000000000000000000000000000000 GIT binary patch literal 0 HcmV?d00001 literal 5880 zcmc&&-EQMV6ppiLlT6a4on^bhETwCeMT%$@J4xH@)v_yu6akf3aZ`jwt~1*@P3)}g zgf6f*6&D}`NU%4c6)Fz^kHTB<2HeA$@$cBFxN^yj=bRJIeE!aypP5PxFHh`abLQUv zYX`lQaD7f5*_P|+mPguhBN|)wkdn4i2^B3Fp}9VuQDR!2`{1AOCtXVA!3wHc`h>bu z-JrNI*hZVvvy)xhvd&$@F{j?HG1gs|oSM|p9pm&FE`IXqU<=8qf3Z1;YuHojvVWkT zCA93>XVh|$028mw$;X!abVljtl&JDrEO6MWQjIiuVf#TsA_`QdwoA;Ri>q_;?$en* zCYoIL2OO2AEO99>L_dno?4+I9NjP&2RU>CM3zt$cFHpZo0xyqG^(iG?x$X}*FI{Dc zOD%+MP;1yM0-RPIEy132ye90g+`zq%fCds^LdaloK-3{=;~J=g?bVZ4!y$*;Q_6vy z7iTF)?xh{Mcau2EguK^x}f=H4m_L)!Lr;x+dwp`}}?ozd_ zl$ey4OTXbR4YDcoZh!$-=Hwk~NK{p>v8NA4Rc?Tg7nW}&6s^kWc%-vz0Xd!zd<2?O zVp6M$ik{sp=a{>G=gb5cgx$0X(MN|fgFc>3h}x0YW61~EQL2%u)j*XqN~c3};^tZi zUPx@4hus8GyYgBj_^`THrL{ck)2|Xr@x?ZNgiP(S2I)PA>K=87+LIgM*vAaTAW-DB z>(adKWqA8v%E5QA{Y>CQ_e|RoDcj@0IfUTchv0;#Pt`sI=P4lXJR6?c8vr|rGY%P!Q?9p6SgwzCGnlMF#Nr|{wiiu}gmRDC$3W}+!)^sNyanLFochTA;*Z=+o3ukoiJ@$U*faIJB(`1736#NYBHopcVede5+EQKNrD6ykaQvaT|{^NThQGTY6V^F?$jQe z#xugQpaWrw|15r?IQEz-_qUWpZX#2M4ow421j^fB;P;dKFa8L0!9^fbhC|Rs4Jp&d z0$3(OyUNTdO)zJ0cM#4SieUA#3Nr*F(DAHiS|?09h0}{;f5reqhUE<&lU;6es7swW z{~Fd_FdcuwpA(sfsgojD!+OC?(HQjf$rEaLOq-xhd43@c(Dr}Q4p*pHVOPgH;2PcB dO-!*9G5fNe&7;#>9{6+eFJ z9f6*STIRrk<$5$&M!=LGsZ@<0$Sa93}x?dbfjp{^Dl?Sd8(XG$^SqE+~%6!U`m zt-v1R-c;ZTnPDd6l?0QBa03<9Lj2aJM1IDn`p(LXuuqL@;^)QqQ-S><7JKiJf#IKq8ulp`-M+8+uu z;+9i4z-$npn(uzbiK92)4Ja2(%ww8ohPqnF3+PgTJJ&!C-G-#;Gr;}w_`|aEA*Mx@LpshFb zPi}uVXe03c)^R`ps!UpfqN?-{)`GXpprwv7A40n2=pXFnA*36z0b?FR<}Gi8jQ6FO ztrRy+IKG=vnF^F7ePJxmOzBOOFpV$wrDsWniriI%!a_NL<5=)JBgzvg>v*oitU0z^ zE`0e1K*g`$7RpO#|2I<320BQh;v^k1oiT?Ie}3LwC=cLhYYsXbQcFa_|EdIq-PT_U CB^YA> diff --git a/tests/fixtures/onnx_genai_workflows/static_cache/policies/token_sampler.onnx b/tests/fixtures/onnx_genai_workflows/static_cache/policies/token_sampler.onnx deleted file mode 100644 index 1a5d5e311eae66f17403b7c4307e7609913be913..0000000000000000000000000000000000000000 GIT binary patch literal 0 HcmV?d00001 literal 58720 zcmd5_ON?a6RqgCwb!UGkJ>xFd_|r~X1`2zW@Bf~Kggs_347Nv_HUf=hrLwDXx@)F8 zvs{_g9)l2QkPYH%jAg)l+_L#vd<3#3SRf%5ED#b)#EKOXOZF^=6Y*aB?v2;wO+;32 zI^U<>yYJq(C+>|KC*snT(Zz?$y9dXs5C7Z0o*RCCcz%Dmn2)v|O+U3 zPcM(A)y`Aak!)$zcxX1XZXh8Ml!O#_J zbh(4`B;C1c@H~kJ=x9F5{c=E2+qophXqj!aEU3}i=F>J6A(Z@;>E;74H_ayxDhX{q z1#3PwggedWWb7a#=eE1?+&0a%+opx>wrLS92K&x{8BF7HTDaa@GVd)d@7*LWr|2@=fGtIrq0YFO z<=Z0~xH1N=9Rt_V4E)_E8Mt;0+%N+-mVq}4&uQtxbK}BuHev9bO*R0}*+ju}nh|); z5)RK<5*MD+jCgyND0t2=LZsv6ER7j2XBqY0tan`KYCP}FGVQ%tZhLPQZ)k|YbCzg$ z&T_|lGh`h)e36}hhO9%K*=X3Dq32Lw+{PstdJc8Qh37n?flFiH$}w=YIe4}WT)75D z@}bPxX%WeXI8lLma@!lLci1O9edvoL< z>dZ#rIcMWldGC1Ddm|dyj92B(cvT*6yoyE7wt?MYtDFs6<@T`EM&UWn9eB|OBSY87dmAvOsg)e zXrm}%MjHjo6^hR*l}bzi55qi;)!M1a^eJ#C5qB2Tx``>qA0D-Y&QKXYA3!Z z6F0snQgO(!ZCwYbfomp$ywKkm50g=Y^sjWYgOX(jrJWsY5|>L>)Gm9A+A$4mirS@9)Gp%{ zwORUS8`!OFSFE;OotoNq>DIO@QroV0ZM)h;ZM%~Bl?abZTHCHXwe2c$ZM&qk?MhF3 zs_fc!6}7fqvD$XkTif1_2hWwMZC6fhyV_uFyGk`YS8i>)Vzuq+wA8jMR@<(6Yun!+ z(ZHs*T{*SwYJ;_H+rVyZyJEHN>eSS>E4Q{?k=l06YuoiEYTLCQD5z*{yY|$!>&Uh3 ziq^JkUE8kh+IAhawq3K@cHLXsekPvx)~2>yJGJe4gSGA242somZM$Z*?fSITwre(C zRriip{a{1`oAIjJ8Lz72jaRYg**35{Y*n*itJ)s6+Gz2lc5B-;scqN1wq0+cwq5Jm zc1>&BjiPFvXL`gz=3oSu}RZq6u5FFvXKbr>u=+#);&$0W(hgFXv!2 zWz6D9!*YeDH&=*QJZVg>&^Wn56E9aV#gnGcSzzNP2@OjU8aqjdI!O1mffz$q)e*4a zkOk-?*|czK6XlcNmW%`~3Zxhf6%TBZ>|C>GX_F*c1?6vbM+5U(@^poWj#)Hm3lZ&6 zr$iK|H&jIVQtQU1j4HCstkx89Van+^*lndCcf zgF+P*x6^VSqm=SnrV*4wCYC7k30zJl7?9EcjiaCl&HT9Je%zZdge0RdKQ1jlu3SHk zmoxnRfE&*8XEkw2xkx*4$?<19Z`dqiBQel0*||Y)!8;@?g;}~WmIlk(M(iYB&HEpo zXz4`NC^}1nk}W9xwxyHJmB0T?vGQ~~?2;QTErD+EwG}r{LpNC4;rWB;V@$#$pgWPl zbBBWzyOYd=-AR-(_*VuZD(mE;7?8;PLtadfe;Dvi!j{4*cqavO>7=(99+Ij;T{Xa0GQl}`r>OyN zaKT+CK^jBiv2_swK|yp;y^aZ8#is5GoC z@jN5z%`-xlY*<^e@{G*QGqQMj28Q(&W@e>{MlzCMfV@qk5pcK7xc)}R0YTlSZ~+2$ z4_t&KzEBf`WU1i-BySro!0?U*7k}qBG3eb+#}1A+GciizDh7?rf1zOrXu(8U5{7{O z&B726z&9`e2rjU8sfvanpn~%-1Z?mKYo}T{WLD82AUBAB=B5zR8$?zLE=t^ji(5=Wj zFbkc6H)!b8rTcm4A(=GPr9nj3APFYAs+}6TsKCoQXK_)O@uG^Buoyi4K zI_*xtOz%nW&qu7`fS9gf8?1C2WI#!emEK!5oXSc9DpNH9E#0aa$Qj;_T9CF%I zfkJM^h&DPJ2nIPdF?i%Xa1pbH1Fokrapsyh-ol1f4aYHY=7yc1nllsIi{Nf{Dhb%h zPA~!vx*Lpuh0cNz@X$9f31$a=?nWaxp0UkVl2fV>pcflLHbq3yGuWtan8&iO< zSpgn2b`IX4v&RMRnib$dZ&y7I`0W;W2PU}F!67i*nM?N;;A1jyz;V~$4U)U#(wi*6 z+jMBq-CZZ)1^81#cuzUu9jE%hJ0b%IQfN!3F5w-A8w795(h$V4bjpeEIAwLyn~l5) z^BtnO;GKGy?>M&s@WxwpsgL@Olaax@p#^w|pujtIxbHZM8F{A@R)EK;K9`Qj zz=0H&OFISlbh1J2I}E4COCShwX;y$wC#R(VpHA427*6$BIwAuHQW#6S-1l^o+;_*) zg!`Uy?mJGM7Vt)+b0TkuM!?&{eNQ*SeaBlBcmw&!uu1CUzT;#J@Ww{nU2qM#`syXA z!+gh4Q0Qhn9H>H9mozsh1tH$#0yqjP6UhbAgm))NZ+CJb-arlnvp~k}PLg=tiK*eh zFs2_#bk!tH?5YWl;#!Y)5lL?saV4hptS$n~bt{ug!CH?sIAey?;p)bEdSCK(W^LN- zOK!0r#q^2$GJip`1jqV)Kaj}$ru7oS;5IO!z+(w502;4 zS05~A>u1iKIrGZn!64KQ?_`zhM36iQ22xD#Du_g+guPK7)L1>N?~klTx~M#ToEQipR+8r55G{rp=28`U4U zi}XR1;l&5bdk5>4NFu}%eV;gZeGVInT8KL83@K_M-KkL4WXj`xf#8@c4nN3Ujz|7u z^`#=9@-QFr%6k6raDFsfA0N#{K&B*t#^iQUpPaohO+hWcDt@$j1F8HRjzX#IkE8>1 z?;Q^JmW$PTn(R;IjH#botUh-Dg8Pjkr;iD zt}CeI1~LSdJRoGyMh4C}g?(qAani^x{r3;OJ76PY(lJtb>X6EbbZ5@|!Vd?O?i+*Q zg@*@=>7m>B!RL0Ip(C6>I5>ov1>s!$>e`$joZ~Q5bAyDUgX&)WB!?(H4pByMh`%vi z7B+S!9{ciq|9EfSx^uNNGKZb^s_|ugZuiFU((Cj2;r_wH)pp1}?Mrg8INNn)$^H|| zdoV9@1@WVk54i%4LKzp?Con`S0B^5E_umPk2Q$G-x3W*r(nC%opbjfI=ek>HwIVp@ z>J;s>5JjE|=&*n6cUWv7{uJ~!U#(s`df{`&vj@`@bo8g~6P@VE_=-mwaDjjy{GLh1GQ zsORSC&?~9fAQ;O)rr@JrQ7cj6V0`waqfoLvmaFo%FV#Uj5Rn9;qbrZ5KXre8gct@a z^y}@%*-Fv!i?SRSPS-PY=(} z-Z)rozoq_#G(T&XJTn2Qd;?(fb*XY~cPng21l-(FDKJRpH`PI>=VH9ooe9+*nAUi4 zqJJavkmme0843c`*HxAa-Qc3bxn7y8BZ)1f}ioyUN)rSQ^*wY-K^THd4vjw775Y2CNzgshc zYJRaj5lXynw5mmA%ZhZg3fyvLMUczGTMF)JQ zgPU6~UJ?nh`1g2lK9nkdQ;rcSE1IeSxa4*z3OsVVYOobBUPjd5*Q98TI%G&p?FH;| z{K1{Q*%}$=3lHWG=Zm#yljY7fJ=v+g9N!>UcCU2dL&$E1_J7#$oDPo&in(ePyU|jV zKbI;eDiIoV`@PalS;VI3Zps1H`sB$iB1T(-@rmF?!}` z1P!|#vl5lJ*W(wR?RfU+llm zf@o0S!w1nAb#@%U?=)h3=Fah5=n~}`@uPDOlxuJl$~V3$$BVl6lIXsYM9VfV6ufvV zx);3o=+Ql>?@n|N-h0I89{Z&GVp34vd!lA z_q0C*37+ZD&X8Dt7_@-=Eiz093aA(+G(o*PaZHlv-HBsj$M(<@$Ns|2kib`Vg>-Wy z5b8~qf_>VZ90`>AusIUeJ+a?TJMeaPgA5(T{_N(k?H<9R?`DADmS{6TXdC4W`qW|2 zlQtT7&eYeeX(65BTsH@=UKNQztG9|%Vi)(Y#VNzYOLa>F;LNl90QK@jY=J8f|MgncX@rja6*X;RQj-? z1Q}e|2g90i@FiGD>J5*>_~wTD+~AxBpi z^EcMhPGn-mz_ZLoJFEG8f4)COF=u{nd33Oc>6ok8!@~#jqum`*6nowp<9YFt>PXnN z#1zk7`93x2e)O*4vzG>=Tf?oZKz}0eZcJ~-lNxtyPyB$u$=*I9G?!K(= zer3Cx9wo1Ad9ECK|3*#{=T~kG&x?=c58E-O5F$7)ddKkMY;S$=XfFTkvMmFix4&HQ z3wNd1a7gy1yYCq~wq`fHw(UgZ!{H@JbAEXC(!jZbLa=oO;y;>xYPlA)3y5#a_?o;X zA%tIh*5b~JKfGpr+-q-_qiAYaIU8LZo`)!BUK{B@%H_FhIOpD5c7jKE=OvfeQg{u#Hn~42HsyIzjEv@`o`KVR{ z;%m+0u@)bFH+j6#h4px&40*R^CV~Vjgc0u`kEwl9HtyZ}F<@h=$81~57IR(OuN5yC zi?GKs1NOsZ-=aSh>_2wVl@JU(u19Rossy<$v~DCsF!*G}rL8JNgENXtTQvyr$@4so zz-~nN^vu%~s1CphVq$3>WF7^%I>?NZ7@0vv@#_FP53-k$=e6y`)j?+3h^vFl<0VuF z9yL)NU~WDWOX~o8v$2b=%-A4}w}agHjw^G1q61-zEzk3e1Ci&Q048V$*zSpa0@Hy< zPg)0`*l}Arn2PpQkcepZFK48Xj-aRs(2+UaCn(Gv$tDfh1Jc;N4MH2Q2R0k*qv3s-* zu-GuW=;{Ck4K6bl()G)XQTdae=kqPBn%MI^k8e>Oc+^C7(0KHub$}s8CbPMCqJpq4 zC09^m6l5Od(T%xyV&lPbgr4{GHJ-+!6&cW&=nbphsRur-M<5a;#L5&((cPABVhuOy zleVJKK_nB8PrZ*W8;r%v*n{3F0?0P(^r(qSZrYG7xoJdHa!)Iwk2DZ%JrXf+k21uX zNGB-+p)k9LDI@V{%qS!AsESCQ7@)Q#H%LvCfk2r_MIULf+J2+~Yw}10&Fqo1WgvWJ z_b_Gn0%@ZRUnET=w?b)4ax0dml6%~aKGOJ|{Yc|@^pPIVqss6F)Jz$^h}tN_7gAeS zJ*|p-R~KB12-*5BMAoXsTTk<7$+~)SgF<}mff{fx)w_reeStI)AK|p;kzEKaf-3q* zk4C6jd?7V63!-Y)kc~1h2a#Q+l6xBzV(UI-cp4OWB;sn1GO%Z7&m+Ztm?xnZ;%ase zQwBDIWH%XQ6dqMk$-Pis^g_8W*rt#4Lb(-gla}F&w`t4p1>8&-zKGi>!xwTB$*q{% za>=q6%6(BceWcg(thk#f!xwl{W%wd*rVL-`ZIt1Qy{)UBRz<$6i^2(dVQ!yl@z&Ek z_+(u@IX1wiD<1Dc?5*BKbm)t?iTH@SeUCJ$EPbR$Bh)Ou*qfOJTQ)qwVG)P3t5kAt zgTf{PJ<9Nf-sF)8z&*;qYW01O^dJ*e1_E&QB}^Gu0-fDtT85~)QjnyQd!Za<#XiYT zgmSSKHCKjM{n`IWxw0~qk|+a99P=-s$`DHmv(MGa5UXpltFAJ{lDj_1`=A`ljGp{R zpXXs&*FI%@4^cUm4CSw7%6MW_j^+FKBaEsLke{R~EGOwvl{JKIJ?&&IC0;FABiSyP zh&B9}fbY^7AQp#{m#^3b4Pt6U;zIs`6FZ8P%@~gTKmV?srw<_t@KKdTSEd652&v91S@!xk_e2!Xt$NRzZ z4v~WiLrPlOdN_m)Rr`vFmF6LlLz z^(Wl=L2Ftc!uu1<`wo|KWET^uN=)}4$J*5dRV<&u)q!%Zp@y~ngk4Let;!1dO$$fF zFPzQ!yY_0!+vQhx%tU6nvVr)T0g9$)hG$Uz#KZ|hVbo<)agBT_@&E;Ba`;e@^;4t( z>H_@}gMqiir)z@x=+f{aKH)6h@*d)yi2koa!Yw=#RbAX*+XAEP#ujH==?!68UlY>P zriTDPzmOvnL1J7p;~ciFN(gOmSU&%eaR>V*fq) z6l4bKQ{-C>8Osru{#Dz|JQ&-h1mcQ?u^5RLj`GB=00Z-0+`A>Pv!kI9nRIIr>;N>DT96!wKU1xR z{#o0|jE;3bj9J;@f#Qnvi>`$sDG-lSElfJInG~HOGZ4-^z~7$fCTTeRewP-sPld{5 zAC9cR{Ud3OZa%K_IwE~HA}bZOK zg;*f9`2a9%MkT_n+Ye)0d{OHkKTznwcwH0T( zUpObX|sDe&Vb>D@|N-j&Xhd>Z|j;wYYKw<^0^6gXf3uKhX};gZhL$KBC;5n(34|OxnK9asU_`vXX??c)*u104+^V#Ro!Qb5ee|UGf>Hq)$ diff --git a/tests/fixtures/onnx_genai_workflows/static_cache/policies/token_state_update.onnx b/tests/fixtures/onnx_genai_workflows/static_cache/policies/token_state_update.onnx deleted file mode 100644 index 7dfd0e7f3b7409692a79ae731043950283511932..0000000000000000000000000000000000000000 GIT binary patch literal 0 HcmV?d00001 literal 1229 zcmcIk&2G~`5U$%au?Hf;a-cP$3Rc9SqE-@aJ<%R>;D(Tb(8^g)6Q}8}+1;cqh)W)U z6VHOg)9@g$yLOt^6mg0#GxLpi=G*yM{{g9Ah*36I$3FtNgjy;%lE#9*6`FQKu!4&e zX*b$2`k)B`PcBg92}80q=)ggKK8^+FSBl9j*D;$VO3_J%GLdX@O}0OL9JC><{AFWM z#YB$E{sBu5>Iu`?0%4Wgk~rZh?S){4Bhia?jJ_!n-#we$G0wE2}&%g^*Y*h)`mzC zqQyCKsi6 zggk<#ACTbY^zdmpw^hgJ8Mw0C8*sXJr^^7op;WU>rr~-Dt04=Y8xx7K; zo=fkS2QP}Hi^s^#rsz!YI2W^wT~jl7TAE!TKfN!J6f=yD+tKYBN3YsZ*>vk#sRk{2 zV8K!|m@99_#}>^|QYkKNfQ%RGW^Lts5*5k?T4&+1XsH>@z6%y{(%5#|tnH++@$3wl yLmb!W>(V=c&VO%7vy|-H=i|s!BFSQR_^6=&roZ%FLi?XFij?jFIjzA~Ywr(oA9$t! diff --git a/tests/fixtures/onnx_genai_workflows/static_cache/policies/token_to_slot.onnx b/tests/fixtures/onnx_genai_workflows/static_cache/policies/token_to_slot.onnx deleted file mode 100644 index eb8e4b6599f5d6095617f9fe0d294eb1d098d216..0000000000000000000000000000000000000000 GIT binary patch literal 0 HcmV?d00001 literal 438 zcmaiwK}*9h7>3un+4PfQjh9(LrVIrSWxIOP9eekpNXW8mjq8%yq_9rB^zN_lhicoE zp@NuW-iMdxd(i>jU#Y3c_07i~yj##SYRM%4Q&Jl;o{O`5qznA(Y=X}G7R5MOS z7w+yS^1zjJ#dS&<$DL>hXX$bgDk;~RWg;~pOQ_Z)7CfUFiwkT&K1Blveo=J?(M+XW z*97(&8DbQv6|Mth`Of5HNf(^F^;z{hY@1>Pkfwhl0?k#NlR@@LYl8d bB4}3Jj!Wwj27h+&`J6M;Ofu<0$M1gu>uQY1 diff --git a/tests/fixtures/onnx_genai_workflows/tts/code_predictor/model.onnx b/tests/fixtures/onnx_genai_workflows/tts/code_predictor/model.onnx deleted file mode 100644 index 058e3992e363801b52d41f054b46728fc16e4450..0000000000000000000000000000000000000000 GIT binary patch literal 0 HcmV?d00001 literal 106338 zcmeHw?QbNRU*16feG&b} z&e3kadwkO09qf#cPTD6&!YWRk4A&>;Aq$#?2ei*e$s`F z-s+9_`X^ndbz6RBlNY>6pBbh#JL@;dk6Z7co5S8=e{|g2={H}Ot?#V;?9=|xfB*e^ zcd^mCd28RsKO6`p-&y-%|3UBcVEo?EaUU1>{PClQ?W5uFv(e7U;CS5LIp~c>-MvBo zq<6Bj_j&V`k8Z=|-u}+oZK1f^+gsm9Ev3M;!QIggu5Hv@`$*|>S^)3$?d_7`9u0=I zhugIsZub%m_g`N`Pxr>-{xC!CVQ=&ZJ-Kr<9F4nfcMt{5uW2~GdwO{9^swu-z9GYJ zll~z?thauH{J0g1h~1D)Oo-ijwG^4p8)B5m2rc!8yQ_=8T3Fn83SA#zULCJ4;GgDeSpJSSi@vB#{<@B-)BD}1^)(ey__rCJqV=VFr}yUzGylfOHNX&ngBD8cR50iV zgPndi*1?zkJ{!asgD?GQzJ~E16he&jRu~h{$k<64V<&UV*#CMFnbxYr_*cLAf;JW3 zI_dYu{gZc2-uU>mchK#$UQ}kK@<4{I&id2xujdP7I=>6ha+`lf8zi>8cU`BorVJ4O zV>ULfGq>Qv;efgC1{KI_Bs(-PFpq`>zB-VD?PWOTDp%=)HulVi2@eR)$yRNV8RLCF5!1?RfvmfVc=>BF=@R0WxoOF6jE=;;W znsoCQgZ~~p=>mJU4olfOoLjd3Tg%Gk4dL{al-N4xecnGAwVigh$LnT44v$U_+n@Fa z5BJ8+?`gyYTDf(b(^{Irc&LeW=!U4w7#PLW#cR8iT?3u zaBwty_K)Mg9R6xyVdJZ4x%b&%wECy{zYNR|#YDu1a!X8!2-eCpah95${3IPEwas_d zXw$YnDotE;jERdZrN%2tGi0=AoL+m!nA32zoCY?|M>#D>8EBku`(I0FQyz?VZqOd} z4<2-Jc4*lC7-v6^_H!7-Fbp)4_fLAm5l)R_B92>kBHF( zLjeTrh=Cf?N$xo|#6WB|MPl&qiP~8b18Kqq!H69(d{Uhls0k+q&$S^2Vx|&<=RR&? zAnlrOr;|RM9WfL@z>XNGVNML54KWZKnHccF3GumD6D=5mue3V$={< zUv=WERR>sgY&{82-w8aKRQF_Zv4R5xe7SY6-`gFnuB5jTzmC4j?rR*hm?7-#^=oJ88lN0mh%GavZy5N_Aax4PIAX+uiN@K}(T9@*q3?J6RP1dzzGI>Rmu1jPAZRv0!<-Q8PQiX3qjngYt zRB)GF!^o1QG!ynF>ql}QxWAxMfXtXbQNHz9#KMl;q#9V*gc{G&Vl%yb z_uEF`+%W_$!=*ns;*UdCg#j9(!_$gdS3;x8?s>ON_{ zWOn1dPY0vDuAj8NZSoy?iY(U&miwhfTfY!z@Rn>o$qdE?m%Up6UtcYIw?Nt+a^pIk zS?PS;-nwg?rQfufrQ7J=FCj&aE>rI5u|&9r06d*W`O^a3c3Po4nDRWL78==_E;LU2 z4m)l+AzRv*v!vRe8lC9aI1vMaHXIY}?@W(qYxugkX`Y`sJ}AisPRp|a2Yo!7u>rrd zc!h&ZrjjxOC0LGwKU_kVT#2gstl+IK`911(N&@L+^+12kNdC+wer7X$hPcyRU&23J zR{AZAcz2IJ9ZCdT67llqIRyuXYk;XQn2k!pfhm7Av-pNJ=PYS>IuS2xIws=(u~cw2 zlg|NRo|MJ2tjxCP(Xx1!Tjf;io8wt#AjPu*PfK0-TaRbipkaUi#U)gp5GjdxS$S>d zwX&pCV3U*zE=f}AA1#;pAP69}r(0R=sat3%s33&Yp028#1Ts&-M46{aBP4gu)t+t| z8&y$zy4BU51ttPhN6_(+TdMX9ObMq+)t;(nnfh(M)t;HdSzl7`Td6&*rZRkzR009` zB&onKXZ{#UQx*V6+8^;0)D=c>=uF>MnL>d^Uj~f8qAt->^yWW3s{l2?1ttP9BsRoA zMSI+g?3wU;nAoqTEd#eAhWS(S=}sqoz>%9t20le%pr~Pv>jsG(F;K&t7&Jp=VGr>;v{*@xg&0V?=G*C{ z4`)XV1rV?!25Lx&A#`nsf!IjI5aKuek}@WKKZModL7H$uFk(^*iU9z9bxcm;v2fep zp!O6?$n~3{2i^7;Bg6dZEMN$|>D~6Xaq;<8dny!fR(p!0uP(w9_zzU;!v340TK9ta z&!o!}{g=;&(1Ub&;?q=bRi?`;Fws&#VF2m!M0b-4fgDg6JWjQz4PfR{?Wu4w{u3gQ=E!jRHsj>Vg_dAXyeol!Bd90IFaCY&_4Dol^m5mPmnok3j`r zhj+6p(@avBzs6!9!2ad{R$C6^0WXECIKku z6IcLBOo+5IiXj0g@tq(5jneNuPDKC;&EWGW0A;h(BmgBNnoIWRz6ult%1<+%>_!PB z042VlI}5@^#!2JiI{4+@XsUwtw^*rd<_dgWqjl$+kph2 zGB~8o6TioJg@aT8T8sme0958mSO5z63>JVAgH+rZJKA^V^Xu2L05o={BLJOK`Ynri zxg?Ada7o0=pXU@D8*T<->VpNKu_>iBv-pN};Vfymz5rC5eyr&@;x!3C#d%T|&$2QF z0ch-&#k1Usa%b$B<5^}P#j`QbR9*R7k7wDSVSk$hpe!K*3P1@WSb4bsH1^7pQn5`^ zD!wF1DU$${_#pNn0jO750IFMPDX2Jr1fZU!05qN`^E7FM*YXPX$RE7_iN+19qE)^K&>oxeEx&q*c z`=gkrpsp}dC;&BJ1U_mOJw?R;Wk<~gCIT^l-vJi=OLctQEIU7pXZamqD`J>Gai8vV z(gz&5nFOFyBnFBa?#no~BL-@i69f1i-~!-C#4x?@0AD_g*bxK03sr>}sEL3I@H@bw ze~FozHH}TL%{!CTfV7M4+~+{`s}lq1!`Tr-0R-%bff~|Do`By0E&z^53<-aw-W5hH zBL;>MlK@l<0AvAZ90>vdg8)=4A-~d}xR8`X^mG1n7BD34^rRd%+<$%r0LpAP3qZxu zR~O+4{09nnVgF50z$e>;zCkA@o7r|Xn~2A0!n;H04loMC;(0T$0-1{ z0nA(qKou^q6@ZFPXQ%>Y#p2Uc4o1U6@XL<{COWD>S+R&JQXz1sEtx2@IH>^C3RR^2 z{94;Q2tajIkzU|<(j)*Coel{=6JJ<$46;G7EQZjqdejhJMGfT4NAan>C+UnYn7$Ku z5|r|aQ$P1RiT|c3<+Yi(a}@uX1fa$9=xhX_Qr(J_Gl?MqsDuE;iXb6Amjcks4yx!J zfC^L z`=X6U=THEu1lD*9PR15*mP`Urc;?%k(HHpLF5^xOY4*GWUaTxYi{J2|ztA0PT3w5r8U=Jy!xy56@4N05l&iNC4{97l0Z} zwOw|{hXkM=sG$UsWu1vqu#*Zv6)b>_=b5r|Dge#+{uucBD+U#O9p0-b09Bvo2ke~y z5`cP&vX9$jAz&w%p#W5MG1DY5p&>g|KLb@#{p;|VhYRF>CIM)*Ic?gN;ED@HznuWS zFs~vdU))3J!n_I-X=M+mCIG$Yh%gC2L7%_^P+~%)ot+30fD+#c0?I03{=8D*%`E`&v6A+qn{eGACLKKr0e0C;)|h4GTbJ zo1A15NB}B>L!wxIkMRlzsQ|PX2POe1@kP#+umBYB87u%L2B`qF(}4t_#Ix)Rpq$v8%m{~|7P;55)pq5oV-!3;Lr&v~`hTobmz|%f` z_n7N=trwvw&R$N2pSm2$sXtmqwxzFeM>1QE@F~b2d&>Mo}+%L7}Rx4*{4bG?W?|J3d;FbgJv#g zeUa3lSSmi@bVjvN&;-}}0{qaekL!d`3|9aialeEu+46mblR|xm0Vj6!K-B>>M7IKl zCI@t#$c7$>ogS#q@@v6X^gtTM!nZk)e(7!}eZZmnSS7W3g@4y|pq#oGDCUo6C)?O&dQuOY zMd17j1eBOy8`-*0TaPDkp$P6DNe!gtz!fFe^(ILVO)BWK!OUFBL={F5Wun!A1*(rl98%R74HvC|1qsbu*S7%{ zF$aM)QQq-PfJHjM@4wB1Qd9>P=>?82|J@R@ayz6H#atjPKT*{PBB-(a1XmFR5sNBr zwRfgr4o~9Du!b|IsQk5A7v?A&{LhtRiD0VQzSt(rcu$djx^4u$hEm!SC}wnT5S>fO z=%vS1cph*CzQxc2S7lP<$DbKu<+cu$WF z7^EWgOI4(A!4Uu?LTu}~S2LAhC|7$}wIO_ZaZR7a}7VSaU_8TCVBk8uUxs`ra^ zq=wlJX$oD>htE{85&R`lY&pYQ^{SJZxM|Lmq?wO#&Vb1X0O0lT?3>F1X&Ix7Ck5(2N)!(M6LhF6Fb07N3?dSPn=S^@B5 zn}VJcKRYc!Dde*AC@5uO5ho#QIu3&WWC_{Lm9h|*6)JpQ)b%=LAucz| zN!NGGAuh9!ub28flXc~D;Q|4kJL$%UjT-j(?-=q{JXul_1+xN~DA;#OlUBY}(#pRi zNh_1Ml={JUA#rI?SzM}H=qW1SgFZ45*osU2i851@hD^bni%Wy5;?khHxU|4TScrWe zI(`bu#ihO}AvNi>f7Q24kU8Jt(#++oFV^BxyXg#{M3tccK8Y%@WQ%tIr7{b^BW;^} zo}#+KNufN%fD`z1TJ#lF7L=V%7nq3j5ZcfK)dzAra)ZI2<(G)9=z%nh#q4t+F4f&m z`hWvDlel!6^gw~bz3KX)9X(J>Mh}q`f!HyP=o?IdZ3n!9ugaRAVvy3@OSxL;lwg}U^p=;6UAUa z{`9Ep2ZEZwATAY)$kn2*-+|NwVxX8mPNRp;^rRs+i@^C+6DTo(SzIcPzqTGv;KBgZ z2g5E5rl=3TU@kOi07Vz(a{^BxNL(sDE`Na(6oVC*=-Hvbg*1So$BmyLbp!Wtic4)U zGne90g%NPZm$(6k0hNIjI5xndDzh?isR%4km==JlWf54!9J&z~Oq6$=R9tEWEYgAb z5|`@0BE7)CGKot?w}ax+z!jDsgSb>IjL~118cxD=R7XCnfXDpypZZV*D^QV9f#83jXlF2$vn9#>~1S~lYf z6fIles!X)}c*La_3|A&`soAw5EG{)0M(~9e(NN1!RTh^9#p2SSOk66qf~h{cZV-#& z(yVYUE~SM+;!+Xryhuu*IQPnacizRN7YzP(`FW}ltCF9mCPIEb0AKznf`dv0xgEvG zEMNX;dMC z4wJZ)Uk(Fr+KTwm917pl3y8QhbRcnQAc#vtXFB3i#kuE7TpHkcY7&>S0fWS)L49$l z!Bo%9hAt#74M2G&lu8RH%F#|LE>++Fww~wF&gn-=GwQD;E)6P*OV#K3!tI70e5RV- zt7nQx?KD30W+*OIoy=rTtUc@y{#@Uy&wNbpShJ>M&M=8fE6!`vMg`Yg_?=ShcKGJK zik5t(4}Iw7eP9xohW^yVr57C)CUGh36P2$q3j9bUuu((vl z2xHp_5|_#Vk+xJJzXo}Qg;ZQxj0KapRN+imTnf1j7MBu>Ok5hqkhqk1my1iocsk&z$>8@9Y@G{m3h%<6`Qi-p+8ATA9}DX&R|hLFp6d(M)O-!&4lXGzFx z{4ph)#HG?)DGPB~p`y4nOv*xBZpKKrV-9hdg?zm<cV@Z6~)F6GG*P+Y1` zdhWQ=q*Y{VTSyti|BcI=-Q)kZ|6~e zw|j8d-Rt*u5gLvA$KAnjxBpr5>zYg6>W%mMCtWx4T1r8i^iO145qax3$S*dYM9aO; z2BXz$zgk$_ipAknUbkZ!&aGEV2k?300B&z#^eHY7Hvmm-)V9$-UPPu5Xg;qEj}QYy ze(ME!sI`n&WKyrr)kyu5MO0z>qLP7Y-Z(gVI2eyKh)H{sh$}efYOH>D5xpb>gU!M8 z*E`+*;r;&Z?qK+^xuX1H)LN4*H-EI-AC3p(&%1Qmh3k#{S93Mjnor9e+ZZ;M8aEn^ zm4(+W2W$W9sBv(#(;E+thRtUVQx48H`|aJ{xVQf_x^eX2!KgoOE{Cq??B7B+4*J7~ zHeP=d6tC+b9cJqL|17aFhPZvBSb;r2;(&n$I*d*^7>-RbS@ z^{wlQyx7}+hV|*X@5NDqdv$10+h{PfsSTG-NASAeaRR5Py05J~@%ocHN5j#$>vjiG zu+><&M*eNI8VmT3i2f0O_}%a~Z>zDG+dYSe>K3tb+D`X@?`hfLB{=t$n44t5lVeu~)78(%hT>!vy!o4N%*$8O?CRp+q6-L|fp z-e|I^TIe?ZY&f2mVz=|Uy~eTxq)>O$pVEI{Qpat4jHdyHiZH6qH}NwvzSCKS*06|C24G;tyLAf`Pg9SHZXYYMunc}+cZ(|VGxCm z7l=O0q1(-H8$)vjI>V&z`_dh*c5>sh#+R#E9p8cKNOL#r0F4%G_V-uV;6H&sO&%nH zG&MSj?*!QIpX4JDd+deyzO*zU z^O}v^>duZ^UFc0T=HV0o8v0HcNefvH&@zG(>#o49Lfx((`++n$a)PT;cTIAQ9sHp+ z5k-03V9QN zd>IMQExW=q1p;&nZ+w{jC=#F_bYgkxTkv)Lv7(hOa#XCqNuv4M}QsAkLBeg2WS}y ziFMZ_K;Op@lEyP9p(=IPB)|x--qK>0*KMXEvDftoFi5wA8Wq{sT8{vO2xoG&TFVGf z>}@>)>?9puPBrBCG7_NS0;TojDa!$ps~@#2;|*;P6uTKNNTPb0eZfV!!QZD zdA(%>s9tudNcmpiyErdjMES30$yUSdj&id5@Bz9t7#^RFN8|u>)cgS5Yb-S1L2FHX zBzJgxG{i@p>JiTU-gsvZ-EMvpd}7r9__ROV>38vYA6jfQKK!b5#o-28ZamSv_TeFV zx;GxDhp9)y?qP5AsQD54CH~6?7S+2=e6j#E@VJK$sBjz7>epUgqyO*29(oFAxd&No zgWXZ{XXqVV*^ezMv+e7LpP^^7${+PV?|#xdIPH(xPO1%WcF{H-*I!!Hg@+4$r-?J! zC7t5~yo3{oAAW(J(`(L}F#Ea-dhI^8Ub|FuKl2eU=(YFQdhI>NUX!n1&wRWK5T1X% zd**c)_?oFuylgDt|K3QK zkNY2^TgNB;QMwipI{gdi&*O(aw4~ElnlFQNI>uK(4${R!=F|Hl^bD_(@aXI-75YgC zboNyW{mhwj_LU0#B$PS(DusULTsiwng?&CGha2^vHNGb-Lpu%P1dBb`u?i$ z=0CTIPy>g~%ebNi=`Wu2m*9kQ$q3t8(L`>#Xq zId!0uJaI@?(&fjl3@INYmeH^;L&~e@Wr#iqg^^y~|Gr%>g?m=bC(%s|nClr}?Chm1 zh-@#mED~C_WH)wFo`dx}-3ehhuk618y?DbZJG7FZ@B?JLr=w;iSUYeVRx=-d5BgZG zlO1fy$pziKq<%F?$_k+4u!`wN-HKDSEOg^X>OMFV#7h@aMw7K5^@172Vi|R3f>=h$ zOaSpNmLcT|WxQJ18}M#5mUU5GdKX1n5U4OHf^OYeY9Mseb}gL)SQ@?t`8X zPtxn6S}db3s>L$uqFOAY4yr}_I;a-y>!P~!f`Q~#5LLx@2y|G2TwQb(KsPU`e*#oP zXY*3`K~#&!qjga&mQfegVi|Q&EtXLS)uMeJREzd?QC)g!EV&g#Rq?^F4oeW#Md!P^ zc}e{fpc*=GmAVh2T0HZri)yiqx~LY*sEcZ`j5??m?dzagw6BZm(xX7htstt3Py2LO zf~YP!wA0N?>Muw2w~QL;bg!H5CL_C7s5hTRPhpz|x`zKlHsdd9x>Z4ZMpV09SF^i~ zzNXpTJw3d4df4S#ZCuHI8+}82@Ppw9XEpm@^t*mm@Do{s@Rl7v{}Ni&p1-%(JMMRb z%(waG9)3>TfG+I6h<;r!@AQ6G+zc)h*LpC(2Zb@KwY)cp(3^C47xrEBf>!%mC;i^I zfAY@B8y}zI3~vXglu_d*&MKrnxwgM)e#x^V)ImmUqhQ4DrU#8KaobH{bZ?{QwD!C< z!~#4|X~nk-{QCZ9p_Wvyuz>U+ie)-!Z5>`;^@GOsD*(kbt2#xBn zvS_0k`<9KOWzj}?4l9T@%6g>ea>N_uxvZ?JjiP1cjq+GQtWj^Ew#|^|AM5b^=C7fz zU`*Xg(S)bu3P~$fEhbQk$?`u7`yZh1+q8zRwi-OWP$`zS&azm#2m7 zj(7&tszqp6Z|wgFMK)vhp7Ny)k3*fcre;xo0odO_;9FrtT3|1!@A@u&6@H%EW-YRu z$Ta-nbH?XTXDO{W?bwBuMdyr9R#d*kH}*ZW1`qgaySrUK&{*=f(VMUhsL0fih`-b; z2OK5zBgkv1pKPvssTb_F>0Jq z9w3OtKmkIi#vSDWf>&Eh-V9+AwVcBLc_W-K#+=10tEab(;X8pe#C(BTrkB4FQ6ieUc&4soD@A7*2^a z{*81(H|$^3gCo8qMIm!OTUZSNf?7~MHX5Lf0Rkz71PGA^WMhCJpI{07;2>uGFa!v4 zF?laE^cn*MQbrme#GIFH0t8~n*LB{DSORNlfS?vc`w`bbYOEvG5FqI1cp5-s9iWB) zK|GG4>0)72G>jSpgo3j2fze=U3=qVEv;ZOIxV2e}3(CqH#gS_hAczHN0Yc0XT(wbh zfRJd2HUgignn=k+XM)5F?laE^cn*MQbrmebT}`w0O7ab*Jl&Mq5#2lB(#zU;fHWOMJAn^uU`d?{)Rle17#9@A z21f(y%D!c@v@pSSIDFZXY}qXNq``GKf?1$o*(@zka2*cis?Cxk1=rEA-PnH5R4zx2Iq35HbMyl!RVcOx3KsBdzKYvtSuPYUxcEp=PNLBE8g?#Xd-?cr z?6VcRav2)xjfpC$9ZJC9x*YVZJ%;GtTvh?>-$dZIUc|C^$bjF|Q4}*+P9H61a9s|1 z7BsvC6S~+a`Pc!!MOPzL0~6D4NaB-1^?2Zy=q(uH_6HBHr<1&quH*v)zTUd{1Tjx5 zu?5!<>Z$joax(TbXs_3{UA1BJ?mPZd_d)fE}9Q$mwHbf7_WoW23Mh{XMYV?557IV#1*3lu^c^m5Dph@qKAUQ^63+FGzv5@G5eB_W@LyS3^zx0O2!9^ zDkU3MD)lz0u;k+7#T)}{$u&d|YDxL{gc`Js(E}-lMi2N5F>iuR^dR4C6}!UW%(`NT z9^`UZM>HH8qX$w*7CnTVt64WneZtp#u_K(TZK4O#23`{DN*&BbddTtJDj!SHVEKyq vqck~+>H7-yozkzZaEnwPdgJwFbZuqz^{=FY*~9kXVCUp$bo5}nvhe=_4Qt}9 diff --git a/tests/fixtures/onnx_genai_workflows/tts/code_predictor/model.onnx.data b/tests/fixtures/onnx_genai_workflows/tts/code_predictor/model.onnx.data deleted file mode 100644 index 7d72fe4f2011031b6b0cc022f53e07a512cab7ae..0000000000000000000000000000000000000000 GIT binary patch literal 0 HcmV?d00001 literal 2529408 zcmeIufdBvi0K=g9Qy=7oP+`D;0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEj zFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r z3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@ z0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VK zfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5 zV8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM z7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b* z1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd z0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwA zz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEj zFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r z3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@ z0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VK zfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5 zV8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM z7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b* z1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd z0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwA zz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEj zFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r z3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@ z0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VK zfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5 zV8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM z7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b* z1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd z0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwA zz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEj zFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r z3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@ z0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VK zfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5 zV8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM z7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b* z1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd z0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwA zz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEj zFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r z3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@ z0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VK zfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5 zV8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM z7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b* z1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd z0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwA zz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEj zFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r z3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@ z0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VK zfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5 zV8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM z7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b* z1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd z0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwA zz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEj zFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r z3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@ z0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VK zfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5 zV8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM z7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b* z1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd z0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwA zz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEj zFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r z3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@ z0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VK zfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5 zV8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM z7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b* z1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd z0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwA zz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEj zFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r z3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@ z0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VK zfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5 zV8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM z7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b* z1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd z0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwA zz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEj zFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r z3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@ z0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VK zfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5 zV8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM z7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b* z1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd z0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwA zz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEj zFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r z3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@ z0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VK zfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5 zV8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM z7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b* z1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd z0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwA zz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEj zFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r z3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@ z0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VK zfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5 zV8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM z7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b* z1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd z0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwA zz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEj zFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r z3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@ z0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VK zfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5 zV8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM z7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b* z1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd z0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwA zz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEj zFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r z3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@ z0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VK zfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5 zV8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM z7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b* z1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd z0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwA zz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEj zFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r z3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@ z0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VK zfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5 zV8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM z7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b* z1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd z0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwA zz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEj zFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r z3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@ z0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VK zfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5 zV8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM z7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b* z1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd z0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwA zz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEj zFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r z3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@ z0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VK zfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5 zV8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM z7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b* z1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd z0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwA zz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEj zFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r z3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@ z0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VK zfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5 zV8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM z7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b* z1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd z0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwA zz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEj zFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r z3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@ z0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VK zfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5 zV8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM z7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b* z1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd z0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwA zz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEj zFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r z3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@ z0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VK zfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5 zV8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM z7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b* z1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd z0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwA zz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEj zFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r z3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@ z0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VK zfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5 zV8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM z7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b* z1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd z0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwA zz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEj zFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r z3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@ z0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VK zfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5 zV8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM z7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b* z1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd z0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwA zz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEj zFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r z3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@ z0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VK zfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5 zV8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM z7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b* z1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd z0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwA zz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEj zFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r z3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@ z0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VK zfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5 bV8DO@0|pEjFkrxd0RsjM7%*VKz{|h@nz{f1 diff --git a/tests/fixtures/onnx_genai_workflows/tts/code_predictor_indices/model.onnx b/tests/fixtures/onnx_genai_workflows/tts/code_predictor_indices/model.onnx deleted file mode 100644 index 22966b7456015175624d153fc5249fdb8b5cde0e..0000000000000000000000000000000000000000 GIT binary patch literal 0 HcmV?d00001 literal 942 zcmcJN!A=4(5QbSqSOy7LmEazXA)Y*JmerFt@g$F+G0oBzsIu**Ed-7{fF67TAJtuD z2{hrjx5?k>}Yi(;GDZ@d}&4b%E z)ECig&M+6I>-;$7VH4^iikVI#!cb{=16Rq%#Fs)WHBmeUM(jE|Mv3SHNws6(;~P86e@ z7^yN2PMCf<+qo`Og^F84rO_qW#M`HH*g3jwdI7tZ^Z31E-Row)m4|u%paP|KcW??# PB4a=1L`f~jrtSO!P1z~A diff --git a/tests/fixtures/onnx_genai_workflows/tts/code_predictor_indices/model.onnx.data b/tests/fixtures/onnx_genai_workflows/tts/code_predictor_indices/model.onnx.data deleted file mode 100644 index e69de29bb..000000000 diff --git a/tests/fixtures/onnx_genai_workflows/tts/code_predictor_prefill/model.onnx b/tests/fixtures/onnx_genai_workflows/tts/code_predictor_prefill/model.onnx deleted file mode 100644 index c6809bd50f5181889c04c5e26595dd5746b95f0b..0000000000000000000000000000000000000000 GIT binary patch literal 0 HcmV?d00001 literal 465 zcmbVIO-lnY5N+04)+vFSgOojph=_Oys|OK}_2m8pMM8Eq?S}3qAxWYB8xJ1*nQrK* z#ZwOh^We?Wp_kQ{RADN|Irg~zigXt)Y!)X7-gJ&0fu2SR-vmU0Yke3=3KG#gh z6h{Pg}SgA%SV5u25Fs;z_vpU*R6B za4`T^QX3YhETT;z^P$mWiFM$gN_{g#yhrvX;`Z9i=gA*QN0&>JSzcf&xKRbu2J?5z zvDgfV>o?jJKK$5)9nzo?gmjCmELRexHT;M2 zVJ|Vf-(C9NcOKMjDyMe>lM>QHIIsYFzpY_G-4b@tn(zmJ*fRApI#{l@7@dbXTZh=l l+b0J>B^p}+z(zfA%AD;wIDV_W2u=IzFcZ8`T3s05eFglO2>Jj3 diff --git a/tests/fixtures/onnx_genai_workflows/tts/code_predictor_step_embedder/model.onnx.data b/tests/fixtures/onnx_genai_workflows/tts/code_predictor_step_embedder/model.onnx.data deleted file mode 100644 index e69de29bb..000000000 diff --git a/tests/fixtures/onnx_genai_workflows/tts/codec/model.onnx b/tests/fixtures/onnx_genai_workflows/tts/codec/model.onnx deleted file mode 100644 index d11c15e787fc2a924e74490253c38f0c82d671e6..0000000000000000000000000000000000000000 GIT binary patch literal 0 HcmV?d00001 literal 1040 zcmc(ePfNov7{;5@vGm>i7f?q9k>P2Fv37!>2;x9+S8s}hEM~5{C8eoVKY$mnegThu zOw+2<7Pi~;mM8DalV9E^!L@^D46lgPQ`%EfCrl=^2BQe6w$Y2xpV)*3*{#`l7T zGSSD@8>;q~z?Ezwbi}wIR(}_^qr027lJZgWIE<1uA2Y4#C=@Y^`RIw*mskBA@Jd~o z2i06fLT4RiD!)wS&!j?sM-`N*f`6fUM^ND=Dh-7o3(7C3fVe3=ABSAfu2=tD*K{(`a=}hi!=!012!`ZC;}@r2(<4W4h^u5D zHpx0ylSoM+&GR*bul7zBYM1!RArV2R(2n zZqNYLq)ByTV%Z(6)hdPPRlN^IOCMs-(G85gqRy@IlM`&!_xGQEJFah!I1Z+Qw#xLg z`T#z2UK~lbWxetY*K~Zz80em-^-S#QuF?AuiObE)g2I~LIpKyo?$!K zL)V^*w9ddYIZyEr{>;O5Y7xQexON{KzG+)p*N~+MTeR^OqEmMJu5Njb?O{z8BpNaT zA!%Bjn1%?;l_DKkUxs-HyTd!&o!lF4QB~uZ@6fR6m?zV<%8yJ7>#lk?6ZPw9s7Fx< z8 z=7AYGs8Hl)u9mi`duTPQ58re;T3uM9p=LHAFbX?RugsC(A1(ZUali%`x){gx>51v} zv`yjC1jT45#B8&&5WcBWoIpOOwx&6UWIV)-+7qzPs-xj1k71nt4JdbZKiXrQjK$M^F4Cy*~I?$Os<_E zm(I0{6msp-MACTajH*;Aqi($~Mm?J#|C)y~Wje_qV$_Jq;gMJ_M&p56J{L3bgo%8q z#>NBLozTVGR2M&pQ|xV;_%w>0j{1Rw@z<$TT#i#Q5yn$96%SF7PGo(^|MRBE^CfOK z>j=`n3jDw}^gunMm0@rU2vxq+(S5yNf`Z-cde}z=xh~hXl>P-M46t?N_t4_zcB9^C zZaduiFL|CkBpZ`5S6}CIpZK$`hHx=RGqs!r6(joNz4?Fib1djbWX*tnm;a*_HZzmJ z(wEv2U20sdzd~2pLi3ia)Ew^MS?sdl5vcqkx(ip)6~i7nwuLQ!a)OoGx^MJAM60tt zhP#i4Umj!2z#3W02L}*i);)MgLf=ady@qbg3LQ~WJa$`y_;!cBhsQjJ)b!}Ag#%%o zwhoEC06l^4NQQ^$8JO!2QuOPY_c*8o=Sfz8`G9vrrzYbq#d$bKJo&&U^sA+72szRa z06g)UXUQ~P_vjrnjnqm!Q$2f=Mk}yPDl5VqSr%b=4W?@YT3rruq1#TDxRqmP2${N% zVQjGtZj9*IadoCje=C0G)RMY=oP~L66FoYedUqajrOUe)V3{mmX=oa* K?b%(w#Qg`Yj8Av~ diff --git a/tests/fixtures/onnx_genai_workflows/tts/embedding/model.onnx.data b/tests/fixtures/onnx_genai_workflows/tts/embedding/model.onnx.data deleted file mode 100644 index 885cdd9a062995e9b753010d104a5fb72e2baa58..0000000000000000000000000000000000000000 GIT binary patch literal 0 HcmV?d00001 literal 4984640 zcmeFtfdBvi0Dz$VsTV1P3IhfV7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd z0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwA zz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEj zFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r z3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@ z0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VK zfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5 zV8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM z7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b* z1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd z0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwA zz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEj zFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r z3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@ z0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VK zfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5 zV8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM z7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b* z1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd z0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwA zz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEj zFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r z3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@ z0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VK zfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5 zV8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM z7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b* z1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd z0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwA zz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEj zFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r z3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@ z0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VK zfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5 zV8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM z7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b* z1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd z0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwA zz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEj zFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r z3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@ z0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VK zfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5 zV8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM z7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b* z1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd z0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwA zz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEj zFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r z3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@ z0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VK zfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5 zV8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM z7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b* z1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd z0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwA zz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEj zFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r z3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@ z0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VK zfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5 zV8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM z7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b* z1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd z0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwA zz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEj zFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r z3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@ z0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VK zfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5 zV8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM z7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b* z1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd z0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwA zz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEj zFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r z3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@ z0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VK zfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5 zV8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM z7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b* z1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd z0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwA zz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEj zFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r z3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@ z0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VK zfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5 zV8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM z7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b* z1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd z0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwA zz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEj zFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r z3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@ z0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VK zfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5 zV8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM z7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b* z1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd z0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwA zz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEj zFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r z3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@ z0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VK zfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5 zV8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM z7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b* z1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd z0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwA zz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEj zFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r z3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@ z0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VK zfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5 zV8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM z7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b* z1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd z0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwA zz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEj zFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r z3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@ z0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VK zfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5 zV8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM z7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b* z1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd z0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwA zz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEj zFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r z3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@ z0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VK zfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5 zV8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM z7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b* z1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd z0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwA zz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEj zFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r z3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@ z0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VK zfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5 zV8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM z7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b* z1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd z0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwA zz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEj zFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r z3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@ z0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VK zfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5 zV8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM z7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b* z1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd z0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwA zz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEj zFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r z3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@ z0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VK zfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5 zV8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM z7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b* z1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd z0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwA zz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEj zFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r z3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@ z0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VK zfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5 zV8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM z7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b* z1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd z0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwA zz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEj zFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r z3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@ z0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VK zfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5 zV8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM z7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b* z1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd z0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwA zz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEj zFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r z3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@ z0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VK zfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5 zV8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM z7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b* z1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd z0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwA zz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEj zFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r z3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@ z0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VK zfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5 zV8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM z7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b* z1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd z0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwA zz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEj zFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r z3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@ z0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VK zfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5 zV8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM z7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b* z1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd z0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwA zz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEj zFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r z3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@ z0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VK zfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5 zV8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM z7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b* z1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd z0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwA zz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEj zFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r z3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@ z0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VK zfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5 zV8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM z7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b* z1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd z0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwA zz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEj zFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r z3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@ z0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VK zfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5 zV8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM z7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b* z1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd z0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwA zz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEj zFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r z3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@ z0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VK zfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5 zV8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM z7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b* z1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd z0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwA zz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEj zFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r z3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@ z0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VK zfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5 zV8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM z7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b* z1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd z0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwA zz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEj zFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r z3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@ z0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VK zfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5 zV8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM z7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b* z1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd z0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwA zz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEj zFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r z3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@ z0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VK zfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5 zV8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM z7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b* z1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd z0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwA zz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEj zFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r z3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@ z0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VK zfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5 zV8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM z7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b* z1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd z0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwA zz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEj zFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r z3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@ z0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VK zfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5 zV8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM z7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b* z1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd z0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwA zz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEj zFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r z3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@ z0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VK zfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5 zV8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM z7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b* z1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd z0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwA zz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEj zFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r z3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@ z0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VK zfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5 zV8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM z7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b* z1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd z0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwA zz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEj zFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r z3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@ z0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VK zfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5 zV8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM z7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b* z1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd z0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwA zz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEj zFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r z3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@ z0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VK zfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5 zV8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM z7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b* z1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd z0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwA zz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEj zFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r z3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@ z0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VK zfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5 zV8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM z7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b* z1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd z0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwA zz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEj zFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r z3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@ z0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VK zfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5 zV8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM z7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b* z1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd z0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwA zz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEj zFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r z3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@ z0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VK zfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5 zV8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM z7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b* z1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd z0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwA zz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEj zFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjxE(qc00000802p~iJgUz z0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwA zz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEj zFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r z3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@ z0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VK zfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5 zV8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM z7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b* z1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd z0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwA zz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEj zFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r z3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@ z0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VK zfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5 zV8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM z7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b* z1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd z0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwA zz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEj zFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r z3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@ z0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VK zfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5 zV8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM z7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b* z1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd z0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwA zz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEj zFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r z3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@ z0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VK zfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5 zV8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM z7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b* z1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd z0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwA zz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEj zFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r z3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@ z0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VK zfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5 zV8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM z7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b* z1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd z0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwA zz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEj zFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r z3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@ z0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VK zfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5 zV8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM z7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b* z1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd z0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwA zz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEj zFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r z3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@ z0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VK zfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5 zV8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM z7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b* z1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd k0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|wrK028zT0RR91 diff --git a/tests/fixtures/onnx_genai_workflows/tts/policies/cache_length_update.onnx b/tests/fixtures/onnx_genai_workflows/tts/policies/cache_length_update.onnx deleted file mode 100644 index 72ad12c9f4e344e870cd12220f36d4f2593cc2b3..0000000000000000000000000000000000000000 GIT binary patch literal 0 HcmV?d00001 literal 386 zcmaiv!A`C9_8vCK3xkAz0u^oRWg3)KzA zxO>Sx@8;a$n4NBHs<-~}GlX~OmzWVmPUi)&h%AY61}mkwu?qS9%}IMF6Brq>L0<`p ztd}g{s(Q_1Ys}6|rz?tODZJ-}MkkysY8F1fBy*S@QTq<=rLEBK2{zV-Fh|FQQtOLO zg`L4_~e(FQ-eEJ2Sf_QEK diff --git a/tests/fixtures/onnx_genai_workflows/tts/policies/code_frame_update.onnx b/tests/fixtures/onnx_genai_workflows/tts/policies/code_frame_update.onnx deleted file mode 100644 index 76ba875d751abeb777abedb6463134ac1114ac56..0000000000000000000000000000000000000000 GIT binary patch literal 0 HcmV?d00001 literal 1580 zcmc&!O^*^W7^Z}e;RV8OO;(&lvpXT0?13y_9z3m@y>Q)l!I-9$vXq6+I4xw@jUG7r z12oaZALIYhGJGt>HJ;F(`hNE5)2FYK$CQQ>Ps0Z7K%Vu*@!Ad27-q)WP=$_zJ1p|26rcG z(Ss=F2V5;O;1hzhkUbv39RLhdauEeo_DIbwu68k+nwMFc`-1Zo;lh->_NCHPYEpMg zn|%zR?m276sS7ZrqX{QMYIXi|z0h}*-50!X$eg+11}vz!I3;hSdOzX(hLeVSaA(BO z8(wbX*%!VY`kK4+HUBStpN-gJwOh6)+?}v-e!OXW)U_Yj9wlGK@ zX%gqVN0*Fh&BKckA9JHb62JCyue9So%|nVtLY0j`lY0^XK3t$*0wa^4M>X$w*2VAx zICQeO2w9x&Qq2aCbW)x2dJ2`GXrj4Iv{PD*2b_VFi6puwuyy%+yA2;&vjiMdmomqo z;nDW*!470hBngS3Yw4kmVXNk`&3ywGOf+lYB;uBKY11mL*rl!e(4skd1}DpTT)pDb zOiwL*Sq=P=U^62IsYA*%PTqc)qnA+mYbFsr{AiW6Y*M9@NQ?vxXXmv!dJ4yX?7{ne O&U8GOU!@>%Q2qt$*!zP3 diff --git a/tests/fixtures/onnx_genai_workflows/tts/policies/code_history_append.onnx b/tests/fixtures/onnx_genai_workflows/tts/policies/code_history_append.onnx deleted file mode 100644 index 41f2f9fe6724529e403e852d55a18346e830dbbc..0000000000000000000000000000000000000000 GIT binary patch literal 0 HcmV?d00001 literal 853 zcmcJN%}xR_6on~47_Sha#)#t|5TizuiR020>W;k|jcH0LP{eke76KYKyn-8-`YJw| zodF?`7+37td%vEXdwQKUoPJVu(bJdTX?RCC+b1_lYD1+V=kEISELDS(^NRCZI4uD5 zX}8CTkVaS0C(6`vuqx>j*8yc5Cu)bVA3U}^CFMY~Pz1(fU8*%{3m(#twTC!)|4=J| zyFf)6L^Bm|9TD(>V5J$_BFve?xjwm-`nkvX3n#vtpIu_Jer0u9S(ddg1evzbMukI| zw}Z2i4~(5PZn`;}8>5EA&MqqH@%3~KQgk46(Z-itl1IG3qOY46Av6fMY}}nped9gY zaWZ%dGI*s<4Qqpg({bDyl1}PwcM?9{fP&?6GtPGFKFrOD%WzspgpyvMx`GV}Q__T#8Jv85 t8=)&W{Ffx6fww2rSj3~ff#?bv4y*nM9YN_&4c=^WX5wAFPC+7H`~mk72`2ym diff --git a/tests/fixtures/onnx_genai_workflows/tts/policies/codec_layout.onnx b/tests/fixtures/onnx_genai_workflows/tts/policies/codec_layout.onnx deleted file mode 100644 index 31ce6b3a9b69ef63df9a7b26d8b599edb29c19b0..0000000000000000000000000000000000000000 GIT binary patch literal 0 HcmV?d00001 literal 428 zcmaiwJx{|h5QdYGHt8Wm<$;ePP(f8e3=ynMfiW`^2w9HnG!}_%*^bl*vE)B6u<=t# z2v7-xSnnRa_uLz9qG2ghUYV=c4R~=O$homn*JwcHf|<$qRuf^AGz^P&fGPHwpj`;iASIUO5TaqGI~VXmKb=GW5BkNm z+cn#`r*Pc&5iOZ$EiIQKQnKK*##D;-HnqrmdT!innEgv)zCin*R38$NYQ+mqbRAzd V?i{9nX0R+{)P59j>m zU9gO*mp+iysemhKpX-P+jtiY#*oq$Rd?n>jvsgsNXCc*^^aPJ-%z7hiUEg%J!TW{E zdJxT2#C3*%cLXy9xH2G@JV+S7phMF3YCl1)s@-sBnPwH72hO8F^mQFAQQ#}WtLciA z)P_n!4!r8@o{KtYINyn84cq|@6HbIQ^}-TGy&+iK6)b4P4pY(ynQ(9Q{4t>+$)z4P zR`O^`e~HUor9w_+`kJ{~pZco`701Ocl<-QI8rJKs!AunNk3AdU?Y-kD5lI{>*;xnPGI+6BoP7L kn4%_4e3}eJC}=!7ZjaF(Z2#GV-`;X&a?79EP^i_v08>02xBvhE diff --git a/tests/fixtures/onnx_genai_workflows/tts/policies/last_token_logits.onnx b/tests/fixtures/onnx_genai_workflows/tts/policies/last_token_logits.onnx deleted file mode 100644 index f674183330680e7f8e3468568778f80f8327ea0b..0000000000000000000000000000000000000000 GIT binary patch literal 0 HcmV?d00001 literal 631 zcmcJM%T59@6ox4ohv5KGyHLPr5Q#)$CJ=Y7yev$31dVArr9fq9$7zQ^j2m9V#Q1Q& zhYk#=kQi5X)ysGOd_Bj>VSA`N5$W@H3tkY;kI1!>+E8gov%LAWm#V`_v+67_8YRe% zs2_16q$z$Z7kWeKb^)@I4!I5}<2co>!9g(Sxk}2hW}yg-%Y3Re=?fmxko6}xeRtok zKzWUtWf09&!1auPR|K}NdctT}vBbGJ&XDA(T)3g8&qFd7ZdL8%*vNrJ$3hqXQ4nv> zJ&YP~xv;rCTPouds5=NBK?ZMl)UZAPb~JyCxMUpi7)FXwFY;+P@lN6AkMKV1#g)6? zl`Hha&RREYu$&A|O;H^VlhL|8F5^}kp`;rqU$6;bN}7-|gVRsXQ*;TnKS?4wxU`^> g5SK<{;fvV#qBTXwQ28~3ce|XK`4q1#NNp8A0X>A#xc~qF diff --git a/tests/fixtures/onnx_genai_workflows/tts/policies/predictor_body_sampler.onnx b/tests/fixtures/onnx_genai_workflows/tts/policies/predictor_body_sampler.onnx deleted file mode 100644 index e2685bfe5abf6f9a185f76a54debcdd4d02e72e5..0000000000000000000000000000000000000000 GIT binary patch literal 0 HcmV?d00001 literal 536 zcmaiw-%7(U6vmga&F<)&1i>;%!WmB$=#;Hv=ESJ1=}s z?Vqg|f|#r1%XiNCozL3Goml!}r6=z#_(IT$WGIZrZ6lXF8T9T|_(0brkpU++Zl`*V z`d}?NPXiHayYY#dqhn~%wb1tW?x0+q^`M*3nCp}>jvKQfoTtmslS#7HOo`NZETUSI zh2V-Rw#abv>3QbB;FqiDK{S&o*M$QghKlnbBbvr(#MR?fO??KZ7Q!d6huc0iYyk%t zUL>RbDxxa$ZK&45edt+xxLf-+=mWt;6@(Y^BRViaWtn^#xKh`STV zE%VtSxYv#w=Mfbj6w&0 jfCoji(Pq5+wnh^e n{2Pym0QU>JqkIJ_yz;NnhSKfL8eM?%=a4E!33$?g+M)de#-y^; diff --git a/tests/fixtures/onnx_genai_workflows/tts/policies/predictor_state_initializer.onnx b/tests/fixtures/onnx_genai_workflows/tts/policies/predictor_state_initializer.onnx deleted file mode 100644 index 5f30409921a77135af5e7155745f39eabe49d79d..0000000000000000000000000000000000000000 GIT binary patch literal 0 HcmV?d00001 literal 14690 zcmd6uTW{P%6oBLHChJTtO(u|}t|;kJ2&5vcJ(o>DLZuZ56T8FqHC1m*+p)DOE}(}yz$cfhUg<43!lAh-ydy>>U+ZJu~x=pMEGp4V$0 zUBc2=Umt7%OS&qS;fBpl*AL?>a&@!?9c@X+lAvR(aZx({YJk*l>7fn5x%;3KG`+}m ztfj$G;!meK*FW7^0A=r@9Tw`tzYM@>0@o6)*wh8lXNUbsM%FshNBAqz6IHKE|L<5} z_S`!;^LLg_W2JIn8Z-UhOQT$re}6GRmdnUp$)X{S!#zG8_pU4+|7450m=|EA5nO9@T9>ZaEXKgT z5(7o9oX-uo#3heivsQ-5($5~d#-+~+TU&?eQ|N;Eg3BHofM&cjx^~O0#i{=An={{k z;=8tOT^}ANYmRMKD}z6ZN^99)j^Zz8oWF{k#y<=oP9qTfBsMzn0`$2TJYp_H)>8kd z??6RcjR9e01{+*A|e zz7E;^vDGc!4%&Ox^5C#Pe|yTZTjZviZK%>21iV^sH6cW|6Sk9mRl5~Rk;VW~Y^42Q z#2SqP9WDyWaN*|ErWJ0EEr*+9j}14oz-^wC29Cl5hv$Kq1}wk|+kUi>G^hGVM#rq~3uPiNe2yWZH)tYBD8RH71=^w4L;SV7YVCDUwtAlEd)cfht;fp4T;H{p8n z>W=kuPyAr(?$+R_@g5LLNZ{+>HP-b;+I1N&4o8c-Ha-Mdb&@dFw!&+>4quVAWy9Ky zw}Dzk1J}SR>%EkYS~88sCm{RfJ$x73WuDpQuQtINTpgO2xF%K>Tr-~l zWCqL}TX8MwuoT86Tz4Z@awM2?En@`NN0-Ec*SWfCe!vZ(*1@j9IZ$TC#~p;K*wOz| z2WQl<9;-L3l!YdkV`55die^aa>~UhH5%;+!c|??F zBg(Upy{JzO5l3XN%F`xg#5@~Oo{j8nW-3-1G0#SvXCuzDk-hL4HzIrGk~YaU;yfF1 zo{j9Ocq&%3kv*4Ao8%Emo{c2WM)u$^6)TR&9^Iu)%7}S3k~|yP)0I@LG-95OG|xtw zXCu1-9XBGo1x%Y{8)=@6G|xu1OHaj$HnM$k+EjjKY=hhVuNZe9`2qW$FfLWz?RHMu z%}cl038J3Yj9}^8@17YSfjj?urQ2@78-uI*LUz5gi}p#|>s{WnpBZm~t+6uv!$ZFr OC7&wZFN2wt&Hn&Kn7?xX diff --git a/tests/fixtures/onnx_genai_workflows/tts/policies/predictor_step_update.onnx b/tests/fixtures/onnx_genai_workflows/tts/policies/predictor_step_update.onnx deleted file mode 100644 index 5c2e73d7a71c4f373a8eaac8f1919cc9f7b06eb7..0000000000000000000000000000000000000000 GIT binary patch literal 0 HcmV?d00001 literal 2039 zcmcIlL2uJA6pp)g&E8tJc`!5Qdd_4(U52gh zE^G&9rww0-v2;TosD?YFQj$JrAr0OBCC+{L*sX$dD=Vr&r0WMvMj7xG;nD?p;|sl& zP^9Us-HZ*gX3b6l^$mavIvg>=g_4IFV=*gK*9fLeK?)!S1J|XB>^WtPEgBi9d$r|w zT3>}c9dnr}alesfBO=Zq!JAVO-<`xayYFlz)FmNgOOXCkAA1V#&Q3N}_G?X$En z4(5ar`70~HU($|h-*!}c#!>59=82<>>9Yd$9i?2Ma^1RsgFKbJBoDp&OUJx>myUGuVR4gE)z@v&CKcDl70~Nc7 z??Mq5dQ`c6Siwct7fOrs_)+?FcoQ~sZh4xUvr+E=7AYBT!$X~0ndX-9>ZES6TbIp)-g)4pw9!octE9TzbX9J!v#@{MEVjU!6i$TIO8whPZQ91m zkk^2$K^#xkbdvZLdICG~)7m&?LO*hmcpHH~4!jlKz zQ}d_oK@j(}`|UR~-~8J9q?L$(FZKAn0bdrhVi9qpNz;f0rv2`{iXQ07Peq9SjoYf8 z!ydF3h%)4fwrihOeRvEFy5idY-X1Ggr(NizG(jyXLsFX#;9M>uPo(KeGsUIxSWLC{ z=NuJPY@U(&)AQ7U{x4V2!`Dnm)P;i_MhY>^d`%M>qk6pBr9Okxwna{0k2C{n*c=WB zya+~vWlU8T*ifxSdeCj}kdlLpj|>>pLlvJn6P diff --git a/tests/fixtures/onnx_genai_workflows/tts/policies/setup_talker_sampler.onnx b/tests/fixtures/onnx_genai_workflows/tts/policies/setup_talker_sampler.onnx deleted file mode 100644 index f609a5b8a6a49d4be6dc363d14cf83440ca7953c..0000000000000000000000000000000000000000 GIT binary patch literal 0 HcmV?d00001 literal 534 zcmaiwPfNov7{=GJ&Fl4PcQ`+ z{j=>Lh&d%sp7+UYNsF|_E4E9-l^j6$9DpqtVJRYn<-#%u`ZSrob=O;?IZo@tlGR4H%4 zQBuhkIcYvU&ukd{a+MjLVj@FTI>=!t5rf=QG|6I=k5{|aXK-p7;DEr3 zVANm6ROY?~^|weLdgdPK?tE+G!+?t#gd1X7^wC_TZYE;Ra*v5rOUg8nuZ|QkIuqNe z_-rdU*S3=sh{}t)T&dJ3OKr0&v}$ZOU%GHUN2W_jpnkt?ej|!r- zKj%ylMdul=KRr(!=>KvRJ$y}N%5>qthmm3|%zRDaR50~;wM%^lrxwB|(8kSx7&?ao z3@?(BU@3^o0voEea1Xjx8+UfTHF`&IQGxIxmgIR1mx-53!D;4GnHWW=!Q$0bQZRSy zxFw%$1ozr;;~Y`(QH3j%3T2^fZn;*L?dl5`o6Q(Cr9{MnsZCG!ZH*={_&1Z!L)&u=3&6wah+Gj-E6b!nR&1Zsz6E5a&@$N9NPsO)Nml*>xH;sQdW$;54^O(!8U zt=dIgdPYcy1!<9x%7F_94%`vqKf|8MWNc5eh{I&JQYCwAKYxDjeec;X>Ka@*ANHIJ z@4hE384OHv-Xnbvu3nnYPOK3zYg!>}xWhi}>B~W{0>Bb+`(@=;$yN$r z+4HTjU(V28cO9&{);aM;mQ7%$dl$SrIygeFJ$pDJ9-Hw>j7Y$UB#c=BBSV}P#$OZ=>n&~~2rkV>L)W%^)6mu@tw5iK zz7lpTYanM`IbNzd{7V5G61WwhqNa2GCr6=?v9?3;2%iF+Sl^JoS2VDEX$>w&rKl>U zd{w1oLDytyt56k%!jsJrJ z3$a4YhBnn$h@_#bkqRrD)DBE{rsMkFA$7@oS}MC0Tv7nYd1X`d%I53y%H}MuY|Zh? z)*IoKFJ^jW3#!x`m>;%ze{3(zAAhGoUE~FrB7$4JVgK3`heaRgBzz#v$oWEn>l}Gd zO)EuY3E6{c9DPbOwnp?xHDLb4u?GQQv(RXwzFA{bfA-Y{APu0)D3_!J@HO7lkgrx@c56bH7NU|PfA;+%oSIA`ER z$2q~xxF}*^xEb80MnlgD-XKoj6Jv={Pdk=;uP{W@PzDPE7%rd~jfGu6We02~h|!!& zj5Hlk#??yTYR$z(b7^W%H5pesfvY_iS1v_d|8RLqVrWw&Y)kC6dH!}J{wCiSY2H(4 z*vssUT${l|!mHM*w4?1tkB+Dw5+o5^OTIIAqUr~r)eohcW3`VID<6_? z;6#$a^OMX1xr7r*2G3723*-{UGM6xZ!?=WzKpt#}UCeO_*F`R2nqI;;z(T5sfAd_z zb-0`4Qb4MCY-+r2Ho>xQ4bI5eq;*c8)!=B_wg%2~GOlQ>midr-=y%mgSu<17)vPdv zx}|6=s(pUmNfgo#!27BK_dy;m_blH&0T7m=TDwQ+^T@Rc$f|1ZBe`4yKL7`69=;bZ zEyJBa*L0mZAnx9Wdy`fV10mznMf-H=J5=yeyu1ht(T;3+?_)5l3z9l23 zTch56ki7W7U9e48Z;Cr-Ri*bB%nrKnKB!O)-HUP3*p1=J+ zL8oQzwOWG!s1$aNZlMSxw01uo$)4_VcU5jc;RpB?fvo9UofbSeS^KkV+0m2?$ee1d zd=jN3v%xJF1vj~o?CDIwt?3jdc!w(m@*D07K?l{~3dm8lumwUPXZl@QU`z*g___d8 z!3q_Vr8~bZ?ogk?Vv4`RC27PTcj;7@-|?^nic}_G77PBFZ1E0EPnM^E=z;0Rf$1hM zY++LgQHL93+X~ly2tkR;EKE0Ls6_#WXn$#tcLDLd^cfLZf|6Di%3Cu+w9Pr%*(oM}3JcxN&z6jHS=uAG5mjjs>UjdBm{&6Popj)?6CzcfC~flR6V EFC?YyR{#J2 diff --git a/tests/fixtures/onnx_genai_workflows/tts/policies/talker_step_update.onnx b/tests/fixtures/onnx_genai_workflows/tts/policies/talker_step_update.onnx deleted file mode 100644 index 5812ac850e80197463e5e98a44ed90a6cb4bd74a..0000000000000000000000000000000000000000 GIT binary patch literal 0 HcmV?d00001 literal 2044 zcmcIlO>fgM7>>Jk&E9oi9&9B-Q>i8(`GTKaXuMik6pt*iTi^z8xe5^3ErHN`2H+D*+XY5p)N@ooQKt0jwkt=0|mouZPiBW#g2`v z_`5;*tSJ74;6BWAijBypQoZ8?Mw(h3No=;#)asZdvcma$&O}BJ35+J*6l|f=Xzos44y6dRcjHA}I%o9f$(`NptmtsSLLwCl$SClrVyF9aCUKx_7x5z~ z;zE}yw+BnO==wrwWuES(w}!Xjv98^KoV_ru&)KMZ1ouf9@4!=CU!K;k;N?m0%8N>) zyASj9Jh(`Id*))r0Jd(O*HBR3NeQHp`d3a}b<=S=9AE>dtk@Q2*JA`vMebrsHi{&I~@VE4b3gnRhO kq$+Mwr=u|+avEM8H?Glhs65z%ot-gPnMd?Z9<259Kjcw>BLDyZ diff --git a/tests/fixtures/onnx_genai_workflows/tts/policies/token_to_slot.onnx b/tests/fixtures/onnx_genai_workflows/tts/policies/token_to_slot.onnx deleted file mode 100644 index eb8e4b6599f5d6095617f9fe0d294eb1d098d216..0000000000000000000000000000000000000000 GIT binary patch literal 0 HcmV?d00001 literal 438 zcmaiwK}*9h7>3un+4PfQjh9(LrVIrSWxIOP9eekpNXW8mjq8%yq_9rB^zN_lhicoE zp@NuW-iMdxd(i>jU#Y3c_07i~yj##SYRM%4Q&Jl;o{O`5qznA(Y=X}G7R5MOS z7w+yS^1zjJ#dS&<$DL>hXX$bgDk;~RWg;~pOQ_Z)7CfUFiwkT&K1Blveo=J?(M+XW z*97(&8DbQv6|Mth`Of5HNf(^F^;z{hY@1>Pkfwhl0?k#NlR@@LYl8d bB4}3Jj!Wwj27h+&`J6M;Ofu<0$M1gu>uQY1 diff --git a/tests/fixtures/onnx_genai_workflows/tts/policies/tts_state_initializer.onnx b/tests/fixtures/onnx_genai_workflows/tts/policies/tts_state_initializer.onnx deleted file mode 100644 index 517b0f069f2496c38f31b74aecf3eaffec9cc871..0000000000000000000000000000000000000000 GIT binary patch literal 0 HcmV?d00001 literal 2194 zcmd5-O>YuG7-qImwlAPC4VJEnX-U&2#1JV6>Z$6DiwAESvl+HS849!O&Okupg}-3D z^CQ&1r?U%%*^R^#ZTB=Y@5erQp7+Dd;N79yVI%SCPaft7a-Q!FJ!wn#nsR~n#`gOj z@u>aG$}SGvy;QeyD!f$yb|}9rp*b=tSx5;#M5E z`9wHA^JLW-5FzXyqdxJS-V_@jKedm*ig^V&Y~i>b6@dqQM>spS-@07=ic7m@JzSkb zWVNz+^H3=OjLBd`ZN{azNX}7uTc{LqYo;Iq5JL#pA<{m#atd2;3Zb5tcNf=629jjL z#8!#>5hbo3C$6p$w;v%caYxkEHR>8s>KeC7-7lp02x=}A5ofx7P?U~p1Yte8skV1r ztv{ZZS7o_;uUL=@i?Q;+MEP{vOf0`SN^YwvW4+RhOWxEZkM@e{=R4??rY3kRDtPNw z1^=lF-og+LLt6}7xk2w2_d@UAI_w6{1{JCXtdOoh-2nXMeE^P&QVLrQbQz3lCK(y9 zFVz3=G`@`y;(aJ6-`v{QFq55Q(6s!Q3+kLf zgg1NuMU|(Z=ShId)j!ZiJ_azAS5g^J--<5)%~py`l6903i&C?16=n4GsD@_QiAAXf zT1BBqlq%wt#h9TpI9_c1>J{Z`0ZZocs^<=vGgWhqOP@Fr8(+W8&}%sPucyr}@!`^I ik)%pS6E*A?$e{0J(cwXV)E>9{J*(XsDUaN; z#&-9nGb|7LV|)0v%BE<4%>rqx!JWxBWt7W4DDfa<84K~z8K&5;?oO@;Wign8zw4*TE#3|qb7WEq|9Gm zBy8i_es337(QmUz1`$Hdes46kYAvhE{kz8fTaCfE+ip5mO^C5u zMM-cWsd9={_xxJ}yVtS|k?PVi7iXh5O&<3; z$U+P~loAj=A=2l|&d7}r%)&xco4rC8W3xCzpZ7G#zsw>zPyrtyl*dclvtlLnAPayil*1M3e(NHI^$J~h>%|%Ryr&iRHxSkm)8*s-W`ATg z?dFCf)@Bc=y4dehaV#fk@Uo}Ddm3Rpafl*3YWLD1EL3BPaIGp*V_N8Hajm*9F#h!f zZdX&Jkk%3*t$C2{AIJOi7<@fF3{G=}lMv5m)P!?5uaqdQx<7m`guMdK;QQVL9OZhf zCF(KE*ogf0wEGoG2RPvU=dx!xLykivA0=-k?J5c$eg+7zv! z1^q8f*653U=o-@cUsJBfONn~SCUE-@FBcEGkNXh&+bOVL5@YY7E5lx`iYC__ z_9|@3?S1V3y=}eM8s=ZQeJ| zB<{~YD&3!7Z9FC2i$`oPeii*Ig9KgTqFzlo&brJt%`8xwWH^wk@c`>8$-+>L!b4QJTv~=|(7cQiLOG|`n$N<9Z0RsM2z%+n@q7kv z&+h1sLz+rR$)i0tg*`WSdC$STBQiFluMFjwuL)agmbz3>JUxI%qUcM-+?B1hoh|9F zXNBe+?p+iCGvER^%-oa)weD+&pGu?FyhtH&58nuJ8~1DB$%cqVy3cc?}~6Q zto8i5c8q&7W8Is%YuwwPa-j$xJxQEcxof}e3`gZ^*_{Ab?qjb%+$`UA+Upx*#bE9) zmh$q&E3ZKfyW9Sf?amgpr9z$hNUu^Q)O;mUi8zs^;#{}{l5A(3!KmHs_m026^cCaX zOs4b@%>ezc_vb2~#ZNFV6Aok(%!h6r+dSK`FyYZrZpJ$JDzweQSnAd zuMX|r2vYeSORE>uuwWT^g|uAp$%#no#aZ^MbOcQeoZ;q~%o#}DooRG;K|T{FNY)437{L;R`Ij8po(ZB! zMXNlWLKiOD2Y2 zO3E3qDl=e09x-6`qcUK^&+Zf^;r;(-zyi_A4VYH&3|NyHFrf_^Fl@x!3j-#f-sLiX zcM44u4N12ftI|?cDa+!n%kEj(^LSgP)$XlZx?ZToM_yrSH}=x?;;~RcX$}H!*tFYl z?AC~E+dP6E@{gm^kG0yH?jgI9f)2q6NAL73Ie2%SkKJFpyRmy?H8EI!cKgBlYU8Xl zM1LYVL|;WoDiR)VrAlyjRm7NFMT`U$!Sn0VeHCHHxKpXf4-TRUZF@p`L++)*BZhd1 z+eF6)cqRL$^S0$4X2z>Q@PHnUzIMx%lFVVTIkBaiR#iDN7&;@!A|P!l$NWdPt-`5z zXXdv=+!!pku~&FeXw<(kZrX$62jCAc3O$Gp-@Ik{Zq3zM@LX?l@7x{gW3WTrw7U+j zc^cOQZ`j7?k+w#5muy?_bLS=xVEgFjT1u-Lrlu5Mp+z`@GG#S52wHkhxp6D{`iL>3G1!-oPpDlbbV zSBjGyFpZmG$brOulS6lCi=vP_UY5$qsO1N#>BwlPg{O&ld?yQpKHh<0(NNj2cwfRn z3I8gK(qqou2IZN28H?x2?e=KHGBkFQ%`2^7nUGR<_%k#}^~1!C*Lh1jcjuPYcq7rS zjP2OfYU5d{HJy-I(<=IQ4$;)@pv&7D1RS0q;KaCkRLoylhF*9^r1z9E;_M)+npF-v zySugj51bNbN13>nuI@#taGsJ02V5Ox7#!U%?%dH$AVXzQ?B+aqdrrtFdR*gUTk@L4FWEVxrJfrjt~zhSadFkYGK`fwQ$k6FyY!+vplagkoD zf4diOd@+%pPshBMWh_M}!+3aMv=QegTL-!8aR6iEUb??Ki}V=F-~ivxaVzeg4!edz z^gAw&V{MdqNyWu+KDFoK>f$(#5eDNpe|h7+{1z4&NFcwT!O6?$phDIo;vzg&zs#f& zx2R#r7B!4Lx~TDCET+aXq6d4*{JO>COakuP$kK#aEzk*Y1rx-qKA3m3`tTedSXWa7 z)|eB2N#WJRFd9-T<(K^WIhs{zbf)mp8Ri1i_^Vlz3?i3NPL9`^HI`A%owqab9liKE z#NB${?vEMeA*NAqnJc+-*HjA+)2NWWb7KrM{;FNr+nL$fM-4vY^XfgRB4PKr(-m-# z&NBVe75I@r0jDc={2@CyU6HKo+yM;IS>*w^bbAb!Zns)a&jOVgJM1Jey+ndQz-~HX z*Ydbo-#hwr#rNT$8PV->m!w;Gnb-auabAHRZ^ zlxd=ArDpz}Og=l8IVZdPt@8-Vcl&UiHr&sz9NWaV<-0fGms|GO?i@wa{Tnw%&RChM z8JdZoyH0O?yaCTuRUc%eA7Pj8r2dF%m{q|?iqMVObFnwiG_nUWrC1trQ3TK&C>8RV zxtVj^4f>7zf9TPXNLj6sr$t|ml&Ot8&Nje*kut53&yb>Mncm0)8}47}%4o zVUStyI9g*0S1O%?g5-l~l=!uL9>`N3@M$c#TuJ^{dM}%YKWH-TsglM!US}oE?(L$| zdzt)H5@<9G>#Dh z)mZx?3Z;*}jxIvgR}!jzN_ipz?iljuyH-f~rqamKZ+j7xt?+`M<&be(lOx$$MDBA6BIh%4g8DKoX|q!nrr z_G3isv0w9DDds4|{(4#?@%2>;-f{0$QCV$A@Ci^5YpDv>2N4XnFh9G%Cam)RMtB_I(I_~OC@CqirW z>Qa1tdVL%t!1@})Dv<(V)nyiZm+bB->>*~EhY&Xk7j@x*f$x7Ydc~DXGaJAtd>xeQ zkNmGa&Vk4-N5Twc7()w@?-*_)KE+9)$TU-#Lr0+@Yi5vJgIF7h&PdglGD~PC4J1v<2d&eJ?cC`Sc_QfI>Y;lp~ELIS_q6o>G+}CsDj0 zkBkF(a1gP%AQ5&i-D!q*q(z|?;&fJ)SmTX*T&K!K zfq^t8FGo4V?wnZKaoM-5w#O_?bl3rGXDeU!i@(2 diff --git a/tests/fixtures/onnx_genai_workflows/tts/talker/model.onnx.data b/tests/fixtures/onnx_genai_workflows/tts/talker/model.onnx.data deleted file mode 100644 index 1c02c23f81ac7e57af9a942ee134311508f73ee0..0000000000000000000000000000000000000000 GIT binary patch literal 0 HcmV?d00001 literal 72704 zcmeIu0Sy2E0K%a6Pi+qe5hx58Fkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@ z0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VK zfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5 zV8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM z7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b* z1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd b0RsjM7%*VKfB^#r3>YwAz<>b*2A&2691j2i diff --git a/tests/fixtures/onnx_genai_workflows/tts/talker_prefill_embedder/model.onnx b/tests/fixtures/onnx_genai_workflows/tts/talker_prefill_embedder/model.onnx deleted file mode 100644 index 0bdbaa3e8a3442fc45a82ff4d89aaf795b580be9..0000000000000000000000000000000000000000 GIT binary patch literal 0 HcmV?d00001 literal 21848 zcmeHPTaP106&~Aj>8ZUrEf6cO6iuAVkb!u*s@rXQ1jxI~is0-@8!||UXleC$dd7A< zZR76EF7N_x{ET@-`2##4B;FvPJc0yBL=oXB34|wZLQs8mb&spMs>7`0W=0yxeW~iI z@0>p0Ij2swbV*)6nH&rsI<+;W0UOHRTW^)Xk1F6L*BrlL&5Y^HIvS40hIMja9UfY< zhHE``!=s^n+<0USj|VQ0{L{ljM}}`2oK3%Cx&!?WVJ8#b^fhnAt#*Mq~} z1F&hECzdlc`<8s4($(J9tB)*O-QRz0-=Ek0U|M&=S*4z|w-sE#6YSjwPfy=CZcJ?Z zvD2Ror*5M^HXX+p46T_t>kqyoH(uXncE0`jt?jtK?FYSAKrPjyH-5wEPo|b5zx4Vx zWe#3HNkq0EfVV3EQSJ{!HJwdHR^J^?Y~!d8LwU)a3@N2r?axfxnNA$bfVD4B9v5Jk zY`8mb=t~XOH{&xqSHSAjnw@k@-IdeQM$aJS_%+Hyxa61|!`{~2p>3J7-n~poevU(u zL1$Qu_nBfu*omx|kN+Gu?X+#|El(DAHM|0%0p({;0kl99~Ep z{GkGJ;tv!G{tMsugQ4ljs+gp29v&KMO{Y9&p_siUzKqp+1^>PfF!*H|<6^v)5SoLWYk!E62vKCiVguX#s1 z08}jd@OtUARNlD^)*RQIxz}+hQ7#7K{iG~ENB0D7Pi#D4+ci2hTD-U$-{a%fsY}s` zof`nWX^tOS#?W@1?%z*0yZ^xde>>gu0a4%i6|o z5cmYrlv^3kOUg5Hv!Pns6Hx=w&xY!GJK9Sm@%JhquLMLqZ5VjknwYRXw;5`+mWZTd zR~u^Q>T55NjNd7PT)@k9y10TI5p33;rwND1BEe>%(nw;Q7pIYgtR$I>fVv3hoyj9X z{mFCQ5=El#Rse^di*a01dSVe)HY9LhmEX^5?Z}RN` zLPQ$%=|md!gKGD~YWJgR_x)=3Pu0R?Qh!qhd=L`)$eBt-Dt9UssfA0WzEc4WiqbXu zxW=u2w#<$zmx@$YxktL^FcpL}`m$K6Ng|tyv^bmkEwawwRm#jV>I`;M&pO{qyt~YL z1YFlgnS@t~ml`O}zzbJwY0lG99srz(U*4p8v<@^)mOcvVz}6UF0-m!LRhke z06X%tOb@1Sx&~_3UZB`2Zxz_PB|x}cUk{FVE`e3^@zC*?&x-*#+y18Pb+iP6G{WX8 z#BcIEk9*A*MxtVck-sYg9_7r|a@3Kx=vwYA45&)&S(=K+3kRwaBvlGg!tnDl$i##h z_xUw5Dl6`rSJkLq+~=3o#HBP|VNEeC1KSk)(>B^nB@LnEFvTk3GsP;dg-@B{;?yp^ z$Dz6$mP17>!0$Ysnl?5JDjo3f7N7f6ymlNM7elGpbmteqBo^SqGCYFf$Ph!2?o^E{ zo1}V0=*0;)*pTyh7$nS1LKGYxhRPS=sl|%$e#;T3s2bCD`>~AL%yckx9tT4&kc07W z4s21AWv62koaw3P^3}vvQipD|$$&-0i=io&ERI{kHYYq5oHQr8pgC#tS#Yo1x;Pg6 zZ61|ZIJ|0BxJP)^UbuJRRa3T-2adPOli<)4K6|0zec>s4Q5>)Ok6g;jve%($)f5iD zT7z`RNxxdd_0&c3t8whYDk=)Ri8c%65FHK5pK@pBaEvQ%jd?*jC-Z4%rSRgJYp5qhFVLB5YEP;}?Of^O zs;_2(OXq9@aVa(sC1rFiwhe1~2Od;ngQxBHWm>BP2dmpZ`%)9nvHQ&zKY^;yTU@Xj?Gm%3q#ckqJ}NCjf@1e{1E zc_{MmwybAAJc{?o=L2kF4}Mqz3yD#l0vnHHBr=kY+{zwF%);rRjBW~k=t!=Dt9Xb2 zta?P{&1UrFTp!@M3NGV*B(Dd*moJHwtA<-L0FZDWtKy@GOr!jXMiCaxlB2jz#GndB zNy1v~jjuE!uR!V*M7^RWuG|!@oRP1AD;TRyP{H=Y2aFt$mgtNEd*9o_!S!-(v05}y_9~+P^&Z$< zI0`|CNKJ=x1LG+()f=aPJTx^5Kqipd{`* i%Un32mtyc+t6-&i{nk_93JyXVC&T`1;!KX*YU#hZzdb1c diff --git a/tests/fixtures/onnx_genai_workflows/tts/talker_prefill_embedder/model.onnx.data b/tests/fixtures/onnx_genai_workflows/tts/talker_prefill_embedder/model.onnx.data deleted file mode 100644 index 885cdd9a062995e9b753010d104a5fb72e2baa58..0000000000000000000000000000000000000000 GIT binary patch literal 0 HcmV?d00001 literal 4984640 zcmeFtfdBvi0Dz$VsTV1P3IhfV7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd z0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwA zz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEj zFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r z3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@ z0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VK zfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5 zV8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM z7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b* z1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd z0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwA zz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEj zFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r z3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@ z0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VK zfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5 zV8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM z7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b* z1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd z0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwA zz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEj zFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r z3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@ z0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VK zfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5 zV8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM z7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b* z1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd z0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwA zz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEj zFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r z3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@ z0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VK zfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5 zV8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM z7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b* z1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd z0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwA zz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEj zFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r z3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@ z0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VK zfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5 zV8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM z7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b* z1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd z0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwA zz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEj zFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r z3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@ z0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VK zfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5 zV8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM z7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b* z1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd z0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwA zz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEj zFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r z3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@ z0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VK zfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5 zV8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM z7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b* z1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd z0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwA zz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEj zFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r z3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@ z0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VK zfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5 zV8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM z7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b* z1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd z0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwA zz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEj zFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r z3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@ z0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VK zfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5 zV8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM z7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b* z1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd z0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwA zz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEj zFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r z3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@ z0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VK zfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5 zV8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM z7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b* z1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd z0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwA zz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEj zFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r z3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@ z0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VK zfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5 zV8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM z7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b* z1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd z0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwA zz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEj zFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r z3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@ z0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VK zfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5 zV8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM z7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b* z1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd z0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwA zz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEj zFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r z3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@ z0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VK zfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5 zV8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM z7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b* z1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd z0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwA zz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEj zFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r z3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@ z0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VK zfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5 zV8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM z7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b* z1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd z0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwA zz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEj zFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r z3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@ z0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VK zfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5 zV8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM z7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b* z1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd z0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwA zz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEj zFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r z3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@ z0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VK zfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5 zV8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM z7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b* z1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd z0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwA zz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEj zFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r z3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@ z0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VK zfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5 zV8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM z7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b* z1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd z0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwA zz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEj zFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r z3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@ z0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VK zfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5 zV8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM z7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b* z1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd z0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwA zz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEj zFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r z3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@ z0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VK zfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5 zV8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM z7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b* z1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd z0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwA zz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEj zFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r z3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@ z0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VK zfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5 zV8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM z7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b* z1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd z0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwA zz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEj zFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r z3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@ z0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VK zfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5 zV8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM z7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b* z1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd z0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwA zz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEj zFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r z3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@ z0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VK zfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5 zV8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM z7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b* z1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd z0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwA zz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEj zFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r z3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@ z0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VK zfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5 zV8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM z7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b* z1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd z0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwA zz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEj zFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r z3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@ z0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VK zfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5 zV8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM z7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b* z1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd z0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwA zz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEj zFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r z3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@ z0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VK zfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5 zV8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM z7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b* z1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd z0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwA zz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEj zFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r z3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@ z0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VK zfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5 zV8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM z7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b* z1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd z0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwA zz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEj zFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r z3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@ z0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VK zfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5 zV8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM z7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b* z1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd z0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwA zz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEj zFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r z3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@ z0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VK zfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5 zV8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM z7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b* z1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd z0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwA zz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEj zFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r z3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@ z0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VK zfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5 zV8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM z7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b* z1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd z0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwA zz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEj zFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r z3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@ z0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VK zfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5 zV8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM z7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b* z1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd z0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwA zz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEj zFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r z3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@ z0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VK zfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5 zV8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM z7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b* z1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd z0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwA zz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEj zFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r z3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@ z0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VK zfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5 zV8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM z7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b* z1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd z0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwA zz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEj zFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r z3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@ z0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VK zfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5 zV8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM z7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b* z1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd z0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwA zz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEj zFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r z3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@ z0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VK zfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5 zV8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM z7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b* z1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd z0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwA zz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEj zFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r z3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@ z0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VK zfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5 zV8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM z7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b* z1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd z0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwA zz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEj zFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r z3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@ z0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VK zfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5 zV8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM z7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b* z1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd z0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwA zz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEj zFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r z3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@ z0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VK zfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5 zV8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM z7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b* z1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd z0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwA zz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEj zFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r z3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@ z0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VK zfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5 zV8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM z7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b* z1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd z0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwA zz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEj zFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r z3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@ z0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VK zfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5 zV8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM z7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b* z1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd z0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwA zz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEj zFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjxE(qc00000802p~iJgUz z0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwA zz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEj zFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r z3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@ z0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VK zfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5 zV8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM z7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b* z1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd z0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwA zz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEj zFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r z3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@ z0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VK zfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5 zV8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM z7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b* z1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd z0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwA zz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEj zFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r z3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@ z0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VK zfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5 zV8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM z7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b* z1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd z0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwA zz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEj zFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r z3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@ z0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VK zfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5 zV8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM z7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b* z1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd z0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwA zz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEj zFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r z3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@ z0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VK zfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5 zV8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM z7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b* z1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd z0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwA zz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEj zFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r z3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@ z0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VK zfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5 zV8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM z7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b* z1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd z0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwA zz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEj zFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r z3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@ z0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VK zfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5 zV8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM z7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b* z1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd k0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|wrK028zT0RR91 diff --git a/tests/fixtures/onnx_genai_workflows/tts/talker_step_embedder/model.onnx b/tests/fixtures/onnx_genai_workflows/tts/talker_step_embedder/model.onnx deleted file mode 100644 index ec133fe847e6578e828f5eba89ca6fca50047e2c..0000000000000000000000000000000000000000 GIT binary patch literal 0 HcmV?d00001 literal 7537 zcmd^E&2G~`5U!n)Zl{!n6^APPm4c{LA=hgsPE!d9A`nt;s6By@&p|r{EFTI7yY%yIH3w7_HlbWsL+pq1=X&uq79fY)F)Al~C)o53H=7e?jPd&fR zGj9#d9`k4G>)HO}<`T?1r;T7$`JkxPhp_0em_f&JDb{&jcGjN1p`N+7x4Rdmu*))d zI!K`U2yS(b4mIENP6M~w=!BZva)Q9FH)z-Cy7jmC*2@j?`i)mr3-aeY^zyL-*YD5( zufE*i<+Sk%{-|(|XVQA;xdVtA-XX}alYvVTfJ1;Yu;W{kM`LmFVpRcy7n=oOMaro=Q^6NhGB_1CsJvE-7VM=__H8 zDh)amq6^?0?O;f}g-JIjnW#*Z1fiCnp@wfxkQ74lYm+sFkd%wcCt=TQN-1X~q1>!> zDnBpYsTfj3%_NA-+`-0pw?by2U12jEIU3R`C?{qy z2{C7_W%;JUWOa}NTbwK}mQ-HS-m|=yKxhGk^fc@b8}V|(>zsr^|0$4yQ%bk?;%mIS z#byir%^{RiD%=`&emL8_4yfh3PT250ywYZq+g3D7)@n}ZG?yXk9~=ZU#F=8rYOX-G zMZLqYjxSkdV(1-Z^B84NHVQY@{B3gBrL(Y#-pgSSw3h=5nMX=)p)xYStD?EVCuSp6 zq=;-r6{#YdSw)$F?ZxA4Mm|=@ zx;y~H(41!$r8#OLSY$DT7IVz9CF_ZFmKd<$S(uB$o3aH2=imz45A+vAyOBmQ&P6H8 zqL?G7_1=3QHAAQs!&^%MUxCXkJ9M^2x}hRxL;B>7h_gZpL1|hDwl4@lIR-&N2;1;# zsB?BiB)ZL6qDupvCxRti!DoV*;Kp_V9C%!i;9$jZsLaD$Znbh1mRW7o+6}ks2mV2r GQ~m*zQ1p`k diff --git a/tests/fixtures/onnx_genai_workflows/tts/talker_step_embedder/model.onnx.data b/tests/fixtures/onnx_genai_workflows/tts/talker_step_embedder/model.onnx.data deleted file mode 100644 index 3faa8a536d481e35e98281cfb05b9b76304cbd76..0000000000000000000000000000000000000000 GIT binary patch literal 0 HcmV?d00001 literal 69696 zcmeIufdBvi0K=g9Qy<|1g-~I@fB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM z7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b* z1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd z0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwA zz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEj zFkrxd0RsjM7%*VKfB^#r3>YwAz<>b*1`HT5V8DO@0|pEjFkrxd0RsjM7%*VKfB^#r P3>YwAz<>b*2L8YR5KjOB diff --git a/tests/fixtures/onnx_genai_workflows/tts/talker_text_step/model.onnx b/tests/fixtures/onnx_genai_workflows/tts/talker_text_step/model.onnx deleted file mode 100644 index 7270cf0aaf845db7e18bf06b1042575fa6b980ef..0000000000000000000000000000000000000000 GIT binary patch literal 0 HcmV?d00001 literal 1860 zcmcIkO>fgM7>@gyzHLWl4JbLFDcS)=Q`B}{CvigMfK-VSLPF!ia#L^3);b|}tHKZ9 z#0eq64NmYkang@DLz{Lq;Ne(#PC4~lfFk4LR^c})xD6-^Nn_bU${RMJ#_5cS zi24j0!vlC6U7mX(=S$(oK_oqYN`)YkfWh16?tOrm_B9+e9j<4eKqr?$AW*bEk_Nbn^qeGZB6R4&L{}Srl@{rFR{W zH=@#?fQJun*|U{p!4!ksXHYQB(-Tbvu?tVkDn7VNvC%Wo-UWA|tzaJB<)kW$lF}FV zv#MKxk}7Uow6w?XU;tA^OUa>3Mzw{Od8jF6Cb6f_A>Ncwl{E$3h96v$jm`Gms{zi0 oRs*yN*;=VSZYX$2-E1VkIw*l*wFbM;RPW=>0zVE#cqT3M4@T-UKmY&$ diff --git a/tests/fixtures/onnx_genai_workflows/tts/talker_text_step/model.onnx.data b/tests/fixtures/onnx_genai_workflows/tts/talker_text_step/model.onnx.data deleted file mode 100644 index e69de29bb..000000000 diff --git a/tests/fixtures/onnx_genai_workflows/video/policies/continue_predicate.onnx b/tests/fixtures/onnx_genai_workflows/video/policies/continue_predicate.onnx deleted file mode 100644 index 3177d8905c97dc120048da84e39b00bcae00c3f8..0000000000000000000000000000000000000000 GIT binary patch literal 0 HcmV?d00001 literal 910 zcmcIiJx{|h5Ur!oCReIf9-ssXL4`o-5Vf!|@gbH>m zU9gO*mp+iysemhKpX-P+jtiY#*oq$Rd?n>jvsgsNXCc*^^aPJ-%z7hiUEg%J!TW{E zdJxT2#C3*%cLXy9xH2G@JV+S7phMF3YCl1)s@-sBnPwH72hO8F^mQFAQQ#}WtLciA z)P_n!4!r8@o{KtYINyn84cq|@6HbIQ^}-TGy&+iK6)b4P4pY(ynQ(9Q{4t>+$)z4P zR`O^`e~HUor9w_+`kJ{~pZco`701Ocl<-QI8rJKs!AunNk3AdU?Y-kD5lI{>*;xnPGI+6BoP7L kn4%_4e3}eJC}=!7ZjaF(Z2#GV-`;X&a?79EP^i_v08>02xBvhE diff --git a/tests/fixtures/onnx_genai_workflows/video/policies/diffusion_schedule.onnx b/tests/fixtures/onnx_genai_workflows/video/policies/diffusion_schedule.onnx deleted file mode 100644 index be69ec803c55d622285a3798f4b93bc2d11de2cf..0000000000000000000000000000000000000000 GIT binary patch literal 0 HcmV?d00001 literal 387 zcmaiv!AiqG7=*JS%KD2K_8_(giBc4b1igB&wchdyBFlEO+bpKLA-j==B1qn%J$d%& zd<0XYJqY6AzYO2ZKZEzl?%YgOWp6%wct-?Vo{6-|MgQ>LXzQ4ERGU=Lt26u~3@>5l zndOxjMmPx3Wa~*nDdlLr>b1AJ^}N^78w!#G*wbt-Y{9r7UNV4_;%OEetrwP;s&FyS znYA=iqGTn{YT`dUCK2?1QJV~ExhaI*5QL;k%E~IETea&vJ%Vi)0)R0h0r)s*&;^`! zGvk@i?H||L;=<&L*OVLWO2!@WmoE*vgTcQfRVnFxQ(Z{Rs)foGE9>!fgU%uPGebxz Nxcizg3BU`FKLOz1f7$>5 diff --git a/tests/fixtures/onnx_genai_workflows/video/policies/diffusion_timesteps.onnx b/tests/fixtures/onnx_genai_workflows/video/policies/diffusion_timesteps.onnx deleted file mode 100644 index 328fdaf6198fa76d1d1ac7110c1233ede3c3b9a6..0000000000000000000000000000000000000000 GIT binary patch literal 0 HcmV?d00001 literal 383 zcmaiv!AiqG6h&v!DC1MH%%W%!sZvzB2zBd9TUYr7kzqP{X$I4okeNt>xXDMjcIDUk z5vEqV5X85`Yhg`XhjM%9&FcCRL98gQ+#nOXG#H35xL<`w>Q$u=OPK z3P;DR)1{M*C*C8ZcSp$Ul(hFo!+0Ndj9j2Aq(V;OA)FMC^T-;rbgI;ak5n$56Pdzk{Fp2OhJ3`|p0{%gV{ I1kyeH1OwK4G5`Po diff --git a/tests/fixtures/onnx_genai_workflows/video/policies/model_input.onnx b/tests/fixtures/onnx_genai_workflows/video/policies/model_input.onnx deleted file mode 100644 index 93a460da6cbb943b3fa24a7d1d6b092f84443eb8..0000000000000000000000000000000000000000 GIT binary patch literal 0 HcmV?d00001 literal 720 zcmcJN%}T>S6os48#&}DqUIlAGQSk>{1nWY?m3AZc5h^lFC%2u^bP_UCZ4o!!_y}%& z5#PVRu_FAfPyy(plNDINE1t#T8Pw#p8f z9;@tR&zlR)GAwJ}BNKB$m}~CCcKp!wBdrHUBr3MP2)Qv#DoS`Fcg-EBeRHK5 z%!nu^lRDtd0nZ3YK|drY7)vq2BfI8SQo`!B^6WLKQ{Qu^<8cWJe4tD<_wMBXsTr=T zIg1RTR$~F~9L_ukPhkm*ZEl4G58Y0ZA_729NUk-7CRHV=uC%G~Nd;E6HfA5KK-pWu z J3Xof=egZ=1-+urA diff --git a/tests/fixtures/onnx_genai_workflows/video/policies/schedule_history_append.onnx b/tests/fixtures/onnx_genai_workflows/video/policies/schedule_history_append.onnx deleted file mode 100644 index de2caace0d541bb215ae412598925e82ec63fd97..0000000000000000000000000000000000000000 GIT binary patch literal 0 HcmV?d00001 literal 829 zcmcJNKTpFj5XEs-+Vn~UHAs{op;8r)7$VBZ#GkD@3kX@Rdk{mgL+n&7h^1eF39&Hq zDM?cZY9(Om<~u(<|L*ymDlOgXu1w6)M+sgCUFx$lt&HW$vIB29KP%lscF=Ub8~ZhY zJ`WRQQdx6M-buORLPha=G!Yj_7dkE2itcWFt<=DXSVqSDglU zGpex%GeSpb#tL{N;8>aZT8vpT7-p_`X^6=25$fW diff --git a/tests/fixtures/onnx_genai_workflows/video/policies/schedule_lookup.onnx b/tests/fixtures/onnx_genai_workflows/video/policies/schedule_lookup.onnx deleted file mode 100644 index fbd67690d437efe41d1a898b5802585ab8260161..0000000000000000000000000000000000000000 GIT binary patch literal 0 HcmV?d00001 literal 607 zcmcJM%}xR_6ou(zgqceSS{M8n4HDFZ#RRsl{8^ar2pH3pQf4Yc+e|wIV%)Ix0o?jP zK9mfgB8hS3t~d9)r{|vQ6>)x|1|l_=?=HL|@C=LiFpW95#&V5eO75tQctSe$(&i=A zkdsa`zb@NFaOq5#YVHj=y%PA6PPoyOah&V5;6#tZKuI|>ED_oUET+bgNbrOvESh8I z{-M`^`VN(85W|$_CL`bl!Il&3S+V?M>91bArN)w#qi)m9inb4qRn@}(xxwYnkVvf5 zIMw~L@7!<=DjvdzaIlSW8QQ2{*@Rk#l2^b!6!6}FS{4nC;A3M;!q4g(8#1Pju)LJB<=}g3eCiAQA0-ZzS SR}DTIab{N+yKy16U;O|hCcUfx diff --git a/tests/fixtures/onnx_genai_workflows/video/policies/solver_step.onnx b/tests/fixtures/onnx_genai_workflows/video/policies/solver_step.onnx deleted file mode 100644 index f53af2ab7ccfd1231e307da8cc93ec1b3e4684ea..0000000000000000000000000000000000000000 GIT binary patch literal 0 HcmV?d00001 literal 6366 zcmc&(TW=dh6!ykW;*66f85EK(w1idRA+(WaeD4dP6p$*)3*sdRjW+fq-qcyU-nB_V zNL3Jsi-bgZKou40Z{d|c#W^$Hn^TmMXu^ZFb7psEzCGWW@0=MIOLXq2zvqpD+kZ3i zhee7(yDK}Ro}>$7@oCo`NKv*Ib;Iv>z;b4lwM9r?`kf8yxiz!3NM?i39flh@uZfA!N6@xn!9&{Tpv6**zWuOanK%mgK)dubAv#1JvnrT?d~Hw^Z6I|HiL-gNo+255 zbRJwd+7sL+cu^p^lZUn`Fi)~Bs!Y46l5)`s;Er6RSTqIdJBt*8P!2NAsi1S->vTlb zCU{ZKsXBRRo5Jy}MWQ*ep_oAw_dQ?KY=ResqUPkGZ3@NHEJgKcis~tf*0WK(GqaQ5 zDHx39-&5;fb4JNneB_2*ITQ_hY0^?&YB+P++$xcRd+Y^9*Z$8IxFSYEi=YE;5Q?T< zQdfnj=@ih|S|BrF-)KamPDThaQC+7?WAT|EJRC{+l@u*|B{3Anma`D|n2NAxX%?gS zrlr_@tZZLNGs()AO64yWNn^o0qc~*DUQEbj?9$l$Qz`fa!xWuY(sVLzug2_Ytzi5T z2nET2u!?e7HKOdU!rY-rciodQ~>+%l2TCD1qO

e*vnK;9Exnnrm698g$wmz%t7#tri>QVi)Hzl~4(besDW;Na#TyAn%?w9S zc3)Dq;>{OjE8a?z)G|sE1Pri%S&G-=R*Mo}9)|*sB|>SJ;@xZo z?(T;f7oUHe^~wHmnJg9O=|YTT_pQ_s(^aw#Fkd07bU8Y95U4$@QUovsGE=k;?M$W5GWXV!^MUj$uV!AU&^c0o;c_rkbgn?BKUd#VCk!vTmp(44 zryvv#uVwbTIINZO;^Au;eQL|nRO4oDmZwXZH^?=3{c21CIw1W!yq?{Xoto}-C616d z1Y@V?W2dC*%n8?3CW?>pv@b32rM?lRC{|TY#vAY87o+a68a;FPn@QJ62moNrF79dhH$j}TsmZs30fx8OrvZ(D!IElhAb?8JJTZ2s#I PdP?+1tv7r)^z8XjBtS5vrPQbS)C1fH+VIap1;;A}@6pmn^9gw`?FT zJ90$`aoaxuLgIhnZ^22|)~-5n>efr2{rp}&fByXZ^hKOI4h~sAdh~YGbWs{K!kfwgSh|ZB_E`E7_|zp;FDt-^@u}psQZhY*__cjC5YSeM5L@U_TFhaYR11ca7z-ot)@yNgqv^s{jMy^{LH#9yK#f zah~~k;rfMOqA^e)P8`<6tYS#E;Sg@`83lnkwvsLEm2$(Oxd<6@!lJnf@i!7I#e^nE zFm?ieU^~QdX*n^OZ>$S%Lig*Abax(>#?)WTo9mEG2sc4TE2x!8iQhTO%!~7{supU~ zTBu1{xT$DCo1+rGNJ_92CDd?!V87~n_Pzn}Uvk>}rL?fGh@$#5MfGc^sH#v@A5qkh zC~C}=qVEV7q>5ryjFpkzhbLrJB$CvclAdc#`6(%>0xC!tNoiy>@Z$IyTb6M{>|tC{ zEbB^3${z$W7QC7i5~-#HeQs=wNLGr9RlV-2f^QF4n+Em-vv@ujXT)P4Y54Z#xTSak zk9C9_kjL3W!W|a?>yAr&pL!82usdPgJ%EmLE@ek9kJo=-ZSEd6VMejVChUl` z+mp1Ko&ily&C2isW>JgF@Ia*4kyE6T6arz@ynwNQ^Hnne(%qOmOBD>D>vlT%VW)hX{s!=s53U${;0^e7r)sVX48;3NM>047*NOEf3O;xfI zx-)bi?hc#E@)7k#o5xw;mwSQ7oYPn)JR}as+NaNF=qYUfccsnRc=KY_FiDy8PngF@ ac-pki&^_3?GzaZ;sKb*E@>vE_E5(0kSXTT1 diff --git a/tests/fixtures/onnx_genai_workflows/video/policies/video_decode_chunk.onnx b/tests/fixtures/onnx_genai_workflows/video/policies/video_decode_chunk.onnx deleted file mode 100644 index b493c53ff4910834340002212a11597128b83700..0000000000000000000000000000000000000000 GIT binary patch literal 0 HcmV?d00001 literal 3115 zcmcguOK;Oa5RQ{3ai$3^3sr-WP>Cv#dWf<~`Uv9Df)Jt-Ck{v;v~ulj?bclfKSC?S zr4k%KNJt0?E~q#DCcEpzv0Wlg=H%Jg^~~pQz8#M`k1EIRo*nvkel+2mg6bpk$Yp+@ zvw(E8h1h1UNlB-rP8@5O0Y9xA*jnv7AJAcn9h_SAQ!6aX?AQh+>%7d;knFAV zGK*DS{C$mz8D+i*>G7LTcgTiTj|>TEqg9rxydZJybA637ry;>aUxKa|7WF9E)asET zv2C`>a+MdhvG)H$S4V0;bjTLhK7U9!TU`60lNXo|isVD&6pW4e5IA;BlMFNv4~b*Y zDsek69N!hxNHlib(H19-JW7vexAr`5;g^)Q(-4swv4pWUAJc?!t~`BKAkLc7bw{8P zqQPBlJ{Aw|=Bl2L5>1qKdaot`dr4Y&3;&Pu>0gk~>?WS)WbiBg2McOl!oN&PmMrM;1B1Zu-Np?M14XGTnQA-+TW0qxbM5(z-TdmD<;TG5ki*ikK7Hk*$n=lq?qX z(C<`%HesIinq}4pXb2s48-I}3??BT!mbk8$XQOS{((GE;m~lZiMhEabzM2L`>zU<= zie11XW-XOUBrM@_PQ1@wM*HBesEQucauW+%DBu?&os6CvonNIL9r!!TxO6t?;oF1$ zV}OiBsh~1b3vMW{|WqdZ99UDT&8fL?hRmVe{Sl{yx}us^T6*x8`|)3l!$oJs-#D!1|>S>vQXiEBdta}jR^^^dR(d?n%TSg489(&Jn(7j zc388#V?#Mn)46A%t!m#gvTuP(1<8t=6Fy)8N!?Gl;EGe9e`L|y^L%8V!l{Mu5ma&c zk}B5&3s;6wyfXl><$6>I=F2$iF+TECT=;vZ}O6n~ta4=QNs~b?4|1_W#<04TsEC M$@;x0fw9&619_9^MgRZ+ diff --git a/tests/fixtures/onnx_genai_workflows/video/policies/video_latent_permute.onnx b/tests/fixtures/onnx_genai_workflows/video/policies/video_latent_permute.onnx deleted file mode 100644 index a21b3400fe68f2e59fda13048d9a89693ee0e66d..0000000000000000000000000000000000000000 GIT binary patch literal 0 HcmV?d00001 literal 488 zcmaiw!Ab)$6h)hMY^P6z5?yE)Ql%)k7}SM`tFgH1+La<9&Ac|jX%doDJBXY9h8us% z|Iu+u7cJuUa^B(I7jM(%)Qr^3-+VRTgFrLq0d=6UMK_%V%(|U>$F;Y{Bhw~Bc7EJm zC1<-3m0(&+>nTYaY4k~2!yRaAK1FZ2Kw2C2;m|%$5~KCp3#V*IM9#fuQlaBc$b!}% zpN2i?u2H3fd0{O2lAwDpDnrAnf^6CEVX&%v4OcOtm(ZeHBOZi=m^R0*ECN7dA-UF= z`!bcNCNh-ZtIC2L-N9-u-LS0PE{&U(1v!PI>X*p`_0P7Tj83e{l_;1nIyf!@t-rl5 r$Sw5$C7H_T!IG*Z37^eXuDC1Giv>A>-k%N}k5Po*w+$kwb@smj0F{~> diff --git a/tests/fixtures/onnx_genai_workflows/video/policies/video_latent_unscale.onnx b/tests/fixtures/onnx_genai_workflows/video/policies/video_latent_unscale.onnx deleted file mode 100644 index d72a843f8e13836a35e3b9ec1ebd3c7b50112f10..0000000000000000000000000000000000000000 GIT binary patch literal 0 HcmV?d00001 literal 824 zcmcJNK}*9h7>1WQYyDJ|;6!(jIYhxj*>(`|qRyR$zd(_YrP&%+lhR~!L%i&wcR{aS z{A2zS+pbk6h#>6t^78WL$@hKs61HwtTlDquy9KWZFAvDMlG;#dNVT;3rKNH?saEXF zI4Hu>fckw-I=-T&`0@CB@`g}r1r{Z}y3q@cY=2NYSCwNFh z=8bUv>bg~i(vMW!K{QhV*Kq<~5KNJQI+}B`Zs}shHihurs+$ zr^{t@FGeWo1j-j|M3|B$q|D&_{lgfY!S0_)B3!(crY0p0?GJ@7XgE5pjnN*If9>F% O4reAg#YO>gYnz`*K?h3! diff --git a/tests/fixtures/onnx_genai_workflows/video/transformer/model.onnx b/tests/fixtures/onnx_genai_workflows/video/transformer/model.onnx deleted file mode 100644 index b28bcff1be8b640711a07365c2d5aa90eddfd970..0000000000000000000000000000000000000000 GIT binary patch literal 0 HcmV?d00001 literal 5197 zcmd5=U2hvj6!kbxycs)*JCF};OB!ukwemycT|2SO3nla^q7q8qg3xHYp2T+Twb|X6 z(jNeV3WPuiA%qkO35mbK6aNbC%zAfc)-muh2>W5@dgtTJxp(e8W25r+?~m}02(N|i zcsgQaeMX=4-6^9Dvl=)2Ne?V*RakY5SD4@1L)WNSuj17(a)W5E0M_XSUiaNG3#YEf zNa6H8-k-j>Je>Ic%g_r3)9BC}xnW5A0~WY}*Z-6hKm6$Q0Gqj9d=46VlPL>%55PvM zGy2}dho$_8+U8c;7RYSNNPejAVmxz4SBwt)D140K?=>px!gh^t#urv#5-dP_ZbTh( z0}%3-0PI+6v01O+lKXNHYNhWBgpFtbZi(2G7=XzQBib}8;wkVntr9ot>$n(AgkyEl z^@HllR^F8@$(7r<61(zZG;yQtLgCeuYlL(m<18iID1K~_@{B$i%;=Gc`9m;vWUWA3 z7nW}kPS%8Q=8-Rgh*W0ujD_NC9-EsoK3gd5lB0)xEC*y6kDVNk9ql}RryhZB{s?rX zBXCQ}W{_jG<*~1j67*f3ai4vB#n>xG?S{FPw1p4thE<;bpv!|lW%US7KwO4w+uV=} zl5OWIQcC|WIeD%e?Oc72$cD%g^ZM8{YmhQxBlcHwX-ol2+G6T! zgjHM$QIPx9plWF+>K75@$*OjPhkR-|iq+!P`^@z(A^eZbjkpz)d1S2$Q&(`oQIyu^ za!gce``D~zCEGs!&uaFsx$tT)gjZV%FH1>T=<%P)N52+>te=skhF+V8UQ#gaj=7bz zW0&gB`k-Nzv(Q~2QO^4d(&vv8jnd~EaV!!>?2bxZmup5Ye4K_^Nd=J8(5irPbUAjO zE=Lyq3v?CU(MZ>qX=zf3a%@i2)2k769J9ioBDNgIS^-02vXE6FNai%mI~>ouy&iQy z?DJ=gr)jN#p$pIFLPp1{azPyjPKu-@BWYeb928w8(j9I<;#3B=tcc=Te?m@!1s(Yxx0Lb zKz9CO$<=|J;SXSI@BhPi{E$4j8Syj#J|Y@2kgC4NV0#D8e}?**Ihp+5_#5si7bc7i zV6X(pi^E;%zXX!KKtodz=eG2&#HW+4h}o#s;hWiwmKb;qh%Vy7a7US|Mb>Z$-V}y! zq<2oK^IbWi9FrQ}guWOTlJC&UU$mq?MFYXM&`_DABXyWhliOZb+8&DGrXsB&;FbZd z*mkPGvu)=`IO&MnO1ie2f3r-?3SVYN(PEs%vS4vrk{S|Uxi#9-8mq%cocXNGl0Ccw zY!c29@HuBoFt)2g?WRt;FXhyKI~)pA;^TcjJWh$QZVgti3QzK^0L~cRNfTxB1~+AI MTZ?Z`V20Y>U!>#9Nt}4Gr^wcNq<4W|4A;+|^ZOn+vcjE)gps_O=MMI<|LWuVY8C z6NWPv1ZXdOX<%VTNQiGCggC%~{h#;?cvb!NUbQPbAj)254r#mG?O*-gtKX~lUb(sd zh4+^1_W=tB^VwB0NC&5*BXc>Pv0Kx0Ihrk#MS6BV9fj;Fd1F2sq|2n!x|Lte=EG6a zY2V0R9`3W9^yzqUfBSi`eSC|(G)vD$i}Q3a3b&3Qu?OexJ~^7tW=|J`i}CsLXfRC| zi{x}Xx=1eur{4~@zwyoEhpc6NNgYc$d4896TW(DTO4WOC_RPlL!FlO(HD6 z>e;mQoV0;(d0Tp*vPeSFb(d`s9(&r+zLpM?ev9cp#?kui23#M$%y!N|bK5(k;giwx zpyA8&-ZwTjAhcC;X$9fO+S=Kexc)c@Zy|AwqSiqnrWsvPZQ z(z4o&GHhmI9jVrS8Mh+kHE`KPb!$X*okaOu2cW@)wo+#otYnD}bkQp4(N)m%yb8Xb zSv6mQW}#-+J#Er24~I$AZ|&>9*{XyP0WW;5$!An@dB-*>?;w-jVf)#R_MPcGUEbN+ z+WO-2An=9l7XdLE?;x^b6DiWg@~h)_M+w)FqIh9w<2q90@A$&^gFw4iW-)`bK)@qy z$~P~k3FobS{WrtJVc&f@^##XTojk6YGDfCcUvhGdjPbaTF)omdaj{Mr<2AkKZoT&@ zV}C_5mg^Q5&PdK>cy~~y_sRScQg8ioJR^F>m1jNsj#&knvsar+4ygj#RPt`Vsr2Ri zzeBN-cNqgJ7Rw|WCLkK+b%d@vRRrKP0&w3Ez=;*(;|{ESYTdqP>pHPweDP|mu}&1; zPx#dFf1w>7?^T-?-mk9X4%yqp;G(dfH;W~X+(X!XFK;>J+xDy}n-s=b%L;VUxyH`d6a?n-OmA^?--$Tc(-MGwTNXfBE#h>^eUqD@rau5yU}Myec848kRNktw1fr>JpA2qI;v zRzeW**kiP4lC@TfbhlD`*4lq%wuz?8NQ1w`;X+#BLasMl*yPm`9_dbad?pruHeoE^ zrH#QrBuWleYbBvjVQexroS+BzGKP{M_cK?b9eT`7s|O8Ie-U8!s}k?v-^ zPuaCSuoe!aI~@4*z&{~-Qh!rs3t(cYlynF+TO=*DN#CV5o5xzTknU*VQwA+us@8T< zxb3o;TD@gP)hXhbg50BwV+wMQ^~Et&YrU(@wBBLYUJ4uab`$tDJ&Goi3U%HU?Yt}B zop;=Gl&B)5Oq9?aECo@b>^Ulk66@9{1V_Khzb!&3pb<4faaQ-)}V&GR9FV0Bb?S%KLKr%WV1oHq8;g#cfoqH(H@975@o}@^8WYFO9oLvVZtmt|Z~p}{ zM0)o6gMhssG=mVTwVw%J8O+bl=QB8(eS~IrbdoLyr);~q!ybp9Tm6yzq3>MibJ-vH zvD4A`$?1|k3ftE{?OWsF^7L(RkQHahLpUq?a}6Ad{x1OB`Bmm%dg253KOB`GNIyT& zz<>4vz#DFk2=LE`k5@Cv4|B`pl)n}}`iKu2(ocR9uw2SRII#IlM8HA6$i=6y3U@kEW{~z3Ba-VjGS+ltp?toGRWCUu(juF_8I!0pWh;WWb zGV)Nzki1aQ2}5=OK712O6k5l&^#92{^os>@Ym1zkrXcTBr}i;+dJDy9Zqad#0;IRV z8FxxW2_ACSa%6i#@9LqkVDd_N1~EKfw_(04o2i*C*b1Vn$&Mp9bx=N1--gRJ5jB@J zNXd9mZB3(i#dlgaQ(CM+O3q!q$JCrflo;zK%uO)CaC$4?yjpBrqyCD~dxQqw32(3+ z;86RC=BL=@!pXW~z&tNMc2p;vLWOf?nDcUe{5n$LfEB(8nULd|@9O1N5*#AQutWtsLl_RM>(^Hl3_hfZwNv6Dznex<n6+&5jUur)Zvuk zrdGKQ6*nbI4Wr_wQ}|rS zcmzN~?7D{qV#nKY>@@WyjoH*BiZbJ|&-wh%`DaERzx(Mj97|9zEoa{ql(*{K5M-=F zUp1J+X+_C*8auq#2QTt>xTrvYz9pYBCrH} zbe*|$K`9`o*YdF!N+aE_Iu1E@V2V`d`FFe&0d!AMsu{FGCOY&bH}-YD#hFDHmFZB= zhg~c|a7mp6d6#p?Fk8Ye)wCr3EJ2J6r`M|lUrsWzXYZJmv#lp|UuwQ7t}Eyt;r2U5u?QvHc|By5bOL5^7tIc6p5n3ezP7-SEM zZG~J2mZ0qq-KBd~nZa80SEb%1-z=-8(-rFEx;g{H{<@@sEKNz(Tn0bKQZB6;{T6pS znT$AGt0*M!bJsXpCG7FE$i7Q(mXZ+XBXWnX8wS-hW!(Gsh1Imu5O&gH`zgUihvAyI z%MYnPU7}nmV+T%QNWp1y!7aBtpiG=BMqRx5z5?}j$+U2PNif@w>vM;ICj(%AD2Cj> z3#EumMh&jxd8e%^XoJy>&=ypcvz{mLQ+pjhG;PAVERi+1L}qon<=B>||4F2Bow<4g zt`MbvNbryQ+|oH@*-P=#eQ8}+hu#SH<_L=_4g2W&pBsk z&P9%405TR#4s%P**aX%rsnYLXdvQ zMs(!&zhLvj$4&#>wOz^w5q=aip?1J`1gjL-nvip<2;H)NYm|Mjpp>Q!W)juW44HWW zeamJG+Y|D(Ph&<7+{#?gI0xQ6t>p!AgK$NrK~;+KgdDo}Q(4D9^sL1~a}!`fuf~jo zJP|J%GgQipY9>>Y<^xs$F=$+$CgjLXBERtIaYL=rj~{Bgg9fL{JTfw8zLN zMw9VKc3E{o-mqZovsa9Aa_nwqG);Ky!E!0D-2Fh1HnuF7s;y3zO5DXXraT}guJc#Y zSSMa_uI7bwgA~XjIIQIfJ9=ga!iHTI2p8HxFt6w{Z$rvCQDIxh3qsy&x)!RH&N2h5 zW0}VMx_q6PY4{x4RtcZNL;S#xhH=ChPg<#;R+lDzAGYzMe;Ean^N;)N^7K~cJcV5g z;U{;yHLU`eR-^kKRuhI_!i7BDzWellpp@r9ePw0?Sn^v{Of-1BK0B*Wkwyy)d8-ZX zvgxq0p^R-S^E2b6S5qEZ`?WI>B}=wX9ZgpulyG0>J4_rO?muG^?f4nPmxkXM$A4+~*~a+4JlnkM**TWH zl!Sb(G)cqV`-JDb)N&3b?u5Fmd*G)*2_ zi@B{ujuyK`Hq=-yGHCALRs#cr=2~_o%RJBF?}e;Zl`%KKHL7{c&U6iP=3T&P`3s_$ zOLMwfC>doc6?BTJk~%lXiMHQAW?m1@{>{XKp`&?L2@;X8oQrPguL+I7#d2gqR7Cq<4^(W52#dAAb z=-T99cYEx4r~Y6%ngrW}L&x{+y^%Y0rh~mRT>9{%olRsYU$Hf8e=wf7KKlYa$p=Ct zAS43PKBRAdcH5{%Esm^RUkJNJ!TL%LvW^;L9U;h80CF5;g(m0s8mdKw?9laggFRmgva)>xKq1p)Z!+$Shrq@`bC)xs++if2CCppR)9a)bEmG|GnT>z2idF7^IlO9A8Rb+ zQAH~`YT2*0YZNCpdq{e`DyVj$9#3yEY&o~!RV-mYsC)2N26(4=ru z@NQ8~{ZWG=Z^wH}^{pmW7DA2Hq@39yHadd}_bp_tj3X^^O_s7zVIy!hh&7H_MHe|Z z8q!1*3`ejGxRo{BERKdpyr3l3QLNZvMdWQ4Xc%c5OS}>Kv|TT=UQyxpB{y2MF|Ij5 z;CjJm?Ab?-e_$-#b>MIffpqF+)@v-HQZU}A6iA7_R_aDVbzaHIvMId2unBWKtoMOgKZtOfumN znPDdVE@qNxVZ`;e?H`VVy!>fBJ%6I~0rgf`{wNI0Y-ilw77FOh2)`PQoeV5dhsbrk zNU{5!P}fbVuA3bu&SQydc8%q@9hOUwAqDGQ)s?-B&3bvREV8itDwAom&7@mw8?!4L zvB#8e$Y^O(V$&4vs0wA%O32ZS7NMdRnO(?%x)f|8nTslXm6{p8O3e!(`qVt1=u;{a zeM*Za`oLLJok1yHnvDGsC#|EQFA@W4&5Rg)A&QTLr_h!Jr!_M;KXEvmm(Z4^g;qI7 zku0X_NZOu3c>N(1U4mD4T@P-b19xagcR@KPEbaU+PFS_BWF!jBmnd+zT$n^rab=DR zDjpZJt9J@6D|1z4iC0xJF1Sb91 z^^#+@NV=-p3tl>6F&|#lSgl-K7n4gCm6OMkyDdTTS#2ek;B91vlJ74mm)dwen@c>Y z^0Mii(`mNZgs%?W0lfZBZMfYD*h?WWBj<4Rz@0V>c7gM{@M5;3$6_=yMt1anqn~@= zJ^8`8jyClQei2pheBTKMduSCeM_>3ZyuWz^_uh5&#RM;-D$H6;X7yiF?^5tB^a_0P zQt}Bh@CNrlyzJ(!&A8czcS(`P;AcPuTb_4EgKwZM*mpDCcMV@w{vg^i}J$w z*%#U05LjdrBc#LMT!!T}X#^flgpqYoF{A^fePJMXhAAy&NN;a2bU>^7YfL|3f5rMm^$PIHaq)u%mQ|Gt?xi{h`05)+E zJj*6>3npg^=_ac(xAP4n21G#OCw?#O*K@9K<8Rp7eAD>iR2+~uh5NvGP~DW>$O)Uu z7B$<6MU@3R!YnQHcY1M=Oj1_sB?E$gXjAZyybFl<6q}}cfSQR1vd#lR8&73+;(-+j zx%};9*C<#-=V*#kfylNE1fDntYBIciqR=~*3Mjd7=fZcwIPfSiytkQGDCx<%OFbj*ZLZHA5R z0qpq~=M9hZsjvgrmkoPFg&+f$1vIb+S2*+GW48*{ovg415w_eLL33nl&>H)Jp)DbTnbau_5_X$hYnVBQ2$lc&v}oQtFrm(!%*`$wXQYF^m>2b2gPE>9eFT zj+cZH`YSEKUy_b$-*r@b!BIO(=CPx+>C*!B9VJ}UhJx=2SNA2?5xND$jCSr`awq*u z{-)q90!mzQ2bc5}FD6)U>Nt_@U)I*tI3?n6*z)^W2PyddILz$9NAOs@%#vw}>}qb*s;+&E8X3 z=4Ryv!}|*z6$6;6Y!=kcLo1<-X5wEZ(c|W;GQ<6q-Q#w#0miIwEUQ-H-)>^lHeQFU z0%Q!Mc)X?)$FI>7*pD98#yR8ai3_FM@V$|1Uy+cIfZ7sgK7GDMr*QD!OyWBD!7M8( h(x8*6J924ob<(^>kD>a{8m!-Ew$!iZ>ns>MmA^m2fFb|@ diff --git a/tests/fixtures/onnx_genai_workflows/vlm/policies/generated_length_update.onnx b/tests/fixtures/onnx_genai_workflows/vlm/policies/generated_length_update.onnx deleted file mode 100644 index af3bc3fca3b73acdb986806a33ddcf35d9032ffc..0000000000000000000000000000000000000000 GIT binary patch literal 0 HcmV?d00001 literal 930 zcmcIjO-}+b5OsxRVU%c^7;z6qiP75z^yW!Ca`0-RF->X9ZgpulyG0>J4_rO?muG^?f4nPmxkXM$A4+~*~a+4JlnkM**TWH zl!Sb(G)cqV`-JDb)N&3b?u5Fmd*G)*2_ zi@B{ujuyK`Hq=-yGHCALRs#cr=2~_o%RJBF?}e;Zl`%KKHL7{c&U6iP=3T&P`3s_$ zOLMwfC>doc6?BTJk~%lXiMHQAW?m1@{>{XKp`&?L2@;XVSA`N5$W@H3tkY;kI1!>+E8gov%LAWm#V`_v+67_8YRe% zs2_16q$z$Z7kWeKb^)@I4!I5}<2co>!9g(Sxk}2hW}yg-%Y3Re=?fmxko6}xeRtok zKzWUtWf09&!1auPR|K}NdctT}vBbGJ&XDA(T)3g8&qFd7ZdL8%*vNrJ$3hqXQ4nv> zJ&YP~xv;rCTPouds5=NBK?ZMl)UZAPb~JyCxMUpi7)FXwFY;+P@lN6AkMKV1#g)6? zl`Hha&RREYu$&A|O;H^VlhL|8F5^}kp`;rqU$6;bN}7-|gVRsXQ*;TnKS?4wxU`^> g5SK<{;fvV#qBTXwQ28~3ce|XK`4q1#NNp8A0X>A#xc~qF diff --git a/tests/fixtures/onnx_genai_workflows/vlm/policies/termination.onnx b/tests/fixtures/onnx_genai_workflows/vlm/policies/termination.onnx deleted file mode 100644 index 3183438a88f3ecce329e4afc84272274af35cdc8..0000000000000000000000000000000000000000 GIT binary patch literal 0 HcmV?d00001 literal 5880 zcmc&&-EQMV6ppiLlT6a4on^bhETwCeMT%$@J4xH@)v_yu6akf3aZ`jwt~1*@P3)}g zgf6f*6&D}`NU%4c6)Fz^kHTB<2HeA$@$cBFxN^yj=bRJIeE!aypP5PxFHh`abLQUv zYX`lQaD7f5*_P|+mPguhBN|)wkdn4i2^B3Fp}9VuQDR!2`{1AOCtXVA!3wHc`h>bu z-JrNI*hZVvvy)xhvd&$@F{j?HG1gs|oSM|p9pm&FE`IXqU<=8qf3Z1;YuHojvVWkT zCA93>XVh|$028mw$;X!abVljtl&JDrEO6MWQjIiuVf#TsA_`QdwoA;Ri>q_;?$en* zCYoIL2OO2AEO99>L_dno?4+I9NjP&2RU>CM3zt$cFHpZo0xyqG^(iG?x$X}*FI{Dc zOD%+MP;1yM0-RPIEy132ye90g+`zq%fCds^LdaloK-3{=;~J=g?bVZ4!y$*;Q_6vy z7iTF)?xh{Mcau2EguK^x}f=H4m_L)!Lr;x+dwp`}}?ozd_ zl$ey4OTXbR4YDcoZh!$-=Hwk~NK{p>v8NA4Rc?Tg7nW}&6s^kWc%-vz0Xd!zd<2?O zVp6M$ik{sp=a{>G=gb5cgx$0X(MN|fgFc>3h}x0YW61~EQL2%u)j*XqN~c3};^tZi zUPx@4hus8GyYgBj_^`THrL{ck)2|Xr@x?ZNgiP(S2I)PA>K=87+LIgM*vAaTAW-DB z>(adKWqA8v%E5QA{Y>CQ_e|RoDcj@0IfUTchv0;#Pt`sI=P4lXJR6?c8vr|rGY%P!Q?9p6SgwzCGnlMF#Nr|{wiiu}gmRDC$3W}+!)^sNyanLFochTA;*Z=+o3ukoiJ@$U*faIJB(`1736#NYBHopcVede5+EQKNrD6ykaQvaT|{^NThQGTY6V^F?$jQe z#xugQpaWrw|15r?IQEz-_qUWpZX#2M4ow421j^fB;P;dKFa8L0!9^fbhC|Rs4Jp&d z0$3(OyUNTdO)zJ0cM#4SieUA#3Nr*F(DAHiS|?09h0}{;f5reqhUE<&lU;6es7swW z{~Fd_FdcuwpA(sfsgojD!+OC?(HQjf$rEaLOq-xhd43@c(Dr}Q4p*pHVOPgH;2PcB dO-!*9G5fNe&7;#>9{6+eFJ z9f6*STIRrk<$5$&M!=LGsZ@<0$Sa93}x?dbfjp{^Dl?Sd8(XG$^SqE+~%6!U`m zt-v1R-c;ZTnPDd6l?0QBa03<9Lj2aJM1IDn`p(LXuuqL@;^)QqQ-S><7JKiJf#IKq8ulp`-M+8+uu z;+9i4z-$npn(uzbiK92)4Ja2(%ww8ohPqnF3+PgTJJ&!C-G-#;Gr;}w_`|aEA*Mx@LpshFb zPi}uVXe03c)^R`ps!UpfqN?-{)`GXpprwv7A40n2=pXFnA*36z0b?FR<}Gi8jQ6FO ztrRy+IKG=vnF^F7ePJxmOzBOOFpV$wrDsWniriI%!a_NL<5=)JBgzvg>v*oitU0z^ zE`0e1K*g`$7RpO#|2I<320BQh;v^k1oiT?Ie}3LwC=cLhYYsXbQcFa_|EdIq-PT_U CB^YA> diff --git a/tests/fixtures/onnx_genai_workflows/vlm/policies/token_sampler.onnx b/tests/fixtures/onnx_genai_workflows/vlm/policies/token_sampler.onnx deleted file mode 100644 index 1a5d5e311eae66f17403b7c4307e7609913be913..0000000000000000000000000000000000000000 GIT binary patch literal 0 HcmV?d00001 literal 58720 zcmd5_ON?a6RqgCwb!UGkJ>xFd_|r~X1`2zW@Bf~Kggs_347Nv_HUf=hrLwDXx@)F8 zvs{_g9)l2QkPYH%jAg)l+_L#vd<3#3SRf%5ED#b)#EKOXOZF^=6Y*aB?v2;wO+;32 zI^U<>yYJq(C+>|KC*snT(Zz?$y9dXs5C7Z0o*RCCcz%Dmn2)v|O+U3 zPcM(A)y`Aak!)$zcxX1XZXh8Ml!O#_J zbh(4`B;C1c@H~kJ=x9F5{c=E2+qophXqj!aEU3}i=F>J6A(Z@;>E;74H_ayxDhX{q z1#3PwggedWWb7a#=eE1?+&0a%+opx>wrLS92K&x{8BF7HTDaa@GVd)d@7*LWr|2@=fGtIrq0YFO z<=Z0~xH1N=9Rt_V4E)_E8Mt;0+%N+-mVq}4&uQtxbK}BuHev9bO*R0}*+ju}nh|); z5)RK<5*MD+jCgyND0t2=LZsv6ER7j2XBqY0tan`KYCP}FGVQ%tZhLPQZ)k|YbCzg$ z&T_|lGh`h)e36}hhO9%K*=X3Dq32Lw+{PstdJc8Qh37n?flFiH$}w=YIe4}WT)75D z@}bPxX%WeXI8lLma@!lLci1O9edvoL< z>dZ#rIcMWldGC1Ddm|dyj92B(cvT*6yoyE7wt?MYtDFs6<@T`EM&UWn9eB|OBSY87dmAvOsg)e zXrm}%MjHjo6^hR*l}bzi55qi;)!M1a^eJ#C5qB2Tx``>qA0D-Y&QKXYA3!Z z6F0snQgO(!ZCwYbfomp$ywKkm50g=Y^sjWYgOX(jrJWsY5|>L>)Gm9A+A$4mirS@9)Gp%{ zwORUS8`!OFSFE;OotoNq>DIO@QroV0ZM)h;ZM%~Bl?abZTHCHXwe2c$ZM&qk?MhF3 zs_fc!6}7fqvD$XkTif1_2hWwMZC6fhyV_uFyGk`YS8i>)Vzuq+wA8jMR@<(6Yun!+ z(ZHs*T{*SwYJ;_H+rVyZyJEHN>eSS>E4Q{?k=l06YuoiEYTLCQD5z*{yY|$!>&Uh3 ziq^JkUE8kh+IAhawq3K@cHLXsekPvx)~2>yJGJe4gSGA242somZM$Z*?fSITwre(C zRriip{a{1`oAIjJ8Lz72jaRYg**35{Y*n*itJ)s6+Gz2lc5B-;scqN1wq0+cwq5Jm zc1>&BjiPFvXL`gz=3oSu}RZq6u5FFvXKbr>u=+#);&$0W(hgFXv!2 zWz6D9!*YeDH&=*QJZVg>&^Wn56E9aV#gnGcSzzNP2@OjU8aqjdI!O1mffz$q)e*4a zkOk-?*|czK6XlcNmW%`~3Zxhf6%TBZ>|C>GX_F*c1?6vbM+5U(@^poWj#)Hm3lZ&6 zr$iK|H&jIVQtQU1j4HCstkx89Van+^*lndCcf zgF+P*x6^VSqm=SnrV*4wCYC7k30zJl7?9EcjiaCl&HT9Je%zZdge0RdKQ1jlu3SHk zmoxnRfE&*8XEkw2xkx*4$?<19Z`dqiBQel0*||Y)!8;@?g;}~WmIlk(M(iYB&HEpo zXz4`NC^}1nk}W9xwxyHJmB0T?vGQ~~?2;QTErD+EwG}r{LpNC4;rWB;V@$#$pgWPl zbBBWzyOYd=-AR-(_*VuZD(mE;7?8;PLtadfe;Dvi!j{4*cqavO>7=(99+Ij;T{Xa0GQl}`r>OyN zaKT+CK^jBiv2_swK|yp;y^aZ8#is5GoC z@jN5z%`-xlY*<^e@{G*QGqQMj28Q(&W@e>{MlzCMfV@qk5pcK7xc)}R0YTlSZ~+2$ z4_t&KzEBf`WU1i-BySro!0?U*7k}qBG3eb+#}1A+GciizDh7?rf1zOrXu(8U5{7{O z&B726z&9`e2rjU8sfvanpn~%-1Z?mKYo}T{WLD82AUBAB=B5zR8$?zLE=t^ji(5=Wj zFbkc6H)!b8rTcm4A(=GPr9nj3APFYAs+}6TsKCoQXK_)O@uG^Buoyi4K zI_*xtOz%nW&qu7`fS9gf8?1C2WI#!emEK!5oXSc9DpNH9E#0aa$Qj;_T9CF%I zfkJM^h&DPJ2nIPdF?i%Xa1pbH1Fokrapsyh-ol1f4aYHY=7yc1nllsIi{Nf{Dhb%h zPA~!vx*Lpuh0cNz@X$9f31$a=?nWaxp0UkVl2fV>pcflLHbq3yGuWtan8&iO< zSpgn2b`IX4v&RMRnib$dZ&y7I`0W;W2PU}F!67i*nM?N;;A1jyz;V~$4U)U#(wi*6 z+jMBq-CZZ)1^81#cuzUu9jE%hJ0b%IQfN!3F5w-A8w795(h$V4bjpeEIAwLyn~l5) z^BtnO;GKGy?>M&s@WxwpsgL@Olaax@p#^w|pujtIxbHZM8F{A@R)EK;K9`Qj zz=0H&OFISlbh1J2I}E4COCShwX;y$wC#R(VpHA427*6$BIwAuHQW#6S-1l^o+;_*) zg!`Uy?mJGM7Vt)+b0TkuM!?&{eNQ*SeaBlBcmw&!uu1CUzT;#J@Ww{nU2qM#`syXA z!+gh4Q0Qhn9H>H9mozsh1tH$#0yqjP6UhbAgm))NZ+CJb-arlnvp~k}PLg=tiK*eh zFs2_#bk!tH?5YWl;#!Y)5lL?saV4hptS$n~bt{ug!CH?sIAey?;p)bEdSCK(W^LN- zOK!0r#q^2$GJip`1jqV)Kaj}$ru7oS;5IO!z+(w502;4 zS05~A>u1iKIrGZn!64KQ?_`zhM36iQ22xD#Du_g+guPK7)L1>N?~klTx~M#ToEQipR+8r55G{rp=28`U4U zi}XR1;l&5bdk5>4NFu}%eV;gZeGVInT8KL83@K_M-KkL4WXj`xf#8@c4nN3Ujz|7u z^`#=9@-QFr%6k6raDFsfA0N#{K&B*t#^iQUpPaohO+hWcDt@$j1F8HRjzX#IkE8>1 z?;Q^JmW$PTn(R;IjH#botUh-Dg8Pjkr;iD zt}CeI1~LSdJRoGyMh4C}g?(qAani^x{r3;OJ76PY(lJtb>X6EbbZ5@|!Vd?O?i+*Q zg@*@=>7m>B!RL0Ip(C6>I5>ov1>s!$>e`$joZ~Q5bAyDUgX&)WB!?(H4pByMh`%vi z7B+S!9{ciq|9EfSx^uNNGKZb^s_|ugZuiFU((Cj2;r_wH)pp1}?Mrg8INNn)$^H|| zdoV9@1@WVk54i%4LKzp?Con`S0B^5E_umPk2Q$G-x3W*r(nC%opbjfI=ek>HwIVp@ z>J;s>5JjE|=&*n6cUWv7{uJ~!U#(s`df{`&vj@`@bo8g~6P@VE_=-mwaDjjy{GLh1GQ zsORSC&?~9fAQ;O)rr@JrQ7cj6V0`waqfoLvmaFo%FV#Uj5Rn9;qbrZ5KXre8gct@a z^y}@%*-Fv!i?SRSPS-PY=(} z-Z)rozoq_#G(T&XJTn2Qd;?(fb*XY~cPng21l-(FDKJRpH`PI>=VH9ooe9+*nAUi4 zqJJavkmme0843c`*HxAa-Qc3bxn7y8BZ)1f}ioyUN)rSQ^*wY-K^THd4vjw775Y2CNzgshc zYJRaj5lXynw5mmA%ZhZg3fyvLMUczGTMF)JQ zgPU6~UJ?nh`1g2lK9nkdQ;rcSE1IeSxa4*z3OsVVYOobBUPjd5*Q98TI%G&p?FH;| z{K1{Q*%}$=3lHWG=Zm#yljY7fJ=v+g9N!>UcCU2dL&$E1_J7#$oDPo&in(ePyU|jV zKbI;eDiIoV`@PalS;VI3Zps1H`sB$iB1T(-@rmF?!}` z1P!|#vl5lJ*W(wR?RfU+llm zf@o0S!w1nAb#@%U?=)h3=Fah5=n~}`@uPDOlxuJl$~V3$$BVl6lIXsYM9VfV6ufvV zx);3o=+Ql>?@n|N-h0I89{Z&GVp34vd!lA z_q0C*37+ZD&X8Dt7_@-=Eiz093aA(+G(o*PaZHlv-HBsj$M(<@$Ns|2kib`Vg>-Wy z5b8~qf_>VZ90`>AusIUeJ+a?TJMeaPgA5(T{_N(k?H<9R?`DADmS{6TXdC4W`qW|2 zlQtT7&eYeeX(65BTsH@=UKNQztG9|%Vi)(Y#VNzYOLa>F;LNl90QK@jY=J8f|MgncX@rja6*X;RQj-? z1Q}e|2g90i@FiGD>J5*>_~wTD+~AxBpi z^EcMhPGn-mz_ZLoJFEG8f4)COF=u{nd33Oc>6ok8!@~#jqum`*6nowp<9YFt>PXnN z#1zk7`93x2e)O*4vzG>=Tf?oZKz}0eZcJ~-lNxtyPyB$u$=*I9G?!K(= zer3Cx9wo1Ad9ECK|3*#{=T~kG&x?=c58E-O5F$7)ddKkMY;S$=XfFTkvMmFix4&HQ z3wNd1a7gy1yYCq~wq`fHw(UgZ!{H@JbAEXC(!jZbLa=oO;y;>xYPlA)3y5#a_?o;X zA%tIh*5b~JKfGpr+-q-_qiAYaIU8LZo`)!BUK{B@%H_FhIOpD5c7jKE=OvfeQg{u#Hn~42HsyIzjEv@`o`KVR{ z;%m+0u@)bFH+j6#h4px&40*R^CV~Vjgc0u`kEwl9HtyZ}F<@h=$81~57IR(OuN5yC zi?GKs1NOsZ-=aSh>_2wVl@JU(u19Rossy<$v~DCsF!*G}rL8JNgENXtTQvyr$@4so zz-~nN^vu%~s1CphVq$3>WF7^%I>?NZ7@0vv@#_FP53-k$=e6y`)j?+3h^vFl<0VuF z9yL)NU~WDWOX~o8v$2b=%-A4}w}agHjw^G1q61-zEzk3e1Ci&Q048V$*zSpa0@Hy< zPg)0`*l}Arn2PpQkcepZFK48Xj-aRs(2+UaCn(Gv$tDfh1Jc;N4MH2Q2R0k*qv3s-* zu-GuW=;{Ck4K6bl()G)XQTdae=kqPBn%MI^k8e>Oc+^C7(0KHub$}s8CbPMCqJpq4 zC09^m6l5Od(T%xyV&lPbgr4{GHJ-+!6&cW&=nbphsRur-M<5a;#L5&((cPABVhuOy zleVJKK_nB8PrZ*W8;r%v*n{3F0?0P(^r(qSZrYG7xoJdHa!)Iwk2DZ%JrXf+k21uX zNGB-+p)k9LDI@V{%qS!AsESCQ7@)Q#H%LvCfk2r_MIULf+J2+~Yw}10&Fqo1WgvWJ z_b_Gn0%@ZRUnET=w?b)4ax0dml6%~aKGOJ|{Yc|@^pPIVqss6F)Jz$^h}tN_7gAeS zJ*|p-R~KB12-*5BMAoXsTTk<7$+~)SgF<}mff{fx)w_reeStI)AK|p;kzEKaf-3q* zk4C6jd?7V63!-Y)kc~1h2a#Q+l6xBzV(UI-cp4OWB;sn1GO%Z7&m+Ztm?xnZ;%ase zQwBDIWH%XQ6dqMk$-Pis^g_8W*rt#4Lb(-gla}F&w`t4p1>8&-zKGi>!xwTB$*q{% za>=q6%6(BceWcg(thk#f!xwl{W%wd*rVL-`ZIt1Qy{)UBRz<$6i^2(dVQ!yl@z&Ek z_+(u@IX1wiD<1Dc?5*BKbm)t?iTH@SeUCJ$EPbR$Bh)Ou*qfOJTQ)qwVG)P3t5kAt zgTf{PJ<9Nf-sF)8z&*;qYW01O^dJ*e1_E&QB}^Gu0-fDtT85~)QjnyQd!Za<#XiYT zgmSSKHCKjM{n`IWxw0~qk|+a99P=-s$`DHmv(MGa5UXpltFAJ{lDj_1`=A`ljGp{R zpXXs&*FI%@4^cUm4CSw7%6MW_j^+FKBaEsLke{R~EGOwvl{JKIJ?&&IC0;FABiSyP zh&B9}fbY^7AQp#{m#^3b4Pt6U;zIs`6FZ8P%@~gTKmV?srw<_t@KKdTSEd652&v91S@!xk_e2!Xt$NRzZ z4v~WiLrPlOdN_m)Rr`vFmF6LlLz z^(Wl=L2Ftc!uu1<`wo|KWET^uN=)}4$J*5dRV<&u)q!%Zp@y~ngk4Let;!1dO$$fF zFPzQ!yY_0!+vQhx%tU6nvVr)T0g9$)hG$Uz#KZ|hVbo<)agBT_@&E;Ba`;e@^;4t( z>H_@}gMqiir)z@x=+f{aKH)6h@*d)yi2koa!Yw=#RbAX*+XAEP#ujH==?!68UlY>P zriTDPzmOvnL1J7p;~ciFN(gOmSU&%eaR>V*fq) z6l4bKQ{-C>8Osru{#Dz|JQ&-h1mcQ?u^5RLj`GB=00Z-0+`A>Pv!kI9nRIIr>;N>DT96!wKU1xR z{#o0|jE;3bj9J;@f#Qnvi>`$sDG-lSElfJInG~HOGZ4-^z~7$fCTTeRewP-sPld{5 zAC9cR{Ud3OZa%K_IwE~HA}bZOK zg;*f9`2a9%MkT_n+Ye)0d{OHkKTznwcwH0T( zUpObX|sDe&Vb>D@|N-j&Xhd>Z|j;wYYKw<^0^6gXf3uKhX};gZhL$KBC;5n(34|OxnK9asU_`vXX??c)*u104+^V#Ro!Qb5ee|UGf>Hq)$ diff --git a/tests/fixtures/onnx_genai_workflows/vlm/policies/token_state_update.onnx b/tests/fixtures/onnx_genai_workflows/vlm/policies/token_state_update.onnx deleted file mode 100644 index 7dfd0e7f3b7409692a79ae731043950283511932..0000000000000000000000000000000000000000 GIT binary patch literal 0 HcmV?d00001 literal 1229 zcmcIk&2G~`5U$%au?Hf;a-cP$3Rc9SqE-@aJ<%R>;D(Tb(8^g)6Q}8}+1;cqh)W)U z6VHOg)9@g$yLOt^6mg0#GxLpi=G*yM{{g9Ah*36I$3FtNgjy;%lE#9*6`FQKu!4&e zX*b$2`k)B`PcBg92}80q=)ggKK8^+FSBl9j*D;$VO3_J%GLdX@O}0OL9JC><{AFWM z#YB$E{sBu5>Iu`?0%4Wgk~rZh?S){4Bhia?jJ_!n-#we$G0wE2}&%g^*Y*h)`mzC zqQyCKsi6 zggk<#ACTbY^zdmpw^hgJ8Mw0C8*sXJr^^7op;WU>rr~-Dt04=Y8xx7K; zo=fkS2QP}Hi^s^#rsz!YI2W^wT~jl7TAE!TKfN!J6f=yD+tKYBN3YsZ*>vk#sRk{2 zV8K!|m@99_#}>^|QYkKNfQ%RGW^Lts5*5k?T4&+1XsH>@z6%y{(%5#|tnH++@$3wl yLmb!W>(V=c&VO%7vy|-H=i|s!BFSQR_^6=&roZ%FLi?XFij?jFIjzA~Ywr(oA9$t! diff --git a/tests/fixtures/onnx_genai_workflows/vlm/policies/token_to_slot.onnx b/tests/fixtures/onnx_genai_workflows/vlm/policies/token_to_slot.onnx deleted file mode 100644 index eb8e4b6599f5d6095617f9fe0d294eb1d098d216..0000000000000000000000000000000000000000 GIT binary patch literal 0 HcmV?d00001 literal 438 zcmaiwK}*9h7>3un+4PfQjh9(LrVIrSWxIOP9eekpNXW8mjq8%yq_9rB^zN_lhicoE zp@NuW-iMdxd(i>jU#Y3c_07i~yj##SYRM%4Q&Jl;o{O`5qznA(Y=X}G7R5MOS z7w+yS^1zjJ#dS&<$DL>hXX$bgDk;~RWg;~pOQ_Z)7CfUFiwkT&K1Blveo=J?(M+XW z*97(&8DbQv6|Mth`Of5HNf(^F^;z{hY@1>Pkfwhl0?k#NlR@@LYl8d bB4}3Jj!Wwj27h+&`J6M;Ofu<0$M1gu>uQY1 diff --git a/tests/fixtures/onnx_genai_workflows/vlm/vision_encoder/model.onnx b/tests/fixtures/onnx_genai_workflows/vlm/vision_encoder/model.onnx deleted file mode 100644 index a6f8be4e8d580dd1099bdbc762a6c31ec8568225..0000000000000000000000000000000000000000 GIT binary patch literal 0 HcmV?d00001 literal 652 zcmcJMO-}+b5QZtNQpV-OCMISNx*?oA5b)@UXuObk@rW^#+8v;?OOA`?f1y*B)& zhIObJIVYP)MWQt9z)|ur?#9NftV(r~b=6E-ixW+$Ox5H`ly7c_Ap`~0x(8YnCyu}= z;FSq~i5Ic4nKT*pgXZR(L-gDHCvkEBkh7yUnH{i!HwZsibz5Zb3%()HxG1D02@g{H-<$fR9rh zVK#Xb4e)tYaCgbvs@!#;@EIYX<;wkUh3EEg8Ws)7`>;(PHVq0E0dV)uy1H`iL8*24 E2@X5KasU7T diff --git a/tests/fixtures/onnx_genai_workflows/vlm/vision_encoder/model.onnx.data b/tests/fixtures/onnx_genai_workflows/vlm/vision_encoder/model.onnx.data deleted file mode 100644 index e69de29bb..000000000 diff --git a/tests/generate_onnx_genai_validation_packages.py b/tests/generate_onnx_genai_validation_packages.py index bbaff149d..fbe7d52b6 100644 --- a/tests/generate_onnx_genai_validation_packages.py +++ b/tests/generate_onnx_genai_validation_packages.py @@ -944,10 +944,15 @@ def _write_adapter_metadata(package: ModelPackage, directory: Path) -> None: yaml.safe_dump(metadata, handle, sort_keys=False) -def main() -> None: - parser = argparse.ArgumentParser() - parser.add_argument("output", type=Path) - args = parser.parse_args() +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 = { @@ -1035,7 +1040,16 @@ def main() -> None: """# ONNX GenAI workflow conformance fixtures Generated by `tests/generate_onnx_genai_validation_packages.py` for semantic -validation and runtime conformance against `justinchuby/onnx-genai@2af34dca`. +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 @@ -1046,6 +1060,13 @@ def main() -> None: """, 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__": diff --git a/tests/onnx_genai_workflow_conformance.rs b/tests/onnx_genai_workflow_conformance.rs index ce5413c16..cf9930895 100644 --- a/tests/onnx_genai_workflow_conformance.rs +++ b/tests/onnx_genai_workflow_conformance.rs @@ -14,12 +14,19 @@ 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) - .unwrap_or_else(|| { - PathBuf::from(env!("CARGO_MANIFEST_DIR")) - .join("../../tests/fixtures/onnx_genai_workflows") - }); + .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)) } From 43ef64a9df76d2ceac78cfd15c521812f1eea688 Mon Sep 17 00:00:00 2001 From: justinchuby Date: Fri, 21 Aug 2026 05:59:31 +0000 Subject: [PATCH 142/151] Restore revision threading lost to the rebase onto GLM-ASR and LFM2.5-VL MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Rebasing 141 commits onto `79b48bc0` conflicted in the five files where this branch's workflow migration overlaps the two new model integrations. The conflicts were resolved in favour of this branch, which is right for the metadata architecture and wrong for everything else main had just added, so the endpoint was then checked against main hunk by hunk rather than trusted. Every `def`/`class` main added was confirmed present, which surfaced what the resolution had silently dropped. `--revision` was threaded through the asset writers by #509 and did not survive: `_write_clip_tokenizer`, `_write_hf_tokenizer` and `_copy_runtime_assets` lost the keyword from their signatures while their bodies still referenced it, and `write_onnx_genai_config` lost the parameter entirely. The last one is the dangerous shape — it has `**kwargs`, so a pinned revision would have been absorbed and ignored rather than raising, and the package would still build and still validate while its tokenizer came from whatever the branch tip happened to be. That is precisely the failure a pin exists to prevent. Restored at all nine sites, plus `_write_text_runtime_assets`, which is this branch's own helper and needed the same parameter to pass it on. Also restored: audio-processor emission for speech-language packages, which #509 added to the multimodal dispatch and this branch's rewrite of that branch did not carry. Expressed as `_has_audio_encoder` rather than main's inline `"audio_encoder" in pkg` so a package object without `keys()` is a False rather than a TypeError. Two of main's tests could not be taken verbatim. `test_revision_is_forwarded_to_detection_and_build` was overwritten wholesale by a test of this branch's and is restored byte-identically alongside it. `test_dispatch_audio_only_multimodal_pipeline` asserts on `pipeline.models`, the legacy composite ABI this branch replaces, so its still-relevant half — that the revision reaches the feature extractor — is covered by a new test against the workflow instead. `test_runtime_onnx_genai_routes_vlm_through_workflow_emitter` pins the writer call exactly and now expects the threaded `revision`. Verified: 4580 passed (up from 4541; the increase is main's new tests now running), 11/11 `validate_metadata` and 11/11 runtime conformance against ONNX GenAI `6e2ddc78`, lintrunner clean. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: justinchuby --- src/mobius/__main__.py | 8 ++- .../integrations/onnx_genai/auto_export.py | 62 +++++++++++++++---- .../onnx_genai/auto_export_test.py | 32 ++++++++++ .../onnx_genai/inference_metadata.py | 5 +- tests/cli_test.py | 30 +++++++++ 5 files changed, 122 insertions(+), 15 deletions(-) diff --git a/src/mobius/__main__.py b/src/mobius/__main__.py index f6f609e5e..47e1482af 100644 --- a/src/mobius/__main__.py +++ b/src/mobius/__main__.py @@ -451,7 +451,13 @@ def _save_package( config = getattr(pkg, "config", None) source = getattr(args, "config", None) or getattr(args, "model", None) try: - artifacts = write_onnx_genai_config(pkg, output_dir, config=config, source=source) + artifacts = write_onnx_genai_config( + pkg, + output_dir, + config=config, + source=source, + revision=getattr(args, "revision", None), + ) except ValueError as error: raise SystemExit(f"Error: {error}") from error for name, path in artifacts.items(): diff --git a/src/mobius/integrations/onnx_genai/auto_export.py b/src/mobius/integrations/onnx_genai/auto_export.py index a4491f015..8dd023a27 100644 --- a/src/mobius/integrations/onnx_genai/auto_export.py +++ b/src/mobius/integrations/onnx_genai/auto_export.py @@ -195,7 +195,12 @@ def _looks_like_image_edit(pkg: Any) -> bool: ) -def _write_clip_tokenizer(output_dir: str, source: str | None) -> str | None: +def _write_clip_tokenizer( + output_dir: str, + source: str | None, + *, + revision: str | None = None, +) -> str | None: """Emit ``tokenizer.json`` for a text-conditioned diffusion package. Classic Stable Diffusion conditions on a CLIP text encoder, and the @@ -253,7 +258,12 @@ def _write_clip_tokenizer(output_dir: str, source: str | None) -> str | None: return path -def _write_text_runtime_assets(output_dir: str, source: str | None) -> dict[str, str]: +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 @@ -269,15 +279,22 @@ def _write_text_runtime_assets(output_dir: str, source: str | None) -> dict[str, Returns: A mapping of asset stem to written path for every asset materialized. """ - artifacts = _copy_runtime_assets(output_dir, source, _TEXT_RUNTIME_ASSET_NAMES) + 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) + 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) -> str | None: +def _write_hf_tokenizer( + output_dir: str, + source: str | None, + *, + revision: str | None = None, +) -> str | None: """Emit ``tokenizer.json`` for a text-producing package from its HF source. Decoder-LM, multimodal (VLM / speech-language ASR), and Whisper-style ASR @@ -486,6 +503,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). @@ -663,6 +688,7 @@ def write_onnx_genai_config( scheduler: SchedulerConfig | None = None, 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, @@ -704,7 +730,7 @@ def write_onnx_genai_config( num_inference_steps=num_inference_steps, ) artifacts = {"inference_metadata": path} - artifacts.update(_write_text_runtime_assets(output_dir, source)) + artifacts.update(_write_text_runtime_assets(output_dir, source, revision=revision)) return artifacts if _looks_like_diffusion(pkg): @@ -735,7 +761,7 @@ def write_onnx_genai_config( guidance_scale=1.0 if guidance_scale is None else guidance_scale, ) artifacts = {"inference_metadata": path} - tokenizer_path = _write_hf_tokenizer(output_dir, source) + tokenizer_path = _write_hf_tokenizer(output_dir, source, revision=revision) if tokenizer_path is not None: artifacts["tokenizer"] = tokenizer_path return artifacts @@ -813,7 +839,7 @@ def write_onnx_genai_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) + tokenizer_path = _write_hf_tokenizer(output_dir, source, revision=revision) if tokenizer_path is not None: artifacts["tokenizer"] = tokenizer_path return artifacts @@ -854,11 +880,21 @@ def write_onnx_genai_config( # 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)) + artifacts.update(_copy_runtime_assets(output_dir, source, revision=revision)) if "tokenizer" not in artifacts: - tokenizer_path = _write_hf_tokenizer(output_dir, source) + 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 + ) + if audio_processor_path is not None: + artifacts["audio_processor"] = audio_processor_path return artifacts if _looks_like_speech_to_text(pkg): @@ -867,7 +903,7 @@ def write_onnx_genai_config( "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) + audio_processor_path = _write_hf_audio_processor(output_dir, source, revision=revision) path = write_speech_to_text_workflow_metadata( pkg, output_dir, @@ -879,7 +915,7 @@ def write_onnx_genai_config( artifacts = {"inference_metadata": path} # 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)) + 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 @@ -932,5 +968,5 @@ def write_onnx_genai_config( sampler=str(getattr(resolved_config, "workflow_sampler", "greedy")), ) artifacts = {"inference_metadata": path} - artifacts.update(_write_text_runtime_assets(output_dir, source)) + 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 593fe2803..a60ea16fe 100644 --- a/src/mobius/integrations/onnx_genai/auto_export_test.py +++ b/src/mobius/integrations/onnx_genai/auto_export_test.py @@ -824,6 +824,38 @@ def test_dispatch_audio_only_multimodal_pipeline(tmp_path): 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]] = [] + + def fake_audio_processor(output_dir, source, *, revision=None): + calls.append((source, revision)) + return str(audio_processor) + + monkeypatch.setattr( + "mobius.integrations.onnx_genai.auto_export._write_hf_audio_processor", + fake_audio_processor, + ) + artifacts = write_onnx_genai_config( + pkg, + str(tmp_path), + source="zai-org/GLM-ASR-Nano-2512", + revision="pinned-revision", + ) + assert artifacts["audio_processor"] == str(audio_processor) + assert calls == [("zai-org/GLM-ASR-Nano-2512", "pinned-revision")] + + def test_dispatch_vision_and_audio_multimodal_pipeline(tmp_path): pkg = _vlm_package(audio=True) artifacts = write_onnx_genai_config(pkg, str(tmp_path)) diff --git a/src/mobius/integrations/onnx_genai/inference_metadata.py b/src/mobius/integrations/onnx_genai/inference_metadata.py index d91fd9b4a..a5f8039ab 100644 --- a/src/mobius/integrations/onnx_genai/inference_metadata.py +++ b/src/mobius/integrations/onnx_genai/inference_metadata.py @@ -1964,12 +1964,14 @@ 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 names: - source_path = _source_asset_path(source, filename) + 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)) @@ -2020,6 +2022,7 @@ def write_native_vlm_package_metadata( *, config: Any, source: str | None = None, + revision: str | None = None, filename: str = "inference_metadata.yaml", ) -> dict[str, str]: """Write native VLM metadata and the runtime's tokenizer/processor assets.""" diff --git a/tests/cli_test.py b/tests/cli_test.py index 155a35785..59ec850dd 100644 --- a/tests/cli_test.py +++ b/tests/cli_test.py @@ -193,6 +193,35 @@ def test_text_only_skips_diffusers_autodetect(self): mock_build.assert_called_once() assert mock_build.call_args.kwargs.get("text_only") is True + def test_revision_is_forwarded_to_detection_and_build(self): + revision = "61ba4e0b3309b6656edea3e93e419f7bd5c61957" + with ( + tempfile.TemporaryDirectory() as tmpdir, + mock.patch( + "mobius.integrations.diffusers._builder._load_diffusers_pipeline_index", + return_value=None, + ) as mock_diffusers, + mock.patch("mobius.__main__.build", return_value=mock.MagicMock()) as mock_build, + mock.patch("mobius.__main__._save_package"), + ): + main( + [ + "build", + "--model", + "zai-org/GLM-ASR-Nano-2512", + tmpdir, + "--revision", + revision, + "--no-weights", + ] + ) + + mock_diffusers.assert_called_once_with( + "zai-org/GLM-ASR-Nano-2512", + revision=revision, + ) + 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. @@ -727,6 +756,7 @@ def test_runtime_onnx_genai_routes_vlm_through_workflow_emitter(self): tmpdir, config=pkg.config, source="/models/vlm", + revision=None, ) def test_runtime_onnx_genai_does_not_fallback_for_unsupported_vlm(self): From fe1829d1aa2b3c93bdb557e67d10beec271f0f57 Mon Sep 17 00:00:00 2001 From: justinchuby Date: Fri, 21 Aug 2026 06:25:37 +0000 Subject: [PATCH 143/151] Repin ONNX GenAI: the pinned commit was orphaned by a rebase MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit `6e2ddc78` was the branch head when it was pinned, and it is now reachable from no ref on that repository — `git ls-remote` returns zero matches for it and no remote ref has it as an ancestor. The branch was rebased onto a newer base, which rewrote every commit; the same change now exists as `fb358d42` with a different hash. The failure this avoids is a slow one. `actions/checkout` keeps succeeding against an unreferenced object until GitHub collects it, and then fails with `reference is not a tree` — at a time nobody controls, on a commit nobody can inspect, in a job unrelated to whatever change is in flight. Repinned to the live head `52339e10`, verified with `ls-remote` rather than an API lookup, which answers just as happily for an orphan. Not a fast-forward, so both suites were re-run rather than assumed: 11/11 `validate_metadata` and 11/11 runtime conformance against `52339e10`, including `mobius_static_cache_workflow_executes`. The rebase carried the sole-decoder sequence-role rule through unchanged and adds a CUDA device-memory sampler fix that our packages do not exercise. The document had the same fragility as the workflow file and less reason for it, so its two commit citations are replaced by descriptions of what changed. `.github/workflows/main.yml` is now the only place naming a hash, because it is the only place that has to fetch an exact tree. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: justinchuby --- .github/workflows/main.yml | 2 +- docs/onnx-genai-performance-conformance.md | 28 ++++++++++++++++++---- 2 files changed, 24 insertions(+), 6 deletions(-) diff --git a/.github/workflows/main.yml b/.github/workflows/main.yml index 01a3be18e..c4b26ce14 100644 --- a/.github/workflows/main.yml +++ b/.github/workflows/main.yml @@ -23,7 +23,7 @@ jobs: - uses: actions/checkout@v7 with: repository: justinchuby/onnx-genai - ref: 6e2ddc78e956f84401d0fc65adbcf4538fb76ecc + ref: 52339e10fa05e3b3c99e19dbdf4c7b3446b0eff7 path: validation/onnx-genai - uses: actions/setup-python@v7 with: diff --git a/docs/onnx-genai-performance-conformance.md b/docs/onnx-genai-performance-conformance.md index 31d73201a..67475f84d 100644 --- a/docs/onnx-genai-performance-conformance.md +++ b/docs/onnx-genai-performance-conformance.md @@ -123,7 +123,7 @@ well-formed and validate; only the local kernel is missing. ### Canonical representation — lowering verified -Against ONNX GenAI `6e2ddc78`, with no `model.io` in any package and no port +Against ONNX GenAI `52339e10`, with no `model.io` in any package and no port contracts on any component that ships an artifact: | Check | Result | @@ -137,15 +137,33 @@ 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. -`7324351a` 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 +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 — +`git ls-remote` showed zero matches and no remote ref had it as an ancestor. +`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. Bumping a pin is therefore checked with `ls-remote` plus an ancestry +test rather than with an API call that answers just as happily for an orphan. + ### What the fixtures commit, and what they do not Only the metadata is committed. The graphs and the adapter weight file are a From 9390138912845d3bd65a00bd9ee33c4508f960d5 Mon Sep 17 00:00:00 2001 From: justinchuby Date: Fri, 21 Aug 2026 06:56:15 +0000 Subject: [PATCH 144/151] Repin ONNX GenAI again and stop the regenerated tree from being committable The pinned commit 52339e10 is no longer reachable from any ref: the GenAI branch was rebased a second time and now heads at 0497c6f4. An unreferenced commit is GC-eligible, and `actions/checkout` fails with "reference is not a tree" whenever the collection happens to run -- a red build at a time nobody controls, on a commit nobody can inspect. Verified with `git ls-remote` (zero refs match) and `git merge-base --is-ancestor` (not a fast-forward, so the move could not be assumed benign). `gh api .../commits/` still answers 200 for such objects, which is why it must not be the check used. Repinned to the live branch head, re-verified rather than assumed: all 11 generated packages validate against it, and all 11 execute under the runtime conformance suite, with TensorScatter firing per layer per step on the static-cache package. Also ignore validation/. CI and the materialized_workflow_packages fixture regenerate the full 14 MB of packages there, and the directory was untracked but not ignored -- one `git add -A` would have restored exactly the megabytes of unreviewable graphs that were just removed from the index. Signed-off-by: Justin Chu Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: justinchuby --- .github/workflows/main.yml | 2 +- .gitignore | 5 +++++ 2 files changed, 6 insertions(+), 1 deletion(-) diff --git a/.github/workflows/main.yml b/.github/workflows/main.yml index c4b26ce14..2834b06f2 100644 --- a/.github/workflows/main.yml +++ b/.github/workflows/main.yml @@ -23,7 +23,7 @@ jobs: - uses: actions/checkout@v7 with: repository: justinchuby/onnx-genai - ref: 52339e10fa05e3b3c99e19dbdf4c7b3446b0eff7 + ref: 0497c6f4fe746f418ac88ebfdd2865284ece7ce2 path: validation/onnx-genai - uses: actions/setup-python@v7 with: diff --git a/.gitignore b/.gitignore index 93550c1a6..ad9983b71 100644 --- a/.gitignore +++ b/.gitignore @@ -221,6 +221,11 @@ __marimo__/ # 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/** From 894527d17848496de465353fa819bfd63d298929 Mon Sep 17 00:00:00 2001 From: justinchuby Date: Fri, 21 Aug 2026 07:15:39 +0000 Subject: [PATCH 145/151] Anchor the ONNX GenAI pin against rebases instead of chasing it The pinned commit has now been rewritten out from under CI twice in a day, because the branch it names is rebased on a schedule this repository does not control. Each time, the pin kept resolving until the unreferenced object was collected, at which point `actions/checkout` would fail with "reference is not a tree" on a commit nobody can inspect. Bumping the pin each time treats the symptom; the cause is that reachability of the pinned object is someone else's to revoke. Anchor it: the tag mobius-pr478-pin in justinchuby/onnx-genai now points at the exact commit this workflow pins, so the object stays reachable across rebases of the branch it came from. The workflow still pins the SHA, not the tag -- the SHA is what makes a run reproducible, and the tag exists only to keep it alive. The pin itself is unchanged and remains the commit both suites were last verified against. Also correct the check the docs recommend. `git ls-remote | grep ` asks whether a commit is a ref *tip*, which a healthy pin stops being as soon as one more commit lands on the branch; on its own it reports a false orphan. The question that decides whether checkout resolves the commit, and whether it can be collected, is `git merge-base --is-ancestor `. Both beat `gh api .../commits/`, which answers 200 for unreferenced objects and so fails open. Signed-off-by: Justin Chu Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: justinchuby --- .github/workflows/main.yml | 11 ++++++++++ docs/onnx-genai-performance-conformance.md | 24 ++++++++++++++++++---- 2 files changed, 31 insertions(+), 4 deletions(-) diff --git a/.github/workflows/main.yml b/.github/workflows/main.yml index 2834b06f2..728e794d0 100644 --- a/.github/workflows/main.yml +++ b/.github/workflows/main.yml @@ -23,6 +23,17 @@ jobs: - uses: actions/checkout@v7 with: repository: justinchuby/onnx-genai + # Pinned by SHA so the contract under test is reproducible. That + # branch is rebased routinely, which twice left this pin naming a + # commit reachable from no ref -- GC-eligible, and `actions/checkout` + # fails with "reference is not a tree" whenever collection happens to + # run. The tag mobius-pr478-pin in that repository anchors exactly + # this commit so it stays reachable across their rebases; do not + # delete it while this pin names this SHA. When bumping, verify + # reachability with `git merge-base --is-ancestor `, + # not with `git ls-remote | grep` (which only lists ref tips) and not + # with `gh api .../commits/` (which answers 200 for unreferenced + # objects and so fails open). ref: 0497c6f4fe746f418ac88ebfdd2865284ece7ce2 path: validation/onnx-genai - uses: actions/setup-python@v7 diff --git a/docs/onnx-genai-performance-conformance.md b/docs/onnx-genai-performance-conformance.md index 67475f84d..833135707 100644 --- a/docs/onnx-genai-performance-conformance.md +++ b/docs/onnx-genai-performance-conformance.md @@ -157,12 +157,28 @@ 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 — -`git ls-remote` showed zero matches and no remote ref had it as an ancestor. +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. Bumping a pin is therefore checked with `ls-remote` plus an ancestry -test rather than with an API call that answers just as happily for an orphan. +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 From dd9eaaa5fb3d180e4b1b18efc8f2d8910389c8a3 Mon Sep 17 00:00:00 2001 From: justinchuby Date: Fri, 21 Aug 2026 09:27:20 +0000 Subject: [PATCH 146/151] Repin ONNX GenAI onto the rebased head, and move the anchor with it The GenAI branch was rebased again -- twice within about two hours, this time because origin/main advanced underneath it -- so the previously pinned commit is no longer reachable from PR #828's head. CI never became at risk: the tag mobius-pr478-pin added in the previous commit kept the pinned object referenced, which is exactly the failure this anchor exists to absorb. The tag has been moved to the new pin, per the rule recorded alongside it. Verified rather than taken on assurance, because a pin move is gated on both suites: all 11 generated packages validate at the new head, and all 11 execute under the runtime conformance suite, with TensorScatter firing per layer per step on the static-cache package. The upstream summary described this push as touching neither the schema, the validator, decoder_abi.rs, nor the workflow recognizer. That is right in substance: diffing the metadata crate between the two pins leaves validation.rs and decoder_abi.rs untouched, and the only change is in parser.rs, where MtpProposerSpec::mtp_state_output becomes Option so a proposal-local head can decline to thread recurrent state. That is the speculator sidecar descriptor, not the package load path, and no package this producer emits reaches it. Signed-off-by: Justin Chu Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: justinchuby --- .github/workflows/main.yml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/.github/workflows/main.yml b/.github/workflows/main.yml index 728e794d0..ad00e4f08 100644 --- a/.github/workflows/main.yml +++ b/.github/workflows/main.yml @@ -34,7 +34,7 @@ jobs: # not with `git ls-remote | grep` (which only lists ref tips) and not # with `gh api .../commits/` (which answers 200 for unreferenced # objects and so fails open). - ref: 0497c6f4fe746f418ac88ebfdd2865284ece7ce2 + ref: 4315a94d114a69d6c1464dc48b4aeb53c16da566 path: validation/onnx-genai - uses: actions/setup-python@v7 with: From 00c748db0f0fa1a189e40e20878ba01acb17cf17 Mon Sep 17 00:00:00 2001 From: justinchuby Date: Fri, 21 Aug 2026 09:36:44 +0000 Subject: [PATCH 147/151] Give encoder embeddings their own workflow metadata, and add ESM-2 Protein language models were the first real users of the encoder path, and exporting them exposed four defects that all shared one cause: encoders were being treated as decoders that happened to stop early. Metadata. An encoder package fell through to the decoder producer and emitted a greedy autoregressive loop -- max_output_tokens, eos ids, a sampler, a KV cache -- for a bidirectional model that cannot generate anything. The new encoder-embedding producer emits what actually happens: one invoke, one emit, an `embedding` profile carrying mask-aware mean pooling, and `batch_invariance: row_independent`. The declared inputs are read from the artifact rather than the task signature, because optimization prunes unused graph inputs -- ESM-2 has no token-type embedding, so its metadata must not promise one. Batching. The rank-2 int64 attention mask was passed straight to op.Attention, which cannot broadcast it, so every encoder was broken for batch > 1 and merely added a 0/1 bias at batch 1. BERT now builds a 4D bool padding mask once and shares it across layers. DistilBert. It discarded `attention_mask` outright: padded rows changed 99.4% of their values. The mask is now threaded through the encoder layers. This was found by the new padded-batch integration test, not by inspection, which is the argument for the test. ESM-2. It was registered to the generic BERT module, which cannot load the checkpoint: ESM-2 uses rotary embeddings, pre-norm blocks, a final embedding LayerNorm, token dropout, and no token-type embeddings. `models/esm.py` implements it, with submodule names arranged so initializer names match HuggingFace directly and renaming reduces to stripping the `esm.` prefix. ProtBert. Its config declares no `model_type`, so `AutoConfig` refuses it and the builder misrouted to the diffusers path. `model_type` is now inferred from `architectures`, and `--config` accepts a local directory. Verified against real weights on H200. ESM-2 (facebook/esm2_t6_8M_UR50D, MIT) and ProtBert (Rostlab/prot_bert) reproduce HuggingFace to a relative L2 error of 1.2e-4 and 1.7e-5 on human haemoglobin, insulin, and lysozyme; the larger CUDA gap is TF32 and disappears with use_tf32=0. Padding invariance, batch-order invariance, and pad-token isolation are bit-exact. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: justinchuby (cherry picked from commit d6c90f199f5b58a253a18414573d7b2999f2d6f5) --- docs/onnx-genai-workflows.md | 35 ++ src/mobius/__main__.py | 18 +- src/mobius/_registry.py | 3 +- .../integrations/onnx_genai/auto_export.py | 40 ++ .../encoder_embedding_metadata_test.py | 171 ++++++++ .../onnx_genai/workflow_metadata.py | 156 +++++++ .../transformers/_config_resolver.py | 64 ++- src/mobius/models/__init__.py | 3 + src/mobius/models/bert.py | 19 +- src/mobius/models/bert_test.py | 21 + src/mobius/models/distilbert.py | 34 +- src/mobius/models/esm.py | 407 ++++++++++++++++++ src/mobius/models/esm_test.py | 125 ++++++ testdata/cases/encoder/esm2-8m.yaml | 7 +- testdata/cases/encoder/protbert.yaml | 23 + tests/_test_configs.py | 14 +- tests/fixtures/onnx_genai_workflows/README.md | 11 +- .../inference_metadata.yaml | 101 +++++ .../inference_metadata.yaml | 119 +++++ ...generate_onnx_genai_validation_packages.py | 62 ++- tests/integration_test.py | 95 ++++ 21 files changed, 1493 insertions(+), 35 deletions(-) create mode 100644 src/mobius/integrations/onnx_genai/encoder_embedding_metadata_test.py create mode 100644 src/mobius/models/bert_test.py create mode 100644 src/mobius/models/esm.py create mode 100644 src/mobius/models/esm_test.py create mode 100644 testdata/cases/encoder/protbert.yaml create mode 100644 tests/fixtures/onnx_genai_workflows/esm2_protein_embeddings/inference_metadata.yaml create mode 100644 tests/fixtures/onnx_genai_workflows/protbert_protein_embeddings/inference_metadata.yaml diff --git a/docs/onnx-genai-workflows.md b/docs/onnx-genai-workflows.md index 34c53e5e0..b8a6feaac 100644 --- a/docs/onnx-genai-workflows.md +++ b/docs/onnx-genai-workflows.md @@ -230,6 +230,41 @@ Two consequences: ## 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 diff --git a/src/mobius/__main__.py b/src/mobius/__main__.py index 47e1482af..ce5e78117 100644 --- a/src/mobius/__main__.py +++ b/src/mobius/__main__.py @@ -317,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"): diff --git a/src/mobius/_registry.py b/src/mobius/_registry.py index eaa6d4be8..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"), diff --git a/src/mobius/integrations/onnx_genai/auto_export.py b/src/mobius/integrations/onnx_genai/auto_export.py index 8dd023a27..2ada88da5 100644 --- a/src/mobius/integrations/onnx_genai/auto_export.py +++ b/src/mobius/integrations/onnx_genai/auto_export.py @@ -31,6 +31,7 @@ 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, @@ -560,6 +561,34 @@ def _looks_like_ctc_asr(pkg: Any) -> bool: 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. @@ -844,6 +873,17 @@ def write_onnx_genai_config( 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( 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..81aa3412a --- /dev/null +++ b/src/mobius/integrations/onnx_genai/encoder_embedding_metadata_test.py @@ -0,0 +1,171 @@ +# 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_mask_aware_pooling_is_declared_against_a_real_input( + self, metadata, package + ) -> None: + profile = metadata["profiles"]["embedding"] + workflow = metadata["pipeline"]["workflow"] + assert profile["pooling"]["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_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/workflow_metadata.py b/src/mobius/integrations/onnx_genai/workflow_metadata.py index 7628183a1..258e43298 100644 --- a/src/mobius/integrations/onnx_genai/workflow_metadata.py +++ b/src/mobius/integrations/onnx_genai/workflow_metadata.py @@ -8097,6 +8097,162 @@ def write_ctc_asr_workflow_metadata( 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 + declaration: dict[str, Any] = { + "contract": _contract(graph_inputs[name]), + "role": {"kind": "runtime", "version": "1.0", "role": role}, + "source": {"kind": "request"}, + # 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": { + "ir_version": "1.0", + "onnx_opsets": {"ai.onnx": OPSET_VERSION}, + "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: + # Mask-aware mean pooling is only well defined when the graph is told + # which positions are padding; a reader can then reduce a row over its + # own valid region instead of over the width the batch happened to be + # padded to. The same fact is what makes rows independent of their + # neighbours, so both claims are made together or neither is. + profile["pooling"] = { + "kind": "mean", + "source": "last_hidden_state", + "mask": "request.attention_mask", + "time_axis": 1, + "feature_axis": 2, + } + 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 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 1f4863fc8..99557e267 100644 --- a/src/mobius/models/__init__.py +++ b/src/mobius/models/__init__.py @@ -36,6 +36,8 @@ "DiTTransformer2DModel", "DiffLlamaCausalLMModel", "DistilBertModel", + "EsmConfig", + "EsmModel", "DogeCausalLMModel", "EncDecRNNTModel", "Ernie45MoECausalLMModel", @@ -209,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/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/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/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/fixtures/onnx_genai_workflows/README.md b/tests/fixtures/onnx_genai_workflows/README.md index 70ea478d6..dbf9c5df2 100644 --- a/tests/fixtures/onnx_genai_workflows/README.md +++ b/tests/fixtures/onnx_genai_workflows/README.md @@ -14,7 +14,10 @@ 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. 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. +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/esm2_protein_embeddings/inference_metadata.yaml b/tests/fixtures/onnx_genai_workflows/esm2_protein_embeddings/inference_metadata.yaml new file mode 100644 index 000000000..f5b2ad61b --- /dev/null +++ b/tests/fixtures/onnx_genai_workflows/esm2_protein_embeddings/inference_metadata.yaml @@ -0,0 +1,101 @@ +schema_version: v1 +profiles: + embedding: + kind: embedding + version: '1.0' + requirement: required + 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 +pipeline: + workflow: + manifest: + ir_version: '1.0' + onnx_opsets: + ai.onnx: 24 + 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: runtime + version: '1.0' + role: attention_mask + source: + kind: request + 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/protbert_protein_embeddings/inference_metadata.yaml b/tests/fixtures/onnx_genai_workflows/protbert_protein_embeddings/inference_metadata.yaml new file mode 100644 index 000000000..4a4bc2323 --- /dev/null +++ b/tests/fixtures/onnx_genai_workflows/protbert_protein_embeddings/inference_metadata.yaml @@ -0,0 +1,119 @@ +schema_version: v1 +profiles: + embedding: + kind: embedding + version: '1.0' + requirement: required + 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 +pipeline: + workflow: + manifest: + ir_version: '1.0' + onnx_opsets: + ai.onnx: 24 + 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: runtime + version: '1.0' + role: attention_mask + source: + kind: request + required: true + request.token_type_ids: + contract: + dtype: int64 + rank: 2 + shape: + - batch + - sequence_len + batch_layout: + kind: request_aligned + axis: 0 + role: + kind: runtime + version: '1.0' + role: token_type_ids + source: + kind: request + required: true + outputs: + last_hidden_state: + contract: + dtype: float32 + rank: 3 + shape: + - batch + - null + - 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/generate_onnx_genai_validation_packages.py b/tests/generate_onnx_genai_validation_packages.py index fbe7d52b6..3d0cd191a 100644 --- a/tests/generate_onnx_genai_validation_packages.py +++ b/tests/generate_onnx_genai_validation_packages.py @@ -12,6 +12,7 @@ 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, @@ -33,13 +34,18 @@ 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 TTSTask +from mobius.tasks import FeatureExtractionTask, TTSTask def _materialize_deterministic_initializers(package: ModelPackage) -> None: @@ -680,6 +686,35 @@ def _executable_masked_package() -> ModelPackage: 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( @@ -1029,6 +1064,20 @@ def generate_packages(output: Path) -> Path: 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) @@ -1053,10 +1102,13 @@ def generate_packages(output: Path) -> Path: 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. 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. +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", ) 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) From ab343c352d4d9e0597fb9ee46c357ba78ad95129 Mon Sep 17 00:00:00 2001 From: justinchuby Date: Fri, 21 Aug 2026 18:07:21 +0000 Subject: [PATCH 148/151] Align metadata producer with final workflow schema Pin validation to the final ONNX GenAI metadata commit, update encoder embedding profiles and auxiliary inputs to the canonical schema, and assign stable symbols to anonymous dynamic dimensions so generated packages validate without null shapes. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: justinchuby --- .github/workflows/main.yml | 18 ++++-------- .../encoder_embedding_metadata_test.py | 16 ++++++++--- .../onnx_genai/inference_metadata.py | 11 +++++--- .../onnx_genai/workflow_metadata.py | 28 +++++++++++-------- .../inference_metadata.yaml | 13 ++++----- .../inference_metadata.yaml | 22 ++++++--------- 6 files changed, 56 insertions(+), 52 deletions(-) diff --git a/.github/workflows/main.yml b/.github/workflows/main.yml index ad00e4f08..cd8cf2e09 100644 --- a/.github/workflows/main.yml +++ b/.github/workflows/main.yml @@ -23,18 +23,12 @@ jobs: - uses: actions/checkout@v7 with: repository: justinchuby/onnx-genai - # Pinned by SHA so the contract under test is reproducible. That - # branch is rebased routinely, which twice left this pin naming a - # commit reachable from no ref -- GC-eligible, and `actions/checkout` - # fails with "reference is not a tree" whenever collection happens to - # run. The tag mobius-pr478-pin in that repository anchors exactly - # this commit so it stays reachable across their rebases; do not - # delete it while this pin names this SHA. When bumping, verify - # reachability with `git merge-base --is-ancestor `, - # not with `git ls-remote | grep` (which only lists ref tips) and not - # with `gh api .../commits/` (which answers 200 for unreferenced - # objects and so fails open). - ref: 4315a94d114a69d6c1464dc48b4aeb53c16da566 + # 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: 857d96cf25008cdf229252b9dbea3e01c23fac7f path: validation/onnx-genai - uses: actions/setup-python@v7 with: diff --git a/src/mobius/integrations/onnx_genai/encoder_embedding_metadata_test.py b/src/mobius/integrations/onnx_genai/encoder_embedding_metadata_test.py index 81aa3412a..67f47a8e8 100644 --- a/src/mobius/integrations/onnx_genai/encoder_embedding_metadata_test.py +++ b/src/mobius/integrations/onnx_genai/encoder_embedding_metadata_test.py @@ -107,12 +107,11 @@ def test_the_profile_is_an_embedding_profile(self, metadata) -> None: assert profile["kind"] == "embedding" assert profile["outputs"] == {"last_hidden_state": "last_hidden_state"} - def test_mask_aware_pooling_is_declared_against_a_real_input( - self, metadata, package - ) -> None: + def test_pooling_uses_the_sequence_axis(self, metadata) -> None: profile = metadata["profiles"]["embedding"] workflow = metadata["pipeline"]["workflow"] - assert profile["pooling"]["mask"] in workflow["inputs"] + 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: @@ -128,6 +127,15 @@ def test_every_declared_input_is_bound(self, metadata) -> None: 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() diff --git a/src/mobius/integrations/onnx_genai/inference_metadata.py b/src/mobius/integrations/onnx_genai/inference_metadata.py index a5f8039ab..6634c5f32 100644 --- a/src/mobius/integrations/onnx_genai/inference_metadata.py +++ b/src/mobius/integrations/onnx_genai/inference_metadata.py @@ -134,15 +134,18 @@ def _port(value: Any) -> _Port: """Symbolic leading dimension mobius uses for per-request batching.""" -def _shape_metadata(port: _Port) -> list[int | str | None]: +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 diff --git a/src/mobius/integrations/onnx_genai/workflow_metadata.py b/src/mobius/integrations/onnx_genai/workflow_metadata.py index 258e43298..ac6092375 100644 --- a/src/mobius/integrations/onnx_genai/workflow_metadata.py +++ b/src/mobius/integrations/onnx_genai/workflow_metadata.py @@ -8153,10 +8153,20 @@ def build_encoder_embedding_workflow_metadata( 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": {"kind": "runtime", "version": "1.0", "role": role}, - "source": {"kind": "request"}, + "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, @@ -8218,17 +8228,13 @@ def build_encoder_embedding_workflow_metadata( "outputs": profile_outputs, } if "attention_mask" in graph_inputs: - # Mask-aware mean pooling is only well defined when the graph is told - # which positions are padding; a reader can then reduce a row over its - # own valid region instead of over the width the batch happened to be - # padded to. The same fact is what makes rows independent of their - # neighbours, so both claims are made together or neither is. + # 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", - "source": "last_hidden_state", - "mask": "request.attention_mask", - "time_axis": 1, - "feature_axis": 2, + "axis": 1, + "normalize": False, } profile["batch_invariance"] = "row_independent" 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 index f5b2ad61b..836532361 100644 --- a/tests/fixtures/onnx_genai_workflows/esm2_protein_embeddings/inference_metadata.yaml +++ b/tests/fixtures/onnx_genai_workflows/esm2_protein_embeddings/inference_metadata.yaml @@ -8,10 +8,8 @@ profiles: last_hidden_state: last_hidden_state pooling: kind: mean - source: last_hidden_state - mask: request.attention_mask - time_axis: 1 - feature_axis: 2 + axis: 1 + normalize: false batch_invariance: row_independent pipeline: workflow: @@ -57,11 +55,10 @@ pipeline: kind: request_aligned axis: 0 role: - kind: runtime - version: '1.0' - role: attention_mask + kind: opaque source: - kind: request + kind: application + name: request.attention_mask required: true outputs: last_hidden_state: 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 index 4a4bc2323..0eeb0be18 100644 --- a/tests/fixtures/onnx_genai_workflows/protbert_protein_embeddings/inference_metadata.yaml +++ b/tests/fixtures/onnx_genai_workflows/protbert_protein_embeddings/inference_metadata.yaml @@ -8,10 +8,8 @@ profiles: last_hidden_state: last_hidden_state pooling: kind: mean - source: last_hidden_state - mask: request.attention_mask - time_axis: 1 - feature_axis: 2 + axis: 1 + normalize: false batch_invariance: row_independent pipeline: workflow: @@ -57,11 +55,10 @@ pipeline: kind: request_aligned axis: 0 role: - kind: runtime - version: '1.0' - role: attention_mask + kind: opaque source: - kind: request + kind: application + name: request.attention_mask required: true request.token_type_ids: contract: @@ -74,11 +71,10 @@ pipeline: kind: request_aligned axis: 0 role: - kind: runtime - version: '1.0' - role: token_type_ids + kind: opaque source: - kind: request + kind: application + name: request.token_type_ids required: true outputs: last_hidden_state: @@ -87,7 +83,7 @@ pipeline: rank: 3 shape: - batch - - null + - last_hidden_state_dim_1 - 64 batch_layout: kind: request_aligned From 6429957fbceb8a24b825be9b101b410b700fb6a1 Mon Sep 17 00:00:00 2001 From: justinchuby Date: Fri, 21 Aug 2026 18:28:25 +0000 Subject: [PATCH 149/151] Stop duplicating workflow version and ONNX opsets Let schema_version govern workflow syntax and let each ONNX artifact remain authoritative for its opset imports. Regenerate every producer fixture and pin cross-repository validation to the corresponding ONNX GenAI schema commit. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: justinchuby --- .github/workflows/main.yml | 2 +- .../onnx_genai/auto_export_test.py | 3 ++- .../onnx_genai/inference_metadata_test.py | 2 +- .../onnx_genai/workflow_metadata.py | 27 ------------------- .../adapter/inference_metadata.yaml | 3 --- .../codec/inference_metadata.yaml | 3 --- .../decoder/inference_metadata.yaml | 3 --- .../diffusion/inference_metadata.yaml | 3 --- .../diffusion_guided/inference_metadata.yaml | 3 --- .../inference_metadata.yaml | 3 --- .../masked/inference_metadata.yaml | 3 --- .../inference_metadata.yaml | 3 --- .../speculative/inference_metadata.yaml | 3 --- .../static_cache/inference_metadata.yaml | 3 --- .../tts/inference_metadata.yaml | 3 --- .../video/inference_metadata.yaml | 3 --- .../vlm/inference_metadata.yaml | 3 --- ...generate_onnx_genai_validation_packages.py | 2 -- 18 files changed, 4 insertions(+), 71 deletions(-) diff --git a/.github/workflows/main.yml b/.github/workflows/main.yml index cd8cf2e09..e286b455a 100644 --- a/.github/workflows/main.yml +++ b/.github/workflows/main.yml @@ -28,7 +28,7 @@ jobs: # rebasing, so this commit remains reachable from its branch history. # When bumping, verify reachability with `git merge-base --is-ancestor # `. - ref: 857d96cf25008cdf229252b9dbea3e01c23fac7f + ref: f26772bd4b29d7d84b9a4af9ff401c8aa9ba48c6 path: validation/onnx-genai - uses: actions/setup-python@v7 with: diff --git a/src/mobius/integrations/onnx_genai/auto_export_test.py b/src/mobius/integrations/onnx_genai/auto_export_test.py index a60ea16fe..c246ad6c2 100644 --- a/src/mobius/integrations/onnx_genai/auto_export_test.py +++ b/src/mobius/integrations/onnx_genai/auto_export_test.py @@ -165,7 +165,8 @@ def test_dispatch_decoder(tmp_path): with open(arts["inference_metadata"]) as handle: meta = yaml.safe_load(handle) workflow = meta["pipeline"]["workflow"] - assert workflow["manifest"]["ir_version"] == "1.0" + 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" ) diff --git a/src/mobius/integrations/onnx_genai/inference_metadata_test.py b/src/mobius/integrations/onnx_genai/inference_metadata_test.py index 8e529fff6..8bd9c345d 100644 --- a/src/mobius/integrations/onnx_genai/inference_metadata_test.py +++ b/src/mobius/integrations/onnx_genai/inference_metadata_test.py @@ -135,7 +135,7 @@ def test_workflow_policy_components_reference_saved_onnx_artifacts(tmp_path): metadata = { "pipeline": { "workflow": { - "manifest": {"ir_version": "1.0"}, + "manifest": {}, "components": {}, "graph": {"kind": "sequence", "nodes": []}, } diff --git a/src/mobius/integrations/onnx_genai/workflow_metadata.py b/src/mobius/integrations/onnx_genai/workflow_metadata.py index ac6092375..3e3928713 100644 --- a/src/mobius/integrations/onnx_genai/workflow_metadata.py +++ b/src/mobius/integrations/onnx_genai/workflow_metadata.py @@ -16,7 +16,6 @@ import yaml from mobius._constants import ( - OPSET_VERSION, STATIC_CACHE_KV_SEQUENCE_LENGTH, STATIC_CACHE_LAYOUT, STATIC_CACHE_SEQUENCE_AXIS, @@ -501,8 +500,6 @@ def build_audio_codec_workflow_metadata(pkg: Any) -> dict[str, Any]: emit_effect = "audio_emit" workflow = { "manifest": { - "ir_version": "1.0", - "onnx_opsets": {"ai.onnx": OPSET_VERSION}, "capabilities": [ "workflow_ssa", "linear_effects", @@ -2045,8 +2042,6 @@ def frame_generation_nodes(prefix: str, hidden: str, logits: str) -> list[dict[s ) workflow = { "manifest": { - "ir_version": "1.0", - "onnx_opsets": {"ai.onnx": OPSET_VERSION}, "capabilities": [ "workflow_ssa", "linear_effects", @@ -2563,8 +2558,6 @@ def bind_outputs(values: Any, bound: dict[str, str], prefix: str) -> dict[str, s ) workflow = { "manifest": { - "ir_version": "1.0", - "onnx_opsets": {"ai.onnx": OPSET_VERSION}, "capabilities": [ "workflow_ssa", "linear_effects", @@ -3207,8 +3200,6 @@ def denoiser_call(conditioning: str | None, estimate: str) -> dict[str, Any]: workflow = { "manifest": { - "ir_version": "1.0", - "onnx_opsets": {"ai.onnx": OPSET_VERSION}, "capabilities": [ "workflow_ssa", "linear_effects", @@ -3505,8 +3496,6 @@ def denoise(prefix: str, output: str) -> dict[str, Any]: latent_effect = "state:latent" workflow = { "manifest": { - "ir_version": "1.0", - "onnx_opsets": {"ai.onnx": OPSET_VERSION}, "capabilities": [ "workflow_ssa", "linear_effects", @@ -3977,8 +3966,6 @@ def build_video_diffusion_workflow_metadata( frames_contract = _request_aligned(_contract(vae_output)) workflow = { "manifest": { - "ir_version": "1.0", - "onnx_opsets": {"ai.onnx": OPSET_VERSION}, "capabilities": [ "workflow_ssa", "linear_effects", @@ -5174,8 +5161,6 @@ def build_vlm_workflow_metadata( } workflow = { "manifest": { - "ir_version": "1.0", - "onnx_opsets": {"ai.onnx": OPSET_VERSION}, "adapter_abis": {"onnx-genai.image-preprocess": "1"}, "capabilities": [ "workflow_ssa", @@ -5943,8 +5928,6 @@ def build_speculative_workflow_metadata( ) workflow = { "manifest": { - "ir_version": "1.0", - "onnx_opsets": {"ai.onnx": OPSET_VERSION}, "capabilities": [ "workflow_ssa", "linear_effects", @@ -7338,8 +7321,6 @@ def _build_autoregressive_workflow_metadata( ) workflow = { "manifest": { - "ir_version": "1.0", - "onnx_opsets": {"ai.onnx": OPSET_VERSION}, **( {"adapter_abis": {"onnx-genai.audio-preprocess": "1"}} if audio_program is not None @@ -7664,8 +7645,6 @@ def update_invoke( workflow = { "manifest": { - "ir_version": "1.0", - "onnx_opsets": {"ai.onnx": OPSET_VERSION}, "capabilities": [ "workflow_ssa", "linear_effects", @@ -7940,8 +7919,6 @@ def build_ctc_asr_workflow_metadata( workflow = { "manifest": { - "ir_version": "1.0", - "onnx_opsets": {"ai.onnx": OPSET_VERSION}, "adapter_abis": {_AUDIO_PREPROCESS_ABI: _AUDIO_PREPROCESS_ABI_VERSION}, "capabilities": ["workflow_ssa", "linear_effects", "typed_emit"], }, @@ -8199,8 +8176,6 @@ def build_encoder_embedding_workflow_metadata( workflow = { "manifest": { - "ir_version": "1.0", - "onnx_opsets": {"ai.onnx": OPSET_VERSION}, "capabilities": ["workflow_ssa", "linear_effects", "typed_emit"], }, "effects": { @@ -9085,8 +9060,6 @@ def session_cell( workflow = { "manifest": { - "ir_version": "1.0", - "onnx_opsets": {"ai.onnx": OPSET_VERSION}, "capabilities": [ "workflow_ssa", "linear_effects", diff --git a/tests/fixtures/onnx_genai_workflows/adapter/inference_metadata.yaml b/tests/fixtures/onnx_genai_workflows/adapter/inference_metadata.yaml index c2e2ff17b..b8e5e5a82 100644 --- a/tests/fixtures/onnx_genai_workflows/adapter/inference_metadata.yaml +++ b/tests/fixtures/onnx_genai_workflows/adapter/inference_metadata.yaml @@ -2,9 +2,6 @@ schema_version: v1 pipeline: workflow: manifest: - ir_version: '1.0' - onnx_opsets: - ai.onnx: 24 adapter_abis: onnx-genai.parameter-overlay: '1' capabilities: diff --git a/tests/fixtures/onnx_genai_workflows/codec/inference_metadata.yaml b/tests/fixtures/onnx_genai_workflows/codec/inference_metadata.yaml index aa6c5e570..09801175b 100644 --- a/tests/fixtures/onnx_genai_workflows/codec/inference_metadata.yaml +++ b/tests/fixtures/onnx_genai_workflows/codec/inference_metadata.yaml @@ -2,9 +2,6 @@ schema_version: v1 pipeline: workflow: manifest: - ir_version: '1.0' - onnx_opsets: - ai.onnx: 24 capabilities: - workflow_ssa - linear_effects diff --git a/tests/fixtures/onnx_genai_workflows/decoder/inference_metadata.yaml b/tests/fixtures/onnx_genai_workflows/decoder/inference_metadata.yaml index 46066a212..7253fb474 100644 --- a/tests/fixtures/onnx_genai_workflows/decoder/inference_metadata.yaml +++ b/tests/fixtures/onnx_genai_workflows/decoder/inference_metadata.yaml @@ -2,9 +2,6 @@ schema_version: '1.0' pipeline: workflow: manifest: - ir_version: '1.0' - onnx_opsets: - ai.onnx: 24 capabilities: - workflow_ssa - linear_effects diff --git a/tests/fixtures/onnx_genai_workflows/diffusion/inference_metadata.yaml b/tests/fixtures/onnx_genai_workflows/diffusion/inference_metadata.yaml index 64d063dfa..638decb07 100644 --- a/tests/fixtures/onnx_genai_workflows/diffusion/inference_metadata.yaml +++ b/tests/fixtures/onnx_genai_workflows/diffusion/inference_metadata.yaml @@ -2,9 +2,6 @@ schema_version: v1 pipeline: workflow: manifest: - ir_version: '1.0' - onnx_opsets: - ai.onnx: 24 capabilities: - workflow_ssa - linear_effects diff --git a/tests/fixtures/onnx_genai_workflows/diffusion_guided/inference_metadata.yaml b/tests/fixtures/onnx_genai_workflows/diffusion_guided/inference_metadata.yaml index c2fd5a1c8..a40fb218d 100644 --- a/tests/fixtures/onnx_genai_workflows/diffusion_guided/inference_metadata.yaml +++ b/tests/fixtures/onnx_genai_workflows/diffusion_guided/inference_metadata.yaml @@ -2,9 +2,6 @@ schema_version: v1 pipeline: workflow: manifest: - ir_version: '1.0' - onnx_opsets: - ai.onnx: 24 capabilities: - workflow_ssa - linear_effects 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 index 836532361..56ef7f831 100644 --- a/tests/fixtures/onnx_genai_workflows/esm2_protein_embeddings/inference_metadata.yaml +++ b/tests/fixtures/onnx_genai_workflows/esm2_protein_embeddings/inference_metadata.yaml @@ -14,9 +14,6 @@ profiles: pipeline: workflow: manifest: - ir_version: '1.0' - onnx_opsets: - ai.onnx: 24 capabilities: - workflow_ssa - linear_effects diff --git a/tests/fixtures/onnx_genai_workflows/masked/inference_metadata.yaml b/tests/fixtures/onnx_genai_workflows/masked/inference_metadata.yaml index b898a308c..51892ae4d 100644 --- a/tests/fixtures/onnx_genai_workflows/masked/inference_metadata.yaml +++ b/tests/fixtures/onnx_genai_workflows/masked/inference_metadata.yaml @@ -2,9 +2,6 @@ schema_version: '1.0' pipeline: workflow: manifest: - ir_version: '1.0' - onnx_opsets: - ai.onnx: 24 capabilities: - workflow_ssa - linear_effects 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 index 0eeb0be18..28ff17289 100644 --- a/tests/fixtures/onnx_genai_workflows/protbert_protein_embeddings/inference_metadata.yaml +++ b/tests/fixtures/onnx_genai_workflows/protbert_protein_embeddings/inference_metadata.yaml @@ -14,9 +14,6 @@ profiles: pipeline: workflow: manifest: - ir_version: '1.0' - onnx_opsets: - ai.onnx: 24 capabilities: - workflow_ssa - linear_effects diff --git a/tests/fixtures/onnx_genai_workflows/speculative/inference_metadata.yaml b/tests/fixtures/onnx_genai_workflows/speculative/inference_metadata.yaml index e818b8b2c..5ca10d33f 100644 --- a/tests/fixtures/onnx_genai_workflows/speculative/inference_metadata.yaml +++ b/tests/fixtures/onnx_genai_workflows/speculative/inference_metadata.yaml @@ -2,9 +2,6 @@ schema_version: v1 pipeline: workflow: manifest: - ir_version: '1.0' - onnx_opsets: - ai.onnx: 24 capabilities: - workflow_ssa - linear_effects diff --git a/tests/fixtures/onnx_genai_workflows/static_cache/inference_metadata.yaml b/tests/fixtures/onnx_genai_workflows/static_cache/inference_metadata.yaml index 73912050c..534530704 100644 --- a/tests/fixtures/onnx_genai_workflows/static_cache/inference_metadata.yaml +++ b/tests/fixtures/onnx_genai_workflows/static_cache/inference_metadata.yaml @@ -2,9 +2,6 @@ schema_version: '1.0' pipeline: workflow: manifest: - ir_version: '1.0' - onnx_opsets: - ai.onnx: 24 capabilities: - workflow_ssa - linear_effects diff --git a/tests/fixtures/onnx_genai_workflows/tts/inference_metadata.yaml b/tests/fixtures/onnx_genai_workflows/tts/inference_metadata.yaml index 8176c4723..113cbbe00 100644 --- a/tests/fixtures/onnx_genai_workflows/tts/inference_metadata.yaml +++ b/tests/fixtures/onnx_genai_workflows/tts/inference_metadata.yaml @@ -2,9 +2,6 @@ schema_version: v1 pipeline: workflow: manifest: - ir_version: '1.0' - onnx_opsets: - ai.onnx: 24 capabilities: - workflow_ssa - linear_effects diff --git a/tests/fixtures/onnx_genai_workflows/video/inference_metadata.yaml b/tests/fixtures/onnx_genai_workflows/video/inference_metadata.yaml index effc1c829..f31277956 100644 --- a/tests/fixtures/onnx_genai_workflows/video/inference_metadata.yaml +++ b/tests/fixtures/onnx_genai_workflows/video/inference_metadata.yaml @@ -2,9 +2,6 @@ schema_version: v1 pipeline: workflow: manifest: - ir_version: '1.0' - onnx_opsets: - ai.onnx: 24 capabilities: - workflow_ssa - linear_effects diff --git a/tests/fixtures/onnx_genai_workflows/vlm/inference_metadata.yaml b/tests/fixtures/onnx_genai_workflows/vlm/inference_metadata.yaml index c4306fbe4..f2dce1a06 100644 --- a/tests/fixtures/onnx_genai_workflows/vlm/inference_metadata.yaml +++ b/tests/fixtures/onnx_genai_workflows/vlm/inference_metadata.yaml @@ -73,9 +73,6 @@ preprocessing: pipeline: workflow: manifest: - ir_version: '1.0' - onnx_opsets: - ai.onnx: 24 adapter_abis: onnx-genai.image-preprocess: '1' capabilities: diff --git a/tests/generate_onnx_genai_validation_packages.py b/tests/generate_onnx_genai_validation_packages.py index 3d0cd191a..05d341829 100644 --- a/tests/generate_onnx_genai_validation_packages.py +++ b/tests/generate_onnx_genai_validation_packages.py @@ -864,8 +864,6 @@ def _write_adapter_metadata(package: ModelPackage, directory: Path) -> None: "pipeline": { "workflow": { "manifest": { - "ir_version": "1.0", - "onnx_opsets": {"ai.onnx": 24}, "adapter_abis": {"onnx-genai.parameter-overlay": "1"}, "capabilities": [ "workflow_ssa", From 2c8ae522d57e7b87cd820af62c7c8faa32eeac1a Mon Sep 17 00:00:00 2001 From: justinchuby Date: Fri, 21 Aug 2026 20:03:27 +0000 Subject: [PATCH 150/151] Update canonical workflow metadata design Correct the producer documentation to keep ONNX port contracts and opsets artifact-owned, retain semantic roles and state aliases, and support independently shaped K/V tensors across layers. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: justinchuby --- docs/onnx-genai-workflows.md | 31 +++++++++++++++++++++++-------- 1 file changed, 23 insertions(+), 8 deletions(-) diff --git a/docs/onnx-genai-workflows.md b/docs/onnx-genai-workflows.md index b8a6feaac..a3fadf438 100644 --- a/docs/onnx-genai-workflows.md +++ b/docs/onnx-genai-workflows.md @@ -17,11 +17,9 @@ 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 has to carry everything such a lowering needs, so -every ONNX component declares: +For that to work the workflow carries what the ONNX artifact cannot say. +Every ONNX component declares: -* `ports.inputs` / `ports.outputs` — a contract (dtype, rank, shape, batch - layout) for every graph input and output, no more and no fewer. * `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 @@ -34,18 +32,35 @@ every ONNX component declares: `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 the same dtype and the same shape. + 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. -Both are emitted together or not at all. A recurrent or convolution cache has no -halves and no layer index to state, and inventing one would corrupt the very -ordering the index exists to fix. +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 From 1219f0ba1653d9dc6b2b1d5550037296272a577b Mon Sep 17 00:00:00 2001 From: justinchuby Date: Fri, 21 Aug 2026 20:03:44 +0000 Subject: [PATCH 151/151] Pin final metadata design revision Advance the cross-repository conformance pin to the ONNX GenAI revision containing the finalized heterogeneous KV design documentation. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: justinchuby --- .github/workflows/main.yml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/.github/workflows/main.yml b/.github/workflows/main.yml index e286b455a..2b4717c8b 100644 --- a/.github/workflows/main.yml +++ b/.github/workflows/main.yml @@ -28,7 +28,7 @@ jobs: # rebasing, so this commit remains reachable from its branch history. # When bumping, verify reachability with `git merge-base --is-ancestor # `. - ref: f26772bd4b29d7d84b9a4af9ff401c8aa9ba48c6 + ref: 509cd4e9c4471f4cbc59fe44b47168f0ae128fe3 path: validation/onnx-genai - uses: actions/setup-python@v7 with: