diff --git a/src/winml/modelkit/models/hf/qwen3/genai.py b/src/winml/modelkit/models/hf/qwen3/genai.py index 735bac21a..1911d5402 100644 --- a/src/winml/modelkit/models/hf/qwen3/genai.py +++ b/src/winml/modelkit/models/hf/qwen3/genai.py @@ -10,10 +10,15 @@ reused by other model families. This module adds the **Qwen3-specific** layer on top: the Qwen3 transformer -stages target the QNN HTP (NPU) backend, so this is where the QNN -``session_options`` are constructed. Keeping the EP-specific logic here lets the -generic utilities stay universal while the Qwen3 bundle keeps emitting the exact -same ``genai_config.json`` as before. +stages run on an NPU backend, so this is where the per-EP ``session_options`` +are constructed. Two NPU execution providers are supported for the +transformer (context/iterator) stages: + +* **QNN HTP** — Qualcomm Snapdragon NPU (``ep="qnn"``). +* **VitisAI** — AMD Ryzen AI NPU (``ep="vitisai"``). + +Keeping the EP-specific logic here lets the generic utilities stay universal +while the Qwen3 bundle emits the correct per-EP ``genai_config.json``. """ from __future__ import annotations @@ -21,6 +26,7 @@ from typing import TYPE_CHECKING from ....onnx import strip_node_attrs +from ....utils.constants import normalize_ep_name from ....utils.genai import ( DEFAULT_CONTEXT_FILENAME, DEFAULT_EMBEDDINGS_FILENAME, @@ -51,7 +57,7 @@ # --------------------------------------------------------------------------- -# Qwen3-specific QNN execution-provider routing +# Qwen3-specific NPU execution-provider routing (QNN / VitisAI) # --------------------------------------------------------------------------- @@ -86,17 +92,61 @@ def qnn_stage_session_options(log_id: str, soc_model: str = "60") -> dict: } +def vitisai_stage_session_options(log_id: str) -> dict: + """Return the ``session_options`` block that routes a stage to the AMD NPU. + + Routes a Qwen3 transformer stage to the AMD Ryzen AI NPU via the VitisAI + execution provider. The provider options match the AMD reference inference + configuration (``waic_target_vaiml_cpp_me`` VAIML C++ backend with the + XMC runner and linear-slice disabled). + + Args: + log_id: ORT log identifier (shown in ORT logs), e.g. + ``"onnxruntime-genai.context"``. + + Returns: + Dict suitable for the ``session_options`` key of a pipeline stage in + ``genai_config.json``. + """ + return { + "log_id": log_id, + "provider_options": [ + { + "vitisai": { + "target": "waic_target_vaiml_cpp_me", + "xmc_runner_config": "1", + "no_linear_slice": "1", + } + } + ], + "intra_op_num_threads": 8, + "inter_op_num_threads": 1, + } + + def _stage_session_options(ep: str, soc_model: str) -> tuple[dict | None, dict | None]: """Return ``(context, iterator)`` session_options for the given EP. - ``ep="qnn"`` routes the transformer stages to the QNN HTP (NPU) backend; any - other value (e.g. ``"cpu"``) leaves them on the default CPU provider. + Routes the Qwen3 transformer (context/iterator) stages to an NPU backend: + + * ``ep="qnn"`` -> Qualcomm QNN HTP (``soc_model`` selects the Snapdragon SoC). + * ``ep="vitisai"`` -> AMD Ryzen AI NPU. + + Any other value (e.g. ``"cpu"``) leaves the stages on the default CPU + provider. Short aliases and full ``*ExecutionProvider`` names are both + accepted (normalized via :func:`normalize_ep_name`). """ - if ep == "qnn": + canonical = normalize_ep_name(ep) + if canonical == "QNNExecutionProvider": return ( qnn_stage_session_options("onnxruntime-genai.context", soc_model=soc_model), qnn_stage_session_options("onnxruntime-genai.iterator", soc_model=soc_model), ) + if canonical == "VitisAIExecutionProvider": + return ( + vitisai_stage_session_options("onnxruntime-genai.context"), + vitisai_stage_session_options("onnxruntime-genai.iterator"), + ) return None, None @@ -141,11 +191,11 @@ def build_qwen3_transformer_only_stages( ep: str = "cpu", soc_model: str = "60", ) -> tuple[list[PipelineStage], DecoderIOMapping]: - """Build the Qwen3 4-stage pipeline, routing ctx/iter to QNN when ``ep="qnn"``. + """Build the Qwen3 4-stage pipeline, routing ctx/iter to the NPU per ``ep``. Qwen3-specific wrapper over :func:`winml.modelkit.utils.genai.build_decoder_pipeline_stages` that injects - the QNN ``session_options`` for the transformer stages. Tensor names are + the NPU ``session_options`` for the transformer stages. Tensor names are still discovered by introspecting the ONNX graphs, so nothing is hardcoded. Args: @@ -156,11 +206,14 @@ def build_qwen3_transformer_only_stages( iterator_filename: Bundle filename for the iterator model. embeddings_filename: Bundle filename for the embeddings model. lm_head_filename: Bundle filename for the lm_head model. - ep: ``"qnn"`` injects QNN HTP ``session_options`` into the ``context`` - and ``iterator`` stages so they run on the NPU while ``embeddings`` - and ``lm_head`` stay on CPU. ``"cpu"`` (default) omits them. + ep: NPU execution provider for the ``context``/``iterator`` stages — + ``"qnn"`` (Qualcomm) or ``"vitisai"`` (AMD) injects that EP's + ``session_options`` so those stages run on the NPU while + ``embeddings`` and ``lm_head`` stay on CPU. ``"cpu"`` (default) + omits them. soc_model: Snapdragon SoC model number forwarded to the QNN backend when - ``ep="qnn"``. Default ``"60"`` targets Snapdragon 8 Gen 3. + ``ep="qnn"``. Default ``"60"`` targets Snapdragon 8 Gen 3. Ignored + for non-QNN EPs. Returns: ``(stages, decoder_io)`` — see @@ -198,18 +251,20 @@ def write_genai_bundle( soc_model: str = "60", transformer_onnx_passes: Sequence[Callable[[onnx.ModelProto], onnx.ModelProto]] | None = None, ) -> Path: - """Assemble a Qwen3 genai bundle, routing ctx/iter to QNN when ``ep="qnn"``. + """Assemble a Qwen3 genai bundle, routing ctx/iter to the NPU per ``ep``. Qwen3-specific wrapper over - :func:`winml.modelkit.utils.genai.write_genai_bundle` that supplies the QNN + :func:`winml.modelkit.utils.genai.write_genai_bundle` that supplies the NPU ``session_options`` for the transformer stages. See the generic function for the description of every other argument. Args: - ep: ``"qnn"`` routes the transformer (context/iterator) stages to the QNN - HTP (NPU) backend; ``"cpu"`` (default) keeps every stage on CPU. + ep: NPU execution provider routing the transformer (context/iterator) + stages — ``"qnn"`` (Qualcomm HTP) or ``"vitisai"`` (AMD Ryzen AI); + ``"cpu"`` (default) keeps every stage on CPU. soc_model: Snapdragon SoC model passed to the QNN backend when ``ep="qnn"``. Default ``"60"`` = Snapdragon 8 Gen 3 / X Elite. + Ignored for non-QNN EPs. transformer_onnx_passes: Optional ONNX graph transforms applied to the copied context/iterator models before ``genai_config.json`` is written. Forwarded verbatim to the generic assembler. @@ -250,6 +305,7 @@ def write_genai_bundle( "build_qwen3_transformer_only_stages", "qnn_stage_session_options", "strip_gqa_default_attrs", + "vitisai_stage_session_options", "write_genai_bundle", ] @@ -288,7 +344,10 @@ def write_genai_bundle( ), ), assemble=write_genai_bundle, - supported_targets=(GenaiTarget(ep="qnn", device="npu"),), + supported_targets=( + GenaiTarget(ep="qnn", device="npu"), # Qualcomm Snapdragon NPU + GenaiTarget(ep="vitisai", device="npu"), # AMD Ryzen AI NPU + ), transformer_onnx_passes=(strip_gqa_default_attrs,), max_cache_len=2048, prefill_seq_len=64, diff --git a/src/winml/modelkit/models/hf/qwen3/qwen_transformer_only.py b/src/winml/modelkit/models/hf/qwen3/qwen_transformer_only.py index 282017fe3..f6e524490 100644 --- a/src/winml/modelkit/models/hf/qwen3/qwen_transformer_only.py +++ b/src/winml/modelkit/models/hf/qwen3/qwen_transformer_only.py @@ -329,7 +329,7 @@ def outputs(self) -> dict[str, dict[int, str]]: QWEN_TRANSFORMER_ONLY_CONFIG = WinMLBuildConfig( - export=WinMLExportConfig(dynamo=False, opset_version=18), + export=WinMLExportConfig(dynamo=False, opset_version=21), # Pure graph (no post-export RMSNorm fusion / matmul-add fusion): the default # WinMLOptimizationConfig() leaves every fusion flag off. optim=WinMLOptimizationConfig(), diff --git a/src/winml/modelkit/session/genai_session.py b/src/winml/modelkit/session/genai_session.py index f584b54e4..e4c5254f1 100644 --- a/src/winml/modelkit/session/genai_session.py +++ b/src/winml/modelkit/session/genai_session.py @@ -136,6 +136,47 @@ def _compile_stage_worker(src: str, dst: str, ep_alias: str, provider_options: d raise RuntimeError(f"Compilation failed: {result.errors}") +def _compile_shared_stages_worker( + srcs: list[str], dsts: list[str], ep_alias: str, provider_options: dict +) -> None: + """Compile several stage ONNX graphs into weight-sharing EPContexts. + + Executed in a subprocess by :meth:`GenaiSession._compile_stages_shared`. A + single :class:`~winml.modelkit.compiler.Compiler` with + ``n_total_models=len(srcs)`` compiles every source in order through one + shared ``SessionOptions``, so ORT sets ``ep.share_ep_contexts`` on all but + the last graph and ``ep.stop_share_ep_contexts`` on the last — the compiled + graphs then reference a single shared weights ``.bin`` instead of embedding + a private copy each. The EP is resolved generically from *ep_alias* and the + provider options (identical across the group) are forwarded unchanged. + + Args: + srcs: Absolute paths to the source ONNX files, in compile order. + dsts: Absolute EPContext output paths, aligned with *srcs*. + ep_alias: EP short name shared by every stage (e.g. ``"qnn"``). + provider_options: EP provider options from ``genai_config.json`` (shared + by the whole group). + + Raises: + RuntimeError: If *ep_alias* is not EPContext-capable, or any stage + compilation reports failure. + """ + from ..compiler import Compiler, WinMLCompileConfig + + config = WinMLCompileConfig.for_provider(ep_alias) # type: ignore[arg-type] + if config is None: + raise RuntimeError(f"EP {ep_alias!r} does not support EPContext pre-compilation") + config.ep_config.provider_options.update(provider_options) + + # One Compiler instance threads a shared SessionOptions across all models so + # the EP context (and its weights) is shared; the last model flushes it. + compiler = Compiler(n_total_models=len(srcs)) + for src, dst in zip(srcs, dsts): + result = compiler.compile(model_path=src, output_path=dst, config=config) + if not result.success: + raise RuntimeError(f"Compilation failed for {src}: {result.errors}") + + def _prepare_derived_bundle_worker( result_queue: Any, bundle_dir: str, @@ -419,6 +460,7 @@ def __init__( self._compile_timeout = compile_timeout # Resolved at load() time. self._context_length: int | None = None + self._is_decoder_pipeline = False # og.* handles — None until load() is called. self._model: Any = None @@ -483,6 +525,7 @@ def load(self) -> None: og = self._import_og() cfg = self._read_genai_config() + self._is_decoder_pipeline = cfg.get("model", {}).get("type") == "decoder-pipeline" # Apply the ``ep`` override (if any) to obtain the *effective* config that # actually drives routing. Precedence is explicit arg > bundle config: @@ -858,17 +901,23 @@ def _encode_prompt(self, prompt: str | list[int]) -> list[int]: def _new_generator(self, cfg: GenerationConfig, prompt_len: int) -> Any: """Build an ``og.Generator`` with search options from *cfg*. - ``max_length`` is set to ``prompt_len + cfg.max_new_tokens``, capped at - the bundle's ``context_length``. This avoids pre-allocating KV cache - for the full context window (which can be 128K+ for DML bundles) when - only a small generation is requested. + Static decoder pipelines use the bundle's full ``context_length``; + their fixed-size KV-cache stages may expand the internal sequence to + the configured search limit during prompt processing. Other model + types use ``prompt_len + cfg.max_new_tokens``, capped at the context + length, to avoid unnecessarily large dynamic allocations. The + generation loops enforce ``max_new_tokens`` as the output-token limit. The prompt is **not** appended — callers decide whether to time ``append_tokens`` separately (see :meth:`generate_timed`). """ og = self._import_og() assert self._context_length is not None, "_new_generator called before load()" - max_length = min(prompt_len + cfg.max_new_tokens, self._context_length) + max_length = ( + self._context_length + if self._is_decoder_pipeline + else min(prompt_len + cfg.max_new_tokens, self._context_length) + ) params = og.GeneratorParams(self._model) params.set_search_options( max_length=max_length, @@ -1214,8 +1263,6 @@ def _prepare_derived_bundle( Path to the derived bundle directory (equals ``bundle_dir`` when the config was not overridden and no stage needed compiling). """ - from ..onnx import is_compiled_onnx - compiled_dir = self._bundle_dir / self._COMPILED_SUBDIR cfg = effective_cfg if effective_cfg is not None else self._read_genai_config() @@ -1263,61 +1310,23 @@ def _prepare_derived_bundle( # the file must physically live there. compiled_stage_filenames: set[str] = set() - for stage_key, onnx_filename, ep_alias, ep_opts in compilable_stages: - src_onnx = self._bundle_dir / onnx_filename - # The EP is part of the cache key: an ``ep`` override can route the - # same stage onto different EPContext providers across runs, so each - # EP gets its own artifact and EP-A's binary is never reused for EP-B. - ctx_onnx = compiled_dir / f"{stage_key}_{ep_alias}_ctx.onnx" - - # Skip recompilation only when the cache is genuinely up-to-date. - if self._epcontext_is_fresh(src_onnx, ctx_onnx, ep_alias, ep_opts): - logger.info("Stage %r: reusing cached EPContext %s", stage_key, ctx_onnx.name) - # Use just the filename — genai_config.json lives in compiled_dir, - # so ort-genai resolves filenames relative to compiled_dir. - self._patch_stage_filename(modified_cfg, stage_key, ctx_onnx.name) - compiled_stage_filenames.add(onnx_filename) - any_compiled = True - continue - - # A stage whose source ONNX is already an EPContext (pre-compiled) - # model needs no recompilation; reference the original file by its - # bundle-relative filename and let :meth:`_mirror_non_onnx_files` - # link it (plus any weights sidecar) into compiled_dir. A parse - # failure is treated as "not compiled" so a malformed source falls - # through to the normal compile path, which has its own error - # handling. - already_compiled = False - if src_onnx.exists(): - try: - already_compiled = is_compiled_onnx(src_onnx) - except (ValueError, OSError): - already_compiled = False - if already_compiled: - logger.info( - "Stage %r: source is already an EPContext model; using as-is", - stage_key, + # Group stages so that multiple stages on the same EPContext EP that + # share weights (e.g. the Qwen decoder's ``context`` prefill and + # ``iterator`` decode graphs, both on QNN) are compiled together through + # one shared EP context — the compiled artifacts then reference a single + # shared weights ``.bin`` instead of duplicating the weights per stage. + # Stages that cannot share (single-stage EP groups) keep the original + # per-stage compile path unchanged. + for group in self._group_compilable_stages(compilable_stages): + if len(group) > 1: + compiled = self._process_shared_group( + group, compiled_dir, modified_cfg, compiled_stage_filenames ) - self._patch_stage_filename(modified_cfg, stage_key, onnx_filename) - continue - - # Attempt compilation. - success = self._compile_stage(src_onnx, ctx_onnx, stage_key, ep_alias, ep_opts) - if success: - self._write_compile_marker(ctx_onnx, ep_alias, ep_opts) - self._patch_stage_filename(modified_cfg, stage_key, ctx_onnx.name) - compiled_stage_filenames.add(onnx_filename) - any_compiled = True else: - logger.warning( - "Stage %r: compilation failed; using original ONNX (JIT fallback)", stage_key + compiled = self._process_single_stage( + group[0], compiled_dir, modified_cfg, compiled_stage_filenames ) - # Fall back to the original ONNX by its bundle-relative filename. - # ort-genai resolves stage filenames relative to compiled_dir, so - # the source ONNX (+ its weights sidecar) is mirrored in by - # :meth:`_mirror_non_onnx_files` below; patching an absolute path - # would be wrongly joined onto compiled_dir into a broken path. - self._patch_stage_filename(modified_cfg, stage_key, onnx_filename) + any_compiled = any_compiled or compiled # A derived bundle is only needed when routing changed: either a stage # was (re)compiled or the ``ep`` override rewrote the config. @@ -1349,6 +1358,285 @@ def _prepare_derived_bundle( ) return compiled_dir + def _group_compilable_stages( + self, compilable_stages: list[tuple[str, str, str, dict]] + ) -> list[list[tuple[str, str, str, dict]]]: + """Partition compilable stages into weight-sharing groups. + + Stages targeting the *same* EPContext-capable EP are candidates for + weight sharing: when more than one such stage exists and they are + :meth:`_stages_shareable` (identical provider options, existing and + not-yet-compiled sources), they are compiled together through one shared + EP context so their weights live in a single shared ``.bin`` (the Qwen + decoder's ``context`` + ``iterator`` graphs are the canonical case). + Every other stage stays in its own single-element group and keeps the + original per-stage compile path. + + Order is preserved: the first-seen EP's group comes first, and stages + within a shared group keep their pipeline order (so the last stage flushes + the shared context via ``ep.stop_share_ep_contexts``). + """ + by_ep: dict[str, list[tuple[str, str, str, dict]]] = {} + ep_order: list[str] = [] + for stage in compilable_stages: + ep_alias = stage[2] + if ep_alias not in by_ep: + by_ep[ep_alias] = [] + ep_order.append(ep_alias) + by_ep[ep_alias].append(stage) + + groups: list[list[tuple[str, str, str, dict]]] = [] + for ep_alias in ep_order: + stages = by_ep[ep_alias] + if len(stages) > 1 and self._stages_shareable(stages): + groups.append(stages) + else: + groups.extend([stage] for stage in stages) + return groups + + def _stages_shareable(self, stages: list[tuple[str, str, str, dict]]) -> bool: + """Return ``True`` when *stages* may be compiled into one shared context. + + Weight sharing requires a single shared ``SessionOptions``, so every + stage must use identical provider options. Each source graph must also + exist and still be an uncompiled ONNX — a source that is already an + EPContext model (or missing) cannot participate in a fresh shared + compile and forces the group back onto the per-stage path. + """ + from ..onnx import is_compiled_onnx + + first_opts = stages[0][3] + for _stage_key, onnx_filename, _ep_alias, ep_opts in stages: + if ep_opts != first_opts: + return False + src = self._bundle_dir / onnx_filename + if not src.exists(): + return False + try: + if is_compiled_onnx(src): + return False + except (ValueError, OSError): + return False + return True + + def _process_single_stage( + self, + stage: tuple[str, str, str, dict], + compiled_dir: Path, + modified_cfg: dict, + compiled_stage_filenames: set[str], + ) -> bool: + """Compile (or reuse) one pipeline stage; returns whether routing changed. + + Encapsulates the single-stage cache-freshness / already-compiled / + compile / JIT-fallback logic and patches *modified_cfg* to the resolved + stage filename. Returns ``True`` when a compiled EPContext is now in use + (fresh cache reuse or a successful compile), matching the ``any_compiled`` + contribution of the original inline loop. + """ + from ..onnx import is_compiled_onnx + + stage_key, onnx_filename, ep_alias, ep_opts = stage + src_onnx = self._bundle_dir / onnx_filename + # The EP is part of the cache key: an ``ep`` override can route the + # same stage onto different EPContext providers across runs, so each + # EP gets its own artifact and EP-A's binary is never reused for EP-B. + ctx_onnx = compiled_dir / f"{stage_key}_{ep_alias}_ctx.onnx" + + # Skip recompilation only when the cache is genuinely up-to-date. + if self._epcontext_is_fresh(src_onnx, ctx_onnx, ep_alias, ep_opts): + logger.info("Stage %r: reusing cached EPContext %s", stage_key, ctx_onnx.name) + # Use just the filename — genai_config.json lives in compiled_dir, + # so ort-genai resolves filenames relative to compiled_dir. + self._patch_stage_filename(modified_cfg, stage_key, ctx_onnx.name) + compiled_stage_filenames.add(onnx_filename) + return True + + # A stage whose source ONNX is already an EPContext (pre-compiled) + # model needs no recompilation; reference the original file by its + # bundle-relative filename and let :meth:`_mirror_non_onnx_files` + # link it (plus any weights sidecar) into compiled_dir. A parse + # failure is treated as "not compiled" so a malformed source falls + # through to the normal compile path, which has its own error + # handling. + already_compiled = False + if src_onnx.exists(): + try: + already_compiled = is_compiled_onnx(src_onnx) + except (ValueError, OSError): + already_compiled = False + if already_compiled: + logger.info( + "Stage %r: source is already an EPContext model; using as-is", + stage_key, + ) + self._patch_stage_filename(modified_cfg, stage_key, onnx_filename) + return False + + # Attempt compilation. + success = self._compile_stage(src_onnx, ctx_onnx, stage_key, ep_alias, ep_opts) + if success: + self._write_compile_marker(ctx_onnx, ep_alias, ep_opts) + self._patch_stage_filename(modified_cfg, stage_key, ctx_onnx.name) + compiled_stage_filenames.add(onnx_filename) + return True + + logger.warning( + "Stage %r: compilation failed; using original ONNX (JIT fallback)", stage_key + ) + # Fall back to the original ONNX by its bundle-relative filename. + # ort-genai resolves stage filenames relative to compiled_dir, so + # the source ONNX (+ its weights sidecar) is mirrored in by + # :meth:`_mirror_non_onnx_files` below; patching an absolute path + # would be wrongly joined onto compiled_dir into a broken path. + self._patch_stage_filename(modified_cfg, stage_key, onnx_filename) + return False + + def _process_shared_group( + self, + group: list[tuple[str, str, str, dict]], + compiled_dir: Path, + modified_cfg: dict, + compiled_stage_filenames: set[str], + ) -> bool: + """Compile a group of stages together with weight sharing. + + The stages are compiled through a single shared EP context so their + weights are stored once in a shared ``.bin`` that every compiled stage + graph references. Cache reuse is all-or-nothing: the group is only + reused when *every* member's EPContext is fresh, otherwise the whole + group is recompiled together (a partial recompile cannot re-establish + the shared context). On failure every stage falls back to its original + ONNX (JIT). Returns ``True`` when the shared EPContext is now in use. + """ + ep_alias = group[0][2] + ep_opts = group[0][3] + stage_keys = [stage[0] for stage in group] + srcs = [self._bundle_dir / onnx_filename for _sk, onnx_filename, _ea, _eo in group] + ctx_outs = [ + compiled_dir / f"{stage_key}_{ep_alias}_ctx.onnx" + for stage_key, _fn, _ea, _eo in group + ] + + # Reuse only when every stage's cached EPContext is fresh; weight sharing + # is all-or-nothing, so a single stale member recompiles the whole group. + if all( + self._epcontext_is_fresh(src, ctx, ep_alias, ep_opts) + for src, ctx in zip(srcs, ctx_outs) + ): + logger.info("Stages %s: reusing cached shared EPContext", stage_keys) + for (stage_key, onnx_filename, _ea, _eo), ctx in zip(group, ctx_outs): + self._patch_stage_filename(modified_cfg, stage_key, ctx.name) + compiled_stage_filenames.add(onnx_filename) + return True + + logger.info( + "Stages %s: compiling together with weight sharing on EP %r", stage_keys, ep_alias + ) + success = self._compile_stages_shared(srcs, ctx_outs, group) + if success: + for (stage_key, onnx_filename, ea, eo), ctx in zip(group, ctx_outs): + self._write_compile_marker(ctx, ea, eo) + self._patch_stage_filename(modified_cfg, stage_key, ctx.name) + compiled_stage_filenames.add(onnx_filename) + return True + + logger.warning( + "Stages %s: shared compilation failed; using original ONNX (JIT fallback)", + stage_keys, + ) + for stage_key, onnx_filename, _ea, _eo in group: + self._patch_stage_filename(modified_cfg, stage_key, onnx_filename) + return False + + def _compile_stages_shared( + self, + srcs: list[Path], + ctx_outs: list[Path], + group: list[tuple[str, str, str, dict]], + ) -> bool: + """Compile *srcs* together into weight-sharing EPContexts at *ctx_outs*. + + Mirrors :meth:`_compile_stage` but drives the multi-model shared-context + compiler (``Compiler(n_total_models=len(srcs))``) in one spawned + subprocess, so a hang or native teardown fault (issue #1087) is bounded + by a timeout and the same crash-during-teardown salvage applies. The + timeout is scaled by the number of stages since they compile + sequentially in the one worker. Returns ``True`` on success (or a + successful post-crash salvage of every stage). + """ + import multiprocessing + + ep_alias = group[0][2] + ep_opts = dict(group[0][3]) + stage_keys = [stage[0] for stage in group] + + logger.info( + "Compiling stages %s for EP %r with weight sharing (options=%s)", + stage_keys, + ep_alias, + ep_opts, + ) + + # Snapshot mtimes of any pre-existing EPContext artifacts before spawning + # the compiler so a teardown-crash salvage accepts only files this run + # actually produced (see :meth:`_compile_stage`). + pre_compile_mtimes: dict[str, int] = {} + for src, ctx in zip(srcs, ctx_outs): + for p in self._epcontext_candidate_paths(src, ctx): + if p.exists(): + pre_compile_mtimes[str(p)] = p.stat().st_mtime_ns + + ctx = multiprocessing.get_context("spawn") + proc = ctx.Process( + target=_compile_shared_stages_worker, + args=( + [str(s) for s in srcs], + [str(c) for c in ctx_outs], + ep_alias, + ep_opts, + ), + ) + proc.start() + proc.join(timeout=self._compile_timeout * len(srcs)) + + if proc.is_alive(): + logger.error( + "Shared compilation of stages %s timed out after %ds — killing subprocess.", + stage_keys, + self._compile_timeout * len(srcs), + ) + proc.kill() + proc.join() + for ctx_out in ctx_outs: + self._discard_compiled_stage(ctx_out) + return False + + if proc.exitcode != 0: + # QNN can fault during teardown after every artifact (the shared + # ``.bin`` and each ``_ctx.onnx``) has already flushed to disk; + # salvage each stage before treating the group as failed. + if all( + self._salvage_epcontext(src, ctx_out, pre_compile_mtimes) + for src, ctx_out in zip(srcs, ctx_outs) + ): + logger.warning( + "Shared compile subprocess exited %d, but valid EPContexts were " + "already written (crash during teardown); salvaging stages %s", + proc.exitcode, + stage_keys, + ) + return True + logger.warning( + "Shared compilation of stages %s failed (exit %d)", stage_keys, proc.exitcode + ) + for ctx_out in ctx_outs: + self._discard_compiled_stage(ctx_out) + return False + + logger.info("Stages %s compiled successfully with weight sharing", stage_keys) + return True + @staticmethod def _resolve_stage_ep(provider_options: list) -> tuple[str | None, dict]: """Resolve a stage's EPContext-capable EP from its ``provider_options``. diff --git a/tests/unit/models/qwen3/test_genai_config.py b/tests/unit/models/qwen3/test_genai_config.py index 66553d589..759ad0f33 100644 --- a/tests/unit/models/qwen3/test_genai_config.py +++ b/tests/unit/models/qwen3/test_genai_config.py @@ -542,6 +542,25 @@ def test_custom_soc_model(self) -> None: ctx = next(s for s in stages if s.name == "context") assert ctx.session_options["provider_options"][0]["qnn"]["soc_model"] == "73" + def test_vitisai_ep_injects_session_options(self) -> None: + """ep='vitisai': context/iterator get AMD NPU session_options; emb/lm_head do not.""" + with self._patch_onnx(): + stages, _ = build_qwen3_transformer_only_stages( + "ctx.onnx", "iter.onnx", num_layers=4, ep="vitisai" + ) + stage_map = {s.name: s for s in stages} + assert stage_map["embeddings"].session_options is None + assert stage_map["lm_head"].session_options is None + ctx_opts = stage_map["context"].session_options + itr_opts = stage_map["iterator"].session_options + assert ctx_opts is not None + assert itr_opts is not None + vitisai_opts = ctx_opts["provider_options"][0]["vitisai"] + assert vitisai_opts["target"] == "waic_target_vaiml_cpp_me" + assert vitisai_opts["xmc_runner_config"] == 1 + assert vitisai_opts["no_linear_slice"] == 1 + assert itr_opts["log_id"] == "onnxruntime-genai.iterator" + # --------------------------------------------------------------------------- # Tests: write_genai_bundle wrapper (ep routing + transformer_onnx_passes)