From 3cd6c74b57816a35bc9071a3a55c52570dc8242e Mon Sep 17 00:00:00 2001
From: "Gabriel A. Devenyi" <gdevenyi@gmail.com>
Date: Wed, 23 Sep 2026 20:29:54 -0400
Subject: [PATCH 01/11] fix(ple): build the token index inside a graph capture
 instead of reusing an evictable cache entry

The eager warmup forward right before capture leaves the decode (req, local) pair in the
64-entry FIFO memo; the capture hit it and baked its addresses into the graph without a
reference. After ~64 distinct prefill shapes the entry is evicted, the allocator reuses the
storage, and every decode replay feeds garbage indices to the fused hash kernel (out-of-
bounds reads). A capture now always builds its own pair, which the graph recomputes into
its pool on each replay.

Assisted-by: Claude Opus 5.5
---
 python/freetoken/models/qwen4_exp/ple.py | 13 +++++++++----
 tests/models/qwen4_exp/test_ple.py       | 13 +++++++++++++
 2 files changed, 22 insertions(+), 4 deletions(-)

diff --git a/python/freetoken/models/qwen4_exp/ple.py b/python/freetoken/models/qwen4_exp/ple.py
index 284e023bdd..44973bc0a6 100644
--- a/python/freetoken/models/qwen4_exp/ple.py
+++ b/python/freetoken/models/qwen4_exp/ple.py
@@ -49,6 +49,10 @@
 _FUSED_HASH_ENV = "FREETOKEN_PLE_FUSED_HASH"
 
 
+def _capturing(device: torch.device) -> bool:
+    return device.type == "cuda" and torch.cuda.is_current_stream_capturing()
+
+
 def _fused_row_ids_enabled() -> bool:
     """The fused hash kernel, on unless ``FREETOKEN_PLE_FUSED_HASH=0`` takes it back to torch ops."""
     return (os.getenv(_FUSED_HASH_ENV) or "1").strip() not in ("0", "false", "False")
@@ -482,12 +486,13 @@ def _token_index(self, meta: PLEMetadata) -> Tuple[torch.Tensor, torch.Tensor]:
 
         The fused kernel addresses the hash window through these instead of materializing the
         ``[B, ctx+max_len]`` packed window. Memoized on the shape (which is all they depend on)
-        so a captured replay reads a stable address instead of re-running the build; a build
-        that happens DURING capture is not cached, since its buffers live in the graph pool.
+        for eager forwards only: a capture always builds its own, which the graph recomputes
+        into its pool on every replay. A cached pair (the eager warmup forward right before
+        capture leaves one) would be baked in by address, then evicted and reused.
         """
         device = meta.input_ids.device
         num_tokens = meta.input_ids.numel()
-        capturing = device.type == "cuda" and torch.cuda.is_current_stream_capturing()
+        capturing = _capturing(device)
         # is_decode is part of the key, not just the shape: a decode of B requests and a
         # prefill of ONE B-token request are the same [T] and mean opposite things (one token
         # per request at offset 0 vs B offsets inside one request).
@@ -496,7 +501,7 @@ def _token_index(self, meta: PLEMetadata) -> Tuple[torch.Tensor, torch.Tensor]:
             (num_tokens,) if meta.is_decode else tuple(meta.seq_lens),
             str(device),
         )
-        cached = self._token_index_cache.get(key)
+        cached = None if capturing else self._token_index_cache.get(key)
         if cached is not None:
             return cached
         if meta.is_decode:  # one token per request, each at offset 0
diff --git a/tests/models/qwen4_exp/test_ple.py b/tests/models/qwen4_exp/test_ple.py
index 74d4fd614e..02abaee1f4 100644
--- a/tests/models/qwen4_exp/test_ple.py
+++ b/tests/models/qwen4_exp/test_ple.py
@@ -201,6 +201,19 @@ def test_token_index_cache_is_keyed_by_shape():
     assert embedding._token_index(other)[0] is not first[0]
 
 
+def test_token_index_capture_never_reuses_an_evictable_cache_entry(monkeypatch):
+    """The eager warmup forward caches the decode pair right before capture; a capture that
+    reused it would bake an address the FIFO later evicts while the graph still reads it."""
+    embedding = _make_layer(_config()).ple_embedding
+    decode = _meta([[5]], [[EOS, EOS]], decode=True)
+    eager = embedding._token_index(decode)  # the warmup forward
+    monkeypatch.setattr(ple_module, "_capturing", lambda device: True)
+    captured = embedding._token_index(decode)
+    assert captured[0] is not eager[0] and captured[1] is not eager[1]
+    assert all(v[0] is not captured[0] for v in embedding._token_index_cache.values())
+    assert captured[0].tolist() == [0] and captured[1].tolist() == [0]
+
+
 def test_token_index_does_not_confuse_a_decode_with_a_one_request_prefill():
     """Same [T], opposite meaning: B decode rows at offset 0 vs B offsets in one request."""
     layer = _make_layer(_config())

From 9458d508471c40d6d7c97f70c472cfcb8d02f4a4 Mon Sep 17 00:00:00 2001
From: "Gabriel A. Devenyi" <gdevenyi@gmail.com>
Date: Wed, 23 Sep 2026 20:30:27 -0400
Subject: [PATCH 02/11] fix(kvcache): revive a gdn snapshot tombstone only from
 an exact host snapshot

HostPrefixTier.lookup is longest-prefix, and _rehydrate_snapshot attached whatever it
returned to the node at the node's full depth. When a prompt diverges inside a node (the
Qwen template re-rendering earlier turns without their thinking) the split node found an
earlier, shorter snapshot and the request resumed the GDN recurrence with the tokens in
between missing -- silently wrong output, and the bad state was kept on the node.

Assisted-by: Claude Opus 5.5
---
 python/freetoken/scheduler/cache.py           |  7 +++-
 .../kvcache/test_host_tier_rehydrate_exact.py | 39 +++++++++++++++++++
 2 files changed, 44 insertions(+), 2 deletions(-)
 create mode 100644 tests/kvcache/test_host_tier_rehydrate_exact.py

diff --git a/python/freetoken/scheduler/cache.py b/python/freetoken/scheduler/cache.py
index 3f833eaf18..3b9dc55746 100644
--- a/python/freetoken/scheduler/cache.py
+++ b/python/freetoken/scheduler/cache.py
@@ -619,8 +619,11 @@ def _rehydrate_snapshot(self, node) -> int | None:
         prefix, copy the snapshot back into a fresh pool slot and return the slot; None =
         keep walking (the match truncates as before)."""
         tier = self.host_tier
