5080deploy: fix 11 bugs found re-reviewing the stack

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)

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,41 @@ 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
+        # 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:
+                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)
@@ -198,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()
@@ -235,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()

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

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:

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
     )

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
@@ -478,7 +480,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.

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)

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
     )

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:

python/freetoken/utils/mp.py

@@ -22,7 +22,9 @@ class _CoalescedUnpacker:
     """
 
     def __init__(self) -> None:
-        self._unpacker = msgpack.Unpacker(raw=False)
+        # 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] = []
 
     def feed(self, frame: bytes) -> None:

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

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