Chore: Run the API tests on torchnative in CI, with capability skips

.github/workflows/torchnative-tests.yml

@@ -0,0 +1,71 @@
+name: torchnative tests
+
+# The same api/tests as api-tests.yml, on torchnative instead of upstream PyTorch (SPEC S1.11,
+# AGENTS 12). Non-blocking: it is neither a required check nor run on pull requests until
+# torchnative is ready. Tests that need a capability the installed torchnative lacks (continuous
+# batching, q8_0) skip with a reason naming it; `-rs` prints every skip reason. Nothing skips on
+# upstream PyTorch.
+on:
+  workflow_dispatch:
+    inputs:
+      torchnative:
+        description: "torchnative version spec, e.g. 0.1.0b4 or >=0.1.0b4 (empty: the latest pre-release)"
+        required: false
+        default: ""
+  schedule:
+    - cron: "0 5 * * 1"  # weekly, Monday 05:00 UTC
+  # workflow_dispatch and schedule only fire once this file is on the default branch (main), which
+  # develop reaches through the release PR. Until then, push to ci/torchnative to run it on demand
+  # (e.g. git push origin develop:ci/torchnative). This work branch triggers it once for a first run.
+  push:
+    branches: [ci/torchnative, work/torchnative-ci]
+
+permissions:
+  contents: read
+
+concurrency:
+  group: torchnative-tests-${{ github.ref }}
+  cancel-in-progress: true
+
+jobs:
+  pytest:
+    runs-on: ubuntu-latest
+    timeout-minutes: 45
+    continue-on-error: true
+    env:
+      HF_HOME: ${{ github.workspace }}/.hf-cache
+      GEMSTONE_TEST_MODEL: HuggingFaceTB/SmolLM2-135M
+      GEMSTONE_TEST_ALLOW_DOWNLOAD: "1"
+      TORCHNATIVE_SPEC: ${{ inputs.torchnative }}
+    steps:
+      - uses: actions/checkout@v4
+
+      - uses: astral-sh/setup-uv@v6
+        with:
+          python-version: "3.13"
+
+      - name: Cache the Hugging Face model
+        uses: actions/cache@v4
+        with:
+          path: .hf-cache
+          key: hf-${{ env.GEMSTONE_TEST_MODEL }}-v1
+
+      - name: Install (torchnative)
+        # torchnative's Linux wheel is manylinux_2_17_x86_64, cp313-abi3. It is a pre-release:
+        # pyproject sets `prerelease = "if-necessary"`, and a pinned spec with a pre-release
+        # version allows it too.
+        run: |
+          if [ -n "$TORCHNATIVE_SPEC" ]; then
+            case "$TORCHNATIVE_SPEC" in
+              [0-9]*) spec="torchnative==$TORCHNATIVE_SPEC" ;;
+              *) spec="torchnative$TORCHNATIVE_SPEC" ;;
+            esac
+            uv sync --extra torchnative --upgrade-package torchnative
+            uv pip install --prerelease=allow "$spec"
+          else
+            uv sync --extra torchnative --upgrade-package torchnative
+          fi
+          uv run --no-sync python -c "import importlib.metadata as m, torch; print('torchnative', m.version('torchnative'), 'torch', torch.__version__)"
+
+      - name: Test
+        run: uv run --no-sync pytest -v -m "real_model or not real_model" -rs

AGENTS.md

