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 thecudafeature, so every module that calls the driver is gated on it too. - Without the
cudafeature,stubssuppliesHighSpeedSwap,with_local,local_installedandinstall_localunder 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.cuandkernels/attention_finalize.cu. - backend
- 2026-09-25: The
StorageBackendtrait, 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_cufilereads with cuFile straight into the device buffer;bench_io_uringandbench_posixread into a pinned host buffer and copy each read to the device withcuMemcpyHtoDAsync. Sequential reads cycle over the firstSEQ_ITERSreads of the file; random reads land onRAND_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 anotherStorageBackend. 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_streamrunsvalidate_and_preparebefore 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.
rankorders 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_expertsrecords at one stride, which the whole model shares. - expert_
arena - 2026-09-25:
ExpertArena: one pinned host allocation (cuMemAllocHost) divided intonum_slabsslabs ofslots_per_slabexpert-record slots.newrequires the allocation’s device address to equal its host address, so a record written into a slot is read by device code atslot_dev_vawithout a copy. The expert tiers decide what fills a slot. - expert_
pack - 2026-09-25: The expert-store format.
pack_recordandunpack_recordbuild and parse one record in memory, with no I/O.ExpertIndexis the store’smanifest.json, from which a reader rebuilds the record spec and layout. The layer files are written and read inexpert_pack_fs.rs. - expert_
peer - 2026-09-25: The expert peer, a process that serves an expert store’s
records to
RdmaTierclients, 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 arePosixTier(read, then copy into a device buffer),UmaArenaTier(read straight into a pinned arena slot) andRdmaTier(expert_tier_rdma.rs, from a peer). It is not aStorageBackend, 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, asUmaArenaTierdoes 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 fromMETRALE_EXPERT_RDMA_DEV/METRALE_EXPERT_RDMA_GID;METRALE_EXPERT_DUAL_RAIL=1adds rail 1 fromMETRALE_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_bytesbytes, padded to a multiple offs_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.
KvPagingBackendkeeps whole KV blocks in a peer’s paging arena (handshake kindPagingKind::KV); the alternative is the one-sidedRdmaKvBackend, whose client addresses a fixed arena atbase + 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}.kvunderHighSpeedSwapConfig::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_projectandpredictor_scorePTX from build.rs and owns the projectionPand the per-token low-rank keysA_g. - predictor_
ref - 2026-09-25: Host reference for the predictor kernels, with the kernels’ tensor
layouts;
tests/predictor_parity.rscompares 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 thebenchread benchmarks (cuFile’s only when its driver opened), and recommends aBackendfrom 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;Predictoruploads 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), andMmapSlotArena(mmap_arena). Re-exports the generictiertypes 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.cuin this crate). A decode step isbegin_step, which sets the FP32 running state (m, l, o) to (-inf, 0, 0); onestep_tileper tile of KV blocks; andfinalize, which writeso / l(0 wherelis 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
RdmaWeightLoaderin model-weights (weight_tier_rdma.rs).
Functions§
- rdma_
verbs_ enabled - 2026-09-25:
truewhen build.rs emittedmetrale_rdma_verbsfor this crate. The verbs shim is built by metrale-gpu-sys, andrustc-cfgdoes not cross crates, so build.rs re-emits the cfg from that crate’slinksmetadata.rdma_verbs_probe_testschecks it matchesmetrale_gpu_sys::verbs_enabled().