sagellm-kv-cache

KV Cache Management + KV Transfer for sageLLM inference engine.

Overview

This package provides efficient KV cache management and transfer for LLM inference.

Key Features:

• KV Pool: Block-based memory management with budget control.
• KV Transfer: Primitives for cross-node KV block migration.
• Observability: Metrics and hooks for cache monitoring.

Architecture

┌─────────────────────────────────────────────────────────────────────┐
│                    sagellm-control-plane                            │
│              (Scheduling: alloc/free/migrate decisions)             │
└────────────────────────────┬────────────────────────────────────────┘
                             │ KVCacheInterface
                             ▼
┌─────────────────────────────────────────────────────────────────────┐
│                      sagellm-kv-cache (This Package)                 │
│  ┌─────────────┐  ┌─────────────┐  ┌─────────────┐  ┌─────────────┐ │
│  │ PrefixCache │  │  KV Pool    │  │  Eviction   │  │ KV Transfer │ │
│  │  (Task2.1)  │  │  (Task2.2)  │  │  (Task2.3)  │  │  (Task1.3)  │ │
│  └─────────────┘  └─────────────┘  └─────────────┘  └──────┬──────┘ │
└────────────────────────────────────────────────────────────┼────────┘
                             ┌───────────────────────────────┘
                             │ Use CommBackend for transport
                             ▼
┌─────────────────────────────────────────────────────────────────────┐
│                         sagellm-comm                                 │
│              (Network Layer: Topology, Collectives)                 │
└─────────────────────────────────────────────────────────────────────┘

Installation

pip install isagellm-kv-cache

Quick Start (CPU-first)

KV Pool

from sagellm_kv_cache.pool import KVPool

# Create a KV pool with budget control
pool = KVPool(max_tokens=1024)

# Allocate KV cache block
handle = pool.alloc(num_tokens=128, device="cpu")
print(f"Allocated handle: {handle.handle_id}, Tokens: {handle.num_tokens}")

# Free the handle
pool.free(handle)

Prefix Cache (Task 2.1)

from sagellm_kv_cache import PrefixCache

# Create cache with block-based hashing
cache = PrefixCache(block_size=16, max_cached_blocks=100, enable_lru=True)

# Insert prefix blocks
tokens = list(range(48))  # 3 blocks
hashes = cache.compute_block_hashes(tokens)
blocks = [{"block_id": i} for i in range(len(hashes))]
cache.insert(hashes, blocks)

# Lookup with prefix overlap
hit_blocks, num_tokens = cache.lookup(hashes)
print(f"Reused {num_tokens} tokens from cache!")

# Check hit rate
stats = cache.get_stats()
print(f"Hit rate: {stats['hit_rate']:.1%}")

See examples/prefix_cache_example.py for comprehensive usage examples.

KV Cache Access Pattern Profiling

from sagellm_kv_cache.profiling import AccessStatsCollector

# Create statistics collector
collector = AccessStatsCollector()

# Record accesses during inference
collector.record_access("block_001", is_hit=True)
collector.record_access("block_002", is_hit=False)

# Export statistics to JSON
collector.export_stats("stats.json")

# Get summary
summary = collector.get_stats_summary()
print(f"Hit rate: {summary['hit_rate']:.2%}")
print(f"Total accesses: {summary['total_accesses']}")

CLI Tool - Generate Demo Data:

# Generate demo statistics
sage-kv-stats demo --num-accesses 1000 --output demo_stats.json

# Or use the Python script
python examples/kv_profiling_demo.py --num-accesses 500 --output stats.json

CLI Tool - Visualize Results:

# Generate heatmap
sage-kv-stats visualize --input stats.json --output heatmap.png

# Generate all visualizations with summary
sage-kv-stats visualize --input stats.json --type all --summary

# Or use the Python script
python scripts/visualize_access_pattern.py --input stats.json --type all --summary

Install visualization dependencies (matplotlib is optional):

pip install isagellm-kv-cache[visualization]

API Reference

Core Components

• PrefixCache (sagellm_kv_cache): Block-hash based prefix caching for cross-request KV reuse. Supports LRU eviction, hit rate tracking, and handle invalidation. See Task 2.1.
• KVPool (sagellm_kv_cache.pool): Main entry point for memory management. Handles allocation, freeing, and budget enforcement.
• KVHandle (sagellm_kv_cache): Represents a reference to allocated KV cache. Contains metadata like handle_id, dtype, layout.
• KVTransferEngine (sagellm_kv_cache): Handles moving KV blocks between nodes using sagellm-comm.
• EvictionManager (sagellm_kv_cache): Eviction policy management with LRU/FIFO strategies.
• SchedulerBridge (sagellm_kv_cache): Bridge between scheduler IR and KV pool operations.

Dependencies

• isagellm-protocol: Common data structures and protocol definitions.
• isagellm-backend: Backend abstraction.
• isagellm-comm: Communication layer for transfer.

Development

推荐使用仓库脚本完成开发环境安装（不创建 venv，复用现有 Python 环境）：

# 开发模式（默认）：standard 基础 + 本地 editable 覆盖依赖（--no-deps）
./quickstart.sh --dev

# 标准模式：依赖优先从 PyPI 安装，本仓库 editable 安装
./quickstart.sh --standard

# 查看帮助
./quickstart.sh --help

脚本会在安装前动态清理已安装的 isagellm-* 包，避免历史安装导致环境漂移； 安装失败时会打印详细日志，便于诊断。

1. Install dev dependencies:

   pip install -e .[dev]
   

2. Run tests:

   pytest
   

3. Linting:

   ruff check .
   

Version

Current version: 0.4.0.11 See CHANGELOG.md for history.