feat(scheduler): --logprob-read-isolate makes logprob reads independent of other traffic

docs/cli.md

@@ -75,6 +75,7 @@ ft serve --model ... --gpu GPU-9e8d7c6b  # the same card by UUID (a unique prefi
 | `--page-size` | 1 | KV page size; DSV4 forces 128, the TRTLLM backend needs 16/32/64, SWA models require 1 |
 | `--cache-type` | radix | `radix` (prefix reuse; SWA/GDN-aware variants picked automatically) or `naive` |
 | `--logprob-read-snapshot-donate` | all | hybrid GDN models: snapshots a finished `max_tokens=1` + `logprobs` request donates to the prefix cache: `all` (x64 prefill boundary and final state), `boundary` (only the former) or `none` |
+| `--logprob-read-isolate` | off | a `max_tokens=1` + `logprobs` request prefills in a pass of its own and neither resumes from nor stores prefix-cache entries, so its logprobs do not depend on other traffic; costs read throughput under concurrency and prefix reuse for reads |
 | `--kv-cache-dtype` | bf16 | `bf16`, `fp8`, or `nvfp4` (see [NVFP4 KV cache](#nvfp4-kv-cache)): FP8 stores the KV cache as e4m3 codes plus one fp32 scale per (token, kv head), roughly doubling the tokens that fit in the same VRAM; see [FP8 KV cache](#fp8-kv-cache) |
 | `--attention-backend`, `--attn` | auto | `trtllm`/`fi`/`fa`/`triton`/`dsv4_sparse`/`dsa`; `prefill,decode` pair allowed; auto picks per model + GPU |
 

python/freetoken/core.py

@@ -71,6 +71,11 @@ def client_token_ids(self) -> list[tuple[str, list[int]]]:
         ]
         return [(name, ids) for name, ids in fields if ids]
 
+    @property
+    def is_logprob_read(self) -> bool:
+        """A scoring read (e.g. AnyJev label logprobs): one token whose logprobs are the answer."""
+        return bool(self.logprobs) and self.max_tokens == 1
+
     @property
     def is_greedy(self) -> bool:
         return self.temperature <= 0.0 or self.top_k == 1

python/freetoken/scheduler/cache.py

@@ -35,7 +35,7 @@ def _swa_eviction_interval() -> int:
 class CacheManager:
     def __init__(self, num_pages: int, page_size: int, page_table: torch.Tensor, type: str,
                  linear_state_pool=None, swa_pool=None, sliding_window_size=None,
-                 logprob_read_snapshot_donate: str = "all"):
+                 logprob_read_snapshot_donate: str = "all", isolate_logprob_reads: bool = False):
         # The `_free_slots` follows a page-aligned manner. For example, if page_size = 2,
         # the `_free_slots` may look like [0, 2, 4, 6, ...], and each slot represents a page.
         device = page_table.device
