From 519b7acc91ced6c8b3e304cdf4efd351b4570946 Mon Sep 17 00:00:00 2001 From: spalne Date: Wed, 29 Jul 2026 13:12:19 -0700 Subject: [PATCH 1/5] Add support for AMD NPU --- src/winml/modelkit/models/hf/qwen3/genai.py | 131 ++++++++++++++++--- tests/unit/models/qwen3/test_genai_config.py | 35 +++++ 2 files changed, 147 insertions(+), 19 deletions(-) diff --git a/src/winml/modelkit/models/hf/qwen3/genai.py b/src/winml/modelkit/models/hf/qwen3/genai.py index 735bac21a..ce5ba2dc7 100644 --- a/src/winml/modelkit/models/hf/qwen3/genai.py +++ b/src/winml/modelkit/models/hf/qwen3/genai.py @@ -10,10 +10,16 @@ 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. Three 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"``). +* **OpenVINO** — Intel NPU (``ep="openvino"``). + +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 +27,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 +58,7 @@ # --------------------------------------------------------------------------- -# Qwen3-specific QNN execution-provider routing +# Qwen3-specific NPU execution-provider routing (QNN / VitisAI / OpenVINO) # --------------------------------------------------------------------------- @@ -86,17 +93,92 @@ 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 openvino_stage_session_options(log_id: str, device_type: str = "NPU") -> dict: + """Return the ``session_options`` block that routes a stage to the Intel NPU. + + Routes a Qwen3 transformer stage to the Intel NPU via the OpenVINO execution + provider. OpenVINO selects hardware through the ``device_type`` provider + option rather than a backend library path. + + Args: + log_id: ORT log identifier (shown in ORT logs), e.g. + ``"onnxruntime-genai.context"``. + device_type: OpenVINO device selector. ``"NPU"`` (default) targets the + Intel NPU; ``"GPU"``/``"CPU"`` target the respective Intel devices. + + Returns: + Dict suitable for the ``session_options`` key of a pipeline stage in + ``genai_config.json``. + """ + return { + "log_id": log_id, + "provider_options": [{"openvino": {"device_type": device_type}}], + "intra_op_num_threads": 2, + "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. + * ``ep="openvino"`` -> Intel 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"), + ) + if canonical == "OpenVINOExecutionProvider": + return ( + openvino_stage_session_options("onnxruntime-genai.context"), + openvino_stage_session_options("onnxruntime-genai.iterator"), + ) return None, None @@ -141,11 +223,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 +238,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), ``"vitisai"`` (AMD) or ``"openvino"`` (Intel) + 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 +283,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), ``"vitisai"`` (AMD Ryzen AI) or + ``"openvino"`` (Intel); ``"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. @@ -248,8 +335,10 @@ def write_genai_bundle( "build_decoder_pipeline_stages", "build_genai_config", "build_qwen3_transformer_only_stages", + "openvino_stage_session_options", "qnn_stage_session_options", "strip_gqa_default_attrs", + "vitisai_stage_session_options", "write_genai_bundle", ] @@ -288,7 +377,11 @@ 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 + GenaiTarget(ep="openvino", device="npu"), # Intel NPU + ), transformer_onnx_passes=(strip_gqa_default_attrs,), max_cache_len=2048, prefill_seq_len=64, diff --git a/tests/unit/models/qwen3/test_genai_config.py b/tests/unit/models/qwen3/test_genai_config.py index 66553d589..5858ec2c6 100644 --- a/tests/unit/models/qwen3/test_genai_config.py +++ b/tests/unit/models/qwen3/test_genai_config.py @@ -542,6 +542,41 @@ 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" + + def test_openvino_ep_injects_session_options(self) -> None: + """ep='openvino': context/iterator get Intel 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="openvino" + ) + 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 + assert ctx_opts["provider_options"][0]["openvino"]["device_type"] == "NPU" + assert itr_opts["log_id"] == "onnxruntime-genai.iterator" + # --------------------------------------------------------------------------- # Tests: write_genai_bundle wrapper (ep routing + transformer_onnx_passes) From 85797b433eda5828325cfd2375f740221a7d5b13 Mon Sep 17 00:00:00 2001 From: spalne Date: Tue, 11 Aug 2026 09:17:02 -0700 Subject: [PATCH 2/5] Enables weight sharing for the Qwen decoder's context+iterator --- src/winml/modelkit/session/genai_session.py | 388 +++++++++++++++++--- 1 file changed, 334 insertions(+), 54 deletions(-) diff --git a/src/winml/modelkit/session/genai_session.py b/src/winml/modelkit/session/genai_session.py index f584b54e4..aea30b90d 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, @@ -1214,8 +1255,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 +1302,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 +1350,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``. From 02b6f21102496ef327fbf73280bd4cf1f4a46b5f Mon Sep 17 00:00:00 2001 From: Supreet Singh Palne Date: Thu, 13 Aug 2026 09:56:02 -0700 Subject: [PATCH 3/5] Fix max length for static decoder pipelines --- src/winml/modelkit/session/genai_session.py | 18 +++++++++++++----- 1 file changed, 13 insertions(+), 5 deletions(-) diff --git a/src/winml/modelkit/session/genai_session.py b/src/winml/modelkit/session/genai_session.py index aea30b90d..e4c5254f1 100644 --- a/src/winml/modelkit/session/genai_session.py +++ b/src/winml/modelkit/session/genai_session.py @@ -460,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 @@ -524,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: @@ -899,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, From 9414a5845f38d5cfb02748d024aafd95fea1ecb8 Mon Sep 17 00:00:00 2001 From: spalne Date: Thu, 13 Aug 2026 10:00:57 -0700 Subject: [PATCH 4/5] Fix VitisAI provider options and bump opset for AMD NPU compiler support --- src/winml/modelkit/models/hf/qwen3/genai.py | 4 ++-- src/winml/modelkit/models/hf/qwen3/qwen_transformer_only.py | 2 +- 2 files changed, 3 insertions(+), 3 deletions(-) diff --git a/src/winml/modelkit/models/hf/qwen3/genai.py b/src/winml/modelkit/models/hf/qwen3/genai.py index ce5ba2dc7..20c01a6db 100644 --- a/src/winml/modelkit/models/hf/qwen3/genai.py +++ b/src/winml/modelkit/models/hf/qwen3/genai.py @@ -115,8 +115,8 @@ def vitisai_stage_session_options(log_id: str) -> dict: { "vitisai": { "target": "waic_target_vaiml_cpp_me", - "xmc_runner_config": 1, - "no_linear_slice": 1, + "xmc_runner_config": "1", + "no_linear_slice": "1", } } ], 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(), From 4300fe73d7184f9c4101c9443ad382d7b0ba6b0b Mon Sep 17 00:00:00 2001 From: spalne Date: Thu, 13 Aug 2026 10:13:21 -0700 Subject: [PATCH 5/5] Remove openvino options --- src/winml/modelkit/models/hf/qwen3/genai.py | 48 +++----------------- tests/unit/models/qwen3/test_genai_config.py | 16 ------- 2 files changed, 7 insertions(+), 57 deletions(-) diff --git a/src/winml/modelkit/models/hf/qwen3/genai.py b/src/winml/modelkit/models/hf/qwen3/genai.py index 20c01a6db..1911d5402 100644 --- a/src/winml/modelkit/models/hf/qwen3/genai.py +++ b/src/winml/modelkit/models/hf/qwen3/genai.py @@ -11,12 +11,11 @@ This module adds the **Qwen3-specific** layer on top: the Qwen3 transformer stages run on an NPU backend, so this is where the per-EP ``session_options`` -are constructed. Three NPU execution providers are supported for the +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"``). -* **OpenVINO** — Intel NPU (``ep="openvino"``). Keeping the EP-specific logic here lets the generic utilities stay universal while the Qwen3 bundle emits the correct per-EP ``genai_config.json``. @@ -58,7 +57,7 @@ # --------------------------------------------------------------------------- -# Qwen3-specific NPU execution-provider routing (QNN / VitisAI / OpenVINO) +# Qwen3-specific NPU execution-provider routing (QNN / VitisAI) # --------------------------------------------------------------------------- @@ -125,31 +124,6 @@ def vitisai_stage_session_options(log_id: str) -> dict: } -def openvino_stage_session_options(log_id: str, device_type: str = "NPU") -> dict: - """Return the ``session_options`` block that routes a stage to the Intel NPU. - - Routes a Qwen3 transformer stage to the Intel NPU via the OpenVINO execution - provider. OpenVINO selects hardware through the ``device_type`` provider - option rather than a backend library path. - - Args: - log_id: ORT log identifier (shown in ORT logs), e.g. - ``"onnxruntime-genai.context"``. - device_type: OpenVINO device selector. ``"NPU"`` (default) targets the - Intel NPU; ``"GPU"``/``"CPU"`` target the respective Intel devices. - - Returns: - Dict suitable for the ``session_options`` key of a pipeline stage in - ``genai_config.json``. - """ - return { - "log_id": log_id, - "provider_options": [{"openvino": {"device_type": device_type}}], - "intra_op_num_threads": 2, - "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. @@ -157,7 +131,6 @@ def _stage_session_options(ep: str, soc_model: str) -> tuple[dict | None, dict | * ``ep="qnn"`` -> Qualcomm QNN HTP (``soc_model`` selects the Snapdragon SoC). * ``ep="vitisai"`` -> AMD Ryzen AI NPU. - * ``ep="openvino"`` -> Intel NPU. Any other value (e.g. ``"cpu"``) leaves the stages on the default CPU provider. Short aliases and full ``*ExecutionProvider`` names are both @@ -174,11 +147,6 @@ def _stage_session_options(ep: str, soc_model: str) -> tuple[dict | None, dict | vitisai_stage_session_options("onnxruntime-genai.context"), vitisai_stage_session_options("onnxruntime-genai.iterator"), ) - if canonical == "OpenVINOExecutionProvider": - return ( - openvino_stage_session_options("onnxruntime-genai.context"), - openvino_stage_session_options("onnxruntime-genai.iterator"), - ) return None, None @@ -239,9 +207,9 @@ def build_qwen3_transformer_only_stages( embeddings_filename: Bundle filename for the embeddings model. lm_head_filename: Bundle filename for the lm_head model. ep: NPU execution provider for the ``context``/``iterator`` stages — - ``"qnn"`` (Qualcomm), ``"vitisai"`` (AMD) or ``"openvino"`` (Intel) - injects that EP's ``session_options`` so those stages run on the NPU - while ``embeddings`` and ``lm_head`` stay on CPU. ``"cpu"`` (default) + ``"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. Ignored @@ -292,8 +260,8 @@ def write_genai_bundle( Args: ep: NPU execution provider routing the transformer (context/iterator) - stages — ``"qnn"`` (Qualcomm HTP), ``"vitisai"`` (AMD Ryzen AI) or - ``"openvino"`` (Intel); ``"cpu"`` (default) keeps every stage on CPU. + 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. @@ -335,7 +303,6 @@ def write_genai_bundle( "build_decoder_pipeline_stages", "build_genai_config", "build_qwen3_transformer_only_stages", - "openvino_stage_session_options", "qnn_stage_session_options", "strip_gqa_default_attrs", "vitisai_stage_session_options", @@ -380,7 +347,6 @@ def write_genai_bundle( supported_targets=( GenaiTarget(ep="qnn", device="npu"), # Qualcomm Snapdragon NPU GenaiTarget(ep="vitisai", device="npu"), # AMD Ryzen AI NPU - GenaiTarget(ep="openvino", device="npu"), # Intel NPU ), transformer_onnx_passes=(strip_gqa_default_attrs,), max_cache_len=2048, diff --git a/tests/unit/models/qwen3/test_genai_config.py b/tests/unit/models/qwen3/test_genai_config.py index 5858ec2c6..759ad0f33 100644 --- a/tests/unit/models/qwen3/test_genai_config.py +++ b/tests/unit/models/qwen3/test_genai_config.py @@ -561,22 +561,6 @@ def test_vitisai_ep_injects_session_options(self) -> None: assert vitisai_opts["no_linear_slice"] == 1 assert itr_opts["log_id"] == "onnxruntime-genai.iterator" - def test_openvino_ep_injects_session_options(self) -> None: - """ep='openvino': context/iterator get Intel 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="openvino" - ) - 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 - assert ctx_opts["provider_options"][0]["openvino"]["device_type"] == "NPU" - assert itr_opts["log_id"] == "onnxruntime-genai.iterator" - # --------------------------------------------------------------------------- # Tests: write_genai_bundle wrapper (ep routing + transformer_onnx_passes)