From c8d67bc9f1155fa0fb8a38ee1d682b31a4db7c44 Mon Sep 17 00:00:00 2001 From: Tianlei Wu Date: Fri, 30 Oct 2026 02:35:37 +0000 Subject: [PATCH 1/6] Add Qwen3.8 QMoE offload package builder Add build_package_offload_v3.py, which converts a v12b package into a QMoE-offload package for GPUs that cannot hold the routed experts: - The 48 QMoE nodes of the target graph become QMoEOffload nodes with an explicit layer index and sizes. Resident and spare expert slots can be set directly or be derived from an expert-popularity profile. - model_s2.onnx is the rewritten graph; model_s3.onnx and model_s4.onnx extend it to 3 and 4 verify rows for two or three chained MTP drafts, and model.onnx is a copy of model_s4.onnx. - weights_pool.txt places the routed experts in CPU mmap placement and the dense weights and the token embedding in host-mapped placement. Everything else is hard-linked from the source package. - genai_config.json gets the CUDA-graph, arena and device-memory-limit options, and the stop token id can be replaced. - --out must not equal, contain or be nested in the source package. - --migrate-shared-onnx moves a single graph to the compact QMoEOffload schema without touching the weight data. The README documents the script and its options. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- tools/qwen38_graph/README.md | 37 + .../qwen38_graph/build_package_offload_v3.py | 907 ++++++++++++++++++ 2 files changed, 944 insertions(+) create mode 100644 tools/qwen38_graph/build_package_offload_v3.py diff --git a/tools/qwen38_graph/README.md b/tools/qwen38_graph/README.md index 8760258b6e..607deaf759 100644 --- a/tools/qwen38_graph/README.md +++ b/tools/qwen38_graph/README.md @@ -175,6 +175,42 @@ private copy that is extended, so the earlier package is never modified. The scr `--out` that is equal to, nested in or around `--src`, `--reuse-hc-bf16-from`, `--plane-lib` or `--provider-lib` (checked before anything is deleted). +### QMoE-offload packages (`build_package_offload_v3.py`) + +`build_package_offload_v3.py --from-package --out --plane-lib --provider-lib +` converts a v12b package (layered, or flat with `model.onnx` in the root) so that the routed experts +stay in host memory and are streamed to the GPU on demand (`stack` is `v12b-engine-qmoe-offload-v3` in `package_manifest.json`). +Everything that does not change is hard-linked from the source (`--copy` copies): + +* the 48 `com.microsoft::QMoE` nodes of the target graph become `com.microsoft::QMoEOffload` nodes (explicit `layer_index`, + `hidden_size` and `inter_size`, plus tile, working-buffer and CUDA-graph attributes). The graph is written as `model_s2.onnx` (in + `target/` for a layered package); `model_s3.onnx` and `model_s4.onnx` extend it to 3 and 4 verify rows (2 or 3 chained MTP drafts), + and `model.onnx` is a copy of `model_s4.onnx`; +* `weights_pool.txt` puts the routed-expert weights and scales (63.28 GiB) in `cpu` mmap placement. `--dense-placement` (default + `host_mapped`) places the `dense.data` weights of the dense prefill session, `--embed-placement` (default `host_mapped`) the token + embedding and `--mtp-moe-placement` (default `device`) the MTP-head experts; +* `dense/model.onnx` keeps the stock GPU `QMoE` on the UVA-mapped expert weights; `--dense-offload-qmoe` rewrites it to `QMoEOffload` + streaming `--prefill-chunk` rows instead; +* `--target-resident-slots` and `--target-spare-slots` set the resident and spare expert slots of every target node. + `--expert-profile ` (an STRP v1 expert-popularity file) chooses the initial resident experts per layer instead, + `--total-resident-slots` in total (default 48 x `--target-resident-slots`, or 48 x 20) with a floor of `--min-slots-per-layer` (default 4) + per layer, or the same count for every layer with `--uniform-profile-slots`; +* `genai_config.json` gets CUDA-graph capture and `kSameAsRequested` arena options for the decoder, MTP-head, dense and head-KV sessions; + `--vram-limit-gib` becomes `gpu_device_mem_limit`, `--stop-token-id` replaces `eos_token_id` and `--prefill-chunk` (default 512) sets + `prefill.chunk`; +* `--plane-lib` and `--provider-lib` are copied to `libs/` and `provider/`; `--ops-dir` adds its `libqwen_*.so` files to `libs/`. + +`--out` must not be equal to or nested in `--from-package` (or contain it). It does not have to be new or empty, so a rebuild replaces +its files, and for a source without `dense/model.onnx` the script removes stale `dense/` and `head_kv/` (flat source: also +`target/` and `head_r2/`) folders from it. The source package is never modified: every graph is saved to a new file, and the +`model_s2.onnx`, `model_s3.onnx` and `model_s4.onnx` of a source that already has them are regenerated instead of hard-linked. With +`--mtp-moe-placement host_mapped` it can also write a 4.7 GiB `_mtp_moe_aligned.bin` next to the source package (Linux only: +`os.copy_file_range`). Use `run_with_ram_budget.py` to run an offload package with a bounded host expert cache. + +`--migrate-shared-onnx --out-onnx ` is a separate mode that rewrites one graph: legacy 26-input `QMoEOffload` nodes +move to the compact schema (inputs 21, 24 and 25 become inputs 2, 3 and 4; `--strip-phase` drops the optional phase input and +`--strip-stats` the statistics outputs) without touching the external weight data. + ## Engine usage ```python @@ -202,6 +238,7 @@ The builders produce a package; every other script needs one, and the ones that | `build_package_engine_v12b.py` | v12b / v12c / v12d package (folded cycle + fast prefill, optional `--kv-quant`, `--weight-quant`, `--mtp-moe-quant`, `--hc-quant`, `--share-dense-weights`, and the flat `--clean-package` layout) from an existing v12b package | | `build_package_v12d_int4_int8.py` | INT4 QMoE + INT8 MatMulNBits variant of the v12d package for SM80 and newer (layered, or flat with `--clean-flat-package`) | | `build_package_contrib_ops.py` | rewrites a flat RC2 package to standard ONNX and `com.microsoft` contrib operators only (no `custom_ops_libraries`) | +| `build_package_offload_v3.py` | QMoE-offload package from a v12b package (`QMoEOffload` nodes, host-resident routed experts, S=2..4 target graphs, optional expert profile and resident slots); `--migrate-shared-onnx` migrates one graph to the compact `QMoEOffload` schema | | `qwen38_engine_util.py` | helpers shared by the scripts: CUDA EP registration and model loading with config overlays (`$QWEN38_OVERLAY_JSON`), the natural-text prompt pools, and `run_turn` (one greedy turn through `og.Engine`); registers the DLL directories on Windows | | `calibrate_qwen38_kv.py` | calibrate per-channel INT8 and FP8 KV-cache scales for the 12 target sparse-attention layers and 1 MTP-head sparse-attention layer (the JSON of the wanted quantization is the `--kv-scale-file` of `build_package_engine_v12b.py`); `--out` receives the `--quant-type` scales and the other quantization is written beside it under the name with the quantization token swapped (`kv_scales_int8_per_channel.json` -> `kv_scales_fp8_per_channel.json`, `scales.json` -> `scales.fp8_per_channel.json`), so the two files never collide (a collision is an error before any work) | | `smoke.py` | minimal example, prints the stream, speed and speculative statistics | diff --git a/tools/qwen38_graph/build_package_offload_v3.py b/tools/qwen38_graph/build_package_offload_v3.py new file mode 100644 index 0000000000..80748d2765 --- /dev/null +++ b/tools/qwen38_graph/build_package_offload_v3.py @@ -0,0 +1,907 @@ +#!/usr/bin/env python3 +# Copyright (c) Microsoft Corporation. All rights reserved. +# Licensed under the MIT License. +"""Zero-copy package builder and graph migrator for QMoEOffload v3 (NVFP4 and SharedExpertBlockV1 INT4). + +Supports two migration modes: +1. NVFP4 v12b package migration (--from-package --out ): + - Hard-links all weight files (0 GiB disk copy). + - Rewrites the 48 com.microsoft::QMoE nodes in target/model.onnx (or model.onnx) and dense/model.onnx to + com.microsoft::QMoEOffload with explicit layer_index, hidden_size, inter_size, and v3 dispatch controls. + - Updates weights_pool.txt so the 48 layers of routed expert weights/scales (63.281 GiB) use "cpu" + mmap placement, and optionally places dense/source_weights/dense.data (5.270 GiB) in "host_mapped" + UVA placement for 24 GiB VRAM targets. + - Configures genai_config.json for 24 GiB or 32 GiB VRAM operation. +2. Standalone SharedExpertBlockV1 v3 compact schema migration (--migrate-shared-onnx --out-onnx ): + - Rewrites legacy 26-input QMoEOffload nodes (21->2 expert_blocks, 24->3 fc2_bias, 25->4 optional phase) + without reading or regenerating the 63.457 GiB external weight payload (Section 5.3 / F8). +""" + +from __future__ import annotations + +import argparse +import hashlib +import json +import os +import shutil +from pathlib import Path + +import onnx +from onnx import helper + +REPLACED = { + Path("genai_config.json"), + Path("package_manifest.json"), + Path("README.md"), + Path("weights_pool.txt"), + Path("model.onnx"), + Path("model_s2.onnx"), + Path("model_s3.onnx"), + Path("model_s4.onnx"), + Path("target/model.onnx"), + Path("target/model_s2.onnx"), + Path("target/model_s3.onnx"), + Path("target/model_s4.onnx"), + Path("dense/model.onnx"), + Path("libs/libqwen38_plane.so"), + Path("provider/libonnxruntime_providers_cuda.so"), +} + + +def link_file(src: Path, dst: Path, copy: bool = False) -> None: + dst.parent.mkdir(parents=True, exist_ok=True) + if dst.exists() or dst.is_symlink(): + dst.unlink() + if src.is_symlink(): + dst.symlink_to(os.readlink(src)) + return + if not copy: + try: + os.link(src, dst) + return + except OSError: + pass # hard links can fail across filesystems; fall back to copying + shutil.copy2(src, dst) + + +def save_model_fresh(model: onnx.ModelProto, path: Path) -> None: + """Saves to a new inode: the output tree hard-links the source package, and a write must never go through.""" + path.parent.mkdir(parents=True, exist_ok=True) + path.unlink(missing_ok=True) + onnx.save_model(model, str(path), save_as_external_data=False) + + +def link_tree(src: Path, dst: Path, copy: bool = False) -> int: + count = 0 + for root, _dirs, files in os.walk(src, followlinks=False): + rel = Path(root).relative_to(src) + (dst / rel).mkdir(parents=True, exist_ok=True) + if rel.name == "_source_weights": + alias = (dst / rel).parent / "source_weights" + if not alias.exists() and not alias.is_symlink(): + alias.symlink_to("_source_weights") + for name in files: + if name.startswith(".azDownload"): + continue + relative = rel / name + if relative in REPLACED: + continue + link_file(Path(root) / name, dst / relative, copy) + count += 1 + fallback_experts = src.parent / "pkg_v12d_bf16lm_engine/target/source_weights/experts.data" + if not fallback_experts.exists(): + fallback_experts = src.parent / "qwen_flash_v12c_engine_int8kv_fp8w_fp4moe/target/source_weights/experts.data" + exp_rels = ( + [Path("source_weights/experts.data")] + if (src / "model.onnx").exists() + else [Path("target/source_weights/experts.data")] + ) + if (src / "dense").exists(): + exp_rels.append(Path("dense/source_weights/experts.data")) + for exp_rel in exp_rels: + dst_exp = dst / exp_rel + if not dst_exp.exists() and fallback_experts.exists(): + link_file(fallback_experts, dst_exp, copy) + count += 1 + return count + + +def sha256(path: Path) -> str: + digest = hashlib.sha256() + with path.open("rb") as handle: + for block in iter(lambda: handle.read(16 << 20), b""): + digest.update(block) + return digest.hexdigest() + + +def migrate_shared_onnx_v3(in_onnx: Path, out_onnx: Path, strip_phase: bool = False, strip_stats: bool = False) -> int: + """Zero-copy v3 compact shared schema migration (21->2, 24->3, 25->4) per Section 5.3 (F8).""" + model = onnx.load(str(in_onnx), load_external_data=False) + migrated = 0 + stat_outputs_removed = set() + for node in model.graph.node: + if node.domain == "com.microsoft" and node.op_type == "QMoEOffload": + inputs = list(node.input) + if len(inputs) >= 22 and inputs[21]: + x_in = inputs[0] + router_in = inputs[1] + blocks_in = inputs[21] + bias_in = inputs[24] if len(inputs) > 24 else "" + phase_in = "" if strip_phase else (inputs[25] if len(inputs) > 25 else "") + new_inputs = [x_in, router_in, blocks_in] + if bias_in or phase_in: + new_inputs.append(bias_in) + if phase_in: + new_inputs.append(phase_in) + del node.input[:] + node.input.extend(new_inputs) + if strip_stats and len(node.output) > 1: + for out_name in list(node.output)[1:]: + if out_name: + stat_outputs_removed.add(out_name) + del node.output[1:] + existing_attrs = {a.name for a in node.attribute} + if "layout_id" not in existing_attrs: + node.attribute.append(helper.make_attribute("layout_id", "ort.qmoe.w4a16.block32.shared.v1")) + if "layer_index" not in existing_attrs: + node.attribute.append(helper.make_attribute("layer_index", migrated)) + migrated += 1 + if stat_outputs_removed: + kept_outputs = [o for o in model.graph.output if o.name not in stat_outputs_removed] + del model.graph.output[:] + model.graph.output.extend(kept_outputs) + save_model_fresh(model, out_onnx) + return migrated + + +def load_expert_profile_plan( + profile_path: Path, + num_layers: int = 48, + num_experts: int = 512, + total_slots: int = 960, + min_slots_per_layer: int = 4, + uniform: bool = False, +) -> dict[int, list[int]]: + """Loads STRP v1 expert-profile.bin and assigns per-layer hot expert lists summing to total_slots.""" + import struct # noqa: PLC0415 + + import numpy as np # noqa: PLC0415 + + data = profile_path.read_bytes() + magic, version, hdr_layers, hdr_experts, _slots, n_ranked = struct.unpack("<6I", data[:24]) + if magic != 1347572819 or version != 1 or hdr_layers != num_layers or hdr_experts != num_experts: + raise ValueError( + f"Incompatible expert profile header in {profile_path}: {(magic, version, hdr_layers, hdr_experts)}" + ) + pairs = np.frombuffer(data[24 : 24 + n_ranked * 4], dtype=np.uint16).reshape(n_ranked, 2) + + by_layer: dict[int, list[int]] = {layer: [] for layer in range(num_layers)} + for l_id, e_id in pairs: + by_layer[int(l_id)].append(int(e_id)) + + if uniform: + per_layer = total_slots // num_layers + return {layer: by_layer[layer][:per_layer] for layer in range(num_layers)} + + selected: dict[int, list[int]] = {layer: [] for layer in range(num_layers)} + selected_set: dict[int, set[int]] = {layer: set() for layer in range(num_layers)} + budget = total_slots + floor = min(min_slots_per_layer, total_slots // num_layers, num_experts) + for layer in range(num_layers): + for e in by_layer[layer][:floor]: + selected[layer].append(e) + selected_set[layer].add(e) + budget -= 1 + + for l_id, e_id in pairs: + if budget <= 0: + break + layer = int(l_id) + e = int(e_id) + if e not in selected_set[layer] and len(selected[layer]) < num_experts: + selected[layer].append(e) + selected_set[layer].add(e) + budget -= 1 + + return selected + + +def rewrite_nvfp4_qmoe_to_offload( + in_onnx: Path, + out_onnx: Path, + hidden_size: int = 2560, + inter_size: int = 640, + token_tile_size: int = 32, + extra_attrs: dict | None = None, + per_layer_initial_experts: dict[int, list[int]] | None = None, +) -> int: + """Rewrites com.microsoft::QMoE (quant_type='nvfp4') nodes to com.microsoft::QMoEOffload.""" + model = onnx.load(str(in_onnx), load_external_data=False) + rewritten = 0 + for node in model.graph.node: + if node.domain == "com.microsoft" and node.op_type == "QMoE": + layer_idx = rewritten + node.op_type = "QMoEOffload" + if len(node.input) > 14 and node.input[14]: + node.input[14] = "" + existing_attrs = {a.name for a in node.attribute} + if "layer_index" not in existing_attrs: + node.attribute.append(helper.make_attribute("layer_index", layer_idx)) + if "hidden_size" not in existing_attrs: + node.attribute.append(helper.make_attribute("hidden_size", hidden_size)) + if "inter_size" not in existing_attrs: + node.attribute.append(helper.make_attribute("inter_size", inter_size)) + if "token_tile_size" not in existing_attrs: + node.attribute.append(helper.make_attribute("token_tile_size", token_tile_size)) + if "gpu_working_bytes" not in existing_attrs: + node.attribute.append(helper.make_attribute("gpu_working_bytes", 512 << 20)) + if "pinned_staging_bytes" not in existing_attrs: + node.attribute.append(helper.make_attribute("pinned_staging_bytes", 64 << 20)) + if "allow_cuda_graph" not in existing_attrs: + node.attribute.append(helper.make_attribute("allow_cuda_graph", 1)) + if per_layer_initial_experts is not None and layer_idx in per_layer_initial_experts: + init_exps = per_layer_initial_experts[layer_idx] + node.attribute.append(helper.make_attribute("resident_slots", len(init_exps))) + node.attribute.append(helper.make_attribute("initial_experts", init_exps)) + existing_attrs.add("resident_slots") + existing_attrs.add("initial_experts") + if extra_attrs: + for k_attr, v_attr in extra_attrs.items(): + if k_attr not in existing_attrs and v_attr is not None: + node.attribute.append(helper.make_attribute(k_attr, v_attr)) + rewritten += 1 + save_model_fresh(model, out_onnx) + return rewritten + + +def build_multidraft_target_onnx(in_onnx: Path, out_onnx: Path, verify_rows: int = 3) -> dict: + """Extends an S=2 folded target/model.onnx to verify_rows in {3, 4} (2 or 3 chained MTP drafts).""" + import copy # noqa: PLC0415 + import re # noqa: PLC0415 + + from onnx import numpy_helper # noqa: PLC0415 + + if verify_rows not in (3, 4): + raise ValueError(f"verify_rows must be 3 or 4, got {verify_rows}") + + model = onnx.load(str(in_onnx), load_external_data=False) + graph = model.graph + output_info = {o.name: o for o in graph.output} + init_names = {i.name for i in graph.initializer} + + def add_graph_output(source_name: str, new_name: str, shape_override: list[int] | None = None) -> None: + if any(o.name == new_name for o in graph.output): + return + info = output_info[source_name] + new_vi = onnx.ValueInfoProto() + new_vi.CopyFrom(info) + new_vi.name = new_name + if shape_override is not None: + del new_vi.type.tensor_type.shape.dim[:] + for d in shape_override: + dim = new_vi.type.tensor_type.shape.dim.add() + dim.dim_value = d + graph.output.append(new_vi) + output_info[new_name] = new_vi + + def layer_of(name: str, suffix: str) -> str: + m = re.fullmatch(r"present\.(\d+)\." + re.escape(suffix), name) + if not m: + raise ValueError(f"Cannot parse layer from {name!r}") + return m.group(1) + + mids_added = 0 + fold_mid_idx = -1 + for idx, node in enumerate(graph.node): + if node.name == "model/hyper_connection_mixer/Slice_last2": + if "const_last_starts" not in init_names: + graph.initializer.append( + numpy_helper.from_array( + __import__("numpy").asarray([-verify_rows], dtype=__import__("numpy").int64), + name="const_last_starts", + ) + ) + init_names.add("const_last_starts") + node.input[1] = "const_last_starts" + if node.domain == "qwen.fold" and node.op_type in ("FoldMid", "FoldMid2"): + keep_a = [a for a in node.attribute if a.name != "head_rows"] + del node.attribute[:] + node.attribute.extend(keep_a) + node.attribute.append(helper.make_attribute("head_rows", verify_rows)) + fold_mid_idx = idx + break + if node.domain == "qwen.fused_attn" and node.op_type in ("GdnProjPrep", "GdnConvPrep"): + state_idx = 4 if node.op_type == "GdnProjPrep" else 3 + src_out = node.output[state_idx] + lid = layer_of(src_out, "conv_state") + for m_tag in ("present_mid2", "present_mid3")[: verify_rows - 2]: + out_name = f"{m_tag}.{lid}.conv_state" + node.output.append(out_name) + add_graph_output(src_out, out_name) + mids_added += 1 + elif node.domain == "qwen.fused_attn" and node.op_type == "GdnRecurrent": + src_out = node.output[1] + lid = layer_of(src_out, "recurrent_state") + for m_tag in ("present_mid2", "present_mid3")[: verify_rows - 2]: + out_name = f"{m_tag}.{lid}.recurrent_state" + node.output.append(out_name) + add_graph_output(src_out, out_name) + mids_added += 1 + elif node.domain == "qwen.fused_attn" and node.op_type == "PleMidState": + for m_tag in ("present_mid2", "present_mid3")[: verify_rows - 2]: + st_name = f"{m_tag}.1.ple_conv_state" + cx_name = f"{m_tag}.1.ple_context" + node.output.extend([st_name, cx_name]) + add_graph_output("present.1.ple_conv_state", st_name) + add_graph_output("present.1.ple_context", cx_name) + mids_added += 2 + elif node.domain == "qwen.capture" and node.op_type == "PackedQSAIndexer": + src_buf, src_len = node.output[3], node.output[5] + lid = layer_of(src_buf, "index_buffer") + for m_tag in ("present_mid2", "present_mid3")[: verify_rows - 2]: + b_name = f"{m_tag}.{lid}.index_buffer" + l_name = f"{m_tag}.{lid}.index_lengths" + node.output.extend([b_name, l_name]) + add_graph_output(src_buf, b_name) + add_graph_output(src_len, l_name) + mids_added += 2 + + if fold_mid_idx < 0: + raise ValueError("FoldMid2 not found in target graph") + + for init in graph.initializer: + if init.name == "head.capture.cumulative_sequence_lengths": + init.CopyFrom( + numpy_helper.from_array( + __import__("numpy").asarray([0, verify_rows], dtype=__import__("numpy").int32), + name="head.capture.cumulative_sequence_lengths", + ) + ) + if "head.capture.cu_seqlens_s1" not in init_names: + graph.initializer.append( + numpy_helper.from_array( + __import__("numpy").asarray([0, 1], dtype=__import__("numpy").int32), + name="head.capture.cu_seqlens_s1", + ) + ) + init_names.add("head.capture.cu_seqlens_s1") + + stage1_nodes = list(graph.node[fold_mid_idx + 1 :]) + del graph.node[fold_mid_idx + 1 :] + + stage_template = [ + copy.deepcopy(n) for n in stage1_nodes if n.op_type not in ("RowSelectBf16x3", "Concat", "FoldTail") + ] + stage1_produced = {o for n in stage1_nodes for o in n.output if o} + + s1_val_inputs = [] + for node in stage1_nodes: + if node.op_type == "Concat" and node.name == "head.validation/aggregate": + s1_val_inputs = list(node.input) + continue + if node.op_type == "FoldTail": + continue + if node.domain == "qwen.capture" and node.op_type == "PackedQSAIndexer": + node.output[2] = "head.present_s1.0.index_key" + add_graph_output("head.present.0.index_key", "head.present_s1.0.index_key") + for m_tag in ("head.present_mid2", "head.present_mid3")[: verify_rows - 2]: + b_name = f"{m_tag}.0.index_buffer" + l_name = f"{m_tag}.0.index_lengths" + node.output.extend([b_name, l_name]) + add_graph_output("head.present.0.index_buffer", b_name) + add_graph_output("head.present.0.index_lengths", l_name) + elif node.domain == "qwen.static" and node.op_type == "GuardedDSA": + node.output[1] = "head.present_s1.0.key" + node.output[2] = "head.present_s1.0.value" + add_graph_output("head.present.0.key", "head.present_s1.0.key") + add_graph_output("head.present.0.value", "head.present_s1.0.value") + elif node.domain == "qwen.mtp" and node.op_type == "Bf16ArgMax": + node.output[0] = "head.s1.next_token" + graph.node.append(node) + + all_val_inputs = list(s1_val_inputs) + num_extra_stages = verify_rows - 2 + + for extra_idx in range(num_extra_stages): + stage_num = extra_idx + 2 + prefix = f"head.s{stage_num}." + is_last_stage = extra_idx == num_extra_stages - 1 + prev_tok = "head.s1.next_token" if stage_num == 2 else "head.s2.next_token" + prev_hidden = "head.hidden_states_out" if stage_num == 2 else "head.s2.head.hidden_states_out" + prev_key = "head.present_s1.0.key" if stage_num == 2 else "head.present_s2.0.key" + prev_val = "head.present_s1.0.value" if stage_num == 2 else "head.present_s2.0.value" + prev_ikey = "head.present_s1.0.index_key" if stage_num == 2 else "head.present_s2.0.index_key" + + if stage_num == 2: + prep_inputs = [ + prev_tok, + "fold.ctl", + "fold.sel", + "head.mtp.rotary_emb.cos_cache", + "head.mtp.rotary_emb.sin_cache", + "head.present_mid.0.index_buffer", + "head.present_mid.0.index_lengths", + "head.present_mid2.0.index_buffer", + "head.present_mid2.0.index_lengths", + ] + if verify_rows == 4: + prep_inputs.extend(["head.present_mid3.0.index_buffer", "head.present_mid3.0.index_lengths"]) + prep_inputs.extend(["head.present.0.index_buffer", "head.present.0.index_lengths"]) + else: + prep_inputs = [ + prev_tok, + "fold.ctl", + "fold.sel", + "head.mtp.rotary_emb.cos_cache", + "head.mtp.rotary_emb.sin_cache", + "head.s2.out_index_buffer", + "head.s2.out_index_lengths", + ] + + prep_outputs = [ + f"{prefix}head_ids", + f"{prefix}head_c", + f"{prefix}slot_ids", + f"{prefix}seqlens_k", + f"{prefix}rope_cos", + f"{prefix}rope_sin", + f"{prefix}in_index_buffer", + f"{prefix}in_index_lengths", + ] + graph.node.append( + helper.make_node( + "FoldDraftPrep", + prep_inputs, + prep_outputs, + domain="qwen.fold", + name=f"fold/FoldDraftPrep_s{stage_num}", + stage=stage_num - 1, + ) + ) + + input_remap = { + "fold.head_ids": f"{prefix}head_ids", + "hidden_states": prev_hidden, + "fold.head_hidden": prev_hidden, + "head.capture.rope_cos": f"{prefix}rope_cos", + "head.capture.rope_sin": f"{prefix}rope_sin", + "head.capture.cumulative_sequence_lengths": "head.capture.cu_seqlens_s1", + "fold.head_c": f"{prefix}head_c", + "head.capture.slot_ids": f"{prefix}slot_ids", + "head.capture.seqlens_k": f"{prefix}seqlens_k", + "head.past_key_values.0.index_key": prev_ikey, + "head.past_key_values.0.index_buffer": f"{prefix}in_index_buffer", + "head.past_key_values.0.index_lengths": f"{prefix}in_index_lengths", + "head.past_key_values.0.key": prev_key, + "head.past_key_values.0.value": prev_val, + "head.v_mtp.layers.0.self_attn.Mul_200.sel": f"{prefix}head.v_mtp.layers.0.self_attn.Mul_200", + "head.v_mtp.Reshape_23.sel": f"{prefix}head.v_mtp.Reshape_23", + "head.v_mtp.layers.0.attn_hyper_connection.Mul_117.sel": f"{prefix}head.v_mtp.layers.0.attn_hyper_connection.Mul_117", + } + + for tmpl in stage_template: + n = copy.deepcopy(tmpl) + n.name = f"{prefix}{n.name}" + for i, inp in enumerate(n.input): + if not inp: + continue + if inp in input_remap: + n.input[i] = input_remap[inp] + elif inp in stage1_produced and inp not in init_names: + n.input[i] = f"{prefix}{inp}" + for i, out in enumerate(n.output): + if out: + n.output[i] = f"{prefix}{out}" + + if n.domain == "qwen.capture" and n.op_type == "PackedQSAIndexer": + del n.output[7:] + if is_last_stage: + n.output[2] = "head.present.0.index_key" + else: + n.output[2] = f"head.present_s{stage_num}.0.index_key" + add_graph_output("head.present.0.index_key", n.output[2]) + n.output[3] = f"{prefix}out_index_buffer" + n.output[5] = f"{prefix}out_index_lengths" + n.output[6] = f"{prefix}validation.qsa" + all_val_inputs.append(n.output[6]) + elif n.domain == "qwen.static" and n.op_type == "GuardedDSA": + if is_last_stage: + n.output[1] = "head.present.0.key" + n.output[2] = "head.present.0.value" + else: + n.output[1] = f"head.present_s{stage_num}.0.key" + n.output[2] = f"head.present_s{stage_num}.0.value" + add_graph_output("head.present.0.key", n.output[1]) + add_graph_output("head.present.0.value", n.output[2]) + n.output[3] = f"{prefix}validation.dsa" + all_val_inputs.append(n.output[3]) + elif n.domain == "qwen.mtp" and n.op_type == "Bf16ArgMax": + n.output[0] = "head.next_token" if is_last_stage else f"{prefix}next_token" + + graph.node.append(n) + + for o in graph.output: + if o.name == "head.validation.aggregate": + del o.type.tensor_type.shape.dim[:] + dim = o.type.tensor_type.shape.dim.add() + dim.dim_value = len(all_val_inputs) + + graph.node.append( + helper.make_node( + "Concat", + all_val_inputs, + ["head.validation.aggregate"], + name="head.validation/aggregate", + axis=0, + ) + ) + tail_ctl_slot = 8 if verify_rows == 3 else 9 + graph.node.append( + helper.make_node( + "FoldTail", + ["head.next_token", "head.validation.aggregate", "fold.ctl"], + ["fold.tail_done"], + domain="qwen.fold", + name="fold/FoldTail", + ctl_slot=tail_ctl_slot, + ) + ) + + del graph.value_info[:] + save_model_fresh(model, out_onnx) + return { + "verify_rows": verify_rows, + "mids_added": mids_added, + "extra_head_stages": num_extra_stages, + "total_nodes": len(graph.node), + "total_outputs": len(graph.output), + } + + +def build_offload_weights_pool( + source_pkg: Path, + out_pkg: Path, + dense_placement: str = "device", + mtp_moe_placement: str = "device", + embed_placement: str = "host_mapped", +) -> dict: + """Rewrites weights_pool.txt with CPU placement for QMoE experts and optional host_mapped for dense.data.""" + src_pool = (source_pkg / "weights_pool.txt").read_text().splitlines() + assert src_pool[0].strip() == "qwen38_weights_pool 1" + + pool_lines = [] + session_lines = [] + qmoe_pool_ids = set() + mtp_moe_pool_ids = set() + embed_pool_ids = set() + dense_bound_names = set() + + for line in src_pool[1:]: + s = line.strip() + if not s: + continue + parts = s.split() + if parts[0] == "session": + sess, name, pid = parts[1], parts[2], int(parts[3]) + if ".mlp.fc1_" in name or ".mlp.fc2_" in name: + qmoe_pool_ids.add(pid) + if "mtp.layers.0.mlp.experts." in name: + mtp_moe_pool_ids.add(pid) + if name.endswith("embed_tokens.weight"): + embed_pool_ids.add(pid) + if sess == "dense": + dense_bound_names.add(name) + session_lines.append((sess, name, pid)) + + aligned_mtp_rel = "target/head_data/source_weights/mtp_moe_aligned.bin" + need_aligned_mtp_copy = False + for line in src_pool[1:]: + parts = line.strip().split() + if parts and parts[0] in ("pool", "host_pool") and int(parts[1]) in mtp_moe_pool_ids: + if int(parts[3]) % 16 != 0: + need_aligned_mtp_copy = True + + if mtp_moe_placement != "device" and need_aligned_mtp_copy: + shared_aligned = source_pkg.parent / f"{source_pkg.name}_mtp_moe_aligned.bin" + src_st = source_pkg / "target/head_data/source_weights/model-bf16-00011.safetensors" + expected_size = 3355443200 + 1677721600 + if not shared_aligned.exists() or shared_aligned.stat().st_size != expected_size: + tmp_aligned = shared_aligned.with_suffix(".tmp") + with open(src_st, "rb") as fin, open(tmp_aligned, "wb") as fout: + in_fd = fin.fileno() + out_fd = fout.fileno() + for src_off, dst_off, total_len in [ + (7243385096, 0, 3355443200), + (5565663496, 3355443200, 1677721600), + ]: + rem = total_len + r_off = src_off + w_off = dst_off + while rem > 0: + chunk = min(rem, 1 << 26) + copied = os.copy_file_range(in_fd, out_fd, chunk, offset_src=r_off, offset_dst=w_off) + if copied <= 0: + raise RuntimeError("os.copy_file_range returned 0") + r_off += copied + w_off += copied + rem -= copied + tmp_aligned.replace(shared_aligned) + link_file(shared_aligned, out_pkg / aligned_mtp_rel, copy=False) + + for line in src_pool[1:]: + s = line.strip() + if not s: + continue + parts = s.split() + if parts[0] in ("pool", "host_pool"): + pid = int(parts[1]) + rank = int(parts[5]) + base_tokens = ["pool", *parts[1 : 7 + rank]] + base_s = " ".join(base_tokens) + default_placement = ( + "host_mapped" if parts[0] == "host_pool" else (parts[7 + rank] if len(parts) > 7 + rank else "device") + ) + if pid in qmoe_pool_ids: + pool_lines.append(f"{base_s} cpu") + elif pid in embed_pool_ids: + if embed_placement != "device": + pool_lines.append(f"{base_s} {embed_placement}") + else: + pool_lines.append(base_s) + elif pid in mtp_moe_pool_ids and mtp_moe_placement != "device": + if not need_aligned_mtp_copy: + pool_lines.append(f"{base_s} {mtp_moe_placement}") + elif pid == 607: + pool_lines.append(f"pool 607 16 0 3355443200 3 512 1280 2560 {aligned_mtp_rel} {mtp_moe_placement}") + elif pid == 608: + pool_lines.append( + f"pool 608 16 3355443200 1677721600 3 512 2560 640 {aligned_mtp_rel} {mtp_moe_placement}" + ) + elif default_placement != "device": + pool_lines.append(f"{base_s} {default_placement}") + else: + pool_lines.append(base_s) + + next_pool_id = len(pool_lines) + target_bound_names = {name for sess, name, _ in session_lines if sess == "target"} + flat_root = (source_pkg / "model.onnx").exists() + target_onnx_src = (source_pkg / "model.onnx") if flat_root else (source_pkg / "target/model.onnx") + target_model = onnx.load(str(target_onnx_src), load_external_data=False) + for init in target_model.graph.initializer: + if init.name in target_bound_names: + continue + if ".mlp.fc1_" in init.name or ".mlp.fc2_" in init.name: + ext = {e.key: e.value for e in init.external_data} + loc = ext.get("location", "") + if not loc: + continue + offset = int(ext.get("offset", "0")) + length = int(ext["length"]) + dims = list(init.dims) + shape_str = f"{len(dims)} " + " ".join(str(d) for d in dims) + pool_file = loc if flat_root else f"target/{loc}" + pool_lines.append(f"pool {next_pool_id} {init.data_type} {offset} {length} {shape_str} {pool_file} cpu") + session_lines.append(("target", init.name, next_pool_id)) + qmoe_pool_ids.add(next_pool_id) + next_pool_id += 1 + + added_dense_entries = 0 + if dense_placement in ("host_mapped", "cpu") and (source_pkg / "dense/model.onnx").exists(): + dense_model = onnx.load(str(source_pkg / "dense/model.onnx"), load_external_data=False) + for init in dense_model.graph.initializer: + if init.name in dense_bound_names: + continue + ext = {e.key: e.value for e in init.external_data} + loc = ext.get("location", "") + if loc == "source_weights/dense.data": + offset = int(ext.get("offset", "0")) + length = int(ext["length"]) + dims = list(init.dims) + shape_str = f"{len(dims)} " + " ".join(str(d) for d in dims) + pool_file = "target/source_weights/dense.data" + pool_lines.append( + f"pool {next_pool_id} {init.data_type} {offset} {length} {shape_str} {pool_file} {dense_placement}" + ) + session_lines.append(("dense", init.name, next_pool_id)) + next_pool_id += 1 + added_dense_entries += 1 + + out_lines = ["qwen38_weights_pool 1", *pool_lines] + for sess, name, pid in session_lines: + out_lines.append(f"session {sess} {name} {pid}") + (out_pkg / "weights_pool.txt").write_text("\n".join(out_lines) + "\n") + return { + "total_pool_entries": next_pool_id, + "cpu_qmoe_entries": len(qmoe_pool_ids), + "added_dense_entries": added_dense_entries, + "dense_placement": dense_placement, + "mtp_moe_placement": mtp_moe_placement, + "embed_placement": embed_placement, + } + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) + parser.add_argument( + "--migrate-shared-onnx", type=Path, default=None, help="Input ONNX for standalone v3 shared migration" + ) + parser.add_argument("--out-onnx", type=Path, default=None, help="Output ONNX for standalone v3 shared migration") + parser.add_argument("--strip-phase", action="store_true") + parser.add_argument("--strip-stats", action="store_true") + + parser.add_argument("--from-package", type=Path, default=None) + parser.add_argument("--out", type=Path, default=None) + parser.add_argument("--plane-lib", type=Path, default=None) + parser.add_argument("--provider-lib", type=Path, default=None) + parser.add_argument( + "--ops-dir", + type=Path, + default=None, + help="Directory containing rebased libqwen_*.so custom op libraries", + ) + parser.add_argument("--stop-token-id", type=int, default=None) + parser.add_argument("--prefill-chunk", type=int, default=512) + parser.add_argument("--target-resident-slots", type=int, default=None) + parser.add_argument("--target-spare-slots", type=int, default=None) + parser.add_argument("--expert-profile", type=Path, default=None, help="Path to STRP v1 expert-profile.bin") + parser.add_argument("--total-resident-slots", type=int, default=None, help="Total resident slots across 48 layers") + parser.add_argument("--min-slots-per-layer", type=int, default=4) + parser.add_argument("--uniform-profile-slots", action="store_true") + parser.add_argument("--dense-stream-min-rows", type=int, default=1) + parser.add_argument("--dense-stream-wave", type=int, default=64) + parser.add_argument( + "--dense-gpu-qmoe", + action="store_true", + default=True, + help="Keep stock GPU QMoE in dense/model.onnx using UVA-mapped expert weights", + ) + parser.add_argument("--dense-offload-qmoe", dest="dense_gpu_qmoe", action="store_false") + parser.add_argument("--dense-placement", choices=["device", "host_mapped", "cpu"], default="host_mapped") + parser.add_argument("--mtp-moe-placement", choices=["device", "host_mapped"], default="device") + parser.add_argument("--embed-placement", choices=["device", "host_mapped"], default="host_mapped") + parser.add_argument("--vram-limit-gib", type=float, default=None) + parser.add_argument("--copy", action="store_true") + args = parser.parse_args() + + if args.migrate_shared_onnx is not None: + if args.out_onnx is None: + raise SystemExit("--out-onnx is required with --migrate-shared-onnx") + n = migrate_shared_onnx_v3(args.migrate_shared_onnx, args.out_onnx, args.strip_phase, args.strip_stats) + print(f"Migrated {n} QMoEOffload nodes to v3 compact shared schema -> {args.out_onnx}") + return + + if args.from_package is None or args.out is None or args.plane_lib is None or args.provider_lib is None: + raise SystemExit("--from-package, --out, --plane-lib, and --provider-lib are required for package build") + + source = args.from_package.resolve() + out = args.out.resolve() + if out == source or out.is_relative_to(source) or source.is_relative_to(out): + raise SystemExit("--out must not be equal to or nested inside --from-package") + out.mkdir(parents=True, exist_ok=True) + + linked = link_tree(source, out, args.copy) + if args.ops_dir is not None and args.ops_dir.exists(): + for so_file in sorted(args.ops_dir.glob("libqwen_*.so")): + link_file(so_file.resolve(), out / "libs" / so_file.name, copy=True) + link_file(args.plane_lib.resolve(), out / "libs/libqwen38_plane.so", copy=True) + link_file(args.provider_lib.resolve(), out / "provider/libonnxruntime_providers_cuda.so", copy=True) + + profile_plan = None + if args.expert_profile is not None: + total_slots = args.total_resident_slots or ((args.target_resident_slots or 20) * 48) + profile_plan = load_expert_profile_plan( + args.expert_profile.resolve(), + num_layers=48, + num_experts=512, + total_slots=total_slots, + min_slots_per_layer=args.min_slots_per_layer, + uniform=args.uniform_profile_slots, + ) + + target_attrs = { + "resident_slots": args.target_resident_slots, + "spare_slots": args.target_spare_slots, + "stream_rows": 0, + } + dense_attrs = { + "resident_slots": 0, + "spare_slots": 0, + "stream_rows": args.prefill_chunk, + "stream_min_tokens": 1, + "stream_min_rows": args.dense_stream_min_rows, + "stream_wave": args.dense_stream_wave, + } + flat_root = (source / "model.onnx").exists() + target_src = (source / "model.onnx") if flat_root else (source / "target/model.onnx") + target_out_dir = out if flat_root else (out / "target") + target_rewritten = rewrite_nvfp4_qmoe_to_offload( + target_src, + target_out_dir / "model_s2.onnx", + extra_attrs=target_attrs, + per_layer_initial_experts=profile_plan, + ) + build_multidraft_target_onnx(target_out_dir / "model_s2.onnx", target_out_dir / "model_s3.onnx", verify_rows=3) + build_multidraft_target_onnx(target_out_dir / "model_s2.onnx", target_out_dir / "model_s4.onnx", verify_rows=4) + link_file(target_out_dir / "model_s4.onnx", target_out_dir / "model.onnx", copy=True) + if (source / "dense/model.onnx").exists(): + if args.dense_gpu_qmoe: + link_file(source / "dense/model.onnx", out / "dense/model.onnx", args.copy) + dense_rewritten = 0 + else: + dense_rewritten = rewrite_nvfp4_qmoe_to_offload( + source / "dense/model.onnx", out / "dense/model.onnx", extra_attrs=dense_attrs + ) + else: + dense_rewritten = 0 + for obsolete_dir in ("target", "head_r2", "dense", "head_kv") if flat_root else ("dense", "head_kv"): + if (out / obsolete_dir).exists(): + shutil.rmtree(out / obsolete_dir) + pool_stats = build_offload_weights_pool( + source, out, args.dense_placement, args.mtp_moe_placement, args.embed_placement + ) + + config = json.loads((source / "genai_config.json").read_text()) + if args.stop_token_id is not None: + config["model"]["eos_token_id"] = [args.stop_token_id] + + cuda_opts = config["model"]["decoder"]["session_options"]["provider_options"][0]["cuda"] + cuda_opts["enable_cuda_graph"] = "1" + cuda_opts["arena_extend_strategy"] = "kSameAsRequested" + cuda_opts["arena.extend_strategy"] = "1" + if args.vram_limit_gib is not None and args.vram_limit_gib > 0: + cuda_opts["gpu_device_mem_limit"] = str(int(args.vram_limit_gib * (1024**3))) + + qwen = config["model"]["qwen38_graph"] + fold_head_key = "mtp" if "mtp" in qwen["fold"] else ("drafter" if "drafter" in qwen["fold"] else "head_r2") + qwen["fold"][fold_head_key]["session_options"] = { + "log_id": "onnxruntime-genai", + "graph_optimization_level": "ORT_DISABLE_ALL", + "provider_options": [ + { + "cuda": { + "enable_cuda_graph": "1", + "min_num_runs_before_cuda_graph_capture": "4", + "arena_extend_strategy": "kSameAsRequested", + "arena.extend_strategy": "1", + } + } + ], + "session.use_device_allocator_for_initializers": "1", + } + qwen["prefill"]["chunk"] = args.prefill_chunk + qwen["prefill"]["warmup_lengths"] = [min(args.prefill_chunk, 512), 128, 37] + if "dense" in qwen["prefill"]: + dense_cuda = qwen["prefill"]["dense"]["session_options"]["provider_options"][0]["cuda"] + dense_cuda["arena_extend_strategy"] = "kSameAsRequested" + dense_cuda["arena.extend_strategy"] = "1" + if "head_kv" in qwen["prefill"]: + head_kv_cuda = qwen["prefill"]["head_kv"]["session_options"]["provider_options"][0]["cuda"] + head_kv_cuda["arena_extend_strategy"] = "kSameAsRequested" + head_kv_cuda["arena.extend_strategy"] = "1" + + (out / "genai_config.json").write_text(json.dumps(config, indent=2) + "\n") + + libs = {p.name: sha256(p) for p in sorted((out / "libs").glob("*.so"))} + providers = {p.name: sha256(p) for p in sorted((out / "provider").glob("*.so"))} + manifest = { + "from_package": str(source), + "stack": "v12b-engine-qmoe-offload-v3", + "hard_linked_files": linked, + "target_qmoe_offload_nodes": target_rewritten, + "dense_qmoe_offload_nodes": dense_rewritten, + "pool_stats": pool_stats, + "eos_token_id": config["model"]["eos_token_id"], + "vram_limit_gib": args.vram_limit_gib, + "plane_library": {"source": str(args.plane_lib.resolve()), "sha256": libs.get("libqwen38_plane.so")}, + "libs_sha256": libs, + "provider_sha256": providers, + } + (out / "package_manifest.json").write_text(json.dumps(manifest, indent=2) + "\n") + print( + f"Built offload v3 package: {out} (linked={linked}, target_qmoe={target_rewritten}, " + f"dense_qmoe={dense_rewritten}, pool={pool_stats})" + ) + + +if __name__ == "__main__": + main() From 6c5cfd887a141282772a00581d60e4db174477ca Mon Sep 17 00:00:00 2001 From: Tianlei Wu Date: Fri, 30 Oct 2026 02:37:42 +0000 Subject: [PATCH 2/6] Emit multi-draft graphs and offload profiles from clean packages Let the flat-package builders write the S=2..4 target graphs that the offload package builder provides: - build_package_engine_v12b.py --clean-package links model_s2.onnx to model.onnx, generates model_s3.onnx and model_s4.onnx with build_multidraft_target_onnx, and writes four runtime profiles for GPUs with 20 to 34 GiB of device memory. Each profile selects model_s4.onnx and sets the QMoE offload session options (placement, device and host expert budgets, expert popularity file). --expert-profile copies an STRP v1 popularity file into the package; without it the builder prints a note and the file has to be copied in by hand. - build_package_v12d_int4_int8.py --clean-flat-package writes the same S=2..4 graphs. The reused config carries the runtime profiles over. The existing check for a consumer-gpu-24g profile in that config is kept as is; it matches nothing because the profiles are stored as a list. Both builders import build_multidraft_target_onnx lazily, so the flat packages need build_package_offload_v3.py in the same folder. The profiles use host-memory eligibility and a decoder session-option overlay, so the packages load only in a runtime whose profile parser accepts them. The README documents the graphs and the profiles. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- tools/qwen38_graph/README.md | 22 +++- .../qwen38_graph/build_package_engine_v12b.py | 117 +++++++++++++++++- .../build_package_v12d_int4_int8.py | 14 ++- 3 files changed, 147 insertions(+), 6 deletions(-) diff --git a/tools/qwen38_graph/README.md b/tools/qwen38_graph/README.md index 607deaf759..dc458f1672 100644 --- a/tools/qwen38_graph/README.md +++ b/tools/qwen38_graph/README.md @@ -122,6 +122,17 @@ root layout instead of the `target/`, `head_r2/`, `dense/` and `head_kv/` folder cache in BF16 while the target KV cache is INT8; `--keep-bf16-router-shared` and `--keep-bf16-shared-down` keep the MoE router and the shared-expert projections in BF16. +The flat package also carries the multi-draft target graphs: `model_s2.onnx` is a link to `model.onnx`, and `model_s3.onnx` and +`model_s4.onnx` (3 and 4 verify rows) are generated from it by `build_multidraft_target_onnx` of `build_package_offload_v3.py`, which the +builder imports. `genai_config.json` gets four `runtime_profiles` (`consumer-gpu-24g`, `consumer-gpu-24g-host-64g`, +`consumer-gpu-30g-ram-48g-32k` and `consumer-gpu-32g`). Each is eligible for a range of total device memory and effective host memory on a +discrete GPU, selects `model_s4.onnx` as the decoder and sets the `session.moe.offload.*` session options (`mode`, `placement`, +`cpu_execution`, `device_expert_bytes`, `host_expert_bytes` in the profiles that limit the host memory, and `popularity_file`). The +eligibility uses `minimum_effective_host_memory_bytes` / `maximum_effective_host_memory_bytes` and the overlay a decoder +`session_options` block: a runtime whose profile parser does not know them rejects the `genai_config.json`. The +profiles name `expert-profile.bin`; pass an STRP v1 expert-popularity file with `--expert-profile ` and the builder +copies it into the package under that name (without the option, copy the file into the package yourself before using the profiles). + ### Config surface of the built packages The engine rejects every `model.qwen38_graph` key it does not know when a package is loaded, so the builders only write the keys of the @@ -153,7 +164,9 @@ existing packages: * the package is self-contained: no `.safetensors` files and no unreferenced bytes of the source weights. `--clean-flat-package` writes the flat root layout; it needs `--reuse-embed-from `, which supplies `embedding.onnx` + -`embedding.onnx.data`, `attn_consts.data`, `genai_config.json` and the tokenizer files. `--reuse-ple-from` and `--reuse-qmoe-from` +`embedding.onnx.data`, `attn_consts.data`, `genai_config.json` (so the flat package inherits its `runtime_profiles`) and the tokenizer +files. Like the clean v12b package it contains `model_s2.onnx`, `model_s3.onnx` and `model_s4.onnx`, generated by +`build_multidraft_target_onnx` of `build_package_offload_v3.py`. `--reuse-ple-from` and `--reuse-qmoe-from` hard-link the PLE shards and the INT4 expert data of an earlier build instead of recomputing them. `--hf-dir`, `--v12b-dir`, `--v12d-dir`, `--ops-dir` and `--out` are required. `--out` must not be equal to, nested in or around any input (`--hf-dir`, `--v12b-dir`, `--v12d-dir`, `--ops-dir`, `--reuse-ple-from`, `--reuse-qmoe-from`, `--reuse-embed-from`, `--plane-lib`, @@ -211,6 +224,9 @@ its files, and for a source without `dense/model.onnx` the script removes stale move to the compact schema (inputs 21, 24 and 25 become inputs 2, 3 and 4; `--strip-phase` drops the optional phase input and `--strip-stats` the statistics outputs) without touching the external weight data. +`build_multidraft_target_onnx` of this script also generates the S=3 and S=4 graphs of the flat packages: `build_package_engine_v12b.py` +(with `--clean-package`) and `build_package_v12d_int4_int8.py` (with `--clean-flat-package`) import it from this folder. + ## Engine usage ```python @@ -235,8 +251,8 @@ The builders produce a package; every other script needs one, and the ones that | script | purpose | | --- | --- | | `build_package.py` | v11 package (S=1 / S=2 chain verify) from its parts, with the PLE table | -| `build_package_engine_v12b.py` | v12b / v12c / v12d package (folded cycle + fast prefill, optional `--kv-quant`, `--weight-quant`, `--mtp-moe-quant`, `--hc-quant`, `--share-dense-weights`, and the flat `--clean-package` layout) from an existing v12b package | -| `build_package_v12d_int4_int8.py` | INT4 QMoE + INT8 MatMulNBits variant of the v12d package for SM80 and newer (layered, or flat with `--clean-flat-package`) | +| `build_package_engine_v12b.py` | v12b / v12c / v12d package (folded cycle + fast prefill, optional `--kv-quant`, `--weight-quant`, `--mtp-moe-quant`, `--hc-quant`, `--share-dense-weights`, and the flat `--clean-package` layout with the S=2..4 target graphs and offload `runtime_profiles`) from an existing v12b package | +| `build_package_v12d_int4_int8.py` | INT4 QMoE + INT8 MatMulNBits variant of the v12d package for SM80 and newer (layered, or flat with `--clean-flat-package` and the S=2..4 target graphs) | | `build_package_contrib_ops.py` | rewrites a flat RC2 package to standard ONNX and `com.microsoft` contrib operators only (no `custom_ops_libraries`) | | `build_package_offload_v3.py` | QMoE-offload package from a v12b package (`QMoEOffload` nodes, host-resident routed experts, S=2..4 target graphs, optional expert profile and resident slots); `--migrate-shared-onnx` migrates one graph to the compact `QMoEOffload` schema | | `qwen38_engine_util.py` | helpers shared by the scripts: CUDA EP registration and model loading with config overlays (`$QWEN38_OVERLAY_JSON`), the natural-text prompt pools, and `run_turn` (one greedy turn through `og.Engine`); registers the DLL directories on Windows | diff --git a/tools/qwen38_graph/build_package_engine_v12b.py b/tools/qwen38_graph/build_package_engine_v12b.py index 8d7baa0d52..cb2c265977 100644 --- a/tools/qwen38_graph/build_package_engine_v12b.py +++ b/tools/qwen38_graph/build_package_engine_v12b.py @@ -1415,7 +1415,7 @@ def clean_and_pack_package(source: Path, out: Path, copy: bool = False, embed_qu """Clean up unreferenced safetensors/dense.data/merged.data/generated.data, compact expert scales, split model.embed_tokens.weight into a standard CPU embedding.onnx + embedding.onnx.data model (INT8 per-row or BF16), unify target and MTP graphs so the same ONNX graph serves both prefill - and decode, and place all ONNX models (model.onnx, mtp.onnx, embedding.onnx) + and decode, and place all ONNX models (model.onnx, model_s2..s4.onnx, mtp.onnx, embedding.onnx) in the root package directory.""" import numpy as np # noqa: PLC0415 import onnx # noqa: PLC0415 @@ -1842,8 +1842,13 @@ def set_ext(tensor, location: str, offset: int, length: int) -> None: set_ext(t, "norm_and_rope.data", norm_rope_map[dedup_key], s_len) packed_tensors += 1 - # 5. Save cleaned unified ONNX models in the root directory and remove legacy subdirectories + # 5. Save cleaned unified ONNX models in the root directory, generate S=2..4 multi-draft QMoE graphs, and remove legacy subdirectories onnx.save_model(models["target"], str(out / "model.onnx"), save_as_external_data=False) + link_file(out / "model.onnx", out / "model_s2.onnx", copy) + from build_package_offload_v3 import build_multidraft_target_onnx # noqa: PLC0415 + + build_multidraft_target_onnx(out / "model_s2.onnx", out / "model_s3.onnx", verify_rows=3) + build_multidraft_target_onnx(out / "model_s2.onnx", out / "model_s4.onnx", verify_rows=4) onnx.save_model(models["head_r2"], str(out / "mtp.onnx"), save_as_external_data=False) for obsolete_dir in ("target", "head_r2", "dense", "head_kv"): if (out / obsolete_dir).exists(): @@ -2096,6 +2101,12 @@ def main() -> None: default=False, help="keep SharedExpertDown projections in BF16", ) + parser.add_argument( + "--expert-profile", + type=Path, + default=None, + help="STRP v1 expert-profile.bin copied into the package for the offload runtime profiles (clean packages)", + ) args = parser.parse_args() source, out = args.from_package.resolve(), args.out.resolve() @@ -2273,6 +2284,108 @@ def main() -> None: qwen["prefill"].pop("head_kv", None) qwen["prefill"].pop("constants", None) + # Copy the STRP v1 expert popularity profile named by the offload runtime profiles + if args.expert_profile is not None: + shutil.copy2(args.expert_profile, out / "expert-profile.bin") + else: + print("note: no --expert-profile given; copy an STRP v1 expert-profile.bin into the package yourself") + config["runtime_profiles"] = [ + { + "id": "consumer-gpu-24g", + "eligibility": { + "minimum_total_device_memory_bytes": 21474836480, + "maximum_total_device_memory_bytes": 26843545600, + "minimum_effective_host_memory_bytes": 68719476736, + "is_integrated": False, + }, + "overlay": { + "model": { + "decoder": { + "filename": "model_s4.onnx", + "session_options": { + "session.moe.offload.mode": "required", + "session.moe.offload.placement": "adaptive", + "session.moe.offload.cpu_execution": "auto", + "session.moe.offload.device_expert_bytes": "10351538176", + "session.moe.offload.popularity_file": "expert-profile.bin", + }, + } + } + }, + }, + { + "id": "consumer-gpu-24g-host-64g", + "eligibility": { + "minimum_total_device_memory_bytes": 21474836480, + "maximum_total_device_memory_bytes": 26843545600, + "maximum_effective_host_memory_bytes": 68719476735, + "is_integrated": False, + }, + "overlay": { + "model": { + "decoder": { + "filename": "model_s4.onnx", + "session_options": { + "session.moe.offload.mode": "required", + "session.moe.offload.placement": "adaptive", + "session.moe.offload.cpu_execution": "auto", + "session.moe.offload.device_expert_bytes": "10351538176", + "session.moe.offload.host_expert_bytes": "17179869184", + "session.moe.offload.popularity_file": "expert-profile.bin", + }, + } + } + }, + }, + { + "id": "consumer-gpu-30g-ram-48g-32k", + "eligibility": { + "minimum_total_device_memory_bytes": 26843545601, + "maximum_total_device_memory_bytes": 32212254720, + "maximum_effective_host_memory_bytes": 68719476735, + "is_integrated": False, + }, + "overlay": { + "model": { + "decoder": { + "filename": "model_s4.onnx", + "session_options": { + "session.moe.offload.mode": "required", + "session.moe.offload.placement": "adaptive", + "session.moe.offload.cpu_execution": "auto", + "session.moe.offload.device_expert_bytes": "18289511424", + "session.moe.offload.host_expert_bytes": "41339060224", + "session.moe.offload.popularity_file": "expert-profile.bin", + }, + } + } + }, + }, + { + "id": "consumer-gpu-32g", + "eligibility": { + "minimum_total_device_memory_bytes": 26843545601, + "maximum_total_device_memory_bytes": 36507222016, + "minimum_effective_host_memory_bytes": 68719476736, + "is_integrated": False, + }, + "overlay": { + "model": { + "decoder": { + "filename": "model_s4.onnx", + "session_options": { + "session.moe.offload.mode": "required", + "session.moe.offload.placement": "adaptive", + "session.moe.offload.cpu_execution": "auto", + "session.moe.offload.device_expert_bytes": "19229035520", + "session.moe.offload.popularity_file": "expert-profile.bin", + }, + } + } + }, + }, + ] + if transform_models or clean_pkg: if clean_pkg: models_map = { diff --git a/tools/qwen38_graph/build_package_v12d_int4_int8.py b/tools/qwen38_graph/build_package_v12d_int4_int8.py index d6475e4b96..1e46c24a9e 100644 --- a/tools/qwen38_graph/build_package_v12d_int4_int8.py +++ b/tools/qwen38_graph/build_package_v12d_int4_int8.py @@ -1070,7 +1070,7 @@ def unify_and_pack_flat_int4_int8_package( reuse_embed_from: Path, ) -> dict: """ - Unify target -> model.onnx and head_r2 -> mtp.onnx for arbitrary + Unify target -> model.onnx (and model_s2..s4.onnx) and head_r2 -> mtp.onnx for arbitrary sequence_length (serving both prefill and decode), wire INT8 CPU embedding.onnx + embedding.onnx.data, pack all remaining BF16 initializers (norms, RoPE, router, shared_expert) into norm_and_rope.data in the root directory, and write weights_pool.txt. @@ -1377,6 +1377,14 @@ def set_ext(tensor: TensorProto, location: str, offset: int, length: int) -> Non set_ext(t, "norm_and_rope.data", norm_rope_map[pkey], ln) onnx.save_model(models["target"], str(out_dir / "model.onnx"), save_as_external_data=False) + dst_s2 = out_dir / "model_s2.onnx" + if dst_s2.exists() or dst_s2.is_symlink(): + dst_s2.unlink() + os.link(out_dir / "model.onnx", dst_s2) + from build_package_offload_v3 import build_multidraft_target_onnx # noqa: PLC0415 + + build_multidraft_target_onnx(dst_s2, out_dir / "model_s3.onnx", verify_rows=3) + build_multidraft_target_onnx(dst_s2, out_dir / "model_s4.onnx", verify_rows=4) onnx.save_model(models["head_r2"], str(out_dir / "mtp.onnx"), save_as_external_data=False) return write_weights_pool( @@ -1701,6 +1709,10 @@ def main() -> int: "libs/libqwen_mtp_head.so", "libs/libqwen_fold.so", ] + if "runtime_profiles" in cfg and "consumer-gpu-24g" in cfg["runtime_profiles"]: + cfg["runtime_profiles"]["consumer-gpu-24g"]["decoder"]["session_options"]["provider_options"][0]["cuda"][ + "qmoe.offload_gpu_layers" + ] = "2" (out / "genai_config.json").write_text(json.dumps(cfg, indent=2) + "\n") manifest = { "variant": "qwen_flash_next_int4_int8_rc2", From e0db65bacbe5a31e12a4b91a72cf46f7e907d093 Mon Sep 17 00:00:00 2001 From: Tianlei Wu Date: Sat, 31 Oct 2026 01:54:33 +0000 Subject: [PATCH 3/6] Address PR #2821 review feedback on offload package builder Keep the standard com.microsoft::QMoE operator and configure expert offload through the session.moe.offload.* entries of the ORT QMoE offload contract (microsoft/onnxruntime#33622) instead of rewriting nodes to QMoEOffload with layer_index/resident_slots attributes. The legacy rewrite flags are rejected with an explanation, legacy QMoEOffload source graphs are refused, and the per-layer cap goes through node_overrides. Relocate every host-mapped MTP pool with an unaligned offset (not only the pool ids 607/608) into a layout-keyed aligned copy, and distribute the remainder of a uniform resident-slot total so the plan sums exactly. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- .../builder/_qwen38_qmoe_offload_contract.py | 338 +++ .../test_qwen38_offload_package_builder.py | 1886 +++++++++++++++++ tools/qwen38_graph/README.md | 86 +- .../qwen38_graph/build_package_offload_v3.py | 897 +++++--- .../build_package_v12d_int4_int8.py | 4 - 5 files changed, 2938 insertions(+), 273 deletions(-) create mode 100644 test/python/builder/_qwen38_qmoe_offload_contract.py create mode 100644 test/python/builder/test_qwen38_offload_package_builder.py diff --git a/test/python/builder/_qwen38_qmoe_offload_contract.py b/test/python/builder/_qwen38_qmoe_offload_contract.py new file mode 100644 index 0000000000..834dd1d497 --- /dev/null +++ b/test/python/builder/_qwen38_qmoe_offload_contract.py @@ -0,0 +1,338 @@ +# Copyright (c) Microsoft Corporation. All rights reserved. +# Licensed under the MIT License. + +"""A model of the ONNX Runtime QMoE offload contract, for the tests of the Qwen3.8 offload package builder. + +It is transcribed from microsoft/onnxruntime#33622 at commit 1d5ae2f3ad7bc9e5ce5579808ee74cc9f93284b9: the ``QMoE`` kernel +(``contrib_ops/cuda/moe/moe_quantization.cc``), the session option parser (``core/framework/moe_offload_plan.h``) and the +planner (``core/providers/cuda/qmoe_session_controller.h``). Released ONNX Runtime has no QMoE offload and ignores the +``session.moe.offload.*`` entries, and the unit tests have no CUDA build of that pull request. ``simulate_session_creation`` +therefore reproduces what decides whether session creation fails there: the operators and attributes of the graph, the +session entries, the per-node eligibility, the layer numbers and the node overrides. The byte budgets, which depend on the free +GPU memory, are not modeled. +""" + +from __future__ import annotations + +import re +import struct +from collections.abc import Mapping +from pathlib import Path + +import onnx +from onnx import helper + +OFFLOAD_PREFIX = "session.moe.offload." +INT32_MAX = 2**31 - 1 +UINT64_MAX = 2**64 - 1 +STRP_MAGIC = 0x50525453 + +# The attributes of the com.microsoft::QMoE schema; the runtime rejects any other attribute at graph load. +QMOE_ATTRIBUTES = frozenset( + { + "activation_type", + "k", + "normalize_routing_weights", + "use_sparse_mixer", + "expert_weight_bits", + "fc1_expert_weight_bits", + "fc2_expert_weight_bits", + "fc3_expert_weight_bits", + "swiglu_fusion", + "swiglu_limit", + "activation_alpha", + "activation_beta", + "accuracy_level", + "block_size", + "weights_prepacked", + "quant_type", + } +) +UNSUPPORTED_SLOTS = (8, 9, 10, 11, 12, 13, 14, 17, 18, 19, 20) +C_WHITESPACE = " \t\n\v\f\r" + + +class ContractError(Exception): + """Session creation fails in the runtime.""" + + +def extract_layer_index(name: str) -> int | None: + """``MoeOffloadSessionPolicy::TryExtractLayerIndexFromName``: the number after "layers.", "layer." or "blocks.".""" + for prefix in ("layers.", "layer.", "blocks."): + for match in re.finditer(re.escape(prefix) + "([0-9]+)", name): + if int(match.group(1)) > INT32_MAX: + return None + following = name[match.end() : match.end() + 1] + if not following or not (following.isascii() and following.isalnum()): + return int(match.group(1)) + return None + + +def _int32(text: str) -> int | None: + if re.fullmatch("-?[0-9]+", text) and -(2**31) <= int(text) <= INT32_MAX: + return int(text) + return None + + +def _uint64(text: str) -> int | None: + if re.fullmatch("[0-9]+", text) and int(text) <= UINT64_MAX: + return int(text) + return None + + +def parse_node_overrides(raw: str) -> list[dict]: + """``TryParseNodeOverrides``: "name:max_resident_slots=N,device_expert_bytes=B;..." -> one dict per entry.""" + trimmed = raw.strip(C_WHITESPACE) + if not trimmed or trimmed == "{}": + return [] + invalid = ContractError( + f"Invalid {OFFLOAD_PREFIX}node_overrides value '{raw}'; expected semicolon-separated " + "'node_name:max_resident_slots=,device_expert_bytes='." + ) + overrides = [] + for chunk in trimmed.split(";"): + entry = chunk.strip(C_WHITESPACE) + if not entry: + continue + colon = entry.find(":") + if colon <= 0 or colon + 1 >= len(entry): + raise invalid + override = {"name": entry[:colon].strip(C_WHITESPACE), "max_resident_slots": None, "device_expert_bytes": None} + if not override["name"]: + raise invalid + fields = entry[colon + 1 :].strip(C_WHITESPACE) + position = 0 + while position < len(fields): + comma = fields.find(",", position) + end = len(fields) if comma < 0 else comma + field = fields[position:end].strip(C_WHITESPACE) + position = len(fields) if comma < 0 else comma + 1 + equals = field.find("=") + if not field or equals <= 0 or equals + 1 >= len(field): + raise invalid + key, value = field[:equals].strip(C_WHITESPACE), field[equals + 1 :].strip(C_WHITESPACE) + if key == "max_resident_slots": + slots = _int32(value) + if slots is None or slots < 0: + raise invalid + override[key] = slots + elif key == "device_expert_bytes": + override[key] = _uint64(value) + if override[key] is None: + raise invalid + else: + raise invalid + if override["max_resident_slots"] is None and override["device_expert_bytes"] is None: + raise invalid + overrides.append(override) + return overrides + + +def _choice(key: str, value: str, choices: Mapping[str, str], expected: str) -> str: + if value.lower() not in choices: + raise ContractError(f"Invalid {key} value '{value}'; expected {expected}.") + return choices[value.lower()] + + +def _auto_or_uint64(key: str, value: str) -> int | None: + if value.lower() == "auto": + return None + parsed = _uint64(value) + if parsed is None: + raise ContractError(f"Invalid {key} value '{value}'; expected 'auto' or non-negative integer bytes.") + return parsed + + +def parse_session_entries(entries: Mapping[str, str], model_dir: Path | None = None) -> dict: + """``MoeOffloadSessionPolicy::ParseFromConfigOptions``: the policy, or ContractError for an entry it rejects.""" + policy = { + "mode": "disabled", + "placement": "adaptive", + "cpu_execution": "auto", + "popularity_file": None, + "node_overrides": [], + } + for key, value in entries.items(): + if key.startswith("moe.cpu_offload."): + raise ContractError( + f"Prototype configuration key '{key}' is discarded and not supported in QMoE Offload v4; " + "use 'session.moe.offload.*'." + ) + if not key.startswith(OFFLOAD_PREFIX): + continue + name = key[len(OFFLOAD_PREFIX) :] + if name == "mode": + choices = {"disabled": "disabled", "0": "disabled", "off": "disabled", "auto": "auto"} + choices.update({"required": "required", "1": "required", "on": "required"}) + policy["mode"] = _choice(key, value, choices, "'disabled', 'auto', or 'required'") + elif name == "cpu_execution": + choices = {"auto": "auto", "enabled": "enabled", "1": "enabled", "disabled": "disabled", "0": "disabled"} + policy["cpu_execution"] = _choice(key, value, choices, "'auto', 'enabled', or 'disabled'") + elif name == "placement": + policy["placement"] = _choice( + key, value, {"static": "static", "adaptive": "adaptive"}, "'static' or 'adaptive'" + ) + elif name in ("device_expert_bytes", "host_expert_bytes", "pinned_bytes", "reserve_device_bytes"): + policy[name] = _auto_or_uint64(key, value) + elif name in ("device_working_bytes", "host_working_bytes"): + if not _uint64(value): + raise ContractError(f"Invalid {key} value '{value}'; expected positive integer bytes.") + elif name == "allow_file_backing": + _choice(key, value, {"1": "", "true": "", "0": "", "false": ""}, "'0' or '1'") + elif name in ("cpu_threads", "async_max_tokens"): + if (_int32(value) or -1) < 0 and value != "0": + raise ContractError(f"Invalid {key} value '{value}'; expected non-negative integer.") + elif name == "prefill_wave_experts": + if (_int32(value) or 0) <= 0: + raise ContractError(f"Invalid {key} value '{value}'; expected positive integer.") + elif name == "skip_zero_cold_doorbell": + if value not in ("0", "1"): + raise ContractError(f"Invalid {key} value '{value}'; expected '0' or '1'.") + elif name in ("popularity_file", "source_root", "derived_cache_dir"): + path = Path(value) if value else None + if path is not None and not path.is_absolute() and model_dir is not None: + path = model_dir / path + policy[name] = path + elif name == "node_overrides": + policy[name] = parse_node_overrides(value) + elif name != "budget_group": + raise ContractError(f"Unknown session.moe.offload.* configuration key: {key}") + return policy + + +def _attribute(node: onnx.NodeProto, name: str, default): + for attribute in node.attribute: + if attribute.name == name: + value = helper.get_attribute_value(attribute) + return value.decode() if isinstance(value, bytes) else value + return default + + +def node_is_eligible(model: onnx.ModelProto, node: onnx.NodeProto) -> bool: + """The ``node_eligible`` expression of the QMoE kernel constructor.""" + quant_type = _attribute(node, "quant_type", "int") + if quant_type not in ("int", "fp4", "nvfp4", "fp8", "wfp4afp8"): + raise ContractError(f"quant_type must be 'int', 'fp4', 'nvfp4', 'fp8', or 'wfp4afp8', but got '{quant_type}'") + bits = _attribute(node, "expert_weight_bits", 4) + fc1_bits = _attribute(node, "fc1_expert_weight_bits", bits) + fc2_bits = _attribute(node, "fc2_expert_weight_bits", bits) + block_size = _attribute(node, "block_size", -1) + prepacked_mode = _attribute(node, "weights_prepacked", -1) + if prepacked_mode not in (-1, 0, 1): + raise ContractError(f"weights_prepacked must be -1 (default), 0, or 1, but got {prepacked_mode}") + if quant_type == "nvfp4": + if block_size not in (-1, 16): + raise ContractError(f"QMoE quant_type='nvfp4' requires block_size=16, but got {block_size}") + block_size = 16 + prepacked = prepacked_mode != 0 + + nvfp4_row_major = quant_type == "nvfp4" and prepacked_mode == 1 + int4_prepacked = quant_type == "int" and bits == fc1_bits == fc2_bits == 4 and prepacked and block_size == 32 + int3_raw = ( + quant_type == "int" and bits == fc1_bits == fc2_bits == 3 and not prepacked and block_size in (32, 64, 128) + ) + + initializers = {initializer.name: initializer for initializer in model.graph.initializer} + graph_inputs = {value.name for value in model.graph.input} + + def constant(slot: int): + name = node.input[slot] if slot < len(node.input) else "" + overridable = model.ir_version >= 4 and name in graph_inputs + return initializers.get(name) if name and not overridable else None + + has_unsupported_slot = any(slot < len(node.input) and node.input[slot] for slot in UNSUPPORTED_SLOTS) + weights = [constant(slot) for slot in (2, 3, 5, 6)] + has_const_weights = all(tensor is not None for tensor in weights) and len(weights[0].dims) == 3 == len( + weights[2].dims + ) + nvfp4_globals_ok = quant_type != "nvfp4" or (constant(15) is not None and constant(16) is not None) + interleaved_swiglu = _attribute(node, "activation_type", "relu") == "swiglu" and _attribute( + node, "swiglu_fusion", 0 + ) in (0, 1) + return bool( + (nvfp4_row_major or int4_prepacked or int3_raw) + and (bits == 4 or int3_raw) + and interleaved_swiglu + and _attribute(node, "use_sparse_mixer", 0) != 1 + and not has_unsupported_slot + and has_const_weights + and nvfp4_globals_ok + ) + + +def node_layer_index(node: onnx.NodeProto) -> int | None: + """The layer number from the node name, else from the name of the fc1 weights input; None when neither has one.""" + layer = extract_layer_index(node.name) + if layer is None and len(node.input) > 2 and node.input[2]: + layer = extract_layer_index(node.input[2]) + return layer + + +def layer_indices(nodes: list[onnx.NodeProto]) -> list[int]: + """The layer index of each registered node, in registration order, as the session controller assigns them.""" + assigned: list[int] = [] + for node in nodes: + layer = node_layer_index(node) + if layer is None: + layer = len(assigned) + while layer in assigned: + layer += 1 + assigned.append(layer) + return assigned + + +def validate_popularity_file(path: Path) -> None: + """The header and size checks the planner applies to an explicitly requested popularity file.""" + invalid = ContractError(f"MoE Offload failed to load or validate explicitly requested popularity_file: '{path}'.") + if not path.is_file(): + raise invalid + data = path.read_bytes() + if len(data) < 24: + raise invalid + magic, version, layers, experts, _slots, ranked = struct.unpack("<6I", data[:24]) + if magic != STRP_MAGIC or version != 1 or layers == 0 or experts == 0 or ranked > layers * experts: + raise invalid + if len(data) < 24 + 4 * ranked: + raise invalid + + +def simulate_session_creation( + model: onnx.ModelProto, entries: Mapping[str, str], model_dir: Path | None = None +) -> list[str]: + """The names of the QMoE nodes the runtime offloads, or ContractError where its session creation fails.""" + qmoe_nodes = [] + for node in model.graph.node: + if node.op_type == "QMoEOffload": + raise ContractError( + f"Type Error: Type of operator ({node.op_type}) is invalid: no such operator is registered" + ) + if node.domain == "com.microsoft" and node.op_type == "QMoE": + unknown = [attribute.name for attribute in node.attribute if attribute.name not in QMOE_ATTRIBUTES] + if unknown: + raise ContractError(f"Unrecognized attribute: {unknown[0]} for operator QMoE") + qmoe_nodes.append(node) + + policy = parse_session_entries(entries, model_dir) + eligible = [node_is_eligible(model, node) for node in qmoe_nodes] + registered = [] + if policy["mode"] != "disabled": + for node, ok in zip(qmoe_nodes, eligible, strict=True): + if ok: + registered.append(node) + elif policy["mode"] == "required": + raise ContractError(f"QMoE offload mode='required' rejected node '{node.name}'") + + layers = layer_indices(registered) + for index, layer in enumerate(layers): + if layer in layers[:index]: + raise ContractError(f"MoE Offload session admission failed: duplicate layer_index={layer}") + for override in policy["node_overrides"]: + matches = [ + override["name"] in (node.name, str(layer), f"layer.{layer}") + for node, layer in zip(registered, layers, strict=True) + ] + if not any(matches): + raise ContractError(f"MoE Offload node_overrides references unknown node '{override['name']}'.") + if registered and entries.get(OFFLOAD_PREFIX + "popularity_file"): + validate_popularity_file(policy["popularity_file"]) + return [node.name for node in registered] diff --git a/test/python/builder/test_qwen38_offload_package_builder.py b/test/python/builder/test_qwen38_offload_package_builder.py new file mode 100644 index 0000000000..b94cd79edc --- /dev/null +++ b/test/python/builder/test_qwen38_offload_package_builder.py @@ -0,0 +1,1886 @@ +# Copyright (c) Microsoft Corporation. All rights reserved. +# Licensed under the MIT License. + +"""Unit tests of the Qwen3.8 QMoE offload package builder (``tools/qwen38_graph/build_package_offload_v3.py``). + +The builder runs on tiny synthetic packages, so no GPU, checkpoint or custom-op library is needed. Released ONNX Runtime has no +QMoE offload, so the operator contract is checked against ``_qwen38_qmoe_offload_contract``, a model of the runtime of +microsoft/onnxruntime#33622, and by creating sessions in the installed onnxruntime where the emitted graph can be loaded. +""" + +from __future__ import annotations + +import argparse +import hashlib +import json +import os +import random +import shutil +import struct +import sys +import types +from pathlib import Path + +import numpy as np +import onnx +import pytest +from _qwen38_qmoe_offload_contract import ( + QMOE_ATTRIBUTES, + ContractError, + extract_layer_index, + node_is_eligible, + parse_node_overrides, + parse_session_entries, + simulate_session_creation, + validate_popularity_file, +) +from _qwen38_tools_test_utils import load_tools +from onnx import TensorProto, helper, numpy_helper + +LAYERS = 3 +EXPERTS = 8 +NODE_NAME = "model/layers.{0}/mlp/QMoE_node_{0}" +STRP_MAGIC = 0x50525453 +PREFIX = "session.moe.offload." + + +@pytest.fixture(scope="module") +def offload(): + with load_tools("package_paths", "build_package_offload_v3") as (_paths, module): + yield module + + +# --------------------------------------------------------------------------- +# Synthetic graphs and packages +# --------------------------------------------------------------------------- + +NVFP4_ATTRIBUTES = { + "activation_type": "swiglu", + "swiglu_fusion": 1, + "k": 2, + "expert_weight_bits": 4, + "block_size": 16, + "quant_type": "nvfp4", + "weights_prepacked": 1, + "normalize_routing_weights": 1, +} +SLOT_NAMES = { + 2: "fc1_experts_weights", + 3: "fc1_scales", + 5: "fc2_experts_weights", + 6: "fc2_scales", + 15: "fc1_global_scales", + 16: "fc2_global_scales", +} + + +def qmoe_node(layer=0, *, name=None, attributes=None, slots=None): + """A standard NVFP4 ``com.microsoft::QMoE`` node of ``layer``. + + ``slots`` replaces input names by slot number and ``attributes`` replaces attributes (``None`` removes one). + """ + names = {0: f"h{layer}", 1: f"router_probs_{layer}", 14: f"router_weights_{layer}"} + names.update({slot: f"model.layers.{layer}.mlp.{suffix}" for slot, suffix in SLOT_NAMES.items()}) + names.update(slots or {}) + inputs = [names.get(slot, "") for slot in range(max(names) + 1)] + merged = {key: value for key, value in {**NVFP4_ATTRIBUTES, **(attributes or {})}.items() if value is not None} + return helper.make_node( + "QMoE", + inputs, + [f"h{layer + 1}"], + name=NODE_NAME.format(layer) if name is None else name, + domain="com.microsoft", + **merged, + ) + + +def external_tensor(name, dims, dtype=TensorProto.UINT8, location="weights.data", offset=0): + """An initializer whose bytes are in an external file that the tests never create.""" + length = int(np.prod(dims)) * (4 if dtype == TensorProto.FLOAT else 1) + tensor = TensorProto(name=name, data_type=dtype, dims=dims, data_location=TensorProto.EXTERNAL) + for key, value in (("location", location), ("offset", str(offset)), ("length", str(length))): + tensor.external_data.add(key=key, value=value) + return tensor + + +def qmoe_initializers(layer=0, dims=None, skip=()): + """The external expert tensors of a node made by ``qmoe_node``; ``dims`` replaces shapes by slot number.""" + tensors = [] + for slot, suffix in SLOT_NAMES.items(): + if slot in skip: + continue + shape = {2: [EXPERTS, 4, 4], 5: [EXPERTS, 4, 4], 3: [EXPERTS, 4, 2], 6: [EXPERTS, 4, 2]}.get(slot, [EXPERTS]) + dtype = TensorProto.FLOAT if slot in (15, 16) else TensorProto.UINT8 + tensors.append( + external_tensor( + f"model.layers.{layer}.mlp.{suffix}", + (dims or {}).get(slot, shape), + dtype, + offset=4096 * (layer * 20 + slot), + ) + ) + return tensors + + +def make_model(nodes, initializers, ir_version=9, overridable=()): + """A graph of ``nodes``; every input that nothing defines is a graph input, and so is each ``overridable`` initializer.""" + known = {tensor.name for tensor in initializers} + graph_inputs = [] + for node in nodes: + for name in node.input: + if name and name not in known: + known.add(name) + graph_inputs.append(helper.make_tensor_value_info(name, TensorProto.FLOAT16, ["T", 4])) + known.update(node.output) + graph_inputs += [helper.make_tensor_value_info(name, TensorProto.UINT8, [EXPERTS, 4, 4]) for name in overridable] + graph = helper.make_graph( + nodes, + "decoder", + graph_inputs, + [helper.make_tensor_value_info(nodes[-1].output[0], TensorProto.FLOAT16, ["T", 4])], + initializer=initializers, + ) + model = helper.make_model( + graph, opset_imports=[helper.make_opsetid("", 17), helper.make_opsetid("com.microsoft", 1)] + ) + model.ir_version = ir_version + return model + + +def decoder_model(layers=LAYERS, **node_arguments): + nodes = [qmoe_node(layer, **node_arguments) for layer in range(layers)] + return make_model(nodes, [tensor for layer in range(layers) for tensor in qmoe_initializers(layer)]) + + +def one_node_model(*, attributes=None, slots=None, dims=None, skip=(), ir_version=9, overridable=()): + return make_model( + [qmoe_node(0, attributes=attributes, slots=slots)], + qmoe_initializers(0, dims, skip), + ir_version=ir_version, + overridable=overridable, + ) + + +def dense_model(): + node = helper.make_node("Identity", ["x"], ["y"]) + graph = helper.make_graph( + [node], + "dense", + [helper.make_tensor_value_info("x", TensorProto.FLOAT, [4, 4])], + [helper.make_tensor_value_info("y", TensorProto.FLOAT, [4, 4])], + initializer=[external_tensor("dense.layers.0.weight", [4, 4], location="source_weights/dense.data")], + ) + return helper.make_model(graph, opset_imports=[helper.make_opsetid("", 17)]) + + +def save_model(model, path): + path.parent.mkdir(parents=True, exist_ok=True) + onnx.save_model(model, str(path), save_as_external_data=False) + + +def genai_config(layered): + return { + "model": { + "eos_token_id": [7], + "decoder": { + "filename": "target/model.onnx" if layered else "model.onnx", + "session_options": { + "log_id": "onnxruntime-genai", + "provider_options": [{"cuda": {"enable_cuda_graph": "0"}}], + }, + }, + "qwen38_graph": { + "fold": {"mtp": {"filename": "mtp/model.onnx"}}, + "prefill": { + "chunk": 1024, + "dense": {"filename": "dense/model.onnx", "session_options": {"provider_options": [{"cuda": {}}]}}, + }, + }, + } + } + + +DEFAULT_POOL = """\ +qwen38_weights_pool 1 +pool 0 1 0 64 2 4 4 weights.data +pool 1 16 64 64 2 4 4 weights.data +session target model.embed_tokens.weight 0 +session target model.norm.weight 1 +""" + + +def make_package(tmp_path, *, model=None, layered=False, pool_text=DEFAULT_POOL, dense=True): + """A source package under ``tmp_path/pkg`` (flat, or with the decoder graph in ``target/``) and the two libraries.""" + root = tmp_path / "pkg" + decoder_dir = root / "target" if layered else root + save_model(decoder_model() if model is None else model, decoder_dir / "model.onnx") + (root / "genai_config.json").write_text(json.dumps(genai_config(layered))) + (root / "weights_pool.txt").write_text(pool_text) + if dense: + save_model(dense_model(), root / "dense" / "model.onnx") + libs = tmp_path / "libs" + libs.mkdir(exist_ok=True) + (libs / "libqwen38_plane.so").write_bytes(b"plane") + (libs / "libonnxruntime_providers_cuda.so").write_bytes(b"provider") + return types.SimpleNamespace( + root=root, + decoder_dir=decoder_dir, + plane=libs / "libqwen38_plane.so", + provider=libs / "libonnxruntime_providers_cuda.so", + ) + + +def write_profile(path, pairs, layers=LAYERS, experts=EXPERTS, magic=STRP_MAGIC, version=1, ranked=None): + """An STRP v1 expert-popularity file ranking the (layer, expert) ``pairs`` from the hottest down.""" + header = struct.pack( + "<6I", magic, version, layers, experts, layers * experts, len(pairs) if ranked is None else ranked + ) + path.write_bytes(header + b"".join(struct.pack("<2H", layer, expert) for layer, expert in pairs)) + return path + + +def exhaustive_pairs(layers=LAYERS, experts=EXPERTS, seed=0): + """Every (layer, expert) pair in a seeded random popularity order.""" + pairs = [(layer, expert) for layer in range(layers) for expert in range(experts)] + random.Random(seed).shuffle(pairs) + return pairs + + +def run_main(monkeypatch, module, *argv): + monkeypatch.setattr(sys, "argv", [f"{module.__name__}.py", *map(str, argv)]) + return module.main() + + +def run_builder(monkeypatch, offload, package, out, *extra): + return run_main( + monkeypatch, + offload, + "--from-package", + package.root, + "--out", + out, + "--plane-lib", + package.plane, + "--provider-lib", + package.provider, + *extra, + ) + + +@pytest.fixture +def multidraft_calls(offload, monkeypatch): + """Replaces the multi-draft graph rewrite (which needs the real folded graph) with a copy, and records its calls.""" + calls = [] + + def copy_graph(in_onnx, out_onnx, verify_rows=3): + calls.append((Path(in_onnx).name, Path(out_onnx).name, verify_rows)) + shutil.copyfile(in_onnx, out_onnx) + return {"verify_rows": verify_rows} + + monkeypatch.setattr(offload, "build_multidraft_target_onnx", copy_graph) + return calls + + +def snapshot(root): + """Every path below ``root`` with the SHA-256 of its content (```` for folders).""" + return { + path.relative_to(root).as_posix(): "" if path.is_dir() else hashlib.sha256(path.read_bytes()).hexdigest() + for path in sorted(root.rglob("*")) + } + + +def offload_entries(config): + return { + key: value for key, value in config["model"]["decoder"]["session_options"].items() if key.startswith(PREFIX) + } + + +def read_output(out, layered): + """The output config, the offload entries of its decoder, and the decoder graph with its directory.""" + config = json.loads((out / "genai_config.json").read_text()) + decoder_dir = out / "target" if layered else out + return ( + config, + offload_entries(config), + onnx.load(str(decoder_dir / "model.onnx"), load_external_data=False), + decoder_dir, + ) + + +def qmoe_nodes_of(model): + return [node for node in model.graph.node if node.op_type == "QMoE"] + + +# --------------------------------------------------------------------------- +# The operator contract: standard QMoE and session.moe.offload.* entries +# --------------------------------------------------------------------------- + + +def test_contract_model_rejects_what_the_reference_test_rejects(): + """The failures of test_discarded_prototype_features_rejected in onnxruntime#33622 (test_qmoe_offload_v4.py).""" + model = decoder_model(2) + with pytest.raises(ContractError, match="Unknown session.moe.offload"): + simulate_session_creation(model, {PREFIX + "unknown_knob": "1"}) + with pytest.raises(ContractError, match="discarded and not supported"): + simulate_session_creation(model, {"moe.cpu_offload.experts": "4"}) + with pytest.raises(ContractError, match="expected 'static' or 'adaptive'"): + simulate_session_creation(model, {PREFIX + "mode": "required", PREFIX + "placement": "basic"}) + legacy = decoder_model(2) + legacy.graph.node[0].op_type = "QMoEOffload" + with pytest.raises(ContractError, match="QMoEOffload"): + simulate_session_creation(legacy, {}) + prototype = decoder_model(2) + prototype.graph.node[0].attribute.append(helper.make_attribute("layer_index", 0)) + with pytest.raises(ContractError, match="Unrecognized attribute: layer_index"): + simulate_session_creation(prototype, {}) + + +def test_contract_model_follows_the_rest_of_the_reference_test(offload, tmp_path): + model = decoder_model(2) + offload.prepare_qmoe_for_offload(model) + required = {PREFIX + "mode": "required"} + assert simulate_session_creation(model, required, tmp_path) == [NODE_NAME.format(0), NODE_NAME.format(1)] + with pytest.raises(ContractError, match="popularity_file"): + simulate_session_creation(model, {**required, PREFIX + "popularity_file": "missing.bin"}, tmp_path) + with pytest.raises(ContractError, match="node_overrides"): + simulate_session_creation(model, {**required, PREFIX + "node_overrides": "no_colon"}, tmp_path) + override = {**required, PREFIX + "node_overrides": "model/layers.0/mlp/QMoE_node_0:max_resident_slots=6"} + assert simulate_session_creation(model, override, tmp_path) + unknown = {**required, PREFIX + "node_overrides": "no_such_node:max_resident_slots=6"} + with pytest.raises(ContractError, match="unknown node 'no_such_node'"): + simulate_session_creation(model, unknown, tmp_path) + model.graph.node[0].attribute[[a.name for a in model.graph.node[0].attribute].index("swiglu_fusion")].i = 2 + with pytest.raises(ContractError, match="rejected node"): + simulate_session_creation(model, required, tmp_path) + assert simulate_session_creation(model, {PREFIX + "mode": "auto"}, tmp_path) == [NODE_NAME.format(1)] + + +# (id, offloadable by the runtime, one_node_model arguments) +INT_ATTRIBUTES = {"quant_type": "int", "block_size": 32, "weights_prepacked": None} +NO_GLOBAL_SCALES = {"skip": (15, 16), "slots": {15: "", 16: ""}} +ELIGIBILITY = [ + ("nvfp4", True, {}), + ("nvfp4 default block_size", True, {"attributes": {"block_size": None}}), + ("nvfp4 swiglu_fusion 0", True, {"attributes": {"swiglu_fusion": 0}}), + ("nvfp4 swiglu_fusion 2", False, {"attributes": {"swiglu_fusion": 2}}), + ("nvfp4 weights_prepacked 0", False, {"attributes": {"weights_prepacked": 0}}), + ("nvfp4 default weights_prepacked", False, {"attributes": {"weights_prepacked": None}}), + ("nvfp4 relu", False, {"attributes": {"activation_type": "relu"}}), + ("nvfp4 default activation", False, {"attributes": {"activation_type": None}}), + ("nvfp4 sparse mixer", False, {"attributes": {"use_sparse_mixer": 1}}), + ("nvfp4 8-bit", False, {"attributes": {"expert_weight_bits": 8}}), + ("nvfp4 fc3 weights", False, {"slots": {8: "fc3_experts_weights"}}), + ("nvfp4 fc1 zero points", False, {"slots": {11: "fc1_zero_points"}}), + ("nvfp4 activation scales", False, {"slots": {17: "fc1_activation_scale"}}), + ("nvfp4 no global scales", False, NO_GLOBAL_SCALES), + ("nvfp4 one global scale", False, {"skip": (16,), "slots": {16: ""}}), + ("nvfp4 fc1 weights are an input", False, {"slots": {2: "fc1_weights_input"}}), + ("nvfp4 no fc2 scales", False, {"skip": (6,), "slots": {6: ""}}), + ("nvfp4 fc1 weights rank 2", False, {"dims": {2: [EXPERTS, 4]}}), + ("nvfp4 fc2 weights rank 4", False, {"dims": {5: [EXPERTS, 4, 4, 1]}}), + ("nvfp4 overridable fc1 weights", False, {"overridable": ["model.layers.0.mlp.fc1_experts_weights"]}), + ( + "nvfp4 overridable fc1 weights, IR 3", + True, + {"overridable": ["model.layers.0.mlp.fc1_experts_weights"], "ir_version": 3}, + ), + ("int4 block 32", True, {"attributes": INT_ATTRIBUTES, **NO_GLOBAL_SCALES}), + ("int4 block 32 prepacked 1", True, {"attributes": {**INT_ATTRIBUTES, "weights_prepacked": 1}, **NO_GLOBAL_SCALES}), + ("int4 block 32 with unused global scales", True, {"attributes": INT_ATTRIBUTES}), + ("int4 block 64", False, {"attributes": {**INT_ATTRIBUTES, "block_size": 64}, **NO_GLOBAL_SCALES}), + ("int4 default block_size", False, {"attributes": {**INT_ATTRIBUTES, "block_size": None}, **NO_GLOBAL_SCALES}), + ("int4 weights_prepacked 0", False, {"attributes": {**INT_ATTRIBUTES, "weights_prepacked": 0}, **NO_GLOBAL_SCALES}), + ("int8", False, {"attributes": {**INT_ATTRIBUTES, "expert_weight_bits": 8}, **NO_GLOBAL_SCALES}), + ( + "int3 raw block 64", + True, + { + "attributes": {**INT_ATTRIBUTES, "expert_weight_bits": 3, "block_size": 64, "weights_prepacked": 0}, + **NO_GLOBAL_SCALES, + }, + ), + ( + "int3 prepacked", + False, + {"attributes": {**INT_ATTRIBUTES, "expert_weight_bits": 3, "block_size": 64}, **NO_GLOBAL_SCALES}, + ), + ("mxfp4", False, {"attributes": {"quant_type": "fp4", "block_size": 32}, **NO_GLOBAL_SCALES}), +] + + +@pytest.mark.parametrize(("expected", "case"), [pytest.param(*case[1:], id=case[0]) for case in ELIGIBILITY]) +def test_qmoe_eligibility_matches_the_runtime(offload, expected, case): + model = one_node_model(**case) + report = offload.prepare_qmoe_for_offload(model) + assert (not report["rejected"]) == expected + assert node_is_eligible(model, model.graph.node[0]) == expected + + +def test_nvfp4_block_size_the_runtime_refuses_is_reported(offload): + model = one_node_model(attributes={"block_size": 32}) + with pytest.raises(ContractError, match="requires block_size=16"): + node_is_eligible(model, model.graph.node[0]) + reasons = offload.prepare_qmoe_for_offload(model)["rejected"][0] + assert any("block_size=32" in reason for reason in reasons) + + +def test_router_weights_make_a_node_ineligible_until_the_builder_clears_them(offload): + model = one_node_model() + assert not node_is_eligible(model, model.graph.node[0]) + report = offload.prepare_qmoe_for_offload(model) + assert report["router_weights_cleared"] == 1 + assert not report["rejected"] + assert node_is_eligible(model, model.graph.node[0]) + + +def test_rejection_reasons_name_the_defect(offload): + for case, fragment in ( + ({"slots": {8: "fc3_experts_weights", 17: "scale"}}, "inputs [8, 17]"), + ({"attributes": {"activation_type": "relu"}}, "activation_type='relu'"), + ({"dims": {2: [EXPERTS, 4]}}, "input 2 has rank 2"), + (NO_GLOBAL_SCALES, "global scales [15, 16]"), + ({"slots": {3: "fc1_scales_input"}}, "inputs [3]"), + ({"attributes": {"weights_prepacked": 0}}, "weights_prepacked=0"), + ): + rejected = offload.prepare_qmoe_for_offload(one_node_model(**case))["rejected"] + assert any(fragment in reason for reason in rejected[0]), (case, rejected) + + +# --------------------------------------------------------------------------- +# prepare_qmoe_for_offload, node overrides, session entries +# --------------------------------------------------------------------------- + + +def test_prepare_keeps_standard_qmoe_and_clears_only_router_weights(offload): + model = decoder_model() + before = onnx.ModelProto() + before.CopyFrom(model) + report = offload.prepare_qmoe_for_offload(model) + assert report == { + "nodes": [NODE_NAME.format(layer) for layer in range(LAYERS)], + "router_weights_cleared": LAYERS, + "rejected": {}, + } + for old, new in zip(before.graph.node, model.graph.node, strict=True): + assert (new.op_type, new.domain, new.name) == ("QMoE", "com.microsoft", old.name) + assert [a.SerializeToString() for a in new.attribute] == [a.SerializeToString() for a in old.attribute] + assert {a.name for a in new.attribute} <= QMOE_ATTRIBUTES + assert list(new.input) == [*old.input[:14], "", *old.input[15:]] + assert list(new.output) == list(old.output) + assert list(model.graph.initializer) == list(before.graph.initializer) + + +def test_prepare_handles_graphs_without_router_weights_or_qmoe(offload): + model = make_model([qmoe_node(0, slots={14: ""})], qmoe_initializers(0)) + assert offload.prepare_qmoe_for_offload(model)["router_weights_cleared"] == 0 + identity = helper.make_node("Identity", ["x"], ["y"]) + graph = helper.make_graph([identity], "g", [], [], initializer=[]) + report = offload.prepare_qmoe_for_offload(helper.make_model(graph)) + assert report == {"nodes": [], "router_weights_cleared": 0, "rejected": {}} + + +def test_prepare_reports_unnamed_and_rejected_nodes_by_position(offload): + nodes = [qmoe_node(0), qmoe_node(1, name="", attributes={"swiglu_fusion": 2})] + report = offload.prepare_qmoe_for_offload(make_model(nodes, qmoe_initializers(0) + qmoe_initializers(1))) + assert report["nodes"] == [NODE_NAME.format(0), ""] + assert list(report["rejected"]) == [1] + assert offload.node_label("", 1) == "#1 (unnamed)" + assert offload.node_label("a b", 1) == "'a b'" + + +@pytest.mark.parametrize( + "name", + [ + "model/layers.0/mlp/QMoE_node_0", + "model.layers.12.mlp", + "has space", + "equals=sign", + "comma,comma", + "{braces}", + "{}", + "ünïcode/日本", + "\u00a0nbsp", + "0", + "a" * 500, + ], +) +def test_node_overrides_round_trip_through_the_runtime_parser(offload, name): + value = offload.format_node_overrides({name: 5, "other": 0}) + parsed = parse_node_overrides(value) + assert [(o["name"], o["max_resident_slots"], o["device_expert_bytes"]) for o in parsed] == [ + (name, 5, None), + ("other", 0, None), + ] + + +def test_node_overrides_round_trip_for_random_names(offload): + rng = random.Random(7) + alphabet = "ab/._-= ,{}()[]éß日09\t" + for _ in range(300): + name = "".join(rng.choice(alphabet) for _ in range(rng.randint(1, 14))) + if name != name.strip(" \t\n\v\f\r"): + continue + cap = rng.choice([0, 1, 479, 2**31 - 1]) + assert [ + (o["name"], o["max_resident_slots"]) + for o in parse_node_overrides(offload.format_node_overrides({name: cap})) + ] == [(name, cap)] + + +@pytest.mark.parametrize("name", ["", " leading", "trailing ", "\ttab", "semi;colon", "co:lon"]) +def test_node_overrides_refuse_names_the_runtime_would_split(offload, name): + with pytest.raises(ValueError, match="cannot appear in session.moe.offload.node_overrides"): + offload.format_node_overrides({name: 4}) + + +@pytest.mark.parametrize("cap", [-1, 2**31, 2**40]) +def test_node_overrides_refuse_caps_outside_the_runtime_range(offload, cap): + with pytest.raises(ValueError, match="outside 0"): + offload.format_node_overrides({"node": cap}) + + +def test_node_overrides_of_no_caps_is_empty_and_valid(offload): + assert offload.format_node_overrides({}) == "" + assert parse_node_overrides("") == [] + + +@pytest.mark.parametrize("mode", ["required", "auto"]) +@pytest.mark.parametrize("placement", ["adaptive", "static"]) +def test_session_options_are_accepted_by_the_runtime_parser(offload, tmp_path, mode, placement): + minimal = offload.offload_session_options(mode, placement) + assert minimal == { + PREFIX + "mode": mode, + PREFIX + "placement": placement, + PREFIX + "cpu_execution": "auto", + } + entries = offload.offload_session_options( + mode, placement, "auto", "123456", "expert-profile.bin", "node:max_resident_slots=4" + ) + assert all(key.startswith(PREFIX) and isinstance(value, str) for key, value in entries.items()) + for options in (minimal, entries): + policy = parse_session_entries(options, tmp_path) + assert (policy["mode"], policy["placement"], policy["cpu_execution"]) == (mode, placement, "auto") + assert policy["device_expert_bytes"] is None + assert policy["host_expert_bytes"] == 123456 + assert policy["popularity_file"] == tmp_path / "expert-profile.bin" + assert [o["name"] for o in policy["node_overrides"]] == ["node"] + + +@pytest.mark.parametrize(("mode", "placement"), [("disabled", "adaptive"), ("Required", "static"), ("auto", "basic")]) +def test_session_options_refuse_what_the_builder_does_not_support(offload, mode, placement): + with pytest.raises(ValueError, match="offload (mode|placement)"): + offload.offload_session_options(mode, placement) + + +def test_budget_bytes_accepts_what_the_runtime_parses(offload): + for text, expected in ( + ("auto", "auto"), + ("AUTO", "auto"), + ("0", "0"), + ("007", "7"), + (str(2**64 - 1), str(2**64 - 1)), + ): + assert offload.budget_bytes(text) == expected + assert parse_session_entries({PREFIX + "device_expert_bytes": expected}) + for text in ("", " 5", "-1", "1.5", "0x10", "1e3", "١٢٣", str(2**64), "autos"): + with pytest.raises(argparse.ArgumentTypeError): + offload.budget_bytes(text) + with pytest.raises(ContractError): + parse_session_entries({PREFIX + "device_expert_bytes": str(2**64)}) + + +# --------------------------------------------------------------------------- +# Layer numbers +# --------------------------------------------------------------------------- + +LAYER_NAMES = [ + ("model.layers.5.mlp", 5), + ("model/layers.12/mlp/QMoE_node_12", 12), + ("layers.7", 7), + ("layers.7a", None), + ("layers.7a.layers.8", 8), + ("layer.3.x", 3), + ("blocks.9.", 9), + ("layers.", None), + ("layers.2147483647", 2147483647), + ("layers.2147483648", None), + ("layers.99999999999999999999", None), + ("xlayers.3", 3), + ("layers.1layers.2", 2), + ("layers.4é", 4), + ("blocks.3.layers.4", 4), + ("layer.3.layers.4", 4), + ("QMoE_node_3", None), + ("", None), +] + + +@pytest.mark.parametrize(("name", "expected"), LAYER_NAMES) +def test_layer_numbers_are_read_like_the_runtime_reads_them(offload, name, expected): + assert offload.parse_layer_index(name) == expected + assert extract_layer_index(name) == expected + + +def test_layer_numbers_agree_with_the_runtime_for_random_names(offload): + rng = random.Random(11) + tokens = ["layers.", "layer.", "blocks.", "7", "12", "a", ".", "/", "_", "é", "99999999999", "0"] + for _ in range(2000): + name = "".join(rng.choice(tokens) for _ in range(rng.randint(1, 9))) + assert offload.parse_layer_index(name) == extract_layer_index(name), name + + +def numbered_model(names, fc1_names=None): + nodes = [ + qmoe_node(layer, name=name, slots={2: (fc1_names or {}).get(layer, f"w{layer}")}) + for layer, name in enumerate(names) + ] + return make_model(nodes, []) + + +def test_layer_numbering_accepts_names_that_match_the_graph_order(offload): + offload.check_layer_numbering(numbered_model([f"model/layers.{i}/mlp" for i in range(3)])) + offload.check_layer_numbering(numbered_model(["a", "b", "c"])) + offload.check_layer_numbering(numbered_model(["", "", ""])) + offload.check_layer_numbering( + numbered_model(["a", "b"], {0: "model.layers.0.mlp.fc1", 1: "model.layers.1.mlp.fc1"}) + ) + + +@pytest.mark.parametrize( + ("names", "fc1_names", "not_offloaded", "message"), + [ + (["model/layers.1/mlp", "model/layers.0/mlp"], None, frozenset(), "is number 0 in the graph but .* layer 1"), + (["model/layers.48/mlp"], None, frozenset(), "numbers it layer 48"), + (["a", "b"], {1: "layers.0.fc1"}, frozenset(), "'b' is number 1 .* layer 0"), + (["a", "b"], None, frozenset({0}), "'b' is number 1 .* layer 0"), + (["", ""], None, frozenset({0}), "#1 \\(unnamed\\) is number 1 .* layer 0"), + ], +) +def test_layer_numbering_refuses_nodes_the_runtime_would_number_differently( + offload, names, fc1_names, not_offloaded, message +): + with pytest.raises(ValueError, match=message): + offload.check_layer_numbering(numbered_model(names, fc1_names), not_offloaded) + + +# --------------------------------------------------------------------------- +# Resident-slot plans: every requested total is met or refused +# --------------------------------------------------------------------------- + + +def ranking_by_layer(pairs): + """{layer: experts in popularity order}, the first mention of a pair counting.""" + ranked = {} + for layer, expert in pairs: + if expert not in ranked.setdefault(layer, []): + ranked[layer].append(expert) + return ranked + + +def check_plan(plan, pairs, total): + """The plan has ``total`` slots and every layer keeps a prefix of its own popularity ranking.""" + ranked = ranking_by_layer(pairs) + assert sum(len(experts) for experts in plan.values()) == total + for layer, experts in plan.items(): + assert experts == ranked.get(layer, [])[: len(experts)], layer + + +@pytest.mark.parametrize("uniform", [False, True]) +def test_exhaustive_profile_plans_meet_the_requested_total(offload, tmp_path, uniform): + pairs = exhaustive_pairs(48, 512) + profile = write_profile(tmp_path / "expert-profile.bin", pairs, 48, 512) + for total in (961, 47, 960, 959, 1, 0, 100, 1000, 48 * 512 - 1, 48 * 512): + plan = offload.load_expert_profile_plan(profile, 48, 512, total, uniform=uniform) + check_plan(plan, pairs, total) + if uniform: + sizes = [len(plan[layer]) for layer in range(48)] + assert max(sizes) - min(sizes) <= 1, total + + +def test_uniform_plan_of_961_slots_gives_one_layer_the_extra_slot(offload, tmp_path): + pairs = exhaustive_pairs(48, 512) + profile = write_profile(tmp_path / "expert-profile.bin", pairs, 48, 512) + plan = offload.load_expert_profile_plan(profile, 48, 512, 961, uniform=True) + sizes = sorted(len(experts) for experts in plan.values()) + assert sizes == [20] * 47 + [21] + + +def test_uniform_plan_of_47_slots_leaves_the_coldest_layer_empty(offload, tmp_path): + pairs = exhaustive_pairs(48, 512) + profile = write_profile(tmp_path / "expert-profile.bin", pairs, 48, 512) + plan = offload.load_expert_profile_plan(profile, 48, 512, 47, uniform=True) + assert sorted(len(experts) for experts in plan.values()) == [0] + [1] * 47 + empty = next(layer for layer, experts in plan.items() if not experts) + ranked = ranking_by_layer(pairs) + rank_of_best = {layer: pairs.index((layer, ranked[layer][0])) for layer in range(48)} + assert empty == max(rank_of_best, key=rank_of_best.get) + + +def test_uniform_remainder_goes_to_the_layers_whose_next_expert_ranks_best(offload, tmp_path): + hottest = [(0, 0), (1, 0), (2, 0), (2, 1), (0, 1), (1, 1)] + pairs = hottest + [(layer, expert) for layer in range(3) for expert in range(2, 4)] + profile = write_profile(tmp_path / "expert-profile.bin", pairs, 3, 4) + plan = offload.load_expert_profile_plan(profile, 3, 4, 5, uniform=True) + assert plan == {0: [0, 1], 1: [0], 2: [0, 1]} + + +@pytest.mark.parametrize("uniform", [False, True]) +def test_every_feasible_total_is_met_exactly(offload, tmp_path, uniform): + pairs = exhaustive_pairs(3, 8, seed=3) + profile = write_profile(tmp_path / "expert-profile.bin", pairs, 3, 8) + for total in range(3 * 8 + 1): + check_plan(offload.load_expert_profile_plan(profile, 3, 8, total, uniform=uniform), pairs, total) + + +def test_partial_profiles_are_filled_from_the_layers_that_rank_more(offload, tmp_path): + pairs = [(0, e) for e in range(6)] + [(1, e) for e in range(5)] + [(2, e) for e in range(2)] + random.Random(1).shuffle(pairs) + profile = write_profile(tmp_path / "expert-profile.bin", pairs, 3, 8) + for total in range(len(pairs) + 1): + check_plan(offload.load_expert_profile_plan(profile, 3, 8, total, min_slots_per_layer=4), pairs, total) + check_plan(offload.load_expert_profile_plan(profile, 3, 8, 8, uniform=True), pairs, 8) + + +@pytest.mark.parametrize( + ("ranks", "total", "message"), + [ + ([8, 8, 8], 25, "25 resident slots requested but .* ranks only 24"), + ([6, 5, 2], 14, "14 resident slots requested but .* ranks only 13"), + ([8, 8, 2], 9, r"3 slots per layer requested but layers \[2\] rank fewer"), + ([5, 2, 2], 8, "leave 2 extra slots, but only 1 layers rank more than 2"), + ], +) +def test_a_total_the_profile_cannot_fill_is_refused(offload, tmp_path, ranks, total, message): + pairs = [(layer, expert) for layer, count in enumerate(ranks) for expert in range(count)] + profile = write_profile(tmp_path / "expert-profile.bin", pairs, 3, 8) + with pytest.raises(ValueError, match=message): + offload.load_expert_profile_plan(profile, 3, 8, total, uniform=True) + + +def test_a_total_beyond_the_ranked_experts_is_refused_without_uniform_too(offload, tmp_path): + pairs = exhaustive_pairs(3, 8) + profile = write_profile(tmp_path / "expert-profile.bin", pairs, 3, 8) + with pytest.raises(ValueError, match="ranks only 24"): + offload.load_expert_profile_plan(profile, 3, 8, 25) + + +def test_duplicate_ranked_pairs_count_once(offload, tmp_path): + pairs = exhaustive_pairs(3, 8) + profile = write_profile(tmp_path / "expert-profile.bin", pairs + pairs[:10], 3, 8) + check_plan(offload.load_expert_profile_plan(profile, 3, 8, 24), pairs, 24) + with pytest.raises(ValueError, match="ranks only 24"): + offload.load_expert_profile_plan(profile, 3, 8, 25) + + +@pytest.mark.parametrize( + ("content", "message"), + [ + (b"STRP", "shorter than the 24-byte STRP header"), + (struct.pack("<6I", 0x12345678, 1, 3, 8, 24, 0), "Incompatible expert profile header"), + (struct.pack("<6I", STRP_MAGIC, 2, 3, 8, 24, 0), "Incompatible expert profile header"), + (struct.pack("<6I", STRP_MAGIC, 1, 4, 8, 24, 0), "Incompatible expert profile header"), + (struct.pack("<6I", STRP_MAGIC, 1, 3, 16, 24, 0), "Incompatible expert profile header"), + (struct.pack("<6I", STRP_MAGIC, 1, 3, 8, 24, 5) + struct.pack("<2H", 0, 0), "ranks 5 experts but holds only 1"), + (struct.pack("<6I", STRP_MAGIC, 1, 3, 8, 24, 1) + struct.pack("<2H", 3, 0), "outside 3 layers x 8 experts"), + (struct.pack("<6I", STRP_MAGIC, 1, 3, 8, 24, 1) + struct.pack("<2H", 0, 8), "outside 3 layers x 8 experts"), + ], +) +def test_a_damaged_profile_is_refused(offload, tmp_path, content, message): + profile = tmp_path / "expert-profile.bin" + profile.write_bytes(content) + with pytest.raises(ValueError, match=message): + offload.load_expert_profile_plan(profile, 3, 8, 0) + + +@pytest.mark.parametrize( + "arguments", + [{"total_slots": -1}, {"min_slots_per_layer": -1}, {"num_layers": 0}, {"num_experts": 0}], +) +def test_nonsensical_plan_arguments_are_refused(offload, tmp_path, arguments): + profile = write_profile(tmp_path / "expert-profile.bin", exhaustive_pairs(3, 8), 3, 8) + with pytest.raises(ValueError): + offload.load_expert_profile_plan(profile, **{"num_layers": 3, "num_experts": 8, "total_slots": 3, **arguments}) + + +def caps_args(**overrides): + values = { + "expert_profile": None, + "total_resident_slots": None, + "target_resident_slots": None, + "min_slots_per_layer": None, + "uniform_profile_slots": False, + } + return argparse.Namespace(**{**values, **overrides}) + + +def test_caps_without_a_profile_apply_to_every_layer(offload): + model = decoder_model() + assert offload.plan_resident_caps(caps_args(), model, LAYERS) == {} + assert offload.plan_resident_caps(caps_args(target_resident_slots=5), model, LAYERS) == dict.fromkeys( + range(LAYERS), 5 + ) + + +def test_a_profile_without_a_slot_count_sets_no_caps(offload, tmp_path): + profile = write_profile(tmp_path / "expert-profile.bin", exhaustive_pairs()) + assert offload.plan_resident_caps(caps_args(expert_profile=profile), decoder_model(), LAYERS) == {} + + +@pytest.mark.parametrize("uniform", [False, True]) +@pytest.mark.parametrize(("total", "target"), [(10, None), (None, 3), (23, None), (24, None)]) +def test_caps_from_a_profile_add_up_to_the_requested_total(offload, tmp_path, uniform, total, target): + pairs = exhaustive_pairs() + profile = write_profile(tmp_path / "expert-profile.bin", pairs) + args = caps_args( + expert_profile=profile, total_resident_slots=total, target_resident_slots=target, uniform_profile_slots=uniform + ) + caps = offload.plan_resident_caps(args, decoder_model(), LAYERS) + assert sorted(caps) == list(range(LAYERS)) + assert sum(caps.values()) == (total if total is not None else target * LAYERS) + + +def test_caps_from_a_profile_refuse_what_it_cannot_provide(offload, tmp_path): + profile = write_profile(tmp_path / "expert-profile.bin", exhaustive_pairs()) + args = caps_args(expert_profile=profile, total_resident_slots=25) + with pytest.raises(ValueError, match="--expert-profile: 25 resident slots requested"): + offload.plan_resident_caps(args, decoder_model(), LAYERS) + other_layers = write_profile(tmp_path / "other.bin", exhaustive_pairs(4, EXPERTS), layers=4) + with pytest.raises(ValueError, match="--expert-profile: Incompatible expert profile header"): + offload.plan_resident_caps( + caps_args(expert_profile=other_layers, total_resident_slots=4), decoder_model(), LAYERS + ) + + +def test_caps_from_a_profile_need_one_expert_count_and_matching_layer_numbers(offload, tmp_path): + profile = write_profile(tmp_path / "expert-profile.bin", exhaustive_pairs()) + args = caps_args(expert_profile=profile, total_resident_slots=4) + mixed = make_model( + [qmoe_node(0), qmoe_node(1)], qmoe_initializers(0) + qmoe_initializers(1, dims={2: [4, 4, 4], 5: [4, 4, 4]}) + ) + with pytest.raises(ValueError, match=r"one expert count, the graph has \[4, 8\]"): + offload.plan_resident_caps(args, mixed, 2) + shifted = make_model([qmoe_node(1), qmoe_node(0)], qmoe_initializers(0) + qmoe_initializers(1)) + with pytest.raises(ValueError, match="is number 0 in the graph but the runtime numbers it layer 1"): + offload.plan_resident_caps(args, shifted, 2) + + +@pytest.mark.parametrize( + "argv", + [ + ["--total-resident-slots", "5"], + ["--min-slots-per-layer", "2"], + ["--uniform-profile-slots"], + ["--target-resident-slots", "-1"], + ["--expert-profile", "profile.bin", "--total-resident-slots", "-3"], + ["--expert-profile", "profile.bin", "--uniform-profile-slots"], + ["--expert-profile", "profile.bin", "--min-slots-per-layer", "2"], + [ + "--expert-profile", + "profile.bin", + "--total-resident-slots", + "9", + "--uniform-profile-slots", + "--min-slots-per-layer", + "2", + ], + ], +) +def test_contradicting_slot_arguments_are_refused(offload, argv): + parser = offload.build_parser() + with pytest.raises(SystemExit) as excinfo: + offload.check_slot_arguments(parser, parser.parse_args(argv)) + assert excinfo.value.code == 2 + + +@pytest.mark.parametrize( + "argv", + [ + [], + ["--target-resident-slots", "4"], + ["--target-resident-slots", "0"], + ["--expert-profile", "profile.bin"], + ["--expert-profile", "profile.bin", "--total-resident-slots", "9"], + ["--expert-profile", "profile.bin", "--target-resident-slots", "3", "--uniform-profile-slots"], + ["--expert-profile", "profile.bin", "--total-resident-slots", "9", "--min-slots-per-layer", "2"], + ], +) +def test_consistent_slot_arguments_are_accepted(offload, argv): + parser = offload.build_parser() + offload.check_slot_arguments(parser, parser.parse_args(argv)) + + +# --------------------------------------------------------------------------- +# weights_pool.txt: placements, and the MTP expert pools that need an aligned copy +# --------------------------------------------------------------------------- + +MTP_FILE = "target/head_data/source_weights/mtp_moe.bin" +ALIGNED = "target/head_data/source_weights/mtp_moe_aligned.bin" +requires_copy_file_range = pytest.mark.skipif(not hasattr(os, "copy_file_range"), reason="needs os.copy_file_range") + + +def pool_text(records, bindings): + return "\n".join(["qwen38_weights_pool 1", *records, *bindings]) + "\n" + + +def read_pool(path): + """The pool records ({id: tokens}) and the session bindings of a weights_pool.txt; pool ids must be unique.""" + lines = [line.split() for line in path.read_text().splitlines()] + assert lines[0] == ["qwen38_weights_pool", "1"] + records = [parts for parts in lines[1:] if parts[0] in ("pool", "host_pool")] + assert len({parts[1] for parts in records}) == len(records), "duplicate pool ids" + bindings = [(parts[1], parts[2], int(parts[3])) for parts in lines[1:] if parts[0] == "session"] + assert len(records) + len(bindings) == len(lines) - 1 + return {int(parts[1]): parts for parts in records}, bindings + + +def mtp_pool_package(tmp_path, ranges, keyword="pool", records=(), bindings=()): + """A layered package whose MTP expert pools are ``ranges`` = [(pool id, offset, length)] of one source file. + + ``records`` and ``bindings`` add other pool and session lines. Returns the package and the bytes of the source file. + """ + data = random.Random(3).randbytes(max(offset + length for _pid, offset, length in ranges) + 1000) + mtp_records = [f"{keyword} {pid} 1 {offset} {length} 1 {length} {MTP_FILE}" for pid, offset, length in ranges] + mtp_bindings = [f"session mtp model.mtp.layers.0.mlp.experts.w{pid} {pid}" for pid, _offset, _length in ranges] + package = make_package( + tmp_path, layered=True, pool_text=pool_text(mtp_records + list(records), mtp_bindings + list(bindings)) + ) + (package.root / MTP_FILE).parent.mkdir(parents=True, exist_ok=True) + (package.root / MTP_FILE).write_bytes(data) + return package, data + + +def build_pool(offload, package, tmp_path, **placements): + out = tmp_path / "out" + out.mkdir(exist_ok=True) + return out, offload.build_offload_weights_pool(package.root, out, **placements) + + +def aligned_caches(tmp_path): + return sorted(tmp_path.glob("pkg_mtp_moe_aligned-*")) + + +def check_bound(records, bindings): + assert {pid for _sess, _name, pid in bindings} <= set(records) + + +def check_aligned_copy(out, data, ranges, records, expected_offsets, placement="host_mapped"): + """Each unaligned range is in the aligned file at its planned offset, byte for byte, with only zeros around.""" + aligned = np.frombuffer((out / ALIGNED).read_bytes(), dtype=np.uint8) + gaps = np.ones(len(aligned), dtype=bool) + end = 0 + for pid, offset, length in ranges: + record = records[pid] + if offset % 16 == 0: + assert record[2:4] == ["1", str(offset)] and record[-2:] == [MTP_FILE, placement] + continue + destination = expected_offsets[pid] + assert destination % 4096 == 0 + assert record == ["pool", str(pid), "1", str(destination), str(length), "1", str(length), ALIGNED, placement] + assert aligned[destination : destination + length].tobytes() == data[offset : offset + length] + gaps[destination : destination + length] = False + end = max(end, destination + length) + assert len(aligned) == end + assert not aligned[gaps].any() + + +# (id, [(pool id, offset in the source file, length)], planned offsets of the unaligned pools in the aligned file) +MTP_LAYOUTS = [ + ("ordinary ids 1 and 2", [(1, 4097, 100), (2, 9001, 5000)], {1: 0, 2: 4096}), + ("ids 607 and 608", [(607, 12296, 3000), (608, 20488, 1500)], {607: 0, 608: 4096}), + ("ids out of file order", [(9, 20008, 700), (4, 8, 5000), (7, 12008, 300)], {4: 0, 7: 8192, 9: 12288}), + ("sparse ids", [(40, 5001, 10), (3, 6003, 20), (12000, 7007, 4096)], {3: 0, 40: 4096, 12000: 8192}), + ("one pool", [(3, 17, 10)], {3: 0}), + ("aligned pools stay", [(5, 4096, 200), (6, 4105, 300), (7, 8200, 4000), (8, 16384, 64)], {6: 0, 7: 4096}), +] + + +@requires_copy_file_range +@pytest.mark.parametrize("keyword", ["pool", "host_pool"]) +@pytest.mark.parametrize(("ranges", "offsets"), [pytest.param(*case[1:], id=case[0]) for case in MTP_LAYOUTS]) +def test_every_unaligned_mtp_pool_is_moved_to_the_aligned_file(offload, tmp_path, ranges, offsets, keyword): + package, data = mtp_pool_package(tmp_path, ranges, keyword) + out, stats = build_pool(offload, package, tmp_path, mtp_moe_placement="host_mapped") + records, bindings = read_pool(out / "weights_pool.txt") + for pid, _offset, _length in ranges: + assert pid in records, f"pool {pid} was dropped" + check_bound(records, bindings) + check_aligned_copy(out, data, ranges, records, offsets) + (cache,) = aligned_caches(tmp_path) + assert os.path.samefile(cache, out / ALIGNED) + assert not list(tmp_path.rglob("*.tmp")) + assert stats["mtp_moe_placement"] == "host_mapped" + assert stats["total_pool_entries"] == len(records) + + +@requires_copy_file_range +def test_mtp_pools_stay_where_they_are_when_they_are_not_host_mapped(offload, tmp_path): + ranges = [(1, 4097, 100), (2, 9001, 5000)] + package, _data = mtp_pool_package(tmp_path, ranges) + out, _stats = build_pool(offload, package, tmp_path, mtp_moe_placement="device") + records, bindings = read_pool(out / "weights_pool.txt") + check_bound(records, bindings) + for pid, offset, length in ranges: + assert records[pid] == ["pool", str(pid), "1", str(offset), str(length), "1", str(length), MTP_FILE] + assert not (out / ALIGNED).exists() + assert not aligned_caches(tmp_path) + + +def test_the_aligned_file_name_is_part_of_the_package_layout(offload): + assert offload.ALIGNED_MTP_RELATIVE == ALIGNED + + +@requires_copy_file_range +def test_a_changed_layout_or_source_gets_its_own_aligned_file(offload, tmp_path, monkeypatch): + ranges = [(1, 4097, 100), (2, 9001, 5000)] + package, data = mtp_pool_package(tmp_path, ranges) + out, _ = build_pool(offload, package, tmp_path, mtp_moe_placement="host_mapped") + (first,) = aligned_caches(tmp_path) + + def must_not_copy(*_args, **_kwargs): + raise AssertionError("an unchanged layout must reuse its aligned file") + + with monkeypatch.context() as patch: + patch.setattr(os, "copy_file_range", must_not_copy) + build_pool(offload, package, tmp_path, mtp_moe_placement="host_mapped") + assert aligned_caches(tmp_path) == [first] + + # The same pool ids and sizes read from other offsets of the file. + shifted = [(1, 4113, 100), (2, 9017, 5000)] + shutil.rmtree(package.root) + package, data = mtp_pool_package(tmp_path, shifted) + out, _ = build_pool(offload, package, tmp_path, mtp_moe_placement="host_mapped") + records, _bindings = read_pool(out / "weights_pool.txt") + check_aligned_copy(out, data, shifted, records, {1: 0, 2: 4096}) + assert len(aligned_caches(tmp_path)) == 2 + assert os.path.samefile(next(path for path in aligned_caches(tmp_path) if path != first), out / ALIGNED) + + # Edited source bytes of the same size and layout. + source = package.root / MTP_FILE + stat = source.stat() + source.write_bytes(bytes(reversed(data))) + os.utime(source, ns=(stat.st_atime_ns, stat.st_mtime_ns + 10**9)) + out, _ = build_pool(offload, package, tmp_path, mtp_moe_placement="host_mapped") + records, _bindings = read_pool(out / "weights_pool.txt") + check_aligned_copy(out, bytes(reversed(data)), shifted, records, {1: 0, 2: 4096}) + assert len(aligned_caches(tmp_path)) == 3 + + +@requires_copy_file_range +def test_a_cached_aligned_file_of_the_wrong_size_is_rebuilt(offload, tmp_path): + ranges = [(1, 4097, 100), (2, 9001, 5000)] + package, data = mtp_pool_package(tmp_path, ranges) + build_pool(offload, package, tmp_path, mtp_moe_placement="host_mapped") + (cache,) = aligned_caches(tmp_path) + cache.write_bytes(b"truncated") + out, _ = build_pool(offload, package, tmp_path, mtp_moe_placement="host_mapped") + records, _bindings = read_pool(out / "weights_pool.txt") + check_aligned_copy(out, data, ranges, records, {1: 0, 2: 4096}) + assert aligned_caches(tmp_path) == [cache] + + +@requires_copy_file_range +def test_a_failed_copy_leaves_no_partial_aligned_file(offload, tmp_path, monkeypatch): + package, _data = mtp_pool_package(tmp_path, [(1, 4097, 100), (2, 9001, 5000)]) + for failure in (OSError("disk full"), KeyboardInterrupt()): + + def fail(*_args, failure=failure, **_kwargs): + raise failure + + with monkeypatch.context() as patch: + patch.setattr(os, "copy_file_range", fail) + with pytest.raises(type(failure)): + build_pool(offload, package, tmp_path, mtp_moe_placement="host_mapped") + assert not aligned_caches(tmp_path) + assert not (tmp_path / "out" / "weights_pool.txt").exists() + + +@requires_copy_file_range +def test_a_copy_that_stops_early_is_an_error_and_leaves_nothing(offload, tmp_path, monkeypatch): + package, _data = mtp_pool_package(tmp_path, [(1, 4097, 100), (2, 9001, 5000)]) + monkeypatch.setattr(os, "copy_file_range", lambda *_args, **_kwargs: 0) + with pytest.raises(RuntimeError, match="stopped early"): + build_pool(offload, package, tmp_path, mtp_moe_placement="host_mapped") + assert not aligned_caches(tmp_path) + + +@requires_copy_file_range +def test_partial_copies_are_continued_until_the_range_is_complete(offload, tmp_path, monkeypatch): + ranges = [(1, 4097, 100), (2, 9001, 5000)] + package, data = mtp_pool_package(tmp_path, ranges) + real = os.copy_file_range + + def seven_bytes(src, dst, count, offset_src=None, offset_dst=None): + return real(src, dst, min(count, 7), offset_src=offset_src, offset_dst=offset_dst) + + monkeypatch.setattr(os, "copy_file_range", seven_bytes) + out, _ = build_pool(offload, package, tmp_path, mtp_moe_placement="host_mapped") + records, _bindings = read_pool(out / "weights_pool.txt") + check_aligned_copy(out, data, ranges, records, {1: 0, 2: 4096}) + + +@requires_copy_file_range +@pytest.mark.parametrize( + ("ranges", "message"), + [ + ([(1, 4097, 100), (2, 10**6 + 1, 5000)], r"pool 2: bytes 1000001\.\.1005001 lie outside .*mtp_moe\.bin"), + ([(1, 4097, 10**9)], r"pool 1: bytes 4097\.\.1000004097 lie outside"), + ], +) +def test_a_pool_outside_its_file_is_refused_by_id_before_anything_is_written(offload, tmp_path, ranges, message): + package, _data = mtp_pool_package(tmp_path, [(1, 4097, 100), (2, 9001, 5000)]) + (package.root / "weights_pool.txt").write_text( + pool_text( + [f"pool {pid} 1 {offset} {length} 1 {length} {MTP_FILE}" for pid, offset, length in ranges], + [f"session mtp model.mtp.layers.0.mlp.experts.w{pid} {pid}" for pid, _offset, _length in ranges], + ) + ) + with pytest.raises(ValueError, match=message): + build_pool(offload, package, tmp_path, mtp_moe_placement="host_mapped") + assert not aligned_caches(tmp_path) + assert not list(tmp_path.rglob("*.tmp")) + assert not (tmp_path / "out" / "weights_pool.txt").exists() + + +@requires_copy_file_range +def test_a_missing_source_file_is_refused_by_pool_id(offload, tmp_path): + package, _data = mtp_pool_package(tmp_path, [(1, 4097, 100)]) + (package.root / MTP_FILE).unlink() + with pytest.raises(ValueError, match=r"pool 1: source file .*mtp_moe\.bin does not exist"): + build_pool(offload, package, tmp_path, mtp_moe_placement="host_mapped") + assert not aligned_caches(tmp_path) + + +@pytest.mark.parametrize( + ("records", "bindings", "message"), + [ + ( + ["pool 1 1 0 64 1 64 w.bin"], + ["session mtp model.mtp.layers.0.mlp.experts.w 7"], + "1 session bindings name a pool", + ), + ( + ["pool 1 1 0 64 1 64 w.bin", "host_pool 1 1 64 64 1 64 w.bin"], + [], + r"pool ids \[1\] are defined more than once", + ), + ], +) +def test_a_source_pool_file_with_dangling_or_duplicate_pools_is_refused(offload, tmp_path, records, bindings, message): + package = make_package(tmp_path, pool_text=pool_text(records, bindings)) + with pytest.raises(ValueError, match=f"the source weights_pool.txt: .*{message}"): + build_pool(offload, package, tmp_path) + assert not (tmp_path / "out" / "weights_pool.txt").exists() + + +def test_pool_binding_checks_report_what_is_wrong(offload): + offload.check_pool_bindings([1, 2], [("a", "x", 1), ("a", "y", 2), ("b", "x", 2)], "where") + with pytest.raises(ValueError, match=r"where: pool ids \[2, 9\] are defined more than once"): + offload.check_pool_bindings([9, 2, 2, 9, 1], [], "where") + with pytest.raises(ValueError, match=r"where: 2 session bindings .* e\.g\. session a x 5, session b y 6"): + offload.check_pool_bindings([1], [("a", "x", 5), ("b", "y", 6), ("c", "z", 1)], "where") + + +def test_aligned_copy_plan_packs_ranges_on_page_boundaries(offload): + # The two MTP expert banks of the real package. + assert offload.plan_aligned_copy([(607, 3355443200), (608, 1677721600)]) == ( + {607: 0, 608: 3355443200}, + 5033164800, + ) + assert offload.plan_aligned_copy([(1, 10), (2, 5000), (3, 1)]) == ({1: 0, 2: 4096, 3: 12288}, 12289) + assert offload.plan_aligned_copy([(1, 4096), (2, 0), (3, 4097)]) == ({1: 0, 2: 4096, 3: 4096}, 8193) + assert offload.plan_aligned_copy([]) == ({}, 0) + + +@requires_copy_file_range +def test_copy_file_ranges_writes_each_range_at_its_offset(offload, tmp_path): + source = tmp_path / "source.bin" + data = random.Random(1).randbytes(300) + source.write_bytes(data) + destination = tmp_path / "packed.bin" + offload.copy_file_ranges(destination, [(1, source, 250, 0, 50), (2, source, 3, 128, 100)]) + packed = destination.read_bytes() + assert len(packed) == 228 + assert packed[:50] == data[250:300] and packed[128:228] == data[3:103] + assert not any(packed[50:128]) + offload.copy_file_ranges(destination, []) + assert destination.read_bytes() == b"" + + +@requires_copy_file_range +def test_copy_file_ranges_checks_every_range_before_creating_the_file(offload, tmp_path): + source = tmp_path / "source.bin" + source.write_bytes(bytes(100)) + destination = tmp_path / "packed.bin" + for ranges, message in ( + ([(1, source, 0, 0, 10), (5, source, 95, 10, 10)], "pool 5: bytes 95..105 lie outside"), + ([(6, source, -1, 0, 10)], "pool 6: bytes -1..9 lie outside"), + ([(7, source, 0, 0, -1)], "pool 7: bytes 0..-1 lie outside"), + ([(8, tmp_path / "absent.bin", 0, 0, 1)], "pool 8: source file .* does not exist"), + ): + with pytest.raises(ValueError, match=message): + offload.copy_file_ranges(destination, ranges) + assert not destination.exists() + + +def test_other_pools_get_the_requested_placements(offload, tmp_path): + records = [ + "pool 0 1 0 64 2 4 4 weights.data", + "pool 1 1 64 64 2 4 4 weights.data", + "host_pool 2 1 128 64 2 4 4 weights.data", + "pool 3 1 192 64 2 4 4 weights.data", + "pool 4 1 256 64 2 4 4 weights.data device", + "pool 5 1 320 64 2 4 4 weights.data cpu", + "pool 6 1 384 64 2 4 4 weights.data host_mapped", + ] + bindings = [ + "session target model.embed_tokens.weight 0", + "session target model.layers.0.mlp.fc1_experts_weights 1", + "session target model.norm.weight 2", + "session target model.a 3", + "session target model.b 4", + "session target model.c 5", + "session target model.d 6", + ] + package = make_package(tmp_path, pool_text=pool_text(records, bindings), dense=False) + out, _stats = build_pool(offload, package, tmp_path) + placements = {pid: parts[7 + int(parts[5]) :] for pid, parts in read_pool(out / "weights_pool.txt")[0].items()} + assert [placements[pid] for pid in range(7)] == [ + ["host_mapped"], + ["cpu"], + ["host_mapped"], + [], + [], + ["cpu"], + ["host_mapped"], + ] + out, _stats = build_pool(offload, package, tmp_path, embed_placement="device") + assert read_pool(out / "weights_pool.txt")[0][0][7 + 2 :] == [] + + +@pytest.mark.parametrize("layered", [False, True]) +def test_unbound_expert_tensors_get_new_cpu_pools_with_fresh_ids(offload, tmp_path, layered): + pool = pool_text( + ["pool 5 1 0 64 2 4 4 weights.data", "pool 40 1 64 64 2 4 4 weights.data"], + ["session target model.embed_tokens.weight 5", "session target model.norm.weight 40"], + ) + package = make_package(tmp_path, pool_text=pool, layered=layered, dense=False) + out, stats = build_pool(offload, package, tmp_path) + records, bindings = read_pool(out / "weights_pool.txt") + check_bound(records, bindings) + fresh = sorted(set(records) - {5, 40}) + assert fresh == list(range(41, 41 + 6 * LAYERS)) + location = "target/weights.data" if layered else "weights.data" + assert {tuple(records[pid][-2:]) for pid in fresh} == {(location, "cpu")} + assert {name for sess, name, pid in bindings if pid in fresh} == { + init.name for layer in range(LAYERS) for init in qmoe_initializers(layer) + } + assert stats["cpu_qmoe_entries"] == 6 * LAYERS + assert stats["total_pool_entries"] == len(records) + + +def test_dense_weights_are_pooled_only_when_asked(offload, tmp_path): + package = make_package(tmp_path) + out, stats = build_pool(offload, package, tmp_path, dense_placement="device") + assert stats["added_dense_entries"] == 0 + assert not [binding for binding in read_pool(out / "weights_pool.txt")[1] if binding[0] == "dense"] + for placement in ("host_mapped", "cpu"): + out, stats = build_pool(offload, package, tmp_path, dense_placement=placement) + records, bindings = read_pool(out / "weights_pool.txt") + (dense,) = [binding for binding in bindings if binding[0] == "dense"] + assert stats["added_dense_entries"] == 1 + assert records[dense[2]][2:] == ["2", "0", "16", "2", "4", "4", "target/source_weights/dense.data", placement] + + +# --------------------------------------------------------------------------- +# main(): the package a user gets +# --------------------------------------------------------------------------- + +NAMES = [NODE_NAME.format(layer) for layer in range(LAYERS)] +REMOVED_FLAGS = [ + "--migrate-shared-onnx", + "--out-onnx", + "--strip-phase", + "--strip-stats", + "--target-spare-slots", + "--dense-offload-qmoe", + "--dense-stream-min-rows", + "--dense-stream-wave", +] + + +def caps_of(entries): + """{node name: max_resident_slots} of the node_overrides entry.""" + overrides = parse_node_overrides(entries.get(PREFIX + "node_overrides", "")) + return {override["name"]: override["max_resident_slots"] for override in overrides} + + +def build(monkeypatch, offload, package, tmp_path, *extra, layered=False): + out = tmp_path / "out" + run_builder(monkeypatch, offload, package, out, *extra) + return (out, *read_output(out, layered)) + + +@pytest.mark.parametrize("layered", [False, True]) +@pytest.mark.parametrize("mode", ["required", "auto"]) +@pytest.mark.parametrize("placement", ["adaptive", "static"]) +def test_the_built_decoder_is_standard_qmoe_that_the_runtime_offloads( + offload, tmp_path, monkeypatch, multidraft_calls, layered, mode, placement +): + package = make_package(tmp_path, layered=layered) + before = snapshot(package.root) + out, config, entries, model, decoder_dir = build( + monkeypatch, + offload, + package, + tmp_path, + "--offload-mode", + mode, + "--offload-placement", + placement, + layered=layered, + ) + + assert entries == {PREFIX + "mode": mode, PREFIX + "placement": placement, PREFIX + "cpu_execution": "auto"} + assert {(node.op_type, node.domain) for node in model.graph.node} == {("QMoE", "com.microsoft")} + assert all({attribute.name for attribute in node.attribute} <= QMOE_ATTRIBUTES for node in model.graph.node) + source = onnx.load(str(package.decoder_dir / "model.onnx"), load_external_data=False) + for old, new in zip(source.graph.node, model.graph.node, strict=True): + assert (new.name, list(new.output)) == (old.name, list(old.output)) + assert list(new.input) == [*old.input[:14], "", *old.input[15:]] + assert [a.SerializeToString() for a in new.attribute] == [a.SerializeToString() for a in old.attribute] + assert list(model.graph.initializer) == list(source.graph.initializer) + + # The runtime offloads every node of the built graph. It would not offload the nodes of the source graph, whose + # router_weights input is not supported by the offload backend. + assert simulate_session_creation(model, entries, decoder_dir) == NAMES + if mode == "required": + with pytest.raises(ContractError, match="rejected node"): + simulate_session_creation(source, entries, package.decoder_dir) + else: + assert simulate_session_creation(source, entries, package.decoder_dir) == [] + + assert multidraft_calls == [("model_s2.onnx", "model_s3.onnx", 3), ("model_s2.onnx", "model_s4.onnx", 4)] + assert {path.name for path in decoder_dir.glob("model*.onnx")} == { + "model.onnx", + "model_s2.onnx", + "model_s3.onnx", + "model_s4.onnx", + } + assert snapshot(package.root) == before + manifest = json.loads((out / "package_manifest.json").read_text()) + assert manifest["stack"] == "v12b-engine-qmoe-offload-v4" + assert manifest["target_qmoe_nodes"] == manifest["router_weights_cleared"] == LAYERS + assert manifest["qmoe_nodes_not_offloaded"] == [] + assert manifest["offload_session_options"] == entries + + +def test_budgets_and_the_other_options_reach_the_config(offload, tmp_path, monkeypatch, multidraft_calls): + package = make_package(tmp_path) + out, config, entries, model, decoder_dir = build( + monkeypatch, + offload, + package, + tmp_path, + "--device-expert-bytes", + "auto", + "--host-expert-bytes", + str(40 << 30), + "--stop-token-id", + "9", + "--vram-limit-gib", + "24", + "--prefill-chunk", + "256", + ) + assert entries == { + PREFIX + "mode": "required", + PREFIX + "placement": "adaptive", + PREFIX + "cpu_execution": "auto", + PREFIX + "device_expert_bytes": "auto", + PREFIX + "host_expert_bytes": str(40 << 30), + } + assert simulate_session_creation(model, entries, decoder_dir) == NAMES + decoder = config["model"]["decoder"]["session_options"] + assert decoder["log_id"] == "onnxruntime-genai" + assert decoder["provider_options"][0]["cuda"] == { + "enable_cuda_graph": "1", + "arena_extend_strategy": "kSameAsRequested", + "arena.extend_strategy": "1", + "gpu_device_mem_limit": str(24 << 30), + } + assert config["model"]["eos_token_id"] == [9] + prefill = config["model"]["qwen38_graph"]["prefill"] + assert (prefill["chunk"], prefill["warmup_lengths"]) == (256, [256, 128, 37]) + assert ( + config["model"]["qwen38_graph"]["fold"]["mtp"]["session_options"][ + "session.use_device_allocator_for_initializers" + ] + == "1" + ) + assert not any(key.startswith(PREFIX) for key in config["model"]["qwen38_graph"]["fold"]["mtp"]["session_options"]) + + +def test_a_package_without_a_dense_graph_builds(offload, tmp_path, monkeypatch, multidraft_calls): + package = make_package(tmp_path, dense=False) + out, _config, entries, model, decoder_dir = build(monkeypatch, offload, package, tmp_path) + assert simulate_session_creation(model, entries, decoder_dir) == NAMES + assert not (out / "dense").exists() + + +def test_copy_makes_independent_files_and_the_default_links_them(offload, tmp_path, monkeypatch, multidraft_calls): + package = make_package(tmp_path) + (package.root / "weights.data").write_bytes(b"expert bytes") + out, *_ = build(monkeypatch, offload, package, tmp_path) + assert os.path.samefile(package.root / "weights.data", out / "weights.data") + shutil.rmtree(out) + out, *_ = build(monkeypatch, offload, package, tmp_path, "--copy") + assert (out / "weights.data").read_bytes() == b"expert bytes" + assert not os.path.samefile(package.root / "weights.data", out / "weights.data") + + +def test_ops_libraries_are_copied_into_the_package(offload, tmp_path, monkeypatch, multidraft_calls): + package = make_package(tmp_path) + ops = tmp_path / "ops" + ops.mkdir() + (ops / "libqwen_a.so").write_bytes(b"a") + (ops / "other.so").write_bytes(b"other") + out, *_ = build(monkeypatch, offload, package, tmp_path, "--ops-dir", ops) + assert sorted(path.name for path in (out / "libs").iterdir()) == ["libqwen38_plane.so", "libqwen_a.so"] + assert not os.path.samefile(ops / "libqwen_a.so", out / "libs/libqwen_a.so") + + +def test_target_resident_slots_caps_every_node_by_name(offload, tmp_path, monkeypatch, multidraft_calls): + package = make_package(tmp_path) + _out, _config, entries, model, decoder_dir = build( + monkeypatch, offload, package, tmp_path, "--target-resident-slots", "5" + ) + assert caps_of(entries) == dict.fromkeys(NAMES, 5) + assert PREFIX + "popularity_file" not in entries + assert simulate_session_creation(model, entries, decoder_dir) == NAMES + + +def test_an_expert_profile_alone_sets_no_caps(offload, tmp_path, monkeypatch, multidraft_calls): + package = make_package(tmp_path) + profile = write_profile(tmp_path / "profile.bin", exhaustive_pairs()) + out, _config, entries, model, decoder_dir = build( + monkeypatch, offload, package, tmp_path, "--expert-profile", profile + ) + assert entries[PREFIX + "popularity_file"] == "expert-profile.bin" + assert PREFIX + "node_overrides" not in entries + assert (decoder_dir / "expert-profile.bin").read_bytes() == profile.read_bytes() + assert simulate_session_creation(model, entries, decoder_dir) == NAMES + validate_popularity_file(decoder_dir / "expert-profile.bin") + + +@pytest.mark.parametrize("layered", [False, True]) +def test_the_profile_is_copied_next_to_the_decoder_without_touching_the_source( + offload, tmp_path, monkeypatch, multidraft_calls, layered +): + package = make_package(tmp_path, layered=layered) + (package.decoder_dir / "expert-profile.bin").write_bytes(b"stale profile of the source package") + profile = write_profile(tmp_path / "profile.bin", exhaustive_pairs()) + out, _config, _entries, _model, decoder_dir = build( + monkeypatch, offload, package, tmp_path, "--expert-profile", profile, layered=layered + ) + assert (package.decoder_dir / "expert-profile.bin").read_bytes() == b"stale profile of the source package" + assert (decoder_dir / "expert-profile.bin").read_bytes() == profile.read_bytes() + assert not os.path.samefile(decoder_dir / "expert-profile.bin", profile) + assert not os.path.samefile(decoder_dir / "expert-profile.bin", package.decoder_dir / "expert-profile.bin") + + +@pytest.mark.parametrize("uniform", [False, True]) +@pytest.mark.parametrize("total", [10, 23, 24]) +def test_a_profile_and_total_spread_exactly_that_many_slots( + offload, tmp_path, monkeypatch, multidraft_calls, uniform, total +): + package = make_package(tmp_path) + profile = write_profile(tmp_path / "profile.bin", exhaustive_pairs()) + extra = ["--uniform-profile-slots"] if uniform else [] + _out, _config, entries, model, decoder_dir = build( + monkeypatch, offload, package, tmp_path, "--expert-profile", profile, "--total-resident-slots", total, *extra + ) + caps = caps_of(entries) + assert list(caps) == NAMES + assert sum(caps.values()) == total + if uniform: + assert max(caps.values()) - min(caps.values()) <= 1 + assert all(0 <= cap <= EXPERTS for cap in caps.values()) + assert simulate_session_creation(model, entries, decoder_dir) == NAMES + + +def test_target_slots_with_a_profile_are_a_per_layer_average(offload, tmp_path, monkeypatch, multidraft_calls): + package = make_package(tmp_path) + profile = write_profile(tmp_path / "profile.bin", exhaustive_pairs()) + _out, _config, entries, _model, _decoder_dir = build( + monkeypatch, offload, package, tmp_path, "--expert-profile", profile, "--target-resident-slots", 5 + ) + assert sum(caps_of(entries).values()) == 5 * LAYERS + + +def test_a_total_the_profile_cannot_fill_stops_the_build_before_it_writes( + offload, tmp_path, monkeypatch, multidraft_calls +): + package = make_package(tmp_path) + profile = write_profile(tmp_path / "profile.bin", exhaustive_pairs()) + before = snapshot(tmp_path) + with pytest.raises(SystemExit, match="error: --expert-profile: 25 resident slots requested but .* ranks only 24"): + run_builder( + monkeypatch, offload, package, tmp_path / "out", "--expert-profile", profile, "--total-resident-slots", 25 + ) + assert snapshot(tmp_path) == before + + +def ineligible_model(layer=1): + model = decoder_model() + attributes = {attribute.name: attribute for attribute in model.graph.node[layer].attribute} + attributes["swiglu_fusion"].i = 2 + return model + + +def test_required_mode_refuses_nodes_the_runtime_cannot_offload(offload, tmp_path, monkeypatch, multidraft_calls): + package = make_package(tmp_path, model=ineligible_model()) + before = snapshot(tmp_path) + with pytest.raises(SystemExit) as excinfo: + run_builder(monkeypatch, offload, package, tmp_path / "out") + message = str(excinfo.value) + assert message.startswith("1 of 3 QMoE nodes cannot be offloaded by the runtime; first, ") + assert repr(NAMES[1]) in message and "swiglu_fusion=2" in message and "--offload-mode auto" in message + assert snapshot(tmp_path) == before + + +@pytest.mark.parametrize("caps", [[], ["--target-resident-slots", "4"]]) +def test_auto_mode_keeps_the_nodes_the_runtime_cannot_offload_on_the_gpu( + offload, tmp_path, monkeypatch, multidraft_calls, capsys, caps +): + package = make_package(tmp_path, model=ineligible_model()) + _out, _config, entries, model, decoder_dir = build( + monkeypatch, offload, package, tmp_path, "--offload-mode", "auto", *caps + ) + assert "warning: 1 of 3 QMoE nodes cannot be offloaded by the runtime" in capsys.readouterr().out + assert simulate_session_creation(model, entries, decoder_dir) == [NAMES[0], NAMES[2]] + if caps: + assert caps_of(entries) == {NAMES[0]: 4, NAMES[2]: 4} + manifest = json.loads((tmp_path / "out/package_manifest.json").read_text()) + assert manifest["qmoe_nodes_not_offloaded"] == [repr(NAMES[1])] + assert manifest["router_weights_cleared"] == LAYERS + + +def test_auto_mode_with_a_profile_keeps_the_layer_numbers_of_the_other_nodes( + offload, tmp_path, monkeypatch, multidraft_calls +): + package = make_package(tmp_path, model=ineligible_model(layer=0)) + profile = write_profile(tmp_path / "profile.bin", exhaustive_pairs()) + _out, _config, entries, model, decoder_dir = build( + monkeypatch, + offload, + package, + tmp_path, + "--offload-mode", + "auto", + "--expert-profile", + profile, + "--total-resident-slots", + 12, + ) + assert list(caps_of(entries)) == [NAMES[1], NAMES[2]] + assert simulate_session_creation(model, entries, decoder_dir) == [NAMES[1], NAMES[2]] + + +def test_unnamed_nodes_cannot_get_caps_but_build_without_them(offload, tmp_path, monkeypatch, multidraft_calls): + unnamed = decoder_model(name="") + package = make_package(tmp_path, model=unnamed) + before = snapshot(tmp_path) + with pytest.raises( + SystemExit, match=r"error: per-node resident-slot caps need uniquely named QMoE nodes, but #0 \(unnamed\)" + ): + run_builder(monkeypatch, offload, package, tmp_path / "out", "--target-resident-slots", "4") + assert snapshot(tmp_path) == before + _out, _config, entries, model, decoder_dir = build(monkeypatch, offload, package, tmp_path) + assert simulate_session_creation(model, entries, decoder_dir) == ["", "", ""] + + +def test_nodes_sharing_a_name_cannot_get_caps(offload, tmp_path, monkeypatch, multidraft_calls): + package = make_package(tmp_path, model=decoder_model(name="same")) + with pytest.raises( + SystemExit, match=r"error: per-node resident-slot caps need uniquely named QMoE nodes, but 'same' " + ): + run_builder(monkeypatch, offload, package, tmp_path / "out", "--target-resident-slots", "4") + assert not (tmp_path / "out").exists() + + +def test_a_legacy_qmoe_offload_graph_is_refused(offload, tmp_path, monkeypatch, multidraft_calls): + legacy = decoder_model() + legacy.graph.node[1].op_type = "QMoEOffload" + package = make_package(tmp_path, model=legacy) + before = snapshot(tmp_path) + with pytest.raises(SystemExit, match=r"model\.onnx has 1 QMoEOffload nodes\. The runtime has no such operator"): + run_builder(monkeypatch, offload, package, tmp_path / "out") + assert snapshot(tmp_path) == before + + +def test_a_legacy_qmoe_offload_dense_graph_is_refused(offload, tmp_path, monkeypatch, multidraft_calls): + package = make_package(tmp_path) + dense = dense_model() + dense.graph.node[0].op_type = "QMoEOffload" + save_model(dense, package.root / "dense/model.onnx") + with pytest.raises(SystemExit, match=r"dense/model\.onnx has 1 QMoEOffload nodes"): + run_builder(monkeypatch, offload, package, tmp_path / "out") + assert not (tmp_path / "out").exists() + + +def test_a_graph_without_qmoe_is_refused(offload, tmp_path, monkeypatch, multidraft_calls): + graph = helper.make_graph( + [helper.make_node("Identity", ["x"], ["y"])], + "g", + [helper.make_tensor_value_info("x", TensorProto.FLOAT, [1])], + [helper.make_tensor_value_info("y", TensorProto.FLOAT, [1])], + ) + package = make_package(tmp_path, model=helper.make_model(graph)) + with pytest.raises(SystemExit, match="has no com.microsoft::QMoE node to offload"): + run_builder(monkeypatch, offload, package, tmp_path / "out") + assert not (tmp_path / "out").exists() + + +@pytest.mark.parametrize("value", [None, "1"]) +@pytest.mark.parametrize("flag", REMOVED_FLAGS) +def test_removed_flags_fail_instead_of_being_ignored(offload, tmp_path, monkeypatch, capsys, flag, value): + package = make_package(tmp_path) + before = snapshot(tmp_path) + with pytest.raises(SystemExit) as excinfo: + run_builder(monkeypatch, offload, package, tmp_path / "out", flag, *([value] if value else [])) + assert excinfo.value.code == 2 + assert f"{flag} was removed" in capsys.readouterr().err + assert snapshot(tmp_path) == before + + +def test_every_removed_flag_is_tested(offload): + assert sorted(offload.REMOVED_OPTIONS) == sorted(REMOVED_FLAGS) + + +@pytest.mark.parametrize( + ("flag", "value"), + [ + ("--offload-mode", "disabled"), + ("--offload-placement", "basic"), + ("--device-expert-bytes", "-1"), + ("--host-expert-bytes", "12.5"), + ("--target-resident-slots", "-2"), + ("--mtp-moe-placement", "cpu"), + ], +) +def test_option_values_the_runtime_would_reject_are_refused_by_the_parser(offload, tmp_path, monkeypatch, flag, value): + package = make_package(tmp_path) + with pytest.raises(SystemExit) as excinfo: + run_builder(monkeypatch, offload, package, tmp_path / "out", flag, value) + assert excinfo.value.code == 2 + assert not (tmp_path / "out").exists() + + +def test_missing_arguments_and_files_are_named(offload, tmp_path, monkeypatch): + with pytest.raises(SystemExit, match="--from-package, --out, --plane-lib, and --provider-lib are required"): + run_main(monkeypatch, offload, "--out", tmp_path / "out") + for missing, label in ( + ("genai_config.json", "genai_config.json of --from-package"), + ("weights_pool.txt", "weights_pool.txt of --from-package"), + ("model.onnx", "the target graph of --from-package"), + ): + package = make_package(tmp_path / missing) + (package.root / missing).unlink() + with pytest.raises(SystemExit, match=f"{label} not found"): + run_builder(monkeypatch, offload, package, tmp_path / missing / "out") + package = make_package(tmp_path / "missing-libs") + out = tmp_path / "missing-libs" / "out" + for kind, label in (("plane", "--plane-lib"), ("provider", "--provider-lib")): + library = getattr(package, kind) + library.unlink() + with pytest.raises(SystemExit, match=f"{label} not found"): + run_builder(monkeypatch, offload, package, out) + library.write_bytes(b"restored") + with pytest.raises(SystemExit, match="--expert-profile not found"): + run_builder(monkeypatch, offload, package, out, "--expert-profile", tmp_path / "absent.bin") + assert not out.exists() + + +OVERLAPS = [ + ("--from-package", "equal"), + ("--from-package", "inside"), + ("--from-package", "contains"), + ("--plane-lib", "equal"), + ("--plane-lib", "contains"), + ("--provider-lib", "contains"), + ("--ops-dir", "equal"), + ("--ops-dir", "inside"), + ("--ops-dir", "contains"), + ("--expert-profile", "equal"), + ("--expert-profile", "contains"), +] + + +@pytest.mark.parametrize(("flag", "relation"), OVERLAPS) +def test_an_output_overlapping_an_input_is_refused_before_anything_changes( + offload, tmp_path, monkeypatch, multidraft_calls, flag, relation +): + package = make_package(tmp_path) + ops = tmp_path / "ops" + ops.mkdir() + (ops / "libqwen_a.so").write_bytes(b"a") + (tmp_path / "profile").mkdir() + profile = write_profile(tmp_path / "profile" / "expert-profile.bin", exhaustive_pairs()) + inputs = { + "--from-package": package.root, + "--plane-lib": package.plane, + "--provider-lib": package.provider, + "--ops-dir": ops, + "--expert-profile": profile, + } + path = inputs[flag] + out = {"equal": path, "inside": path / "out", "contains": path.parent}[relation] + before = snapshot(tmp_path) + with pytest.raises(SystemExit, match="must not be equal to, nested inside or contain"): + run_builder(monkeypatch, offload, package, out, "--ops-dir", ops, "--expert-profile", profile) + assert snapshot(tmp_path) == before + + +@requires_copy_file_range +def test_host_mapped_mtp_pools_of_any_id_are_kept_end_to_end(offload, tmp_path, monkeypatch, multidraft_calls): + ranges = [(1, 4097, 100), (2, 9001, 5000)] + package, data = mtp_pool_package(tmp_path, ranges) + out, _config, _entries, _model, _decoder_dir = build( + monkeypatch, offload, package, tmp_path, "--mtp-moe-placement", "host_mapped", layered=True + ) + records, bindings = read_pool(out / "weights_pool.txt") + check_bound(records, bindings) + check_aligned_copy(out, data, ranges, records, {1: 0, 2: 4096}) + assert {pid for _sess, _name, pid in bindings} >= {1, 2} + manifest = json.loads((out / "package_manifest.json").read_text()) + assert manifest["pool_stats"]["mtp_moe_placement"] == "host_mapped" + assert manifest["pool_stats"]["total_pool_entries"] == len(records) + + +def test_the_default_mtp_placement_leaves_the_pools_alone(offload, tmp_path, monkeypatch, multidraft_calls): + ranges = [(1, 4097, 100), (2, 9001, 5000)] + package, _data = mtp_pool_package(tmp_path, ranges) + out, *_ = build(monkeypatch, offload, package, tmp_path, layered=True) + records, bindings = read_pool(out / "weights_pool.txt") + check_bound(records, bindings) + assert records[1][3:5] == ["4097", "100"] and records[2][3:5] == ["9001", "5000"] + assert not (out / ALIGNED).exists() + + +# --------------------------------------------------------------------------- +# Sessions in the installed runtime +# --------------------------------------------------------------------------- + + +def stock_decoder(layers=2, op_type="QMoE", extra_attributes=None): + """A tiny INT4 (block 32) swiglu QMoE decoder with embedded weights that the released CPU runtime can load.""" + experts, hidden, inter, block = 4, 64, 64, 32 + rng = np.random.default_rng(0) + nodes, initializers = [], [] + graph_inputs = [helper.make_tensor_value_info("x0", TensorProto.FLOAT16, [3, hidden])] + for layer in range(layers): + for suffix, array in ( + ("fc1_w", rng.integers(0, 255, (experts, 2 * inter, hidden // 2), dtype=np.uint8)), + ("fc1_s", rng.random((experts, 2 * inter, hidden // block), dtype=np.float32).astype(np.float16)), + ("fc2_w", rng.integers(0, 255, (experts, hidden, inter // 2), dtype=np.uint8)), + ("fc2_s", rng.random((experts, hidden, inter // block), dtype=np.float32).astype(np.float16)), + ): + initializers.append(numpy_helper.from_array(array, f"l{layer}_{suffix}")) + graph_inputs.append(helper.make_tensor_value_info(f"router{layer}", TensorProto.FLOAT16, [3, experts])) + nodes.append( + helper.make_node( + op_type, + [ + f"x{layer}", + f"router{layer}", + f"l{layer}_fc1_w", + f"l{layer}_fc1_s", + "", + f"l{layer}_fc2_w", + f"l{layer}_fc2_s", + ], + [f"x{layer + 1}"], + name=f"model.layers.{layer}.mlp", + domain="com.microsoft", + activation_type="swiglu", + swiglu_fusion=1, + k=2, + expert_weight_bits=4, + block_size=block, + normalize_routing_weights=1, + **(extra_attributes or {}), + ) + ) + graph = helper.make_graph( + nodes, + "decoder", + graph_inputs, + [helper.make_tensor_value_info(f"x{layers}", TensorProto.FLOAT16, [3, hidden])], + initializer=initializers, + ) + model = helper.make_model( + graph, opset_imports=[helper.make_opsetid("", 17), helper.make_opsetid("com.microsoft", 1)] + ) + model.ir_version = 9 + return model + + +def create_session(ort, model, entries=None): + options = ort.SessionOptions() + for key, value in (entries or {}).items(): + options.add_session_config_entry(key, value) + return ort.InferenceSession( + model.SerializeToString() if isinstance(model, onnx.ModelProto) else str(model), + options, + providers=["CPUExecutionProvider"], + ) + + +@pytest.fixture(scope="module") +def ort(): + runtime = pytest.importorskip("onnxruntime") + try: + create_session(runtime, stock_decoder()) + except Exception as error: + pytest.skip(f"the installed onnxruntime cannot load a tiny QMoE graph: {error}") + return runtime + + +@pytest.mark.parametrize("mode", ["required", "auto"]) +@pytest.mark.parametrize("placement", ["adaptive", "static"]) +def test_the_built_package_creates_a_session_in_the_installed_runtime( + ort, offload, tmp_path, monkeypatch, multidraft_calls, mode, placement +): + package = make_package(tmp_path, model=stock_decoder()) + _out, _config, entries, model, decoder_dir = build( + monkeypatch, + offload, + package, + tmp_path, + "--offload-mode", + mode, + "--offload-placement", + placement, + "--target-resident-slots", + "2", + ) + assert set(entries) == {PREFIX + key for key in ("mode", "placement", "cpu_execution", "node_overrides")} + # The released runtime ignores session.moe.offload.*; a build with QMoE offload validates every entry here. + for graph in ("model.onnx", "model_s2.onnx", "model_s3.onnx", "model_s4.onnx"): + create_session(ort, decoder_dir / graph, entries) + assert simulate_session_creation(model, entries, decoder_dir) == [f"model.layers.{layer}.mlp" for layer in range(2)] + + +@pytest.mark.parametrize( + ("graph", "message"), + [ + (lambda: stock_decoder(op_type="QMoEOffload"), "QMoEOffload"), + (lambda: stock_decoder(extra_attributes={"layer_index": 0}), "layer_index"), + (lambda: stock_decoder(extra_attributes={"resident_slots": 2}), "resident_slots"), + ], + ids=["QMoEOffload operator", "layer_index attribute", "resident_slots attribute"], +) +def test_what_the_previous_builder_emitted_does_not_load(ort, graph, message): + with pytest.raises(Exception, match=message) as excinfo: + create_session(ort, graph()) + assert type(excinfo.value).__name__ == "InvalidGraph" diff --git a/tools/qwen38_graph/README.md b/tools/qwen38_graph/README.md index dc458f1672..e9bc4e6fa6 100644 --- a/tools/qwen38_graph/README.md +++ b/tools/qwen38_graph/README.md @@ -192,37 +192,67 @@ anything is deleted). `build_package_offload_v3.py --from-package --out --plane-lib --provider-lib ` converts a v12b package (layered, or flat with `model.onnx` in the root) so that the routed experts -stay in host memory and are streamed to the GPU on demand (`stack` is `v12b-engine-qmoe-offload-v3` in `package_manifest.json`). -Everything that does not change is hard-linked from the source (`--copy` copies): - -* the 48 `com.microsoft::QMoE` nodes of the target graph become `com.microsoft::QMoEOffload` nodes (explicit `layer_index`, - `hidden_size` and `inter_size`, plus tile, working-buffer and CUDA-graph attributes). The graph is written as `model_s2.onnx` (in - `target/` for a layered package); `model_s3.onnx` and `model_s4.onnx` extend it to 3 and 4 verify rows (2 or 3 chained MTP drafts), - and `model.onnx` is a copy of `model_s4.onnx`; +stay in host memory and are streamed to the GPU on demand (`stack` is `v12b-engine-qmoe-offload-v4` in `package_manifest.json`). +Everything that does not change is hard-linked from the source (`--copy` copies). + +The package targets the ONNX Runtime CUDA provider with QMoE offload (microsoft/onnxruntime#33450 and #33622). That runtime has no +`QMoEOffload` operator and the `com.microsoft::QMoE` schema has no offload attribute, so the graphs keep standard `QMoE` nodes and +the decoder session is configured with `session.moe.offload.*` session options, which a runtime without the offload ignores. A source +whose graphs have `QMoEOffload` nodes (a package of the earlier builder) is refused: rebuild it with standard `QMoE` nodes. The +contract was checked against the runtime sources (commit `1d5ae2f3` of #33622, mirrored by +`test/python/builder/_qwen38_qmoe_offload_contract.py`) and against the installed onnxruntime in the tests. No package built here has +been loaded by a CUDA build of #33622. + +* the 48 `com.microsoft::QMoE` nodes of the target graph keep their operator and attributes. The only edit is clearing the optional + `router_weights` input (14), which the offload backend does not take; the nodes then weight the selected experts by + `router_probs` (`router_weights_cleared` in `package_manifest.json`). The graph is written as `model_s2.onnx` (in `target/` for a + layered package); `model_s3.onnx` and `model_s4.onnx` extend it to 3 and 4 verify rows (2 or 3 chained MTP drafts), and + `model.onnx` is a copy of `model_s4.onnx`; +* `--offload-mode` (`required`, the default, or `auto`) is `session.moe.offload.mode`. The runtime offloads a node that is NVFP4 with + `weights_prepacked=1` (or INT4 with block size 32 and prepacked weights, or raw INT3), uses swiglu with `swiglu_fusion` 0 or 1, + has constant rank-3 expert weights and scales (and constant NVFP4 global scales) and none of the fc3, zero-point, `router_weights` + and activation-scale inputs. With `required` the build stops, before it writes anything, at a node that breaks a rule (the session + creation would fail); with `auto` it warns and the runtime keeps those nodes on the GPU (`qmoe_nodes_not_offloaded` in the + manifest). `--offload-placement` (`adaptive`, the default, or `static`: no spare slots) is `session.moe.offload.placement`. + `--device-expert-bytes` and `--host-expert-bytes` (`auto` or a byte count) set the matching budgets; an option that is not given + writes no entry; * `weights_pool.txt` puts the routed-expert weights and scales (63.28 GiB) in `cpu` mmap placement. `--dense-placement` (default `host_mapped`) places the `dense.data` weights of the dense prefill session, `--embed-placement` (default `host_mapped`) the token - embedding and `--mtp-moe-placement` (default `device`) the MTP-head experts; -* `dense/model.onnx` keeps the stock GPU `QMoE` on the UVA-mapped expert weights; `--dense-offload-qmoe` rewrites it to `QMoEOffload` - streaming `--prefill-chunk` rows instead; -* `--target-resident-slots` and `--target-spare-slots` set the resident and spare expert slots of every target node. - `--expert-profile ` (an STRP v1 expert-popularity file) chooses the initial resident experts per layer instead, - `--total-resident-slots` in total (default 48 x `--target-resident-slots`, or 48 x 20) with a floor of `--min-slots-per-layer` (default 4) - per layer, or the same count for every layer with `--uniform-profile-slots`; -* `genai_config.json` gets CUDA-graph capture and `kSameAsRequested` arena options for the decoder, MTP-head, dense and head-KV sessions; - `--vram-limit-gib` becomes `gpu_device_mem_limit`, `--stop-token-id` replaces `eos_token_id` and `--prefill-chunk` (default 512) sets - `prefill.chunk`; + embedding and `--mtp-moe-placement` (default `device`) the MTP-head experts. Every pool keeps its id, and the build refuses a source + or an output whose sessions name a pool that is not defined, or whose pool ids repeat; +* `dense/model.onnx` keeps the stock GPU `QMoE` on the UVA-mapped expert weights (`--dense-gpu-qmoe`, always on); +* `--target-resident-slots N` caps the resident experts of every offloaded node at N (`max_resident_slots` in + `session.moe.offload.node_overrides`, which names the nodes, so they must have unique names). + `--expert-profile ` (an STRP v1 expert-popularity file) is copied next to the decoder graph as a new file and + named by `session.moe.offload.popularity_file`; alone it sets no caps and the byte budgets decide. With `--total-resident-slots` + (or `--target-resident-slots`, meaning N per node) the profile also decides the cap of every layer: each layer first gets + `--min-slots-per-layer` (default 4, at most total / layers), then the most popular remaining experts of any layer follow; with + `--uniform-profile-slots` every layer gets the total divided by the layers, rounded down, and the remaining slots (the total modulo + the layers) go to different layers, the ones whose next expert ranks best. The caps always add up to the requested total; a total + the profile cannot fill, or a layer that ranks fewer experts than its uniform share, is an error before anything is written. The + profile ranks experts by layer, so the build also checks that the runtime numbers the QMoE nodes like their positions in the graph + (the number after `layers.`, `layer.` or `blocks.` in the node name, or in the name of its fc1 weights, else the registration + order); +* `genai_config.json` gets the `session.moe.offload.*` options on the decoder, and CUDA-graph capture and `kSameAsRequested` arena + options for the decoder, MTP-head, dense and head-KV sessions; `--vram-limit-gib` becomes `gpu_device_mem_limit`, + `--stop-token-id` replaces `eos_token_id` and `--prefill-chunk` (default 512) sets `prefill.chunk`; * `--plane-lib` and `--provider-lib` are copied to `libs/` and `provider/`; `--ops-dir` adds its `libqwen_*.so` files to `libs/`. -`--out` must not be equal to or nested in `--from-package` (or contain it). It does not have to be new or empty, so a rebuild replaces -its files, and for a source without `dense/model.onnx` the script removes stale `dense/` and `head_kv/` (flat source: also -`target/` and `head_r2/`) folders from it. The source package is never modified: every graph is saved to a new file, and the -`model_s2.onnx`, `model_s3.onnx` and `model_s4.onnx` of a source that already has them are regenerated instead of hard-linked. With -`--mtp-moe-placement host_mapped` it can also write a 4.7 GiB `_mtp_moe_aligned.bin` next to the source package (Linux only: -`os.copy_file_range`). Use `run_with_ram_budget.py` to run an offload package with a bounded host expert cache. - -`--migrate-shared-onnx --out-onnx ` is a separate mode that rewrites one graph: legacy 26-input `QMoEOffload` nodes -move to the compact schema (inputs 21, 24 and 25 become inputs 2, 3 and 4; `--strip-phase` drops the optional phase input and -`--strip-stats` the statistics outputs) without touching the external weight data. +`--out` must not be equal to, nested in or contain `--from-package`, `--plane-lib`, `--provider-lib`, `--ops-dir` or +`--expert-profile` (checked before anything is written). It does not have to be new or empty, so a rebuild replaces its files, and +for a source without `dense/model.onnx` the script removes stale `dense/` and `head_kv/` (flat source: also `target/` and `head_r2/`) +folders from it. The source package is never modified: every graph and the profile are saved to new files, and the `model_s2.onnx`, +`model_s3.onnx` and `model_s4.onnx` of a source that already has them are regenerated instead of hard-linked. With +`--mtp-moe-placement host_mapped`, the MTP-head pools of the source that do not start on a 16-byte boundary (whatever their pool ids) +are copied (Linux only: `os.copy_file_range`) into one aligned file next to the source package, +`_mtp_moe_aligned-.bin` (4.7 GiB for the real MTP head). The digest covers the pool layout and the size and +modification time of the source files, so a changed source gets a new file, and a failed copy leaves no partial file. Use +`run_with_ram_budget.py` to run an offload package with a bounded host expert cache. + +The flags of the earlier builder that only made sense for `QMoEOffload` (`--migrate-shared-onnx`, `--out-onnx`, `--strip-phase`, +`--strip-stats`, `--target-spare-slots`, `--dense-offload-qmoe`, `--dense-stream-min-rows` and `--dense-stream-wave`) were removed; +they stop with a message that names the flag and the reason. A legacy 26-input `QMoEOffload` graph cannot be converted to standard +`QMoE` without losing data: its `expert_blocks` input has no standard counterpart. `build_multidraft_target_onnx` of this script also generates the S=3 and S=4 graphs of the flat packages: `build_package_engine_v12b.py` (with `--clean-package`) and `build_package_v12d_int4_int8.py` (with `--clean-flat-package`) import it from this folder. @@ -254,7 +284,7 @@ The builders produce a package; every other script needs one, and the ones that | `build_package_engine_v12b.py` | v12b / v12c / v12d package (folded cycle + fast prefill, optional `--kv-quant`, `--weight-quant`, `--mtp-moe-quant`, `--hc-quant`, `--share-dense-weights`, and the flat `--clean-package` layout with the S=2..4 target graphs and offload `runtime_profiles`) from an existing v12b package | | `build_package_v12d_int4_int8.py` | INT4 QMoE + INT8 MatMulNBits variant of the v12d package for SM80 and newer (layered, or flat with `--clean-flat-package` and the S=2..4 target graphs) | | `build_package_contrib_ops.py` | rewrites a flat RC2 package to standard ONNX and `com.microsoft` contrib operators only (no `custom_ops_libraries`) | -| `build_package_offload_v3.py` | QMoE-offload package from a v12b package (`QMoEOffload` nodes, host-resident routed experts, S=2..4 target graphs, optional expert profile and resident slots); `--migrate-shared-onnx` migrates one graph to the compact `QMoEOffload` schema | +| `build_package_offload_v3.py` | QMoE-offload package from a v12b package (standard `QMoE` nodes plus the `session.moe.offload.*` session options, host-resident routed experts, S=2..4 target graphs, optional expert profile and resident-slot caps); refuses a graph with legacy `QMoEOffload` nodes | | `qwen38_engine_util.py` | helpers shared by the scripts: CUDA EP registration and model loading with config overlays (`$QWEN38_OVERLAY_JSON`), the natural-text prompt pools, and `run_turn` (one greedy turn through `og.Engine`); registers the DLL directories on Windows | | `calibrate_qwen38_kv.py` | calibrate per-channel INT8 and FP8 KV-cache scales for the 12 target sparse-attention layers and 1 MTP-head sparse-attention layer (the JSON of the wanted quantization is the `--kv-scale-file` of `build_package_engine_v12b.py`); `--out` receives the `--quant-type` scales and the other quantization is written beside it under the name with the quantization token swapped (`kv_scales_int8_per_channel.json` -> `kv_scales_fp8_per_channel.json`, `scales.json` -> `scales.fp8_per_channel.json`), so the two files never collide (a collision is an error before any work) | | `smoke.py` | minimal example, prints the stream, speed and speculative statistics | diff --git a/tools/qwen38_graph/build_package_offload_v3.py b/tools/qwen38_graph/build_package_offload_v3.py index 80748d2765..7b891626a9 100644 --- a/tools/qwen38_graph/build_package_offload_v3.py +++ b/tools/qwen38_graph/build_package_offload_v3.py @@ -1,25 +1,28 @@ #!/usr/bin/env python3 # Copyright (c) Microsoft Corporation. All rights reserved. # Licensed under the MIT License. -"""Zero-copy package builder and graph migrator for QMoEOffload v3 (NVFP4 and SharedExpertBlockV1 INT4). +"""Zero-copy QMoE-offload package builder: standard com.microsoft::QMoE nodes plus session.moe.offload.* options. -Supports two migration modes: -1. NVFP4 v12b package migration (--from-package --out ): +The target runtime is the ONNX Runtime CUDA provider with QMoE offload (microsoft/onnxruntime#33450 -> #33622). It has no +QMoEOffload operator (creating a session with one fails) and the QMoE schema has no offload attribute: exporters emit +standard com.microsoft::QMoE nodes and the offload is configured with `session.moe.offload.*` session options. + +NVFP4 v12b package conversion (--from-package --out ): - Hard-links all weight files (0 GiB disk copy). - - Rewrites the 48 com.microsoft::QMoE nodes in target/model.onnx (or model.onnx) and dense/model.onnx to - com.microsoft::QMoEOffload with explicit layer_index, hidden_size, inter_size, and v3 dispatch controls. + - Checks that the runtime can offload the 48 com.microsoft::QMoE nodes of target/model.onnx (or model.onnx) and keeps + them as they are, apart from clearing the router_weights input (input 14), which the offload backend does not take. + The graph is written as model_s2.onnx; model_s3.onnx and model_s4.onnx extend it to 3 and 4 verify rows. - Updates weights_pool.txt so the 48 layers of routed expert weights/scales (63.281 GiB) use "cpu" mmap placement, and optionally places dense/source_weights/dense.data (5.270 GiB) in "host_mapped" UVA placement for 24 GiB VRAM targets. - - Configures genai_config.json for 24 GiB or 32 GiB VRAM operation. -2. Standalone SharedExpertBlockV1 v3 compact schema migration (--migrate-shared-onnx --out-onnx ): - - Rewrites legacy 26-input QMoEOffload nodes (21->2 expert_blocks, 24->3 fc2_bias, 25->4 optional phase) - without reading or regenerating the 63.457 GiB external weight payload (Section 5.3 / F8). + - Configures genai_config.json for 24 GiB or 32 GiB VRAM operation and sets the decoder's session.moe.offload.* + options: mode, placement, the byte budgets, the expert-popularity file and per-layer resident-slot caps. """ from __future__ import annotations import argparse +import collections import hashlib import json import os @@ -28,6 +31,33 @@ import onnx from onnx import helper +from package_paths import reject_output_overlap + +OFFLOAD_MODES = ("required", "auto") +OFFLOAD_PLACEMENTS = ("adaptive", "static") +OFFLOAD_OPTION_PREFIX = "session.moe.offload." +EXPERT_PROFILE_NAME = "expert-profile.bin" +ALIGNED_MTP_RELATIVE = "target/head_data/source_weights/mtp_moe_aligned.bin" +ROUTER_WEIGHTS_INPUT = 14 +# Inputs of the QMoE schema that the offload backend does not take: fc3 and zero-point banks, router_weights and the +# activation scales. +UNSUPPORTED_OFFLOAD_INPUTS = (8, 9, 10, 11, 12, 13, 14, 17, 18, 19, 20) +OFFLOAD_WEIGHT_INPUTS = (2, 3, 5, 6) +NVFP4_GLOBAL_SCALE_INPUTS = (15, 16) + +REMOVED_OPTIONS = { + "--migrate-shared-onnx": "it rewrites legacy 26-input QMoEOffload graphs, and the runtime has no QMoEOffload operator. " + "No data-preserving conversion of their expert_blocks input to standard QMoE exists; INT4 packages come from " + "build_package_v12d_int4_int8.py", + "--out-onnx": "it belongs to --migrate-shared-onnx", + "--strip-phase": "it belongs to --migrate-shared-onnx", + "--strip-stats": "it belongs to --migrate-shared-onnx", + "--target-spare-slots": "spare slots are the runtime's decision; choose --offload-placement static (none) or adaptive", + "--dense-offload-qmoe": "it rewrote dense/model.onnx to the QMoEOffload operator, which the runtime does not have; " + "the dense graph keeps the stock QMoE on UVA-mapped expert weights", + "--dense-stream-min-rows": "it set attributes of the QMoEOffload operator, which the runtime does not have", + "--dense-stream-wave": "it set attributes of the QMoEOffload operator, which the runtime does not have", +} REPLACED = { Path("genai_config.json"), @@ -114,46 +144,6 @@ def sha256(path: Path) -> str: return digest.hexdigest() -def migrate_shared_onnx_v3(in_onnx: Path, out_onnx: Path, strip_phase: bool = False, strip_stats: bool = False) -> int: - """Zero-copy v3 compact shared schema migration (21->2, 24->3, 25->4) per Section 5.3 (F8).""" - model = onnx.load(str(in_onnx), load_external_data=False) - migrated = 0 - stat_outputs_removed = set() - for node in model.graph.node: - if node.domain == "com.microsoft" and node.op_type == "QMoEOffload": - inputs = list(node.input) - if len(inputs) >= 22 and inputs[21]: - x_in = inputs[0] - router_in = inputs[1] - blocks_in = inputs[21] - bias_in = inputs[24] if len(inputs) > 24 else "" - phase_in = "" if strip_phase else (inputs[25] if len(inputs) > 25 else "") - new_inputs = [x_in, router_in, blocks_in] - if bias_in or phase_in: - new_inputs.append(bias_in) - if phase_in: - new_inputs.append(phase_in) - del node.input[:] - node.input.extend(new_inputs) - if strip_stats and len(node.output) > 1: - for out_name in list(node.output)[1:]: - if out_name: - stat_outputs_removed.add(out_name) - del node.output[1:] - existing_attrs = {a.name for a in node.attribute} - if "layout_id" not in existing_attrs: - node.attribute.append(helper.make_attribute("layout_id", "ort.qmoe.w4a16.block32.shared.v1")) - if "layer_index" not in existing_attrs: - node.attribute.append(helper.make_attribute("layer_index", migrated)) - migrated += 1 - if stat_outputs_removed: - kept_outputs = [o for o in model.graph.output if o.name not in stat_outputs_removed] - del model.graph.output[:] - model.graph.output.extend(kept_outputs) - save_model_fresh(model, out_onnx) - return migrated - - def load_expert_profile_plan( profile_path: Path, num_layers: int = 48, @@ -162,96 +152,268 @@ def load_expert_profile_plan( min_slots_per_layer: int = 4, uniform: bool = False, ) -> dict[int, list[int]]: - """Loads STRP v1 expert-profile.bin and assigns per-layer hot expert lists summing to total_slots.""" + """Loads an STRP v1 expert-profile.bin and picks the hottest experts of every layer, `total_slots` in total. + + The profile ranks (layer, expert) pairs from the most to the least popular, and every layer keeps a prefix of its + own ranking. Uniform: each layer gets total_slots // num_layers experts, and each of the total_slots % num_layers + remaining slots goes to a different layer, the ones whose next expert ranks highest. Otherwise every layer first + gets min_slots_per_layer (at most total_slots // num_layers), then the globally most popular experts take the rest. + A request the profile cannot fill raises ValueError, so the plan always sums to exactly `total_slots`. + """ import struct # noqa: PLC0415 import numpy as np # noqa: PLC0415 + if num_layers <= 0 or num_experts <= 0: + raise ValueError(f"num_layers and num_experts must be positive, got {num_layers} and {num_experts}") + if total_slots < 0 or min_slots_per_layer < 0: + raise ValueError( + f"total_slots and min_slots_per_layer must not be negative: {total_slots}, {min_slots_per_layer}" + ) + data = profile_path.read_bytes() + if len(data) < 24: + raise ValueError(f"{profile_path} is {len(data)} bytes, shorter than the 24-byte STRP header") magic, version, hdr_layers, hdr_experts, _slots, n_ranked = struct.unpack("<6I", data[:24]) if magic != 1347572819 or version != 1 or hdr_layers != num_layers or hdr_experts != num_experts: raise ValueError( f"Incompatible expert profile header in {profile_path}: {(magic, version, hdr_layers, hdr_experts)}" ) + if len(data) < 24 + n_ranked * 4: + raise ValueError(f"{profile_path} ranks {n_ranked} experts but holds only {(len(data) - 24) // 4}") pairs = np.frombuffer(data[24 : 24 + n_ranked * 4], dtype=np.uint16).reshape(n_ranked, 2) - - by_layer: dict[int, list[int]] = {layer: [] for layer in range(num_layers)} - for l_id, e_id in pairs: - by_layer[int(l_id)].append(int(e_id)) + if n_ranked and (int(pairs[:, 0].max()) >= num_layers or int(pairs[:, 1].max()) >= num_experts): + raise ValueError(f"{profile_path} ranks a layer or expert outside {num_layers} layers x {num_experts} experts") + + ranked: list[list[int]] = [[] for _ in range(num_layers)] + position: list[list[int]] = [[] for _ in range(num_layers)] + order: list[tuple[int, int]] = [] + seen: set[tuple[int, int]] = set() + for rank, (layer, expert) in enumerate(pairs.tolist()): + if (layer, expert) not in seen: + seen.add((layer, expert)) + ranked[layer].append(expert) + position[layer].append(rank) + order.append((layer, expert)) + + if total_slots > len(order): + raise ValueError(f"{total_slots} resident slots requested but {profile_path} ranks only {len(order)} experts") if uniform: - per_layer = total_slots // num_layers - return {layer: by_layer[layer][:per_layer] for layer in range(num_layers)} - - selected: dict[int, list[int]] = {layer: [] for layer in range(num_layers)} - selected_set: dict[int, set[int]] = {layer: set() for layer in range(num_layers)} - budget = total_slots - floor = min(min_slots_per_layer, total_slots // num_layers, num_experts) - for layer in range(num_layers): - for e in by_layer[layer][:floor]: - selected[layer].append(e) - selected_set[layer].add(e) - budget -= 1 - - for l_id, e_id in pairs: - if budget <= 0: - break - layer = int(l_id) - e = int(e_id) - if e not in selected_set[layer] and len(selected[layer]) < num_experts: - selected[layer].append(e) - selected_set[layer].add(e) - budget -= 1 - + base, remainder = divmod(total_slots, num_layers) + short = [layer for layer in range(num_layers) if len(ranked[layer]) < base] + if short: + raise ValueError( + f"{base} slots per layer requested but layers {short[:8]} rank fewer experts in the profile" + ) + candidates = sorted( + (layer for layer in range(num_layers) if len(ranked[layer]) > base), key=lambda layer: position[layer][base] + ) + if len(candidates) < remainder: + raise ValueError( + f"{total_slots} slots over {num_layers} layers leave {remainder} extra slots, but only " + f"{len(candidates)} layers rank more than {base} experts" + ) + extra = set(candidates[:remainder]) + selected = {layer: ranked[layer][: base + (layer in extra)] for layer in range(num_layers)} + else: + floor = min(min_slots_per_layer, total_slots // num_layers) + selected = {layer: ranked[layer][:floor] for layer in range(num_layers)} + chosen = {layer: set(experts) for layer, experts in selected.items()} + budget = total_slots - sum(len(experts) for experts in selected.values()) + for layer, expert in order: + if budget <= 0: + break + if expert not in chosen[layer]: + selected[layer].append(expert) + chosen[layer].add(expert) + budget -= 1 + + planned = sum(len(experts) for experts in selected.values()) + if planned != total_slots: + raise RuntimeError(f"expert profile plan holds {planned} slots, not the {total_slots} requested") return selected -def rewrite_nvfp4_qmoe_to_offload( - in_onnx: Path, - out_onnx: Path, - hidden_size: int = 2560, - inter_size: int = 640, - token_tile_size: int = 32, - extra_attrs: dict | None = None, - per_layer_initial_experts: dict[int, list[int]] | None = None, -) -> int: - """Rewrites com.microsoft::QMoE (quant_type='nvfp4') nodes to com.microsoft::QMoEOffload.""" - model = onnx.load(str(in_onnx), load_external_data=False) - rewritten = 0 - for node in model.graph.node: - if node.domain == "com.microsoft" and node.op_type == "QMoE": - layer_idx = rewritten - node.op_type = "QMoEOffload" - if len(node.input) > 14 and node.input[14]: - node.input[14] = "" - existing_attrs = {a.name for a in node.attribute} - if "layer_index" not in existing_attrs: - node.attribute.append(helper.make_attribute("layer_index", layer_idx)) - if "hidden_size" not in existing_attrs: - node.attribute.append(helper.make_attribute("hidden_size", hidden_size)) - if "inter_size" not in existing_attrs: - node.attribute.append(helper.make_attribute("inter_size", inter_size)) - if "token_tile_size" not in existing_attrs: - node.attribute.append(helper.make_attribute("token_tile_size", token_tile_size)) - if "gpu_working_bytes" not in existing_attrs: - node.attribute.append(helper.make_attribute("gpu_working_bytes", 512 << 20)) - if "pinned_staging_bytes" not in existing_attrs: - node.attribute.append(helper.make_attribute("pinned_staging_bytes", 64 << 20)) - if "allow_cuda_graph" not in existing_attrs: - node.attribute.append(helper.make_attribute("allow_cuda_graph", 1)) - if per_layer_initial_experts is not None and layer_idx in per_layer_initial_experts: - init_exps = per_layer_initial_experts[layer_idx] - node.attribute.append(helper.make_attribute("resident_slots", len(init_exps))) - node.attribute.append(helper.make_attribute("initial_experts", init_exps)) - existing_attrs.add("resident_slots") - existing_attrs.add("initial_experts") - if extra_attrs: - for k_attr, v_attr in extra_attrs.items(): - if k_attr not in existing_attrs and v_attr is not None: - node.attribute.append(helper.make_attribute(k_attr, v_attr)) - rewritten += 1 - save_model_fresh(model, out_onnx) - return rewritten +def qmoe_nodes(model: onnx.ModelProto) -> list[onnx.NodeProto]: + return [node for node in model.graph.node if node.domain == "com.microsoft" and node.op_type == "QMoE"] + + +def node_label(name: str, ordinal: int) -> str: + return repr(name) if name else f"#{ordinal} (unnamed)" + + +def qmoe_attribute(node: onnx.NodeProto, name: str, default=None): + """Reads a node attribute like the runtime does: strings come back as str and an absent one as `default`.""" + for attribute in node.attribute: + if attribute.name == name: + value = helper.get_attribute_value(attribute) + return value.decode() if isinstance(value, bytes) else value + return default + + +def parse_layer_index(name: str) -> int | None: + """Mirrors the runtime's layer-number extraction: digits after "layers.", "layer." or "blocks." that end the name or + are followed by a non-alphanumeric character, e.g. 5 in "model.layers.5.mlp". Returns None when there is none.""" + for prefix in ("layers.", "layer.", "blocks."): + search_from = 0 + while search_from < len(name): + position = name.find(prefix, search_from) + if position < 0: + break + end = position + len(prefix) + digits_end = end + while digits_end < len(name) and name[digits_end] in "0123456789": + digits_end += 1 + if digits_end > end: + value = int(name[end:digits_end]) + if value > 2**31 - 1: + return None + following = name[digits_end] if digits_end < len(name) else "" + if not following or not (following.isascii() and following.isalnum()): + return value + search_from = end + return None + + +def qmoe_offload_problems(node: onnx.NodeProto, constants: dict[str, onnx.TensorProto]) -> list[str]: + """Lists the reasons the runtime's QMoE offload (onnxruntime#33622, moe_quantization.cc) rejects `node`. + + `constants` maps the names of the graph's constant initializers to them. An empty list means the node is eligible; + `session.moe.offload.mode=required` fails session creation for a node with reasons, and `auto` runs that node + without offload. + """ + problems = [] + + def slot(index: int) -> str: + return node.input[index] if index < len(node.input) else "" + + quant_type = qmoe_attribute(node, "quant_type", "int") + bits = qmoe_attribute(node, "expert_weight_bits", 4) + fc1_bits = qmoe_attribute(node, "fc1_expert_weight_bits", bits) + fc2_bits = qmoe_attribute(node, "fc2_expert_weight_bits", bits) + block_size = qmoe_attribute(node, "block_size", -1) + prepacked = qmoe_attribute(node, "weights_prepacked", -1) + nvfp4_row_major = quant_type == "nvfp4" and prepacked == 1 + int4_prepacked = quant_type == "int" and bits == fc1_bits == fc2_bits == 4 and prepacked != 0 and block_size == 32 + int3_raw = ( + quant_type == "int" and bits == fc1_bits == fc2_bits == 3 and prepacked == 0 and block_size in (32, 64, 128) + ) + if quant_type == "nvfp4" and block_size not in (-1, 16): + problems.append(f"block_size={block_size}: quant_type='nvfp4' takes block_size 16") + if not (nvfp4_row_major or int4_prepacked or int3_raw): + problems.append( + f"quant_type={quant_type!r} expert_weight_bits={bits} block_size={block_size} weights_prepacked={prepacked} " + "is not NVFP4 with weights_prepacked=1, INT4 with block_size=32 and prepacked weights, or raw INT3" + ) + elif not (bits == 4 or int3_raw): + problems.append(f"expert_weight_bits={bits} must be 4 for NVFP4") + + activation = qmoe_attribute(node, "activation_type", "relu") + swiglu_fusion = qmoe_attribute(node, "swiglu_fusion", 0) + if activation != "swiglu" or swiglu_fusion not in (0, 1): + problems.append( + f"activation_type={activation!r} swiglu_fusion={swiglu_fusion}: only swiglu with swiglu_fusion 0 or 1 is offloaded" + ) + if qmoe_attribute(node, "use_sparse_mixer", 0) == 1: + problems.append("use_sparse_mixer=1 is not offloaded") + + unsupported = [index for index in UNSUPPORTED_OFFLOAD_INPUTS if slot(index)] + if unsupported: + problems.append(f"inputs {unsupported} are not supported by the offload backend") + + not_constant = [index for index in OFFLOAD_WEIGHT_INPUTS if slot(index) not in constants] + if not_constant: + problems.append(f"inputs {not_constant} (expert weights and scales) must be constant graph initializers") + else: + for index in (2, 5): + rank = len(constants[slot(index)].dims) + if rank != 3: + problems.append(f"input {index} has rank {rank}, expected 3 (experts x rows x columns)") + if quant_type == "nvfp4": + missing = [index for index in NVFP4_GLOBAL_SCALE_INPUTS if slot(index) not in constants] + if missing: + problems.append( + f"NVFP4 needs the constant global scales {list(NVFP4_GLOBAL_SCALE_INPUTS)}, missing {missing}" + ) + return problems + + +def prepare_qmoe_for_offload(model: onnx.ModelProto) -> dict: + """Readies the com.microsoft::QMoE nodes of `model` (edited in place) for the runtime's offload backend. + + The runtime offloads standard QMoE nodes and has no offload operator or attribute, so the nodes keep their op type and + attributes. The only edit is clearing the optional router_weights input (14), which the backend does not take: the + nodes then weight the selected experts by router_probs, like the fresh nodes of build_package_v12d_int4_int8.py. + Returns {"nodes": the names of the QMoE nodes in graph order, "router_weights_cleared": count, + "rejected": {position in "nodes": reasons}} for the nodes the backend would still reject. + """ + graph_inputs = {value.name for value in model.graph.input} + # An initializer that is also a graph input can be overridden at run time, so the runtime does not treat it as constant. + constants = { + init.name: init for init in model.graph.initializer if model.ir_version < 4 or init.name not in graph_inputs + } + nodes = qmoe_nodes(model) + cleared = 0 + for node in nodes: + if len(node.input) > ROUTER_WEIGHTS_INPUT and node.input[ROUTER_WEIGHTS_INPUT]: + node.input[ROUTER_WEIGHTS_INPUT] = "" + cleared += 1 + rejected = {} + for ordinal, node in enumerate(nodes): + problems = qmoe_offload_problems(node, constants) + if problems: + rejected[ordinal] = problems + return {"nodes": [node.name for node in nodes], "router_weights_cleared": cleared, "rejected": rejected} + + +def format_node_overrides(caps: dict[str, int]) -> str: + """Formats per-node resident-slot caps as the value of session.moe.offload.node_overrides. + + The runtime splits entries at ';', takes the name up to the first ':' and trims ASCII whitespace around it, so a name + that has one of those cannot be written. + """ + entries = [] + for name, cap in caps.items(): + if not name or name != name.strip(" \t\n\v\f\r") or ";" in name or ":" in name: + raise ValueError(f"node name {name!r} cannot appear in session.moe.offload.node_overrides") + if not 0 <= cap <= 2**31 - 1: + raise ValueError(f"resident slot cap {cap} of node {name!r} is outside 0..{2**31 - 1}") + entries.append(f"{name}:max_resident_slots={cap}") + return ";".join(entries) + + +def offload_session_options( + mode: str, + placement: str, + device_expert_bytes: str | None = None, + host_expert_bytes: str | None = None, + popularity_file: str | None = None, + node_overrides: str | None = None, +) -> dict[str, str]: + """The session.moe.offload.* entries (all values are strings) that configure the runtime's QMoE offload.""" + if mode not in OFFLOAD_MODES: + raise ValueError(f"offload mode {mode!r} is not one of {OFFLOAD_MODES}") + if placement not in OFFLOAD_PLACEMENTS: + raise ValueError(f"offload placement {placement!r} is not one of {OFFLOAD_PLACEMENTS}") + options = {"mode": mode, "placement": placement, "cpu_execution": "auto"} + optional = { + "device_expert_bytes": device_expert_bytes, + "host_expert_bytes": host_expert_bytes, + "popularity_file": popularity_file, + "node_overrides": node_overrides, + } + options.update({key: value for key, value in optional.items() if value}) + return {OFFLOAD_OPTION_PREFIX + key: value for key, value in options.items()} + + +def budget_bytes(text: str) -> str: + """argparse type of the byte budgets: "auto" or a non-negative integer, as the runtime parses them.""" + if text.lower() == "auto": + return "auto" + if text.isascii() and text.isdigit() and int(text) < 2**64: + return str(int(text)) + raise argparse.ArgumentTypeError(f"{text!r} is not 'auto' or a number of bytes between 0 and {2**64 - 1}") def build_multidraft_target_onnx(in_onnx: Path, out_onnx: Path, verify_rows: int = 3) -> dict: @@ -558,6 +720,112 @@ def layer_of(name: str, suffix: str) -> str: } +def check_pool_bindings(pool_ids: list[int], sessions: list[tuple[str, str, int]], where: str) -> None: + """Raises ValueError unless every pool id is defined once and every session binding names a defined pool.""" + duplicated = sorted(pid for pid, count in collections.Counter(pool_ids).items() if count > 1) + if duplicated: + raise ValueError(f"{where}: pool ids {duplicated} are defined more than once") + defined = set(pool_ids) + dangling = [(sess, name, pid) for sess, name, pid in sessions if pid not in defined] + if dangling: + shown = ", ".join(f"session {sess} {name} {pid}" for sess, name, pid in dangling[:5]) + raise ValueError(f"{where}: {len(dangling)} session bindings name a pool that is not defined, e.g. {shown}") + + +def plan_aligned_copy(lengths: list[tuple[int, int]], alignment: int = 4096) -> tuple[dict[int, int], int]: + """Packs ranges back to back, each starting on an `alignment` boundary. Takes (key, length) pairs in order and + returns ({key: destination offset}, size of the packed file).""" + offsets: dict[int, int] = {} + cursor = 0 + for key, length in lengths: + cursor = -(-cursor // alignment) * alignment + offsets[key] = cursor + cursor += length + return offsets, cursor + + +def check_file_ranges(ranges: list[tuple[int, Path, int, int, int]]) -> None: + """Raises ValueError naming the pool whose (pool id, source file, source offset, _, length) range is not in its file.""" + for pid, src, src_offset, _dst_offset, length in ranges: + if not src.is_file(): + raise ValueError(f"pool {pid}: source file {src} does not exist") + size = src.stat().st_size + if src_offset < 0 or length < 0 or src_offset + length > size: + raise ValueError(f"pool {pid}: bytes {src_offset}..{src_offset + length} lie outside {src} ({size} bytes)") + + +def copy_file_ranges(dst: Path, ranges: list[tuple[int, Path, int, int, int]]) -> None: + """Writes `dst` from (pool id, source file, source offset, destination offset, length) ranges. + + Every range is checked against its source file before `dst` is created, so an out-of-bounds pool fails with its id + instead of leaving a short file. + """ + check_file_ranges(ranges) + with open(dst, "wb") as fout: + for _pid, src, src_offset, dst_offset, length in ranges: + with open(src, "rb") as fin: + remaining = length + while remaining > 0: + copied = os.copy_file_range( + fin.fileno(), + fout.fileno(), + min(remaining, 1 << 26), + offset_src=src_offset + length - remaining, + offset_dst=dst_offset + length - remaining, + ) + if copied <= 0: + raise RuntimeError(f"os.copy_file_range stopped early copying {src}") + remaining -= copied + fout.truncate(max((dst_offset + length for _pid, _src, _off, dst_offset, length in ranges), default=0)) + + +def relocate_unaligned_mtp_pools( + source_pkg: Path, out_pkg: Path, pool_records: list[list[str]], mtp_pool_ids: set[int], placement: str +) -> dict[int, str]: + """Moves every MTP MoE pool that starts at a non-16-byte-aligned file offset into one aligned file. + + `pool_records` holds the split `pool`/`host_pool` lines of weights_pool.txt. The ids are whatever the source package + uses: each moved pool gets a record derived from its own line, pointing at its range of the aligned file, and the + result maps pool id -> that record. The file sits next to the source package and is named after the layout, so a + changed layout never reuses a stale copy; the output package hard-links it. + """ + unaligned = sorted( + (parts for parts in pool_records if int(parts[1]) in mtp_pool_ids and int(parts[3]) % 16 != 0), + key=lambda parts: int(parts[1]), + ) + if not unaligned: + return {} + offsets, total_size = plan_aligned_copy([(int(parts[1]), int(parts[4])) for parts in unaligned]) + ranges = [ + (int(parts[1]), source_pkg / parts[6 + int(parts[5])], int(parts[3]), offsets[int(parts[1])], int(parts[4])) + for parts in unaligned + ] + check_file_ranges(ranges) + # The file name identifies the layout and the source files, so a changed package never reuses a stale copy. + layout = [] + for (_pid, src, offset, dst_offset, length), parts in zip(ranges, unaligned, strict=True): + stat = src.stat() + layout.append([parts[6 + int(parts[5])], offset, length, dst_offset, stat.st_size, stat.st_mtime_ns]) + digest = hashlib.sha256(json.dumps(sorted(layout)).encode()).hexdigest()[:12] + shared_aligned = source_pkg.parent / f"{source_pkg.name}_mtp_moe_aligned-{digest}.bin" + if not shared_aligned.exists() or shared_aligned.stat().st_size != total_size: + tmp_aligned = shared_aligned.with_suffix(".tmp") + try: + copy_file_ranges(tmp_aligned, ranges) + except BaseException: + tmp_aligned.unlink(missing_ok=True) + raise + tmp_aligned.replace(shared_aligned) + link_file(shared_aligned, out_pkg / ALIGNED_MTP_RELATIVE, copy=False) + records = {} + for parts in unaligned: + pid = int(parts[1]) + rank = int(parts[5]) + head = ["pool", parts[1], parts[2], str(offsets[pid]), parts[4], *parts[5 : 6 + rank]] + records[pid] = " ".join([*head, ALIGNED_MTP_RELATIVE, placement]) + return records + + def build_offload_weights_pool( source_pkg: Path, out_pkg: Path, @@ -593,40 +861,15 @@ def build_offload_weights_pool( dense_bound_names.add(name) session_lines.append((sess, name, pid)) - aligned_mtp_rel = "target/head_data/source_weights/mtp_moe_aligned.bin" - need_aligned_mtp_copy = False - for line in src_pool[1:]: - parts = line.strip().split() - if parts and parts[0] in ("pool", "host_pool") and int(parts[1]) in mtp_moe_pool_ids: - if int(parts[3]) % 16 != 0: - need_aligned_mtp_copy = True - - if mtp_moe_placement != "device" and need_aligned_mtp_copy: - shared_aligned = source_pkg.parent / f"{source_pkg.name}_mtp_moe_aligned.bin" - src_st = source_pkg / "target/head_data/source_weights/model-bf16-00011.safetensors" - expected_size = 3355443200 + 1677721600 - if not shared_aligned.exists() or shared_aligned.stat().st_size != expected_size: - tmp_aligned = shared_aligned.with_suffix(".tmp") - with open(src_st, "rb") as fin, open(tmp_aligned, "wb") as fout: - in_fd = fin.fileno() - out_fd = fout.fileno() - for src_off, dst_off, total_len in [ - (7243385096, 0, 3355443200), - (5565663496, 3355443200, 1677721600), - ]: - rem = total_len - r_off = src_off - w_off = dst_off - while rem > 0: - chunk = min(rem, 1 << 26) - copied = os.copy_file_range(in_fd, out_fd, chunk, offset_src=r_off, offset_dst=w_off) - if copied <= 0: - raise RuntimeError("os.copy_file_range returned 0") - r_off += copied - w_off += copied - rem -= copied - tmp_aligned.replace(shared_aligned) - link_file(shared_aligned, out_pkg / aligned_mtp_rel, copy=False) + source_records = [ + parts for parts in (line.split() for line in src_pool[1:]) if parts and parts[0] in ("pool", "host_pool") + ] + check_pool_bindings([int(parts[1]) for parts in source_records], session_lines, "the source weights_pool.txt") + aligned_mtp_records = {} + if mtp_moe_placement != "device": + aligned_mtp_records = relocate_unaligned_mtp_pools( + source_pkg, out_pkg, source_records, mtp_moe_pool_ids, mtp_moe_placement + ) for line in src_pool[1:]: s = line.strip() @@ -649,20 +892,13 @@ def build_offload_weights_pool( else: pool_lines.append(base_s) elif pid in mtp_moe_pool_ids and mtp_moe_placement != "device": - if not need_aligned_mtp_copy: - pool_lines.append(f"{base_s} {mtp_moe_placement}") - elif pid == 607: - pool_lines.append(f"pool 607 16 0 3355443200 3 512 1280 2560 {aligned_mtp_rel} {mtp_moe_placement}") - elif pid == 608: - pool_lines.append( - f"pool 608 16 3355443200 1677721600 3 512 2560 640 {aligned_mtp_rel} {mtp_moe_placement}" - ) + pool_lines.append(aligned_mtp_records.get(pid, f"{base_s} {mtp_moe_placement}")) elif default_placement != "device": pool_lines.append(f"{base_s} {default_placement}") else: pool_lines.append(base_s) - next_pool_id = len(pool_lines) + next_pool_id = max((int(parts[1]) for parts in source_records), default=-1) + 1 target_bound_names = {name for sess, name, _ in session_lines if sess == "target"} flat_root = (source_pkg / "model.onnx").exists() target_onnx_src = (source_pkg / "model.onnx") if flat_root else (source_pkg / "target/model.onnx") @@ -709,9 +945,10 @@ def build_offload_weights_pool( out_lines = ["qwen38_weights_pool 1", *pool_lines] for sess, name, pid in session_lines: out_lines.append(f"session {sess} {name} {pid}") + check_pool_bindings([int(line.split()[1]) for line in pool_lines], session_lines, "the rewritten weights_pool.txt") (out_pkg / "weights_pool.txt").write_text("\n".join(out_lines) + "\n") return { - "total_pool_entries": next_pool_id, + "total_pool_entries": len(pool_lines), "cpu_qmoe_entries": len(qmoe_pool_ids), "added_dense_entries": added_dense_entries, "dense_placement": dense_placement, @@ -720,14 +957,106 @@ def build_offload_weights_pool( } -def main() -> None: +def check_layer_numbering(model: onnx.ModelProto, not_offloaded: frozenset[int] = frozenset()) -> None: + """Refuses a graph whose QMoE nodes the runtime would number differently from their position in the graph. + + The expert profile ranks experts by layer number, and the runtime matches them to nodes by the layer it reads from a + node's name (digits after "layers.", "layer." or "blocks.", else the same in the name of its fc1 weights input). A node + without one is numbered in registration order, counting only the nodes the runtime offloads. + """ + registered: list[int] = [] + for ordinal, node in enumerate(qmoe_nodes(model)): + if ordinal in not_offloaded: + continue + number = parse_layer_index(node.name) + if number is None and len(node.input) > 2: + number = parse_layer_index(node.input[2]) + if number is None: + number = len(registered) + while number in registered: + number += 1 + if number != ordinal: + raise ValueError( + f"QMoE node {node_label(node.name, ordinal)} is number {ordinal} in the graph but the runtime " + f"numbers it layer {number}" + ) + registered.append(number) + + +def qmoe_expert_counts(model: onnx.ModelProto) -> set[int]: + """The number of experts (leading dimension of the fc1 weights) of the QMoE nodes.""" + initializers = {init.name: init for init in model.graph.initializer} + counts = set() + for node in qmoe_nodes(model): + weights = initializers.get(node.input[2]) if len(node.input) > 2 else None + if weights is not None and len(weights.dims) == 3: + counts.add(weights.dims[0]) + return counts + + +def plan_resident_caps( + args: argparse.Namespace, model: onnx.ModelProto, node_count: int, not_offloaded: frozenset[int] = frozenset() +) -> dict[int, int]: + """The resident-slot caps asked for on the command line, keyed by position in the graph's QMoE nodes. + + `--target-resident-slots N` caps every layer at N experts. With `--expert-profile` and a slot count, the profile + decides how the slots (`--total-resident-slots`, else N per layer) are spread over the layers, and each layer is capped + at its share; the runtime then keeps the hottest experts of the profile resident up to that cap. Without a slot count + there is no cap and the runtime plans the residency from its byte budgets. + """ + if args.expert_profile is None or (args.total_resident_slots is None and args.target_resident_slots is None): + return ( + dict.fromkeys(range(node_count), args.target_resident_slots) + if args.target_resident_slots is not None + else {} + ) + counts = qmoe_expert_counts(model) + if len(counts) != 1: + raise ValueError(f"--expert-profile needs QMoE nodes with one expert count, the graph has {sorted(counts)}") + check_layer_numbering(model, not_offloaded) + total = args.total_resident_slots + if total is None: + total = args.target_resident_slots * node_count + try: + plan = load_expert_profile_plan( + args.expert_profile.resolve(), + num_layers=node_count, + num_experts=counts.pop(), + total_slots=total, + min_slots_per_layer=4 if args.min_slots_per_layer is None else args.min_slots_per_layer, + uniform=args.uniform_profile_slots, + ) + except ValueError as error: + raise ValueError(f"--expert-profile: {error}") from error + return {layer: len(experts) for layer, experts in plan.items()} + + +def check_slot_arguments(parser: argparse.ArgumentParser, args: argparse.Namespace) -> None: + for flag in ("target_resident_slots", "total_resident_slots", "min_slots_per_layer"): + if getattr(args, flag) is not None and getattr(args, flag) < 0: + parser.error(f"--{flag.replace('_', '-')} must not be negative") + if args.expert_profile is None: + for flag, given in ( + ("--total-resident-slots", args.total_resident_slots is not None), + ("--min-slots-per-layer", args.min_slots_per_layer is not None), + ("--uniform-profile-slots", args.uniform_profile_slots), + ): + if given: + parser.error(f"{flag} needs --expert-profile") + elif (args.uniform_profile_slots or args.min_slots_per_layer is not None) and ( + args.total_resident_slots is None and args.target_resident_slots is None + ): + parser.error( + "--uniform-profile-slots and --min-slots-per-layer need --total-resident-slots or --target-resident-slots" + ) + if args.uniform_profile_slots and args.min_slots_per_layer is not None: + parser.error("--min-slots-per-layer does not apply with --uniform-profile-slots") + + +def build_parser() -> argparse.ArgumentParser: parser = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) - parser.add_argument( - "--migrate-shared-onnx", type=Path, default=None, help="Input ONNX for standalone v3 shared migration" - ) - parser.add_argument("--out-onnx", type=Path, default=None, help="Output ONNX for standalone v3 shared migration") - parser.add_argument("--strip-phase", action="store_true") - parser.add_argument("--strip-stats", action="store_true") + for flag in REMOVED_OPTIONS: + parser.add_argument(flag, nargs="?", const="", default=None, help=argparse.SUPPRESS) parser.add_argument("--from-package", type=Path, default=None) parser.add_argument("--out", type=Path, default=None) @@ -741,42 +1070,161 @@ def main() -> None: ) parser.add_argument("--stop-token-id", type=int, default=None) parser.add_argument("--prefill-chunk", type=int, default=512) - parser.add_argument("--target-resident-slots", type=int, default=None) - parser.add_argument("--target-spare-slots", type=int, default=None) - parser.add_argument("--expert-profile", type=Path, default=None, help="Path to STRP v1 expert-profile.bin") - parser.add_argument("--total-resident-slots", type=int, default=None, help="Total resident slots across 48 layers") - parser.add_argument("--min-slots-per-layer", type=int, default=4) - parser.add_argument("--uniform-profile-slots", action="store_true") - parser.add_argument("--dense-stream-min-rows", type=int, default=1) - parser.add_argument("--dense-stream-wave", type=int, default=64) + parser.add_argument( + "--offload-mode", + choices=OFFLOAD_MODES, + default="required", + help="session.moe.offload.mode. required: offload every QMoE node, and fail the build if the runtime would reject " + "one. auto: the runtime offloads what it can and keeps the other nodes on the GPU", + ) + parser.add_argument( + "--offload-placement", + choices=OFFLOAD_PLACEMENTS, + default="adaptive", + help="session.moe.offload.placement. static: a fixed resident set. adaptive: spare slots promote hot experts", + ) + parser.add_argument( + "--device-expert-bytes", + type=budget_bytes, + default=None, + help="session.moe.offload.device_expert_bytes: GPU bytes for resident experts ('auto' or a byte count)", + ) + parser.add_argument( + "--host-expert-bytes", + type=budget_bytes, + default=None, + help="session.moe.offload.host_expert_bytes: host bytes for cached experts ('auto' or a byte count)", + ) + parser.add_argument( + "--target-resident-slots", + type=int, + default=None, + help="Cap on the resident experts of every layer (a max_resident_slots node override). With --expert-profile it is " + "the per-layer average of --total-resident-slots", + ) + parser.add_argument( + "--expert-profile", + type=Path, + default=None, + help="STRP v1 expert-profile.bin; copied next to the decoder graph and set as session.moe.offload.popularity_file", + ) + parser.add_argument( + "--total-resident-slots", + type=int, + default=None, + help="Total resident experts over all layers, spread by --expert-profile (needs --expert-profile)", + ) + parser.add_argument( + "--min-slots-per-layer", + type=int, + default=None, + help="Floor of the per-layer share of --total-resident-slots when spreading by --expert-profile (default 4)", + ) + parser.add_argument( + "--uniform-profile-slots", + action="store_true", + help="Share --total-resident-slots evenly over the layers instead of by popularity", + ) parser.add_argument( "--dense-gpu-qmoe", action="store_true", default=True, - help="Keep stock GPU QMoE in dense/model.onnx using UVA-mapped expert weights", + help="Keep stock GPU QMoE in dense/model.onnx using UVA-mapped expert weights (always on)", ) - parser.add_argument("--dense-offload-qmoe", dest="dense_gpu_qmoe", action="store_false") parser.add_argument("--dense-placement", choices=["device", "host_mapped", "cpu"], default="host_mapped") parser.add_argument("--mtp-moe-placement", choices=["device", "host_mapped"], default="device") parser.add_argument("--embed-placement", choices=["device", "host_mapped"], default="host_mapped") parser.add_argument("--vram-limit-gib", type=float, default=None) parser.add_argument("--copy", action="store_true") - args = parser.parse_args() + return parser - if args.migrate_shared_onnx is not None: - if args.out_onnx is None: - raise SystemExit("--out-onnx is required with --migrate-shared-onnx") - n = migrate_shared_onnx_v3(args.migrate_shared_onnx, args.out_onnx, args.strip_phase, args.strip_stats) - print(f"Migrated {n} QMoEOffload nodes to v3 compact shared schema -> {args.out_onnx}") - return + +def main() -> None: + parser = build_parser() + args = parser.parse_args() + for flag, reason in REMOVED_OPTIONS.items(): + if getattr(args, flag.lstrip("-").replace("-", "_")) is not None: + parser.error(f"{flag} was removed: {reason}") if args.from_package is None or args.out is None or args.plane_lib is None or args.provider_lib is None: raise SystemExit("--from-package, --out, --plane-lib, and --provider-lib are required for package build") + check_slot_arguments(parser, args) source = args.from_package.resolve() - out = args.out.resolve() - if out == source or out.is_relative_to(source) or source.is_relative_to(out): - raise SystemExit("--out must not be equal to or nested inside --from-package") + out = reject_output_overlap( + args.out, + { + "--from-package": args.from_package, + "--plane-lib": args.plane_lib, + "--provider-lib": args.provider_lib, + "--ops-dir": args.ops_dir, + "--expert-profile": args.expert_profile, + }, + ) + flat_root = (source / "model.onnx").exists() + target_src = (source / "model.onnx") if flat_root else (source / "target/model.onnx") + target_out_dir = out if flat_root else (out / "target") + required_files = { + "genai_config.json of --from-package": source / "genai_config.json", + "weights_pool.txt of --from-package": source / "weights_pool.txt", + "the target graph of --from-package": target_src, + "--plane-lib": args.plane_lib, + "--provider-lib": args.provider_lib, + "--expert-profile": args.expert_profile, + } + for label, path in required_files.items(): + if path is not None and not path.is_file(): + raise SystemExit(f"{label} not found: {path}") + + # Everything that can reject the request runs before the output folder is created. + model = onnx.load(str(target_src), load_external_data=False) + graphs = {target_src: model} + if (source / "dense/model.onnx").exists(): + graphs[source / "dense/model.onnx"] = onnx.load(str(source / "dense/model.onnx"), load_external_data=False) + for path, graph_model in graphs.items(): + legacy = sum(node.op_type == "QMoEOffload" for node in graph_model.graph.node) + if legacy: + raise SystemExit( + f"{path} has {legacy} QMoEOffload nodes. The runtime has no such operator, so the graph cannot load: " + "rebuild the package with standard com.microsoft::QMoE nodes" + ) + qmoe = prepare_qmoe_for_offload(model) + names, rejected = qmoe["nodes"], qmoe["rejected"] + if not names: + raise SystemExit(f"{target_src} has no com.microsoft::QMoE node to offload") + if rejected: + first = next(iter(rejected)) + summary = ( + f"{len(rejected)} of {len(names)} QMoE nodes cannot be offloaded by the runtime; " + f"first, {node_label(names[first], first)}: {'; '.join(rejected[first])}" + ) + if args.offload_mode == "required": + raise SystemExit(f"{summary}. Fix the graph, or use --offload-mode auto to keep those nodes on the GPU.") + print(f"warning: {summary}. They stay on the GPU.") + try: + overrides = {} + for ordinal, cap in plan_resident_caps(args, model, len(names), frozenset(rejected)).items(): + # A node the runtime does not offload is not registered with it, and an override naming it is an error. + if ordinal in rejected: + continue + if not names[ordinal] or names[ordinal] in overrides: + raise ValueError( + "per-node resident-slot caps need uniquely named QMoE nodes, " + f"but {node_label(names[ordinal], ordinal)} is unnamed or shares its name" + ) + overrides[names[ordinal]] = cap + node_overrides = format_node_overrides(overrides) + except ValueError as error: + raise SystemExit(f"error: {error}") from error + offload_options = offload_session_options( + args.offload_mode, + args.offload_placement, + args.device_expert_bytes, + args.host_expert_bytes, + EXPERT_PROFILE_NAME if args.expert_profile is not None else None, + node_overrides, + ) + out.mkdir(parents=True, exist_ok=True) linked = link_tree(source, out, args.copy) @@ -786,53 +1234,16 @@ def main() -> None: link_file(args.plane_lib.resolve(), out / "libs/libqwen38_plane.so", copy=True) link_file(args.provider_lib.resolve(), out / "provider/libonnxruntime_providers_cuda.so", copy=True) - profile_plan = None - if args.expert_profile is not None: - total_slots = args.total_resident_slots or ((args.target_resident_slots or 20) * 48) - profile_plan = load_expert_profile_plan( - args.expert_profile.resolve(), - num_layers=48, - num_experts=512, - total_slots=total_slots, - min_slots_per_layer=args.min_slots_per_layer, - uniform=args.uniform_profile_slots, - ) - - target_attrs = { - "resident_slots": args.target_resident_slots, - "spare_slots": args.target_spare_slots, - "stream_rows": 0, - } - dense_attrs = { - "resident_slots": 0, - "spare_slots": 0, - "stream_rows": args.prefill_chunk, - "stream_min_tokens": 1, - "stream_min_rows": args.dense_stream_min_rows, - "stream_wave": args.dense_stream_wave, - } - flat_root = (source / "model.onnx").exists() - target_src = (source / "model.onnx") if flat_root else (source / "target/model.onnx") - target_out_dir = out if flat_root else (out / "target") - target_rewritten = rewrite_nvfp4_qmoe_to_offload( - target_src, - target_out_dir / "model_s2.onnx", - extra_attrs=target_attrs, - per_layer_initial_experts=profile_plan, - ) + save_model_fresh(model, target_out_dir / "model_s2.onnx") build_multidraft_target_onnx(target_out_dir / "model_s2.onnx", target_out_dir / "model_s3.onnx", verify_rows=3) build_multidraft_target_onnx(target_out_dir / "model_s2.onnx", target_out_dir / "model_s4.onnx", verify_rows=4) link_file(target_out_dir / "model_s4.onnx", target_out_dir / "model.onnx", copy=True) + if args.expert_profile is not None: + # A copy: the output may hard-link the source package, and nothing here may write through to it. + link_file(args.expert_profile.resolve(), target_out_dir / EXPERT_PROFILE_NAME, copy=True) if (source / "dense/model.onnx").exists(): - if args.dense_gpu_qmoe: - link_file(source / "dense/model.onnx", out / "dense/model.onnx", args.copy) - dense_rewritten = 0 - else: - dense_rewritten = rewrite_nvfp4_qmoe_to_offload( - source / "dense/model.onnx", out / "dense/model.onnx", extra_attrs=dense_attrs - ) + link_file(source / "dense/model.onnx", out / "dense/model.onnx", args.copy) else: - dense_rewritten = 0 for obsolete_dir in ("target", "head_r2", "dense", "head_kv") if flat_root else ("dense", "head_kv"): if (out / obsolete_dir).exists(): shutil.rmtree(out / obsolete_dir) @@ -850,6 +1261,7 @@ def main() -> None: cuda_opts["arena.extend_strategy"] = "1" if args.vram_limit_gib is not None and args.vram_limit_gib > 0: cuda_opts["gpu_device_mem_limit"] = str(int(args.vram_limit_gib * (1024**3))) + config["model"]["decoder"]["session_options"].update(offload_options) qwen = config["model"]["qwen38_graph"] fold_head_key = "mtp" if "mtp" in qwen["fold"] else ("drafter" if "drafter" in qwen["fold"] else "head_r2") @@ -885,10 +1297,12 @@ def main() -> None: providers = {p.name: sha256(p) for p in sorted((out / "provider").glob("*.so"))} manifest = { "from_package": str(source), - "stack": "v12b-engine-qmoe-offload-v3", + "stack": "v12b-engine-qmoe-offload-v4", "hard_linked_files": linked, - "target_qmoe_offload_nodes": target_rewritten, - "dense_qmoe_offload_nodes": dense_rewritten, + "target_qmoe_nodes": len(names), + "router_weights_cleared": qmoe["router_weights_cleared"], + "qmoe_nodes_not_offloaded": [node_label(names[ordinal], ordinal) for ordinal in sorted(rejected)], + "offload_session_options": offload_options, "pool_stats": pool_stats, "eos_token_id": config["model"]["eos_token_id"], "vram_limit_gib": args.vram_limit_gib, @@ -898,8 +1312,9 @@ def main() -> None: } (out / "package_manifest.json").write_text(json.dumps(manifest, indent=2) + "\n") print( - f"Built offload v3 package: {out} (linked={linked}, target_qmoe={target_rewritten}, " - f"dense_qmoe={dense_rewritten}, pool={pool_stats})" + f"Built offload package: {out} (linked={linked}, target_qmoe={len(names)}, " + f"router_weights_cleared={qmoe['router_weights_cleared']}, not_offloaded={len(rejected)}, " + f"offload={offload_options}, pool={pool_stats})" ) diff --git a/tools/qwen38_graph/build_package_v12d_int4_int8.py b/tools/qwen38_graph/build_package_v12d_int4_int8.py index 1e46c24a9e..9032f7b5b6 100644 --- a/tools/qwen38_graph/build_package_v12d_int4_int8.py +++ b/tools/qwen38_graph/build_package_v12d_int4_int8.py @@ -1709,10 +1709,6 @@ def main() -> int: "libs/libqwen_mtp_head.so", "libs/libqwen_fold.so", ] - if "runtime_profiles" in cfg and "consumer-gpu-24g" in cfg["runtime_profiles"]: - cfg["runtime_profiles"]["consumer-gpu-24g"]["decoder"]["session_options"]["provider_options"][0]["cuda"][ - "qmoe.offload_gpu_layers" - ] = "2" (out / "genai_config.json").write_text(json.dumps(cfg, indent=2) + "\n") manifest = { "variant": "qwen_flash_next_int4_int8_rc2", From 943f077c8beb4a77c6635c2c7c386675670cf5e3 Mon Sep 17 00:00:00 2001 From: Tianlei Wu Date: Sat, 31 Oct 2026 02:20:38 +0000 Subject: [PATCH 4/6] Align offload package builder with the reworked qwen38_flash_graph config surface The qwen38_flash_graph config parser rejects every model.qwen38_graph key it does not know. The offload builder picked the head session as fold.mtp, fold.drafter or fold.head_r2, whichever the source config had, and copied the rest of the section as it was. - Move the section of the source package to the current surface with package_config.migrate_qwen38_graph_config and print one line per change; the head session is fold.head_r2 only. - Refuse a source without fold.head_r2 or prefill, and a --prefill-chunk the engine rejects (not positive, above kv_capacity with prefill.enabled, odd without it), before anything is written to --out. - Test the migration of each removed key, the written config against the keys the engine accepts, the unchanged weights_pool.txt session labels and the chunk limits. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- .../test_qwen38_offload_package_builder.py | 226 +++++++++++++++++- tools/qwen38_graph/README.md | 7 +- .../qwen38_graph/build_package_offload_v3.py | 51 +++- 3 files changed, 272 insertions(+), 12 deletions(-) diff --git a/test/python/builder/test_qwen38_offload_package_builder.py b/test/python/builder/test_qwen38_offload_package_builder.py index b94cd79edc..a8a29a1136 100644 --- a/test/python/builder/test_qwen38_offload_package_builder.py +++ b/test/python/builder/test_qwen38_offload_package_builder.py @@ -34,7 +34,12 @@ simulate_session_creation, validate_popularity_file, ) -from _qwen38_tools_test_utils import load_tools +from _qwen38_tools_test_utils import ( + REMOVED_QWEN38_GRAPH_KEYS, + load_tools, + prefill_chunk_problem, + unknown_qwen38_graph_keys, +) from onnx import TensorProto, helper, numpy_helper LAYERS = 3 @@ -190,8 +195,10 @@ def genai_config(layered): }, }, "qwen38_graph": { - "fold": {"mtp": {"filename": "mtp/model.onnx"}}, + "kv_capacity": 4096, + "fold": {"enabled": True, "head_r2": {"filename": "head_r2/model.onnx"}}, "prefill": { + "enabled": True, "chunk": 1024, "dense": {"filename": "dense/model.onnx", "session_options": {"provider_options": [{"cuda": {}}]}}, }, @@ -209,12 +216,12 @@ def genai_config(layered): """ -def make_package(tmp_path, *, model=None, layered=False, pool_text=DEFAULT_POOL, dense=True): +def make_package(tmp_path, *, model=None, layered=False, pool_text=DEFAULT_POOL, dense=True, config=None): """A source package under ``tmp_path/pkg`` (flat, or with the decoder graph in ``target/``) and the two libraries.""" root = tmp_path / "pkg" decoder_dir = root / "target" if layered else root save_model(decoder_model() if model is None else model, decoder_dir / "model.onnx") - (root / "genai_config.json").write_text(json.dumps(genai_config(layered))) + (root / "genai_config.json").write_text(json.dumps(genai_config(layered) if config is None else config)) (root / "weights_pool.txt").write_text(pool_text) if dense: save_model(dense_model(), root / "dense" / "model.onnx") @@ -1412,12 +1419,14 @@ def test_budgets_and_the_other_options_reach_the_config(offload, tmp_path, monke prefill = config["model"]["qwen38_graph"]["prefill"] assert (prefill["chunk"], prefill["warmup_lengths"]) == (256, [256, 128, 37]) assert ( - config["model"]["qwen38_graph"]["fold"]["mtp"]["session_options"][ + config["model"]["qwen38_graph"]["fold"]["head_r2"]["session_options"][ "session.use_device_allocator_for_initializers" ] == "1" ) - assert not any(key.startswith(PREFIX) for key in config["model"]["qwen38_graph"]["fold"]["mtp"]["session_options"]) + assert not any( + key.startswith(PREFIX) for key in config["model"]["qwen38_graph"]["fold"]["head_r2"]["session_options"] + ) def test_a_package_without_a_dense_graph_builds(offload, tmp_path, monkeypatch, multidraft_calls): @@ -1767,6 +1776,211 @@ def test_the_default_mtp_placement_leaves_the_pools_alone(offload, tmp_path, mon assert not (out / ALIGNED).exists() +# --------------------------------------------------------------------------- +# genai_config.json: only the model.qwen38_graph keys the engine accepts +# --------------------------------------------------------------------------- + +HEAD_SESSION_OPTIONS = { + "log_id": "onnxruntime-genai", + "graph_optimization_level": "ORT_DISABLE_ALL", + "provider_options": [ + { + "cuda": { + "enable_cuda_graph": "1", + "min_num_runs_before_cuda_graph_capture": "4", + "arena_extend_strategy": "kSameAsRequested", + "arena.extend_strategy": "1", + } + } + ], + "session.use_device_allocator_for_initializers": "1", +} + + +def config_with_removed_key(key, layered=False): + """The source config of a package built for an earlier config surface: ``key`` is a path below model.qwen38_graph.""" + config = genai_config(layered) + qwen = config["model"]["qwen38_graph"] + if key in ("fold.mtp", "fold.drafter"): + qwen["fold"][key.removeprefix("fold.")] = qwen["fold"].pop("head_r2") + return config + *parents, leaf = key.split(".") + section = qwen + for part in parents: + section = section.setdefault(part, {}) + section[leaf] = 1 + return config + + +def printed_config_changes(capsys): + return sorted(line for line in capsys.readouterr().out.splitlines() if line.startswith("genai_config:")) + + +@pytest.mark.parametrize("layered", [False, True]) +def test_the_built_config_has_only_keys_the_engine_accepts(offload, tmp_path, monkeypatch, multidraft_calls, layered): + package = make_package(tmp_path, layered=layered) + _out, config, *_ = build(monkeypatch, offload, package, tmp_path, layered=layered) + qwen = config["model"]["qwen38_graph"] + assert unknown_qwen38_graph_keys(qwen) == [] + assert prefill_chunk_problem(qwen) is None + assert qwen["fold"] == { + "enabled": True, + "head_r2": {"filename": "head_r2/model.onnx", "session_options": HEAD_SESSION_OPTIONS}, + } + assert qwen["prefill"]["chunk"] == 512 + + +@pytest.mark.parametrize("key", REMOVED_QWEN38_GRAPH_KEYS) +def test_a_removed_key_of_the_source_config_is_not_carried_into_the_package( + offload, tmp_path, monkeypatch, multidraft_calls, key +): + assert unknown_qwen38_graph_keys(config_with_removed_key(key)["model"]["qwen38_graph"]) != [] + package = make_package(tmp_path, config=config_with_removed_key(key)) + _out, config, *_ = build(monkeypatch, offload, package, tmp_path) + qwen = config["model"]["qwen38_graph"] + assert unknown_qwen38_graph_keys(qwen) == [] + assert qwen["fold"]["head_r2"] == {"filename": "head_r2/model.onnx", "session_options": HEAD_SESSION_OPTIONS} + + +@pytest.mark.parametrize("layered", [False, True]) +@pytest.mark.parametrize("alias", ["mtp", "drafter"]) +def test_a_source_of_the_earlier_config_surface_is_migrated( + offload, tmp_path, monkeypatch, multidraft_calls, capsys, layered, alias +): + config = genai_config(layered) + qwen = config["model"]["qwen38_graph"] + qwen.update( + target_s1_graph_id=1, target_s2_graph_id=2, head_s1_graph_id=3, head_s2_graph_id=4, prefill_chunk_tokens=2048 + ) + qwen["mtp"] = {"enabled": True, "prefill_head": False, "skip_refeed": True} + head = { + **qwen["fold"].pop("head_r2"), + "graph_id": 9, + "session_options": {"custom_ops_libraries": ["libs/earlier.so"]}, + } + qwen["fold"].update({alias: head, "force_reject": False, "target_graph_id": 8}) + package = make_package(tmp_path, layered=layered, config=config) + before = snapshot(package.root) + + _out, built, *_ = build(monkeypatch, offload, package, tmp_path, layered=layered) + + qwen = built["model"]["qwen38_graph"] + assert unknown_qwen38_graph_keys(qwen) == [] + assert qwen["mtp"] == {"enabled": True, "prefill_head": False} + assert qwen["fold"] == { + "enabled": True, + "head_r2": {"filename": "head_r2/model.onnx", "session_options": HEAD_SESSION_OPTIONS}, + } + assert printed_config_changes(capsys) == sorted( + f"genai_config: {change}" + for change in ( + "dropped target_s1_graph_id", + "dropped target_s2_graph_id", + "dropped head_s1_graph_id", + "dropped head_s2_graph_id", + "dropped mtp.skip_refeed", + "dropped fold.force_reject", + "dropped fold.target_graph_id", + f"renamed fold.{alias} to fold.head_r2", + "dropped fold.head_r2.graph_id", + "dropped prefill_chunk_tokens (prefill.chunk is the one chunk setting)", + ) + ) + assert snapshot(package.root) == before + + +def test_a_source_that_is_already_on_the_current_surface_is_not_reported_as_changed( + offload, tmp_path, monkeypatch, multidraft_calls, capsys +): + package = make_package(tmp_path) + build(monkeypatch, offload, package, tmp_path) + assert printed_config_changes(capsys) == [] + + +def test_prefill_chunk_tokens_of_a_source_without_the_dense_prefill_becomes_the_chunk_the_option_replaces( + offload, tmp_path, monkeypatch, multidraft_calls, capsys +): + config = genai_config(False) + qwen = config["model"]["qwen38_graph"] + qwen["prefill"]["enabled"] = False + qwen["prefill_chunk_tokens"] = 4096 + package = make_package(tmp_path, config=config) + _out, built, *_ = build(monkeypatch, offload, package, tmp_path, "--prefill-chunk", "256") + assert printed_config_changes(capsys) == ["genai_config: renamed prefill_chunk_tokens to prefill.chunk"] + qwen = built["model"]["qwen38_graph"] + assert "prefill_chunk_tokens" not in qwen + assert (qwen["prefill"]["chunk"], prefill_chunk_problem(qwen)) == (256, None) + + +def test_the_session_labels_of_the_weights_pool_are_not_config_keys_and_stay( + offload, tmp_path, monkeypatch, multidraft_calls +): + extra = [("mtp", "model.mtp.norm.weight", 0), ("drafter", "model.drafter.norm.weight", 1)] + extra.append(("head_r2", "model.head.norm.weight", 0)) + pool = DEFAULT_POOL + "".join(f"session {label} {name} {pid}\n" for label, name, pid in extra) + package = make_package(tmp_path, pool_text=pool) + out, *_ = build(monkeypatch, offload, package, tmp_path) + _records, bindings = read_pool(out / "weights_pool.txt") + assert [binding for binding in bindings if binding[0] in {label for label, _name, _pid in extra}] == extra + + +@pytest.mark.parametrize( + ("edit", "message"), + [ + (lambda model: model.pop("qwen38_graph"), "has no model.qwen38_graph section"), + (lambda model: model["qwen38_graph"].pop("fold"), r"has no model.qwen38_graph.fold.head_r2 section"), + (lambda model: model["qwen38_graph"]["fold"].pop("head_r2"), r"has no model.qwen38_graph.fold.head_r2 section"), + (lambda model: model["qwen38_graph"].pop("prefill"), "has no model.qwen38_graph.prefill section"), + ], + ids=["no qwen38_graph", "no fold", "no head_r2", "no prefill"], +) +def test_a_source_config_without_what_the_builder_rewrites_is_refused_before_anything_is_written( + offload, tmp_path, monkeypatch, multidraft_calls, edit, message +): + config = genai_config(False) + edit(config["model"]) + package = make_package(tmp_path, config=config) + out = tmp_path / "out" + with pytest.raises(SystemExit, match=message): + run_builder(monkeypatch, offload, package, out) + assert not out.exists() + + +@pytest.mark.parametrize( + ("enabled", "chunk", "message"), + [ + (True, 0, "--prefill-chunk must be positive"), + (True, -2, "--prefill-chunk must be positive"), + (True, 4098, r"--prefill-chunk \(4098\) must not exceed model.qwen38_graph.kv_capacity \(4096\)"), + (False, 0, "--prefill-chunk must be positive"), + (False, 513, r"--prefill-chunk \(513\) must be even"), + ], +) +def test_a_prefill_chunk_the_engine_rejects_is_refused_before_anything_is_written( + offload, tmp_path, monkeypatch, multidraft_calls, enabled, chunk, message +): + config = genai_config(False) + config["model"]["qwen38_graph"]["prefill"]["enabled"] = enabled + package = make_package(tmp_path, config=config) + out = tmp_path / "out" + with pytest.raises(SystemExit, match=message): + run_builder(monkeypatch, offload, package, out, "--prefill-chunk", chunk) + assert not out.exists() + + +@pytest.mark.parametrize(("enabled", "chunk"), [(True, 4096), (True, 513), (True, 1), (False, 2), (False, 4098)]) +def test_a_prefill_chunk_the_engine_accepts_is_written( + offload, tmp_path, monkeypatch, multidraft_calls, enabled, chunk +): + config = genai_config(False) + config["model"]["qwen38_graph"]["prefill"]["enabled"] = enabled + package = make_package(tmp_path, config=config) + _out, built, *_ = build(monkeypatch, offload, package, tmp_path, "--prefill-chunk", chunk) + qwen = built["model"]["qwen38_graph"] + assert qwen["prefill"]["chunk"] == chunk + assert prefill_chunk_problem(qwen) is None + + # --------------------------------------------------------------------------- # Sessions in the installed runtime # --------------------------------------------------------------------------- diff --git a/tools/qwen38_graph/README.md b/tools/qwen38_graph/README.md index e9bc4e6fa6..1934659f1a 100644 --- a/tools/qwen38_graph/README.md +++ b/tools/qwen38_graph/README.md @@ -234,8 +234,11 @@ been loaded by a CUDA build of #33622. (the number after `layers.`, `layer.` or `blocks.` in the node name, or in the name of its fc1 weights, else the registration order); * `genai_config.json` gets the `session.moe.offload.*` options on the decoder, and CUDA-graph capture and `kSameAsRequested` arena - options for the decoder, MTP-head, dense and head-KV sessions; `--vram-limit-gib` becomes `gpu_device_mem_limit`, - `--stop-token-id` replaces `eos_token_id` and `--prefill-chunk` (default 512) sets `prefill.chunk`; + options for the decoder, the MTP head (`fold.head_r2`), dense and head-KV sessions; `--vram-limit-gib` becomes + `gpu_device_mem_limit`, `--stop-token-id` replaces `eos_token_id` and `--prefill-chunk` (default 512) sets `prefill.chunk`, which + the engine accepts up to `kv_capacity` with `prefill.enabled` and as an even number without it (checked before anything is + written). The `model.qwen38_graph` section of the source is migrated to the current config surface first (see "Config surface of + the built packages"), and a source without `fold.head_r2` or `prefill` is refused; * `--plane-lib` and `--provider-lib` are copied to `libs/` and `provider/`; `--ops-dir` adds its `libqwen_*.so` files to `libs/`. `--out` must not be equal to, nested in or contain `--from-package`, `--plane-lib`, `--provider-lib`, `--ops-dir` or diff --git a/tools/qwen38_graph/build_package_offload_v3.py b/tools/qwen38_graph/build_package_offload_v3.py index 7b891626a9..22b8415bbd 100644 --- a/tools/qwen38_graph/build_package_offload_v3.py +++ b/tools/qwen38_graph/build_package_offload_v3.py @@ -17,6 +17,8 @@ UVA placement for 24 GiB VRAM targets. - Configures genai_config.json for 24 GiB or 32 GiB VRAM operation and sets the decoder's session.moe.offload.* options: mode, placement, the byte budgets, the expert-popularity file and per-layer resident-slot caps. + The model.qwen38_graph section of the source is moved to the config surface the engine accepts (see + package_config.py): the MTP head session is fold.head_r2, and there are no CUDA graph ids. """ from __future__ import annotations @@ -31,6 +33,7 @@ import onnx from onnx import helper +from package_config import migrate_qwen38_graph_config from package_paths import reject_output_overlap OFFLOAD_MODES = ("required", "auto") @@ -1069,7 +1072,13 @@ def build_parser() -> argparse.ArgumentParser: help="Directory containing rebased libqwen_*.so custom op libraries", ) parser.add_argument("--stop-token-id", type=int, default=None) - parser.add_argument("--prefill-chunk", type=int, default=512) + parser.add_argument( + "--prefill-chunk", + type=int, + default=512, + help="model.qwen38_graph.prefill.chunk: at most the kv_capacity of the source package when it has prefill.enabled, " + "an even number otherwise", + ) parser.add_argument( "--offload-mode", choices=OFFLOAD_MODES, @@ -1139,6 +1148,38 @@ def build_parser() -> argparse.ArgumentParser: return parser +def migrate_source_config(config: dict, prefill_chunk: int) -> list[str]: + """Moves `model.qwen38_graph` of the source config to the current config surface and returns what it changed. + + The builder rewrites `fold.head_r2.session_options` and `prefill.chunk`, so the source needs both sections. The engine + accepts a chunk of 1 to kv_capacity with `prefill.enabled`, and an even chunk without it. + """ + qwen = config["model"].get("qwen38_graph") + if not isinstance(qwen, dict): + raise SystemExit("genai_config.json of --from-package has no model.qwen38_graph section") + changes = migrate_qwen38_graph_config(qwen) + fold = qwen.get("fold") + if not isinstance(fold, dict) or not isinstance(fold.get("head_r2"), dict): + raise SystemExit("genai_config.json of --from-package has no model.qwen38_graph.fold.head_r2 section") + prefill = qwen.get("prefill") + if not isinstance(prefill, dict): + raise SystemExit("genai_config.json of --from-package has no model.qwen38_graph.prefill section") + if prefill_chunk <= 0: + raise SystemExit("--prefill-chunk must be positive") + kv_capacity = qwen.get("kv_capacity") + if prefill.get("enabled"): + if isinstance(kv_capacity, int) and prefill_chunk > kv_capacity: + raise SystemExit( + f"--prefill-chunk ({prefill_chunk}) must not exceed model.qwen38_graph.kv_capacity ({kv_capacity})" + ) + elif prefill_chunk % 2: + raise SystemExit( + f"--prefill-chunk ({prefill_chunk}) must be even: prefill.enabled is false in --from-package, " + "so the decode graphs ingest the prompt" + ) + return changes + + def main() -> None: parser = build_parser() args = parser.parse_args() @@ -1176,6 +1217,10 @@ def main() -> None: if path is not None and not path.is_file(): raise SystemExit(f"{label} not found: {path}") + config = json.loads((source / "genai_config.json").read_text()) + for change in migrate_source_config(config, args.prefill_chunk): + print(f"genai_config: {change}") + # Everything that can reject the request runs before the output folder is created. model = onnx.load(str(target_src), load_external_data=False) graphs = {target_src: model} @@ -1251,7 +1296,6 @@ def main() -> None: source, out, args.dense_placement, args.mtp_moe_placement, args.embed_placement ) - config = json.loads((source / "genai_config.json").read_text()) if args.stop_token_id is not None: config["model"]["eos_token_id"] = [args.stop_token_id] @@ -1264,8 +1308,7 @@ def main() -> None: config["model"]["decoder"]["session_options"].update(offload_options) qwen = config["model"]["qwen38_graph"] - fold_head_key = "mtp" if "mtp" in qwen["fold"] else ("drafter" if "drafter" in qwen["fold"] else "head_r2") - qwen["fold"][fold_head_key]["session_options"] = { + qwen["fold"]["head_r2"]["session_options"] = { "log_id": "onnxruntime-genai", "graph_optimization_level": "ORT_DISABLE_ALL", "provider_options": [ From c744cd9bc4845732742f85429267c8bdeddb0d81 Mon Sep 17 00:00:00 2001 From: Tianlei Wu Date: Sat, 31 Oct 2026 02:25:47 +0000 Subject: [PATCH 5/6] Address PR #2821 review feedback on code scanning alerts in the offload tests The build helper returned (out, *read_output(...)), a tuple whose length code scanning cannot see, so every call site that unpacks its five results was reported as a mismatched multiple assignment. It now returns the five values explicitly. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- test/python/builder/test_qwen38_offload_package_builder.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/test/python/builder/test_qwen38_offload_package_builder.py b/test/python/builder/test_qwen38_offload_package_builder.py index a8a29a1136..08daaf2927 100644 --- a/test/python/builder/test_qwen38_offload_package_builder.py +++ b/test/python/builder/test_qwen38_offload_package_builder.py @@ -1322,9 +1322,11 @@ def caps_of(entries): def build(monkeypatch, offload, package, tmp_path, *extra, layered=False): + """Runs the builder and returns (output folder, config, offload entries, decoder graph, decoder folder).""" out = tmp_path / "out" run_builder(monkeypatch, offload, package, out, *extra) - return (out, *read_output(out, layered)) + config, entries, model, decoder_dir = read_output(out, layered) + return out, config, entries, model, decoder_dir @pytest.mark.parametrize("layered", [False, True]) From eed7415db08ab155653611e1b2be3111ca0c6216 Mon Sep 17 00:00:00 2001 From: Tianlei Wu Date: Sat, 31 Oct 2026 03:41:51 +0000 Subject: [PATCH 6/6] Fix the offload builder tests on the released onnxruntime and on Windows The released onnxruntime 1.31.0 (the CPU wheel of CI) reports a QMoEOffload node as Fail (not a registered function/op); InvalidGraph is what a build that registers its own QMoEOffload schema reports. The test of the graphs that the previous builder emitted pinned InvalidGraph, so it only passed on such a build. It now accepts the failure categories of a failed load (Fail, InvalidGraph, NotImplemented) and still requires the message to name the operator or the attribute. The refusal of a legacy QMoEOffload dense graph names the path with the separator of the OS, and the test matched a forward slash. The expected path is now built with pathlib. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- .../builder/test_qwen38_offload_package_builder.py | 9 +++++++-- 1 file changed, 7 insertions(+), 2 deletions(-) diff --git a/test/python/builder/test_qwen38_offload_package_builder.py b/test/python/builder/test_qwen38_offload_package_builder.py index 08daaf2927..52fe0e99f5 100644 --- a/test/python/builder/test_qwen38_offload_package_builder.py +++ b/test/python/builder/test_qwen38_offload_package_builder.py @@ -15,6 +15,7 @@ import json import os import random +import re import shutil import struct import sys @@ -1634,7 +1635,8 @@ def test_a_legacy_qmoe_offload_dense_graph_is_refused(offload, tmp_path, monkeyp dense = dense_model() dense.graph.node[0].op_type = "QMoEOffload" save_model(dense, package.root / "dense/model.onnx") - with pytest.raises(SystemExit, match=r"dense/model\.onnx has 1 QMoEOffload nodes"): + dense_graph = re.escape(str(Path("dense") / "model.onnx")) + with pytest.raises(SystemExit, match=rf"{dense_graph} has 1 QMoEOffload nodes"): run_builder(monkeypatch, offload, package, tmp_path / "out") assert not (tmp_path / "out").exists() @@ -2097,6 +2099,9 @@ def test_the_built_package_creates_a_session_in_the_installed_runtime( ids=["QMoEOffload operator", "layer_index attribute", "resident_slots attribute"], ) def test_what_the_previous_builder_emitted_does_not_load(ort, graph, message): + # The status category of a failed load depends on the runtime build: the released runtime reports an unregistered operator as + # Fail and an unknown attribute as InvalidGraph, while a build that registers its own QMoEOffload schema can report InvalidGraph + # or NotImplemented. The behaviour that matters is that session creation fails and names the operator or the attribute. with pytest.raises(Exception, match=message) as excinfo: create_session(ort, graph()) - assert type(excinfo.value).__name__ == "InvalidGraph" + assert type(excinfo.value).__name__ in {"Fail", "InvalidGraph", "NotImplemented"}