Add Qwen3.8-Flash Engine qualification and benchmark tools by tianleiwu · Pull Request #2729 · microsoft/onnxruntime-genai

Tianlei Wu (tianleiwu) · 2026-10-13T22:53:08Z

Independent PR, based on main. The scripts are Python-only, so the PR does not conflict with the runtime stack (#2726 to #2824, #2825, #2784), but they drive the qwen38_flash_graph model of that stack and need it to run, and the README describes the final state of the stack. #2820 (package builders) is stacked on this PR, and #2821 (the offload package builder) on that one. Series map: #2726.

This PR was re-cut. It used to carry the package builders as well (and to be stacked on #2728); they are in #2820 and #2821 now, so that each PR has a reviewable size. The benchmark tweaks that #2824 used to carry are here.

Summary

Adds the scripts that qualify a qwen38_flash_graph package through the public Engine API (og.Engine / og.Request), and a README of the backend (build, external artifacts, package layout, the SpeculativePrefixCommit contract of the rc4/rc5 DFlash2 packages, usage, scripts, limits).

The scripts depend only on this tree, on the package they are given and on the external artifacts the README lists. The research helper modules that the MMLU-Pro and LiveCodeBench drivers import (eval_lib, eval_mcq_genai) are not in this repository: they are found through $QWEN38_EVAL_PATHS, and the prompts, samples, tests, checkpoint and grader are command-line arguments without defaults.

Summary of Changes

File	Change
`tools/qwen38_graph/README.md`	Backend overview: build (the plane is compiled into `onnxruntime-genai-cuda`; `-DENABLE_QWEN38_PLANE=ON` additionally builds a standalone library), the external artifacts a package needs, the package layout, the v12b stack, the `SpeculativePrefixCommit` contract (graph I/O and bindings), Engine usage, the scripts, the qualification pass and the limitations.
`tools/qwen38_graph/qwen38_engine_util.py`, `smoke.py`	Shared helpers (registering the plugin CUDA EP and loading a package, running a turn, the natural-text pools from `--text-pool` / `$QWEN38_TEXT_POOL` and `--prose-pool` / `$QWEN38_PROSE_POOL`) and a minimal usage example. The pool is repeated when it is shorter than the requested context.
`tools/qwen38_graph/bench_engine_mtp.py`, `bench_ttft_engine.py`	vLLM-protocol decode benchmark (exact prompt lengths, 256 new tokens, EOS ignored, one warm-up and three measured runs; the acceptance histogram has any number of drafts per round; `QWEN38_SKIP_LONG_WARMUP=1` skips the warm-up at 32768 tokens and above) and time to first token.
`tools/qwen38_graph/calibrate_qwen38_kv.py`	Calibrates the symmetric per-channel INT8 and FP8 KV-cache scales for the 12 target sparse-attention layers and the MTP-head layer, with the partial-rotary (`rotary_dim=64`) pair envelope.
`tools/qwen38_graph/engine_gate.py`	Greedy token-identity gate against a `--reference` stream (exact for `--prefill steps`; prompt length and a minimum prefix agreement for `--prefill fast`). With `--force-reject` every draft is replaced, so MTP must equal no-MTP: the self-contained gate.
`tools/qwen38_graph/engine_behaviours.py`	27 checks of the Engine contract on the backend: stop tokens, limits, stop strings, cancellation during prefill and decode, queueing, continuation, rewind and replay, rejected sampling options, session limits, speculative statistics, one Engine per model.
`tools/qwen38_graph/engine_prompt_lengths.py`	Prompt-length and dense chunk-boundary check (1 token, odd and even lengths, 8191 / 8192 / 8193, two and three chunks), `steps` against `fast` prefill.
`tools/qwen38_graph/mmlu_engine_eval.py`, `lcb_engine_eval.py`	MMLU-Pro (800 items) and LiveCodeBench v6 (175 problems, sandbox grading in the `lcb-grader` container at 16K and 32K caps, paired exact McNemar test) through the Engine API. They import the research helper modules from `$QWEN38_EVAL_PATHS`.
`tools/qwen38_graph/run_with_ram_budget.py`, `test_ram_budget.py`	Runs a command with a bounded expert cache and optional cgroup RAM enforcement (for the offload packages), with two unit tests of the budget plan.
`tools/qwen38_graph/test_plane_kernels.py`	Kernel-level checks of the fast-prefill plane entry points (batched PLE rows, pooled index keys, GDN state transpose) against the research implementations.

Testing

ruff check and ruff format --check pass on the files of this PR; every script byte-compiles; python -m unittest test_ram_budget: 2 tests, OK.
End-to-end on the H200 (CUDA 13.0, sm_90) with the Qwen3.8-Flash-Next rc5 INT4 DFlash2 (S=8) package on the merged tip of the runtime series: smoke.py decodes 75 tokens at 385.9 tok/s (first token 0.050 s, EOS), and engine_behaviours.py --text-pool <the repository's docs, about 500 KB> ends with BEHAVIOURS PASS (27 PASS, 0 FAIL).
Earlier measurements with these scripts on the stack this series was cut from (attributed to that head, not re-run here): MMLU-Pro 800 666/800 (83.25%) at a median 361.14 tok/s with the DFlash2 S=8 package; engine_behaviours.py 27/27 PASS.
Not run: lcb_engine_eval.py (needs the Docker grader and the LiveCodeBench data), test_plane_kernels.py (needs the research environment), the cgroup enforcement of run_with_ram_budget.py.

Checklist

Tests added/updated (test_ram_budget.py; the rest are drivers that are run against a package)
No breaking changes (new files only)
Documentation updated (tools/qwen38_graph/README.md)

…th auxiliary sessions (#2725) ## Description Adds `session_options.custom_ops_libraries`, a list of custom-op libraries registered in order after the existing single `custom_ops_library`, and exports `InheritSessionOptionsFromParent` so components that create auxiliary sessions (for example MTP/draft heads) can inherit the decoder's providers, provider options, config entries and custom ops. Models whose graphs use several custom-op domains (fused attention, hyper-connection, GEMM and head ops shipped next to the model package) currently cannot name more than one library. This change is generic and has no dependency on any model type; it is the first step of the Qwen3.8-Flash static-graph backend series (see the end of this description), but stands alone and targets `main`. ## Summary of Changes | File | Change | |------|--------| | `src/config.h`, `src/config.cpp` | New `session_options.custom_ops_libraries` (JSON string array; scalar values such as `"custom_ops_libraries": "ops.so"` are rejected during parsing). Overlays append to `custom_ops_libraries` (and `[]` leaves existing entries unchanged). `InheritSessionOptionsFromParent` exposes the existing inheritance helper; `custom_ops_library` and `custom_ops_libraries` are inherited together when the child sets neither (note that `"custom_ops_libraries": []` leaves the vector empty and therefore still inherits from the parent unless `custom_ops_library` is set). | | `src/models/model.cpp` | `custom_ops_library` and then every `custom_ops_libraries` entry go through the same path-traversal and allow-list checks, with duplicate resolved paths registered once. All paths are validated against the allow-list before loading starts; validated libraries are then loaded and registered sequentially in order (so a later load/registration error can still leave earlier libraries registered on `session_options`). | | `test/CMakeLists.txt`, `test/cpp/engine/custom_ops_test_lib/custom_ops_test_lib.cpp`, `test/cpp/engine/custom_ops_libraries_tests.cpp` | Builds two small test custom-op shared libraries (`custom_ops_test_first` and `custom_ops_test_second`) and tests parsing/ordering, overlay rules, rejection of non-string entries and scalar values, inheritance (including `[]` and `CreateMtpDecoderConfig`), allow-list validation, and actual library loading, ordering, and deduplication. | ### Library Location Requirement Every entry in `custom_ops_library` and `custom_ops_libraries` must resolve to an existing file inside one of the allowed directories: 1. The model folder (alongside `genai_config.json`), 2. A registered Execution Provider (EP) library directory, or 3. The current working directory. Previously, an unresolved bare library name in `custom_ops_library` was passed directly to `RegisterCustomOpsLibrary` and could be resolved through the platform loader search path (`LD_LIBRARY_PATH` / `PATH`). Unresolved entries are now rejected before loading so the directory allow-list cannot be bypassed; users should place custom-op libraries in one of the allowed directories above (typically the model folder). ## Testing - CPU-only Release build (`-DUSE_CUDA=OFF`) on `main`: `engine_unit_tests` **929 passed, 0 failed, 2 skipped** (including all 9 tests in `CustomOpsLibrariesConfigTest`, which load real test custom-op shared libraries to verify registration, ordering, and deduplication). - `lintrunner`, `clang-format` (20.1.8) and `git diff --check` are clean. ## Motivation and Context Found while hosting a model that uses seven custom-op libraries. Registering them through one `custom_ops_library` entry was not possible, and auxiliary sessions did not inherit the decoder's custom ops, so head graphs failed to load without repeating the whole session configuration. This is the first PR of a stack that adds an Engine-hosted CUDA-graph MTP backend for Qwen3.8-Flash-Next (greedy only, one resident request, Linux/CUDA, no prefix cache). The follow-ups are drafts stacked on `kvaishnavi/qwen-38-flash-run`, so this commit is also carried as the first commit of the next PR (#2726) until this one merges. The stack: #2726 (CUDA data plane library) -> #2727 (model type and device engine) -> #2728 (Engine hosting) -> #2729 (qualification tools). ## Checklist - [x] Tests added/updated - [x] No breaking changes - [ ] Documentation updated (the option is described in the PR; no dedicated config documentation exists for `session_options` yet) --------- Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

- Reject duplicate custom-op library basenames, nested --out paths inside input directories, and non-empty --out directories in build_package.py and build_package_engine_v12b.py. - Enforce a minimum prefix agreement floor in engine_gate.py --prefill fast mode and compare delivered token prefixes plus EOS termination across all EOS_IDS without fabricating EOS_IDS[0]. - Add v12c-opt and v12d builder options to build_package_engine_v12b.py: --lm-head-quant {none,fp8,head_fp8}, --hc-quant {none,fp8}, --host-map-embed, --prefill-chunk, --fuse-dense-ops, MTP head HC block merging, and rotary/attn_consts/head_kv initializer deduplication. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

Tianlei Wu (tianleiwu) · 2026-10-30T05:26:21Z

Re-cut and force-pushed (52635f2d6 -> 8c2f4a1c8). The PR is now based on main (it was stacked on #2728) and carries the qualification and benchmark tools only. The package builders moved to #2820 and #2821, which are stacked on this PR as a second GitHub stack. The head includes small fixes for CodeQL findings in the Python tools: an explicit return after parser.error, a with-block for a file, no sys.exit on a function that returns None, and a comment on an intentionally empty except.

Add tools/qwen38_graph with the first half of the scripts that drive a qwen38_flash_graph package through the public Engine API, and a README of the backend (build, external artifacts, package layout, the v12b stack, usage, limits). - qwen38_engine_util.py: CUDA EP registration, model loading with config overlays, natural-text prompt pools supplied by the caller (no hard-coded locations), and the run_turn/pump helpers the other scripts share. Registers the DLL directories on Windows. - smoke.py: minimal usage example. - bench_engine_mtp.py: vLLM-protocol decode benchmark with per-position acceptance histograms. - bench_ttft_engine.py: time to first token at several prompt lengths. - calibrate_qwen38_kv.py: per-channel INT8/FP8 KV-cache scale calibration. - run_with_ram_budget.py, test_ram_budget.py: bound the host expert cache, and optionally the whole command through a systemd scope, for QMoE-offload runs. The scripts depend only on this tree, on the package they are given and on the external artifacts the README lists. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

Add the scripts that qualify a qwen38_flash_graph package through the Engine API, and document them in the README. - engine_gate.py: greedy token-identity gate against a reference stream; a --no-mtp or --force-reject recording makes the gate MTP == no-MTP. The steps prefill must be identical. The fast prefill is tolerance level, so it must agree on the prompt length, the first token and an identical prefix of at least --min-fast-prefix tokens. Streams are compared as the delivered tokens plus the termination reason across all stop ids, without assuming the first stop id. - engine_behaviours.py: stop tokens, limits, stop strings, cancellation during prefill and decode, queueing, continuation, rewind, capability errors and speculative statistics. - engine_prompt_lengths.py: prompt length and dense chunk boundary check, run once per prefill mode and then compared. - mmlu_engine_eval.py, lcb_engine_eval.py: resumable MMLU-Pro 800 and LiveCodeBench v6 accuracy runs through the Engine, with the MTP == no-MTP stream check for MMLU-Pro and sandboxed grading with a paired McNemar comparison for LiveCodeBench. - test_plane_kernels.py: kernel-level tests of the fast prefill kernels of the plane library against the research implementations. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

Tianlei Wu (tianleiwu) added this pull request to stack #2731 October 13, 2026 23:32