Module offload

Expand description

§Offload Module

The offload module manages the asynchronous transfer of KV cache blocks between storage tiers. It provides a pipeline-based architecture for evaluating, batching, and executing block transfers with full cancellation support.

§Overview

Offloading moves blocks from a source tier (e.g., GPU memory) to a destination tier (e.g., host memory, remote storage, or object storage). The pipeline ensures:

Policy-based filtering: Only blocks meeting criteria are transferred
Batched execution: Blocks are grouped for efficient transfer
Cancellation support: Transfers can be cancelled at any point before commitment
Precondition synchronization: Transfers wait for forward pass completion

§Pipeline Architecture

┌─────────────────┐     ┌─────────────────────┐     ┌─────────────────────┐     ┌──────────────────┐
│ PolicyEvaluator │────►│ PreconditionAwaiter │────►│       Batcher       │────►│ TransferExecutor │
└─────────────────┘     └─────────────────────┘     └─────────────────────┘     └──────────────────┘
                                                             ▲                          ▲
                                                             │                          │
                                                    CancellableQueue          CancellableQueue
                                                             │                          │
                                                             └──────── CancelSweeper ───┘

§Stages

Stage	Purpose
PolicyEvaluator	Filters blocks based on configured policies (frequency, presence, etc.)
PreconditionAwaiter	Waits for forward pass completion before proceeding
Batcher	Groups containers into batches based on total block count
TransferExecutor	Upgrades blocks and executes the actual transfer

§Container Data Model

The fundamental unit flowing through the pipeline is an OffloadContainer:

struct OffloadContainer<T: BlockMetadata> {
    /// The blocks to offload
    blocks: Vec<SourceBlock<T>>,
    /// Precondition event (forward pass completion)
    precondition: Option<EventHandle>,
    /// Cancellation token
    cancel_token: CancellationToken,
}

Containers are grouped into batches for efficient transfer:

struct OffloadBatch<T: BlockMetadata> {
    /// Multiple containers, each independently cancellable
    containers: Vec<OffloadContainer<T>>,
}

§P1: Container is the Unit of Cancellation

Individual blocks within a container are not independently cancellable. When a container is cancelled, all its blocks are cancelled together.

§P2: Token Travels with Container

Each container carries its own CancellationToken, cloned from the TransferHandle at enqueue time. The token travels with the container through all pipeline stages until upgrade.

§P3: Upgrade is the Commitment Boundary

The upgrade step (Weak → Strong) is the point of no return:

Before upgrade: Containers can be cancelled via sweep or token check
After upgrade: We own the blocks; cancellation no longer applies

§P4: Sweep Before Upgrade

The last cancellation check occurs immediately before upgrade. The TransferExecutor calls batch.sweep_cancelled() to remove cancelled containers before committing.

§P5: Flat Map After Upgrade

After upgrade, all blocks from all containers are consolidated into a single Vec<ImmutableBlock<T>> for efficient batch transfer. Per-container identity is lost at this point.

§P6: PreconditionAwaiter Uses Select

The precondition awaiter can be cancelled via select! on both the precondition event and the cancellation token. If cancelled while waiting, the container is dropped immediately.

§Configuration

Pipeline behavior is controlled via PipelineConfig:

Option	Default	Description
`batch_config.max_batch_size`	64	Maximum blocks per batch
`batch_config.min_batch_size`	8	Minimum blocks before flush
`batch_config.flush_interval`	10ms	Time before flushing partial batch
`policy_timeout`	100ms	Timeout for policy evaluation
`sweep_interval`	10ms	Interval for cancel sweeper
`max_concurrent_transfers`	1	Concurrent transfer batches

§Usage

§Enqueueing Blocks

let handle = pipeline.enqueue(source_blocks, precondition_event);

// Track progress
println!("Status: {:?}", handle.status());

// Wait for completion
let result = handle.wait().await?;

§Cancelling a Transfer

// Request cancellation and wait for confirmation
handle.cancel().await;
// All blocks are now released

offload-developer.md - Implementation details and extension rules

Offload Engine for asynchronous block transfers between storage tiers.

The offload engine provides a policy-based, cancellable pipeline for moving blocks from higher-performance tiers (G1/G2) to lower-cost tiers (G3/G4).

§Architecture

┌─────────────────────────────────────────────────────────────────┐
│                        OffloadEngine                            │
│                                                                 │
│  ┌───────────────┐    ┌───────────────┐    ┌───────────────┐    │
│  │G1→G2 Pipeline │────│ G2→G3 Pipeline│    │ G2→G4 Pipeline│    │
│  └───────────────┘    └───────────────┘    └───────────────┘    │
│         │                     │                     │           │
│         └─────────auto_chain──┘                     │           │
│                                                                 │
└─────────────────────────────────────────────────────────────────┘

Pipeline stages:
┌─────────────┐    ┌────────────────┐    ┌──────────────────┐
│   Policy    │───▶│     Batch      │───▶│    Transfer      │
│  Evaluator  │    │   Collector    │    │    Executor      │
└─────────────┘    └────────────────┘    └──────────────────┘
      │                   │                      │
      ▼                   ▼                      ▼
  cancel check       cancel check          wait for in-flight

§Features

Policy-based filtering: Blocks pass through configurable policies (presence checks, LFU thresholds) before transfer
Batched transfers: Blocks are accumulated into batches for efficient bulk transfers
Cancellation: Clean cancellation with confirmation that all blocks are released and no outstanding operations remain
Pipeline chaining: G1→G2 completions can automatically feed G2→G3

