feat(models): add cross-platform onnxruntime-genai backend + refactor winml (#3960)

* feat(models): add cross-platform onnxruntime-genai backend + refactor 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]>

* docs: keep model_guide subclassing tip generic

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]>

* test: split winml tests into their own module

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]>

* fix(onnxruntime-genai): honor genai_config.json providers by default

execution_provider now defaults to None, meaning "use the providers the export
declares in genai_config.json" — og.Config(model_dir) seeds its provider list
from that file, and _select_ep no longer calls clear_providers() unless the
user explicitly asks for a provider. Previously the "cpu" default cleared the
list and appended nothing, silently downgrading a CUDA/DML/NPU export to CPU.
Pass execution_provider=cpu to force CPU, or a provider name to override.

winml pins its historical max_length/max_gen_toks (4096) so existing runs are
numerically unchanged, while the cross-platform backend keeps model-aware
defaults (context_length from the config; max_gen_toks=256).

Add regression tests for the three _select_ep branches and winml's defaults.

Co-Authored-By: Claude <[email protected]>

* pacify pre-commmit

---------

Co-authored-by: Claude <[email protected]>
Co-authored-by: Baber Abbasi <[email protected]>

.github/workflows/unit_tests.yml

@@ -82,7 +82,7 @@ jobs:
       - name: Test with pytest
         env:
           HF_TOKEN: ${{ secrets.HF_TOKEN }}
-        run: pytest -x --showlocals -s -vv -n=auto --ignore=tests/models/test_openvino.py --ignore=tests/models/test_hf_steered.py --ignore=tests/scripts/test_zeno_visualize.py
+        run: pytest -x --showlocals -s -vv -n=auto --ignore=tests/models/test_openvino.py --ignore=tests/models/test_hf_steered.py --ignore=tests/models/test_onnxruntime_genai.py --ignore=tests/models/test_winml.py --ignore=tests/scripts/test_zeno_visualize.py
 
       # Save test artifacts
       - name: Archive test artifacts

README.md

@@ -485,6 +485,34 @@ lm_eval --model sglang \
 > 2. Lower KV cache pool memory usage by adjusting `mem_fraction_static` - Add to your model arguments for example `--model_args pretrained=...,mem_fraction_static=0.7`.
 > 3. Increase tensor parallel size `tp_size` (if using multiple GPUs).
 