-        hit = tier.lookup(self.prefix_cache._collect_key(node), self.page_size)
-        if hit is None or hit.gdn_slot is None:
+        key = self.prefix_cache._collect_key(node)
+        hit = tier.lookup(key, self.page_size)
+        # lookup is longest-prefix; a shorter entry's state would resume this node's full
+        # depth with the tokens in between missing from the recurrence
+        if hit is None or hit.gdn_slot is None or hit.host_len != len(key):
             return None
         pool = self.linear_state_pool
         if pool.num_free_slots < 1:
diff --git a/tests/kvcache/test_host_tier_rehydrate_exact.py b/tests/kvcache/test_host_tier_rehydrate_exact.py
new file mode 100644
index 0000000000..f67402041f
--- /dev/null
+++ b/tests/kvcache/test_host_tier_rehydrate_exact.py
@@ -0,0 +1,39 @@
+"""The GDN snapshot tier (upstream #525) may only revive a tombstone from a snapshot of
+exactly that node's prefix: the host index lookup is longest-prefix, and a shorter
+snapshot would resume the recurrence with the tokens in between silently missing."""
+
+from types import SimpleNamespace
+
+import torch
+
+from freetoken.kvcache.host_prefix_tier import HostMatch
+from freetoken.scheduler.cache import CacheManager
+
+
+def _manager(host_len: int, key_len: int = 128):
+    reads = []
+    tier = SimpleNamespace(
+        lookup=lambda key, page_size: HostMatch(host_len, None, 7),
+        read_gdn=lambda host_slot, gpu_slot: reads.append((host_slot, gpu_slot)),
+    )
+    pool = SimpleNamespace(num_free_slots=4, alloc=lambda n: [3])
+    manager = SimpleNamespace(
+        host_tier=tier,
+        prefix_cache=SimpleNamespace(_collect_key=lambda node: torch.arange(key_len)),
+        page_size=64,
+        linear_state_pool=pool,
+        ensure_mamba_slots=lambda n: None,
+    )
+    return manager, reads
+
+
+def test_a_shorter_host_snapshot_does_not_revive_a_deeper_node():
+    manager, reads = _manager(host_len=64)
+    assert CacheManager._rehydrate_snapshot(manager, node=object()) is None
+    assert reads == []
+
+
+def test_an_exact_host_snapshot_revives_the_node():
+    manager, reads = _manager(host_len=128)
+    assert CacheManager._rehydrate_snapshot(manager, node=object()) == 3
+    assert reads == [(7, 3)]

From d0bafe4a68421599b20300f6850601b7a72abd94 Mon Sep 17 00:00:00 2001
From: "Gabriel A. Devenyi" <gdevenyi@gmail.com>
Date: Wed, 23 Sep 2026 20:31:21 -0400
Subject: [PATCH 03/11] fix(kvcache): copy gdn snapshots against the live pool
 tensors, ordered after the engine stream

Two defects in upstream #525's HostPrefixTier:
- it kept the LinearStatePool tensors from construction; a runtime rebuild with
  num_mamba_slots replaces them, so snapshots were read from / restored into orphaned
  tensors (wrong state, IndexError when the pool grew, and the old pool never freed);
- its private copy stream never waited on the engine stream, so a restore into a slot
  freed while overlap's already-launched extra decode step still writes it could be
  overwritten.

Assisted-by: Claude Opus 5.5
---
 python/freetoken/kvcache/host_prefix_tier.py | 50 +++++++++---------
 tests/kvcache/test_host_tier_live_pool.py    | 53 ++++++++++++++++++++
 2 files changed, 79 insertions(+), 24 deletions(-)
 create mode 100644 tests/kvcache/test_host_tier_live_pool.py