See also: Developer Guide for implementation details and extension rules.

§Example

use kvbm::v2::distributed::offload::{
    OffloadEngine, PipelineBuilder, PresenceFilter, PresenceAndLFUFilter,
};

// Build engine with pipelines
let engine = OffloadEngine::builder(leader.clone())
    .with_registry(registry.clone())
    .with_g2_manager(g2_manager.clone())
    .with_g3_manager(g3_manager.clone())
    .with_g2_to_g3_pipeline(
        PipelineBuilder::<G2, G3>::new()
            .policy(Arc::new(PresenceAndLFUFilter::with_default_threshold(registry.clone())))
            .batch_size(64)
            .build()
    )
    .build()?;

// Enqueue blocks for offload
let handle = engine.enqueue_g2_to_g3(blocks)?;

// Wait for completion or cancel
tokio::select! {
    result = handle.wait() => {
        println!("Completed: {:?}", result?.completed_blocks);
    }
    _ = shutdown_signal => {
        handle.cancel().wait().await;
        println!("Cancelled");
    }
}

Structs§

AllOfPolicy: Composite policy that requires ALL sub-policies to pass (AND logic).
AnyOfPolicy: Composite policy that requires ANY sub-policy to pass (OR logic).
BatchConfig: Configuration for batch collection.
CancelConfirmation: Future that resolves when cancellation is fully confirmed.
CancellableQueue: A lock-free queue that supports active cancellation via sweeping.
CancellationToken: Token for requesting and tracking cancellation.
EvalContext: Context provided to policies for block evaluation.
ExternalBlock: External block reference with sequence hash for registration.
ObjectLockPresenceFilter: G2→G4 filter with distributed locking: check meta, acquire lock, track acquired locks.
ObjectPipeline: A running pipeline instance for object storage destinations.
ObjectPipelineBuilder: Builder for object pipeline configuration.
ObjectPipelineConfig: Configuration for an object storage pipeline.
ObjectPresenceFilter: G2→G4 filter: async presence check for object storage destinations.
OffloadEngine: Central coordinator for offload pipelines.
OffloadEngineBuilder: Builder for OffloadEngine.
PassAllPolicy: A pass-all policy (no filtering).
PendingGuard: RAII guard that removes a sequence hash from the pending set on drop.
PendingTracker: Tracks sequence hashes that are currently pending transfer.
Pipeline: A running pipeline instance.
PipelineBuilder: Builder for pipeline configuration.
PipelineConfig: Configuration for a pipeline.
PipelineFailure: Terminal executor failure observed while waiting for causal settlement.
PresenceAndLFUFilter: G2→G3 filter: presence check + LFU count threshold.
PresenceFilter: G1→G2 filter: skip blocks already present in destination tier.
ResolvedBatch: A batch of resolved blocks ready for transfer.
ResolvedBlock: A resolved block ready for transfer execution.
S3PresenceChecker: S3/Object storage presence checker.
SettlementTarget: Expected completed transfer batches for each pipeline lane.
SettlementToken: Opaque checkpoint captured before an external completion source is fired.
TimingTrace: Timing trace for tracking block progression through pipeline stages.
TransferHandle: Handle for tracking and controlling an offload transfer.
TransferId: Unique identifier for a transfer operation.
TransferProgressCounts: Monotonic block counts for a transfer.
TransferProgressCursor: Per-consumer cursor for incrementally reading transfer progress.
TransferProgressDelta: Blocks appended since a TransferProgressCursor was last consumed.
TransferResult: Result of a completed transfer.

Enums§

CancelState: State of a cancellation request.
PipelineFailureKind: Stable failure categories for pipeline settlement.
PipelineLane: An offload pipeline whose executor can participate in causal settlement.
SettlementError: Error returned when a causal settlement boundary cannot be established.
SourceBlock: Represents a single block source for offloading.
SourceBlocks: Collection of source blocks for batch operations.
TransferStatus: Status of a transfer operation.

Traits§

OffloadPolicy: Trait for offload policies that filter blocks.
PresenceChecker: Async presence checker for object storage or other external destinations.

Functions§

async_batch_result: Create an async batch policy result (boxes the future).
async_result: Create an async policy result (boxes the future).
create_policy_from_config: Create a composite policy from tier configuration.
sync_batch_result: Create a synchronous batch policy result (zero allocation).
sync_result: Create a synchronous policy result (zero allocation).
upgrade_batch: Upgrade a batch of queued blocks by resolving weak references.

Type Aliases§

BoxFuture: Boxed future type for async policy evaluation.
PolicyBatchFuture: Future type for batch policy evaluation.
PolicyFuture: Future type for single-block policy evaluation.

Module offload

Module offload Copy item path

§Offload Module

§Overview

§Pipeline Architecture

§Stages

§Container Data Model

§P1: Container is the Unit of Cancellation

§P2: Token Travels with Container

§P3: Upgrade is the Commitment Boundary

§P4: Sweep Before Upgrade

§P5: Flat Map After Upgrade

§P6: PreconditionAwaiter Uses Select

§Configuration

§Usage

§Enqueueing Blocks

§Cancelling a Transfer

§Related Documentation

§Architecture

§Features

§Example

Structs§

Enums§

Traits§

Functions§

Type Aliases§

Module offload