From 0a499af80183f9d5ed1e680cdab7b4a2c5d38d32 Mon Sep 17 00:00:00 2001 From: 0z5a Date: Sun, 25 Oct 2026 17:32:23 +0100 Subject: [PATCH] [Perf] Host-path costs: prefix-key caching, O(1) free-block count, per-run KV bookkeeping, one argmax per batch Four host-side fixes found by profiling multi-turn session workloads (cProfile inside the benchmark; the GPU was 13% busy and Python 100% in one cell): - Scheduler._plan_admission rehashed every waiting prompt (one chained SHA-256 per block) on every admission attempt, and num_free_blocks rescanned every prefix entry on each capacity check. Prompt keys are now computed once per request, and the manager keeps the count of cache-only blocks (refs == 1) current at every reference-count change, so num_free_blocks is O(1). A randomized test checks the counter against the old scan. - Growth reservations are computed once per admission tick instead of once per admission candidate. - Rows of a host KV batch are grouped once into runs of consecutive positions per (allocation, depth); per-layer write/attend bookkeeping advances one written record per run (add_run, identical to sequential add) and checks the strongest prefix requirement per run. A 2048-token prefill chunk costs one call per layer instead of 2048. The per-layer written contract is unchanged; a randomized test checks add_run against add. In the session profile this bookkeeping was 31M calls, about 20 s of 146 s. - An all-greedy CODA batch samples with one argmax launch (Sampler.sample_batch) instead of one per row. RTX 5090, Ouro-1.4B, 32 agent sessions, 6 turns, prefix cache on, one run per arm (prefix-key and free-block changes vs main): per-row prefill: +2.6% tok/s, TTFT p50 8.4 -> 7.1 s, Python calls -19% packed prefill: +4.7% tok/s, TTFT p50 2.57 -> 1.80 s The other two changes were not measured on their own. Signed-off-by: 0z5a Co-Authored-By: Claude Opus 5.5 --- tests/test_async_pipeline.py | 6 +- tests/test_kv_cache.py | 24 ++++++ tests/test_prefix_growth.py | 40 ++++++++++ tests/test_sampler.py | 22 ++++++ vllm_rlt/core/kv_cache_manager.py | 120 +++++++++++++++++++++++++----- vllm_rlt/core/scheduler.py | 43 +++++++++-- vllm_rlt/engine/stage_engine.py | 4 +- vllm_rlt/request.py | 2 + vllm_rlt/worker/cuda_graph.py | 13 ++-- vllm_rlt/worker/model_runner.py | 9 +-- vllm_rlt/worker/sampler.py | 15 ++++ 11 files changed, 255 insertions(+), 43 deletions(-) diff --git a/tests/test_async_pipeline.py b/tests/test_async_pipeline.py index 76c5dc7c..2126681c 100644 --- a/tests/test_async_pipeline.py +++ b/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() diff --git a/tests/test_kv_cache.py b/tests/test_kv_cache.py index d49c59f2..cf162ae1 100644 --- a/tests/test_kv_cache.py +++ b/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 diff --git a/tests/test_prefix_growth.py b/tests/test_prefix_growth.py index c31d4aeb..5e929870 100644 --- a/tests/test_prefix_growth.py +++ b/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) diff --git a/tests/test_sampler.py b/tests/test_sampler.py index a9595f67..fcf67851 100644 --- a/tests/test_sampler.py +++ b/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)]), + ) diff --git a/vllm_rlt/core/kv_cache_manager.py b/vllm_rlt/core/kv_cache_manager.py index 3bb6bf02..4eb80935 100644 --- a/vllm_rlt/core/kv_cache_manager.py +++ b/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, diff --git a/vllm_rlt/core/scheduler.py b/vllm_rlt/core/scheduler.py index f443d742..fb205805 100644 --- a/vllm_rlt/core/scheduler.py +++ b/vllm_rlt/core/scheduler.py @@ -372,7 +372,9 @@ def _reserved_growth_blocks(self, *, reserve_outputs: bool) -> int: reserved += max(0, cache.required_blocks(tokens) - allocated) return reserved - def _plan_admission(self, request: Request, active_count: int) -> AdmissionPlan: + def _plan_admission( + self, request: Request, active_count: int, *, reserved_growth=None + ) -> AdmissionPlan: """Describe prefix reuse, physical allocation and logical admission cost. With 8 prompt tokens and 5 outputs, capacity is 12: the last sampled @@ -385,7 +387,9 @@ def _plan_admission(self, request: Request, active_count: int) -> AdmissionPlan: """ cache = self.cache_manager capacity = len(request.prompt_token_ids) + request.sampling_params.max_tokens - 1 - prefix = cache.lookup_prefix(request.prompt_token_ids) + if request.prefix_keys is None and cache.enable_prefix_caching: + request.prefix_keys = cache.prefix_keys(request.prompt_token_ids) + prefix = cache.lookup_prefix(request.prompt_token_ids, keys=request.prefix_keys) cached_tokens = len(prefix) * cache.block_size initial_tokens = min( len(request.prompt_token_ids), cached_tokens + self.config.prefill_chunk_size @@ -399,7 +403,14 @@ def _plan_admission(self, request: Request, active_count: int) -> AdmissionPlan: else capacity ) required = cache.required_blocks(budget_tokens) - len(prefix) * cache.storage_depths - reserved = self._reserved_growth_blocks(reserve_outputs=reserve_outputs) + # Growth reservations only change when a request is admitted, resumed or + # preempted; one admission tick may pass a value computed once for all + # candidates instead of rescanning every request per candidate. + reserved = ( + self._reserved_growth_blocks(reserve_outputs=reserve_outputs) + if reserved_growth is None + else reserved_growth + ) # Allow an otherwise idle engine to start a large request. admission_headroom_blocks = cache.watermark_blocks if active_count else 0 # num_free_blocks includes evictable prefix-only blocks. Reusing those @@ -490,6 +501,10 @@ def _admit(self): waiting = self._order_waiting_requests() deferred = [] scan_count = min(len(waiting), self.config.admission_scan_limit) + reserve_outputs = ( + self.cache_manager.incremental_allocation and self.suspend_callback is None + ) + reserved_growth = None # valid until an admission, resume or preemption changes it for _ in range(scan_count): if not waiting: break @@ -507,27 +522,37 @@ def _admit(self): # Imported and local requests compete in this same admission queue. self.cache_manager.unpin_transfer(request_id, request_id) self.enqueue(request, RequestPhase.CODA) + reserved_growth = None elif self.cache_manager.has_allocation( request_id ) and self.cache_manager.has_transfer_lease(request_id): # Outgoing transfer reservation already allocated the prompt KV. self._enter_prefill(request) + reserved_growth = None else: resume_result = self._try_resume(request) if resume_result == ResumeResult.RESTORED: active_count += 1 + reserved_growth = None continue if resume_result == ResumeResult.BLOCKED: deferred.append(request_id) continue - plan = self._plan_admission(request, active_count) + if reserved_growth is None: + reserved_growth = self._reserved_growth_blocks(reserve_outputs=reserve_outputs) + plan = self._plan_admission(request, active_count, reserved_growth=reserved_growth) admitted = self._commit_admission(request, plan) while not admitted and self.preempt(request, waiting_only=True): # Recompute prefix claims and growth reservations after # releasing a resident allocation; the previous plan is stale. - plan = self._plan_admission(request, active_count) + reserved_growth = self._reserved_growth_blocks(reserve_outputs=reserve_outputs) + plan = self._plan_admission( + request, active_count, reserved_growth=reserved_growth + ) admitted = self._commit_admission(request, plan) + if admitted: + reserved_growth = None if not admitted: deferred.append(request_id) if request.admission_bypasses >= self.config.max_admission_bypasses: @@ -1020,6 +1045,7 @@ def _apply_split_prefill_progress( request.prompt_token_ids, request.num_prefilled_tokens, event, + keys=request.prefix_keys, ) request.prefill_reclaimable_tokens = end prompt_complete = end == len(request.prompt_token_ids) @@ -1045,7 +1071,11 @@ def _apply_prefill_progress( request.num_prefilled_tokens += scheduled_request.token_count final = request.num_prefilled_tokens == len(request.prompt_token_ids) self.cache_manager.publish_prefix( - request.request_id, request.prompt_token_ids, request.num_prefilled_tokens, event + request.request_id, + request.prompt_token_ids, + request.num_prefilled_tokens, + event, + keys=request.prefix_keys, ) if final: request.loops_done = self.total_ut_steps @@ -1128,6 +1158,7 @@ def _apply_wavefront_prefill_batch_progress( request.prompt_token_ids, request.num_prefilled_tokens, event, + keys=request.prefix_keys, ) if prompt_complete: request.loops_done = self.total_ut_steps diff --git a/vllm_rlt/engine/stage_engine.py b/vllm_rlt/engine/stage_engine.py index 40a6c9a8..47b5e64c 100644 --- a/vllm_rlt/engine/stage_engine.py +++ b/vllm_rlt/engine/stage_engine.py @@ -584,7 +584,9 @@ def accept_import(self, request, hidden, *, cached_tokens=0): raise ValueError("request is not awaiting an import") count = len(request.prompt_token_ids) self.cache_manager.mark_imported_prefix(request.request_id, count, start=cached_tokens) - self.cache_manager.publish_prefix(request.request_id, request.prompt_token_ids, count) + self.cache_manager.publish_prefix( + request.request_id, request.prompt_token_ids, count, keys=request.prefix_keys + ) request.num_prefilled_tokens = request.num_preluded_tokens = count request.prefill_window_origin = count request.prefill_reclaimable_tokens = count diff --git a/vllm_rlt/request.py b/vllm_rlt/request.py index ef784437..4872bc72 100644 --- a/vllm_rlt/request.py +++ b/vllm_rlt/request.py @@ -67,6 +67,8 @@ class Request: compute_released: bool = False finish_reason: FinishReason | None = None exit_trace: tuple[int, ...] = field(default=(), repr=False) + # Prefix-cache block keys of the prompt, computed once on first admission attempt. + prefix_keys: tuple[bytes, ...] | None = field(default=None, repr=False) @staticmethod def prefill_boundaries( diff --git a/vllm_rlt/worker/cuda_graph.py b/vllm_rlt/worker/cuda_graph.py index db1a7519..5647d837 100644 --- a/vllm_rlt/worker/cuda_graph.py +++ b/vllm_rlt/worker/cuda_graph.py @@ -168,12 +168,12 @@ def prepare_inputs(self, call, hidden): seen = set() # Device-owned progress is ordered by the runner's stream. The CPU # written frontier intentionally lags in-flight speculative rounds. - host_rows = () if isinstance(batch, DeviceKVMetadata) else batch.rows - for allocation, depth, pos in host_rows: + runs = () if isinstance(batch, DeviceKVMetadata) else cache.row_runs(batch) + for allocation, depth, start, _ in runs: key = (id(allocation), depth) if key not in seen: for layer in cache.recurrent_layers: - cache._require_prefix(allocation, layer, depth, pos) + cache._require_prefix(allocation, layer, depth, start) seen.add(key) stream = torch.cuda.current_stream(cache.device) if self.last_event is not None: @@ -206,10 +206,11 @@ def execute(self, call, hidden): output = tuple(t.clone() if t is not None else None for t in entry.output) self.last_event = torch.cuda.Event() self.last_event.record(torch.cuda.current_stream(self.cache.device)) - host_rows = () if isinstance(batch, DeviceKVMetadata) else batch.rows - for allocation, depth, pos in host_rows: + runs = () if isinstance(batch, DeviceKVMetadata) else self.cache.row_runs(batch) + for allocation, depth, start, count in runs: + plane = allocation.written[self.cache._plane(depth)] for layer in self.cache.recurrent_layers: - allocation.written[self.cache._plane(depth)][layer].add(pos) + plane[layer].add_run(start, count) self.replays += 1 return output diff --git a/vllm_rlt/worker/model_runner.py b/vllm_rlt/worker/model_runner.py index ada23d20..b760d019 100644 --- a/vllm_rlt/worker/model_runner.py +++ b/vllm_rlt/worker/model_runner.py @@ -1123,13 +1123,8 @@ def _execute_batch(self, batch: ExecutionBatch, prepared): ) values = selected.float().sigmoid() elif forward.compute_logits: - values = torch.stack( - [ - self.sampler.sample(row, params, generator) - for row, params, generator in zip( - result.logits, prepared.sampling_params, prepared.sampling_generators - ) - ] + values = self.sampler.sample_batch( + result.logits, prepared.sampling_params, prepared.sampling_generators ) if prepared.routing is not None: prepared.routing.scatter(values, tokens=True) diff --git a/vllm_rlt/worker/sampler.py b/vllm_rlt/worker/sampler.py index 62bc7846..21989736 100644 --- a/vllm_rlt/worker/sampler.py +++ b/vllm_rlt/worker/sampler.py @@ -59,6 +59,21 @@ def draw(probs: torch.Tensor, generator: torch.Generator) -> torch.Tensor: """Draw one token from a normalized distribution as a 0-dim long tensor.""" return torch.multinomial(probs, 1, generator=generator).squeeze(0) + def sample_batch(self, logits, params, generators): + """Sample one token per row; all-greedy batches use a single argmax launch. + + ``logits`` has at least ``len(params)`` rows. Row ``i`` is sampled with + ``params[i]`` and ``generators[i]`` exactly as ``sample`` would; greedy rows + are independent of each other, so a batched argmax is the same result in + one kernel instead of one per row plus a stack. + """ + rows = len(params) + if all(p.temperature == 0 for p in params): + return logits[:rows].argmax(-1) + return torch.stack( + [self.sample(row, p, g) for row, p, g in zip(logits[:rows], params, generators)] + ) + def sample( self, logits: torch.Tensor,