[presentyourlove]

Description

This PR implements Frontier Dimension 2.1: Distributed PagedAttention & Speculative KV Prefetching, eliminating memory fragmentation and masking cross-node/NVMe I/O transfer latency through a pipelined double-buffering architecture.

Background & Motivation

In distributed edge inference and long-context generation (8K+ tokens), KV cache memory fragmentation severely degrades memory efficiency. Furthermore, reloading historical KV blocks from NVMe SSDs or remote peer nodes introduces noticeable multi-millisecond decoding stalls.
By decoupling contiguous KV tensors into fixed-size virtual pages (PagedKVBlock & BlockTable) and executing speculative background prefetching concurrently with forward compute layers, I/O transfer latencies are completely hidden behind compute time, achieving "Zero-Latency Perception".

Key Changes

1. Virtual Paged KV Types & Telemetry (src/exo/shared/types/paged_kv.py):

   • PagedKVBlock: Immutable page unit holding discrete KV projections with device location tracking (local_ram, nvme_ssd, remote_node).
   • BlockTable: Manages logical token index to physical/remote page mappings.
   • PrefetchRequest & PrefetchTelemetry: Tracks hit rates, masked I/O latency, and pipeline masking efficiency.

2. Speculative Predictor & Double-Buffering Pipeline (src/exo/worker/engines/paged_attention/prefetch_manager.py):

   • SpeculativePrefetchPredictor: Predicts future blocks along sequential decoding trajectories while retaining critical initial attention sink pages.
   • DoubleBufferPipeline: Orchestrates concurrent background I/O fetching while the main thread executes forward layer computation, enabling zero-wait handover.

3. Paged Cache Engine (src/exo/worker/engines/paged_attention/paged_cache_engine.py):

   • PagedCacheEngine: High-level orchestrator handling token appending, cross-tier page eviction, speculative prefetching execution, and transparent fallback.

4. Module Exports (src/exo/worker/engines/paged_attention/__init__.py):

   • Exported PagedCacheEngine, DoubleBufferPipeline, and SpeculativePrefetchPredictor.

5. Unit Tests (src/exo/shared/tests/test_paged_kv_prefetch.py):

   • Verified page block allocation and token capacity metrics.
   • Verified speculative predictor candidates and attention sink retention.
   • Verified concurrent double-buffering pipeline execution and compute overlap.
   • Verified end-to-end block eviction, speculative prefetch hit, and telemetry calculation.

Verification

• Ran unit tests:

  .venv/Scripts/pytest.exe src/exo/shared/tests/test_paged_kv_prefetch.py
  # 5 passed in 0.12s

• Code styling and linting:

  .venv/Scripts/ruff.exe check src/exo/shared/types/paged_kv.py src/exo/worker/engines/paged_attention/prefetch_manager.py src/exo/worker/engines/paged_attention/paged_cache_engine.py src/exo/worker/engines/paged_attention/__init__.py src/exo/shared/tests/test_paged_kv_prefetch.py
  # All checks passed!

Related Issues

• Resolves: Frontier Dimension 2.1 in 人人都能跑llm.md (Distributed PagedAttention & Speculative KV Prefetching)