Add explicit reset points to page cache, decode state etc. (#1291)

Instead of trying to unwind the allocated sequences, just reset to the
initial state explicitly.

Author: rjpower

src/levanter/inference/engine.py

@@ -272,6 +272,12 @@ class GenState(eqx.Module):
     cache: PageCache
     decode_state: DecodeState
 
+    def reset(self):
+        return GenState(
+            cache=self.cache.reset(),
+            decode_state=self.decode_state.reset(),
+        )
+
     def clone_sequence(
         self, parent_local_id: int, child_local_id: int | None = None, seq_params: SeqDecodingParams | None = None
     ) -> tuple["GenState", int]:
@@ -797,22 +803,6 @@ def __init__(
         # Results by request id -> choice -> DecodeResult
         self.results: dict[int, dict[int, DecodeResult]] = {}
 
-    def _verify_free_slot_view(self, *, context: str) -> None:
-        """Ensure host free-list matches the device page-table used mask."""
-
-        used_mask = np.asarray(jax.device_get(self.gen_state.decode_state.sequences.used_mask.array)).astype(bool)
-        free_set = set(self.free_slots)
-
-        for slot_id, is_used in enumerate(used_mask):
-            if is_used and slot_id in free_set:
-                raise RuntimeError(
-                    f"[free slot invariant] slot {slot_id} marked used but present in free list during {context}"
-                )
-            if not is_used and slot_id not in free_set:
-                raise RuntimeError(
-                    f"[free slot invariant] slot {slot_id} free in page table but missing from free list during {context}"
-                )
-
     @classmethod
     def from_model_with_config(
         cls,
@@ -853,26 +843,12 @@ def reset(self) -> None:
 
         Keeps the KV cache memory allocated. Reuses current `PageTable` object with pages freed.
         """
-        decode_state = self.gen_state.decode_state
-        page_table = decode_state.page_table
-        sequences = decode_state.sequences
-        for slot_id in range(page_table.max_seqs):
-            sequences, page_table = sequences.free_pages(page_table, slot_id)
-
-        new_decode_state = DecodeState.init(
-            page_table,
-            max_stop_seqs=self.config.max_stop_seqs,
-            max_stop_tokens=self.config.max_stop_tokens,
-            max_queued_tokens=self.config.max_queued_tokens,
-        )
-        self.gen_state = dataclasses.replace(self.gen_state, decode_state=new_decode_state)
-        self.free_slots = list(range(int(page_table.max_seqs)))
+        self.gen_state = self.gen_state.reset()
+        self.free_slots = list(range(int(self.gen_state.decode_state.max_seqs)))
         self.local_map.clear()
         self.sequences.clear()
         self.results = {}
 
-        self._verify_free_slot_view(context="reset")
-
     def _prefill_batch(self, batch: Sequence[Request]) -> _DecodeOutputs | None:
         """Admit a batch from the head of the queue that fits in free slots/pages.
 

src/levanter/inference/jit_scheduler.py

@@ -564,6 +564,8 @@ class DecodeState(eqx.Module):
     # Page table for KV page allocation and per-sequence lengths/usage
     page_table: PageTable
 
+    pad_token_id: int
+
     # Per sequence sampling parameters
     max_num_tokens: ht.i32[NamedArray, "seq"]
     """
@@ -583,6 +585,15 @@ class DecodeState(eqx.Module):
     # Cached finished flags per sequence (updated when tokens are enqueued)
     finished: ht.bool_[NamedArray, "seq"]
 
+    def reset(self):
+        return DecodeState.init(
+            page_table=self.page_table.reset(),
+            pad_token_id=self.pad_token_id,
+            max_stop_seqs=self.stop_tokens.shape["stop_seq"] if self.stop_tokens is not None else 0,
+            max_stop_tokens=self.stop_tokens.shape["position"] if self.stop_tokens is not None else 0,
+            max_queued_tokens=self.tqueue.max_queued_tokens,
+        )
+
     @staticmethod
     def init(
         page_table: PageTable,
@@ -605,6 +616,7 @@ def init(
             sequences=sequence_table,
             page_size=page_size,
             page_table=page_table,
+            pad_token_id=pad_token_id,
             tokens=hax.full({"seq": max_seqs, "position": max_seq_len}, pad_token_id, dtype=jnp.int32),
             logprobs=hax.full({"seq": max_seqs, "position": max_seq_len}, jnp.nan, dtype=jnp.float32),
             max_num_tokens=hax.full({"seq": max_seqs}, 0, dtype=jnp.int32),

src/levanter/inference/page_table.py

@@ -4,9 +4,9 @@
 import dataclasses
 
 import equinox as eqx
-import jax.numpy as jnp
 import haliax as hax
 import haliax.haxtyping as ht
+import jax.numpy as jnp
 from haliax import NamedArray
 
 from levanter.inference.utils import INVALID, is_valid
@@ -35,6 +35,10 @@ def init(max_pages: int, max_seqs: int, page_size: int, max_pages_per_seq: int)
         ref_counts = hax.full({"page": max_pages}, 0, dtype=jnp.int32)
         return PageTable(ref_counts, page_size, max_seqs, max_pages_per_seq)
 
+    def reset(self) -> "PageTable":
+        ref_counts = hax.full_like(self.page_ref_counts, 0)
+        return PageTable(ref_counts, self.page_size, self._max_seqs, self._pages_per_seq)
+
     @property
     def num_pages(self) -> int:
         return self.page_ref_counts.axis_size("page")

src/levanter/layers/kv_cache.py

@@ -5,16 +5,15 @@
 
 import dataclasses
 import functools
-from typing import Generic, Iterable, Iterator, TypeVar, Self
+from typing import Generic, Iterable, Iterator, Self, TypeVar
 
 import equinox as eqx
+import haliax as hax
 import jax
 import jax.numpy as jnp
-from jax import lax
-
-import haliax as hax
 from haliax import Axis, NamedArray
 from haliax.jax_utils import named_call
+from jax import lax
 
 from levanter.inference.page_table import PageBatchInfo, PageTableSpec
 
@@ -26,6 +25,10 @@ def copy_page(self, src_page: int, dst_page: int) -> Self:
         """Return a copy of this cache with ``src_page`` cloned into ``dst_page``."""
         raise NotImplementedError
 
+    def reset(self) -> Self:
+        """Return a reset version of this cache."""
+        raise NotImplementedError
+
 
 class KvPageCache(PageCache):
     """Concrete KV cache storing interleaved key/value pages for paged attention."""
@@ -54,6 +57,11 @@ def init(spec: PageTableSpec, kv_heads: Axis, head_size: Axis, dtype=jnp.float32
         )
         return KvPageCache(kv_pages)
 
+    def reset(self) -> "KvPageCache":
+        """Return a reset version of this cache."""
+        reset_pages = jnp.zeros_like(self.kv_pages.array)
+        return dataclasses.replace(self, kv_pages=NamedArray(reset_pages, self.kv_pages.axes))
+
     @named_call
     def update(
         self,
@@ -103,6 +111,9 @@ class ListCache(PageCache, Generic[PageCacheT]):
     def __post_init__(self):
         object.__setattr__(self, "caches", tuple(self.caches))
 
+    def reset(self) -> "ListCache[PageCacheT]":
+        return ListCache(tuple(cache.reset() for cache in self.caches))
+
     @staticmethod
     def from_iterable(caches: Iterable[PageCacheT]) -> "ListCache[PageCacheT]":
         return ListCache(tuple(caches))

tests/inference/test_inference_server.py

@@ -38,7 +38,7 @@ def trainer_config():
 def baby_llama_config():
     return InferenceServerConfig(
         service=InferenceEngineConfig(
-            max_seq_len=16,
+            max_seq_len=32,
             max_seqs=2,
             page_size=4,
             max_queued_tokens=32,
@@ -424,6 +424,35 @@ def test_logprobs_deterministic_behavior(test_client):
     print("Deterministic logprobs test passed!")
 
 
+def test_many_requests_threaded(test_client):
+    executor = ThreadPoolExecutor(max_workers=8)
+    client, server = test_client
+    futures = []
+    num_requests = 20
+    for i in range(num_requests):
+        futures.append(
+            executor.submit(
+                client.post,
+                "/v1/completions",
+                json={
+                    "model": "timinar/baby-llama-58m",
+                    "prompt": "The quick brown fox",
+                    "max_tokens": 16,
+                    "temperature": 0.0,
+                    "seed": i,
+                },
+            )
+        )
+
+    for i, future in enumerate(futures):
+        response = future.result()
+        assert response.status_code == 200
+        completion = Completion.model_validate(response.json())
+        choice = completion.choices[0]
+        assert choice.text
+        print(f"Request {i} generated text: '{choice.text}'")
+
+
 def test_reload_with_zeros_clears_outputs(test_client):
     """Test that reloading with a zeroed-out model properly clears outputs."""
     client, server = test_client