test: refactor tests and fix inference edge cases

- Convert protocol and MoE test classes to plain functions
- Add real server/engine integration and generate_async tests
- Isolate test model per test and use pytest tmp_path
- Reset FastAPI engine state after inference tests
- Fix generate_async StopIteration handling on Python 3.12
- Fix HF adapter MoE dense/shared and Gemma qk_norm mapping
- Correct dev dependency httpx2 to httpx

astrai/inference/engine.py

@@ -156,9 +156,8 @@ class InferenceEngine:
         async def _agen():
             loop = asyncio.get_event_loop()
             while True:
-                try:
+                token = await loop.run_in_executor(None, next, sync_gen, None)
-                    token = await loop.run_in_executor(None, next, sync_gen)
+                if token is None:
-                except StopIteration:
                     break
                 yield token
 

astrai/serialization/hf_adapter.py

@@ -44,6 +44,8 @@ HF_MODEL_TYPES = frozenset(
 
 _EMBED = re.compile(r"^model\.embed_tokens\.weight$")
 _ATTN = re.compile(r"^model\.layers\.(\d+)\.self_attn\.(q|k|v|o)_proj\.(weight|bias)$")
+_Q_NORM = re.compile(r"^model\.layers\.(\d+)\.self_attn\.q_norm\.weight$")
+_K_NORM = re.compile(r"^model\.layers\.(\d+)\.self_attn\.k_norm\.weight$")
 _INPUT_NORM = re.compile(r"^model\.layers\.(\d+)\.input_layernorm\.weight$")
 _POST_NORM = re.compile(r"^model\.layers\.(\d+)\.post_attention_layernorm\.weight$")
 _FINAL_NORM = re.compile(r"^model\.norm\.weight$")
@@ -56,7 +58,7 @@ _MOE_EXPERTS = re.compile(
     r"^model\.layers\.(\d+)\.mlp\.experts\.(\d+)\.(gate|up|down)_proj\.(weight|bias)$"
 )
 _MOE_SHARED = re.compile(
-    r"^model\.layers\.(\d+)\.mlp\.shared_experts\.(\d+)\."
+    r"^model\.layers\.(\d+)\.mlp\.shared_expert(?:s)?\.(\d+)\."
     r"(gate|up|down)_proj\.(weight|bias)$"
 )
 
@@ -74,6 +76,17 @@ def looks_like_hf_state_dict(state_dict: Mapping[str, Any]) -> bool:
     )
 
 
+def _is_dense_mlp_layer(config: BaseConfig, layer_id: int) -> bool:
+    """Return whether a layer uses dense MLP instead of routed experts."""
+    if getattr(config, "ffn_type", "mlp") != "moe":
+        return True
+    mlp_only = getattr(config, "mlp_only_layers", None) or []
+    if layer_id in mlp_only:
+        return True
+    step = getattr(config, "decoder_sparse_step", 1) or 1
+    return step > 1 and (layer_id + 1) % step != 0
+
+
 def adapt_config(raw: Dict[str, Any]) -> Dict[str, Any]:
     """Translate *raw* for AstrAI if it looks like an HF model config."""
     if raw.get("model_type") in HF_MODEL_TYPES:
@@ -112,6 +125,8 @@ def convert_hf_config(raw: Dict[str, Any]) -> Dict[str, Any]:
         "topk_method",
         "norm_topk_prob",
         "moe_aux_loss_coef",
+        "decoder_sparse_step",
+        "mlp_only_layers",
         "neftune_alpha",
     ):
         if key in raw:
@@ -119,6 +134,13 @@ def convert_hf_config(raw: Dict[str, Any]) -> Dict[str, Any]:
 
     if "qk_norm" in raw and "use_qk_norm" not in cfg:
         cfg["use_qk_norm"] = raw["qk_norm"]
+    if (
+        raw.get("model_type") in ("gemma", "gemma2")
+        and "use_qk_norm" not in cfg
+        and "qk_norm" not in raw
+    ):
+        # Gemma/Gemma2 always apply RMSNorm to Q and K before attention.
+        cfg["use_qk_norm"] = True
 
     n_heads = raw.get("num_attention_heads")
     if cfg.get("num_key_value_heads") is None and n_heads is not None:
@@ -205,6 +227,10 @@ def convert_hf_weights(
                             f"layers.{m.group(1)}.mlp.shared_experts.{m.group(2)}."
                             f"{m.group(3)}.{m.group(4)}"
                         )
+            if new_key is None:
+                m = _DENSE_MLP.match(key)
+                if m and _is_dense_mlp_layer(config, int(m.group(1))):
+                    new_key = f"layers.{m.group(1)}.mlp.{m.group(2)}.{m.group(3)}"
         else:
             m = _DENSE_MLP.match(key)
             if m:
