From e1aa52bb63a65b07981e8ff2d3afde012eaf81d8 Mon Sep 17 00:00:00 2001 From: justinchuby Date: Wed, 29 Jul 2026 19:47:12 +0000 Subject: [PATCH 1/3] Emit explicit ONNX GenAI inference metadata Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- .../integrations/onnx_genai/auto_export.py | 22 ++ .../onnx_genai/auto_export_test.py | 25 +- .../onnx_genai/inference_metadata.py | 249 +++++++++++++++++- .../onnx_genai/inference_metadata_test.py | 131 +++++++++ 4 files changed, 406 insertions(+), 21 deletions(-) diff --git a/src/mobius/integrations/onnx_genai/auto_export.py b/src/mobius/integrations/onnx_genai/auto_export.py index b5fef98a..9c2a0c3c 100644 --- a/src/mobius/integrations/onnx_genai/auto_export.py +++ b/src/mobius/integrations/onnx_genai/auto_export.py @@ -16,12 +16,15 @@ import os from typing import Any +import yaml + from mobius.integrations.onnx_genai.decoder_metadata import ( decoder_metadata_from_config, write_decoder_metadata, ) from mobius.integrations.onnx_genai.inference_metadata import ( SchedulerConfig, + add_explicit_package_io, load_diffusers_scheduler_config, write_audio_codec_pipeline_metadata, write_diffusion_pipeline_metadata, @@ -35,6 +38,21 @@ _DENOISER_KEYS = ("denoiser", "transformer", "unet") +def _add_explicit_io_to_file(path: str, pkg: Any, config: Any) -> None: + """Augment an emitted sidecar with roles derived from the actual ONNX ports.""" + 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) + 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) -> str | None: """Emit ``tokenizer.json`` for a text-conditioned diffusion package. @@ -451,6 +469,7 @@ def write_onnx_genai_config( activation_dtype=_activation_dtype_tag(resolved_config), **kwargs, ) + _add_explicit_io_to_file(path, pkg, resolved_config) artifacts = {"inference_metadata": path} tokenizer_path = _write_hf_tokenizer(output_dir, source) if tokenizer_path is not None: @@ -467,6 +486,7 @@ def write_onnx_genai_config( activation_dtype=_activation_dtype_tag(resolved_config), **kwargs, ) + _add_explicit_io_to_file(path, pkg, resolved_config) artifacts = {"inference_metadata": path} tokenizer_path = _write_hf_tokenizer(output_dir, source) if tokenizer_path is not None: @@ -495,6 +515,7 @@ def write_onnx_genai_config( 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) if tokenizer_path is not None: @@ -520,6 +541,7 @@ def write_onnx_genai_config( path = write_decoder_metadata( output_dir, config=resolved_config, kv_native_dtype=kv_native_dtype ) + _add_explicit_io_to_file(path, pkg, resolved_config) artifacts = {"inference_metadata": path} tokenizer_path = _write_hf_tokenizer(output_dir, source) 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 e29f550a..71f17499 100644 --- a/src/mobius/integrations/onnx_genai/auto_export_test.py +++ b/src/mobius/integrations/onnx_genai/auto_export_test.py @@ -297,7 +297,9 @@ def test_dispatch_speech_to_text_pipeline(tmp_path): pkg = _EncoderDecoderPkg( { "encoder": _FakeModel(["input_features"], ["encoder_hidden_states"]), - "decoder": _FakeModel(["decoder_input_ids", "encoder_hidden_states"]), + "decoder": _FakeModel( + ["decoder_input_ids", "encoder_hidden_states"], ["logits"] + ), } ) artifacts = write_onnx_genai_config(pkg, str(tmp_path), kv_native_dtype="bf16") @@ -307,14 +309,14 @@ def test_dispatch_speech_to_text_pipeline(tmp_path): assert metadata["kv_cache"] == {"native_dtype": "bfloat16"} pipeline = metadata["pipeline"] - assert pipeline["models"] == { - "encoder": {"filename": "encoder/model.onnx", "type": "encoder"}, - "decoder": { - "filename": "decoder/model.onnx", - "type": "decoder", - "tokenizer": "tokenizer.json", - }, - } + 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", @@ -403,7 +405,9 @@ def test_dispatch_multi_decoder_tts_with_pre_embedder(tmp_path): pkg = _TTSPkg( { "talker": _FakeModel(["inputs_embeds"], ["logits", "last_hidden_state"]), - "code_predictor": _FakeModel(["inputs_embeds"], ["logits"]), + "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"] @@ -425,6 +429,7 @@ def test_dispatch_multi_decoder_tts_with_pre_embedder(tmp_path): } 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 diff --git a/src/mobius/integrations/onnx_genai/inference_metadata.py b/src/mobius/integrations/onnx_genai/inference_metadata.py index a8c98bc2..56df33b2 100644 --- a/src/mobius/integrations/onnx_genai/inference_metadata.py +++ b/src/mobius/integrations/onnx_genai/inference_metadata.py @@ -578,8 +578,12 @@ def _processor_values(source: str | None, config: Any) -> dict[str, Any]: _STATE_INPUT = re.compile( r"^past_key_values\.(?P\d+)\." + r"(?:(?Pself|cross)\.)?" r"(?Pkey|value|conv_state|recurrent_state|ssm_state)$" ) +_STATIC_CACHE_PORT = re.compile( + r"^(?Pupdated_)?(?Pkey|value)_cache\.(?P\d+)$" +) _REPLACE_ROLES = { "lightning_attention": {"recurrent_state"}, "linear_attention": {"conv_state", "recurrent_state"}, @@ -594,12 +598,14 @@ def _state_and_kv_pairs( decoder_inputs: list[_Port], decoder_outputs: list[_Port], config: Any, -) -> tuple[list[str], list[str], list[dict[str, str]]]: +) -> tuple[list[str], list[str], list[str], list[str], list[dict[str, str]]]: """Pair decoder state by declared port role and config layer type.""" outputs = {port.name: port for port in decoder_outputs} layer_types = getattr(config, "layer_types", None) kv_inputs: list[str] = [] kv_outputs: list[str] = [] + cross_kv_inputs: list[str] = [] + cross_kv_outputs: list[str] = [] state_pairs: list[dict[str, str]] = [] consumed_outputs: set[str] = set() for input_port in decoder_inputs: @@ -612,8 +618,9 @@ def _state_and_kv_pairs( "with declared state port roles or register this decoder signature." ) layer = int(match.group("layer")) + scope = match.group("scope") role = match.group("role") - output_name = f"present.{layer}.{role}" + output_name = f"present.{layer}.{scope + '.' if scope else ''}{role}" output_port = outputs.get(output_name) if output_port is None: raise ValueError( @@ -647,8 +654,12 @@ def _state_and_kv_pairs( "append state. Regenerate the graph/config pair or register the " "decoder state contract explicitly." ) - kv_inputs.append(input_port.name) - kv_outputs.append(output_name) + if scope == "cross": + cross_kv_inputs.append(input_port.name) + cross_kv_outputs.append(output_name) + else: + kv_inputs.append(input_port.name) + kv_outputs.append(output_name) else: allowed = _REPLACE_ROLES.get(layer_type or "", set()) if role not in allowed: @@ -679,7 +690,52 @@ def _state_and_kv_pairs( f"{unpaired}. Regenerate the package with present.. outputs " "matching declared past_key_values inputs, or register an explicit mapping." ) - return kv_inputs, kv_outputs, state_pairs + return kv_inputs, kv_outputs, cross_kv_inputs, cross_kv_outputs, state_pairs + + +def _static_cache_io( + decoder_inputs: list[_Port], + decoder_outputs: list[_Port], +) -> dict[str, Any] | None: + """Return the explicit TensorScatter static-cache ABI from exported ports.""" + inputs: dict[tuple[int, str], str] = {} + outputs: dict[tuple[int, str], str] = {} + for port in decoder_inputs: + match = _STATIC_CACHE_PORT.fullmatch(port.name) + if match is not None and match.group("updated") is None: + inputs[(int(match.group("layer")), match.group("role"))] = port.name + for port in decoder_outputs: + match = _STATIC_CACHE_PORT.fullmatch(port.name) + if match is not None and match.group("updated") is not None: + outputs[(int(match.group("layer")), match.group("role"))] = port.name + if not inputs and not outputs: + return None + + layers = sorted({layer for layer, _ in inputs} | {layer for layer, _ in outputs}) + missing = [ + f"{kind}.{layer}.{role}" + for layer in layers + for role in ("key", "value") + for kind, ports in (("input", inputs), ("output", outputs)) + if (layer, role) not in ports + ] + input_names = {port.name for port in decoder_inputs} + for control in ("write_indices", "nonpad_kv_seqlen"): + if control not in input_names: + missing.append(f"input.{control}") + if missing: + raise ValueError( + "Cannot emit model.io.static_cache because the exported TensorScatter " + f"ABI is incomplete: {missing}" + ) + return { + "write_indices_input": "write_indices", + "kv_sequence_length_input": "nonpad_kv_seqlen", + "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], + "value_cache_outputs": [outputs[(layer, "value")] for layer in layers], + } @dataclasses.dataclass(frozen=True) @@ -760,10 +816,21 @@ def _decoder_io( io: dict[str, Any] = { "inputs": [_port_metadata(port) for port in inputs], "outputs": [_port_metadata(port) for port in outputs], + "kv_ownership": "owned", } routed = [port for port in inputs if port.name in routed_inputs] embedded = next((port for port in routed if _is_float(port) and port.rank == 3), None) + if embedded is None: + embedded = _select_one( + inputs, + lambda port: ( + _is_float(port) + and port.rank == 3 + and port.name != "encoder_hidden_states" + and _STATIC_CACHE_PORT.fullmatch(port.name) is None + ), + ) if embedded is not None: io["inputs_embeds_input"] = embedded.name io["sequence_source"] = "inputs_embeds" @@ -778,7 +845,15 @@ def _decoder_io( if position is not None: io["position_ids_input"] = position.name - token = input_by_name.get("input_ids") + token = input_by_name.get("input_ids") or _select_one( + inputs, + lambda port: ( + _is_integer(port) + and port.rank == 2 + and port.name not in {"attention_mask", "position_ids"} + and _STATE_INPUT.fullmatch(port.name) is None + ), + ) if token is not None: io["token_input"] = token.name io.setdefault("sequence_source", "token_ids") @@ -791,21 +866,73 @@ def _decoder_io( "logits role or register an explicit decoder I/O contract." ) io["logits_output"] = logits.name - + encoder_hidden_states = input_by_name.get("encoder_hidden_states") + if encoder_hidden_states is not None: + io["encoder_hidden_states_input"] = encoder_hidden_states.name + + static_cache = _static_cache_io(inputs, outputs) + if static_cache is not None: + io["static_cache"] = static_cache + static_names = ( + { + static_cache["write_indices_input"], + static_cache["kv_sequence_length_input"], + *static_cache["key_cache_inputs"], + *static_cache["value_cache_inputs"], + *static_cache["key_cache_outputs"], + *static_cache["value_cache_outputs"], + } + if static_cache is not None + else set() + ) core_inputs = routed_inputs | { - port.name for port in (attention_mask, position, token) if port is not None + port.name + for port in (attention_mask, position, token, embedded, encoder_hidden_states) + if port is not None } - state_inputs = [port for port in inputs if port.name not in core_inputs] - state_outputs = [port for port in outputs if port.name != logits.name] - kv_inputs, kv_outputs, state_pairs = _state_and_kv_pairs( + state_inputs = [ + port + for port in inputs + if port.name not in core_inputs + and port.name not in static_names + and _STATE_INPUT.fullmatch(port.name) is not None + ] + state_outputs = [ + port + for port in outputs + if port.name != logits.name + and port.name not in static_names + and port.name.startswith("present.") + ] + ( + kv_inputs, + kv_outputs, + cross_kv_inputs, + cross_kv_outputs, + state_pairs, + ) = _state_and_kv_pairs( state_inputs, state_outputs, config ) if kv_inputs: io["kv_inputs"] = kv_inputs io["kv_outputs"] = kv_outputs io["kv_update"] = "append" + if cross_kv_inputs: + io["cross_kv_inputs"] = cross_kv_inputs + io["cross_kv_outputs"] = cross_kv_outputs if state_pairs: io["state_pairs"] = state_pairs + consumed_outputs = set(kv_outputs) | set(cross_kv_outputs) + hidden_outputs = [ + port + for port in outputs + if port.name != logits.name + and port.name not in static_names + and port.name not in consumed_outputs + and _is_float(port) + ] + if len(hidden_outputs) == 1: + io["hidden_output"] = hidden_outputs[0].name positions = _positions_from_registry(position, config) if position is not None else None return io, positions @@ -898,6 +1025,25 @@ def _input_source_map( "from": f"{decoder_name}.{pair['output']}", "update": pair["update"], } + static_cache = decoder_io.get("static_cache") + if static_cache is not None: + sources[f"{decoder_name}.{static_cache['write_indices_input']}"] = { + "kind": "generated", + "generator": "static_cache_write_indices", + } + sources[f"{decoder_name}.{static_cache['kv_sequence_length_input']}"] = { + "kind": "generated", + "generator": "kv_sequence_length", + } + for input_name, output_name in zip( + static_cache["key_cache_inputs"] + static_cache["value_cache_inputs"], + static_cache["key_cache_outputs"] + static_cache["value_cache_outputs"], + ): + sources[f"{decoder_name}.{input_name}"] = { + "kind": "stateful", + "from": f"{decoder_name}.{output_name}", + "update": "shared_buffer", + } return sources @@ -1082,6 +1228,85 @@ 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": + inner_name = strategy.get("inner") + inner_io = component_ios.get(inner_name, {}) + hidden_output = inner_io.get("hidden_output") + if not hidden_output: + raise ValueError( + "Cannot emit pipeline.strategy.inner_embedding_output: the inner " + f"decoder {inner_name!r} has no unique non-logits float output" + ) + strategy["inner_embedding_output"] = hidden_output + for stage in strategy.get("stages", []): + nested = stage.get("strategy") + if isinstance(nested, dict): + annotate_strategy(nested) + + strategy = pipeline.get("strategy") + if isinstance(strategy, dict): + annotate_strategy(strategy) + if "model" in metadata: + decoder_names = [ + name for name, model in models.items() if model.get("type") == "decoder" + ] + if decoder_names: + metadata["model"]["io"] = component_ios[decoder_names[0]] + return metadata + + def _topological_order( names: Iterable[str], edges: list[dict[str, Any]], @@ -2046,6 +2271,7 @@ def build_tts_pipeline_metadata( 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. @@ -2128,6 +2354,7 @@ def build_tts_pipeline_metadata( "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", diff --git a/src/mobius/integrations/onnx_genai/inference_metadata_test.py b/src/mobius/integrations/onnx_genai/inference_metadata_test.py index 35c6392b..5cc6da77 100644 --- a/src/mobius/integrations/onnx_genai/inference_metadata_test.py +++ b/src/mobius/integrations/onnx_genai/inference_metadata_test.py @@ -24,6 +24,7 @@ ) from mobius.integrations.onnx_genai.inference_metadata import ( SchedulerConfig, + add_explicit_package_io, build_diffusion_pipeline_metadata, build_language_diffusion_pipeline_metadata, build_multimodal_pipeline_metadata, @@ -231,6 +232,135 @@ def _decoder_model( return _model("decoder", inputs, output_specs) +def _static_cache_decoder_model() -> ir.Model: + inputs = [ + _value("input_ids", ir.DataType.INT64, ["batch", "sequence"]), + _value("attention_mask", ir.DataType.INT64, ["batch", 32]), + _value("position_ids", ir.DataType.INT64, ["batch", "sequence"]), + _value("key_cache.1", ir.DataType.FLOAT, ["batch", 32, 16]), + _value("value_cache.1", ir.DataType.FLOAT, ["batch", 32, 16]), + _value("key_cache.3", ir.DataType.FLOAT, ["batch", 32, 16]), + _value("value_cache.3", ir.DataType.FLOAT, ["batch", 32, 16]), + _value("write_indices", ir.DataType.INT64, ["batch"]), + _value("nonpad_kv_seqlen", ir.DataType.INT64, ["batch"]), + ] + outputs = [ + ("logits", ir.DataType.FLOAT, ["batch", "sequence", 128]), + ("updated_key_cache.1", ir.DataType.FLOAT, ["batch", 32, 16]), + ("updated_value_cache.1", ir.DataType.FLOAT, ["batch", 32, 16]), + ("updated_key_cache.3", ir.DataType.FLOAT, ["batch", 32, 16]), + ("updated_value_cache.3", ir.DataType.FLOAT, ["batch", 32, 16]), + ] + 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" + + def _native_package( vision_encoder: ir.Model, config: _VlmConfig, @@ -1651,6 +1781,7 @@ def test_minimal_nested_autoregressive_with_pre_embedder(self): 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 From 50eb36b8330a2296583b50ba2a14eef69a22e603 Mon Sep 17 00:00:00 2001 From: justinchuby Date: Wed, 29 Jul 2026 23:52:58 +0000 Subject: [PATCH 2/3] Fix explicit metadata loop-state emission Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- .../onnx_genai/auto_export_test.py | 8 +- .../onnx_genai/inference_metadata.py | 28 ++-- .../onnx_genai/inference_metadata_test.py | 152 ++++++++++++++++++ 3 files changed, 170 insertions(+), 18 deletions(-) diff --git a/src/mobius/integrations/onnx_genai/auto_export_test.py b/src/mobius/integrations/onnx_genai/auto_export_test.py index 71f17499..45c85fd9 100644 --- a/src/mobius/integrations/onnx_genai/auto_export_test.py +++ b/src/mobius/integrations/onnx_genai/auto_export_test.py @@ -297,9 +297,7 @@ def test_dispatch_speech_to_text_pipeline(tmp_path): pkg = _EncoderDecoderPkg( { "encoder": _FakeModel(["input_features"], ["encoder_hidden_states"]), - "decoder": _FakeModel( - ["decoder_input_ids", "encoder_hidden_states"], ["logits"] - ), + "decoder": _FakeModel(["decoder_input_ids", "encoder_hidden_states"], ["logits"]), } ) artifacts = write_onnx_genai_config(pkg, str(tmp_path), kv_native_dtype="bf16") @@ -405,9 +403,7 @@ def test_dispatch_multi_decoder_tts_with_pre_embedder(tmp_path): pkg = _TTSPkg( { "talker": _FakeModel(["inputs_embeds"], ["logits", "last_hidden_state"]), - "code_predictor": _FakeModel( - ["inputs_embeds"], ["logits", "codec_embeddings"] - ), + "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"] diff --git a/src/mobius/integrations/onnx_genai/inference_metadata.py b/src/mobius/integrations/onnx_genai/inference_metadata.py index 56df33b2..bc801d5a 100644 --- a/src/mobius/integrations/onnx_genai/inference_metadata.py +++ b/src/mobius/integrations/onnx_genai/inference_metadata.py @@ -910,9 +910,7 @@ def _decoder_io( cross_kv_inputs, cross_kv_outputs, state_pairs, - ) = _state_and_kv_pairs( - state_inputs, state_outputs, config - ) + ) = _state_and_kv_pairs(state_inputs, state_outputs, config) if kv_inputs: io["kv_inputs"] = kv_inputs io["kv_outputs"] = kv_outputs @@ -1010,15 +1008,19 @@ def _input_source_map( "generator": generator, } - for input_name, output_name in zip( - decoder_io.get("kv_inputs", []), - decoder_io.get("kv_outputs", []), + for input_field, output_field in ( + ("kv_inputs", "kv_outputs"), + ("cross_kv_inputs", "cross_kv_outputs"), ): - sources[f"{decoder_name}.{input_name}"] = { - "kind": "stateful", - "from": f"{decoder_name}.{output_name}", - "update": decoder_io.get("kv_update", "append"), - } + for input_name, output_name in zip( + decoder_io.get(input_field, []), + decoder_io.get(output_field, []), + ): + sources[f"{decoder_name}.{input_name}"] = { + "kind": "stateful", + "from": f"{decoder_name}.{output_name}", + "update": decoder_io.get("kv_update", "append"), + } for pair in decoder_io.get("state_pairs", []): sources[f"{decoder_name}.{pair['input']}"] = { "kind": "stateful", @@ -1280,7 +1282,9 @@ def add_explicit_package_io( component_ios[name] = io def annotate_strategy(strategy: dict[str, Any]) -> None: - if strategy.get("kind") == "nested_autoregressive": + 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") diff --git a/src/mobius/integrations/onnx_genai/inference_metadata_test.py b/src/mobius/integrations/onnx_genai/inference_metadata_test.py index 5cc6da77..247f8571 100644 --- a/src/mobius/integrations/onnx_genai/inference_metadata_test.py +++ b/src/mobius/integrations/onnx_genai/inference_metadata_test.py @@ -24,6 +24,9 @@ ) from mobius.integrations.onnx_genai.inference_metadata import ( SchedulerConfig, + _decoder_io, + _input_source_map, + _port, add_explicit_package_io, build_diffusion_pipeline_metadata, build_language_diffusion_pipeline_metadata, @@ -360,6 +363,155 @@ def test_emits_explicit_encoder_prompt_role(self, encoder_input, role_field): assert decoder_io["encoder_hidden_states_input"] == "encoder_hidden_states" assert decoder_io["logits_output"] == "logits" + def test_cross_attention_cache_inputs_are_loop_state(self): + inputs = [ + _value("input_ids", ir.DataType.INT64, ["batch", "sequence"]), + _value( + "encoder_hidden_states", + ir.DataType.FLOAT, + ["batch", "encoder_sequence", 64], + ), + _value( + "past_key_values.0.self.key", + ir.DataType.FLOAT, + ["batch", 2, "past_sequence", 8], + ), + _value( + "past_key_values.0.self.value", + ir.DataType.FLOAT, + ["batch", 2, "past_sequence", 8], + ), + _value( + "past_key_values.0.cross.key", + ir.DataType.FLOAT, + ["batch", 2, "encoder_sequence", 8], + ), + _value( + "past_key_values.0.cross.value", + ir.DataType.FLOAT, + ["batch", 2, "encoder_sequence", 8], + ), + ] + outputs = [ + ("logits", ir.DataType.FLOAT, ["batch", "sequence", 128]), + ( + "present.0.self.key", + ir.DataType.FLOAT, + ["batch", 2, "total_sequence", 8], + ), + ( + "present.0.self.value", + ir.DataType.FLOAT, + ["batch", 2, "total_sequence", 8], + ), + ( + "present.0.cross.key", + ir.DataType.FLOAT, + ["batch", 2, "encoder_sequence", 8], + ), + ( + "present.0.cross.value", + ir.DataType.FLOAT, + ["batch", 2, "encoder_sequence", 8], + ), + ] + decoder = _model("decoder", inputs, outputs) + decoder_io, _ = _decoder_io(decoder, {"encoder_hidden_states"}, _VlmConfig()) + ports = { + "decoder": { + "inputs": [_port(value) for value in decoder.graph.inputs], + "outputs": [_port(value) for value in decoder.graph.outputs], + } + } + models = {"decoder": {"io": decoder_io}} + sources = _input_source_map( + ports=ports, + dataflow=[], + models=models, + decoder_name="decoder", + image_endpoints=set(), + ) + assert sources["decoder.past_key_values.0.cross.key"] == { + "kind": "stateful", + "from": "decoder.present.0.cross.key", + "update": "append", + } + assert sources["decoder.past_key_values.0.cross.value"] == { + "kind": "stateful", + "from": "decoder.present.0.cross.value", + "update": "append", + } + + @staticmethod + def _nested_package(inner_outputs): + talker = _model( + "talker", + [_value("inputs_embeds", ir.DataType.FLOAT, ["batch", "sequence", 64])], + [("logits", ir.DataType.FLOAT, ["batch", "sequence", 128])], + ) + code_predictor = _model( + "code_predictor", + [_value("inputs_embeds", ir.DataType.FLOAT, ["batch", "sequence", 64])], + inner_outputs, + ) + return {"talker": talker, "code_predictor": code_predictor} + + @staticmethod + def _nested_metadata(inner_embedding_output=None): + strategy = { + "kind": "nested_autoregressive", + "outer": "talker", + "inner": "code_predictor", + } + if inner_embedding_output is not None: + strategy["inner_embedding_output"] = inner_embedding_output + return { + "pipeline": { + "models": { + "talker": {"type": "decoder"}, + "code_predictor": {"type": "decoder"}, + }, + "dataflow": [], + "strategy": strategy, + } + } + + 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, From e609609a4c896fc8cf69b6b15e6ac641d48a413f Mon Sep 17 00:00:00 2001 From: Justin Chu Date: Wed, 29 Jul 2026 21:50:42 -0700 Subject: [PATCH 3/3] Potential fix for pull request finding Co-authored-by: Copilot Autofix powered by AI <175728472+Copilot@users.noreply.github.com> Signed-off-by: Justin Chu --- .../integrations/onnx_genai/inference_metadata.py | 12 +++++++++++- 1 file changed, 11 insertions(+), 1 deletion(-) diff --git a/src/mobius/integrations/onnx_genai/inference_metadata.py b/src/mobius/integrations/onnx_genai/inference_metadata.py index bc801d5a..9be8a2d9 100644 --- a/src/mobius/integrations/onnx_genai/inference_metadata.py +++ b/src/mobius/integrations/onnx_genai/inference_metadata.py @@ -820,7 +820,17 @@ def _decoder_io( } routed = [port for port in inputs if port.name in routed_inputs] - embedded = next((port for port in routed if _is_float(port) and port.rank == 3), None) + embedded = next( + ( + port + for port in routed + if _is_float(port) + and port.rank == 3 + and port.name != "encoder_hidden_states" + and _STATIC_CACHE_PORT.fullmatch(port.name) is None + ), + None, + ) if embedded is None: embedded = _select_one( inputs,