diff --git a/README.md b/README.md index fa0118d..b3e6c6d 100644 --- a/README.md +++ b/README.md @@ -42,7 +42,8 @@ output = generate(model, **inputs, max_new_tokens=4096) | Greedy and sampled generation | Supported | | Beam search and batched generation | Not yet supported | | Sliding, chunked, and linear attention | Not yet supported | -| vLLM backend and optimized GPU kernels | Not yet implemented | +| vLLM integration | CPU-tested compaction planner; runtime patch pending | +| Optimized GPU kernels | Not yet implemented | ## Policy quickstart @@ -198,3 +199,6 @@ the launch criteria and current work packages. New contributors can start with a [`good first issue`](https://github.com/DaBestCode/randkv/issues?q=is%3Aissue+is%3Aopen+label%3A%22good+first+issue%22). Read [`CONTRIBUTING.md`](CONTRIBUTING.md) before opening a pull request. + +The pinned vLLM boundary, runtime lifecycle, and unsupported-mode decisions are +documented in [`docs/vllm-integration.md`](docs/vllm-integration.md). diff --git a/ROADMAP.md b/ROADMAP.md index 5021214..98ad6c9 100644 --- a/ROADMAP.md +++ b/ROADMAP.md @@ -13,7 +13,7 @@ vLLM throughput result. - [ ] Validate benchmark artifacts against a versioned schema. - [ ] Add an offline end-to-end Transformers integration test. - [ ] Measure task quality at matched dense and RandKV token budgets. -- [ ] Implement and test a vLLM-compatible policy/backend boundary. +- [x] Implement and test a vLLM-compatible policy/backend boundary. - [ ] Benchmark dense, random eviction, and at least one scored selector on an NVIDIA GPU under the same model, request distribution, and cache budget. - [ ] Publish peak-memory, throughput, latency, and quality results with raw diff --git a/docs/vllm-integration.md b/docs/vllm-integration.md new file mode 100644 index 0000000..0e07e0d --- /dev/null +++ b/docs/vllm-integration.md @@ -0,0 +1,98 @@ +# vLLM integration boundary + +This design targets vLLM `v0.28.0`, commit +`2cf0a6915ce544dc493a0990f2ea38d81601128a`. The integration uses internal +V1 scheduler and model-runner APIs, so this pin is part of the contract rather +than a statement of compatibility with later vLLM releases. + +## Call site + +The intended upstream call site adds one option and keeps RandKV defaults: + +```python +llm = LLM(model="Qwen/Qwen3-4B", kv_cache_policy="randkv") +``` + +An explicit experiment can still use one structured option: + +```python +llm = LLM( + model="Qwen/Qwen3-4B", + kv_cache_policy={"type": "randkv", "budget": 2048, "buffer_size": 64}, +) +``` + +This option does not exist in upstream vLLM `v0.28.0`; adding it requires a +small upstream patch. RandKV must not monkey-patch a running engine. + +## Why block eviction is insufficient + +vLLM allocates paged KV blocks per cache group, and one block table is shared +by every KV head in that group. RandKV selects tokens independently per KV +head. Selecting or freeing whole blocks would therefore force every head to +retain the same tokens and would no longer implement the paper's policy. + +The runtime must instead allocate dense destination blocks and compact each +head independently into those blocks. Source blocks can be released only after +every head has been copied. The CPU-only `VLLMCompactionPlanner` produces the +logical retained positions and exact source offsets for that operation. A +PyTorch fallback and a CUDA kernel must consume the same plan. + +## vLLM mapping + +| RandKV concept | vLLM `v0.28.0` mapping | +| --- | --- | +| Stable request identity | `Request.request_id`; never scheduler order or a Python hash | +| Prompt boundary | Original prompt-token count captured before chunked prefill | +| Absolute position | `Request.num_computed_tokens`; remains monotonic after compaction | +| Physical context length | New per-request retained-token count used by attention metadata | +| Eviction cadence | Plan after decode append when physical length exceeds `K + r`; subsequent rounds occur after `r` appended tokens | +| Per-head identity | Global KV-head index: tensor-parallel head offset plus local head index | +| Block allocation | `KVCacheManager` reserves exclusive destination blocks before source blocks are released | +| Worker operation | `SchedulerOutput` carries a compaction descriptor; the model runner applies copies before the next attention step | +| Observability | Scheduler counts plans/blocks; worker reports copied and evicted token-head pairs | + +The narrow hook spans `KVCacheManager.allocate_slots`, `SchedulerOutput`, and +the model runner's block-table/attention-metadata preparation. These are +internal APIs and therefore high upgrade risk. A scheduler-only plugin cannot +work because current vLLM uses `num_computed_tokens` for allocation progress +while attention also needs the smaller physical retained length. + +## Lifecycle and unsupported modes + +- **Prefill:** never compact an incomplete prompt. With chunked prefill, wait + until the entire prompt is computed and reject prompts larger than `budget`. +- **Prefix caching:** shared prompt blocks may remain read-only. The compacted + generated suffix must use copy-on-write destination blocks and must not enter + the prefix hash cache. A partial prompt-boundary block must be copied before + packing generated tokens. +- **Cancellation:** discard planner state and free destination plus remaining + source blocks through the normal request cleanup path. A partially applied + plan must be idempotently recoverable or fail the engine step. +- **Multi-request scheduling:** keep request-local absolute length, eviction + index, positions, and destination blocks. Determinism comes from the stable + request ID, not execution order. +- **Tensor parallelism:** pass each rank's global KV-head offset to the planner. + Pipeline stages must use global layer indices for the same reason. +- **Preemption/recompute:** the first runtime milestone should reject or disable + preemption for compacted requests. Reconstructing a non-contiguous cache from + token IDs would otherwise require replaying the retention history. +- **Speculative decoding, beam search, hybrid/sliding attention, KV transfer, + and disaggregated serving:** explicitly unsupported in the first integration. + +## Runtime implementation sequence + +1. Add the single `kv_cache_policy` configuration field and validation. +2. Add request-local physical-length and retention metadata without changing + monotonic `num_computed_tokens`. +3. Reserve exclusive destination blocks and send a typed compaction descriptor + in `SchedulerOutput`. +4. Implement the PyTorch copy fallback and verify it against planner offsets. +5. Update block tables and attention sequence lengths atomically after copies. +6. Add cancellation and two-concurrent-request integration tests. +7. Add the CUDA compactor behind the same descriptor, then run Issue #8's + matched serving benchmark. + +The upstream source anchors for this design are +`vllm/v1/core/kv_cache_manager.py`, `vllm/v1/core/sched/output.py`, and the V1 +GPU model runner at the pinned commit. diff --git a/src/randkv/__init__.py b/src/randkv/__init__.py index 4a1d6ba..47cf1a5 100644 --- a/src/randkv/__init__.py +++ b/src/randkv/__init__.py @@ -4,6 +4,14 @@ from .metrics import PolicyStats from .policy import RandomEvictionPolicy from .protocols import KVPolicy +from .vllm import VLLMCompactionPlan, VLLMCompactionPlanner -__all__ = ["KVPolicy", "PolicyStats", "RandKVConfig", "RandomEvictionPolicy"] +__all__ = [ + "KVPolicy", + "PolicyStats", + "RandKVConfig", + "RandomEvictionPolicy", + "VLLMCompactionPlan", + "VLLMCompactionPlanner", +] __version__ = "0.1.0" diff --git a/src/randkv/vllm.py b/src/randkv/vllm.py new file mode 100644 index 0000000..aad2bd7 --- /dev/null +++ b/src/randkv/vllm.py @@ -0,0 +1,194 @@ +"""Dependency-free planning boundary for a future vLLM runtime adapter.""" + +from __future__ import annotations + +from dataclasses import dataclass +from math import ceil +from typing import TypeAlias + +from .config import RandKVConfig +from .policy import RandomEvictionPolicy + +RequestId: TypeAlias = str | int | bytes + +VLLM_TARGET_VERSION = "0.28.0" +VLLM_TARGET_COMMIT = "2cf0a6915ce544dc493a0990f2ea38d81601128a" + + +@dataclass(frozen=True, slots=True) +class VLLMCompactionPlan: + """Backend-neutral description of one head-wise KV compaction. + + vLLM owns physical blocks shared by every KV head. The runtime adapter can + use ``source_offsets_by_head`` to copy each head's independently selected + tokens into a shared set of dense destination blocks. + """ + + request_id: RequestId + layer: int + head_offset: int + eviction_index: int + absolute_tokens_seen: int + block_size: int + source_tokens: int + retained_positions_by_head: tuple[tuple[int, ...], ...] + source_offsets_by_head: tuple[tuple[int, ...], ...] + + @property + def num_kv_heads(self) -> int: + return len(self.retained_positions_by_head) + + @property + def retained_tokens(self) -> int: + return len(self.retained_positions_by_head[0]) + + @property + def evicted_tokens_per_head(self) -> int: + return self.source_tokens - self.retained_tokens + + @property + def destination_blocks(self) -> int: + return ceil(self.retained_tokens / self.block_size) + + +class VLLMCompactionPlanner: + """Translate logical RandKV selections into head-wise copy offsets. + + This prototype deliberately imports no vLLM or tensor library. The future + runtime shim owns block allocation, copies, and attention metadata; this + class owns deterministic policy semantics only. + """ + + def __init__(self, config: RandKVConfig | None = None) -> None: + self.config = config or RandKVConfig() + self.policy = RandomEvictionPolicy(self.config) + + def plan( + self, + positions_by_head: tuple[tuple[int, ...], ...], + *, + prompt_length: int, + request_id: RequestId, + layer: int, + eviction_index: int, + absolute_tokens_seen: int, + block_size: int, + head_offset: int = 0, + ) -> VLLMCompactionPlan | None: + """Return a compaction plan, or ``None`` while under capacity. + + ``head_offset`` is the tensor-parallel rank's first global KV-head ID. + Using global IDs keeps a draw stable when the same request is sharded. + """ + + source_tokens = self._validate_inputs( + positions_by_head, + prompt_length=prompt_length, + request_id=request_id, + layer=layer, + eviction_index=eviction_index, + absolute_tokens_seen=absolute_tokens_seen, + block_size=block_size, + head_offset=head_offset, + ) + capacity = self.config.budget + self.config.buffer_size + if source_tokens <= capacity: + return None + + retained_positions = tuple( + self.policy.select( + positions, + prompt_length=prompt_length, + layer=layer, + kv_head=head_offset + local_head, + request_id=request_id, + eviction_index=eviction_index, + ) + for local_head, positions in enumerate(positions_by_head) + ) + retained_lengths = {len(positions) for positions in retained_positions} + if len(retained_lengths) != 1: + raise RuntimeError("all KV heads must retain the same physical length") + + source_offsets = tuple( + self._source_offsets(source, retained) + for source, retained in zip( + positions_by_head, retained_positions, strict=True + ) + ) + return VLLMCompactionPlan( + request_id=request_id, + layer=layer, + head_offset=head_offset, + eviction_index=eviction_index, + absolute_tokens_seen=absolute_tokens_seen, + block_size=block_size, + source_tokens=source_tokens, + retained_positions_by_head=retained_positions, + source_offsets_by_head=source_offsets, + ) + + def _validate_inputs( + self, + positions_by_head: tuple[tuple[int, ...], ...], + *, + prompt_length: int, + request_id: RequestId, + layer: int, + eviction_index: int, + absolute_tokens_seen: int, + block_size: int, + head_offset: int, + ) -> int: + if not positions_by_head: + raise ValueError("positions_by_head must contain at least one KV head") + if block_size <= 0: + raise ValueError("block_size must be greater than zero") + if absolute_tokens_seen < 0: + raise ValueError("absolute_tokens_seen must be non-negative") + if prompt_length < 0: + raise ValueError("prompt_length must be non-negative") + if prompt_length > self.config.budget: + raise ValueError( + f"prompt requires {prompt_length} cache slots but budget is " + f"{self.config.budget}; increase the budget" + ) + if layer < 0 or eviction_index < 0 or head_offset < 0: + raise ValueError( + "layer, eviction_index, and head_offset must be non-negative" + ) + if not isinstance(request_id, (str, int, bytes)): + raise TypeError("request_id must be a string, integer, or bytes") + + lengths = {len(positions) for positions in positions_by_head} + if len(lengths) != 1: + raise ValueError("all KV heads must have the same physical length") + for positions in positions_by_head: + if any(not isinstance(position, int) for position in positions): + raise TypeError("positions must contain only integers") + if any( + left >= right + for left, right in zip(positions, positions[1:], strict=False) + ): + raise ValueError("positions must be unique and strictly increasing") + if any( + position < 0 or position >= absolute_tokens_seen + for position in positions + ): + raise ValueError( + "retained positions must be non-negative and below " + "absolute_tokens_seen" + ) + prompt = tuple( + position for position in positions if position < prompt_length + ) + if prompt != tuple(range(prompt_length)): + raise ValueError("every KV head must contain the protected prompt") + return len(positions_by_head[0]) + + @staticmethod + def _source_offsets( + source: tuple[int, ...], retained: tuple[int, ...] + ) -> tuple[int, ...]: + offsets = {position: offset for offset, position in enumerate(source)} + return tuple(offsets[position] for position in retained) diff --git a/tests/test_vllm.py b/tests/test_vllm.py new file mode 100644 index 0000000..08d9c09 --- /dev/null +++ b/tests/test_vllm.py @@ -0,0 +1,160 @@ +import unittest + +from randkv import RandKVConfig +from randkv.vllm import VLLMCompactionPlanner + + +class VLLMCompactionPlannerTests(unittest.TestCase): + def setUp(self) -> None: + self.config = RandKVConfig(budget=4, buffer_size=2, seed=17) + + def test_returns_none_while_cache_is_within_capacity(self) -> None: + planner = VLLMCompactionPlanner(self.config) + positions = (tuple(range(6)),) * 2 + + plan = planner.plan( + positions, + prompt_length=2, + request_id="request-1", + layer=0, + eviction_index=0, + absolute_tokens_seen=6, + block_size=4, + ) + + self.assertIsNone(plan) + + def test_emits_independent_head_copy_offsets(self) -> None: + planner = VLLMCompactionPlanner(self.config) + source = tuple(range(12)) + + plan = planner.plan( + (source,) * 3, + prompt_length=2, + request_id="request-1", + layer=4, + eviction_index=0, + absolute_tokens_seen=12, + block_size=4, + ) + + self.assertIsNotNone(plan) + assert plan is not None + self.assertEqual(plan.retained_tokens, 6) + self.assertEqual(plan.evicted_tokens_per_head, 6) + self.assertEqual(plan.destination_blocks, 2) + self.assertEqual(plan.num_kv_heads, 3) + self.assertEqual(len(set(plan.retained_positions_by_head)), 3) + for positions, offsets in zip( + plan.retained_positions_by_head, + plan.source_offsets_by_head, + strict=True, + ): + self.assertEqual(positions[:2], (0, 1)) + self.assertEqual(positions[-2:], (10, 11)) + self.assertEqual(tuple(source[offset] for offset in offsets), positions) + + def test_plan_is_deterministic_across_planners(self) -> None: + arguments = { + "prompt_length": 2, + "request_id": "request-9", + "layer": 3, + "eviction_index": 2, + "absolute_tokens_seen": 20, + "block_size": 16, + "head_offset": 8, + } + positions = (tuple(range(20)),) * 2 + + first = VLLMCompactionPlanner(self.config).plan(positions, **arguments) + second = VLLMCompactionPlanner(self.config).plan(positions, **arguments) + + self.assertEqual(first, second) + + def test_global_head_offset_changes_the_draw(self) -> None: + planner = VLLMCompactionPlanner(self.config) + positions = (tuple(range(30)),) + arguments = { + "prompt_length": 2, + "request_id": "request-1", + "layer": 0, + "eviction_index": 0, + "absolute_tokens_seen": 30, + "block_size": 16, + } + + first_rank = planner.plan(positions, head_offset=0, **arguments) + second_rank = planner.plan(positions, head_offset=1, **arguments) + + self.assertIsNotNone(first_rank) + self.assertIsNotNone(second_rank) + assert first_rank is not None and second_rank is not None + self.assertNotEqual( + first_rank.retained_positions_by_head, + second_rank.retained_positions_by_head, + ) + + def test_supports_positions_retained_by_an_earlier_round(self) -> None: + planner = VLLMCompactionPlanner(self.config) + initial = tuple(range(12)) + first = planner.plan( + (initial,) * 2, + prompt_length=2, + request_id="request-1", + layer=0, + eviction_index=0, + absolute_tokens_seen=12, + block_size=4, + ) + assert first is not None + extended = tuple( + positions + (12, 13, 14) for positions in first.retained_positions_by_head + ) + + second = planner.plan( + extended, + prompt_length=2, + request_id="request-1", + layer=0, + eviction_index=1, + absolute_tokens_seen=15, + block_size=4, + ) + + self.assertIsNotNone(second) + assert second is not None + for positions in second.retained_positions_by_head: + self.assertEqual(positions[:2], (0, 1)) + self.assertEqual(positions[-2:], (13, 14)) + + def test_rejects_unequal_physical_head_lengths(self) -> None: + planner = VLLMCompactionPlanner(self.config) + + with self.assertRaisesRegex(ValueError, "same physical length"): + planner.plan( + (tuple(range(8)), tuple(range(7))), + prompt_length=2, + request_id="request-1", + layer=0, + eviction_index=0, + absolute_tokens_seen=8, + block_size=4, + ) + + def test_validates_state_even_when_no_compaction_is_due(self) -> None: + planner = VLLMCompactionPlanner(self.config) + + with self.assertRaisesRegex(ValueError, "protected prompt"): + planner.plan( + ((0, 2, 3),), + prompt_length=2, + request_id="request-1", + layer=0, + eviction_index=0, + absolute_tokens_seen=4, + block_size=4, + ) + + +if __name__ == "__main__": + unittest.main()