From a01aedfe94b0799e4bd13e8c90621805c97fbf49 Mon Sep 17 00:00:00 2001
From: BREW <81044362+b-re-w@users.noreply.github.com>
Date: Sat, 3 Oct 2026 18:30:09 +0900
Subject: [PATCH] Test: Let the real-model tests run on another device

GEMSTONE_TEST_DEVICE (default cpu) picks the device of the real-model
test engines, so the mps half of #84 can be checked locally; the
greedy-vs-generate test moves its inputs to the model's device.
---
 AGENTS.md                      | 1 +
 api/tests/conftest.py          | 3 +++
 api/tests/test_engine.py       | 2 +-
 api/tests/test_quantization.py | 3 ++-
 4 files changed, 7 insertions(+), 2 deletions(-)

diff --git a/AGENTS.md b/AGENTS.md
index d15ec971..15448c3e 100644
--- a/AGENTS.md
+++ b/AGENTS.md
@@ -171,6 +171,7 @@ uv sync --extra torch                     # upstream PyTorch; or --extra torchna
 uv run --extra torch python -m api run server   # serve on 127.0.0.1:23100 (GEMSTONE_HOST=host[:port] overrides)
 uv run --extra torch pytest                     # api/tests without the real-model tests (marker real_model)
 uv run --extra torch pytest -m "real_model or not real_model"   # all, as CI runs them (SmolLM2-135M)
+GEMSTONE_TEST_DEVICE=mps uv run --extra torch pytest -m real_model   # the real-model tests on another device
 ```
 
 Run on torchnative by the non-blocking `.github/workflows/torchnative-tests.yml` (manual, with a
diff --git a/api/tests/conftest.py b/api/tests/conftest.py
index 88896cb4..dbf07481 100644
--- a/api/tests/conftest.py
+++ b/api/tests/conftest.py
@@ -23,6 +23,8 @@
 
 TEST_MODEL = os.environ.get("GEMSTONE_TEST_MODEL", "HuggingFaceTB/SmolLM2-135M")
 ALLOW_DOWNLOAD = os.environ.get("GEMSTONE_TEST_ALLOW_DOWNLOAD") == "1"
+# Device of the real-model engines (`cpu` by default; e.g. `mps` to check Apple GPUs locally).
+TEST_DEVICE = os.environ.get("GEMSTONE_TEST_DEVICE", "cpu")
 
 # SmolLM2-135M is a base model and ships without a chat template.
 CHAT_TEMPLATE = (
@@ -109,6 +111,7 @@ def engine():
             # stated in float32. In the checkpoint's bfloat16, top logits tie exactly often enough
             # that a one-ulp difference between batch shapes flips the argmax.
             dtype=torch.float32,
+            device=TEST_DEVICE,
             chat_template=CHAT_TEMPLATE,
             local_files_only=not ALLOW_DOWNLOAD,
         )
diff --git a/api/tests/test_engine.py b/api/tests/test_engine.py
index b0908686..2446411b 100644
--- a/api/tests/test_engine.py
+++ b/api/tests/test_engine.py
@@ -29,7 +29,7 @@ def test_greedy_output_equals_transformers_generate(engine):
     prompt = tokenizer.apply_chat_template(
         MESSAGES, chat_template=CHAT_TEMPLATE, add_generation_prompt=True, tokenize=False
     )
-    inputs = tokenizer(prompt, return_tensors="pt", add_special_tokens=False)
+    inputs = tokenizer(prompt, return_tensors="pt", add_special_tokens=False).to(engine.model.device)
     output = engine.model.generate(**inputs, max_new_tokens=16, do_sample=False)
     expected = tokenizer.decode(output[0, inputs["input_ids"].shape[1]:], skip_special_tokens=True)
 
diff --git a/api/tests/test_quantization.py b/api/tests/test_quantization.py
index 360c765e..59e8873e 100644
--- a/api/tests/test_quantization.py
+++ b/api/tests/test_quantization.py
@@ -6,7 +6,7 @@
 """
 import pytest
 
-from api.tests.conftest import ALLOW_DOWNLOAD, CHAT_TEMPLATE, TEST_MODEL
+from api.tests.conftest import ALLOW_DOWNLOAD, CHAT_TEMPLATE, TEST_DEVICE, 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]
@@ -20,6 +20,7 @@ def q8_engine():
     engine = Engine(
         TEST_MODEL,
         dtype=torch.float32,
+        device=TEST_DEVICE,
         quantization="q8_0",
         chat_template=CHAT_TEMPLATE,
         local_files_only=not ALLOW_DOWNLOAD,