misc

tests/test_real_model_backend.py

@@ -1,5 +1,6 @@
 """US-012 tests for the real PyTorch node backend."""
 
+from collections import OrderedDict
 import json
 import os
 from pathlib import Path
@@ -14,6 +15,7 @@ import pytest
 from meshnet_node.model_backend import (
     InsufficientVRAMError,
     PartialModelLoadUnsupported,
+    ShardCacheMiss,
     TensorPayload,
     TorchModelShard,
     _call_layer,
@@ -43,7 +45,15 @@ class _FakeBackend:
             position_ids_header=None,
         )
 
-    def forward_bytes(self, body, shape, attention_mask_header, position_ids_header, start_layer=None):
+    def forward_bytes(
+        self,
+        body,
+        shape,
+        attention_mask_header,
+        position_ids_header,
+        start_layer=None,
+        **kwargs,  # noqa: ARG002
+    ):
         assert shape == [1, 6, 8]
         return TensorPayload(
             body=body,
@@ -57,7 +67,15 @@ class _FakeTailBackend(_FakeBackend):
     is_head = False
     is_tail = True
 
-    def forward_bytes(self, body, shape, attention_mask_header, position_ids_header, start_layer=None):
+    def forward_bytes(
+        self,
+        body,
+        shape,
+        attention_mask_header,
+        position_ids_header,
+        start_layer=None,
+        **kwargs,  # noqa: ARG002
+    ):
         assert len(body) == 1 * 6 * 8 * 2
         return " Paris"
 
@@ -114,7 +132,15 @@ class _FakePipelineTailBackend(_FakeTailBackend):
     def __init__(self) -> None:
         self.start_layers: list[int | None] = []
 
-    def forward_bytes(self, body, shape, attention_mask_header, position_ids_header, start_layer=None):
+    def forward_bytes(
+        self,
+        body,
+        shape,
+        attention_mask_header,
+        position_ids_header,
+        start_layer=None,
+        **kwargs,  # noqa: ARG002
+    ):
         self.start_layers.append(start_layer)
         assert len(body) == 1 * 6 * 8 * 2
         return " token"
@@ -125,7 +151,15 @@ class _BlockingStreamingTailBackend(_FakeTailBackend):
         self._release = second_token_release
         self.calls = 0
 
-    def forward_bytes(self, body, shape, attention_mask_header, position_ids_header, start_layer=None):
+    def forward_bytes(
+        self,
+        body,
+        shape,
+        attention_mask_header,
+        position_ids_header,
+        start_layer=None,
+        **kwargs,  # noqa: ARG002
+    ):
         self.calls += 1
         if self.calls == 1:
             return " first"
@@ -488,13 +522,118 @@ def test_call_layer_passes_rotary_position_embeddings():
             assert kwargs["position_embeddings"] == "rotary"
             return hidden_states
 
-    assert _call_layer(
+    hidden, cache_state = _call_layer(
         NeedsPositionEmbeddings(),
         "hidden",
         attention_mask=None,
         position_ids="positions",
         position_embeddings="rotary",
-    ) == "hidden"
+    )
+
+    assert hidden == "hidden"
+    assert cache_state is None
+
+
+def _fake_cache_shard(torch, *, max_sessions=16, ttl=600.0):
+    class RecordingLayer:
+        def __init__(self, index):
+            self.index = index
+            self.calls = []
+
+        def __call__(self, hidden_states, **kwargs):
+            self.calls.append({
+                "shape": tuple(hidden_states.shape),
+                "use_cache": kwargs.get("use_cache"),
+                "past_key_value": kwargs.get("past_key_value"),
+            })
+            present = {
+                "layer": self.index,
+                "shape": tuple(hidden_states.shape),
+                "opaque": object(),
+            }
+            return hidden_states + (self.index + 1), present
+
+    shard = object.__new__(TorchModelShard)
+    shard.shard_start = 0
+    shard.shard_end = 1
+    shard.torch = torch
+    shard.model = types.SimpleNamespace(model=types.SimpleNamespace(layers=[]))
+    shard.layers = [RecordingLayer(0), RecordingLayer(1)]
+    shard._session_cache = OrderedDict()
+    shard._cache_max_sessions = max_sessions
+    shard._cache_ttl_seconds = ttl
+    return shard
+
+
+def test_shard_cache_prefill_then_decode_reuses_opaque_layer_state():
+    torch = pytest.importorskip("torch")
+    shard = _fake_cache_shard(torch)
+
+    prefill_hidden = torch.zeros((1, 4, 2), dtype=torch.bfloat16)
+    prefill_mask = torch.ones((1, 4), dtype=torch.long)
+    prefill_positions = torch.arange(4, dtype=torch.long).reshape(1, 4)
+    shard._run_layers(
+        prefill_hidden,
+        prefill_mask,
+        prefill_positions,
+        session_id="session-1",
+        cache_mode="prefill",
+        seq_len=4,
+    )
+
+    assert len(shard._session_cache) == 1
+    cached_states = next(iter(shard._session_cache.values())).layer_states
+    assert len(cached_states) == 2
+    assert cached_states[0]["shape"] == (1, 4, 2)
+
+    decode_hidden = torch.zeros((1, 1, 2), dtype=torch.bfloat16)
+    decode_mask = torch.ones((1, 5), dtype=torch.long)
+    decode_positions = torch.tensor([[4]], dtype=torch.long)
+    shard._run_layers(
+        decode_hidden,
+        decode_mask,
+        decode_positions,
+        session_id="session-1",
+        cache_mode="decode",
+        seq_len=5,
+    )
+
+    assert shard.layers[0].calls[-1]["shape"] == (1, 1, 2)
+    assert shard.layers[0].calls[-1]["past_key_value"] is cached_states[0]
+    assert shard.layers[1].calls[-1]["past_key_value"] is cached_states[1]
+    assert next(iter(shard._session_cache.values())).seq_len == 5
+
+
+def test_shard_cache_decode_miss_is_explicit():
+    torch = pytest.importorskip("torch")
+    shard = _fake_cache_shard(torch)
+
+    with pytest.raises(ShardCacheMiss):
+        shard._run_layers(
+            torch.zeros((1, 1, 2), dtype=torch.bfloat16),
+            torch.ones((1, 5), dtype=torch.long),
+            torch.tensor([[4]], dtype=torch.long),
+            session_id="missing",
+            cache_mode="decode",
+            seq_len=5,
+        )
+
+
+def test_shard_cache_lru_bounds_sessions():
+    torch = pytest.importorskip("torch")
+    shard = _fake_cache_shard(torch, max_sessions=1)
+
+    for session in ("old", "new"):
+        shard._run_layers(
+            torch.zeros((1, 2, 2), dtype=torch.bfloat16),
+            torch.ones((1, 2), dtype=torch.long),
+            torch.arange(2, dtype=torch.long).reshape(1, 2),
+            session_id=session,
+            cache_mode="prefill",
+            seq_len=2,
+        )
+
+    assert list(shard._session_cache.keys()) == [("new", 0, 1)]
 
 
 def test_partial_materialize_guard_requires_local_non_full_non_quantized_snapshot(tmp_path):