[Perf] Host-path costs: prefix-key caching, O(1) free-block count, per-run KV bookkeeping, one argmax per batch

tests/test_async_pipeline.py

@@ -132,11 +132,11 @@ def test_delayed_eos_discards_next_work_and_reclaims_after_gpu_completion(
         eos = eos[0]
     samples = []
 
-    def sample(logits, params, generator):
+    def sample_batch(logits, params, generators):
         samples.append("a")
-        return logits.new_tensor(eos, dtype=torch.long)
+        return logits.new_full((len(params),), eos, dtype=torch.long)
 
-    monkeypatch.setattr(engine.model_runner.sampler, "sample", sample)
+    monkeypatch.setattr(engine.model_runner.sampler, "sample_batch", sample_batch)
     engine.add_request("a", [2], SamplingParams(max_tokens=5, exit_threshold=1))
     engine.step()
     engine.step()

tests/test_kv_cache.py

@@ -356,3 +356,27 @@ def test_snapshot_refuses_leased_allocations():
     cache.pin_transfer("a", "t")
     with pytest.raises(RuntimeError, match="transfer"):
         cache.snapshot("a")
+
+
+def test_add_run_matches_sequential_add_and_row_runs_group_consecutive_positions():
+    import random
+
+    from vllm_rlt.core.kv_cache_manager import KVCacheManager, _WrittenPositions
+
+    rng = random.Random(3)
+    for _ in range(200):
+        a, b = _WrittenPositions(), _WrittenPositions()
+        for _ in range(rng.randrange(1, 8)):
+            start, count = rng.randrange(0, 12), rng.randrange(1, 6)
+            for position in range(start, start + count):
+                a.add(position)
+            b.add_run(start, count)
+            # Observable state: the contiguous prefix and membership of every position.
+            assert a.prefix == b.prefix
+            assert [p in a for p in range(24)] == [p in b for p in range(24)]
+    cache = make_cache()
+    assert cache.allocate("a", 6) and cache.allocate("b", 2)
+    batch = cache._prepare_batch(["a", "a", "a", "b", "a"], [0, 0, 1, 0, 0], [0, 1, 2, 0, 5])
+    runs = KVCacheManager.row_runs(batch)
+    assert [(d, s, c) for _, d, s, c in runs] == [(0, 0, 2), (1, 2, 1), (0, 0, 1), (0, 5, 1)]
+    assert KVCacheManager.row_runs(batch) is runs

tests/test_prefix_growth.py

@@ -66,6 +66,46 @@ def test_prefix_refcounts_and_eviction():
     assert c.num_free_blocks == 12
 
 
+def _scanned_free_blocks(c):
+    # The pre-incremental definition: free list plus prefix-cache blocks held only by the cache.
+    return len(c._free_blocks) + sum(
+        c._refs[b] == 1 for blocks in c._prefixes.values() for b in blocks
+    )
+
+
+def test_incremental_evictable_count_matches_scan_under_random_operations():
+    import random
+
+    rng = random.Random(7)
+    c = KVCacheManager(1, 1, 4, 64, 2, 2, enable_prefix_caching=True)
+    live = {}
+    prompts = [[rng.randrange(1, 6) for _ in range(rng.randrange(3, 12))] for _ in range(6)]
+    for step in range(400):
+        action = rng.random()
+        if action < 0.45 and len(live) < 10:
+            prompt = rng.choice(prompts)
+            rid = f"r{step}"
+            hit = c.lookup_prefix(prompt)
+            if c.allocate(rid, len(prompt) + 2, prefix=hit):
+                live[rid] = prompt
+                c.mark_imported_prefix(rid, len(prompt), start=len(hit) * c.block_size)
+                c.publish_prefix(rid, prompt, len(prompt), keys=c.prefix_keys(prompt))
+        elif action < 0.8 and live:
+            rid = rng.choice(sorted(live))
+            del live[rid]
+            c.free(rid)
+        elif live:
+            rid = rng.choice(sorted(live))
+            c.ensure_capacity(rid, len(live[rid]) + 2)
+        assert c.num_free_blocks == _scanned_free_blocks(c), step
+        assert c._evictable_blocks >= 0
+    for rid in list(live):
+        c.free(rid)
+    assert c.num_free_blocks == _scanned_free_blocks(c) == 64
+    assert c.allocate("all", 64)
+    assert c.num_free_blocks == _scanned_free_blocks(c) == 0
+
+
 def test_incremental_growth_matches_reserved_generation():
     m = model()
     p = SamplingParams(max_tokens=9, exit_threshold=0, ignore_eos=True)

tests/test_sampler.py

@@ -250,3 +250,25 @@ def test_rng_registry_drop_removes_only_that_requests_slot():
     registry.drop("dropped")
     assert registry.peek("kept") is kept
     assert dropped is not kept
+
+
+def test_sample_batch_matches_per_row_sampling():
+    import torch
+
+    from vllm_rlt import SamplingParams
+    from vllm_rlt.worker.sampler import Sampler
+
+    sampler = Sampler()
+    logits = torch.randn(5, 11)
+    greedy = [SamplingParams(temperature=0)] * 4
+    assert torch.equal(
+        sampler.sample_batch(logits, greedy, [None] * 4),
+        torch.stack([sampler.sample(row, p, None) for row, p in zip(logits[:4], greedy)]),
+    )
+    mixed = [SamplingParams(temperature=0), SamplingParams(temperature=0.7, seed=3)]
+    generators = [None, torch.Generator().manual_seed(3)]
+    reference = [None, torch.Generator().manual_seed(3)]
+    assert torch.equal(
+        sampler.sample_batch(logits, mixed, generators),
+        torch.stack([sampler.sample(r, p, g) for r, p, g in zip(logits[:2], mixed, reference)]),
+    )

vllm_rlt/core/kv_cache_manager.py

@@ -46,6 +46,19 @@ def add(self, position: int) -> None:
             self.pending.remove(self.prefix)
             self.prefix += 1
 
+    def add_run(self, start: int, count: int) -> None:
+        """Same result as ``add`` for ``start, ..., start + count - 1`` in order."""
+        end = start + count
+        if end <= self.prefix:
+            return
+        if start > self.prefix:
+            self.pending.update(range(start, end))
+            return
+        self.prefix = end
+        while self.prefix in self.pending:
+            self.pending.remove(self.prefix)
+            self.prefix += 1
+
     def __contains__(self, position: int) -> bool:
         return position < self.prefix or position in self.pending
 
@@ -153,6 +166,11 @@ def __init__(
         self.watermark_blocks = int(watermark_ratio * num_blocks) if incremental_allocation else 0
         self._refs = [0] * num_blocks
         self._prefixes = OrderedDict()
+        # Blocks referenced by a prefix entry, and how many of them only the cache
+        # holds (refs == 1). Maintained at every reference-count change so that
+        # num_free_blocks does not rescan the cache on each admission check.
+        self._prefix_blocks = set()
+        self._evictable_blocks = 0
         self._pending_prefixes = []
         self.prefix_hits = self.prefix_queries = 0
         if enable_prefix_caching and layout != "last_exited":
@@ -164,9 +182,14 @@ def attention_capabilities(self) -> BackendCapabilities:
 
     @property
     def num_free_blocks(self) -> int:
-        return len(self._free_blocks) + sum(
-            self._refs[b] == 1 for blocks in self._prefixes.values() for b in blocks
-        )
+        return len(self._free_blocks) + self._evictable_blocks
+
+    def _retain(self, b):
+        """Increment a block's reference count, tracking cache-only blocks."""
+        before = self._refs[b]
+        self._refs[b] = before + 1
+        if b in self._prefix_blocks:
+            self._evictable_blocks += (before + 1 == 1) - (before == 1)
 
     @property
     def num_used_blocks(self) -> int:
@@ -188,27 +211,63 @@ def required_blocks(self, max_tokens: int) -> int:
             raise ValueError("max_tokens must be a positive integer")
         return ((max_tokens + self.block_size - 1) // self.block_size) * self.storage_depths
 
+    @staticmethod
+    def row_runs(batch):
+        """Rows of a host batch grouped into runs of consecutive positions.
+
+        Each run is ``(allocation, depth, start, count)``. Prefill chunks are one
+        run per request and depth, so per-layer bookkeeping touches one record
+        per run instead of one per token. The grouping is memoized on the batch,
+        which is immutable for its traversal.
+        """
+        runs = getattr(batch, "_runs", None)
+        if runs is None:
+            grouped = []
+            for allocation, depth, position in batch.rows:
+                last = grouped[-1] if grouped else None
+                if (
+                    last is not None
+                    and last[0] is allocation
+                    and last[1] == depth
+                    and last[2] + last[3] == position
+                ):
+                    last[3] += 1
+                else:
+                    grouped.append([allocation, depth, position, 1])
+            runs = tuple(tuple(run) for run in grouped)
+            object.__setattr__(batch, "_runs", runs)
+        return runs
+
     def _drop_refs(self, blocks):
         for b in blocks:
-            self._refs[b] -= 1
+            before = self._refs[b]
+            self._refs[b] = before - 1
             assert self._refs[b] >= 0
+            if b in self._prefix_blocks:
+                self._evictable_blocks += (before - 1 == 1) - (before == 1)
             if not self._refs[b]:
                 self._free_blocks.append(b)
 
+    def _evict_entry(self, blocks):
+        for b in blocks:
+            self._prefix_blocks.discard(b)
+            if self._refs[b] == 1:
+                self._evictable_blocks -= 1
+        self._drop_refs(blocks)
+
     def _claim(self, count):
         """Take ``count`` free blocks, evicting prefix entries only if that is enough.
 
         ``num_free_blocks`` already includes every block eviction could
         release. If it is still short, return ``None`` without evicting, so a
         failed admission, growth, resumption, or PD reservation does not
-        throw away prefixes it could not use anyway. ``num_free_blocks`` scans
-        every prefix entry, so it is only consulted when eviction is needed.
+        throw away prefixes it could not use anyway.
         """
         if len(self._free_blocks) < count and self.num_free_blocks < count:
             return None
         while len(self._free_blocks) < count and self._prefixes:
             _, blocks = self._prefixes.popitem(last=False)
-            self._drop_refs(blocks)
+            self._evict_entry(blocks)
         if len(self._free_blocks) < count:
             return None
         blocks = [self._free_blocks.pop() for _ in range(count)]
@@ -224,17 +283,26 @@ def _prefix_keys(self, tokens):
             digest = hashlib.sha256(digest + struct.pack(f"<{len(values)}q", *values)).digest()
             yield digest
 
+    def prefix_keys(self, tokens):
+        """Chained block keys of a prompt, excluding its final token.
+
+        A prompt's keys never change, so callers cache the tuple on the request
+        and pass a prefix of it to lookup_prefix/publish_prefix instead of
+        rehashing the prompt on every admission attempt.
+        """
+        return tuple(self._prefix_keys(tokens[:-1]))
+
     def poll_prefixes(self):
         pending, self._pending_prefixes = self._pending_prefixes, []
-        for request_id, allocation, tokens, length, event in pending:
+        for request_id, allocation, tokens, length, event, keys in pending:
             if self._allocations.get(request_id) is not allocation:
                 continue
             if event is not None and not event.query():
-                self._pending_prefixes.append((request_id, allocation, tokens, length, event))
+                self._pending_prefixes.append((request_id, allocation, tokens, length, event, keys))
                 continue
-            self.publish_prefix(request_id, tokens, length)
+            self.publish_prefix(request_id, tokens, length, keys=keys)
 
-    def publish_prefix(self, request_id, tokens, length, event=None):
+    def publish_prefix(self, request_id, tokens, length, event=None, *, keys=None):
         """Index complete prompt pages for reuse without copying their KV.
 
         length is the prefilled extent; len(tokens) is the full prompt length.
@@ -244,13 +312,19 @@ def publish_prefix(self, request_id, tokens, length, event=None):
             return
         allocation = self._get_allocation(request_id)
         if event is not None:
-            self._pending_prefixes.append((request_id, allocation, tuple(tokens), length, event))
+            self._pending_prefixes.append(
+                (request_id, allocation, tuple(tokens), length, event, keys)
+            )
             return
         # Prefix entries retain KV, not the final hidden state needed by CODA.
         # Leave at least the last prompt token to recompute on a cache hit;
         # full-page alignment may leave additional tokens to recompute.
         publishable_tokens = min(length, len(tokens) - 1)
-        for index, key in enumerate(self._prefix_keys(tokens[:publishable_tokens])):
+        if keys is None:
+            keys = self._prefix_keys(tokens[:publishable_tokens])
+        else:
+            keys = keys[: publishable_tokens // self.block_size]
+        for index, key in enumerate(keys):
             if key in self._prefixes:
                 continue
             end = (index + 1) * self.block_size
@@ -260,15 +334,20 @@ def publish_prefix(self, request_id, tokens, length, event=None):
                 break
             blocks = tuple(t[index] for t in allocation.block_tables)
             for b in blocks:
+                # The cache's own reference; the block becomes evictable once the
+                # publishing allocation (and any later sharer) releases it.
                 self._refs[b] += 1
+                self._prefix_blocks.add(b)
+                if self._refs[b] == 1:
+                    self._evictable_blocks += 1
             self._prefixes[key] = blocks
 
-    def lookup_prefix(self, tokens):
+    def lookup_prefix(self, tokens, *, keys=None):
         if not self.enable_prefix_caching:
             return ()
         self.poll_prefixes()
         found = []
-        for key in self._prefix_keys(tokens[:-1]):
+        for key in self._prefix_keys(tokens[:-1]) if keys is None else keys:
             blocks = self._prefixes.get(key)
             if blocks is None:
                 break
@@ -295,7 +374,7 @@ def allocate(self, request_id: str, max_tokens: int, *, initial_tokens=None, pre
             raise ValueError("prefix exceeds initial KV frontier")
         claimed = [b for group in prefix for b in group]
         for b in claimed:
-            self._refs[b] += 1
+            self._retain(b)
         fresh = self._claim((pages - len(prefix)) * self.storage_depths)
         if fresh is None:
             self._drop_refs(claimed)
@@ -680,8 +759,8 @@ def _write_prepared(
             return
         self.key_cache[batch.write_blocks, layer, batch.write_offsets] = k[: len(batch.rows)]
         self.value_cache[batch.write_blocks, layer, batch.write_offsets] = v[: len(batch.rows)]
-        for allocation, depth, position in batch.rows:
-            allocation.written[self._plane(depth)][layer].add(position)
+        for allocation, depth, start, count in self.row_runs(batch):
+            allocation.written[self._plane(depth)][layer].add_run(start, count)
 
     def _require_prefix(self, allocation, layer, depth, length):
         written = allocation.written[self._plane(depth)][layer]
@@ -723,8 +802,9 @@ def _attend_prepared(
             )
         if not batch.rows:
             return torch.empty_like(q)
-        for allocation, depth, position in batch.rows:
-            self._require_prefix(allocation, layer, depth, position + 1)
+        for allocation, depth, start, count in self.row_runs(batch):
+            # The run's last row carries the strongest requirement.
+            self._require_prefix(allocation, layer, depth, start + count)
         if batch.cu_seqlens_q is not None:
             return self.attention.prefill(
                 q,