Crate kvbm_engine

Expand description

§kvbm-engine

kvbm-engine provides distributed coordination primitives for KV Block Management (KVBM). It implements a tiered storage model where KV cache blocks flow between GPU memory, host DRAM, local disk, and object storage. The crate coordinates leaders (which own block metadata and make placement decisions) with workers (which execute data transfers via RDMA, NVMe, or object storage APIs).

§Storage Tier Model

Tier	Medium	Latency	Capacity	Description
G1	GPU HBM	~ns	Smallest	Active KV cache used by attention kernels
G2	Pinned DRAM	~us	Medium	Staging area for RDMA transfers and tier promotion
G3	NVMe/SSD	~ms	Large	Persistent warm-block storage
G4	S3/MinIO	~100ms	Unlimited	Cold/archival object storage

§Architecture

                    +-----------------+
                    | InstanceLeader  |
                    |  (find_matches, |
                    |   BlockAccessor)|
                    +--------+--------+
                             |
               +-------------+-------------+
               |                           |
      +--------v--------+        +--------v--------+
      | CoordinatedWorker|       | CoordinatedWorker|
      |   (rank 0)       |       |   (rank 1)       |
      +--------+---------+       +--------+---------+
               |                           |
      +--------v--------+        +--------v--------+
      | PhysicalWorker   |       | PhysicalWorker   |
      | (TransferManager)|       | (TransferManager)|
      +-----------------+        +-----------------+

The leader drives workers through the ParallelWorkers trait (SpmdParallelWorkers for SPMD execution). For onboarding, the leader creates sessions that progress through stages: search, hold, prepare (G3->G2), and pull (remote G2->local G2 via RDMA).

§Modules

Module	Purpose
`leader`	Block coordination: matching, onboarding sessions, policy-based scanning
`worker`	Transfer execution: local, RDMA, and object storage data movement
`object`	G4 storage: S3/MinIO client for cold-tier block persistence
`offload`	Tier demotion pipeline: batched G2->G3 and G2->G4 offloading
`runtime`	Shared infrastructure: `KvbmRuntime`, tokio handle, NIXL agent
`pubsub`	Event pub/sub: block-level notifications for cross-instance coordination
`collectives`	NCCL collectives for multi-GPU synchronization (feature-gated)
`testing`	Test utilities: mock workers, in-memory block managers (feature-gated)

§Feature Flags

Flag	Dependencies	Description
`default`	`["s3"]`	Default features
`s3`	`aws-sdk-s3`, `aws-config`, `rayon`, `tokio-rayon`, `chrono`	S3/MinIO object storage support
`collectives`	`nixl-sys`, `nccl`	NIXL + NCCL multi-GPU collectives
`nccl`	`cudarc`	NCCL support via cudarc
`testing-nccl`	`collectives`	Enable collectives for tests
`nats`	`async-nats`, `flume`	NATS-based pub/sub transport
`testing`	`kvbm-logical/testing`, `kvbm-physical/testing`	Test utilities and mock infrastructure
`nvtx`	`kvbm-config/nvtx`	NVIDIA Tools Extension profiling markers

§Quick Start

use kvbm_engine::{KvbmRuntime, leader::InstanceLeader};

// Build runtime from environment
let runtime = KvbmRuntime::from_env_leader().await?;

// Create a leader instance
let leader = InstanceLeader::new(/* ... */);

// Search for cached blocks
let result = leader.find_matches(&sequence_hashes)?;

Re-exports§

pub use runtime::KvbmRuntime;
pub use runtime::KvbmRuntimeBuilder;
pub use runtime::RuntimeHandle;

Modules§

leader: Leader Module
object: Object Storage Module
offload: Offload Module
pubsub: PubSub abstraction for distributed messaging.
runtime: Runtime
worker: Worker Module

Structs§

G1: GPU/device tier – HBM KV cache. Fastest access, smallest capacity. Blocks here are actively used by attention kernels.
G2: CPU/host tier – pinned DRAM cache. Microsecond-latency staging area for RDMA transfers and G3/G4 promotion.
G3: Disk tier – NVMe/SSD cache. Millisecond-latency persistent storage for warm blocks.
G4: Object store tier – S3/MinIO. Highest latency but unlimited capacity for cold/archival blocks.
InstanceId: Unique identifier for a runtime instance.
PeerInfo: Peer information combining instance ID and worker address.
WorkerAddress: Opaque worker address for discovery.

Enums§

LogicalLayoutHandle: Logical layout handle type encoding the layout ID.

Type Aliases§

BlockId
SequenceHash