@@ -173,6 +173,14 @@ uv run --extra torch pytest                     # api/tests without the real-mod
 uv run --extra torch pytest -m "real_model or not real_model"   # all, as CI runs them (SmolLM2-135M)
 ```
 
+Run on torchnative by the non-blocking `.github/workflows/torchnative-tests.yml` (manual, with a
+torchnative version input, and weekly; `-rs` lists skips). There a test marked `needs_batching`
+(`engine.batching` is False; the reason is `engine.batching_unavailable`) or `needs_quantization`
+(`torchnative.quant` lacks `TorchnativeConfig` / q8_0) skips with a reason naming that capability and
+the torchnative version, and `torchnative_only` tests (q8_0 engine) exist only there. On upstream
+PyTorch nothing skips and `torchnative_only` tests are not collected (`api/tests/test_substrate.py`).
+Any other failure on torchnative is a real failure.
+
 The tests load a real model. `GEMSTONE_TEST_MODEL` picks it (default `HuggingFaceTB/SmolLM2-135M`);
 it is read from the local Hugging Face cache and never downloaded unless
 `GEMSTONE_TEST_ALLOW_DOWNLOAD=1` (CI sets it). Model inference is heavy: run at most one such job

api/tests/conftest.py

@@ -8,7 +8,14 @@
 the local Hugging Face cache (`HF_HOME`, default `~/.cache/huggingface`) and never downloaded,
 unless `GEMSTONE_TEST_ALLOW_DOWNLOAD=1` is set. CI sets that flag and caches `HF_HOME`, so CI and
 local runs use the same checkpoint.
+
+Substrate and capabilities (SPEC S1.11): `import torch` is upstream PyTorch or torchnative. On
+torchnative a test that needs a capability the installed build lacks is skipped with a reason that
+names the capability (markers `needs_batching`, `needs_quantization`); on upstream PyTorch nothing
+is ever skipped for a capability. Tests marked `torchnative_only` are not collected on upstream.
+Anything else fails loudly on either substrate.
 """
+import importlib.metadata
 import os
 
 import pytest
@@ -26,6 +33,70 @@
 )
 
 
+def substrate():
+    """ `("torchnative", version)` when `import torch` is torchnative, else `("torch", torch version)`. """
+    try:
+        return "torchnative", importlib.metadata.version("torchnative")
+    except importlib.metadata.PackageNotFoundError:
+        pass
+    try:
+        return "torch", importlib.metadata.version("torch")
+    except importlib.metadata.PackageNotFoundError:
+        return "torch", "unknown"
+
+
+def is_torchnative():
+    return substrate()[0] == "torchnative"
+
+
+def quantization_unavailable(fmt="q8_0"):
+    """ Why torchnative cannot quantise to `fmt` (None: it can). """
+    try:
+        from torchnative.quant import FORMATS, TorchnativeConfig  # noqa: F401
+    except ImportError as e:
+        return f"torchnative.quant.TorchnativeConfig is missing ({e})"
+    if fmt not in FORMATS():
+        return f"torchnative.quant does not list the {fmt} format"
+    return None
+
+
+def capability_skip_reason(capability, engine=None):
+    """
+    The reason to skip a test that needs `capability`, or None to run it. Always None on upstream
+    PyTorch: there a missing capability is a failure, not a skip.
+    """
+    kind, version = substrate()
+    if kind != "torchnative":
+        return None
+    if capability == "batching":
+        if engine is not None and not engine.batching:
+            return f"continuous batching unavailable on torchnative {version}: {engine.batching_unavailable}"
+    elif capability == "quantization":
+        why = quantization_unavailable("q8_0")
+        if why:
+            return f"q8_0 quantisation unavailable on torchnative {version}: {why}"
+    return None
+
+
+@pytest.fixture(autouse=True)
+def _capabilities(request):
+    """ Skip a `needs_*` test, naming the missing capability, when torchnative lacks it. """
+    if not is_torchnative():
+        return
+    if request.node.get_closest_marker("needs_quantization"):
+        reason = capability_skip_reason("quantization")
+        if reason:
+            pytest.skip(reason)
+    if request.node.get_closest_marker("needs_batching"):
+        # Batching is a property of the engine under test, so the engine is built first.
+        for name in ("q8_engine", "engine"):
+            if name in request.fixturenames:
+                reason = capability_skip_reason("batching", request.getfixturevalue(name))
+                if reason:
+                    pytest.skip(reason)
+                break
+
+
 @pytest.fixture(scope="session")
 def engine():
     import torch
@@ -49,7 +120,15 @@ def engine():
 
 
 def pytest_collection_modifyitems(config, items):
-    """ Mark every test that loads the real model, so it is deselected by default (see pyproject). """
+    """
+    Mark every test that loads the real model, so it is deselected by default (see pyproject), and
+    drop the `torchnative_only` tests when `import torch` is upstream PyTorch.
+    """
     for item in items:
         if "engine" in getattr(item, "fixturenames", ()):
             item.add_marker(pytest.mark.real_model)
+    if not is_torchnative():
+        dropped = [item for item in items if item.get_closest_marker("torchnative_only")]
+        if dropped:
+            config.hook.pytest_deselected(items=dropped)
+            items[:] = [item for item in items if item not in dropped]

api/tests/test_batching.py

@@ -9,6 +9,11 @@
 import threading
 import time
 
+import pytest
+
+# On torchnative these skip, naming the reason, when the build cannot batch (conftest.py).
+pytestmark = pytest.mark.needs_batching
+
 PROMPTS = [
     [{"role": "user", "content": "The capital of France is"}],
     [{"role": "user", "content": "Once upon a time, there was"}],

api/tests/test_quantization.py

@@ -0,0 +1,48 @@
+"""
+q8_0 weights on torchnative (SPEC S1.11): the engine loads the model through torchnative's
+`TorchnativeConfig("q8_0")`, streams, and its output does not depend on how requests are scheduled.
+Not collected on upstream PyTorch (`torchnative_only`); skipped on a torchnative build without
+q8_0 (`needs_quantization`), or, for the concurrent test, without batching (`needs_batching`).
+"""
+import pytest
+
+from api.tests.conftest import ALLOW_DOWNLOAD, CHAT_TEMPLATE, TEST_MODEL
+from api.tests.test_batching import PROMPTS, generate, run_concurrently
+
+pytestmark = [pytest.mark.real_model, pytest.mark.needs_quantization, pytest.mark.torchnative_only]
+
+
+@pytest.fixture(scope="module")
+def q8_engine():
+    import torch
+    from api.src.main.engine import Engine
+
+    engine = Engine(
+        TEST_MODEL,
+        dtype=torch.float32,
+        quantization="q8_0",
+        chat_template=CHAT_TEMPLATE,
+        local_files_only=not ALLOW_DOWNLOAD,
+    )
+    yield engine
+    engine.close()
+
+
+def test_q8_0_engine_streams_text_in_several_chunks(q8_engine):
+    chunks = list(q8_engine(PROMPTS[0], temperature=0, max_new_tokens=16))
+
+    assert len(chunks) > 1
+    assert "".join(chunks).strip()
+
+
+def test_q8_0_greedy_output_is_repeatable(q8_engine):
+    assert generate(q8_engine, PROMPTS[0]) == generate(q8_engine, PROMPTS[0])
+
+
+@pytest.mark.needs_batching
+def test_q8_0_concurrent_greedy_requests_equal_sequential(q8_engine):
+    expected = [generate(q8_engine, messages) for messages in PROMPTS[:2]]
+
+    chunks, _ = run_concurrently(q8_engine, [(messages, {}) for messages in PROMPTS[:2]])
+
+    assert ["".join(c) for c in chunks] == expected

api/tests/test_substrate.py

@@ -0,0 +1,56 @@
+"""
+Capability-based skips (SPEC S1.11): on torchnative a missing capability skips with a reason that
+names it; on upstream PyTorch nothing is ever skipped for a capability. No model is loaded.
+"""
+import pytest
+
+from api.tests import conftest
+from api.tests.conftest import capability_skip_reason, substrate
+
+
+class Engine:
+    def __init__(self, batching, why=None):
+        self.batching = batching
+        self.batching_unavailable = why
+
+
+def test_substrate_is_named_with_a_version():
+    kind, version = substrate()
+
+    assert kind in ("torch", "torchnative")
+    assert version
+
+
+@pytest.mark.parametrize("capability", ["batching", "quantization"])
+def test_nothing_skips_on_upstream_pytorch(monkeypatch, capability):
+    monkeypatch.setattr(conftest, "substrate", lambda: ("torch", "2.9.0"))
+    monkeypatch.setattr(conftest, "quantization_unavailable", lambda fmt="q8_0": "missing")
+
+    assert capability_skip_reason(capability, Engine(False, "no batching")) is None
+
+
+def test_missing_batching_on_torchnative_skips_with_the_reason_and_version(monkeypatch):
+    monkeypatch.setattr(conftest, "substrate", lambda: ("torchnative", "0.1.0b4"))
+
+    reason = capability_skip_reason("batching", Engine(False, "psutil is not installed"))
+
+    assert "continuous batching" in reason and "0.1.0b4" in reason and "psutil is not installed" in reason
+    assert capability_skip_reason("batching", Engine(True)) is None
+
+
+def test_missing_quantization_on_torchnative_skips_naming_q8_0(monkeypatch):
+    monkeypatch.setattr(conftest, "substrate", lambda: ("torchnative", "0.1.0b4"))
+    monkeypatch.setattr(conftest, "quantization_unavailable", lambda fmt="q8_0": "no TorchnativeConfig")
+
+    reason = capability_skip_reason("quantization")
+
+    assert "q8_0" in reason and "0.1.0b4" in reason and "no TorchnativeConfig" in reason
+
+
+def test_capability_markers_are_registered():
+    # Every `needs_*` marker is registered, so a typo cannot silently turn a skip into a pass.
+    import tomllib
+    from pathlib import Path
+
+    markers = tomllib.loads((Path(__file__).parents[2] / "pyproject.toml").read_text())["tool"]["pytest"]["ini_options"]["markers"]
+    assert {m.split(":")[0] for m in markers} >= {"needs_batching", "needs_quantization", "torchnative_only"}

docs/SPEC.md

@@ -304,6 +304,15 @@ Implemented (#84), `api/src/main/engine.py`:
 - The substrate is chosen at install time: `uv sync --extra torch` (upstream PyTorch) or
   `--extra torchnative`.
 
+Capability-based skips: the CI job `.github/workflows/torchnative-tests.yml` (manual with a
+torchnative version input, weekly, non-blocking) runs `api/tests` on torchnative. `conftest.py`
+detects the substrate (torchnative and its version, or upstream PyTorch). On torchnative only,
+`needs_batching` skips when `engine.batching` is False (reason: `engine.batching_unavailable` and the
+torchnative version; all of `test_batching.py` and the concurrent q8_0 test) and `needs_quantization`
+skips when `torchnative.quant` has no `TorchnativeConfig` or q8_0 (`test_quantization.py`, which is
+`torchnative_only`: not collected on upstream). On upstream PyTorch nothing skips for a capability
+(`test_substrate.py` guards it); every other failure is loud.
+
 Tests: `api/tests/test_queue_budget.py` (fake tokenizer and batcher: the budget rule, the queue
 timeout and its 503 / 1013 mapping), `api/tests/test_engine.py`, `api/tests/test_server.py`, on upstream PyTorch with
 SmolLM2-135M. Still open for M1: the same tests on torchnative (cpu, mps) once torchnative TN-M1

pyproject.toml

@@ -68,4 +68,9 @@ pythonpath = ["."]
 # Real-model tests are heavy and the machine is shared: they run in CI (`-m "real_model or not real_model"`),
 # not by default. `api/tests/conftest.py` marks every test that uses the `engine` fixture.
 addopts = ["-m", "not real_model"]
-markers = ["real_model: loads or downloads a real model (deselected by default; CI runs it)"]
+markers = [
+    "real_model: loads or downloads a real model (deselected by default; CI runs it)",
+    "needs_batching: needs continuous batching; skipped (reason given) on torchnative when engine.batching is False, never skipped on upstream PyTorch",
+    "needs_quantization: needs torchnative's q8_0 quantisation; skipped (reason given) on torchnative when it is missing, never skipped on upstream PyTorch",
+    "torchnative_only: not collected when `import torch` is upstream PyTorch",
+]