diff --git a/python/freetoken/kvcache/host_prefix_tier.py b/python/freetoken/kvcache/host_prefix_tier.py
index 3b2d2d5dc3..314e0790c5 100644
--- a/python/freetoken/kvcache/host_prefix_tier.py
+++ b/python/freetoken/kvcache/host_prefix_tier.py
@@ -59,13 +59,15 @@ def __init__(self, state_pool, kv_pool=None, *, gdn_slots_host: int = 64,
         tensors = {"conv": state_pool.conv_states, "rec": state_pool.recurrent_states}
         tensors.update(state_pool.slot_states)
         self._order = list(tensors)
-        self._banks: dict[str, tuple[torch.Tensor, HostBank]] = {}
+        # banks only: the pool tensors are re-read on every copy, since
+        # LinearStatePool.rebuild replaces them (a runtime cache rebuild)
+        self._banks: dict[str, HostBank] = {}
         for name, t in tensors.items():
             # pool: [layers, slots, *rest] -> bank: [host_slots, layers, *rest]
             bank = HostBank((gdn_slots_host, t.shape[0], *t.shape[2:]), t.dtype)
             if self.device.type == "cuda" and torch.cuda.is_available():
                 bank.pin()
-            self._banks[name] = (t, bank)
+            self._banks[name] = bank
         self.num_slots = gdn_slots_host
         self._free = list(range(gdn_slots_host))
         self._in_use: set[int] = set()
@@ -102,40 +104,40 @@ def __init__(self, state_pool, kv_pool=None, *, gdn_slots_host: int = 64,
             torch.cuda.Stream(self.device) if self.device.type == "cuda" else None)
 
     # ---------------- GDN snapshot storage ----------------
+    def _pairs(self):
+        """(live pool tensor [layers, slots, ...], host bank) per state tensor."""
+        pool = self.pool
+        live = {"conv": pool.conv_states, "rec": pool.recurrent_states, **pool.slot_states}
+        return [(live[name], self._banks[name]) for name in self._order]
+
+    def _copy(self, copies) -> None:
+        if self.copy_stream is None:
+            for dst, src in copies:
+                dst.copy_(src)
+            return
+        # The engine stream may still have work queued on these slots (overlap launches the
+        # next batch before the previous one is drained): order the copies after it.
+        self.copy_stream.wait_stream(torch.cuda.current_stream(self.device))
+        with torch.cuda.stream(self.copy_stream):
+            for dst, src in copies:
+                dst.copy_(src, non_blocking=True)
+        self.copy_stream.synchronize()
+
     def write_gdn(self, gpu_slot: int) -> int:
         """D2H do snapshot do gpu_slot pra um slot host livre. Retorna o host_slot.
-        Deve ser enfileirado ANTES de state_pool.free(gpu_slot): a ordem de stream do
-        engine garante leitura antes de qualquer reescrita do slot (ver spec)."""
+        Deve ser enfileirado ANTES de state_pool.free(gpu_slot)."""
         if not self._free:
             raise RuntimeError("HostPrefixTier: sem slots host GDN livres (chame evict_host)")
         hs = self._free.pop()
         self._in_use.add(hs)
-        if self.copy_stream is not None:
-            with torch.cuda.stream(self.copy_stream):
-                for name in self._order:
-                    t, bank = self._banks[name]
-                    bank.tensor[hs].copy_(t[:, gpu_slot], non_blocking=True)
-            self.copy_stream.synchronize()
-        else:
-            for name in self._order:
-                t, bank = self._banks[name]
-                bank.tensor[hs].copy_(t[:, gpu_slot])
+        self._copy([(bank.tensor[hs], t[:, gpu_slot]) for t, bank in self._pairs()])
         return hs
 
     def read_gdn(self, host_slot: int, dst_gpu_slot: int) -> None:
         """H2D host_slot -> dst_gpu_slot do pool (já alocado pelo caller)."""
         if host_slot not in self._in_use:
             raise KeyError(f"HostPrefixTier: slot host {host_slot} não está em uso")
-        if self.copy_stream is not None:
-            with torch.cuda.stream(self.copy_stream):
-                for name in self._order:
-                    t, bank = self._banks[name]
-                    t[:, dst_gpu_slot].copy_(bank.tensor[host_slot], non_blocking=True)
-            self.copy_stream.synchronize()
-        else:
-            for name in self._order:
-                t, bank = self._banks[name]
-                t[:, dst_gpu_slot].copy_(bank.tensor[host_slot])
+        self._copy([(t[:, dst_gpu_slot], bank.tensor[host_slot]) for t, bank in self._pairs()])
 
     def free_gdn(self, host_slot: int) -> None:
         self._in_use.discard(host_slot)
diff --git a/tests/kvcache/test_host_tier_live_pool.py b/tests/kvcache/test_host_tier_live_pool.py
new file mode 100644
index 0000000000..bd12836893
--- /dev/null
+++ b/tests/kvcache/test_host_tier_live_pool.py
@@ -0,0 +1,53 @@
+"""The GDN snapshot tier (upstream #525) must copy against the pool's CURRENT tensors, which
+LinearStatePool.rebuild replaces, and must order its copies after the engine stream."""
+
+import contextlib
+from types import SimpleNamespace
+
+import torch
+
+import freetoken.kvcache.host_prefix_tier as tier_module
+from freetoken.kvcache.host_prefix_tier import HostPrefixTier
+
+
+def _pool(slots: int):
+    return SimpleNamespace(
+        conv_states=torch.zeros(2, slots, 4, 3),
+        recurrent_states=torch.zeros(2, slots, 2, 4),
+        slot_states={},
+    )
+
+
+def test_read_after_a_pool_rebuild_lands_in_the_new_tensors():
+    pool = _pool(3)
+    tier = HostPrefixTier(pool, gdn_slots_host=2, device="cpu")
+    pool.conv_states[:, 1] = 5.0
+    pool.recurrent_states[:, 1] = 7.0
+    host = tier.write_gdn(1)
+    # a runtime rebuild swaps in fresh (larger) state tensors
+    pool.conv_states = torch.zeros(2, 6, 4, 3)
+    pool.recurrent_states = torch.zeros(2, 6, 2, 4)
+    tier.read_gdn(host, 5)
+    assert torch.all(pool.conv_states[:, 5] == 5.0)
+    assert torch.all(pool.recurrent_states[:, 5] == 7.0)
+
+
+def test_copies_wait_for_the_engine_stream(monkeypatch):
+    events = []
+    engine = object()
+
+    class FakeStream:
+        def wait_stream(self, other):
+            events.append(("wait", other))
+
+        def synchronize(self):
+            events.append(("sync",))
+
+    monkeypatch.setattr(tier_module.torch.cuda, "current_stream", lambda device=None: engine)
+    monkeypatch.setattr(tier_module.torch.cuda, "stream", lambda s: contextlib.nullcontext())
+    pool = _pool(3)
+    tier = HostPrefixTier(pool, gdn_slots_host=2, device="cpu")
+    tier.copy_stream = FakeStream()
+    host = tier.write_gdn(1)
+    tier.read_gdn(host, 2)
+    assert events == [("wait", engine), ("sync",), ("wait", engine), ("sync",)]

From 6d1f0addb5736830252f847da6b6b1e22c9703ad Mon Sep 17 00:00:00 2001
From: "Gabriel A. Devenyi" <gdevenyi@gmail.com>
Date: Wed, 23 Sep 2026 20:32:06 -0400
Subject: [PATCH 04/11] fix(tokenizer): merge developer messages into the
 leading system message after mapping them

#487 hoists and merges system messages in the frontend, but #391 maps developer -> system
later, in the tokenizer worker. [system, developer, user] (or a developer message mid-
conversation) reached the Qwen template as a non-first system message and failed with
"System message must be at the beginning". The mapping now re-applies the merge.

Assisted-by: Claude Opus 5.5
---
 python/freetoken/tokenizer/tokenize.py | 10 +++++++++-
 tests/tokenizer/test_tokenize.py       | 18 ++++++++++++++++++
 2 files changed, 27 insertions(+), 1 deletion(-)

diff --git a/python/freetoken/tokenizer/tokenize.py b/python/freetoken/tokenizer/tokenize.py
index 6a2cad2e2c..8085dd9d96 100644
--- a/python/freetoken/tokenizer/tokenize.py
+++ b/python/freetoken/tokenizer/tokenize.py
@@ -65,7 +65,15 @@ def _map_developer_role(
         m.get("role") == "developer" for m in messages
     ):
         return messages
-    return [{**m, "role": "system"} if m.get("role") == "developer" else m for m in messages]
+    mapped = [{**m, "role": "system"} if m.get("role") == "developer" else m for m in messages]
+    # The frontend hoisted and merged system messages before this mapping ran (templates such
+    # as Qwen's raise on a system message that is not first), so do it again for the new ones.
+    system = [m for m in mapped if m.get("role") == "system"]
+    if len(system) == 1 and mapped[0] is system[0]:
+        return mapped
+    rest = [m for m in mapped if m.get("role") != "system"]
+    content = "\n\n".join(str(m.get("content") or "") for m in system if m.get("content"))
+    return [{"role": "system", "content": content}] + rest
 
 
 class TokenizeManager:
diff --git a/tests/tokenizer/test_tokenize.py b/tests/tokenizer/test_tokenize.py
index ae3af3a1bd..5aea33855f 100644
--- a/tests/tokenizer/test_tokenize.py
+++ b/tests/tokenizer/test_tokenize.py
@@ -296,3 +296,21 @@ def test_developer_role_maps_to_system_unless_the_template_knows_it():
     assert mapped[0]["content"] == "be brief" and msgs[0]["role"] == "developer"  # input untouched
     assert _map_developer_role(msgs, "{% if message.role == 'developer' %}") is msgs
     assert _map_developer_role(msgs[1:], None) == msgs[1:]  # nothing to map
+
+
+def test_mapped_developer_messages_are_merged_into_one_leading_system_message():
+    """The frontend hoists system messages before the tokenizer maps developer -> system; a
+    template that rejects a non-first system message must still get exactly one, first."""
+    from freetoken.tokenizer.tokenize import _map_developer_role
+
+    template = "{% if message.role == 'system' %}"
+    msgs = [
+        {"role": "system", "content": "be brief"},
+        {"role": "developer", "content": "answer in German"},
+        {"role": "user", "content": "hi"},
+    ]
+    mapped = _map_developer_role(msgs, template)
+    assert [m["role"] for m in mapped] == ["system", "user"]
+    assert mapped[0]["content"] == "be brief\n\nanswer in German"
+    late = [{"role": "user", "content": "hi"}, {"role": "developer", "content": "be brief"}]
+    assert [m["role"] for m in _map_developer_role(late, template)] == ["system", "user"]

From 59df75a163d9ffef3befdd0c27c521f298466df8 Mon Sep 17 00:00:00 2001
From: "Gabriel A. Devenyi" <gdevenyi@gmail.com>
Date: Wed, 23 Sep 2026 20:32:37 -0400
Subject: [PATCH 05/11] fix(server): parse a qwen3 continue_final_message
 continuation as content

The qwen3 reasoning parser starts inside reasoning unless thinking is off, but the Qwen
template renders the final assistant turn as <think>...</think> + content, so a
continue_final_message continuation (fork #47) never emits </think> and the whole reply
came back under reasoning_content with empty content.

Assisted-by: Claude Opus 5.5
---
 python/freetoken/server/generation.py |  7 ++++++-
 tests/server/test_openai_api.py       | 20 ++++++++++++++++++++
 2 files changed, 26 insertions(+), 1 deletion(-)

diff --git a/python/freetoken/server/generation.py b/python/freetoken/server/generation.py
index 43d18d5723..7d3d1c06e7 100644
--- a/python/freetoken/server/generation.py
+++ b/python/freetoken/server/generation.py
@@ -478,7 +478,12 @@ def _make_reasoning_parser(spec: GenSpec, state: Any) -> ReasoningParser | None:
         # The qwen3 chat template opens an implicit <think> (thinking on) unless
         # enable_thinking is explicitly false, so the model emits only the closing
         # </think>. Mirror that default here, else the chain-of-thought leaks into content.
-        force_reasoning = (spec.chat_template_kwargs or {}).get("enable_thinking") is not False
+        # continue_final_message is the exception: the template renders the final assistant
+        # turn as <think>...</think> + content, so the continuation starts in content.
+        ctk = spec.chat_template_kwargs or {}
+        force_reasoning = ctk.get("enable_thinking") is not False and not ctk.get(
+            "continue_final_message"
+        )
     elif parser_name == "glm":
         # GLM's template honors enable_thinking (default on) even with tools; the
         # generic fallback would force thinking and mislabel disabled output as reasoning.
diff --git a/tests/server/test_openai_api.py b/tests/server/test_openai_api.py
index ce675a2508..7d5b3cd114 100644
--- a/tests/server/test_openai_api.py
+++ b/tests/server/test_openai_api.py
@@ -223,6 +223,26 @@ def test_glm_reasoning_parser_honors_disabled_thinking_with_tools():
     assert parser is not None and parser.detector.force_reasoning is True
 
 
+def test_qwen3_continuation_of_a_final_assistant_message_is_content_not_reasoning():
+    # The Qwen template closes the think block before the assistant's content, so a
+    # continue_final_message continuation never emits </think>; starting the parser inside
+    # reasoning filed the whole continuation under reasoning_content.
+    from freetoken.server.generation import _make_reasoning_parser
+
+    state = FakeState([], reasoning_parser="qwen3")
+    req = chat_request(
+        messages=[
+            {"role": "user", "content": "Count to five."},
+            {"role": "assistant", "content": "One, two,"},
+        ],
+        continue_final_message=True,
+    )
+    parser = _make_reasoning_parser(chat_request_to_genspec(req, {}), state)
+    assert parser is not None and parser.detector.force_reasoning is False
+    plain = _make_reasoning_parser(chat_request_to_genspec(chat_request(), {}), state)
+    assert plain.detector.force_reasoning is True
+
+
 def test_non_stream_chat_completion_returns_openai_tool_calls_and_sends_tools():
     output = '[TOOL_CALLS] [{"name":"get_weather","arguments":{"city":"Paris"}}]'
     state = FakeState(

From 5c4be62e15041f068f800f75311d21fefafd99e9 Mon Sep 17 00:00:00 2001
From: "Gabriel A. Devenyi" <gdevenyi@gmail.com>
Date: Wed, 23 Sep 2026 20:33:10 -0400
Subject: [PATCH 06/11] fix(server): reject a non-finite repetition_penalty

The check was `<= 0`, which NaN passes; NaN or inf then poisons the logits of every
token seen so far (fork #47's sampler plan) and the request silently returns garbage.

Assisted-by: Claude Opus 5.5
---
 python/freetoken/server/generation.py    |  6 ++++--
 tests/server/test_sampling_validation.py | 23 +++++++++++++++++++++++
 2 files changed, 27 insertions(+), 2 deletions(-)
 create mode 100644 tests/server/test_sampling_validation.py

diff --git a/python/freetoken/server/generation.py b/python/freetoken/server/generation.py
index 7d3d1c06e7..5c0e60f557 100644
--- a/python/freetoken/server/generation.py
+++ b/python/freetoken/server/generation.py
@@ -14,6 +14,7 @@
 from __future__ import annotations
 
 import asyncio
+import math
 import json
 import time
 from collections.abc import AsyncIterator
@@ -229,8 +230,9 @@ def pick(value, key, framework):
     for name, value in (("presence_penalty", presence_penalty), ("frequency_penalty", frequency_penalty)):
         if not -2.0 <= float(value) <= 2.0:
             raise ValueError(f"{name} must be in [-2, 2], got {value}")
-    if resolved_rep <= 0.0:
-        raise ValueError(f"repetition_penalty must be positive, got {resolved_rep}")
+    # not `<= 0`: NaN compares False and would poison every seen token's logit
+    if not (math.isfinite(resolved_rep) and resolved_rep > 0.0):
+        raise ValueError(f"repetition_penalty must be a positive finite number, got {resolved_rep}")
     if min_tokens < 0:
         raise ValueError(f"min_tokens must be >= 0, got {min_tokens}")
     bias: list[list[float]] | None = None
diff --git a/tests/server/test_sampling_validation.py b/tests/server/test_sampling_validation.py
new file mode 100644
index 0000000000..3b2e1d6473
--- /dev/null
+++ b/tests/server/test_sampling_validation.py
@@ -0,0 +1,23 @@
+"""resolve_sampling rejects non-finite penalties: NaN passes a `<= 0` check and would
+turn every already-seen token's logit into NaN."""
+
+import pytest
+
+from freetoken.server.generation import resolve_sampling
+
+
+def _resolve(**extra):
+    return resolve_sampling(
+        temperature=None, top_k=None, top_p=None, max_tokens=8, ignore_eos=False,
+        model_sampling={}, **extra,
+    )
+
+
+@pytest.mark.parametrize("value", [float("nan"), float("inf"), 0.0, -1.0])
+def test_repetition_penalty_must_be_positive_and_finite(value):
+    with pytest.raises(ValueError, match="repetition_penalty"):
+        _resolve(repetition_penalty=value)
+
+
+def test_a_normal_repetition_penalty_is_kept():
+    assert _resolve(repetition_penalty=1.1).repetition_penalty == pytest.approx(1.1)

From dc060d900b5749fb62bc3d2980fb585e6c5986a4 Mon Sep 17 00:00:00 2001
From: "Gabriel A. Devenyi" <gdevenyi@gmail.com>
Date: Wed, 23 Sep 2026 20:33:46 -0400
Subject: [PATCH 07/11] fix(server): cancel the other choices when one n>1 chat
 choice fails

asyncio.gather passes the first GenerationError up without cancelling its siblings, and
the aborts that follow drop the events those siblings wait on, so they stayed pending
forever (no accounting row, 'Task was destroyed but it is pending' at shutdown).

Assisted-by: Claude Opus 5.5
---
 python/freetoken/server/openai_api.py    | 15 ++++++----
 tests/server/test_stream_cancellation.py | 36 ++++++++++++++++++++++++
 2 files changed, 45 insertions(+), 6 deletions(-)

diff --git a/python/freetoken/server/openai_api.py b/python/freetoken/server/openai_api.py
index b67825bfaf..9ce63060ab 100644
--- a/python/freetoken/server/openai_api.py
+++ b/python/freetoken/server/openai_api.py
@@ -401,14 +401,17 @@ async def handle_chat_completion(
             chunks = state.stream_with_cancellation(chunks, request, uids)
         return StreamingResponse(chunks, media_type="text/event-stream")
 
+    tasks = [
+        asyncio.ensure_future(generate_full(u, spec, state, source="/v1/chat/completions"))
+        for u in uids
+    ]
     try:
-        results = await _await_watching_disconnect(
-            asyncio.gather(*(generate_full(u, spec, state, source="/v1/chat/completions") for u in uids)),
-            request,
-            state,
-            uids,
-        )
+        results = await _await_watching_disconnect(asyncio.gather(*tasks), request, state, uids)
     except GenerationError as exc:
+        # gather does not cancel the other choices; the aborts below drop the events they
+        # wait on, so without the cancel they would stay pending forever
+        for task in tasks:
+            task.cancel()
         for u in uids:
             await state.abort_user(u)
         return create_error_response(str(exc), code=exc.code)
diff --git a/tests/server/test_stream_cancellation.py b/tests/server/test_stream_cancellation.py
index ba7da7c17b..9eae8a780d 100644
--- a/tests/server/test_stream_cancellation.py
+++ b/tests/server/test_stream_cancellation.py
@@ -316,3 +316,39 @@ def new_user(self):
 
     assert resp.status_code == 499
     assert sorted(state.aborted) == [8, 9]
+
+
+def test_a_failed_choice_cancels_the_other_choices(monkeypatch):
+    # n > 1: one choice errors while another is still generating. gather passes the error up
+    # without cancelling the sibling, and the aborts drop the event it waits on.
+    monkeypatch.setattr(openai_api, "_DISCONNECT_POLL_SECONDS", 0.01)
+
+    class _State(_ApiState):
+        def __init__(self):
+            super().__init__(acks=None)
+            self._next = 7
+
+        def new_user(self):
+            self._next += 1
+            return self._next
+
+        async def wait_for_ack(self, uid):
+            if uid == 8:
+                yield SimpleNamespace(**{**vars(_ack("")), "error": "boom"})
+                return
+            await asyncio.sleep(3600)
+            yield _ack("never")
+
+    async def scenario():
+        state = _State()
+        resp = await openai_api.handle_chat_completion(
+            _chat_req(n=2), _Request(disconnected=False), state, {}
+        )
+        await asyncio.sleep(0)
+        pending = [t for t in asyncio.all_tasks() if t is not asyncio.current_task()]
+        return resp, state, pending
+
+    resp, state, pending = asyncio.run(scenario())
+    assert resp.status_code == 400
+    assert sorted(state.aborted) == [8, 9]
+    assert pending == []

From 655d491a6c99f45bc401560c24fc76a9db09a25e Mon Sep 17 00:00:00 2001
From: "Gabriel A. Devenyi" <gdevenyi@gmail.com>
Date: Wed, 23 Sep 2026 20:34:56 -0400
Subject: [PATCH 08/11] fix(server): watch for client disconnects on
 non-streaming /v1/messages and /v1/responses

#222 only wrapped the OpenAI handlers, so an abandoned non-streaming Anthropic or
Responses request kept decoding to max_tokens, blocking a single-slot server meanwhile.
Both now use the same watcher and answer 499.

Assisted-by: Claude Opus 5.5
---
 python/freetoken/server/anthropic_api.py |  8 +++++++-
 python/freetoken/server/responses_api.py |  8 +++++++-
 tests/server/test_stream_cancellation.py | 22 ++++++++++++++++++++++
 3 files changed, 36 insertions(+), 2 deletions(-)

diff --git a/python/freetoken/server/anthropic_api.py b/python/freetoken/server/anthropic_api.py
index 48c296ffdc..7f0d61c255 100644
--- a/python/freetoken/server/anthropic_api.py
+++ b/python/freetoken/server/anthropic_api.py
@@ -55,6 +55,7 @@
     submit_generation,
     with_keepalive,
 )
+from .openai_api import _await_watching_disconnect
 from .request_logger import log_request
 
 # Emit a protocol-native `ping` event after this many seconds of stream silence,
@@ -140,9 +141,14 @@ async def handle_anthropic_messages(
         return StreamingResponse(events, media_type="text/event-stream")
 
     try:
-        result = await generate_full(uid, spec, state, source="/v1/messages")
+        # an abandoned request must not keep decoding to max_tokens (#222's watcher)
+        result = await _await_watching_disconnect(
+            generate_full(uid, spec, state, source="/v1/messages"), request, state, [uid]
+        )
     except GenerationError as exc:
         return _anthropic_error_response(400, "invalid_request_error", str(exc))
+    if result is None:
+        return _anthropic_error_response(499, "api_error", "client disconnected before the response was ready")
     response = anthropic_full_response(
         result, req.model, uid, cache_report=cache_report, metrics=metrics
     )
diff --git a/python/freetoken/server/responses_api.py b/python/freetoken/server/responses_api.py
index eea41722bb..9fd063a992 100644
--- a/python/freetoken/server/responses_api.py
+++ b/python/freetoken/server/responses_api.py
@@ -79,6 +79,7 @@
     submit_generation,
     with_keepalive,
 )
+from .openai_api import _await_watching_disconnect
 from .request_logger import log_request
 
 # Seconds of event silence before a keep-alive frame is emitted on the stream.
@@ -177,9 +178,14 @@ async def handle_responses(
         return StreamingResponse(events, media_type="text/event-stream")
 
     try:
-        result = await generate_full(uid, spec, state, source="/v1/responses")
+        # an abandoned request must not keep decoding to max_tokens (#222's watcher)
+        result = await _await_watching_disconnect(
+            generate_full(uid, spec, state, source="/v1/responses"), request, state, [uid]
+        )
     except GenerationError as exc:
         return _error_response(400, str(exc), exc.code)
+    if result is None:
+        return _error_response(499, "client disconnected before the response was ready")
     response = build_responses_response(
         result, req, response_id, created, cache_report=cache_report, metrics=metrics
     )
diff --git a/tests/server/test_stream_cancellation.py b/tests/server/test_stream_cancellation.py
index 9eae8a780d..71d5ac8b85 100644
--- a/tests/server/test_stream_cancellation.py
+++ b/tests/server/test_stream_cancellation.py
@@ -352,3 +352,25 @@ async def scenario():
     assert resp.status_code == 400
     assert sorted(state.aborted) == [8, 9]
     assert pending == []
+
+
+def test_messages_and_responses_non_stream_disconnect_delivers_abort(monkeypatch):
+    # #222 only wrapped the OpenAI handlers; the Anthropic and Responses endpoints kept
+    # decoding an abandoned non-streaming request to max_tokens.
+    from freetoken.server import anthropic_api, responses_api
+    from freetoken.server.anthropic_models import AnthropicMessagesRequest
+    from freetoken.server.responses_api import ResponsesRequest
+
+    monkeypatch.setattr(openai_api, "_DISCONNECT_POLL_SECONDS", 0.01)
+
+    state = _ApiState(acks=None)
+    req = AnthropicMessagesRequest(model="m", max_tokens=8, messages=[{"role": "user", "content": "hi"}])
+    handler = anthropic_api.handle_anthropic_messages(req, _Request(disconnected=True), state, {})
+    resp = asyncio.run(asyncio.wait_for(handler, timeout=5))  # unwatched: hangs, then times out
+    assert resp.status_code == 499 and state.aborted == [7]
+
+    state = _ApiState(acks=None)
+    req = ResponsesRequest(model="m", input="hi", max_output_tokens=8)
+    handler = responses_api.handle_responses(req, _Request(disconnected=True), state, {})
+    resp = asyncio.run(asyncio.wait_for(handler, timeout=5))
+    assert resp.status_code == 499 and state.aborted == [7]

From 63c18127bd1d8cd45c3f7748aa5c1ee8c4a4c413 Mon Sep 17 00:00:00 2001
From: "Gabriel A. Devenyi" <gdevenyi@gmail.com>
Date: Wed, 23 Sep 2026 20:35:40 -0400
Subject: [PATCH 09/11] fix(engine): keep multimodal placeholder ids out of the
 repetition-penalty prompt mask

Fork #47's sampler plan scattered the raw prompt ids into a [rows, vocab+1] mask; image
tokens carry pseudo-ids >= 1,000,000, so repetition_penalty with an image indexed out of
bounds -- a device-side assert that takes the server down. They now land in the scratch
column.

Assisted-by: Claude Opus 5.5
---
 python/freetoken/engine/sample.py      |  4 +++-
 tests/engine/test_logits_processors.py | 10 ++++++++++
 2 files changed, 13 insertions(+), 1 deletion(-)

diff --git a/python/freetoken/engine/sample.py b/python/freetoken/engine/sample.py
index f1984f7576..bca85c56fa 100644
--- a/python/freetoken/engine/sample.py
+++ b/python/freetoken/engine/sample.py
@@ -187,7 +187,9 @@ def _plan(self, batch: Batch, params) -> LogitsPlan | None:
                 hist.append(ids[prompt_len:].tolist() if penalized else [])
                 if p.repetition_penalty != 1.0:
                     want_prompt = True
-                    prompts.append(ids[:prompt_len].tolist())
+                    # multimodal placeholders are pseudo-ids >= vocab_size (MM_PAD_SHIFT_VALUE);
+                    # send them to the scratch column, not past the scatter target
+                    prompts.append(ids[:prompt_len].clamp(max=pad).tolist())
                 else:
                     prompts.append([])
             plan.hist_ids = _padded_rows(hist, pad, self.device)
diff --git a/tests/engine/test_logits_processors.py b/tests/engine/test_logits_processors.py
index f051dec51f..a9645f4210 100644
--- a/tests/engine/test_logits_processors.py
+++ b/tests/engine/test_logits_processors.py
@@ -153,3 +153,13 @@ def test_prepare_builds_no_plan_without_processors_and_a_plan_with_them():
     assert torch.isfinite(out[3]).all()
     # greedy path: argmax over the processed logits
     assert sampler.sample(logits, args)[1].item() == 5
+
+
+def test_repetition_penalty_ignores_multimodal_placeholder_ids_in_the_prompt():
+    # image tokens carry pseudo-ids >= MM_PAD_SHIFT_VALUE; scattering them into the
+    # [rows, V+1] prompt mask indexed out of bounds (a device-side assert on CUDA)
+    sampler = Sampler(CPU, V)
+    batch = SimpleNamespace(reqs=[_req([1, 1_000_000, 2], [3], 4, repetition_penalty=2.0)])
+    args = sampler.prepare(batch)
+    out = apply_logits_processors(torch.ones(1, V), args.plan, V)
+    assert out[0, 1] == 0.5 and out[0, 2] == 0.5 and out[0, 4] == 1.0

From 1ac87c729d0d46e71dc827c8d09d03cb2cebd1ca Mon Sep 17 00:00:00 2001
From: "Gabriel A. Devenyi" <gdevenyi@gmail.com>
Date: Wed, 23 Sep 2026 20:35:40 -0400
Subject: [PATCH 10/11] fix(utils): lift the 100 MiB buffer cap #466's
 streaming unpacker introduced

msgpack.Unpacker defaults to max_buffer_size=100 MiB where the unpackb it replaced sized
the buffer to the message, so a larger message (a big image's pixel values) raised
BufferFull and killed the receiving worker.

Assisted-by: Claude Opus 5.5
---
 python/freetoken/utils/mp.py       |  4 +++-
 tests/utils/test_mp_large_frame.py | 13 +++++++++++++
 2 files changed, 16 insertions(+), 1 deletion(-)
 create mode 100644 tests/utils/test_mp_large_frame.py

diff --git a/python/freetoken/utils/mp.py b/python/freetoken/utils/mp.py
index f91498fbb2..14d421157f 100644
--- a/python/freetoken/utils/mp.py
+++ b/python/freetoken/utils/mp.py
@@ -22,7 +22,9 @@ class _CoalescedUnpacker:
     """
 
     def __init__(self) -> None:
-        self._unpacker = msgpack.Unpacker(raw=False)
+        # 0 = no size cap (2**32-1), as unpackb had; the Unpacker default of 100 MiB would
+        # kill the worker on a large message (e.g. an image's pixel values)
+        self._unpacker = msgpack.Unpacker(raw=False, max_buffer_size=0)
         self._pending: list[Any] = []
 
     def feed(self, frame: bytes) -> None:
diff --git a/tests/utils/test_mp_large_frame.py b/tests/utils/test_mp_large_frame.py
new file mode 100644
index 0000000000..f91ee5604d
--- /dev/null
+++ b/tests/utils/test_mp_large_frame.py
@@ -0,0 +1,13 @@
+"""The coalesced-frame decoder (#466) must accept a message larger than msgpack's
+Unpacker default buffer (100 MiB), which unpackb never limited."""
+
+import msgpack
+
+from freetoken.utils.mp import _CoalescedUnpacker
+
+
+def test_a_message_over_100_mib_decodes():
+    payload = {"pixels": b"x" * (101 * 1024 * 1024)}
+    unpacker = _CoalescedUnpacker()
+    unpacker.feed(msgpack.packb(payload))
+    assert len(unpacker) == 1 and len(unpacker.take()["pixels"]) == len(payload["pixels"])

From d14f06d73ed67cc9a80911228dfddc47781669c0 Mon Sep 17 00:00:00 2001
From: "Gabriel A. Devenyi" <gdevenyi@gmail.com>
Date: Wed, 23 Sep 2026 20:48:20 -0400
Subject: [PATCH 11/11] fix(kvcache): order the host tier's kv-page copies
 after the caller's stream too

Same race the GDN copies had (the KV-page tier is off here and refused on nvfp4 pools);
also corrects the ordering comment and the msgpack buffer-cap comment.

Assisted-by: Claude Opus 5.5
---
 python/freetoken/kvcache/host_prefix_tier.py | 7 +++++--
 python/freetoken/utils/mp.py                 | 2 +-
 2 files changed, 6 insertions(+), 3 deletions(-)

diff --git a/python/freetoken/kvcache/host_prefix_tier.py b/python/freetoken/kvcache/host_prefix_tier.py
index 314e0790c5..39c96dfd6d 100644
--- a/python/freetoken/kvcache/host_prefix_tier.py
+++ b/python/freetoken/kvcache/host_prefix_tier.py
@@ -115,8 +115,9 @@ def _copy(self, copies) -> None:
             for dst, src in copies:
                 dst.copy_(src)
             return
-        # The engine stream may still have work queued on these slots (overlap launches the
-        # next batch before the previous one is drained): order the copies after it.
+        # Order the copies after the caller's stream: the scheduler stream (itself made to wait
+        # on the engine stream) or the engine stream. Overlap launches the next batch before
+        # the previous one is drained, so a just-freed slot may still be written by it.
         self.copy_stream.wait_stream(torch.cuda.current_stream(self.device))
         with torch.cuda.stream(self.copy_stream):
             for dst, src in copies:
@@ -200,6 +201,7 @@ def _go():
                     del staged_s
 
         if self.copy_stream is not None:
+            self.copy_stream.wait_stream(torch.cuda.current_stream(self.device))
             with torch.cuda.stream(self.copy_stream):
                 _go()
             self.copy_stream.synchronize()
@@ -237,6 +239,7 @@ def _go():
                     del staged_s
 
         if self.copy_stream is not None:
+            self.copy_stream.wait_stream(torch.cuda.current_stream(self.device))
             with torch.cuda.stream(self.copy_stream):
                 _go()
             self.copy_stream.synchronize()
diff --git a/python/freetoken/utils/mp.py b/python/freetoken/utils/mp.py
index 14d421157f..ee35c057c1 100644
--- a/python/freetoken/utils/mp.py
+++ b/python/freetoken/utils/mp.py
@@ -22,7 +22,7 @@ class _CoalescedUnpacker:
     """
 
     def __init__(self) -> None:
-        # 0 = no size cap (2**32-1), as unpackb had; the Unpacker default of 100 MiB would
+        # 0 = effectively uncapped (2**31-1), as unpackb had; the Unpacker default of 100 MiB would
         # kill the worker on a large message (e.g. an image's pixel values)
         self._unpacker = msgpack.Unpacker(raw=False, max_buffer_size=0)
         self._pending: list[Any] = []