Test: Let the real-model tests run on another device

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

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,
         )

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)
 

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,