diff --git a/src/mobius/integrations/onnx_genai/inference_metadata.py b/src/mobius/integrations/onnx_genai/inference_metadata.py index c2eb0a5dc..872269d3d 100644 --- a/src/mobius/integrations/onnx_genai/inference_metadata.py +++ b/src/mobius/integrations/onnx_genai/inference_metadata.py @@ -1395,16 +1395,19 @@ def validate_executable_closure(pkg: Any, metadata: dict[str, Any]) -> None: ) -#: 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. +#: Symbolic dimension Mobius emits for every request-aligned ONNX port. 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} + axes = [ + axis + for axis, dimension in enumerate(shape or []) + if str(dimension) in _BATCH_DIMENSION_NAMES + ] + if len(axes) == 1: + return {"kind": "request_aligned", "axis": axes[0]} return None @@ -1519,11 +1522,11 @@ def tensor_contract(value: Any) -> dict[str, Any]: def declare_request_alignment(workflow: dict[str, Any]) -> None: - """Stamp the request-aligned row axis onto every batch-leading contract. + """Stamp the request axis named by exactly one batch symbol. 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, + permutation to every request-aligned tensor. A contract whose batch axis + 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 @@ -1534,8 +1537,8 @@ 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} + if layout := request_batch_layout(shape): + contract["batch_layout"] = layout for section in ("inputs", "outputs", "state"): for declaration in (workflow.get(section) or {}).values(): diff --git a/src/mobius/integrations/onnx_genai/inference_metadata_test.py b/src/mobius/integrations/onnx_genai/inference_metadata_test.py index 087def81d..c667bdc0c 100644 --- a/src/mobius/integrations/onnx_genai/inference_metadata_test.py +++ b/src/mobius/integrations/onnx_genai/inference_metadata_test.py @@ -39,6 +39,7 @@ load_diffusers_scheduler_config, load_diffusers_vae_scaling_factor, published_value_references, + request_batch_layout, validate_executable_closure, write_diffusion_pipeline_metadata, write_mtp_speculator_metadata, @@ -49,6 +50,19 @@ ) +def test_request_batch_layout_follows_the_unique_batch_symbol(): + assert request_batch_layout(["batch", "sequence"]) == { + "kind": "request_aligned", + "axis": 0, + } + assert request_batch_layout([3, "batch", "sequence"]) == { + "kind": "request_aligned", + "axis": 1, + } + assert request_batch_layout(["batch", "batch"]) is None + assert request_batch_layout([3, "sequence"]) is None + + def test_ort_extensions_processor_config_supplies_structural_values(tmp_path): (tmp_path / "processor_config.json").write_text( json.dumps( diff --git a/src/mobius/integrations/onnx_genai/workflow_metadata.py b/src/mobius/integrations/onnx_genai/workflow_metadata.py index 8c2442aa9..db6cd3b37 100644 --- a/src/mobius/integrations/onnx_genai/workflow_metadata.py +++ b/src/mobius/integrations/onnx_genai/workflow_metadata.py @@ -239,6 +239,157 @@ def _component( return declaration +def _declare_batch_capacities( + metadata: dict[str, Any], + component_names: Sequence[str], +) -> None: + """Assert independent-row execution for explicitly audited components. + + Presence is the semantic permission; it is never inferred from a dynamic + leading axis. For synthesized policy graphs, whose port contracts are part + of the workflow dataflow, every non-request symbolic input extent is also + declared uniform. Exported model graphs remain authoritative for their own + port geometry and therefore contribute an empty capacity declaration. + """ + workflow = metadata["pipeline"]["workflow"] + components = workflow["components"] + declared = False + for name in component_names: + component = components.get(name) + if component is None: + continue + uniform_dimensions: list[str] = [] + for contract in component.get("ports", {}).get("inputs", {}).values(): + layout = contract.get("batch_layout", {}) + request_axis = ( + layout.get("axis") + if layout.get("kind") in {"request_aligned", "request_expanded"} + else None + ) + for axis, dimension in enumerate(contract.get("shape") or []): + if ( + axis != request_axis + and isinstance(dimension, str) + and dimension not in uniform_dimensions + ): + uniform_dimensions.append(dimension) + capacity: dict[str, Any] = {} + if uniform_dimensions: + capacity["uniform_dimensions"] = uniform_dimensions + component["batch_capacity"] = capacity + declared = True + if declared: + metadata["schema_version"] = "v1.1" + + +def _preserve_request_expansion( + metadata: dict[str, Any], + component_names: Sequence[str], + *, + factor: int, +) -> None: + """Keep a fixed internal row expansion distinct from request batching.""" + workflow = metadata["pipeline"]["workflow"] + for name in component_names: + component = workflow["components"][name] + for side in ("inputs", "outputs"): + for contract in component.get("ports", {}).get(side, {}).values(): + shape = contract.get("shape") or [] + if shape and shape[0] == "batch": + contract["batch_layout"] = { + "kind": "request_expanded", + "axis": 0, + "factor": factor, + } + + +_DECODER_BATCH_COMPONENTS = ( + "model", + "token_sampler", + "termination", + "token_state_update", + "last_token_logits", + "decoder_state_initializer", + "decoder_step_update", + "cache_length_update", + "termination_batch_initializer", + "token_to_slot", + "generated_length_update", + "state_init", + "step_update", + "next_token", + "loop_continue", +) + +_DIFFUSION_BATCH_COMPONENTS = ( + "text_encoder", + "denoiser", + "vae_decoder", + "image_output_clamp", + "solver_step", + "continue_predicate", + "model_input_scale", + "diffusion_schedule", + "diffusion_timesteps", + "schedule_lookup", + "tensor_scale", + "initial_state_scale", + "decoder_input_scale", + "history_initializer", + "guidance_combine", + "latent_row_shape", + "latent_noise", +) + +_TTS_BATCH_COMPONENTS = ( + "talker", + "code_predictor", + "embedding", + "talker_step_embedder", + "talker_prefill_embedder", + "code_predictor_prefill", + "code_predictor_step_embedder", + "code_predictor_indices", + "talker_text_step", + "codec", + "last_token_logits", + "setup_talker_sampler", + "setup_predictor_sampler", + "talker_sampler", + "predictor_prefill_sampler", + "predictor_body_sampler", + "continue_predicate", + "tts_state_initializer", + "token_to_slot", + "code_frame_update", + "code_history_append", + "cache_length_update", + "talker_state_initializer", + "predictor_state_initializer", + "talker_step_update", + "predictor_step_update", + "codec_layout", +) + +_VIDEO_BATCH_COMPONENTS = ( + "transformer", + "vae_decoder", + "model_input", + "solver_step", + "continue_predicate", + "video_latent_init", + "schedule_history_append", + "video_latent_permute", + "video_latent_unscale", + "video_decode_chunks", + "video_decode_chunk", + "video_conv_cache_init", + "diffusion_schedule", + "diffusion_timesteps", + "schedule_lookup", +) + + def _grammar_adapter_component(action: str) -> dict[str, Any]: """Declare one action of the versioned grammar-guidance adapter ABI.""" @@ -2013,6 +2164,14 @@ def invoke_model( "pipeline": {"workflow": _publish_workflow_v1(workflow)}, } add_policy_components_to_workflow(metadata, pkg) + expansion = metadata["pipeline"]["workflow"]["inputs"]["request.prompt_tokens"][ + "contract" + ]["batch_layout"]["factor"] + _preserve_request_expansion( + metadata, + ("global_initializer", "global_step_update"), + factor=expansion, + ) return metadata @@ -3794,7 +3953,9 @@ def build_tts_workflow_metadata(pkg: Any, config: Any) -> dict[str, Any]: "talker_text_step", } if real_transition_components <= set(pkg.keys()): - return _build_real_tts_workflow_metadata(pkg, config) + metadata = _build_real_tts_workflow_metadata(pkg, config) + _declare_batch_capacities(metadata, _TTS_BATCH_COMPONENTS) + return metadata required = { "talker", "code_predictor", @@ -4243,6 +4404,7 @@ def bind_outputs(values: Any, bound: dict[str, str], prefix: str) -> dict[str, s "pipeline": {"workflow": _publish_workflow_v1(workflow)}, } add_policy_components_to_workflow(metadata, pkg) + _declare_batch_capacities(metadata, _TTS_BATCH_COMPONENTS) return metadata @@ -4838,6 +5000,7 @@ def denoiser_call(conditioning: str | None, estimate: str) -> dict[str, Any]: "value": "denoiser.estimate", "output": "noise_estimate", "mode": "append", + "axis": 3, "effect_name": "emit", "effect": _effect("emit.0", "emit.1"), }, @@ -4846,6 +5009,7 @@ def denoiser_call(conditioning: str | None, estimate: str) -> dict[str, Any]: "value": "latent.body", "output": "latent_trajectory", "mode": "append", + "axis": 3, "effect_name": "emit", "effect": _effect("emit.1", "emit.2"), }, @@ -4951,6 +5115,7 @@ def denoiser_call(conditioning: str | None, estimate: str) -> dict[str, Any]: "pipeline": {"workflow": _publish_workflow_v1(workflow)}, } add_policy_components_to_workflow(metadata, pkg) + _declare_batch_capacities(metadata, _DIFFUSION_BATCH_COMPONENTS) return metadata @@ -6063,6 +6228,7 @@ def build_video_diffusion_workflow_metadata( "pipeline": {"workflow": _publish_workflow_v1(workflow)}, } add_policy_components_to_workflow(metadata, pkg) + _declare_batch_capacities(metadata, _VIDEO_BATCH_COMPONENTS) return metadata @@ -8071,7 +8237,11 @@ 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") - return _build_autoregressive_workflow_metadata(pkg, config, sampler=sampler, source=source) + metadata = _build_autoregressive_workflow_metadata( + pkg, config, sampler=sampler, source=source + ) + _declare_batch_capacities(metadata, _DECODER_BATCH_COMPONENTS) + return metadata def _build_autoregressive_workflow_metadata( diff --git a/tests/canonical_workflow_contract_test.py b/tests/canonical_workflow_contract_test.py index 1e1ffe077..d31ecda85 100644 --- a/tests/canonical_workflow_contract_test.py +++ b/tests/canonical_workflow_contract_test.py @@ -831,6 +831,114 @@ def _fixture_packages() -> list[str]: ) +_MULTI_REQUEST_COMPONENTS = { + "adapter": {"overlay"}, + "decoder": { + "model", + "token_sampler", + "termination", + "token_state_update", + "last_token_logits", + "decoder_state_initializer", + "decoder_step_update", + "cache_length_update", + "termination_batch_initializer", + "token_to_slot", + "generated_length_update", + }, + "diffusion": { + "text_encoder", + "denoiser", + "vae_decoder", + "image_output_clamp", + "solver_step", + "continue_predicate", + "model_input_scale", + "diffusion_schedule", + "diffusion_timesteps", + "schedule_lookup", + "tensor_scale", + "initial_state_scale", + }, + "diffusion_guided": { + "text_encoder", + "denoiser", + "vae_decoder", + "image_output_clamp", + "solver_step", + "continue_predicate", + "diffusion_schedule", + "diffusion_timesteps", + "schedule_lookup", + "tensor_scale", + "decoder_input_scale", + "history_initializer", + "guidance_combine", + "latent_row_shape", + "latent_noise", + }, + "static_cache": { + "model", + "token_sampler", + "termination", + "token_state_update", + "last_token_logits", + "decoder_state_initializer", + "decoder_step_update", + "cache_length_update", + "termination_batch_initializer", + "token_to_slot", + "generated_length_update", + }, + "tts": { + "talker", + "code_predictor", + "embedding", + "talker_step_embedder", + "talker_prefill_embedder", + "code_predictor_prefill", + "code_predictor_step_embedder", + "code_predictor_indices", + "talker_text_step", + "codec", + "last_token_logits", + "setup_talker_sampler", + "setup_predictor_sampler", + "talker_sampler", + "predictor_prefill_sampler", + "predictor_body_sampler", + "continue_predicate", + "tts_state_initializer", + "token_to_slot", + "code_frame_update", + "code_history_append", + "cache_length_update", + "talker_state_initializer", + "predictor_state_initializer", + "talker_step_update", + "predictor_step_update", + "codec_layout", + }, + "video": { + "transformer", + "vae_decoder", + "model_input", + "solver_step", + "continue_predicate", + "video_latent_init", + "schedule_history_append", + "video_latent_permute", + "video_latent_unscale", + "video_decode_chunks", + "video_decode_chunk", + "video_conv_cache_init", + "diffusion_schedule", + "diffusion_timesteps", + "schedule_lookup", + }, +} + + @pytest.mark.parametrize( "directory", _fixture_packages(), ids=lambda path: os.path.basename(path) ) @@ -859,6 +967,46 @@ def test_no_fixture_contains_retired_batching_declarations(self, directory): def test_fixture_validates_against_current_onnx_genai_schema(self, directory): jsonschema.validate(self._metadata(directory), ONNX_GENAI_SCHEMA) + def test_fixture_round_trips_through_yaml_and_schema(self, directory): + metadata = self._metadata(directory) + round_tripped = yaml.safe_load(yaml.safe_dump(metadata, sort_keys=False)) + assert round_tripped == metadata + jsonschema.validate(round_tripped, ONNX_GENAI_SCHEMA) + + def test_multi_request_components_declare_capacity(self, directory): + package = os.path.basename(directory) + expected = _MULTI_REQUEST_COMPONENTS.get(package) + if expected is None: + return + components = self._metadata(directory)["pipeline"]["workflow"]["components"] + assert { + name for name, component in components.items() if "batch_capacity" in component + } == expected + + def test_unproven_encoder_capacity_remains_absent(self, directory): + if os.path.basename(directory) not in { + "esm2_protein_embeddings", + "protbert_protein_embeddings", + }: + return + encoder = self._metadata(directory)["pipeline"]["workflow"]["components"]["encoder"] + assert "batch_capacity" not in encoder + + def test_hierarchical_audio_preserves_its_internal_row_expansion(self, directory): + if os.path.basename(directory) != "hierarchical_audio": + return + components = self._metadata(directory)["pipeline"]["workflow"]["components"] + for name in ("global_initializer", "global_step_update"): + contracts = components[name]["ports"] + for contract in (*contracts["inputs"].values(), *contracts["outputs"].values()): + shape = contract.get("shape") or [] + if shape and shape[0] == "batch": + assert contract["batch_layout"] == { + "kind": "request_expanded", + "axis": 0, + "factor": 2, + } + def test_fixture_matches_regenerated_metadata( self, directory, materialized_workflow_packages ): diff --git a/tests/fixtures/onnx_genai_workflows/adapter/inference_metadata.yaml b/tests/fixtures/onnx_genai_workflows/adapter/inference_metadata.yaml index 86759a2da..6674454d2 100644 --- a/tests/fixtures/onnx_genai_workflows/adapter/inference_metadata.yaml +++ b/tests/fixtures/onnx_genai_workflows/adapter/inference_metadata.yaml @@ -1,4 +1,4 @@ -schema_version: v1 +schema_version: v1.1 pipeline: workflow: manifest: @@ -115,6 +115,7 @@ pipeline: kind: adapter abi: onnx-genai.parameter-overlay version: '1' + batch_capacity: {} ports: inputs: input: diff --git a/tests/fixtures/onnx_genai_workflows/decoder/inference_metadata.yaml b/tests/fixtures/onnx_genai_workflows/decoder/inference_metadata.yaml index c053a0fd8..f17cd0911 100644 --- a/tests/fixtures/onnx_genai_workflows/decoder/inference_metadata.yaml +++ b/tests/fixtures/onnx_genai_workflows/decoder/inference_metadata.yaml @@ -1,4 +1,4 @@ -schema_version: '1.0' +schema_version: v1.1 pipeline: workflow: manifest: @@ -343,6 +343,7 @@ pipeline: attention_mask: attention_mask position_ids: position_ids logits: logits + batch_capacity: {} token_sampler: implementation: kind: onnx @@ -459,6 +460,9 @@ pipeline: batching: per_row inactive_rows: preserve application_overridable: true + batch_capacity: + uniform_dimensions: + - vocabulary termination: implementation: kind: onnx @@ -549,6 +553,9 @@ pipeline: parameters: batching: per_row inactive_rows: preserve + batch_capacity: + uniform_dimensions: + - num_eos token_state_update: implementation: kind: onnx @@ -611,6 +618,7 @@ pipeline: parameters: batching: per_row inactive_rows: preserve + batch_capacity: {} last_token_logits: implementation: kind: onnx @@ -637,6 +645,10 @@ pipeline: batch_layout: kind: request_aligned axis: 0 + batch_capacity: + uniform_dimensions: + - sequence + - vocabulary decoder_state_initializer: implementation: kind: onnx @@ -736,6 +748,9 @@ pipeline: batch_layout: kind: request_aligned axis: 0 + batch_capacity: + uniform_dimensions: + - prompt_sequence decoder_step_update: implementation: kind: onnx @@ -779,6 +794,9 @@ pipeline: batch_layout: kind: request_aligned axis: 0 + batch_capacity: + uniform_dimensions: + - context cache_length_update: implementation: kind: onnx @@ -826,6 +844,7 @@ pipeline: batch_layout: kind: request_aligned axis: 0 + batch_capacity: {} termination_batch_initializer: implementation: kind: onnx @@ -896,6 +915,9 @@ pipeline: batch_layout: kind: request_aligned axis: 0 + batch_capacity: + uniform_dimensions: + - num_eos token_to_slot: implementation: kind: onnx @@ -920,6 +942,7 @@ pipeline: batch_layout: kind: request_aligned axis: 0 + batch_capacity: {} generated_length_update: implementation: kind: onnx @@ -967,6 +990,7 @@ pipeline: batch_layout: kind: request_aligned axis: 0 + batch_capacity: {} state: token: contract: diff --git a/tests/fixtures/onnx_genai_workflows/diffusion/inference_metadata.yaml b/tests/fixtures/onnx_genai_workflows/diffusion/inference_metadata.yaml index d3842b72c..ac63ec27c 100644 --- a/tests/fixtures/onnx_genai_workflows/diffusion/inference_metadata.yaml +++ b/tests/fixtures/onnx_genai_workflows/diffusion/inference_metadata.yaml @@ -1,4 +1,4 @@ -schema_version: v1 +schema_version: v1.1 pipeline: workflow: manifest: @@ -154,6 +154,7 @@ pipeline: roles: input_ids: token_ids encoder_hidden_states: encoder_hidden_states + batch_capacity: {} denoiser: implementation: kind: onnx @@ -161,10 +162,12 @@ pipeline: ports: roles: encoder_hidden_states: encoder_hidden_states + batch_capacity: {} vae_decoder: implementation: kind: onnx artifact: vae_decoder/model.onnx + batch_capacity: {} solver_step: implementation: kind: onnx @@ -227,6 +230,12 @@ pipeline: step: step schedule: schedule next_state: next_state + batch_capacity: + uniform_dimensions: + - channels + - height + - width + - schedule_length continue_predicate: implementation: kind: onnx @@ -247,6 +256,7 @@ pipeline: rank: 1 shape: - 1 + batch_capacity: {} model_input_scale: implementation: kind: onnx @@ -289,6 +299,12 @@ pipeline: batch_layout: kind: request_aligned axis: 0 + batch_capacity: + uniform_dimensions: + - channels + - height + - width + - schedule_length diffusion_schedule: implementation: kind: onnx @@ -301,6 +317,7 @@ pipeline: rank: 1 shape: - 31 + batch_capacity: {} diffusion_timesteps: implementation: kind: onnx @@ -313,6 +330,7 @@ pipeline: rank: 1 shape: - 30 + batch_capacity: {} schedule_lookup: implementation: kind: onnx @@ -341,6 +359,9 @@ pipeline: batch_layout: kind: request_aligned axis: 0 + batch_capacity: + uniform_dimensions: + - schedule_length tensor_scale: implementation: kind: onnx @@ -375,6 +396,11 @@ pipeline: batch_layout: kind: request_aligned axis: 0 + batch_capacity: + uniform_dimensions: + - channels + - height + - width initial_state_scale: implementation: kind: onnx @@ -387,6 +413,7 @@ pipeline: rank: 1 shape: - 1 + batch_capacity: {} image_output_clamp: implementation: kind: onnx @@ -416,6 +443,11 @@ pipeline: batch_layout: kind: request_aligned axis: 0 + batch_capacity: + uniform_dimensions: + - axis_1 + - axis_2 + - axis_3 state: latent_state: contract: @@ -517,10 +549,12 @@ pipeline: value: denoiser.estimate output: noise_estimate mode: append + axis: 3 - kind: emit value: latent.body output: latent_trajectory mode: append + axis: 3 - kind: invoke component: continue_predicate inputs: 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 1f187e726..f8d50ecd0 100644 --- a/tests/fixtures/onnx_genai_workflows/diffusion_guided/inference_metadata.yaml +++ b/tests/fixtures/onnx_genai_workflows/diffusion_guided/inference_metadata.yaml @@ -1,4 +1,4 @@ -schema_version: v1 +schema_version: v1.1 pipeline: workflow: manifest: @@ -214,6 +214,7 @@ pipeline: roles: input_ids: token_ids encoder_hidden_states: encoder_hidden_states + batch_capacity: {} denoiser: implementation: kind: onnx @@ -221,10 +222,12 @@ pipeline: ports: roles: encoder_hidden_states: encoder_hidden_states + batch_capacity: {} vae_decoder: implementation: kind: onnx artifact: vae_decoder/model.onnx + batch_capacity: {} solver_step: implementation: kind: onnx @@ -311,6 +314,12 @@ pipeline: schedule: schedule next_state: next_state next_history: next_history + batch_capacity: + uniform_dimensions: + - channels + - height + - width + - schedule_length continue_predicate: implementation: kind: onnx @@ -331,6 +340,7 @@ pipeline: rank: 1 shape: - 1 + batch_capacity: {} diffusion_schedule: implementation: kind: onnx @@ -343,6 +353,7 @@ pipeline: rank: 1 shape: - 4 + batch_capacity: {} diffusion_timesteps: implementation: kind: onnx @@ -355,6 +366,7 @@ pipeline: rank: 1 shape: - 3 + batch_capacity: {} schedule_lookup: implementation: kind: onnx @@ -383,6 +395,9 @@ pipeline: batch_layout: kind: request_aligned axis: 0 + batch_capacity: + uniform_dimensions: + - schedule_length tensor_scale: implementation: kind: onnx @@ -417,6 +432,11 @@ pipeline: batch_layout: kind: request_aligned axis: 0 + batch_capacity: + uniform_dimensions: + - channels + - height + - width decoder_input_scale: implementation: kind: onnx @@ -429,6 +449,7 @@ pipeline: rank: 1 shape: - 1 + batch_capacity: {} history_initializer: implementation: kind: onnx @@ -458,6 +479,11 @@ pipeline: batch_layout: kind: request_aligned axis: 0 + batch_capacity: + uniform_dimensions: + - channels + - height + - width guidance_combine: implementation: kind: onnx @@ -514,6 +540,11 @@ pipeline: conditional: conditional scale: scale estimate: estimate + batch_capacity: + uniform_dimensions: + - channels + - height + - width image_output_clamp: implementation: kind: onnx @@ -543,6 +574,11 @@ pipeline: batch_layout: kind: request_aligned axis: 0 + batch_capacity: + uniform_dimensions: + - axis_1 + - axis_2 + - axis_3 latent_row_shape: implementation: kind: onnx @@ -555,6 +591,7 @@ pipeline: rank: 1 shape: - 3 + batch_capacity: {} latent_noise: implementation: kind: onnx @@ -611,6 +648,7 @@ pipeline: row_shape: row_shape noise: noise next_offset: next_offset + batch_capacity: {} state: latent_state: contract: @@ -757,10 +795,12 @@ pipeline: value: denoiser.estimate output: noise_estimate mode: append + axis: 3 - kind: emit value: latent.body output: latent_trajectory mode: append + axis: 3 - kind: invoke component: continue_predicate inputs: diff --git a/tests/fixtures/onnx_genai_workflows/hierarchical_audio/inference_metadata.yaml b/tests/fixtures/onnx_genai_workflows/hierarchical_audio/inference_metadata.yaml index 742c0f988..bdd8401d4 100644 --- a/tests/fixtures/onnx_genai_workflows/hierarchical_audio/inference_metadata.yaml +++ b/tests/fixtures/onnx_genai_workflows/hierarchical_audio/inference_metadata.yaml @@ -384,8 +384,9 @@ pipeline: - batch - prompt_sequence batch_layout: - kind: request_aligned + kind: request_expanded axis: 0 + factor: 2 outputs: attention_mask: dtype: int64 @@ -394,8 +395,9 @@ pipeline: - batch - past_sequence_length + sequence_length batch_layout: - kind: request_aligned + kind: request_expanded axis: 0 + factor: 2 position_ids: dtype: int64 rank: 2 @@ -403,8 +405,9 @@ pipeline: - batch - prompt_sequence batch_layout: - kind: request_aligned + kind: request_expanded axis: 0 + factor: 2 body_attention_mask: dtype: int64 rank: 2 @@ -412,8 +415,9 @@ pipeline: - batch - prompt_sequence + 1 batch_layout: - kind: request_aligned + kind: request_expanded axis: 0 + factor: 2 body_position_ids: dtype: int64 rank: 2 @@ -421,8 +425,9 @@ pipeline: - batch - 1 batch_layout: - kind: request_aligned + kind: request_expanded axis: 0 + factor: 2 token_slot: dtype: int64 rank: 2 @@ -430,8 +435,9 @@ pipeline: - batch - 1 batch_layout: - kind: request_aligned + kind: request_expanded axis: 0 + factor: 2 past_key_values.0.key: dtype: float32 rank: 4 @@ -441,8 +447,9 @@ pipeline: - past_sequence_length - 4 batch_layout: - kind: request_aligned + kind: request_expanded axis: 0 + factor: 2 past_key_values.0.value: dtype: float32 rank: 4 @@ -452,8 +459,9 @@ pipeline: - past_sequence_length - 4 batch_layout: - kind: request_aligned + kind: request_expanded axis: 0 + factor: 2 global_step_update: implementation: kind: onnx @@ -467,8 +475,9 @@ pipeline: - batch - context batch_layout: - kind: request_aligned + kind: request_expanded axis: 0 + factor: 2 position_ids: dtype: int64 rank: 2 @@ -476,8 +485,9 @@ pipeline: - batch - 1 batch_layout: - kind: request_aligned + kind: request_expanded axis: 0 + factor: 2 outputs: next_attention_mask: dtype: int64 @@ -486,8 +496,9 @@ pipeline: - batch - context + 1 batch_layout: - kind: request_aligned + kind: request_expanded axis: 0 + factor: 2 next_position_ids: dtype: int64 rank: 2 @@ -495,8 +506,9 @@ pipeline: - batch - 1 batch_layout: - kind: request_aligned + kind: request_expanded axis: 0 + factor: 2 ar_initializer: implementation: kind: onnx diff --git a/tests/fixtures/onnx_genai_workflows/shared_state_pixel_flow/inference_metadata.yaml b/tests/fixtures/onnx_genai_workflows/shared_state_pixel_flow/inference_metadata.yaml index edf3e62ad..ce0f0f3be 100644 --- a/tests/fixtures/onnx_genai_workflows/shared_state_pixel_flow/inference_metadata.yaml +++ b/tests/fixtures/onnx_genai_workflows/shared_state_pixel_flow/inference_metadata.yaml @@ -458,6 +458,9 @@ pipeline: - 3 - batch - prompt_sequence + batch_layout: + kind: request_aligned + axis: 1 body_attention_mask: dtype: int64 rank: 2 @@ -474,6 +477,9 @@ pipeline: - 3 - batch - 1 + batch_layout: + kind: request_aligned + axis: 1 token_slot: dtype: int64 rank: 2 @@ -553,6 +559,9 @@ pipeline: - 3 - batch - image_tokens + batch_layout: + kind: request_aligned + axis: 1 token_grid: dtype: int64 rank: 1 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 3799f4da9..3d101cf37 100644 --- a/tests/fixtures/onnx_genai_workflows/static_cache/inference_metadata.yaml +++ b/tests/fixtures/onnx_genai_workflows/static_cache/inference_metadata.yaml @@ -1,4 +1,4 @@ -schema_version: '1.0' +schema_version: v1.1 pipeline: workflow: manifest: @@ -354,6 +354,7 @@ pipeline: input_ids: token_ids position_ids: position_ids logits: logits + batch_capacity: {} token_sampler: implementation: kind: onnx @@ -470,6 +471,9 @@ pipeline: batching: per_row inactive_rows: preserve application_overridable: true + batch_capacity: + uniform_dimensions: + - vocabulary termination: implementation: kind: onnx @@ -560,6 +564,9 @@ pipeline: parameters: batching: per_row inactive_rows: preserve + batch_capacity: + uniform_dimensions: + - num_eos token_state_update: implementation: kind: onnx @@ -622,6 +629,7 @@ pipeline: parameters: batching: per_row inactive_rows: preserve + batch_capacity: {} last_token_logits: implementation: kind: onnx @@ -648,6 +656,10 @@ pipeline: batch_layout: kind: request_aligned axis: 0 + batch_capacity: + uniform_dimensions: + - sequence + - vocabulary decoder_state_initializer: implementation: kind: onnx @@ -743,6 +755,9 @@ pipeline: batch_layout: kind: request_aligned axis: 0 + batch_capacity: + uniform_dimensions: + - prompt_sequence decoder_step_update: implementation: kind: onnx @@ -768,6 +783,7 @@ pipeline: batch_layout: kind: request_aligned axis: 0 + batch_capacity: {} cache_length_update: implementation: kind: onnx @@ -815,6 +831,7 @@ pipeline: batch_layout: kind: request_aligned axis: 0 + batch_capacity: {} termination_batch_initializer: implementation: kind: onnx @@ -885,6 +902,9 @@ pipeline: batch_layout: kind: request_aligned axis: 0 + batch_capacity: + uniform_dimensions: + - num_eos token_to_slot: implementation: kind: onnx @@ -909,6 +929,7 @@ pipeline: batch_layout: kind: request_aligned axis: 0 + batch_capacity: {} generated_length_update: implementation: kind: onnx @@ -956,6 +977,7 @@ pipeline: batch_layout: kind: request_aligned axis: 0 + batch_capacity: {} state: token: contract: diff --git a/tests/fixtures/onnx_genai_workflows/tts/inference_metadata.yaml b/tests/fixtures/onnx_genai_workflows/tts/inference_metadata.yaml index 113cbbe00..cf6992633 100644 --- a/tests/fixtures/onnx_genai_workflows/tts/inference_metadata.yaml +++ b/tests/fixtures/onnx_genai_workflows/tts/inference_metadata.yaml @@ -1,4 +1,4 @@ -schema_version: v1 +schema_version: v1.1 pipeline: workflow: manifest: @@ -255,6 +255,7 @@ pipeline: position_ids: position_ids logits: logits last_hidden_state: hidden_states + batch_capacity: {} code_predictor: implementation: kind: onnx @@ -265,10 +266,12 @@ pipeline: attention_mask: attention_mask position_ids: position_ids logits: logits + batch_capacity: {} embedding: implementation: kind: onnx artifact: embedding/model.onnx + batch_capacity: {} talker_step_embedder: implementation: kind: onnx @@ -276,10 +279,12 @@ pipeline: ports: roles: inputs_embeds: inputs_embeds + batch_capacity: {} talker_prefill_embedder: implementation: kind: onnx artifact: talker_prefill_embedder/model.onnx + batch_capacity: {} code_predictor_prefill: implementation: kind: onnx @@ -287,6 +292,7 @@ pipeline: ports: roles: inputs_embeds: inputs_embeds + batch_capacity: {} code_predictor_step_embedder: implementation: kind: onnx @@ -294,18 +300,22 @@ pipeline: ports: roles: inputs_embeds: inputs_embeds + batch_capacity: {} code_predictor_indices: implementation: kind: onnx artifact: code_predictor_indices/model.onnx + batch_capacity: {} talker_text_step: implementation: kind: onnx artifact: talker_text_step/model.onnx + batch_capacity: {} codec: implementation: kind: onnx artifact: codec/model.onnx + batch_capacity: {} last_token_logits: implementation: kind: onnx @@ -332,6 +342,10 @@ pipeline: batch_layout: kind: request_aligned axis: 0 + batch_capacity: + uniform_dimensions: + - sequence + - vocabulary setup_talker_sampler: implementation: kind: onnx @@ -365,6 +379,9 @@ pipeline: parameters: mode: greedy application_overridable: true + batch_capacity: + uniform_dimensions: + - vocabulary setup_predictor_sampler: implementation: kind: onnx @@ -398,6 +415,9 @@ pipeline: parameters: mode: greedy application_overridable: true + batch_capacity: + uniform_dimensions: + - vocabulary talker_sampler: implementation: kind: onnx @@ -431,6 +451,9 @@ pipeline: parameters: mode: greedy application_overridable: true + batch_capacity: + uniform_dimensions: + - vocabulary predictor_prefill_sampler: implementation: kind: onnx @@ -464,6 +487,9 @@ pipeline: parameters: mode: greedy application_overridable: true + batch_capacity: + uniform_dimensions: + - vocabulary predictor_body_sampler: implementation: kind: onnx @@ -497,6 +523,9 @@ pipeline: parameters: mode: greedy application_overridable: true + batch_capacity: + uniform_dimensions: + - vocabulary continue_predicate: implementation: kind: onnx @@ -517,6 +546,7 @@ pipeline: rank: 1 shape: - 1 + batch_capacity: {} tts_state_initializer: implementation: kind: onnx @@ -561,6 +591,9 @@ pipeline: batch_layout: kind: request_aligned axis: 0 + batch_capacity: + uniform_dimensions: + - sequence token_to_slot: implementation: kind: onnx @@ -585,6 +618,7 @@ pipeline: batch_layout: kind: request_aligned axis: 0 + batch_capacity: {} code_frame_update: implementation: kind: onnx @@ -622,6 +656,7 @@ pipeline: batch_layout: kind: request_aligned axis: 0 + batch_capacity: {} code_history_append: implementation: kind: onnx @@ -658,6 +693,9 @@ pipeline: batch_layout: kind: request_aligned axis: 0 + batch_capacity: + uniform_dimensions: + - frames cache_length_update: implementation: kind: onnx @@ -689,6 +727,7 @@ pipeline: batch_layout: kind: request_aligned axis: 0 + batch_capacity: {} talker_state_initializer: implementation: kind: onnx @@ -722,6 +761,9 @@ pipeline: - 3 - batch - sequence_len + batch_layout: + kind: request_aligned + axis: 1 body_attention_mask: dtype: int64 rank: 2 @@ -738,6 +780,9 @@ pipeline: - 3 - batch - 1 + batch_layout: + kind: request_aligned + axis: 1 past_key_values.0.key: dtype: float32 rank: 4 @@ -760,6 +805,9 @@ pipeline: batch_layout: kind: request_aligned axis: 0 + batch_capacity: + uniform_dimensions: + - prefill_sequence predictor_state_initializer: implementation: kind: onnx @@ -923,6 +971,9 @@ pipeline: batch_layout: kind: request_aligned axis: 0 + batch_capacity: + uniform_dimensions: + - prefill_sequence talker_step_update: implementation: kind: onnx @@ -945,6 +996,9 @@ pipeline: - 3 - batch - 1 + batch_layout: + kind: request_aligned + axis: 1 outputs: next_attention_mask: dtype: int64 @@ -962,6 +1016,12 @@ pipeline: - 3 - batch - 1 + batch_layout: + kind: request_aligned + axis: 1 + batch_capacity: + uniform_dimensions: + - context predictor_step_update: implementation: kind: onnx @@ -1005,6 +1065,9 @@ pipeline: batch_layout: kind: request_aligned axis: 0 + batch_capacity: + uniform_dimensions: + - context codec_layout: implementation: kind: onnx @@ -1032,6 +1095,9 @@ pipeline: batch_layout: kind: request_aligned axis: 0 + batch_capacity: + uniform_dimensions: + - frames state: last_frame: contract: @@ -1090,6 +1156,9 @@ pipeline: - 3 - batch - 1 + batch_layout: + kind: request_aligned + axis: 1 scope: invocation initializer: talker.initializer.body_position_ids recurrence: diff --git a/tests/fixtures/onnx_genai_workflows/video/inference_metadata.yaml b/tests/fixtures/onnx_genai_workflows/video/inference_metadata.yaml index 8d4c0fc2d..19988ce0b 100644 --- a/tests/fixtures/onnx_genai_workflows/video/inference_metadata.yaml +++ b/tests/fixtures/onnx_genai_workflows/video/inference_metadata.yaml @@ -1,4 +1,4 @@ -schema_version: v1 +schema_version: v1.1 pipeline: workflow: manifest: @@ -159,10 +159,12 @@ pipeline: ports: roles: encoder_hidden_states: encoder_hidden_states + batch_capacity: {} vae_decoder: implementation: kind: onnx artifact: vae_decoder/model.onnx + batch_capacity: {} model_input: implementation: kind: onnx @@ -207,6 +209,13 @@ pipeline: batch_layout: kind: request_aligned axis: 0 + batch_capacity: + uniform_dimensions: + - frames + - channels + - height + - width + - schedule_length solver_step: implementation: kind: onnx @@ -272,6 +281,13 @@ pipeline: step: step schedule: schedule next_state: next_state + batch_capacity: + uniform_dimensions: + - frames + - channels + - height + - width + - schedule_length continue_predicate: implementation: kind: onnx @@ -292,6 +308,7 @@ pipeline: rank: 1 shape: - 1 + batch_capacity: {} video_latent_init: implementation: kind: onnx @@ -332,6 +349,12 @@ pipeline: batch_layout: kind: request_aligned axis: 0 + batch_capacity: + uniform_dimensions: + - frames + - channels + - height + - width schedule_history_append: implementation: kind: onnx @@ -365,6 +388,9 @@ pipeline: batch_layout: kind: request_aligned axis: 0 + batch_capacity: + uniform_dimensions: + - history video_latent_permute: implementation: kind: onnx @@ -396,6 +422,12 @@ pipeline: batch_layout: kind: request_aligned axis: 0 + batch_capacity: + uniform_dimensions: + - frames + - channels + - height + - width video_latent_unscale: implementation: kind: onnx @@ -427,6 +459,12 @@ pipeline: batch_layout: kind: request_aligned axis: 0 + batch_capacity: + uniform_dimensions: + - channels + - frames + - height + - width video_decode_chunks: implementation: kind: onnx @@ -451,6 +489,12 @@ pipeline: rank: 1 shape: - 1 + batch_capacity: + uniform_dimensions: + - channels + - latent_frames + - height + - width video_decode_chunk: implementation: kind: onnx @@ -490,6 +534,12 @@ pipeline: batch_layout: kind: request_aligned axis: 0 + batch_capacity: + uniform_dimensions: + - channels + - latent_frames + - height + - width video_conv_cache_init: implementation: kind: onnx @@ -533,6 +583,12 @@ pipeline: batch_layout: kind: request_aligned axis: 0 + batch_capacity: + uniform_dimensions: + - channels + - latent_frames + - height + - width diffusion_schedule: implementation: kind: onnx @@ -545,6 +601,7 @@ pipeline: rank: 1 shape: - 4 + batch_capacity: {} diffusion_timesteps: implementation: kind: onnx @@ -557,6 +614,7 @@ pipeline: rank: 1 shape: - 3 + batch_capacity: {} schedule_lookup: implementation: kind: onnx @@ -585,6 +643,9 @@ pipeline: batch_layout: kind: request_aligned axis: 0 + batch_capacity: + uniform_dimensions: + - schedule_length state: latent: contract: diff --git a/tests/generate_onnx_genai_validation_packages.py b/tests/generate_onnx_genai_validation_packages.py index 30f2337be..44e874cd6 100644 --- a/tests/generate_onnx_genai_validation_packages.py +++ b/tests/generate_onnx_genai_validation_packages.py @@ -1326,7 +1326,7 @@ def _adapter_package(source_root: Path) -> ModelPackage: def _write_adapter_metadata(package: ModelPackage, directory: Path) -> None: metadata = { - "schema_version": "v1", + "schema_version": "v1.1", "pipeline": { "workflow": { "manifest": { @@ -1392,6 +1392,7 @@ def _write_adapter_metadata(package: ModelPackage, directory: Path) -> None: "abi": "onnx-genai.parameter-overlay", "version": "1", }, + "batch_capacity": {}, "ports": { "inputs": { "input": {