+### ONNX Runtime GenAI
+
+We support **ONNX Runtime GenAI** for cross-platform evaluation of ONNX LLMs produced by the [ONNX Runtime GenAI *Model Builder*](https://onnxruntime.ai/docs/genai/howto/build-model.html). Unlike the [Windows ML](#windows-ml) backend, this backend runs on Linux, macOS, and Windows and selects execution providers with the cross-platform `og.Config` API (CPU, CUDA, DirectML, WebGPU, and AMD NPU via VitisAI/RyzenAI).
+
+Install the backend along with the execution-provider wheel that matches your hardware:
+
+```bash
+# CPU
+pip install "lm_eval[onnxruntime-genai]"
+# CUDA / DirectML instead of the CPU wheel (mutually exclusive):
+#   pip install onnxruntime-genai-cuda
+#   pip install onnxruntime-genai-directml
+```
+
+Evaluate a Model Builder ONNX model on a chosen execution provider:
+
+```bash
+lm_eval --model onnxruntime-genai \
+    --model_args pretrained=/path/to/model_builder_output,execution_provider=cuda \
+    --tasks hellaswag \
+    --batch_size 1
+```
+
+The `pretrained` path is a Model Builder output directory (containing `genai_config.json`, the ONNX graph(s), and an HF tokenizer) or a `.onnx` file inside one. `execution_provider` defaults to `cpu` (which runs with no provider appended); pass `cuda`, `dml`, `VitisAI`, etc. for accelerators. Extra provider settings can be passed via `provider_options`.
+
+> [!Note]
+> Supported architectures are whatever the Model Builder supports (Llama, Phi, Qwen, Gemma, Mistral, Granite, ChatGLM, …). Inference is batch-size 1 and runs one execution provider per run.
+
 ### Windows ML
 
 We support **Windows ML** for hardware-accelerated inference on Windows platforms. This enables evaluation on CPU, GPU, and **NPU (Neural Processing Unit)** devices.
@@ -556,6 +584,7 @@ Note that for externally hosted models, configs such as `--device` which relate
 | NVIDIA NeMo                                                                                                               | :heavy_check_mark:                                                                                      | `nemo_lm`                                             | [All supported models](https://docs.nvidia.com/nemo-framework/user-guide/24.09/nemotoolkit/core/core.html#nemo-models)                                          | `generate_until`, `loglikelihood`, `loglikelihood_rolling`                     |
 | NVIDIA Megatron-LM                                                                                                        | :heavy_check_mark:                                                                                      | `megatron_lm`                                         | [Megatron-LM GPT models](https://github.com/NVIDIA/Megatron-LM) (standard and distributed checkpoints)                                                          | `generate_until`, `loglikelihood`, `loglikelihood_rolling`                     |
 | Watsonx.ai                                                                                                                | :heavy_check_mark:                                                                                      | `watsonx_llm`                                         | [Supported Watsonx.ai Engines](https://dataplatform.cloud.ibm.com/docs/content/wsj/analyze-data/fm-models.html?context=wx)                                      | `generate_until` `loglikelihood`                                               |
+| ONNX Runtime GenAI                                                                                                        | :heavy_check_mark:                                                                                      | `onnxruntime-genai`                                   | [ONNX models in GenAI format](https://onnxruntime.ai/docs/genai/howto/build-model.html) (cross-platform: CPU/CUDA/DirectML/NPU)                                  | `generate_until`, `loglikelihood`, `loglikelihood_rolling`                     |
 | Windows ML                                                                                                                | :heavy_check_mark:                                                                                      | `winml`                                               | [ONNX models in GenAI format](https://code.visualstudio.com/docs/intelligentapps/modelconversion)                                                               | `generate_until`, `loglikelihood`, `loglikelihood_rolling`                     |
 | [Your local inference server!](docs/API_guide.md)                                                                         | :heavy_check_mark:                                                                                      | `local-completions` or `local-chat-completions`       | Support for OpenAI API-compatible servers, with easy customization for other APIs.                                                                              | `generate_until`, `loglikelihood`, `loglikelihood_rolling`                     |
 
@@ -801,6 +830,7 @@ These extras install dependencies required to run specific model backends:
 | habana         | Intel Gaudi backend                              |
 | optimum        | Intel OpenVINO models                            |
 | neuronx        | AWS Inferentia2 instances                        |
+| onnxruntime-genai | ONNX Runtime GenAI (cross-platform) - CPU/CUDA/DirectML/NPU |
 | winml          | Windows ML (ONNX Runtime GenAI) - CPU/GPU/NPU    |
 | sparsify       | Sparsify model steering                          |
 | sae_lens       | SAELens model steering                           |

docs/model_guide.md

@@ -187,6 +187,8 @@ If not implemented for a given model type, the flags `--apply_chat_template` , `
 
 **Pro tip**: In order to make the Evaluation Harness overestimate total runtimes rather than underestimate it, HuggingFace models come in-built with the ability to provide responses on data points in *descending order by total input length* via `lm_eval.utils.Reorderer`. Take a look at `lm_eval.models.hf_causal.HFLM` to see how this is done, and see if you can implement it in your own model!
 
+**Subclassing an existing backend**: If your backend shares most of its logic with an existing one and differs only in a narrow concern (e.g. how the underlying model is created or which device it targets), prefer subclassing over duplication. For example, `lm_eval.models.optimum_lm.OptimumLM` subclasses `HFLM` and overrides only `_create_model`.
+
 ## Conclusion
 
 After reading this guide, you should be able to add new model APIs or implementations to the Eval Harness library!

lm_eval/models/__init__.py

@@ -45,6 +45,7 @@ MODEL_MAPPING = {
     "megatron_lm": "lm_eval.models.megatron_lm:MegatronLMEval",
     "nemo_lm": "lm_eval.models.nemo_lm:NeMoLM",
     "neuronx": "lm_eval.models.neuron_optimum:NeuronModelForCausalLM",
+    "onnxruntime-genai": "lm_eval.models.onnxruntime_genai:ONNXRuntimeGenAILM",
     "openai-chat-completions": "lm_eval.models.openai_completions:OpenAIChatCompletion",
     "openai-completions": "lm_eval.models.openai_completions:OpenAICompletionsAPI",
     "openvino": "lm_eval.models.optimum_lm:OptimumLM",
@@ -75,4 +76,4 @@ def _register_all_models():
 # Call registration on module import
 _register_all_models()
 
-__all__ = ["MODEL_MAPPING"]
+__all__ = ["MODEL_MAPPING"]

lm_eval/models/onnxruntime_genai.py

@@ -0,0 +1,405 @@
+"""Cross-platform ONNX Runtime GenAI backend for lm-eval-harness.
+
+Runs models exported by the ONNX Runtime GenAI *Model Builder* through the
+``onnxruntime-genai`` runtime (``og.Model`` / ``og.Generator``). Unlike the
+``winml`` backend -- which selects execution providers via the Windows ML
+catalog and therefore only runs on Windows -- this backend selects providers
+with the cross-platform ``og.Config`` API and runs on CPU, CUDA, DirectML,
+WebGPU, and AMD NPU (VitisAI/RyzenAI).
+
+Example usage:
+    lm_eval --model onnxruntime-genai --model_args pretrained=path/to/model_builder_output,execution_provider=cuda --tasks hellaswag --limit 10
+
+The ``pretrained`` path is a Model Builder output directory (containing
+``genai_config.json``, the ONNX graph(s), and an HF tokenizer) or a ``.onnx``
+file inside such a directory.
+"""
+
+import json
+import logging
+from pathlib import Path
+from typing import TYPE_CHECKING, Any
+
+import numpy as np
+from tqdm import tqdm
+
+from lm_eval import utils
+from lm_eval.api.model import TemplateLM
+from lm_eval.api.registry import register_model
+
+
+if TYPE_CHECKING:
+    from lm_eval.api.instance import Instance
+
+
+eval_logger = logging.getLogger(__name__)
+
+# ONNX Runtime GenAI treats CPU as "no execution provider appended" rather than
+# an appendable provider name, so these aliases are handled specially.
+_CPU_EP_ALIASES = {"cpu", "cpuexecutionprovider"}
+
+
+@register_model("onnxruntime-genai")
+class ONNXRuntimeGenAILM(TemplateLM):
+    """Cross-platform backend for Model Builder ONNX models via onnxruntime-genai.
+
+    All lm-eval scoring logic lives here (tokenization, single-pass logits,
+    log-likelihood, rolling perplexity, and generation). Subclasses need only
+    override :meth:`_select_ep` to change how execution providers are chosen
+    (see :class:`lm_eval.models.winml.WindowsML`).
+    """
+
+    _DEFAULT_MAX_LENGTH = 2048
+
+    def __init__(
+        self,
+        pretrained: str,
+        execution_provider: str | None = None,
+        max_length: int | None = None,
+        batch_size: int = 1,
+        max_gen_toks: int = 256,
+        provider_options: dict[str, Any] | None = None,
+        **kwargs,
+    ) -> None:
+        """Initialize the backend.
+
+        Args:
+            pretrained: Path to a Model Builder output directory or a ``.onnx``
+                file inside one.
+            execution_provider: EP to run on (e.g. ``cpu``, ``cuda``, ``dml``,
+                ``VitisAI``). When ``None`` (default), the providers declared in
+                ``genai_config.json`` are used as-is — matching how a Model
+                Builder export is meant to run on its target device. Pass
+                ``cpu`` to force CPU, or a provider name to override the export.
+            max_length: Maximum sequence length. Defaults to the
+                ``context_length`` from ``genai_config.json`` when available.
+            batch_size: Only 1 is supported; other values are coerced to 1.
+            max_gen_toks: Default maximum number of tokens to generate.
+            provider_options: Extra key/value options forwarded to the EP.
+        """
+        super().__init__()
+
+        self.og = self._import_genai()
+        self.pretrained = pretrained
+        self.model_dir = self._resolve_model_dir(pretrained)
+        self.execution_provider = execution_provider
+        self.provider_options = provider_options or {}
+        self._max_gen_toks = max_gen_toks
+
+        if batch_size != 1:
+            eval_logger.warning(
+                f"{type(self).__name__} currently supports batch size 1 only; "
+                f"overriding requested batch_size={batch_size}."
+            )
+        self.batch_size = 1
+
+        self._genai_config = self._read_genai_config(self.model_dir)
+        model_cfg = self._genai_config.get("model", {})
+        context_length = model_cfg.get("context_length")
+        self.max_length = max_length or context_length or self._DEFAULT_MAX_LENGTH
+
+        self._eot_token_id = self._resolve_eot_token_id(model_cfg)
+        self._bos_token_id = model_cfg.get("bos_token_id")
+
+        self._load()
+
+    @staticmethod
+    def _import_genai():
+        try:
+            import onnxruntime_genai as og
+        except ImportError as e:
+            raise ImportError(
+                "onnxruntime-genai is required for the onnxruntime-genai backend. "
+                "Install a matching EP wheel, e.g. `pip install onnxruntime-genai` "
+                "(CPU) or `pip install onnxruntime-genai-cuda` (CUDA)."
+            ) from e
+        eval_logger.info(f"ONNX Runtime GenAI version: {og.__version__}")
+        return og
+
+    @staticmethod
+    def _resolve_model_dir(pretrained: str) -> Path:
+        path = Path(pretrained)
+        if path.is_file() and path.suffix == ".onnx":
+            return path.parent
+        if path.is_dir():
+            return path
+        raise FileNotFoundError(
+            f"Model path {pretrained} not found or is not a directory / .onnx file"
+        )
+
+    @staticmethod
+    def _read_genai_config(model_dir: Path) -> dict[str, Any]:
+        config_path = model_dir / "genai_config.json"
+        if not config_path.exists():
+            eval_logger.warning(
+                f"No genai_config.json found in {model_dir}; falling back to defaults."
+            )
+            return {}
+        with open(config_path, encoding="utf-8") as f:
+            return json.load(f)
+
+    def _resolve_eot_token_id(self, model_cfg: dict[str, Any]) -> int:
+        eos = model_cfg.get("eos_token_id")
+        if isinstance(eos, list):
+            eos = eos[0] if eos else None
+        if eos is not None:
+            return int(eos)
+        # Fall back to the tokenizer once it is loaded (see _load).
+        return None  # type: ignore[return-value]
+
+    # ------------------------------------------------------------------ #
+    # Engine hooks (override in subclasses to change engine / EP selection)
+    # ------------------------------------------------------------------ #
+    def _select_ep(self, config) -> None:
+        """Configure execution providers on an ``og.Config``.
+
+        ``og.Config(model_dir)`` seeds its provider list from
+        ``genai_config.json``, so we only touch it when the user explicitly
+        asks for a provider:
+
+        * ``execution_provider is None`` — leave the config untouched, honoring
+          the providers the export was built with (a CUDA/DML/NPU export runs
+          on its target device without extra flags).
+        * ``execution_provider`` is a CPU alias — clear providers so it runs on
+          CPU (an empty provider list means CPU).
+        * otherwise — clear and append the named provider with any
+          ``provider_options``.
+        """
+        ep = self.execution_provider
+        if ep is None:
+            eval_logger.info(
+                "No execution_provider given; using the providers declared in "
+                "genai_config.json."
+            )
+            return
+        config.clear_providers()
+        if ep.lower() in _CPU_EP_ALIASES:
+            eval_logger.info("Using CPU execution provider (no provider appended).")
+            return
+        config.append_provider(ep)
+        for key, value in self.provider_options.items():
+            config.set_provider_option(ep, str(key), str(value))
+        eval_logger.info(f"Using execution provider: {ep}")
+
+    def _load(self) -> None:
+        """Load the tokenizer and the onnxruntime-genai model."""
+        from transformers import AutoTokenizer
+
+        self.tokenizer = AutoTokenizer.from_pretrained(str(self.model_dir))
+        if self._eot_token_id is None:
+            self._eot_token_id = (
+                self.tokenizer.eos_token_id
+                if self.tokenizer.eos_token_id is not None
+                else 0
+            )
+        if self._bos_token_id is None:
+            self._bos_token_id = self.tokenizer.bos_token_id
+
+        config = self.og.Config(str(self.model_dir))
+        self._select_ep(config)
+        self.model = self.og.Model(config)
+        eval_logger.info(f"Loaded onnxruntime-genai model from {self.model_dir}")
+
+    def _forward_logits(self, tokens: list[int]) -> np.ndarray:
+        """Run one forward pass and return per-position logits.
+
+        Args:
+            tokens: The full input token sequence.
+
+        Returns:
+            A ``[seq_len, vocab]`` array where row ``i`` is the distribution over
+            the token at position ``i + 1`` given ``tokens[: i + 1]``.
+        """
+        params = self.og.GeneratorParams(self.model)
+        generator = self.og.Generator(self.model, params)
+        generator.append_tokens(np.asarray(tokens, dtype=np.int32))
+        logits = np.asarray(generator.get_output("logits"), dtype=np.float32)
+        if logits.ndim == 3:  # (batch, seq, vocab)
+            logits = logits[0]
+        elif logits.ndim != 2:
+            raise ValueError(f"Unexpected logits shape from GenAI: {logits.shape}")
+        return logits
+
+    def _generate(self, tokens: list[int], gen_kwargs: dict[str, Any]) -> list[int]:
+        """Greedily (or with sampling) generate token ids from a prompt.
+
+        Args:
+            tokens: The prompt token ids.
+            gen_kwargs: Generation options (``max_gen_toks``, ``until``,
+                ``temperature``, ``top_p``, ``top_k``, ``do_sample``).
+
+        Returns:
+            The newly generated token ids (excluding the prompt).
+        """
+        max_gen_toks = int(gen_kwargs.get("max_gen_toks", self.max_gen_toks))
+        temperature = float(gen_kwargs.get("temperature", 0.0))
+        top_p = float(gen_kwargs.get("top_p", 1.0))
+        top_k = int(gen_kwargs.get("top_k", 50))
+        do_sample = bool(gen_kwargs.get("do_sample", False)) and temperature > 0
+
+        params = self.og.GeneratorParams(self.model)
+        # max_length is total (prompt + generated) length for the genai search.
+        total_max = min(len(tokens) + max_gen_toks, self.max_length)
+        if do_sample:
+            params.set_search_options(
+                max_length=total_max,
+                do_sample=True,
+                temperature=temperature,
+                top_p=top_p,
+                top_k=top_k,
+            )
+        else:
+            params.set_search_options(max_length=total_max, do_sample=False)
+
+        generator = self.og.Generator(self.model, params)
+        generator.append_tokens(np.asarray(tokens, dtype=np.int32))
+
+        generated: list[int] = []
+        while not generator.is_done() and len(generated) < max_gen_toks:
+            generator.generate_next_token()
+            seq = generator.get_sequence(0)
+            if len(seq) > len(tokens) + len(generated):
+                generated.append(int(seq[len(tokens) + len(generated)]))
+        return generated
+
+    # ------------------------------------------------------------------ #
+    # TemplateLM interface
+    # ------------------------------------------------------------------ #
+    @property
+    def eot_token_id(self) -> int:
+        return self._eot_token_id
+
+    @property
+    def prefix_token_id(self) -> int:
+        if self._bos_token_id is not None:
+            return int(self._bos_token_id)
+        return self._eot_token_id
+
+    @property
+    def max_gen_toks(self) -> int:
+        return self._max_gen_toks
+
+    def tok_encode(
+        self, string: str, add_special_tokens: bool | None = None, **kwargs
+    ) -> list[int]:
+        if add_special_tokens is None:
+            add_special_tokens = False
+        return self.tokenizer.encode(string, add_special_tokens=add_special_tokens)
+
+    def tok_decode(self, tokens: list[int], **kwargs) -> str:
+        return self.tokenizer.decode(tokens, skip_special_tokens=True)
+
+    def _loglikelihood_tokens(
+        self,
+        requests: list[tuple[tuple[str, str] | None, list[int], list[int]]],
+        disable_tqdm: bool = False,
+        override_bs: int | None = None,
+    ) -> list[tuple[float, bool]]:
+        results: list[tuple[float, bool]] = []
+
+        for cache_key, context_enc, continuation_enc in tqdm(
+            requests, disable=disable_tqdm, desc="Computing log-likelihoods"
+        ):
+            if not continuation_enc:
+                results.append((0.0, True))
+                continue
+
+            # Feed the full sequence; keep the last max_length tokens so the
+            # continuation is always scored, mirroring HFLM's left-truncation.
+            # At least one preceding position is required to score the first
+            # continuation token, so cap contlen to leave room for context.
+            full = (context_enc + continuation_enc)[-self.max_length :]
+            contlen = min(len(continuation_enc), len(full) - 1)
+            ctxlen = len(full) - contlen
+
+            logits = self._forward_logits(full)
+            # Logits at position i predict token i+1, so the continuation is
+            # scored by rows [ctxlen-1 : ctxlen-1+contlen].
+            cont_logits = logits[ctxlen - 1 : ctxlen - 1 + contlen]
+            log_probs = _log_softmax(cont_logits)
+
+            # Left-truncation may have dropped leading continuation tokens.
+            targets = np.asarray(continuation_enc[-contlen:], dtype=np.int64)
+            token_ll = log_probs[np.arange(contlen), targets]
+            total_ll = float(token_ll.sum())
+            is_greedy = bool(np.all(log_probs.argmax(axis=-1) == targets))
+
+            results.append((total_ll, is_greedy))
+
+            if cache_key is not None:
+                self.cache_hook.add_partial(
+                    "loglikelihood", cache_key, (total_ll, is_greedy)
+                )
+
+        return results
+
+    def loglikelihood_rolling(
+        self, requests: list["Instance"], disable_tqdm: bool = False
+    ) -> list[float]:
+        loglikelihoods: list[float] = []
+
+        for (string,) in tqdm(
+            [req.args for req in requests],
+            disable=disable_tqdm,
+            desc="Computing rolling log-likelihoods",
+        ):
+            rolling_token_windows = [
+                (None,) + window
+                for window in map(
+                    utils.make_disjoint_window,
+                    utils.get_rolling_token_windows(
+                        token_list=self.tok_encode(string),
+                        prefix_token=self.prefix_token_id,
+                        max_seq_len=self.max_length,
+                        context_len=1,
+                    ),
+                )
+            ]
+
+            window_lls = self._loglikelihood_tokens(
+                rolling_token_windows, disable_tqdm=True
+            )
+            string_ll = sum(ll for ll, _ in window_lls)
+            loglikelihoods.append(string_ll)
+            self.cache_hook.add_partial("loglikelihood_rolling", (string,), string_ll)
+
+        return loglikelihoods
+
+    def generate_until(
+        self, requests: list["Instance"], disable_tqdm: bool = False
+    ) -> list[str]:
+        results: list[str] = []
+
+        for context, gen_kwargs in tqdm(
+            [req.args for req in requests],
+            disable=disable_tqdm,
+            desc="Generating text",
+        ):
+            gen_kwargs = dict(gen_kwargs)
+            until = gen_kwargs.pop("until", None)
+            if isinstance(until, str):
+                until = [until]
+            elif until is None:
+                until = []
+
+            prompt_tokens = self.tok_encode(context)
+            new_tokens = self._generate(prompt_tokens, gen_kwargs)
+            text = self.tok_decode(new_tokens)
+
+            for stop in until:
+                if stop:
+                    idx = text.find(stop)
+                    if idx != -1:
+                        text = text[:idx]
+
+            results.append(text)
+            self.cache_hook.add_partial("generate_until", (context, gen_kwargs), text)
+
+        return results
+
+
+def _log_softmax(logits: np.ndarray) -> np.ndarray:
+    """Numerically stable log-softmax over the last axis."""
+    shifted = logits - logits.max(axis=-1, keepdims=True)
+    log_sum_exp = np.log(np.exp(shifted).sum(axis=-1, keepdims=True))
+    return shifted - log_sum_exp

lm_eval/models/winml.py

@@ -1,172 +1,101 @@
-"""
-WinML backend for lm-eval-harness with NPU/GPU/CPU support.
+"""WinML backend for lm-eval-harness with NPU/GPU/CPU support.
 
-This backend leverages Windows Machine Learning (WinML) to run models on various
-hardware backends including NPUs, GPUs, and CPUs. It's particularly useful for
-running inference on Windows devices with dedicated Neural Processing Units.
+This backend runs Model Builder ONNX models through the Windows Machine Learning
+(WinML) execution-provider catalog. It is a thin specialization of
+:class:`~lm_eval.models.onnxruntime_genai.ONNXRuntimeGenAILM`: all scoring and
+generation logic is inherited, and only execution-provider selection is
+Windows-specific.
 
 Example usage:
-    lm_eval --model winml --model_args pretrained=path/to/onnx/model.onnx --tasks hellaswag
+    lm_eval --model winml --model_args pretrained=path/to/onnx/model --tasks hellaswag
 
 """
 
 import logging
 from pathlib import Path
-from typing import TYPE_CHECKING, Any
 
-import numpy as np
-from tqdm import tqdm
-
-from lm_eval import utils
-from lm_eval.api.model import TemplateLM
 from lm_eval.api.registry import register_model
+from lm_eval.models.onnxruntime_genai import ONNXRuntimeGenAILM
 
 
-if TYPE_CHECKING:
-    from lm_eval.api.instance import Instance
-
 eval_logger = logging.getLogger(__name__)
 
 
 @register_model("winml")
-class WindowsML(TemplateLM):
-    """
-    WindowsML backend for lm-eval-harness with NPU/GPU/CPU support.
+class WindowsML(ONNXRuntimeGenAILM):
+    """WindowsML backend: onnxruntime-genai with Windows ML provider selection.
 
-    This model class provides integration with Windows Machine Learning (WindowsML)
-    to enable evaluation on NPUs and other Windows-optimized hardware.
+    Inherits all lm-eval logic (tokenization, single-pass log-likelihood,
+    rolling perplexity, and generation) from
+    :class:`~lm_eval.models.onnxruntime_genai.ONNXRuntimeGenAILM`, overriding
+    only :meth:`_select_ep` to register execution providers through the Windows
+    ML catalog.
+
+    The default ``max_length`` / ``max_gen_toks`` are pinned to the historical
+    winml values so numeric behavior is unchanged; the cross-platform backend
+    picks its own (more model-aware) defaults.
     """
 
-    _DEFAULT_MAX_LENGTH = 2048
-
-    @classmethod
-    def create_from_arg_obj(
-        cls, arg_dict: dict[str, Any], additional_config: dict[str, Any] | None = None
-    ) -> "WindowsML":
-        """
-        Override to properly merge dictionaries and avoid duplicate keyword arguments.
-
-        Args:
-            arg_dict: Dictionary containing model arguments
-            additional_config: Optional dictionary containing additional configuration
-
-        Returns:
-            Instance of WindowsML class
-        """
-        # Merge the dictionaries, with additional_config taking precedence
-        merged_config = {**(arg_dict or {})}
-        if additional_config:
-            # Filter out None values and merge
-            filtered_additional = {
-                k: v for k, v in additional_config.items() if v is not None
-            }
-            merged_config.update(filtered_additional)
-
-        return cls(**merged_config)
+    # Historical winml defaults, kept so existing winml runs are unaffected.
+    _WINML_MAX_LENGTH = 4096
+    _WINML_MAX_GEN_TOKS = 4096
 
     def __init__(
         self,
         pretrained: str,
-        max_length: int | None = 4096,
-        batch_size: int = 1,
-        max_batch_size: int = 64,
+        max_length: int | None = None,
+        max_gen_toks: int | None = None,
         **kwargs,
     ) -> None:
+        super().__init__(
+            pretrained,
+            max_length=max_length or self._WINML_MAX_LENGTH,
+            max_gen_toks=max_gen_toks or self._WINML_MAX_GEN_TOKS,
+            **kwargs,
+        )
+
+    def _select_ep(self, config) -> None:
+        """Register Windows ML execution providers on the ``og.Config``.
+
+        Discovers providers via the Windows ML catalog and registers them with
+        onnxruntime-genai. If discovery fails (e.g. running outside Windows), we
+        fall back to the cross-platform selection in the base class.
         """
-        Initialize WindowsML model.
-
-        Args:
-            pretrained: Path to ONNX model file or directory containing model files
-            max_length: Maximum sequence length
-            batch_size: Batch size for inference
-            max_batch_size: Maximum batch size for auto-batching
-        """
-        super().__init__()
-
-        # Validate and import dependencies
-        self._validate_dependencies()
-
-        # Store configuration
-        self.pretrained = pretrained
-        self.max_length = max_length or self._DEFAULT_MAX_LENGTH
-        self.batch_size = batch_size
-        self.max_batch_size = max_batch_size
-
-        # Warn about batch size limitations
-        if batch_size != 1 or max_batch_size != 1:
-            eval_logger.warning(
-                f"WindowsML backend currently only supports batch size 1. "
-                f"Requested batch_size={batch_size}, max_batch_size={max_batch_size}. "
-                f"Setting both to 1."
-            )
-            self.batch_size = 1
-            self.max_batch_size = 1
-
         self._fix_winrt_runtime()
+        if self._register_winml_providers_to_genai():
+            self._log_winml_devices()
+        else:
+            eval_logger.warning(
+                "Windows ML provider registration failed; falling back to "
+                "cross-platform execution-provider selection."
+            )
+        super()._select_ep(config)
 
-        # Initialize Windows ML execution providers
-        self._register_winml_providers_to_genai()
+    def _fix_winrt_runtime(self) -> None:
+        """Remove the bundled msvcp140.dll from winrt-runtime to avoid clashes.
 
-        # Setup device and execution providers using new Windows ML API
-        self._setup_winml_devices_and_providers()
-
-        # Load and compile ONNX model
-        self._load_and_compile_model(pretrained)
-
-        eval_logger.info(
-            f"Available EP devices: {len(self.ep_device_map)} execution providers"
-        )
-
-    def _validate_dependencies(self) -> None:
-        """
-        Validate that required dependencies are available.
-
-        Raises:
-            ImportError: If required dependencies are not installed
+        This DLL, shipped by the ``winrt-runtime`` package, can conflict with
+        other libraries loaded in the same process.
         """
         try:
-            import onnxruntime_genai as og
-
-            self.og = og
-            eval_logger.info(f"ONNX Runtime GenAI version: {og.__version__}")
-        except ImportError as e:
-            raise ImportError(
-                "ONNX Runtime GenAI is required for WinML backend. "
-                "Install with: pip install onnxruntime-genai"
-            ) from e
-
-        # Also import regular ONNX Runtime for EP registration
+            from importlib import metadata
+        except ImportError:
+            return
         try:
-            import onnxruntime as ort
-
-            self.ort = ort
-            eval_logger.info(f"ONNX Runtime version: {ort.__version__}")
-        except ImportError as e:
-            raise ImportError(
-                "ONNX Runtime is also required for execution provider registration. "
-                "Install with: pip install onnxruntime"
-            ) from e
-
-    def _fix_winrt_runtime(self):
-        """
-        This function removes the msvcp140.dll from the winrt-runtime package.
-        So it does not cause issues with other libraries.
-        """
-        from importlib import metadata
-
-        site_packages_path = Path(
-            str(metadata.distribution("winrt-runtime").locate_file(""))
-        )
+            site_packages_path = Path(
+                str(metadata.distribution("winrt-runtime").locate_file(""))
+            )
+        except metadata.PackageNotFoundError:
+            return
         dll_path = site_packages_path / "winrt" / "msvcp140.dll"
         if dll_path.exists():
             dll_path.unlink()
 
     def _register_winml_providers_to_genai(self) -> bool:
-        """
-        Register Windows ML execution providers to ONNX Runtime GenAI.
+        """Discover Windows ML providers and register them with onnxruntime-genai.
 
         Returns:
-            True if registration was successful, False otherwise
+            True if registration succeeded, False otherwise.
         """
         try:
             import winui3.microsoft.windows.ai.machinelearning as winml
@@ -180,610 +109,44 @@ class WindowsML(TemplateLM):
                 providers = catalog.find_all_providers()
                 for provider in providers:
                     provider.ensure_ready_async().get()
-                    # Register to GenAI instead of regular ONNX Runtime
                     self.og.register_execution_provider_library(
                         provider.name, provider.library_path
                     )
                     eval_logger.info(
-                        f"Registered {provider.name} to ONNX Runtime GenAI"
+                        "Registered %s to ONNX Runtime GenAI", provider.name
                     )
-
             return True
         except ImportError as e:
-            eval_logger.warning(f"Windows ML import error: {e}")
+            eval_logger.warning("Windows ML import error: %s", e)
             return False
-        except Exception as e:
-            eval_logger.warning(f"Error registering providers to GenAI: {e}")
+        except Exception as e:  # noqa: BLE001
+            eval_logger.warning("Error registering providers to GenAI: %s", e)
             return False
 
-    def _setup_winml_devices_and_providers(self) -> None:
-        """Setup execution providers using Windows ML device enumeration API.
-
-        This method queries available devices and builds a mapping of execution providers.
-        """
+    def _log_winml_devices(self) -> None:
+        """Log available execution-provider devices via the Windows ML API."""
         try:
-            # Get available EP devices using Windows ML API
-            ep_devices = self.ort.get_ep_devices()
-            self.ep_device_map = {}
+            import onnxruntime as ort
 
-            # Build device map
+            ep_devices = ort.get_ep_devices()
+            ep_device_map: dict[str, list] = {}
             for device in ep_devices:
-                ep_name = device.ep_name
-                if ep_name not in self.ep_device_map:
-                    self.ep_device_map[ep_name] = []
-                self.ep_device_map[ep_name].append(device)
+                ep_device_map.setdefault(device.ep_name, []).append(device)
 
-            # Log available devices
-            eval_logger.info("Available execution provider devices:")
-            for name, devices in self.ep_device_map.items():
-                eval_logger.info(f"Execution Provider: {name}")
+            eval_logger.info(
+                "Available EP devices: %s execution providers", len(ep_device_map)
+            )
+            for name, devices in ep_device_map.items():
+                eval_logger.info("Execution Provider: %s", name)
                 for device in devices:
                     try:
-                        device_type = self.ort.OrtHardwareDeviceType(
-                            device.device.type
-                        ).name
-                        eval_logger.info(
-                            f" | Vendor: {device.ep_vendor:<16} | Device Type: {device_type:<8}"
-                        )
-                    except Exception:
-                        eval_logger.info(
-                            f" | Vendor: {device.ep_vendor:<16} | Device Type: Unknown"
-                        )
-
-        except Exception as e:
-            eval_logger.warning(f"Windows ML device enumeration failed: {e}")
-            eval_logger.info("Falling back to legacy provider selection")
-            self.ep_device_map = {}
-
-    def _load_and_compile_model(self, model_path: str) -> None:
-        """
-        Load and optionally compile ONNX model with ONNX Runtime GenAI.
-
-        Args:
-            model_path: Path to ONNX model file or directory
-
-        Raises:
-            FileNotFoundError: If model path is not found or invalid
-            Exception: If model loading fails
-        """
-        model_path = Path(model_path)
-
-        # Handle different input formats
-        if model_path.is_file() and model_path.suffix == ".onnx":
-            input_model_path = model_path.parent  # GenAI expects directory
-        elif model_path.is_dir():
-            input_model_path = model_path
-        else:
-            raise FileNotFoundError(f"Model path {model_path} not found or invalid")
-
-        # Load model using ONNX Runtime GenAI
-        try:
-            eval_logger.info(
-                f"Loading model with ONNX Runtime GenAI from: {input_model_path}"
-            )
-
-            # Load model and tokenizer using GenAI
-            self.genai_model = self.og.Model(str(input_model_path))
-            self.genai_tokenizer = self.og.Tokenizer(self.genai_model)
-
-            eval_logger.info(
-                "Model and tokenizer loaded successfully with ONNX Runtime GenAI"
-            )
-
-            # Store model info
-            self.model_path = input_model_path
-
-        except Exception as e:
-            eval_logger.error(
-                f"Failed to load model with ONNX Runtime GenAI from {input_model_path}: {e}"
-            )
-            raise
-
-    @property
-    def eot_token_id(self) -> int:
-        """
-        Get the end-of-text token ID.
-
-        Returns:
-            End-of-text token ID from the GenAI tokenizer
-        """
-        try:
-            # For GenAI tokenizers, try to access EOS from model config or encoding
-            if hasattr(self.genai_model, "config"):
-                config = self.genai_model.config
-                if hasattr(config, "eos_token_id") and config.eos_token_id:
-                    if isinstance(config.eos_token_id, list):
-                        return config.eos_token_id[0]
-                    else:
-                        return config.eos_token_id
-
-        except Exception as e:
-            eval_logger.warning(
-                f"Error getting EOS token ID: {e}, using fallback value 2"
-            )
-
-        # Fallback: return a common EOS token ID
-        eval_logger.warning("Could not determine EOS token ID, using fallback value 2")
-        return 2
-
-    @property
-    def prefix_token_id(self) -> int | None:
-        """
-        Get the prefix token ID (typically BOS token).
-
-        Returns:
-            BOS token ID if available, otherwise None
-        """
-        try:
-            # For GenAI tokenizers, try to access BOS from model config or encoding
-
-            # If we have access to the model config
-            if hasattr(self.genai_model, "config"):
-                config = self.genai_model.config
-                if hasattr(config, "bos_token_id") and config.bos_token_id is not None:
-                    return config.bos_token_id
-
-            # Check if encoding empty string gives BOS token
-            try:
-                empty_enc = self.genai_tokenizer.encode("")
-                if len(empty_enc) > 0:
-                    return empty_enc[0]
-            except Exception:
-                pass
-
-            return None
-
-        except Exception:
-            return None
-
-    @property
-    def max_gen_toks(self) -> int:
-        """
-        Get the maximum number of tokens to generate.
-
-        Returns:
-            Maximum generation tokens (default: 4096)
-        """
-        return 4096
-
-    def tok_encode(
-        self,
-        string: str,
-        left_truncate_len: int | None = None,
-        add_special_tokens: bool = True,
-    ) -> list[int]:
-        """
-        Tokenize string and return token IDs.
-
-        Args:
-            string: Input string to tokenize
-            left_truncate_len: If provided, truncate from the left to this length
-            add_special_tokens: Whether to add special tokens (note: GenAI tokenizer handles this automatically)
-
-        Returns:
-            List of token IDs
-        """
-        # Use GenAI tokenizer for consistency with model inference
-        encoding = self.genai_tokenizer.encode(string)
-
-        # Handle left truncation if requested
-        if left_truncate_len is not None and len(encoding) > left_truncate_len:
-            encoding = encoding[-left_truncate_len:]
-
-        return encoding
-
-    def tok_decode(self, tokens: list[int]) -> str:
-        """
-        Decode token IDs back to text.
-
-        Args:
-            tokens: List of token IDs to decode
-
-        Returns:
-            Decoded text string
-        """
-        return self.genai_tokenizer.decode(tokens)
-
-    def _run_genai_inference_for_full_logits(self, input_text: str) -> np.ndarray:
-        """
-        Run inference using ONNX Runtime GenAI to get full logits sequence.
-
-        Args:
-            input_text: Input text string to compute logits for
-
-        Returns:
-            Logits matrix of shape (seq_len, vocab_size) where logits[i] contains
-            predictions for the token at position i+1 given tokens[0:i+1]
-
-        Raises:
-            Exception: If inference fails
-        """
-        try:
-            # Encode input text to tokens
-            input_tokens = self.genai_tokenizer.encode(input_text)
-
-            if len(input_tokens) == 0:
-                eval_logger.warning("No tokens to process; returning empty array")
-                return np.empty((0, 0), dtype=np.float32)
-
-            # Create generator and get full logits in single pass
-            params = self.og.GeneratorParams(self.genai_model)
-            params.set_search_options(max_length=4096, do_sample=False)
-            generator = self.og.Generator(self.genai_model, params)
-
-            # Append all tokens at once
-            generator.append_tokens(input_tokens)
-
-            # Get FULL logits using get_output("logits") - this gives all positions!
-            full_logits_tensor = generator.get_output("logits")
-            logits_array = np.array(full_logits_tensor, dtype=np.float32)
-
-            # Handle different tensor shapes
-            if len(logits_array.shape) == 3:  # (batch_size, seq_len, vocab_size)
-                logits_matrix = logits_array[0]  # Remove batch dimension
-            elif len(logits_array.shape) == 2:  # (seq_len, vocab_size)
-                logits_matrix = logits_array
-            else:
-                raise ValueError(f"Unexpected logits shape: {logits_array.shape}")
-
-            eval_logger.debug(
-                f"Full logits shape: {logits_matrix.shape} for {len(input_tokens)} input tokens"
-            )
-            return logits_matrix
-
-        except Exception as e:
-            eval_logger.error(f"GenAI inference failed: {e}")
-            raise
-
-    def _loglikelihood_tokens(
-        self,
-        requests: list[tuple[tuple[str, str], list[int], list[int]]],
-        disable_tqdm: bool = False,
-        override_bs: int | None = None,
-    ) -> list[tuple[float, bool]]:
-        """
-        Stub implementation - not used since we override loglikelihood directly.
-        WindowsML uses the GenAI tokenizer and overrides loglikelihood to work
-        with text inputs directly, avoiding tokenization round-trip issues.
-
-        Args:
-            requests: List of tokenized requests
-            disable_tqdm: Whether to disable progress bar
-            override_bs: Optional batch size override
-
-        Returns:
-            Empty list (method not used)
-        """
-        raise NotImplementedError(
-            "WindowsML overrides loglikelihood() directly and does not use _loglikelihood_tokens()"
-        )
-
-    def loglikelihood(
-        self, requests: list["Instance"], disable_tqdm: bool = False
-    ) -> list[tuple[float, bool]]:
-        """
-        Compute log-likelihood using ONNX Runtime GenAI with teacher forcing.
-
-        Args:
-            requests: List of instances containing (context, continuation) text pairs
-            disable_tqdm: Whether to disable progress bar
-
-        Returns:
-            List of tuples containing (log_likelihood, is_greedy) for each request
-        """
-        results = []
-
-        for request in tqdm(
-            requests, disable=disable_tqdm, desc="Computing log-likelihoods"
-        ):
-            context, continuation = request.args
-
-            if len(continuation) == 0:
-                results.append((0.0, True))
-                continue
-
-            try:
-                full_text = context + continuation
-
-                # Tokenize to get continuation tokens
-                if context:
-                    context_enc = self.tok_encode(context)
-                    if (
-                        len(context_enc) > 0
-                        and self.prefix_token_id is not None
-                        and context_enc[0] == self.prefix_token_id
-                    ):
-                        context_enc = context_enc[1:]
-                    context_len = len(context_enc)
-                else:
-                    context_enc = (
-                        [self.prefix_token_id]
-                        if self.prefix_token_id is not None
-                        else []
+                        device_type = ort.OrtHardwareDeviceType(device.device.type).name
+                    except Exception:  # noqa: BLE001
+                        device_type = "Unknown"
+                    eval_logger.info(
+                        " | Vendor: %-16s | Device Type: %-8s",
+                        device.ep_vendor,
+                        device_type,
                     )
-                    context_len = len(context_enc)
-
-                full_tokens = self.tok_encode(full_text)
-                continuation_tokens = full_tokens[context_len:]
-
-                if len(continuation_tokens) == 0:
-                    results.append((0.0, True))
-                    continue
-
-                # Create generator for teacher forcing
-                params = self.og.GeneratorParams(self.genai_model)
-                params.set_search_options(max_length=4096, do_sample=False)
-                gen = self.og.Generator(self.genai_model, params)
-
-                # Initialize context: append context tokens to build the initial state
-                context_tokens = full_tokens[:context_len]
-
-                if len(context_tokens) > 0:
-                    # Append context tokens as 1D numpy array
-                    gen.append_tokens(np.array(context_tokens, dtype=np.int32))
-                else:
-                    # If no context, append BOS if available
-                    if self.prefix_token_id is not None:
-                        gen.append_tokens(
-                            np.array([self.prefix_token_id], dtype=np.int32)
-                        )
-
-                # Teacher forcing: iterate through continuation tokens
-                total_ll = 0.0
-                greedy_ok = True
-
-                for tok in continuation_tokens:
-                    # Get logits for next token prediction (only last position)
-                    logits = gen.get_logits()
-                    logits_array = np.array(logits, dtype=np.float32)
-
-                    # Extract logits vector (get_logits returns last position only)
-                    if len(logits_array.shape) == 3:  # (batch, 1, vocab)
-                        next_logits = logits_array[0, -1, :]
-                    elif len(logits_array.shape) == 2:  # (1, vocab)
-                        next_logits = logits_array[-1, :]
-                    else:  # (vocab,)
-                        next_logits = logits_array
-
-                    # Compute log probability using numerically stable log_softmax
-                    max_logit = np.max(next_logits)
-                    exp_logits = np.exp(next_logits - max_logit)
-                    log_sum_exp = max_logit + np.log(np.sum(exp_logits))
-                    log_prob = next_logits[int(tok)] - log_sum_exp
-
-                    total_ll += float(log_prob)
-
-                    # Check if greedy (argmax matches actual token)
-                    if int(np.argmax(next_logits)) != int(tok):
-                        greedy_ok = False
-
-                    # Teacher forcing: append the actual token to continue generation
-                    gen.append_tokens(np.array([int(tok)], dtype=np.int32))
-
-                results.append((total_ll, greedy_ok))
-
-            except Exception as e:
-                eval_logger.warning(f"Failed to compute loglikelihood: {e}")
-                import traceback
-
-                eval_logger.debug(traceback.format_exc())
-                results.append((0.0, False))
-
-        return results
-
-    def loglikelihood_rolling(
-        self, requests: list["Instance"], disable_tqdm: bool = False
-    ) -> list[float]:
-        """
-        Compute rolling log-likelihood for perplexity using ONNX Runtime GenAI.
-        Uses sliding windows to handle sequences longer than max_length.
-
-        Args:
-            requests: List of instances containing text sequences
-            disable_tqdm: Whether to disable progress bar
-
-        Returns:
-            List of sum of log-likelihood values for each request
-        """
-        loglikelihoods = []
-
-        for request in tqdm(
-            requests, disable=disable_tqdm, desc="Computing rolling log-likelihoods"
-        ):
-            string = request.args[0]
-
-            # Use sliding window approach for long sequences
-            rolling_token_windows = list(
-                map(
-                    utils.make_disjoint_window,
-                    utils.get_rolling_token_windows(
-                        token_list=self.tok_encode(string),
-                        prefix_token=self.prefix_token_id,
-                        max_seq_len=self.max_length,
-                        context_len=1,
-                    ),
-                )
-            )
-
-            # Prepare windows for loglikelihood (convert to context/continuation text pairs)
-            window_requests = []
-            for context_tokens, continuation_tokens in rolling_token_windows:
-                context_text = self.tok_decode(context_tokens)
-                continuation_text = self.tok_decode(continuation_tokens)
-                window_requests.append((context_text, continuation_text))
-
-            # Use loglikelihood to compute log-likelihoods for all windows
-            # Create Instance objects for the windows
-            from lm_eval.api.instance import Instance
-
-            window_instances = [
-                Instance(
-                    request_type="loglikelihood",
-                    doc={},
-                    arguments=(ctx, cont),
-                    idx=0,
-                    metadata={},
-                )
-                for ctx, cont in window_requests
-            ]
-
-            string_nll = self.loglikelihood(
-                window_instances,
-                disable_tqdm=True,
-            )
-
-            # Extract log-likelihoods (discard is_greedy boolean)
-            string_nll = [x[0] for x in string_nll]
-
-            # Sum all window log-likelihoods
-            total_nll = sum(string_nll)
-            loglikelihoods.append(total_nll)
-
-            # Cache this loglikelihood_rolling request
-            self.cache_hook.add_partial("loglikelihood_rolling", (string,), total_nll)
-
-        return loglikelihoods
-
-    def generate_until(
-        self, requests: list["Instance"], disable_tqdm: bool = False
-    ) -> list[str]:
-        """
-        Generate text until stopping criteria using ONNX Runtime GenAI.
-
-        Args:
-            requests: List of generation requests with context and generation kwargs
-            disable_tqdm: Whether to disable progress bar
-
-        Returns:
-            List of generated text strings
-        """
-        if not requests:
-            return []
-
-        results = []
-
-        for request in tqdm(requests, disable=disable_tqdm, desc="Generating text"):
-            context, gen_kwargs = request.args
-
-            max_gen_toks = gen_kwargs.get("max_gen_toks", self.max_gen_toks)
-            until = gen_kwargs.get("until", [])
-            temperature = gen_kwargs.get("temperature", 0.0)
-            top_p = gen_kwargs.get("top_p", 1.0)
-            top_k = gen_kwargs.get("top_k", 50)
-            do_sample = gen_kwargs.get("do_sample", False)
-
-            try:
-                # Use GenAI generation API
-                generated_text = self._run_genai_generation(
-                    context,
-                    max_gen_toks,
-                    until,
-                    temperature=temperature,
-                    top_p=top_p,
-                    top_k=top_k,
-                    do_sample=do_sample,
-                )
-                results.append(generated_text)
-
-            except Exception as e:
-                eval_logger.warning(f"Generation failed for request: {e}")
-                results.append("")
-
-        return results
-
-    def _run_genai_generation(
-        self,
-        prompt: str,
-        max_tokens: int = 4096,
-        stop_sequences: list[str] | None = None,
-        temperature: float = 0.0,
-        top_p: float = 1.0,
-        top_k: int = 50,
-        do_sample: bool = False,
-    ) -> str:
-        """
-        Run text generation using ONNX Runtime GenAI.
-
-        Args:
-            prompt: Input text to generate from
-            max_tokens: Maximum number of tokens to generate
-            stop_sequences: List of sequences that will stop generation
-            temperature: Sampling temperature (0 = greedy, higher = more random)
-            top_p: Nucleus sampling threshold
-            top_k: Top-k sampling parameter
-            do_sample: Whether to use sampling (if False, uses greedy decoding)
-
-        Returns:
-            Generated text string
-        """
-        try:
-            # Create generator parameters
-            params = self.og.GeneratorParams(self.genai_model)
-            # Set generation parameters based on mode
-            if do_sample and temperature > 0:
-                # Sampling mode: use provided parameters
-                params.set_search_options(
-                    max_length=int(max_tokens),
-                    top_p=float(top_p),
-                    top_k=int(top_k),
-                    temperature=float(temperature),
-                    repetition_penalty=1.0,
-                )
-            else:
-                # Greedy mode: force temperature to 0 for deterministic output
-                params.set_search_options(
-                    max_length=int(max_tokens),
-                    top_p=1.0,
-                    top_k=1,
-                    temperature=0.0,
-                    repetition_penalty=1.0,
-                )
-            # Create generator
-            generator = self.og.Generator(self.genai_model, params)
-            # Encode prompt and append tokens to the generator
-            input_tokens = self.genai_tokenizer.encode(prompt)
-            generator.append_tokens(input_tokens)
-
-            # Generate tokens
-            generated_tokens = []
-            while not generator.is_done():
-                generator.generate_next_token()
-
-                if generator.is_done():
-                    break
-
-                # Get the latest generated token
-                output_tokens = generator.get_sequence(0)
-                base = len(input_tokens) + len(generated_tokens)
-                if len(output_tokens) > base:
-                    new_token = output_tokens[base]
-                    generated_tokens.append(new_token)
-
-                    # Check stopping criteria
-                    if len(generated_tokens) >= max_tokens:
-                        break
-
-                    # Decode current generation to check stop sequences
-                    if stop_sequences:
-                        current_text = self.genai_tokenizer.decode(generated_tokens)
-                        if any(stop_seq in current_text for stop_seq in stop_sequences):
-                            break
-
-            # Decode generated tokens
-            if generated_tokens:
-                generated_text = self.genai_tokenizer.decode(generated_tokens)
-
-                # Remove stop sequences from the end
-                if stop_sequences:
-                    for stop_seq in stop_sequences:
-                        if generated_text.endswith(stop_seq):
-                            generated_text = generated_text[: -len(stop_seq)]
-                            break
-
-                return generated_text
-            else:
-                return ""
-
-        except Exception as e:
-            eval_logger.error(f"GenAI generation error: {e}")
-            return ""
+        except Exception as e:  # noqa: BLE001
+            eval_logger.warning("Windows ML device enumeration failed: %s", e)

pyproject.toml

@@ -59,6 +59,10 @@ gptq = ["auto-gptq[triton]>=0.6.0"]
 gptqmodel = ["gptqmodel>=1.0.9"]
 ipex = ["optimum-intel"]
 ibm_watsonx_ai = ["ibm_watsonx_ai>=1.1.22", "python-dotenv"]
+# ONNX Runtime GenAI backend. For non-CPU execution providers install the
+# matching wheel instead of `onnxruntime-genai`, e.g. `onnxruntime-genai-cuda`
+# or `onnxruntime-genai-directml` (these are mutually exclusive).
+onnxruntime-genai = ["onnxruntime-genai", "transformers", "numpy"]
 litellm = ["litellm>=1.60,<1.85", "aiohttp", "requests", "tenacity", "tqdm"]
 # mamba requires CUDA (nvcc) - cannot build on macOS/CPU-only systems
 # mamba = ["mamba_ssm", "causal-conv1d==1.0.2", "torch"]

tests/models/test_onnxruntime_genai.py

@@ -0,0 +1,239 @@
+"""Tests for the cross-platform onnxruntime-genai backend.
+
+The heavy smoke tests require ``onnxruntime-genai`` plus its Model Builder and
+build a tiny random Llama model on the fly (CPU execution provider), mirroring
+the export-on-the-fly pattern in ``tests/models/test_openvino.py``. They are
+excluded from CI (see the ``--ignore`` in ``.github/workflows/unit_tests.yml``)
+because the ONNX runtimes are not installed there; run them locally with:
+
+    python -m pytest tests/models/test_onnxruntime_genai.py -vv
+"""
+
+import tempfile
+
+import numpy as np
+import pytest
+
+from lm_eval.api.model import CacheHook
+from lm_eval.models.onnxruntime_genai import ONNXRuntimeGenAILM, _log_softmax
+
+
+pytest.importorskip("onnxruntime_genai")
+
+
+def _build_tiny_model_dir():
+    """Build a tiny random Llama Model Builder export (fp32, CPU).
+
+    head_size must be a multiple of 16 for the genai CPU GroupQueryAttention
+    kernel, so hidden_size=128 / num_heads=8 -> head_size=16.
+    """
+    builder = pytest.importorskip("onnxruntime_genai.models.builder")
+    from transformers import AutoTokenizer, LlamaConfig, LlamaForCausalLM
+
+    src = tempfile.mkdtemp(prefix="ort_genai_src_")
+    out = tempfile.mkdtemp(prefix="ort_genai_out_")
+    cache = tempfile.mkdtemp(prefix="ort_genai_cache_")
+
+    tokenizer = AutoTokenizer.from_pretrained(
+        "hf-internal-testing/tiny-random-LlamaForCausalLM"
+    )
+    config = LlamaConfig(
+        vocab_size=tokenizer.vocab_size,
+        hidden_size=128,
+        intermediate_size=256,
+        num_hidden_layers=2,
+        num_attention_heads=8,
+        num_key_value_heads=8,
+        max_position_embeddings=512,
+        tie_word_embeddings=True,
+    )
+    LlamaForCausalLM(config).save_pretrained(src)
+    tokenizer.save_pretrained(src)
+
+    try:
+        builder.create_model(None, src, out, "fp32", "cpu", cache)
+    except Exception as e:  # noqa: BLE001
+        pytest.skip(f"Model Builder could not create a CPU fixture: {e}")
+    return out
+
+
+@pytest.fixture(scope="session")
+def tiny_model_dir():
+    return _build_tiny_model_dir()
+
+
+# --------------------------------------------------------------------------- #
+# Unit tests: scoring math, no model load required.
+# --------------------------------------------------------------------------- #
+def _make_bare_lm(forward_logits, max_length=2048):
+    """Build an instance without loading a model, wiring only what the token scoring path needs."""
+    lm = ONNXRuntimeGenAILM.__new__(ONNXRuntimeGenAILM)
+    lm.max_length = max_length
+    lm.cache_hook = CacheHook(None)
+    lm._forward_logits = forward_logits
+    return lm
+
+
+def test_log_softmax_matches_reference():
+    rng = np.random.default_rng(0)
+    logits = rng.normal(size=(4, 7)).astype(np.float32)
+    got = _log_softmax(logits)
+    # Reference: log(exp(x) / sum(exp(x)))
+    ref = logits - np.log(np.exp(logits).sum(axis=-1, keepdims=True))
+    np.testing.assert_allclose(got, ref, rtol=1e-5, atol=1e-5)
+    # Each row is a valid log-prob distribution (sums to 1 in prob space).
+    np.testing.assert_allclose(np.exp(got).sum(axis=-1), np.ones(4), atol=1e-5)
+
+
+def test_loglikelihood_tokens_math():
+    # vocab=5; craft logits so the continuation tokens are the argmax everywhere.
+    vocab = 5
+    context_enc = [1, 2]
+    continuation_enc = [3, 4]
+    full_len = len(context_enc) + len(continuation_enc)
+
+    # Deterministic logits: position i strongly predicts continuation target.
+    logits = np.full((full_len, vocab), -10.0, dtype=np.float32)
+    # Row (ctxlen-1 + j) should predict continuation_enc[j].
+    ctxlen = len(context_enc)
+    for j, tok in enumerate(continuation_enc):
+        logits[ctxlen - 1 + j, tok] = 10.0
+
+    lm = _make_bare_lm(lambda toks: logits)
+    ((ll, is_greedy),) = lm._loglikelihood_tokens(
+        [(None, context_enc, continuation_enc)], disable_tqdm=True
+    )
+
+    # Manually compute expected loglikelihood from the same logits.
+    lp = _log_softmax(logits[ctxlen - 1 : ctxlen - 1 + len(continuation_enc)])
+    expected = float(lp[np.arange(len(continuation_enc)), continuation_enc].sum())
+    assert ll == pytest.approx(expected)
+    assert is_greedy is True
+
+
+def test_loglikelihood_tokens_not_greedy():
+    vocab = 5
+    context_enc = [1]
+    continuation_enc = [2]
+    # Predict a *different* token as argmax, so is_greedy must be False.
+    logits = np.full((2, vocab), -10.0, dtype=np.float32)
+    logits[0, 4] = 10.0  # argmax is 4, but target is 2
+    lm = _make_bare_lm(lambda toks: logits)
+    ((ll, is_greedy),) = lm._loglikelihood_tokens(
+        [(None, context_enc, continuation_enc)], disable_tqdm=True
+    )
+    assert is_greedy is False
+    assert ll < 0.0
+
+
+def test_loglikelihood_tokens_empty_continuation():
+    lm = _make_bare_lm(lambda toks: np.zeros((1, 3), dtype=np.float32))
+    (result,) = lm._loglikelihood_tokens([(None, [1], [])], disable_tqdm=True)
+    assert result == (0.0, True)
+
+
+class _FakeConfig:
+    """Records provider mutations made by _select_ep."""
+
+    def __init__(self):
+        self.cleared = False
+        self.appended = []
+        self.options = []
+
+    def clear_providers(self):
+        self.cleared = True
+
+    def append_provider(self, ep):
+        self.appended.append(ep)
+
+    def set_provider_option(self, ep, key, value):
+        self.options.append((ep, key, value))
+
+
+def _make_ep_lm(execution_provider, provider_options=None):
+    lm = ONNXRuntimeGenAILM.__new__(ONNXRuntimeGenAILM)
+    lm.execution_provider = execution_provider
+    lm.provider_options = provider_options or {}
+    return lm
+
+
+def test_select_ep_none_leaves_config_untouched():
+    # Regression: with no execution_provider, the providers Model Builder wrote
+    # into genai_config.json must survive (clear_providers must NOT be called),
+    # else a CUDA/DML/NPU export would be silently downgraded to CPU.
+    cfg = _FakeConfig()
+    _make_ep_lm(None)._select_ep(cfg)
+    assert cfg.cleared is False
+    assert cfg.appended == []
+
+
+def test_select_ep_cpu_clears_without_appending():
+    cfg = _FakeConfig()
+    _make_ep_lm("cpu")._select_ep(cfg)
+    assert cfg.cleared is True
+    assert cfg.appended == []
+
+
+def test_select_ep_named_provider_clears_and_appends():
+    cfg = _FakeConfig()
+    _make_ep_lm("cuda", {"device_id": 0})._select_ep(cfg)
+    assert cfg.cleared is True
+    assert cfg.appended == ["cuda"]
+    assert cfg.options == [("cuda", "device_id", "0")]
+
+
+# --------------------------------------------------------------------------- #
+# Smoke tests: build a tiny model and run real tasks end-to-end on CPU.
+# --------------------------------------------------------------------------- #
+@pytest.fixture
+def lm(tiny_model_dir):
+    from lm_eval.api.registry import get_model
+
+    return get_model("onnxruntime-genai").create_from_arg_string(
+        f"pretrained={tiny_model_dir},execution_provider=cpu",
+        {"batch_size": 1},
+    )
+
+
+def test_backend_loads(lm):
+    assert lm.eot_token_id is not None
+    assert lm.max_length > 0
+    ids = lm.tok_encode("hello world")
+    assert isinstance(ids, list) and len(ids) > 0
+    # Tokenization round-trips (up to special-token stripping).
+    assert "hello" in lm.tok_decode(ids)
+
+
+@pytest.mark.parametrize("task", ["hellaswag", "wikitext", "gsm8k"])
+def test_simple_evaluate(lm, task):
+    from lm_eval import evaluator
+
+    results = evaluator.simple_evaluate(
+        model=lm,
+        tasks=[task],
+        limit=2,
+        num_fewshot=0,
+        bootstrap_iters=0,
+    )
+    assert task in results["results"]
+    metrics = results["results"][task]
+    # Some real metric must be present and finite (no silent 0.0-on-error).
+    numeric = [v for v in metrics.values() if isinstance(v, (int, float))]
+    assert numeric, f"no numeric metrics for {task}: {metrics}"
+    assert all(np.isfinite(v) for v in numeric)
+
+
+def test_loglikelihood_deterministic(lm):
+    """The same request scored twice yields identical log-probs (greedy path)."""
+    from lm_eval.api.instance import Instance
+
+    req = Instance(
+        request_type="loglikelihood",
+        doc={},
+        arguments=("The capital of France is", " Paris"),
+        idx=0,
+    )
+    (first,) = lm.loglikelihood([req], disable_tqdm=True)
+    (second,) = lm.loglikelihood([req], disable_tqdm=True)
+    assert first[0] == pytest.approx(second[0])
+    assert first[0] < 0.0  # a real, non-zero log-prob

tests/models/test_winml.py

@@ -0,0 +1,125 @@
+"""Tests for the winml backend.
+
+``winml`` (:class:`lm_eval.models.winml.WindowsML`) is a thin subclass of
+:class:`lm_eval.models.onnxruntime_genai.ONNXRuntimeGenAILM` that overrides only
+execution-provider selection. The shared scoring/generation logic is covered by
+``test_onnxruntime_genai.py``; this module covers the winml-specific wiring:
+registration, the subclass relationship, and that Windows-provider registration
+degrades gracefully to the cross-platform CPU path when the Windows ML APIs are
+unavailable (e.g. off-Windows), so the inherited scoring core still runs.
+
+Like the genai tests, these are excluded from CI (see the ``--ignore`` in
+``.github/workflows/unit_tests.yml``); run locally with:
+
+    python -m pytest tests/models/test_winml.py -vv
+"""
+
+import sys
+import tempfile
+
+import numpy as np
+import pytest
+
+from lm_eval.models.onnxruntime_genai import ONNXRuntimeGenAILM
+
+
+pytest.importorskip("onnxruntime_genai")
+
+
+def test_winml_registered_and_is_subclass():
+    from lm_eval.api.registry import get_model
+    from lm_eval.models.winml import WindowsML
+
+    assert get_model("winml") is WindowsML
+    assert issubclass(WindowsML, ONNXRuntimeGenAILM)
+    # winml should specialize only provider selection + defaults, not the
+    # scoring core.
+    assert "_select_ep" in vars(WindowsML)
+    for inherited in ("_loglikelihood_tokens", "loglikelihood", "generate_until"):
+        assert inherited not in vars(WindowsML)
+
+
+def test_winml_pins_historical_defaults(tiny_model_dir):
+    # winml keeps its historical max_length / max_gen_toks so existing runs are
+    # numerically unchanged, independent of the base backend's modern defaults.
+    from lm_eval.api.registry import get_model
+
+    lm = get_model("winml").create_from_arg_string(
+        f"pretrained={tiny_model_dir}", {"batch_size": 1}
+    )
+    assert lm.max_length == 4096
+    assert lm.max_gen_toks == 4096
+    # explicit overrides still win
+    lm2 = get_model("winml").create_from_arg_string(
+        f"pretrained={tiny_model_dir},max_length=1024", {"batch_size": 1}
+    )
+    assert lm2.max_length == 1024
+
+
+def _build_tiny_model_dir():
+    """Build a tiny random Llama Model Builder export (fp32, CPU).
+
+    head_size must be a multiple of 16 for the genai CPU GroupQueryAttention
+    kernel, so hidden_size=128 / num_heads=8 -> head_size=16.
+    """
+    builder = pytest.importorskip("onnxruntime_genai.models.builder")
+    from transformers import AutoTokenizer, LlamaConfig, LlamaForCausalLM
+
+    src = tempfile.mkdtemp(prefix="winml_src_")
+    out = tempfile.mkdtemp(prefix="winml_out_")
+    cache = tempfile.mkdtemp(prefix="winml_cache_")
+
+    tokenizer = AutoTokenizer.from_pretrained(
+        "hf-internal-testing/tiny-random-LlamaForCausalLM"
+    )
+    config = LlamaConfig(
+        vocab_size=tokenizer.vocab_size,
+        hidden_size=128,
+        intermediate_size=256,
+        num_hidden_layers=2,
+        num_attention_heads=8,
+        num_key_value_heads=8,
+        max_position_embeddings=512,
+        tie_word_embeddings=True,
+    )
+    LlamaForCausalLM(config).save_pretrained(src)
+    tokenizer.save_pretrained(src)
+
+    try:
+        builder.create_model(None, src, out, "fp32", "cpu", cache)
+    except Exception as e:  # noqa: BLE001
+        pytest.skip(f"Model Builder could not create a CPU fixture: {e}")
+    return out
+
+
+@pytest.fixture(scope="session")
+def tiny_model_dir():
+    return _build_tiny_model_dir()
+
+
+@pytest.mark.skipif(
+    sys.platform == "win32",
+    reason="Exercises the non-Windows fallback path; on Windows the Windows ML "
+    "catalog drives provider selection and requires a real NPU/GPU device.",
+)
+def test_winml_falls_back_to_cpu_off_windows(tiny_model_dir):
+    """Off-Windows, winml should load via the base CPU path and score correctly."""
+    from lm_eval.api.registry import get_model
+
+    lm = get_model("winml").create_from_arg_string(
+        f"pretrained={tiny_model_dir},execution_provider=cpu",
+        {"batch_size": 1},
+    )
+    from lm_eval import evaluator
+
+    results = evaluator.simple_evaluate(
+        model=lm,
+        tasks=["hellaswag"],
+        limit=2,
+        num_fewshot=0,
+        bootstrap_iters=0,
+    )
+    metrics = results["results"]["hellaswag"]
+    numeric = [v for v in metrics.values() if isinstance(v, (int, float))]
+    assert numeric
+    assert all(np.isfinite(v) for v in numeric)