From fc30523843d79f9cf9c3d7ce497c31d33fd4810f Mon Sep 17 00:00:00 2001 From: "Gabriel A. Devenyi" Date: Tue, 20 Oct 2026 15:54:11 -0400 Subject: [PATCH 1/3] feat(scheduler): add the --logprob-read-isolate flag and the is_logprob_read predicate Assisted-by: Claude Opus 5.5 Co-Authored-By: Claude Opus 5.5 (1M context) Claude-Session: https://claude.ai/code/session_01Td2gieja9Da2rPTN17Z846 --- docs/cli.md | 1 + python/freetoken/core.py | 5 ++++ python/freetoken/scheduler/cache.py | 6 ++--- python/freetoken/scheduler/config.py | 2 ++ python/freetoken/scheduler/prefill.py | 5 ++++ python/freetoken/scheduler/scheduler.py | 2 ++ python/freetoken/server/args.py | 10 ++++++++ tests/scheduler/test_hybrid_cache_manager.py | 25 ++++++++++++++++++++ 8 files changed, 53 insertions(+), 3 deletions(-) diff --git a/docs/cli.md b/docs/cli.md index b5af7aa2d3..d9ec62a0d0 100644 --- a/docs/cli.md +++ b/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 | diff --git a/python/freetoken/core.py b/python/freetoken/core.py index 541c1de761..a375dc9dcf 100644 --- a/python/freetoken/core.py +++ b/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 diff --git a/python/freetoken/scheduler/cache.py b/python/freetoken/scheduler/cache.py index c29d578b7b..384e2e9d87 100644 --- a/python/freetoken/scheduler/cache.py +++ b/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 @@ -394,10 +395,9 @@ 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 + if req.sampling_params.is_logprob_read else "all" ) free_upto = old_handle.cached_len diff --git a/python/freetoken/scheduler/config.py b/python/freetoken/scheduler/config.py index 75502210ca..18981135e2 100644 --- a/python/freetoken/scheduler/config.py +++ b/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) diff --git a/python/freetoken/scheduler/prefill.py b/python/freetoken/scheduler/prefill.py index 9019604707..25437060e3 100644 --- a/python/freetoken/scheduler/prefill.py +++ b/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 @@ -282,6 +284,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 +312,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] = [] diff --git a/python/freetoken/scheduler/scheduler.py b/python/freetoken/scheduler/scheduler.py index 78d2edf25e..d2fb4c3c8f 100644 --- a/python/freetoken/scheduler/scheduler.py +++ b/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 diff --git a/python/freetoken/server/args.py b/python/freetoken/server/args.py index 4fe0a7a8e5..6ed30ff073 100644 --- a/python/freetoken/server/args.py +++ b/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", diff --git a/tests/scheduler/test_hybrid_cache_manager.py b/tests/scheduler/test_hybrid_cache_manager.py index 6b4f164b52..c5373959bc 100644 --- a/tests/scheduler/test_hybrid_cache_manager.py +++ b/tests/scheduler/test_hybrid_cache_manager.py @@ -344,3 +344,28 @@ 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 From fed01568cffe56f5594a61e388d4abe7f6bed896 Mon Sep 17 00:00:00 2001 From: "Gabriel A. Devenyi" Date: Tue, 20 Oct 2026 15:55:07 -0400 Subject: [PATCH 2/3] feat(scheduler): prefill an isolated logprob read in a pass of its own Assisted-by: Claude Opus 5.5 Co-Authored-By: Claude Opus 5.5 (1M context) Claude-Session: https://claude.ai/code/session_01Td2gieja9Da2rPTN17Z846 --- python/freetoken/scheduler/prefill.py | 6 ++ .../test_scheduler_chunked_prefill.py | 59 +++++++++++++++++++ 2 files changed, 65 insertions(+) diff --git a/python/freetoken/scheduler/prefill.py b/python/freetoken/scheduler/prefill.py index 25437060e3..e367d1e935 100644 --- a/python/freetoken/scheduler/prefill.py +++ b/python/freetoken/scheduler/prefill.py @@ -324,6 +324,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): @@ -346,6 +350,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: diff --git a/tests/scheduler/test_scheduler_chunked_prefill.py b/tests/scheduler/test_scheduler_chunked_prefill.py index a4c7e0659a..f0fe3a7a1b 100644 --- a/tests/scheduler/test_scheduler_chunked_prefill.py +++ b/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)]] From cc13fecce538bc7f38efa459b0f10cc313dfa588 Mon Sep 17 00:00:00 2001 From: "Gabriel A. Devenyi" Date: Tue, 20 Oct 2026 15:56:42 -0400 Subject: [PATCH 3/3] feat(scheduler): keep isolated logprob reads out of the prefix cache An isolated read matches with max_len 0 (no host-tier lookup) and donates nothing. The prefill-commit branch also skips the snapshot insert for a non-donating read; the scheduler only reaches it for requests that keep decoding, never a read. Assisted-by: Claude Opus 5.5 Co-Authored-By: Claude Opus 5.5 (1M context) Claude-Session: https://claude.ai/code/session_01Td2gieja9Da2rPTN17Z846 --- python/freetoken/scheduler/cache.py | 17 +++-- python/freetoken/scheduler/prefill.py | 6 +- tests/scheduler/test_hybrid_cache_manager.py | 68 +++++++++++++++++++- 3 files changed, 82 insertions(+), 9 deletions(-) diff --git a/python/freetoken/scheduler/cache.py b/python/freetoken/scheduler/cache.py index 384e2e9d87..a9ef9dca11 100644 --- a/python/freetoken/scheduler/cache.py +++ b/python/freetoken/scheduler/cache.py @@ -146,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): @@ -375,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 @@ -395,11 +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. - donate = ( - self.logprob_read_snapshot_donate - if req.sampling_params.is_logprob_read - 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) diff --git a/python/freetoken/scheduler/prefill.py b/python/freetoken/scheduler/prefill.py index e367d1e935..7536e7a786 100644 --- a/python/freetoken/scheduler/prefill.py +++ b/python/freetoken/scheduler/prefill.py @@ -75,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 diff --git a/tests/scheduler/test_hybrid_cache_manager.py b/tests/scheduler/test_hybrid_cache_manager.py index c5373959bc..d24de93d5c 100644 --- a/tests/scheduler/test_hybrid_cache_manager.py +++ b/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, @@ -369,3 +368,68 @@ def test_logprob_read_isolate_is_a_server_arg(): 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)