Crate metrale_storage

Crate metrale_storage 

Source
Expand description

2026-09-25: metrale-storage: the KV, SSM-snapshot, expert and weight storage tiers (NVMe files, host RAM, RDMA peers), the peers that serve them, and their CUDA helpers.

Owner: storage. Invariants:

  • The CUDA driver bindings (cuda_min, cuda_module, cuda_graph) exist only with the cuda feature, so every module that calls the driver is gated on it too.
  • Without the cuda feature, stubs supplies HighSpeedSwap, with_local, local_installed and install_local under the same names.

Re-exports§

pub use cuda_module::CudaEvent;
pub use cuda_module::CudaModule;
pub use cuda_module::launch_kernel;
pub use model_dims::ModelDims;
pub use backend::PosixBackend;
pub use backend::ReadRequest;
pub use backend::StorageBackend;
pub use backend::IoUringBackend;
pub use config::HighSpeedSwapConfig;
pub use eviction::EvictionPolicy;
pub use expert::ExpertKey;
pub use expert::ExpertLayout;
pub use expert::ExpertRecordHeader;
pub use expert::ExpertRecordId;
pub use expert::ExpertRecordSpec;
pub use expert::Proj;
pub use expert::ProjBytes;
pub use expert_arena::ExpertArena;
pub use expert_pack::ExpertFileReader;
pub use expert_pack::ExpertFileWriter;
pub use expert_pack::ExpertIndex;
pub use expert_pack::ProjData;
pub use expert_pack::ProjView;
pub use expert_pack::pack_record;
pub use expert_pack::unpack_record;
pub use expert_tier::ArenaSlot;
pub use expert_tier::ExpertResidency;
pub use expert_tier::ExpertTier;
pub use expert_tier::PosixTier;
pub use expert_tier::TierKind;
pub use expert_tier::UmaArenaTier;
pub use expert_tier::open_tier;
pub use expert_tier_rdma::RdmaTier;
pub use high_speed_swap::HighSpeedSwap;
pub use high_speed_swap::install_local;
pub use high_speed_swap::local_installed;
pub use high_speed_swap::with_local;
pub use ngram_cache::NgramRowCache;
pub use rdma_snapshot::RdmaSnapshotArena;
pub use predictor::Predictor;
pub use predictor::PredictorDims;
pub use probe::Backend;
pub use probe::ProbeConfig;
pub use probe::ProbeResult;
pub use probe::run_probe;
pub use projection::PredictorShape;
pub use projection::build_projection;
pub use tiled_attention::TiledAttention;
pub use tiled_attention::TiledAttentionDims;
pub use weight_peer::WeightManifest;
pub use weight_peer::WeightTensorRecord;

Modules§