@@ -216,6 +242,10 @@ def convert_hf_weights(
                 new_key = (
                     f"layers.{m.group(1)}.attention.{m.group(2)}_proj.{m.group(3)}"
                 )
+            elif (m := _Q_NORM.match(key)) is not None:
+                new_key = f"layers.{m.group(1)}.attention.q_norm.weight"
+            elif (m := _K_NORM.match(key)) is not None:
+                new_key = f"layers.{m.group(1)}.attention.k_norm.weight"
             elif (m := _INPUT_NORM.match(key)) is not None:
                 new_key = f"layers.{m.group(1)}.input_norm.weight"
             elif (m := _POST_NORM.match(key)) is not None:

pyproject.toml

@@ -31,7 +31,7 @@ classifiers = [
 urls = { Homepage = "https://github.com/ViperEkura/AstrAI" }
 
 [project.optional-dependencies]
-dev = ["pytest==9.0.2", "ruff", "httpx2"]
+dev = ["pytest==9.0.2", "ruff", "httpx"]
 flash = ["flash-attn>=2.6"]
 
 [tool.setuptools.packages.find]

tests/conftest.py

@@ -1,7 +1,5 @@
 import json
 import os
-import shutil
-import tempfile
 
 import pytest
 import torch
@@ -44,20 +42,18 @@ def test_tokenizer():
     return create_test_tokenizer()
 
 
-@pytest.fixture(scope="session")
+@pytest.fixture
 def test_model(device):
-    """Session-scoped small AutoRegressiveLM model, created once."""
+    """Function-scoped small AutoRegressiveLM model, isolated per test."""
     config = make_tiny_config()
     model = AutoRegressiveLM(config).to(device=device)
     return {"model": model, "device": device, "config": config}
 
 
 @pytest.fixture
-def temp_dir():
+def temp_dir(tmp_path):
-    """Function-scoped temporary directory, cleaned up after each test."""
+    """Function-scoped temporary directory, cleaned up by pytest."""
-    d = tempfile.mkdtemp()
+    return str(tmp_path)
-    yield d
-    shutil.rmtree(d, ignore_errors=True)
 
 
 @pytest.fixture

tests/inference/conftest.py

@@ -8,6 +8,13 @@ from fastapi.testclient import TestClient
 from astrai.inference import get_app
 
 
+@pytest.fixture(autouse=True)
+def _cleanup_app_engine():
+    """Reset the lazy FastAPI singleton engine after each inference test."""
+    yield
+    get_app().state.engine = None
+
+
 @pytest.fixture
 def client():
     """Provide a test client for the FastAPI app."""

tests/inference/test_engine.py

@@ -1,5 +1,6 @@
 """Unit tests for GenerateResult accumulator and InferenceEngine.generate()."""
 
+import asyncio
 import threading
 from unittest.mock import MagicMock, patch
 
@@ -156,6 +157,35 @@ def test_engine_generate_streaming_yields_tokens():
         assert tokens == ["t1", "t2"]
 
 
+def test_engine_generate_async_yields_tokens_until_stop():
+    mock_model, mock_tokenizer = _make_engine_mocks(decode="tok")
+    callbacks_saved = []
+
+    def capture_cb(prompt, **kw):
+        callbacks_saved.append(kw.get("stream_callback"))
+
+    with patch("astrai.inference.engine.InferenceScheduler") as MockSched:
+        instance = MockSched.return_value
+        instance.add_task.side_effect = capture_cb
+        instance.remove_task.return_value = []
+
+        eng = InferenceEngine(mock_model, mock_tokenizer, max_batch_size=1)
+        agen = eng.generate_async("hello")
+
+        async def collect():
+            out = []
+            async for token in agen:
+                out.append(token)
+            return out
+
+        cb = callbacks_saved[0]
+        cb("t1")
+        cb("t2")
+        cb(STOP)
+
+        assert asyncio.run(collect()) == ["t1", "t2"]
+
+
 def test_engine_generate_non_streaming_batch():
     mock_model, mock_tokenizer = _make_engine_mocks(decode="r")
 

tests/inference/test_protocol.py

@@ -3,8 +3,6 @@
 import json
 from unittest.mock import MagicMock
 
-import pytest
-
 from astrai.inference.network.anthropic import AnthropicResponseBuilder
 from astrai.inference.network.openai import OpenAIResponseBuilder
 from astrai.inference.network.protocol import GenContext, StopChecker, StopInfo
@@ -34,223 +32,247 @@ def _sse_payloads(events):
     return payloads
 
 
-class TestStopChecker:
+def _make_openai_builder():
-    def test_check_finds_match(self):
+    builder = OpenAIResponseBuilder()
-        sc = StopChecker(["stop", "end"])
+    req = MagicMock()
-        assert sc.check("hello stop world") == "stop"
+    req.messages = [MagicMock(role="user", content="Hello")]
-
+    req.stop = None
-    def test_check_returns_none_when_no_match(self):
+    req.model = "astrai"
-        sc = StopChecker(["stop"])
+    engine = MagicMock()
-        assert sc.check("hello world") is None
+    engine.tokenizer.apply_chat_template.return_value = "Hello"
-
+    builder.prepare(req, engine)
-    def test_check_empty_sequences(self):
+    return builder
-        sc = StopChecker([])
-        assert sc.check("hello") is None
 
 
-class TestGenContext:
+def _make_anthropic_builder():
-    def test_defaults(self):
+    builder = AnthropicResponseBuilder()
-        ctx = GenContext(resp_id="a", created=1, model="m", prompt_tokens=10)
+    req = MagicMock()
-        assert ctx.completion_tokens == 0
+    req.messages = [MagicMock(role="user", content="Hello")]
-
+    req.model = "claude"
-    def test_fields_mutable(self):
+    req.system = None
-        ctx = GenContext(resp_id="a", created=1, model="m", prompt_tokens=10)
+    engine = MagicMock()
-        ctx.completion_tokens = 42
+    engine.tokenizer.apply_chat_template.return_value = "Hello"
-        assert ctx.completion_tokens == 42
+    builder.prepare(req, engine)
+    return builder
 
 
-class TestStopInfo:
+def test_check_finds_match():
-    def test_defaults(self):
+    sc = StopChecker(["stop", "end"])
-        s = StopInfo()
+    assert sc.check("hello stop world") == "stop"
-        assert s.matched is None
-        assert s.body == ""
-        assert s.yielded == ""
-
-    def test_with_values(self):
-        s = StopInfo(matched="stop", body="hello stop", yielded="hello ")
-        assert s.matched == "stop"
-        assert s.body == "hello stop"
-        assert s.yielded == "hello "
 
 
-class TestOpenAIResponseBuilder:
+def test_check_returns_none_when_no_match():
-    @pytest.fixture
+    sc = StopChecker(["stop"])
-    def builder(self):
+    assert sc.check("hello world") is None
-        builder = OpenAIResponseBuilder()
-        req = MagicMock()
-        req.messages = [MagicMock(role="user", content="Hello")]
-        req.stop = None
-        req.model = "astrai"
-        engine = MagicMock()
-        engine.tokenizer.apply_chat_template.return_value = "Hello"
-        builder.prepare(req, engine)
-        return builder
-
-    def test_prepare_returns_prompt_ctx_stops(self, builder):
-        req = MagicMock()
-        req.messages = [MagicMock(role="user", content="Hi")]
-        req.stop = ["END"]
-        req.model = "gpt"
-        engine = MagicMock()
-        engine.tokenizer.apply_chat_template.return_value = "Hi"
-        prompt, ctx, stops = builder.prepare(req, engine)
-        assert prompt == "Hi"
-        assert ctx.model == "gpt"
-        assert ctx.prompt_tokens == 0
-        assert stops == ["END"]
-
-    def test_prepare_no_stop_returns_empty_list(self, builder):
-        req = MagicMock()
-        req.messages = []
-        req.stop = None
-        req.model = "x"
-        engine = MagicMock()
-        engine.tokenizer.apply_chat_template.return_value = ""
-        _, _, stops = builder.prepare(req, engine)
-        assert stops == []
-
-    def test_format_stream_start(self, builder):
-        ctx = _make_ctx()
-        events = builder.format_stream_start(ctx)
-        payloads = _sse_payloads(events)
-        assert len(payloads) == 1
-        p = payloads[0]
-        assert p["object"] == "chat.completion.chunk"
-        assert p["choices"][0]["delta"]["role"] == "assistant"
-        assert p["choices"][0]["finish_reason"] is None
-
-    def test_format_chunk(self, builder):
-        events = builder.format_chunk("hello", body="hello")
-        payload = json.loads(events[0].split("data: ", 1)[1])
-        assert payload["choices"][0]["delta"]["content"] == "hello"
-        assert payload["choices"][0]["finish_reason"] is None
-
-    def test_format_stream_end(self, builder):
-        ctx = _make_ctx(completion_tokens=5)
-        stop = StopInfo(matched="stop")
-        events = builder.format_stream_end(ctx, stop)
-        payloads = _sse_payloads(events)
-        finish = payloads[0]
-        assert finish["choices"][0]["finish_reason"] == "stop"
-        usage = payloads[1]
-        assert usage["completion_tokens"] == 5
-        assert usage["total_tokens"] == 15
-
-    def test_format_response(self, builder):
-        ctx = _make_ctx()
-        stop = StopInfo()
-        resp = builder.format_response(ctx, "hello", stop)
-        assert resp["object"] == "chat.completion"
-        assert resp["choices"][0]["message"]["content"] == "hello"
-        assert resp["usage"]["prompt_tokens"] == 10
 
 
-class TestAnthropicResponseBuilder:
+def test_check_empty_sequences():
-    @pytest.fixture
+    sc = StopChecker([])
-    def builder(self):
+    assert sc.check("hello") is None
-        builder = AnthropicResponseBuilder()
-        req = MagicMock()
-        req.messages = [MagicMock(role="user", content="Hello")]
-        req.model = "claude"
-        engine = MagicMock()
-        engine.tokenizer.apply_chat_template.return_value = "Hello"
-        req.system = None
-        builder.prepare(req, engine)
-        return builder
 
-    def test_prepare_messages(self, builder):
-        req = MagicMock()
-        req.messages = [MagicMock(role="user", content="Hi")]
-        req.model = "claude"
-        req.system = None
-        req.stop_sequences = None
-        engine = MagicMock()
-        engine.tokenizer.apply_chat_template.return_value = "Hi"
-        prompt, ctx, stops = builder.prepare(req, engine)
-        assert prompt == "Hi"
-        assert stops == []
 
-    def test_prepare_with_stop_sequences(self, builder):
+def test_gen_context_defaults():
-        req = MagicMock()
+    ctx = GenContext(resp_id="a", created=1, model="m", prompt_tokens=10)
-        req.messages = []
+    assert ctx.completion_tokens == 0
-        req.model = "x"
-        req.stop_sequences = ["stop", "end"]
-        req.system = None
-        engine = MagicMock()
-        engine.tokenizer.apply_chat_template.return_value = ""
-        _, _, stops = builder.prepare(req, engine)
-        assert stops == ["stop", "end"]
 
-    def test_format_stream_start(self, builder):
-        ctx = _make_ctx(prompt_tokens=3)
-        events = builder.format_stream_start(ctx)
-        payloads = _sse_payloads(events)
-        assert len(payloads) == 2
-        assert payloads[0]["type"] == "message_start"
-        assert payloads[0]["message"]["usage"]["input_tokens"] == 3
-        assert payloads[1]["type"] == "content_block_start"
 
-    def test_format_chunk(self, builder):
+def test_gen_context_fields_mutable():
-        events = builder.format_chunk("tok", body="tok")
+    ctx = GenContext(resp_id="a", created=1, model="m", prompt_tokens=10)
-        payload = json.loads(events[0].split("data: ", 1)[1])
+    ctx.completion_tokens = 42
-        assert payload["type"] == "content_block_delta"
+    assert ctx.completion_tokens == 42
-        assert payload["delta"]["text"] == "tok"
 
-    def test_format_stream_end_no_stop(self, builder):
-        ctx = _make_ctx(completion_tokens=3)
-        stop = StopInfo()
-        events = builder.format_stream_end(ctx, stop)
-        payloads = _sse_payloads(events)
-        # content_block_stop, message_delta, message_stop
-        types = [p["type"] for p in payloads]
-        assert types == ["content_block_stop", "message_delta", "message_stop"]
-        assert payloads[1]["delta"]["stop_reason"] == "end_turn"
 
-    def test_format_stream_end_with_stop_trims_and_emits_remaining(self, builder):
+def test_stop_info_defaults():
-        ctx = _make_ctx(completion_tokens=7)
+    s = StopInfo()
-        stop = StopInfo(
+    assert s.matched is None
-            matched="END",
+    assert s.body == ""
-            body="Hello world END extra",
+    assert s.yielded == ""
-            yielded="Hello ",
-        )
-        events = builder.format_stream_end(ctx, stop)
-        payloads = _sse_payloads(events)
-        # unyielded delta, content_block_stop, message_delta, message_stop
-        types = [p["type"] for p in payloads]
-        assert types == [
-            "content_block_delta",
-            "content_block_stop",
-            "message_delta",
-            "message_stop",
-        ]
-        assert payloads[0]["delta"]["text"] == "world "
-        assert payloads[2]["delta"]["stop_reason"] == "stop_sequence"
-        assert payloads[2]["delta"]["stop_sequence"] == "END"
 
-    def test_format_stream_end_stop_trimmed_already_yielded(self, builder):
-        ctx = _make_ctx()
-        stop = StopInfo(
-            matched="END",
-            body="Hello END",
-            yielded="Hello ",
-        )
-        events = builder.format_stream_end(ctx, stop)
-        payloads = _sse_payloads(events)
-        # No unyielded delta (everything already sent)
-        types = [p["type"] for p in payloads]
-        assert types == ["content_block_stop", "message_delta", "message_stop"]
 
-    def test_format_response_with_stop_trims_content(self, builder):
+def test_stop_info_with_values():
-        ctx = _make_ctx()
+    s = StopInfo(matched="stop", body="hello stop", yielded="hello ")
-        stop = StopInfo(matched="STOP", body="text STOP extra", yielded="text ")
+    assert s.matched == "stop"
-        resp = builder.format_response(ctx, "text STOP extra", stop)
+    assert s.body == "hello stop"
-        assert resp["content"][0]["text"] == "text "
+    assert s.yielded == "hello "
-        assert resp["stop_reason"] == "stop_sequence"
-        assert resp["stop_sequence"] == "STOP"
 
-    def test_format_response_no_stop(self, builder):
+
-        ctx = _make_ctx()
+def test_openai_prepare_returns_prompt_ctx_stops():
-        stop = StopInfo()
+    builder = _make_openai_builder()
-        resp = builder.format_response(ctx, "full text", stop)
+    req = MagicMock()
-        assert resp["content"][0]["text"] == "full text"
+    req.messages = [MagicMock(role="user", content="Hi")]
-        assert resp["stop_reason"] == "end_turn"
+    req.stop = ["END"]
+    req.model = "gpt"
+    engine = MagicMock()
+    engine.tokenizer.apply_chat_template.return_value = "Hi"
+    prompt, ctx, stops = builder.prepare(req, engine)
+    assert prompt == "Hi"
+    assert ctx.model == "gpt"
+    assert ctx.prompt_tokens == 0
+    assert stops == ["END"]
+
+
+def test_openai_prepare_no_stop_returns_empty_list():
+    builder = _make_openai_builder()
+    req = MagicMock()
+    req.messages = []
+    req.stop = None
+    req.model = "x"
+    engine = MagicMock()
+    engine.tokenizer.apply_chat_template.return_value = ""
+    _, _, stops = builder.prepare(req, engine)
+    assert stops == []
+
+
+def test_openai_format_stream_start():
+    builder = _make_openai_builder()
+    ctx = _make_ctx()
+    events = builder.format_stream_start(ctx)
+    payloads = _sse_payloads(events)
+    assert len(payloads) == 1
+    p = payloads[0]
+    assert p["object"] == "chat.completion.chunk"
+    assert p["choices"][0]["delta"]["role"] == "assistant"
+    assert p["choices"][0]["finish_reason"] is None
+
+
+def test_openai_format_chunk():
+    builder = _make_openai_builder()
+    events = builder.format_chunk("hello", body="hello")
+    payload = json.loads(events[0].split("data: ", 1)[1])
+    assert payload["choices"][0]["delta"]["content"] == "hello"
+    assert payload["choices"][0]["finish_reason"] is None
+
+
+def test_openai_format_stream_end():
+    builder = _make_openai_builder()
+    ctx = _make_ctx(completion_tokens=5)
+    stop = StopInfo(matched="stop")
+    events = builder.format_stream_end(ctx, stop)
+    payloads = _sse_payloads(events)
+    finish = payloads[0]
+    assert finish["choices"][0]["finish_reason"] == "stop"
+    usage = payloads[1]
+    assert usage["completion_tokens"] == 5
+    assert usage["total_tokens"] == 15
+
+
+def test_openai_format_response():
+    builder = _make_openai_builder()
+    ctx = _make_ctx()
+    stop = StopInfo()
+    resp = builder.format_response(ctx, "hello", stop)
+    assert resp["object"] == "chat.completion"
+    assert resp["choices"][0]["message"]["content"] == "hello"
+    assert resp["usage"]["prompt_tokens"] == 10
+
+
+def test_anthropic_prepare_messages():
+    builder = _make_anthropic_builder()
+    req = MagicMock()
+    req.messages = [MagicMock(role="user", content="Hi")]
+    req.model = "claude"
+    req.system = None
+    req.stop_sequences = None
+    engine = MagicMock()
+    engine.tokenizer.apply_chat_template.return_value = "Hi"
+    prompt, ctx, stops = builder.prepare(req, engine)
+    assert prompt == "Hi"
+    assert stops == []
+
+
+def test_anthropic_prepare_with_stop_sequences():
+    builder = _make_anthropic_builder()
+    req = MagicMock()
+    req.messages = []
+    req.model = "x"
+    req.stop_sequences = ["stop", "end"]
+    req.system = None
+    engine = MagicMock()
+    engine.tokenizer.apply_chat_template.return_value = ""
+    _, _, stops = builder.prepare(req, engine)
+    assert stops == ["stop", "end"]
+
+
+def test_anthropic_format_stream_start():
+    builder = _make_anthropic_builder()
+    ctx = _make_ctx(prompt_tokens=3)
+    events = builder.format_stream_start(ctx)
+    payloads = _sse_payloads(events)
+    assert len(payloads) == 2
+    assert payloads[0]["type"] == "message_start"
+    assert payloads[0]["message"]["usage"]["input_tokens"] == 3
+    assert payloads[1]["type"] == "content_block_start"
+
+
+def test_anthropic_format_chunk():
+    builder = _make_anthropic_builder()
+    events = builder.format_chunk("tok", body="tok")
+    payload = json.loads(events[0].split("data: ", 1)[1])
+    assert payload["type"] == "content_block_delta"
+    assert payload["delta"]["text"] == "tok"
+
+
+def test_anthropic_format_stream_end_no_stop():
+    builder = _make_anthropic_builder()
+    ctx = _make_ctx(completion_tokens=3)
+    stop = StopInfo()
+    events = builder.format_stream_end(ctx, stop)
+    payloads = _sse_payloads(events)
+    types = [p["type"] for p in payloads]
+    assert types == ["content_block_stop", "message_delta", "message_stop"]
+    assert payloads[1]["delta"]["stop_reason"] == "end_turn"
+
+
+def test_anthropic_format_stream_end_with_stop_trims_and_emits_remaining():
+    builder = _make_anthropic_builder()
+    ctx = _make_ctx(completion_tokens=7)
+    stop = StopInfo(
+        matched="END",
+        body="Hello world END extra",
+        yielded="Hello ",
+    )
+    events = builder.format_stream_end(ctx, stop)
+    payloads = _sse_payloads(events)
+    types = [p["type"] for p in payloads]
+    assert types == [
+        "content_block_delta",
+        "content_block_stop",
+        "message_delta",
+        "message_stop",
+    ]
+    assert payloads[0]["delta"]["text"] == "world "
+    assert payloads[2]["delta"]["stop_reason"] == "stop_sequence"
+    assert payloads[2]["delta"]["stop_sequence"] == "END"
+
+
+def test_anthropic_format_stream_end_stop_trimmed_already_yielded():
+    builder = _make_anthropic_builder()
+    ctx = _make_ctx()
+    stop = StopInfo(
+        matched="END",
+        body="Hello END",
+        yielded="Hello ",
+    )
+    events = builder.format_stream_end(ctx, stop)
+    payloads = _sse_payloads(events)
+    types = [p["type"] for p in payloads]
+    assert types == ["content_block_stop", "message_delta", "message_stop"]
+
+
+def test_anthropic_format_response_with_stop_trims_content():
+    builder = _make_anthropic_builder()
+    ctx = _make_ctx()
+    stop = StopInfo(matched="STOP", body="text STOP extra", yielded="text ")
+    resp = builder.format_response(ctx, "text STOP extra", stop)
+    assert resp["content"][0]["text"] == "text "
+    assert resp["stop_reason"] == "stop_sequence"
+    assert resp["stop_sequence"] == "STOP"
+
+
+def test_anthropic_format_response_no_stop():
+    builder = _make_anthropic_builder()
+    ctx = _make_ctx()
+    stop = StopInfo()
+    resp = builder.format_response(ctx, "full text", stop)
+    assert resp["content"][0]["text"] == "full text"
+    assert resp["stop_reason"] == "end_turn"

tests/inference/test_server.py

@@ -1,8 +1,15 @@
 """Unit tests for the inference HTTP server."""
 
+from pathlib import Path
+
 import pytest
+import torch
 
 from astrai.inference import get_app
+from astrai.inference.network.app import _create_engine
+from astrai.model.transformer import AutoRegressiveLM
+from astrai.serialization import save_model
+from tests.helpers import CHAT_TEMPLATE, build_test_tokenizer, make_tiny_config
 
 
 def test_health_no_model(client):
@@ -212,5 +219,42 @@ def test_chat_completions_stop_sequence_stream(client, loaded_model):
     )
 
 
+def test_chat_completions_real_engine(tmp_path, client):
+    """POST /v1/chat/completions with a real tiny model and tokenizer."""
+    cfg = make_tiny_config(vocab_size=256)
+    model = AutoRegressiveLM(cfg).eval()
+    save_model(cfg.to_dict(), model.state_dict(), str(tmp_path))
+
+    tokenizer = build_test_tokenizer(vocab_size=256, chat_template=CHAT_TEMPLATE)
+    tokenizer.save_pretrained(str(tmp_path))
+
+    engine = _create_engine(
+        Path(tmp_path),
+        device="cpu",
+        dtype=torch.float32,
+        max_batch_size=1,
+        max_seq_len=64,
+    )
+    try:
+        get_app().state.engine = engine
+        response = client.post(
+            "/v1/chat/completions",
+            json={
+                "messages": [{"role": "user", "content": "Hello"}],
+                "max_tokens": 4,
+                "temperature": 0.0,
+                "stream": False,
+            },
+        )
+        assert response.status_code == 200
+        data = response.json()
+        content = data["choices"][0]["message"]["content"]
+        assert isinstance(content, str)
+        assert data["usage"]["completion_tokens"] > 0
+    finally:
+        engine.shutdown()
+        get_app().state.engine = None
+
+
 if __name__ == "__main__":
     pytest.main([__file__, "-v"])

tests/module/test_model_forward.py

@@ -54,6 +54,20 @@ def _make_batch(config, batch_size=2, seq_len=8, with_extra=False):
     return batch
 
 
+def _make_seq_moe_fixture(device):
+    config = _make_tiny_moe_config()
+    model = _make_model(config).to(device)
+    model.train()
+    return config, model
+
+
+def _make_sft_moe_fixture(device):
+    config = _make_tiny_moe_config()
+    model = _make_model(config).to(device)
+    model.train()
+    return config, model
+
+
 def test_model_forward_contract_uses_dense_training_and_packed_inference():
     from astrai.inference.cache import PagePool, TaskCacheManager
     from astrai.inference.workspace import InferenceWorkspace
@@ -185,172 +199,166 @@ def test_moe_metrics_flow_through_wrapped_model(device):
     assert "router_entropy" in strategy._moe_metrics
 
 
-class TestSEQStrategyMoE:
+def test_seq_compute_loss_returns_scalar(device):
-    """End‑to‑end tests for SEQStrategy with MoE aux loss."""
+    """compute_loss should return a scalar tensor."""
-
+    config, model = _make_seq_moe_fixture(device)
-    @pytest.fixture(autouse=True)
+    strategy = SEQStrategy(
-    def setup(self, device):
+        model,
-        self.device = device
+        device,
-        self.config = _make_tiny_moe_config()
+        moe_aux_loss_coef=0.01,
-        self.model = _make_model(self.config).to(device)
+    )
-        self.model.train()
+    loss = strategy.compute_loss(_make_batch(config))
-
+    assert loss.ndim == 0
-    def test_compute_loss_returns_scalar(self):
+    assert loss.requires_grad
-        """compute_loss should return a scalar tensor."""
-        strategy = SEQStrategy(
-            self.model,
-            self.device,
-            moe_aux_loss_coef=0.01,
-        )
-        loss = strategy.compute_loss(_make_batch(self.config))
-        assert loss.ndim == 0
-        assert loss.requires_grad
-
-    def test_compute_loss_output_has_metrics(self):
-        """compute_loss_output dict with moe_aux_loss_coef > 0 includes MoE metrics."""
-        strategy = SEQStrategy(
-            self.model,
-            self.device,
-            moe_aux_loss_coef=0.01,
-        )
-        output = strategy.compute_loss_output(_make_batch(self.config))
-
-        assert "loss" in output
-        assert "metrics" in output
-        assert output["loss"].ndim == 0
-        assert output["loss"].requires_grad
-
-        metrics = output["metrics"]
-        # MoE metrics should appear when coef > 0 and model has MoE layers
-        for key in ("moe_aux_loss", "moe_aux_loss_weighted", "task_loss", "loss"):
-            assert key in metrics, f"Missing metric: {key}"
-            assert isinstance(metrics[key], float)
-
-    def test_moe_metrics_populated_after_forward(self):
-        """strategy._moe_metrics populated after compute_loss_output."""
-        strategy = SEQStrategy(
-            self.model,
-            self.device,
-            moe_aux_loss_coef=0.01,
-        )
-        strategy.compute_loss_output(_make_batch(self.config))
-
-        moe_metrics = strategy._moe_metrics
-        assert moe_metrics, "_moe_metrics should not be empty for MoE model"
-        for key in (
-            "aux_loss",
-            "router_entropy",
-            "dead_expert_fraction",
-            "load_imbalance_mean",
-            "load_imbalance_max",
-        ):
-            assert key in moe_metrics, f"Missing _moe_metrics key: {key}"
-            assert isinstance(moe_metrics[key], float)
-
-    def test_zero_coef_zeroes_weighted_aux(self):
-        """moe_aux_loss_coef=0 → weighted_aux_loss is zero, task_loss == loss."""
-        strategy = SEQStrategy(
-            self.model,
-            self.device,
-            moe_aux_loss_coef=0.0,
-        )
-        output = strategy.compute_loss_output(_make_batch(self.config))
-        metrics = output["metrics"]
-
-        # task_loss and loss should be equal (aux weighted by zero)
-        assert "task_loss" in metrics
-        assert "loss" in metrics
-        assert metrics["loss"] == pytest.approx(metrics["task_loss"], abs=1e-6)
-
-        # weighted aux loss is zero
-        assert metrics.get("moe_aux_loss_weighted") == pytest.approx(0.0, abs=1e-6)
-
-        # MoE diagnostics are still collected (monitoring purposes)
-        assert strategy._moe_metrics
-        assert "router_entropy" in strategy._moe_metrics
-
-    def test_aux_loss_added_to_total_loss(self):
-        """Total loss > task_loss when moe_aux_loss_coef > 0."""
-        strategy = SEQStrategy(
-            self.model,
-            self.device,
-            moe_aux_loss_coef=0.01,
-        )
-        output = strategy.compute_loss_output(_make_batch(self.config))
-        assert output["metrics"]["loss"] > output["metrics"]["task_loss"] + 1e-12
-
-    def test_factory_creates_strategy_with_coef(self):
-        """StrategyFactory.create passes moe_aux_loss_coef to strategy."""
-        strategy = StrategyFactory.create(
-            "seq",
-            model=self.model,
-            device=self.device,
-            moe_aux_loss_coef=0.02,
-        )
-        assert strategy.moe_aux_loss_coef == 0.02
-
-    def test_no_aux_loss_for_mlp_model(self):
-        """Pure MLP model: model outputs no aux_loss → no MoE metrics."""
-        from astrai.config.model_config import AutoRegressiveLMConfig
-
-        mlp_config = AutoRegressiveLMConfig(**{**TINY_CONFIG, "ffn_type": "mlp"})
-        mlp_model = AutoRegressiveLM(mlp_config).to(self.device)
-        mlp_model.train()
-
-        strategy = SEQStrategy(
-            mlp_model,
-            self.device,
-            moe_aux_loss_coef=0.01,
-        )
-        output = strategy.compute_loss_output(_make_batch(self.config))
-        metrics = output["metrics"]
-
-        assert "moe_aux_loss" not in metrics
-        assert "moe_aux_loss_weighted" not in metrics
-        assert metrics["loss"] == pytest.approx(metrics["task_loss"], abs=1e-6)
-        assert strategy._moe_metrics == {}
 
 
-class TestSFTStrategyMoE:
+def test_seq_compute_loss_output_has_metrics(device):
-    """End‑to‑end tests for SFTStrategy with MoE aux loss."""
+    """compute_loss_output dict with moe_aux_loss_coef > 0 includes MoE metrics."""
+    config, model = _make_seq_moe_fixture(device)
+    strategy = SEQStrategy(
+        model,
+        device,
+        moe_aux_loss_coef=0.01,
+    )
+    output = strategy.compute_loss_output(_make_batch(config))
 
-    @pytest.fixture(autouse=True)
+    assert "loss" in output
-    def setup(self, device):
+    assert "metrics" in output
-        self.device = device
+    assert output["loss"].ndim == 0
-        self.config = _make_tiny_moe_config()
+    assert output["loss"].requires_grad
-        self.model = _make_model(self.config).to(device)
-        self.model.train()
 
-    def test_compute_loss_output_with_aux_loss(self):
+    metrics = output["metrics"]
-        """SFTStrategy produces MoE metrics when coef > 0."""
+    # MoE metrics should appear when coef > 0 and model has MoE layers
-        strategy = SFTStrategy(
+    for key in ("moe_aux_loss", "moe_aux_loss_weighted", "task_loss", "loss"):
-            self.model,
+        assert key in metrics, f"Missing metric: {key}"
-            self.device,
+        assert isinstance(metrics[key], float)
-            moe_aux_loss_coef=0.01,
-        )
-        output = strategy.compute_loss_output(_make_batch(self.config, with_extra=True))
 
-        metrics = output["metrics"]
-        assert "moe_aux_loss" in metrics
-        assert "moe_aux_loss_weighted" in metrics
-        assert metrics["loss"] > metrics["task_loss"] + 1e-12
 
-        moe_metrics = strategy._moe_metrics
+def test_seq_moe_metrics_populated_after_forward(device):
-        assert "router_entropy" in moe_metrics
+    """strategy._moe_metrics populated after compute_loss_output."""
-        assert "dead_expert_fraction" in moe_metrics
+    config, model = _make_seq_moe_fixture(device)
+    strategy = SEQStrategy(
+        model,
+        device,
+        moe_aux_loss_coef=0.01,
+    )
+    strategy.compute_loss_output(_make_batch(config))
 
-    def test_sft_zero_coef_zeroes_weighted_aux(self):
+    moe_metrics = strategy._moe_metrics
-        """SFTStrategy with zero coef: weighted aux is zero, loss == task_loss."""
+    assert moe_metrics, "_moe_metrics should not be empty for MoE model"
-        strategy = SFTStrategy(
+    for key in (
-            self.model,
+        "aux_loss",
-            self.device,
+        "router_entropy",
-            moe_aux_loss_coef=0.0,
+        "dead_expert_fraction",
-        )
+        "load_imbalance_mean",
-        output = strategy.compute_loss_output(_make_batch(self.config, with_extra=True))
+        "load_imbalance_max",
-        metrics = output["metrics"]
+    ):
+        assert key in moe_metrics, f"Missing _moe_metrics key: {key}"
+        assert isinstance(moe_metrics[key], float)
 
-        assert metrics["loss"] == pytest.approx(metrics["task_loss"], abs=1e-6)
+
-        assert metrics.get("moe_aux_loss_weighted") == pytest.approx(0.0, abs=1e-6)
+def test_seq_zero_coef_zeroes_weighted_aux(device):
-        # Diagnostics still collected
+    """moe_aux_loss_coef=0 → weighted_aux_loss is zero, task_loss == loss."""
-        assert strategy._moe_metrics
+    config, model = _make_seq_moe_fixture(device)
-        assert "router_entropy" in strategy._moe_metrics
+    strategy = SEQStrategy(
+        model,
+        device,
+        moe_aux_loss_coef=0.0,
+    )
+    output = strategy.compute_loss_output(_make_batch(config))
+    metrics = output["metrics"]
+
+    # task_loss and loss should be equal (aux weighted by zero)
+    assert "task_loss" in metrics
+    assert "loss" in metrics
+    assert metrics["loss"] == pytest.approx(metrics["task_loss"], abs=1e-6)
+
+    # weighted aux loss is zero
+    assert metrics.get("moe_aux_loss_weighted") == pytest.approx(0.0, abs=1e-6)
+
+    # MoE diagnostics are still collected (monitoring purposes)
+    assert strategy._moe_metrics
+    assert "router_entropy" in strategy._moe_metrics
+
+
+def test_seq_aux_loss_added_to_total_loss(device):
+    """Total loss > task_loss when moe_aux_loss_coef > 0."""
+    config, model = _make_seq_moe_fixture(device)
+    strategy = SEQStrategy(
+        model,
+        device,
+        moe_aux_loss_coef=0.01,
+    )
+    output = strategy.compute_loss_output(_make_batch(config))
+    assert output["metrics"]["loss"] > output["metrics"]["task_loss"] + 1e-12
+
+
+def test_seq_factory_creates_strategy_with_coef(device):
+    """StrategyFactory.create passes moe_aux_loss_coef to strategy."""
+    _, model = _make_seq_moe_fixture(device)
+    strategy = StrategyFactory.create(
+        "seq",
+        model=model,
+        device=device,
+        moe_aux_loss_coef=0.02,
+    )
+    assert strategy.moe_aux_loss_coef == 0.02
+
+
+def test_seq_no_aux_loss_for_mlp_model(device):
+    """Pure MLP model: model outputs no aux_loss → no MoE metrics."""
+    config, _ = _make_seq_moe_fixture(device)
+    mlp_config = AutoRegressiveLMConfig(**{**TINY_CONFIG, "ffn_type": "mlp"})
+    mlp_model = AutoRegressiveLM(mlp_config).to(device)
+    mlp_model.train()
+
+    strategy = SEQStrategy(
+        mlp_model,
+        device,
+        moe_aux_loss_coef=0.01,
+    )
+    output = strategy.compute_loss_output(_make_batch(config))
+    metrics = output["metrics"]
+
+    assert "moe_aux_loss" not in metrics
+    assert "moe_aux_loss_weighted" not in metrics
+    assert metrics["loss"] == pytest.approx(metrics["task_loss"], abs=1e-6)
+    assert strategy._moe_metrics == {}
+
+
+def test_sft_compute_loss_output_with_aux_loss(device):
+    """SFTStrategy produces MoE metrics when coef > 0."""
+    config, model = _make_sft_moe_fixture(device)
+    strategy = SFTStrategy(
+        model,
+        device,
+        moe_aux_loss_coef=0.01,
+    )
+    output = strategy.compute_loss_output(_make_batch(config, with_extra=True))
+
+    metrics = output["metrics"]
+    assert "moe_aux_loss" in metrics
+    assert "moe_aux_loss_weighted" in metrics
+    assert metrics["loss"] > metrics["task_loss"] + 1e-12
+
+    moe_metrics = strategy._moe_metrics
+    assert "router_entropy" in moe_metrics
+    assert "dead_expert_fraction" in moe_metrics
+
+
+def test_sft_zero_coef_zeroes_weighted_aux(device):
+    """SFTStrategy with zero coef: weighted aux is zero, loss == task_loss."""
+    config, model = _make_sft_moe_fixture(device)
+    strategy = SFTStrategy(
+        model,
+        device,
+        moe_aux_loss_coef=0.0,
+    )
+    output = strategy.compute_loss_output(_make_batch(config, with_extra=True))
+    metrics = output["metrics"]
+
+    assert metrics["loss"] == pytest.approx(metrics["task_loss"], abs=1e-6)
+    assert metrics.get("moe_aux_loss_weighted") == pytest.approx(0.0, abs=1e-6)
+    # Diagnostics still collected
+    assert strategy._moe_metrics
+    assert "router_entropy" in strategy._moe_metrics

tests/serialization/test_hf_adapter.py

@@ -183,6 +183,68 @@ def test_convert_hf_weights_rejects_mla():
         convert_hf_weights(sd, cfg)
 
 
+def test_convert_hf_config_qwen2_moe_preserves_sparse_fields():
+    raw = {
+        **LLAMA_RAW,
+        "model_type": "qwen2_moe",
+        "num_local_experts": 2,
+        "num_experts_per_tok": 1,
+        "n_shared_experts": 1,
+        "decoder_sparse_step": 2,
+        "mlp_only_layers": [0],
+    }
+    cfg = ConfigFactory.load(convert_hf_config(raw))
+    assert cfg.decoder_sparse_step == 2
+    assert cfg.mlp_only_layers == [0]
+
+
+def test_convert_hf_config_gemma_enables_qk_norm():
+    raw = {**LLAMA_RAW, "model_type": "gemma"}
+    cfg = ConfigFactory.load(convert_hf_config(raw))
+    assert cfg.use_qk_norm is True
+
+
+def test_convert_hf_weights_moe_with_dense_layers_roundtrip():
+    cfg = make_tiny_config(
+        ffn_type="moe",
+        n_routed_experts=2,
+        n_shared_experts=1,
+        n_activated_experts=1,
+        moe_intermediate_size=16,
+        shared_expert_intermediate_size=16,
+        mlp_only_layers=[0],
+        decoder_sparse_step=1,
+    )
+    model = AutoRegressiveLM(cfg)
+    converted = convert_hf_weights(to_hf_keys(model.state_dict()), cfg)
+    assert_state_dicts_equal(converted, model.state_dict())
+
+
+def test_convert_hf_weights_qwen2_moe_singular_shared_expert_roundtrip():
+    cfg = make_tiny_config(
+        ffn_type="moe",
+        n_routed_experts=2,
+        n_shared_experts=1,
+        n_activated_experts=1,
+        moe_intermediate_size=16,
+        shared_expert_intermediate_size=16,
+    )
+    model = AutoRegressiveLM(cfg)
+    hf_sd = to_hf_keys(model.state_dict())
+    hf_sd = {
+        k.replace("shared_experts.", "shared_expert.", 1): v for k, v in hf_sd.items()
+    }
+    converted = convert_hf_weights(hf_sd, cfg)
+    assert_state_dicts_equal(converted, model.state_dict())
+
+
+def test_convert_hf_weights_gemma_qk_norm_roundtrip():
+    cfg = make_tiny_config(use_qk_norm=True)
+    model = AutoRegressiveLM(cfg)
+    converted = convert_hf_weights(to_hf_keys(model.state_dict()), cfg)
+    assert_state_dicts_equal(converted, model.state_dict())
+
+
 def test_from_pretrained_hf_directory(tmp_path):
     cfg = make_tiny_config()
     model = AutoRegressiveLM(cfg).eval()

tests/trainer/test_callbacks.py

@@ -31,13 +31,11 @@ def test_gradient_checkpointing_empty_modules_noop(test_model):
     model = test_model["model"]
     callback = GradientCheckpointingCallback()
 
-    originals = [layer.forward for layer in model.layers]
-
     for layer in model.layers:
         callback._enable(layer)
 
-    for layer, orig in zip(model.layers, originals):
+    for layer in model.layers:
-        assert layer.forward is orig
+        assert not hasattr(layer, "_original_forward")
 
 
 def test_gradient_checkpointing_forward_unchanged(test_model):