@@ -54,6 +54,7 @@ def __init__(self, num_pages: int, page_size: int, page_table: torch.Tensor, typ
         if logprob_read_snapshot_donate not in LOGPROB_READ_SNAPSHOT_DONATE:
             raise ValueError(f"unknown snapshot donate policy {logprob_read_snapshot_donate!r}")
         self.logprob_read_snapshot_donate = logprob_read_snapshot_donate
+        self.isolate_logprob_reads = isolate_logprob_reads
         self.is_swa = type == "swa_radix"
         # swa_paged: this SWA model drives the global-paged swa pool -- true for BOTH the naive
         # (NaivePrefixCache, no reuse) and radix (SWARadixCache) paths. Gates the swa slot
@@ -145,7 +146,7 @@ def match_req(self, req: PendingReq, max_len: int | None = None) -> MatchResult:
             # (KV bytes + GDN snapshot). If the tree match fell short, rehydrate the deeper
             # host entry into fresh pages, re-insert it into the tree and re-match.
             tier = self.host_tier
-            if tier is not None and tier.has_kv:
+            if tier is not None and tier.has_kv and len(ids) > 0:
                 hit = tier.lookup(ids, self.page_size)
                 if (hit is not None and hit.kv_pages is not None and hit.gdn_slot is not None
                         and hit.host_len > m.cached_len):
@@ -374,6 +375,11 @@ def cache_req(self, req: Req, *, finished: bool) -> None:
             req.cache_handle = new_handle
             self.lock(new_handle)
 
+    def _donate_policy(self, req: Req) -> str:
+        if not req.sampling_params.is_logprob_read:
+            return "all"
+        return "none" if self.isolate_logprob_reads else self.logprob_read_snapshot_donate
+
     def _cache_req_hybrid(self, req: Req, *, finished: bool) -> None:
         """Hybrid (GDN) cache_req: commit KV like radix AND manage the GDN state snapshot.
         Prefill chunk commit: DONATE the frozen ping-pong slot (the snapshot the forward wrote
@@ -394,12 +400,7 @@ def _cache_req_hybrid(self, req: Req, *, finished: bool) -> None:
             # is this request's dup to free. The frozen slot is consumed either way (taken by
             # the tree or freed here) and both ping-pong refs are dropped before
             # _free_req_slots so nothing double-frees.
-            sp = req.sampling_params
-            donate = (
-                self.logprob_read_snapshot_donate
-                if sp.logprobs and sp.max_tokens == 1
-                else "all"
-            )
+            donate = self._donate_policy(req)
             free_upto = old_handle.cached_len
             L = req.mamba_last_track_seqlen
             if (
@@ -439,6 +440,10 @@ def _cache_req_hybrid(self, req: Req, *, finished: bool) -> None:
             return
 
         # Prefill chunk commit: donate the frozen snapshot at the tracked ×64 boundary.
+        # A non-donating read keeps its own pages and frozen slot until the finish frees them.
+        if self._donate_policy(req) == "none":
+            req.mamba_last_track_seqlen = None
+            return
         L = req.mamba_last_track_seqlen
         if L is None:
             return  # no ×64 boundary crossed this chunk; req keeps its pages (committed later)

python/freetoken/scheduler/config.py

@@ -20,6 +20,8 @@ class SchedulerConfig(EngineConfig):
     special_token_ckpt: bool = False
     # GDN snapshots a finished logprob read donates (CacheManager): all | boundary | none
     logprob_read_snapshot_donate: str = "all"
+    # logprob reads prefill alone and bypass the prefix cache, so their logprobs do not depend on other traffic
+    logprob_read_isolate: bool = False
 
     # networking config
     _unique_suffix: str = field(default_factory=_get_pid_suffix)

python/freetoken/scheduler/prefill.py

@@ -49,6 +49,8 @@ class PrefillAdder:
     encoder_cache: EncoderCache | None = None
     # end a chunk before an image it would cut; only models whose image spans attend in both directions need it
     keep_images_whole: bool = False
+    # a logprob read (max_tokens=1 + logprobs) neither resumes from nor feeds the prefix cache
+    isolate_logprob_reads: bool = False
     # the whole budget of this pass; token_budget shrinks as requests are admitted
     pass_budget: int = 0
     # SWA-pool tokens charged to reqs admitted so far this pass. Mirrors reserved_size: swa is
@@ -73,7 +75,11 @@ def _try_allocate_one(self, req: PendingReq):
 
         # TODO: consider host cache match case (it needs the same replay cap)
         replay = self.cache_manager.prefix_replay_tokens
-        mr = self.cache_manager.match_req(req, max_len=max(0, req.input_len - replay) if replay else None)
+        if self.isolate_logprob_reads and req.sampling_params.is_logprob_read:
+            max_len = 0  # a resumed prefill differs numerically from a cold one
+        else:
+            max_len = max(0, req.input_len - replay) if replay else None
+        mr = self.cache_manager.match_req(req, max_len=max_len)
         handle = mr.cuda_handle
         cached_len = handle.cached_len
         # TODO: better estimate policy
@@ -282,6 +288,8 @@ class PrefillManager:
     decode_manager: DecodeManager
     encoder_cache: EncoderCache | None = None
     keep_images_whole: bool = False
+    # a logprob read prefills in a pass of its own (see PrefillAdder for the cache side)
+    isolate_logprob_reads: bool = False
     pending_list: List[PendingReq] = field(default_factory=list)
 
     def add_one_req(self, req: UserMsg) -> None:
@@ -308,6 +316,7 @@ def schedule_next_batch(self, prefill_budget: int) -> Batch | None:
             table_manager=self.table_manager,
             encoder_cache=self.encoder_cache,
             keep_images_whole=self.keep_images_whole,
+            isolate_logprob_reads=self.isolate_logprob_reads,
         )
         reqs: List[Req] = []
         chunked_list: List[PendingReq] = []
@@ -319,6 +328,10 @@ def schedule_next_batch(self, prefill_budget: int) -> Batch | None:
         log_cached_tokens = 0
         for pending_req in self.pending_list:
             is_continuation = pending_req.chunked_req is not None
+            # break, not continue: the slicing below assumes the admitted requests are a prefix
+            read = self.isolate_logprob_reads and pending_req.sampling_params.is_logprob_read
+            if read and reqs:
+                break  # a read never shares a pass
             if req := adder.try_add_one(pending_req):
                 pending_req.chunked_req = None
                 if isinstance(req, ChunkedReq):
@@ -341,6 +354,8 @@ def schedule_next_batch(self, prefill_budget: int) -> Batch | None:
                 log_new_tokens += req.extend_len
                 if not is_continuation:
                     log_cached_tokens += req.cache_handle.cached_len
+                if read:
+                    break  # nobody joins a read's pass
             else:
                 break  # We cannot add more requests
         if len(reqs) == 0:

python/freetoken/scheduler/scheduler.py

@@ -100,6 +100,7 @@ def __init__(self, config: SchedulerConfig):
                 None,
             ) or getattr(self.engine.kv_cache, "sliding_window_size", None),
             logprob_read_snapshot_donate=getattr(config, "logprob_read_snapshot_donate", "all"),
+            isolate_logprob_reads=getattr(config, "logprob_read_isolate", False),
         )
         self.decode_manager = DecodeManager(config.page_size)
         # Bound how long a run of prefill steps may starve in-flight decodes. A long
@@ -117,6 +118,7 @@ def __init__(self, config: SchedulerConfig):
             self.decode_manager,
             encoder_cache=self.engine.encoder_cache,
             keep_images_whole=self._bidirectional_mm,
+            isolate_logprob_reads=getattr(config, "logprob_read_isolate", False),
         )
 
         # some alias for easy access

python/freetoken/server/args.py

@@ -556,6 +556,16 @@ def _infer_reasoning_parser(model_path: str) -> str | None:
         "neither, so a burst of reads cannot evict other conversations' snapshots.",
     )
 
+    parser.add_argument(
+        "--logprob-read-isolate",
+        action="store_true",
+        default=ServerArgs.logprob_read_isolate,
+        help="Make a logprob read's (max_tokens=1 with logprobs, e.g. AnyJev label scoring) "
+        "logprobs independent of other traffic: it prefills in a pass of its own and neither "
+        "resumes from nor stores prefix-cache entries. Costs read throughput under concurrency "
+        "and all prefix reuse for reads; other requests are unchanged.",
+    )
+
     parser.add_argument(
         "--text-model-only",
         action="store_true",

tests/scheduler/test_hybrid_cache_manager.py

@@ -256,11 +256,10 @@ def test_host_snapshot_tier_lowers_the_gpu_snapshot_floor(monkeypatch):
     raise SystemExit(pytest.main([__file__]))
 
 
-def _finish_read(policy: str, sp: SamplingParams, *, warm_prefix: bool = False):
+def _finish_read(policy: str, sp: SamplingParams, *, warm_prefix: bool = False, num_pages: int = 256):
     """A 130-token prompt finishing on its prefill forward, with a frozen x64 snapshot at 128
     (the deepest boundary strictly inside it) and its live state at 130. ``warm_prefix`` first
     donates a 64-token prefix so the read resumes from a locked prefix-cache hit."""
-    num_pages = 256
     pool = _pool()
     pt = torch.zeros(4, 256, dtype=torch.int32)
     cm = CacheManager(num_pages, 1, pt, "hybrid_radix", linear_state_pool=pool,
@@ -344,3 +343,93 @@ def test_snapshot_donate_policy_is_a_server_arg():
         assert parse_args(["--model", "/models/anon"])[0].logprob_read_snapshot_donate == "all"
         args, _ = parse_args(["--model", "/models/anon", "--logprob-read-snapshot-donate", "none"])
     assert args.logprob_read_snapshot_donate == "none"
+
+
+@pytest.mark.parametrize(
+    ("sp", "read"),
+    [
+        (READ, True),
+        (SamplingParams(max_tokens=1, logprobs=True), True),
+        (SamplingParams(max_tokens=1), False),
+        (SamplingParams(max_tokens=2, logprobs=True, top_logprobs=4), False),
+    ],
+)
+def test_is_logprob_read(sp, read):
+    assert sp.is_logprob_read is read
+
+
+def test_logprob_read_isolate_is_a_server_arg():
+    from unittest.mock import patch
+
+    from freetoken.server.args import parse_args
+
+    config = SimpleNamespace(to_dict=lambda: {"architectures": ["LlamaForCausalLM"], "torch_dtype": "bfloat16"})
+    with patch("freetoken.utils.cached_load_hf_config", lambda _path: config):
+        assert parse_args(["--model", "/models/anon"])[0].logprob_read_isolate is False
+        args, _ = parse_args(["--model", "/models/anon", "--logprob-read-isolate"])
+    assert args.logprob_read_isolate is True
+
+
+def _assert_no_leak(cm, num_pages=256):
+    pc, pool = cm.prefix_cache, cm.linear_state_pool
+    assert len(cm.free_slots) + pc.full_evictable + pc.full_protected == num_pages
+    assert pc.full_protected == 0 and pc.mamba_protected == 0
+    assert pool.num_free_slots + pc.mamba_evictable_size == pool.num_slots - 1
+    cm.check_integrity()
+
+
+@pytest.mark.parametrize("isolate", [False, True])
+def test_isolated_logprob_read_ignores_and_keeps_out_of_the_prefix_cache(isolate):
+    """An identical earlier read donated under "all"; an isolated read still prefills cold and stores nothing."""
+    from freetoken.scheduler.decode import DecodeManager
+    from freetoken.scheduler.prefill import PrefillManager
+    from freetoken.scheduler.table import TableManager
+    from freetoken.scheduler.utils import PendingReq
+
+    cm, prompt = _finish_read("all", READ, num_pages=512)  # room for both reads: no eviction
+    cm.isolate_logprob_reads = isolate
+    tree = (cm.prefix_cache.full_evictable, cm.prefix_cache.mamba_evictable_size)
+    pm = PrefillManager(cm, TableManager(max_running_reqs=2, page_table=cm.page_table),
+                        DecodeManager(page_size=1), isolate_logprob_reads=isolate)
+    pm.pending_list = [PendingReq(1, torch.tensor(prompt, dtype=torch.int32), READ)]
+    if isolate:  # nor may the host tier rehydrate a prefix for it
+        cm.host_tier = SimpleNamespace(has_kv=True, lookup=lambda *_: pytest.fail("host tier lookup"))
+    batch = pm.schedule_next_batch(4096)
+    (req,) = batch.reqs
+    assert req.cache_handle.cached_len == batch.log_cached_tokens == (0 if isolate else 128)
+
+    cm.allocate_paged(batch.reqs)
+    req.complete_one()
+    if isolate:  # a cold 130-token extend tracks the x64 boundary at 128 (build_fla_metadata)
+        req.mamba_next_track_idx, req.mamba_last_track_seqlen = 1, 128
+    cm.cache_req(req, finished=True)
+    if isolate:
+        assert (cm.prefix_cache.full_evictable, cm.prefix_cache.mamba_evictable_size) == tree
+    _assert_no_leak(cm, num_pages=512)
+
+
+@pytest.mark.parametrize(("policy", "isolate"), [("none", False), ("all", True), ("boundary", True)])
+def test_non_donating_read_chunk_commit_leaves_the_tree_alone(policy, isolate):
+    """The prefill-commit branch donated the x64 snapshot whatever the read policy said."""
+    num_pages, pool = 256, _pool()
+    pt = torch.zeros(4, 256, dtype=torch.int32)
+    cm = CacheManager(num_pages, 1, pt, "hybrid_radix", linear_state_pool=pool,
+                      logprob_read_snapshot_donate=policy, isolate_logprob_reads=isolate)
+    prompt = list(range(1, 131))
+    mr = cm.match_req(_pend(prompt))
+    pt[0, :130] = cm._allocate(130)
+    row = pt[0, :130].clone()
+    req = Req(input_ids=torch.tensor(prompt + [7], dtype=torch.int32), table_idx=0, cached_len=130,
+              output_len=1, uid=0, sampling_params=READ, cache_handle=mr.cuda_handle)
+    req.linear_slot_idx, req.mamba_ping_pong = pool.alloc(1)[0], tuple(pool.alloc(2))
+    req.mamba_next_track_idx, req.mamba_last_track_seqlen = 1, 128
+    cm.lock(mr.cuda_handle)
+
+    cm.cache_req(req, finished=False)
+    assert cm.prefix_cache.full_evictable + cm.prefix_cache.full_protected == 0
+    assert cm.prefix_cache.mamba_evictable_size + cm.prefix_cache.mamba_protected == 0
+    assert torch.equal(pt[0, :130], row)  # the request still owns its pages
+
+    cm.cache_req(req, finished=True)
+    assert len(cm.free_slots) == num_pages
+    _assert_no_leak(cm)

tests/scheduler/test_scheduler_chunked_prefill.py

@@ -14,6 +14,7 @@
 
 from __future__ import annotations
 
+import pytest
 import torch
 
 CHUNK = 8
@@ -237,3 +238,61 @@ def test_the_sliding_window_pool_cap_cannot_recut_an_image():
         for r in batch.reqs:
             r.complete_one()
     assert ends[0] == 4900 and ends[-1] == 6000
+
+
+def _passes(pm, cm, prompts, budget, isolate=True):
+    """Drain ``prompts`` ([(uid, prompt_len, read)]) and return each pass's (uid, extend_len) list."""
+    from freetoken.core import SamplingParams
+    from freetoken.scheduler.utils import PendingReq
+
+    pm.isolate_logprob_reads = isolate
+    pm.pending_list = [
+        PendingReq(uid, torch.arange(n, dtype=torch.int32),
+                   SamplingParams(max_tokens=1, logprobs=True) if read else SamplingParams(max_tokens=4))
+        for uid, n, read in prompts
+    ]
+    passes = []
+    while pm.runnable:
+        batch = pm.schedule_next_batch(budget)
+        passes.append([(r.uid, r.extend_len) for r in batch.reqs])
+        cm.allocate_paged(batch.reqs)
+        for r in batch.reqs:
+            r.complete_one()
+    return passes
+
+
+READ, CHAT = True, False
+
+
+@pytest.mark.parametrize(
+    ("prompts", "expected"),
+    [
+        ([(1, 5, READ), (2, 6, READ)], [[(1, 5)], [(2, 6)]]),
+        ([(1, 5, CHAT), (2, 6, READ)], [[(1, 5)], [(2, 6)]]),
+        ([(1, 5, READ), (2, 6, CHAT)], [[(1, 5)], [(2, 6)]]),
+        ([(1, 5, CHAT), (2, 6, CHAT)], [[(1, 5), (2, 6)]]),
+        ([(1, 5, READ), (2, 6, READ), (3, 4, CHAT), (4, 3, CHAT)], [[(1, 5)], [(2, 6)], [(3, 4), (4, 3)]]),
+    ],
+)
+def test_isolated_logprob_read_prefills_alone(prompts, expected):
+    cm, _tm, _dm, pm = _build_managers(num_pages=64)
+    assert _passes(pm, cm, prompts, budget=32) == expected
+
+
+def test_isolation_off_keeps_reads_batched():
+    cm, _tm, _dm, pm = _build_managers(num_pages=64)
+    prompts = [(1, 5, READ), (2, 6, READ), (3, 4, CHAT)]
+    assert _passes(pm, cm, prompts, budget=32, isolate=False) == [[(1, 5), (2, 6), (3, 4)]]
+
+
+def test_isolated_chunked_read_continuations_stay_alone():
+    """Every chunk of a read runs alone and on the same cuts as a solo run (budget multiples)."""
+    cm, _tm, _dm, pm = _build_managers(num_pages=64)
+    prompts = [(1, 2 * CHUNK + 3, READ), (2, 4, CHAT)]
+    assert _passes(pm, cm, prompts, budget=CHUNK) == [[(1, CHUNK)], [(1, CHUNK)], [(1, 3)], [(2, 4)]]
+
+
+def test_isolated_read_waits_behind_a_chunked_chat_prompt():
+    cm, _tm, _dm, pm = _build_managers(num_pages=64)
+    prompts = [(1, CHUNK + 2, CHAT), (2, 4, READ)]
+    assert _passes(pm, cm, prompts, budget=CHUNK) == [[(1, CHUNK)], [(1, 2)], [(2, 4)]]