attention_ref
2026-09-25: Host reference for the tiled online-softmax decode attention of kernels/paged_decode_attn_tiled.cu and kernels/attention_finalize.cu.
backend
2026-09-25: The StorageBackend trait, which moves KV groups and blocks between a storage tier and device memory, and the two file backends.
bench
2026-09-26: Read micro-benchmarks for the storage backends, in MiB/s, from a file opened with O_DIRECT (open_test_file, which also advises the kernel to drop the file’s cached pages) into device memory. bench_cufile reads with cuFile straight into the device buffer; bench_io_uring and bench_posix read into a pinned host buffer and copy each read to the device with cuMemcpyHtoDAsync. Sequential reads cycle over the first SEQ_ITERS reads of the file; random reads land on RAND_IO_BYTES-aligned offsets.
cache_peer
2026-09-25: The cache peer: a process that lends RAM to KV and snapshot clients over one-sided RDMA. A connection gets either its own arena (RAW mode), which the client allocates in, or the shared paging arena of its (kind, blob_bytes), backed by a swap file (registry).
cascade_backend
2026-09-25: CascadeBackend: a write-back cache of groups in pinned host memory (T1) in front of another StorageBackend. Writes land in T1; when T1 is full, the least recently used group is written down to the backing. Reads copy T1 hits to the device and forward misses to the backing. No tier transforms bytes.
cascade_policy
2026-09-25: Slot bookkeeping for the cascade’s first tier (T1), the write-back group cache of CascadeBackend: placement, LRU eviction and the hit/miss split of a read. No I/O and no CUDA.
config
2026-09-25: HighSpeedSwapConfig, the settings of the high-speed swap, and their range checks. HighSpeedSwap::new_on_stream runs validate_and_prepare before it builds anything.
cuda_graph
2026-09-25: CUDA graph capture of a stream’s work, replayed with one cuGraphLaunch.
cuda_min
2026-09-25: Minimal CUDA driver-API FFI for metrale-storage: context and stream, device and pinned host buffers, async copies, stream sync.
cuda_module
2026-09-25: CUDA driver-API module loading, kernel launch and events.
eviction
2026-09-25: Eviction order for the high-speed-swap scratch slots: which resident slot a newly needed block overwrites. rank orders the unpinned slots by (recorded predictor score, last access epoch), ascending.
expert
2026-09-25: Expert identity and on-disk record geometry for MoE expert streaming. A record holds one (moe_layer, expert)’s gate, up and down NVFP4 projections behind a fixed header; each MoE layer’s file holds num_experts records at one stride, which the whole model shares.
expert_arena
2026-09-25: ExpertArena: one pinned host allocation (cuMemAllocHost) divided into num_slabs slabs of slots_per_slab expert-record slots. new requires the allocation’s device address to equal its host address, so a record written into a slot is read by device code at slot_dev_va without a copy. The expert tiers decide what fills a slot.
expert_pack
2026-09-25: The expert-store format. pack_record and unpack_record build and parse one record in memory, with no I/O. ExpertIndex is the store’s manifest.json, from which a reader rebuilds the record spec and layout. The layer files are written and read in expert_pack_fs.rs.
expert_peer
2026-09-25: The expert peer, a process that serves an expert store’s records to RdmaTier clients, and its wire format (little-endian):
expert_tier
2026-09-25: ExpertTier: fetch one expert record into a slot and return the device addresses of its six sub-buffers, with its scalars. The tiers are PosixTier (read, then copy into a device buffer), UmaArenaTier (read straight into a pinned arena slot) and RdmaTier (expert_tier_rdma.rs, from a peer). It is not a StorageBackend, whose reads take KV group keys.
expert_tier_rdma
2026-09-25: RdmaTier: fetches expert records from an expert peer (expert_peer) straight into the pinned arena and returns residency addresses inside it, as UmaArenaTier does from local files. Transports: TCP, where the peer reads each record and sends it, and one-sided RDMA READ from the peer’s registered store (cfg(metrale_rdma_verbs) only). Verbs rail 0 comes from METRALE_EXPERT_RDMA_DEV / METRALE_EXPERT_RDMA_GID; METRALE_EXPERT_DUAL_RAIL=1 adds rail 1 from METRALE_EXPERT_RAIL2_DEV / METRALE_EXPERT_RAIL2_GID.
group
2026-09-25: Group addressing for the high-speed-swap layer. A group is the K or V stripe of one (layer, block, kv_head): block_size × head_dim × elem_bytes bytes, padded to a multiple of fs_block_size.
high_speed_swap
2026-09-25: HighSpeedSwap, the high-speed-swap orchestrator. It writes KV blocks to per-layer files (offload_block*) and streams a sequence’s blocks back through a scratch pool into tiled attention (attend_layer*), using the predictor, the file backend and the eviction policy. It also keeps the disk block ids and the per-thread slot (install_local, with_local) through which model code reaches it.
kv_paging
2026-09-25: KV blocks as a paging kind. KvPagingBackend keeps whole KV blocks in a peer’s paging arena (handshake kind PagingKind::KV); the alternative is the one-sided RdmaKvBackend, whose client addresses a fixed arena at base + group_id × group_stride.
layout
2026-09-25: On-disk layout of the high-speed-swap KV tier: one pre-allocated file per layer, layer_{:05}.kv under HighSpeedSwapConfig::dir (--high-speed-swap-dir).
model_dims
2026-09-25: ModelDims, the model geometry the high-speed-swap tier is built for.
ngram_cache
2026-09-25: NVMe-backed row cache for n-gram embedding tables. Rows fault from the checkpoint file into a pinned arena, and the gather kernel reads the arena by slot.
predictor
2026-09-25: GPU side of the high-speed-swap block predictor. Loads the q_lowrank_project, kv_lowrank_project and predictor_score PTX from build.rs and owns the projection P and the per-token low-rank keys A_g.
predictor_ref
2026-09-25: Host reference for the predictor kernels, with the kernels’ tensor layouts; tests/predictor_parity.rs compares the GPU results against it.
probe
2026-09-26: Storage capability probe (run_probe): loads cuFile and opens its driver when it can, checks for the nvidia-fs module, runs the bench read benchmarks (cuFile’s only when its driver opened), and recommends a Backend from the results (decide). A benchmark that fails is logged and left out of the result.
projection
2026-09-25: The high-speed-swap predictor’s random Gaussian projection P, built on the host from a seed; Predictor uploads it once.
rdma_snapshot
2026-09-25: RdmaSnapshotArena, the synchronous RDMA transport of the SSM snapshot spill tier, addressed by byte offset or, in paging mode, by key.
scratch_pool
2026-09-25: Device scratch pool of the high-speed-swap tier: slots that hold one block’s K and V for every KV head, filled from the storage backend.
snapshot_swap
2026-09-25: Paging core of the SSM snapshot tier: the peer’s TCP control protocol and paging loops, the client side of that protocol (wire), and MmapSlotArena (mmap_arena). Re-exports the generic tier types it builds on.
tier
2026-09-25: The generic tiered cache: a bounded hot arena over a cold record store, written without CUDA or verbs calls.
tiled_attention
2026-09-25: Host side of the tiled decode attention (kernels/paged_decode_attn_tiled.cu, kernels/attention_finalize.cu in this crate). A decode step is begin_step, which sets the FP32 running state (m, l, o) to (-inf, 0, 0); one step_tile per tile of KV blocks; and finalize, which writes o / l (0 where l is 0) as BF16.
weight_peer
2026-09-26: The weight peer: a daemon that holds models’ safetensors shards mapped in its RAM and lets a client read any tensor with one-sided RDMA. The client is RdmaWeightLoader in model-weights (weight_tier_rdma.rs).

Functions§

rdma_verbs_enabled
2026-09-25: true when build.rs emitted metrale_rdma_verbs for this crate. The verbs shim is built by metrale-gpu-sys, and rustc-cfg does not cross crates, so build.rs re-emits the cfg from that crate’s links metadata. rdma_verbs_probe_tests checks it matches metrale_gpu_sys::verbs_enabled().