Skip to content

Worker decouple — Phase 2 / Slice 5: KVCacheManager (biggest cohesive subsystem) #175

Description

@Andrewxu313

Phase 2 of the worker decouple initiative. Largest single-slice extraction.

Scope

Extract the KV Cache Helper Methods section (worker:2578–3908, 27 methods, ~1330 LOC) into batchgen/worker/kv_manager.py. Handles GPU + host KV allocation, eviction, migration, async load lifecycle.

Design

class GpuKvBackend(Protocol):
    def allocate_pages(self, uuid: str, n: int) -> list[int]: ...
    def release_pages(self, uuid: str) -> None: ...
    def free_pages(self) -> int: ...

class HostKvBackend(Protocol):
    def allocate_pages(self, uuid: str, n: int) -> list[int]: ...
    def load_to_gpu_async(self, uuid: str, page_ids: list[int]) -> AsyncHandle: ...

class KVCacheManager:
    def __init__(self, *, gpu_kv: GpuKvBackend, host_kv: HostKvBackend): ...
    def plan_allocation(self, req: AllocRequest) -> AllocPlan: ...
    def plan_eviction(self, req: EvictionRequest) -> EvictionPlan: ...
    def plan_migration(self, req: MigrationRequest) -> MigrationPlan: ...

Handler returns plans; worker applies them. No self.state.X writes.

3-PR plan

  • PR-5.1 Port (≤500 LOC; may split into PR-5.1a/b/c if larger).
  • PR-5.2 Dual-path gate + compare-mode (BATCHGEN_WORKER_KV_*). H20 validate L2-2048 + L2-4096 back-to-back (KV slices specifically must validate decode→prefill cleanup).
  • PR-5.3 Cleanup.

Pre-merge checklist (from scheduler-split Step 5 fix-up commits)

  • No self.host_paged_kv_worker_view reference drift — handler receives backend by injection.
  • No self.model_config access — passed via frozen snapshot.
  • host_kv_eviction_watermark is computed from snapshot, not read from instance state.
  • _pending_kv_append_tensors lifecycle stays on worker (handler is stateless).

Design source reference

origin/tairan/scheduler-split:batchgen/worker/kv_manager.py (637 LOC) — read for shape only.

Design: /Users/andrew/.claude/plans/the-task-for-this-declarative-melody.md (POIS-local). Builds on the Phase A/B/C cuda-graph contract pattern (PR #162/#163/#164, merged 2026-05-27).

Metadata

Metadata

Assignees

Labels

enhancementNew feature or requestworker-decoupleWorker monolith decouple initiative (May 2026)

Type

No type

Projects

No projects

Relationships

None yet

Development

No branches or pull requests

Issue actions