feat(models): add cross-platform onnxruntime-genai backend + refactor winml - #3960
Open
thiagocrepaldi wants to merge 3 commits into
Open
feat(models): add cross-platform onnxruntime-genai backend + refactor winml#3960thiagocrepaldi wants to merge 3 commits into
thiagocrepaldi wants to merge 3 commits into
Conversation
… winml
Add a cross-platform `onnxruntime-genai` backend that runs ONNX LLMs exported
by the ONNX Runtime GenAI Model Builder on CPU/CUDA/DirectML/WebGPU/AMD-NPU via
the `og.Config` provider API. All lm-eval logic (single-pass log-likelihood,
multiple-choice, rolling perplexity, generation) lives in
`ONNXRuntimeGenAILM(TemplateLM)`, following the `OptimumLM(HFLM)` single-file
subclass precedent.
Refactor `winml` into a thin `WindowsML(ONNXRuntimeGenAILM)` subclass that
overrides only execution-provider selection, inheriting the shared scoring
core. This removes the previous winml quality debt: the per-token
teacher-forcing loop (now a single `get_output("logits")` pass), the direct
`loglikelihood` override that bypassed `TemplateLM` token handling, and the
silent `(0.0, False)`-on-exception path that corrupted scores.
Add a shared parametrized test suite (unit tests for the scoring math with a
fake `_forward_logits`, plus CPU smoke tests over hellaswag/wikitext/gsm8k
using a tiny Model Builder fixture), packaging extra, and docs. The test is
added to the CI `--ignore` list since the ONNX runtimes are not installed in CI.
Part of QUARK-606 Phase 1. The raw `onnxruntime` backend, cross-backend parity
CI, OpenAI-compatible server, and Quark migration are follow-up stages.
Co-Authored-By: Claude <[email protected]>
Trim the subclassing note to the pre-existing OptimumLM example instead of documenting the onnxruntime-genai backend internals, which belong in the backend's own docstrings/README rather than the general model-authoring guide. Co-Authored-By: Claude <[email protected]>
Move winml-specific coverage (registration, subclass relationship, non-Windows CPU fallback) into tests/models/test_winml.py, leaving test_onnxruntime_genai.py focused on the genai backend and its scoring math. Keeps each backend's tests self-contained and adds the new file to the CI ignore list. Co-Authored-By: Claude <[email protected]>
thiagocrepaldi
marked this pull request as draft
July 29, 2026 17:52
thiagocrepaldi
marked this pull request as ready for review
August 4, 2026 19:57
7 tasks
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
Makes the existing
winmlbackend cross-platform. The newonnxruntime-genaibackend is essentially the currentwinmlcode with the Windows-specific execution-provider selection factored out — the onnxruntime-genai runtime core (og.Model/og.Generator, tokenization, logits, generation) was already platform-agnostic; only three methods inwinmlwere Windows-locked. This PR lifts the shared logic into a base backend and reduceswinmlto just its Windows provider selection.Concretely:
onnxruntime-genaibackend (lm_eval/models/onnxruntime_genai.py):ONNXRuntimeGenAILM(TemplateLM)runs Model Builder ONNX exports via the cross-platformog.Configprovider API (CPU, CUDA, DirectML, WebGPU, AMD-NPU via VitisAI/RyzenAI). It carries the same lm-eval surface the winml backend implemented —loglikelihood(+multiple_choice),loglikelihood_rolling, andgenerate_until.winmlbecomes a thin subclass (WindowsML(ONNXRuntimeGenAILM),lm_eval/models/winml.py) overriding only execution-provider selection (_select_ep) — the Windows ML catalog methods are kept verbatim; if they're unavailable (e.g. off-Windows) it falls back to the cross-platform CPU path. This follows the existingOptimumLM(HFLM)single-file subclass precedent.Because the two share one code path, the winml logic also gets the correctness cleanups that were made while lifting it into the base:
get_output("logits")pass;loglikelihoodoverride that bypassedTemplateLMtoken handling → uses the stockTemplateLMmachinery (so the space-convention / empty-context handling is inherited);(0.0, False)-on-exception path that corrupted scores → removed.Tests
tests/models/test_onnxruntime_genai.py— unit tests for the log-softmax/gather scoring math (fake_forward_logits, no model load) + CPU smoke tests overhellaswag/wikitext/gsm8kusing a tiny Model Builder fixture built on the fly (mirrorstest_openvino.py's export-on-the-fly pattern).tests/models/test_winml.py— winml-specific wiring: registration, the subclass relationship, and graceful fallback to the cross-platform CPU path off-Windows.--ignorelist (liketest_openvino.py) because the ONNX runtimes are not installed in CI.onnxruntime-genai
winml (using --limit 200 for brevity)
Packaging & docs
pip install "lm_eval[onnxruntime-genai]"(EP-specific wheels likeonnxruntime-genai-cuda/-directmlare user-installed and mutually exclusive).Scope
First stage of a larger effort to standardize ONNX evaluation in lm-eval. Out of scope here (follow-up PRs): a raw
onnxruntime.InferenceSessionbackend (matches the exact runtime a customer deploys, and reaches EPs GenAI can't, e.g. ROCm/MIGraphX), a cross-backend genai-vs-raw numerical parity CI test, and an OpenAI-compatible completions server.Test plan
pip install -e ".[onnxruntime-genai]"get_model("onnxruntime-genai"),get_model("winml")python -m pytest tests/models/test_onnxruntime_genai.py tests/models/test_winml.py -vv(11 passed on Linux/CPU)hellaswag/wikitext/gsm8kproduce sane, finite, non-zero-on-error metricsruff check/ruff format --checkcleanwinmlon Windows 11 (testing CPU with --limit 200 for phi 3.5 mini instruct)onnxruntime-genaion Linux (testing CPU and CUDA evaluation for phi 3.5 mini instruct)winmlon Windows for NPU (checking whether NPU testing is actually needed for this PR)