Loaderx

Zrecord is a rebuildable, typed, ordered record sequence built from authoritative source data and scripts. A creator consumes records in append order, close publishes one immutable container. An opened Store is bound to an explicit placement view before integer indexing, natural iteration, slices, or indexed gather. Dense integer indexing unwraps the leading batch axis; integer indexing of Ragged returns a one-record RaggedBatch. To change content or order, rebuild it at a new path.

Zrecord is the typed on-disk container; Loaderx is the sampler and data loader that consumes Zrecord streams. They currently ship together while both layers mature, but their public responsibilities remain separate.

pip install loaderx

Wheels are published for CPython 3.10+ on glibc Linux (x86-64 and ARM64), Apple Silicon macOS, and Windows AMD64. Free-threaded CPython 3.14 is also supported.

Design Philosophy

loaderx is built around several core principles:

1. A pragmatic approach that prioritizes minimal memory overhead and minimal dependencies.
2. A strong focus on single-machine training workflows.
3. NumPy-native typed records with explicit schemas.
4. An immortal (endless) step-based data loader, rather than the traditional epoch-based design—better aligned with modern ML training practices.
5. Dense and Ragged are explicit physical layouts. Applications may regularize records into one fixed shape or preserve their variable geometry. Loaderx stores and delivers that choice without converting between layouts.
6. Logical IDs are stable sequence positions. Append order defines 0..N-1; published containers are immutable.
7. Single-controller SPMD placement is explicit. One sampler draws a global batch, and Loader splits it across caller-provided placement views without constructing and redistributing a global tensor. Process-per-device DDP is not part of the Loader contract.

Usage

Quick Start

import numpy as np

from loaderx import Allocator, Loader
from loaderx.zrecord import DenseStore, DenseWriter
from loaderx.sampler import Sampler

data = np.load('data.npy', mmap_mode='r')
label = np.load('label.npy', mmap_mode='r')
with DenseWriter('train_data', data.dtype, data.shape[1:], codec='zstd') as ds:
    ds.append(data)
with DenseWriter('train_label', label.dtype, label.shape[1:], codec='zstd') as ds:
    ds.append(label)

with DenseStore('train_data') as data_store, \
        DenseStore('train_label') as label_store, \
        Allocator('cpu') as allocator:
    data_view = data_store.bind(allocator)
    label_view = label_store.bind(allocator)
    sampler = Sampler(len(data_view), 256, Sampler.Mode.CYCLIC, seed=42)
    loader = Loader({'data': [data_view], 'label': [label_view]}, sampler,
                    transform=lambda shards: shards)
    try:
        for i, shards in enumerate(loader):
            if i >= 256:
                break
            shard = shards[0]

        print(shard['data'].shape)
        print(shard['label'].shape)
    finally:
        loader.close()

A Loader step is a list of shard dicts [{name: values}, ...]. A Dense value is a DLTensor; a Ragged value contains values, cu_seqlens, and shapes. Every named stream at one view position is gathered with the same index slice, so records stay aligned. Singleton view lists deliberately produce a one-element output list. The transform callback receives the complete list as the collate step and its return value is passed to the model unchanged.

Records

The Dense or Ragged contract follows the geometry an application chooses to persist. Fixed-shape records have one schema-known stride and stack directly; records whose shapes vary require explicit payload boundaries and per-record shapes. Resizing, padding, truncating, or otherwise converting between these layouts is application policy, not loader behavior.

Dense stores one fixed-shape array per record. An integer-index read or one step of iteration returns the record itself; a multi-record read returns one stacked DLTensor:

import numpy as np
from loaderx import Allocator, dtypes
from loaderx.zrecord import DenseStore, DenseWriter

data = np.arange(64, dtype=dtypes.float32).reshape(8, 2, 4)
with DenseWriter('data', data.dtype, data.shape[1:], codec='zstd') as ds:
    ds.append(data)
with DenseStore('data') as store, Allocator('cpu') as allocator:
    view = store.bind(allocator)
    record = view[0]                  # DLTensor (2, 4), no leading batch axis
    batch = view[0, 5, 2]             # DLTensor (3, 2, 4), requested order retained
    for record in view:                # each record is a DLTensor (2, 4)
        consume(record)

A Store is an ordered sequence of records, not an ndarray. Open Stores are unbound and cannot be indexed directly; DenseStore() returns a DenseStore, RaggedStore() returns a RaggedStore, and store.bind(allocator) returns the placement DenseView or RaggedView that owns the index API. DenseWriter() and RaggedWriter() return DenseWriter and RaggedWriter. A view expression such as view[0, 5, 2] selects sequence positions 0, 5 and 2, never view[0][5][2]. An integer index selects one unwrapped Dense record, while any collection of indices selects a batch. Reads return loaderx.zrecord.DLTensor, a DLPack-only interchange object. It has no to_numpy() or __array__; consume a CPU value with np.from_dlpack(record) or any value with a compatible framework such as torch.from_dlpack(record). Indices must be in 0..len(view)-1.

Ragged stores variable-shape arrays with one shared dtype and ndim >= 1. It does not support zero-dimensional records. Scalar-valued records have a fixed 0-D item shape and therefore belong in Dense with item_shape=(). Reads return a loaderx.zrecord.RaggedBatch: contiguous 1-D values in the store dtype, contiguous int32 cu_seqlens record offsets, and contiguous uint64 shapes with shape (records, ndim). The three fields use one contract and one DLPack device on CPU and GPU. An empty read batch has cu_seqlens == [0]; stored records themselves must be nonempty. Integer indexing returns a single-record RaggedBatch (shapes.shape[0] == 1). CPU batches provide to_tensors() for explicit per-record reconstruction; device batches must be consumed in packed form or converted by the device framework. Padding remains an application policy:

from loaderx import Allocator, dtypes
from loaderx.zrecord import RaggedBatch, RaggedStore, RaggedWriter

seqs = [np.arange(L, dtype=dtypes.int32) for L in (3, 1, 4, 1, 5)]
with RaggedWriter('tokens', dtypes.int32, ndim=1, codec='zstd') as rs:
    rs.append(RaggedBatch.from_records(seqs))
with RaggedStore('tokens') as store, Allocator('cpu') as allocator:
    view = store.bind(allocator)
    batch = view[0, 2, 4]             # RaggedBatch: values + cu_seqlens + shapes
    values = batch.values             # packed data, ready for a ragged-aware consumer
    tensors = batch.to_tensors()      # explicit per-record DLTensor views
    for tensor in tensors:            # each is one shape-restored DLTensor
        consume(tensor)

Logical ID is the stable sequence position. Each append preserves input order, and successive calls extend the sequence. After close, the sequence is immutable: there is no delete, update, compaction, or reopen-for-append.

Creating containers

Constructing DenseWriter and RaggedWriter creates a new container. The required codec is "raw", "zstd", or "zstd_dict". Schema, shape, and dtype are explicit and are never inferred from the input. Both append methods synchronously resolve the input metadata and CPU address; DLPack is the normal input protocol:

• DenseWriter.append(batch) requires one C-contiguous CPU or ROCm tensor with the store dtype and shape (batch_size, *item_shape).
• RaggedWriter.append(batch) requires RaggedBatch(values, cu_seqlens, shapes). values is a contiguous 1-D tensor in the store dtype, cu_seqlens is an aligned contiguous int32[batch_size + 1] tensor, and shapes is an aligned contiguous uint64[batch_size, ndim] tensor. All three fields must report the same exact DLPack device.

Ragged offsets start at zero, strictly increase, span all packed values, and agree with each record's shape. RaggedBatch.from_records(records) explicitly adapts a finite iterable of NumPy arrays; Arrow or kernel producers can construct the three packed fields directly. Pass the complete input to one append; the native Zig writer bounds scheduling and scratch memory internally, so callers do not need to slice the source into memory-control batches.

A creator consumes a finite build stream and publishes an immutable sequence; placement views do not tail an active writer. Append is a synchronous boundary and returns after the complete input has been persisted. The input Allocation must remain alive, stable, and exclusively used by that call; append does not provide an asynchronous or concurrent cross-framework access contract. Framework-owned ROCm producers use the DLPack stream handoff, and Loaderx synchronizes that handoff before reading the allocation. If a framework modifies a Loaderx allocation through an imported view but the original Loaderx DLTensor is passed back to append, synchronize the framework first: the original producer cannot describe work submitted by a later consumer.

import numpy as np
from loaderx import dtypes
from loaderx.zrecord import DenseStore, DenseWriter, RaggedBatch, RaggedStore, RaggedWriter

with DenseWriter('mnist/x', dtype=dtypes.uint8, item_shape=(28, 28),
                  codec='zstd', data_shards=4) as ds:
    ds.append(images)

with RaggedWriter('tokens', dtype=dtypes.int32, ndim=1, codec='zstd') as tok:
    tok.append(RaggedBatch.from_records(sequences))

with DenseStore('mnist/x') as store, Allocator('cpu') as allocator:
    view = store.bind(allocator)
    first_four = view[:4]              # the Store is read-only; the view materializes

data_shards controls write parallelism, accepts 1..255, and defaults to four. Opening a Store discovers it automatically.

Append errors are reported by the current call. A writer cannot be read and a read-only Store cannot be appended to. Writer close() publishes the container; Store close() releases it and invalidates its placement views. Dense, Ragged, and both loaders support with.

Schemas use the canonical catalog in loaderx.dtypes. It includes native-endian bool, integer, floating-point, complex, bfloat16, and supported float8 dtypes. Use dtypes.float32, dtypes.bfloat16, or another catalog entry when choosing a schema; structured, subarray, object, metadata-bearing, non-native-endian, and zero-itemsize dtypes are not supported. Append does not use NumPy's general array coercion. The sole non-DLPack fallback is a read-only CPU buffer that can be viewed without copying, such as a read-only NumPy mmap whose DLPack export is unavailable. Writable buffer-only objects and external DLPack devices other than CPU or ROCm are rejected. Build a packed RaggedBatch directly or convert record objects explicitly with RaggedBatch.from_records. Raw bytes and encoded files can be represented as dtypes.uint8 records.

Codec notes

"zstd" compresses each record independently with plain zstd (level 3). Use it for general-purpose compression; it is fast and must be selected explicitly.

"zstd_dict" uses a shared dictionary trained from representative settled data. It is useful for large corpora of small, similar records such as token sequences and image tiles. Plain "zstd" or "raw" is usually better for large records and small corpora.

Train the dictionary manually with train_dict on a representative Zrecord sample, then pass its bytes to dict_bytes. zstd_dict never trains automatically:

from loaderx import Allocator, dtypes
from loaderx.zrecord import DenseStore, DenseWriter, RaggedBatch, RaggedStore, RaggedWriter, train_dict

# train once on a representative settled Zrecord sample
with DenseStore('dict-sample') as sample_store, Allocator('cpu') as allocator:
    d = train_dict(sample_store.bind(allocator))

# then any new store can install it and append explicitly
with RaggedWriter('tokens', dtypes.int32, ndim=1,
                   codec='zstd_dict', dict_bytes=d) as ds:
    ds.append(RaggedBatch.from_records(token_generator))
with DenseWriter('data', data.dtype, data.shape[1:],
                       codec='zstd_dict', dict_bytes=d) as ds:
    ds.append(data)

Changing an existing store's codec requires writing a new store:

from loaderx import Allocator
from loaderx.zrecord import train_dict
from loaderx.zrecord import DenseStore, DenseWriter

with DenseStore("src") as source_store, Allocator("cpu") as allocator:
    source = source_store.bind(allocator)
    dictionary = train_dict(source)
    with DenseWriter(
        "dst", dtype=source.dtype, item_shape=source.item_shape,
        codec="zstd_dict", dict_bytes=dictionary,
    ) as destination:
        destination.append(source[:])             # DLTensor -> DLTensor

For both geometries, a batch read is already the exact append input. Dense uses DLTensor; Ragged uses the packed RaggedBatch. Integer indexing unwraps one record and therefore omits Dense's leading batch axis. Create the destination with the source schema. The destination path must be new.

Important: Train the dictionary from settled authoritative input before building the container. Training it before preprocessing is complete wastes compression and does not describe the final records.

Offline Hugging Face conversion

The optional converter accepts one Arrow-backed Hugging Face Dataset and writes it directly into typed Zrecord streams. Handle splits in the caller. Load and preprocess the dataset with Hugging Face, then pass the materialized Arrow dataset to Loaderx:

pip install 'loaderx[converter]'

from datasets import Dataset, Features, Sequence, Value
from loaderx.converter import convert

dataset = Dataset.from_dict(
    {"tokens": [[1, 2], [3], [4, 5, 6]]},
    features=Features({"tokens": Sequence(Value("int32"))}),
)
convert(dataset, "tokens", codec="zstd")

Load the dataset directly through Hugging Face. For a mirror or private Hub, configure its endpoint using the Hugging Face configuration or environment before calling load_dataset:

import datasets

datasets.config.HF_ENDPOINT = "https://hf-mirror.com"
dataset = datasets.load_dataset("ylecun/mnist", split="train")
convert(dataset, "mnist-train", codec="zstd")

Pass name=..., revision=..., and token="hf_..." directly to datasets.load_dataset when needed.

convert consumes standard Arrow physical batches, not decoded Python records. Primitive and fixed-size-list columns become Dense stores. List<primitive> becomes one-dimensional Ragged, and a multidimensional Ragged column uses an explicit Arrow struct with values: List<primitive> and shape: FixedSizeList<uint64, ndim>. Binary, null, empty, nested dynamic, decoded Image and Audio columns fail explicitly. Materialize preprocessing first with Dataset.map; DatasetDict splits and IterableDataset are intentionally outside the converter contract.

The result publishes one dataset's aligned streams under one root:

dataset/
  tokens/
  label/

The output is published only after all columns have been written and their record counts agree. Each column remains an ordinary Zrecord store.

Loaders and multi-stream stores

An opened Zrecord Store supplies one stream; a training sample is usually several named streams (skeleton + label + id, tokens + label, ...). Loader accepts a dict[str, list[view]]: names identify aligned streams and each list position is one explicit placement view. All view lists must be nonempty and equally long, all views must have the same record count, and the sampled global batch must divide evenly by the number of views. There is no persistent wrapper, manifest, directory convention or bundle mutation API.

from loaderx import Allocator, Loader, dtypes
from loaderx.zrecord import DenseStore, DenseWriter, RaggedBatch, RaggedStore, RaggedWriter
from loaderx.sampler import Sampler

root = "xsub/train"
with DenseWriter(root + "/joint", joint.dtype, joint.shape[1:], codec="zstd") as s:
    s.append(joint)
with DenseWriter(root + "/label", label.dtype, label.shape[1:], codec="zstd") as s:
    s.append(label)
with RaggedWriter(root + "/token", dtypes.int32, ndim=1, codec="zstd") as s:
    s.append(RaggedBatch.from_records(seqs))

stores = {
    "joint": DenseStore(root + "/joint"),
    "label": DenseStore(root + "/label"),
    "token": RaggedStore(root + "/token"),
}
with Allocator("cpu") as allocator:
    streams = {
        name: [store.bind(allocator)] for name, store in stores.items()
    }
    streams["joint"][0][0, 5, 2]     # each view has its own index API
    sampler = Sampler(len(streams["joint"][0]), 256, Sampler.Mode.CYCLIC, seed=42)
    loader = Loader(streams, sampler)
    try:
        shards = next(loader)             # [{name: values}], index-aligned
        shard = shards[0]
    finally:
        loader.close()
for store in stores.values():
    store.close()

The Store remains one global logical sequence. A stream names the role that Store plays in a sample; binding a view chooses where batches from that stream are materialized. Multiple devices therefore appear as multiple views of the same Store, not as multiple logical Stores:

streams = {
    "data": [
        data_store.bind(allocator_0),
        data_store.bind(allocator_1),
    ],
    "label": [
        label_store.bind(allocator_0),
        label_store.bind(allocator_1),
    ],
}

The sampler draws one global index array. Loader divides that array into contiguous, zero-copy views and gathers each view into its allocation; each materialized local batch is a shard. It does not divide the Store into fixed ranges. Randomness and cyclic order therefore come entirely from the global sampler. Fixed-shape and variable-shape records are both just stores: a dense shard value is a DLTensor of shape (B / shard_count, *item_shape) and a ragged shard value is a loaderx.zrecord.RaggedBatch: values (packed data), cu_seqlens (record start offsets), and shapes, the same contract on CPU and GPU, only the location differs. No padding is imposed. Densify to a fixed shape however the model needs, or reshape/stack in the transform collate.

Collation is the transform: the complete shard list in, an arbitrary result out. workers=0 runs sampler draw, gather, and transform synchronously in the caller thread. Positive workers starts complete pipeline jobs: each worker serializes sampler draw and gather under one lock, then runs transform outside the lock and puts its result directly into the output queue. One worker preserves sampling order; multiple workers deliver in completion order, so transform must be thread-safe. Allocator-backed device streams support both policies. workers is the throughput control: increase it when gather and thread-safe transform work can overlap. prefetch is the jitter buffer: it spends live batch memory to absorb variation between producer and consumer latency, but does not create concurrency or increase steady-state processing capacity. It is always a positive integer and is never constrained or rewritten by workers. Each worker owns one possible active or blocked batch, while the ready-result queue holds at most prefetch batches. The default workers=4, prefetch=4 retains overlap for blocking or native transforms; lighter workloads should select a smaller policy explicitly.

Choose the execution policy from the workload, not from the allocation device:

# Light GPU handoff: synchronous delivery is usually the lower-overhead path.
loader = Loader(streams, sampler, workers=0, transform=to_torch)

# Expensive thread-safe transform: tune concurrency and ahead-batch memory
# independently from an end-to-end profile.
loader = Loader(streams, sampler, workers=4, prefetch=4,
                transform=tokenize_and_collate)

With positive workers, the output queue capacity is exactly prefetch. A full queue blocks each worker at its current result, so excluding batches retained by the consumer, at most workers active or blocked batches plus prefetch queued batches remain live. Raise workers for measured throughput and prefetch for measured latency jitter; changing either value never changes the other resource. This applies equally to CPU and device allocators. A Python queue copies no tensor bytes: it holds object references, and each DLTensor or RaggedBatch retains its exclusive Allocation. An AMDGPU batch therefore remains in the same device-consumable allocation while queued and is imported by torch.from_dlpack or JAX without an implicit D2H transfer. Only an explicit transform such as .cpu() requests a readback.

def collate(batch):
    return [
        {'input_ids': shard['tokens'], 'label': shard['label']}
        for shard in batch
    ]

sampler = Sampler(len(dense_tokens), 32, Sampler.Mode.CYCLIC, seed=42)
loader = Loader({'tokens': [dense_tokens], 'label': [labelset]}, sampler,
                 transform=collate)
shards = next(loader)
loader.close()

The transform runs once per gathered batch; its return value is passed to the consumer unchanged, and exceptions propagate to the consumer.

Numba can optionally accelerate a CPU-heavy Dense transform. Compile it before timing, then call it from the ordinary transform:

import numba
import numpy as np
from loaderx import dtypes

@numba.njit(nogil=True, parallel=False)
def normalize_u8(x):
    out = np.empty(x.shape, dtype=dtypes.float32)
    for i in range(x.size):
        out.flat[i] = x.flat[i] / 255.0
    return out

normalize_u8(np.zeros((1, 3, 224, 224), dtype=dtypes.uint8))  # compile warmup

def transform(batch):
    for shard in batch:
        shard["image"] = normalize_u8(np.from_dlpack(shard["image"]))
    return batch

Numba is optional. Compile it before measuring loader throughput.

CPU → GPU transfer

Data reaches the GPU through one of two placement paths. With CPU allocation, the model receives a CPU batch that the framework copies to the GPU; the AMDGPU allocator instead materializes Store output in a device-consumable allocation. Loaderx's job in both is to hand the consumer a correct representation — it moves and reconstructs records, it does not run GPU operators.

The CPU path produces CPU batches. Device transfer belongs in transform, where the training framework is already available:

import torch
from loaderx import Allocator

device = "cuda:0"
def to_device(shards):
    return [
        {k: torch.from_dlpack(v).to(device, non_blocking=True)
         for k, v in shard.items()}
        for shard in shards
    ]

with DenseStore(root + "/joint") as joint_store, \
        DenseStore(root + "/label") as label_store, \
        Allocator("cpu") as allocator:
    streams = {
        "joint": [joint_store.bind(allocator)],
        "label": [label_store.bind(allocator)],
    }
    sampler = Sampler(len(streams["joint"][0]), 256, Sampler.Mode.CYCLIC, seed=42)
    loader = Loader(streams, sampler, transform=to_device)
    try:
        for shards in loader:
            model(shards[0])                # already on device
    finally:
        loader.close()

The caller owns both the allocator and Stores. Call loader.close() when the training loop exits; do not close any of them from transform.

A non_blocking=True copy is asynchronous only from pinned memory. Loaderx does not pin memory; use the framework's API in the transform when needed:

def to_device(shards):
    return [
        {k: torch.from_dlpack(v).pin_memory().to(device, non_blocking=True)
         for k, v in shard.items()}
        for shard in shards
    ]

For JAX, use jax.device_put:

import jax
import jax.dlpack

def to_device(shards):
    return [
        {k: jax.device_put(jax.dlpack.from_dlpack(v)) for k, v in shard.items()}
        for shard in shards
    ]

For complete read and compute integrations, see data2latent, which runs the same prepared Dense and Ragged images through Torch or JAX, and Word2Vec, which trains equivalent padded and packed CBOW models on WikiText-103. The focused DLPack round trip demonstrates the opposite direction by appending framework-owned CPU or GPU Dense and Ragged tensors.

Loaderx targets Torch and JAX through the standard DLPack interchange protocol. It does not vendor either framework or ROCm libraries, and TensorFlow is not a maintained target. The workload examples expose CPU and AMDGPU read allocations; the AMDGPU read path and ROCm append path require a compatible Linux ROCm configuration. Protocol compatibility does not imply support for every DLPack framework, device type, or allocation.

Storage → GPU transfer

The CPU path above asks the framework to create a separate device tensor. An AMDGPU allocation instead lets Store produce a device-consumable batch directly. Zrecord computes each output's exact layout and asks an allocator for an aligned region; live batches retain exclusive allocations and matching-size-class released regions return to their allocator pool:

from loaderx import Allocator
from loaderx.zrecord import DenseStore, DenseWriter
import torch

allocator = Allocator("amdgpu", 0)
try:
    with DenseStore(path) as store:
        view = store.bind(allocator)
        batch = view[idxs]                           # lazily sized allocation
        gpu_tensor = torch.from_dlpack(batch)
        consume(gpu_tensor)
finally:
    allocator.close()

JAX consumes the same device object directly with jax.dlpack.from_dlpack(batch). For gather/read output, Loaderx accepts DLPack's stream argument for framework compatibility, but the synchronous gather does not create a producer stream or insert framework-specific events. Every live batch has an exclusive allocation, so a later gather cannot overwrite it.

loaderx.Allocator defines the common Allocator / Allocation contract; an opened Store is unbound and must be explicitly bound before indexing. CPU and AMDGPU use the same bounded pool policy. Each placement view records its own recent request size classes, while the shared allocator retains all view classes in one global idle pool. pool_size limits the allocator's aggregate retained bytes; it is not a per-view budget. Callers select a placement with a readable name, for example Allocator("cpu", 0) or Allocator("amdgpu", 0); DLPack device codes remain an internal protocol detail. Dense computes exact bytes directly; Ragged reads physical lengths, computes aligned shapes/values/cu_seqlens layout, then acquires the final region. Users never provide a batch capacity.

Concrete placement implementations live under loaderx.backend; importing Loaderx does not discover optional GPU runtimes. An implementation supplies writable regions to the common allocator contract and does not add a Store or Loader execution path.

Dense views return DLTensor; Ragged views return three packed DLPack fields in RaggedBatch. DenseStore() and RaggedStore() return only storage objects; store.bind(allocator) creates a DenseView or RaggedView. Each allocation's immutable (DLPack device type, logical device id) is inherited by its DLTensors.

Loaderx imports the process's already-loaded HIP runtime when present. Otherwise it checks LOADERX_HIP_LIBRARY, ROCM_PATH/ROCM_HOME, an installed package-owned runtime from rocm-sdk-core or legacy Torch, /opt/rocm, and finally the system loader. Multiple distinct package-owned runtimes are rejected rather than guessed; set LOADERX_HIP_LIBRARY to resolve the choice. No Torch or JAX module is imported during discovery. HIP logical devices and DRM render nodes are matched by PCI BDF, so visibility remapping does not depend on render-node order.

Raw stores read directly into the selected allocation. Compressed stores use CPU decoding scratch but place the final Dense/Ragged layout in that same allocation; codec and memory location are independent. Use Loader(..., workers=0) for synchronous delivery and minimal live device memory, or positive workers with an explicit ahead-batch bound when overlap is worth the allocator pool growth. See the benchmark section for the end-to-end comparison.

Loaderx supports computation, it does not implement it: allocator-selected delivery preserves the same data layout on CPU or GPU, never the operators themselves. Ragged delivers tightly packed values, element boundaries in cu_seqlens, and exact per-record shapes. Ragged-aware or custom kernels can consume this native layout without first materializing a Python list or padding it into Dense; kernel-specific views and metadata remain the consumer's responsibility. The examples apply this distinction to two representative workloads. data2latent checks equivalent image-to-latent computation, while word2vec checks equivalent embedding training from padded and packed contexts; both expose Torch and JAX through the same Loader over CPU or AMDGPU allocations.

GPU → Storage append

DenseWriter.append and RaggedWriter.append consume framework-owned GPU tensors directly through DLPack. The tensors remain on the GPU until append resolves their backing allocation; there is no .cpu(), .numpy(), or framework-specific branch in Loaderx:

import numpy as np
import torch
from loaderx import Allocator, dtypes
from loaderx.zrecord import DenseStore, DenseWriter, RaggedBatch, RaggedStore, RaggedWriter

device = torch.device("cuda", 0)
producer_stream = torch.cuda.Stream(device=device)
with torch.cuda.stream(producer_stream):
    dense = torch.arange(
        12, dtype=torch.int32, device=device,
    ).reshape(3, 2, 2)
    values = torch.tensor(
        [3, 5, 7, 11, 13, 17, 19], dtype=torch.int32, device=device,
    )
    cu_seqlens = torch.tensor(
        [0, 2, 6, 7], dtype=torch.int32, device=device,
    )
    shapes = torch.tensor(
        [[1, 2], [2, 2], [1, 1]], dtype=torch.uint64, device=device,
    )

# No producer_stream.synchronize(): DLPack carries the stream dependency.
with DenseWriter("dense", dtypes.int32, (2, 2), codec="zstd") as writer:
    writer.append(dense)

with RaggedWriter("ragged", dtypes.int32, ndim=2, codec="zstd") as writer:
    writer.append(RaggedBatch(values, cu_seqlens, shapes))

# Closing the writers publishes the stores; reopening proves persistence.
with DenseStore("dense") as dense_store, RaggedStore("ragged") as ragged_store, \
        Allocator("cpu") as allocator:
    dense_view = dense_store.bind(allocator)
    ragged_view = ragged_store.bind(allocator)
    np.testing.assert_array_equal(
        np.from_dlpack(dense_view[:]), np.arange(12, dtype=dtypes.int32).reshape(3, 2, 2),
    )
    batch = ragged_view[:]
    np.testing.assert_array_equal(
        np.from_dlpack(batch.values),
        np.asarray([3, 5, 7, 11, 13, 17, 19], dtype=dtypes.int32),
    )

Loaderx asks DLPack for a consumer-stream handoff, synchronizes that stream, exports the backing allocation as a dma-buf, and reads a temporary read-only CPU mapping during append. This avoids allocating and filling a complete intermediate CPU tensor. It is not a GPU-direct storage path: Store still reads the mapping and performs normal encoding and file I/O, and the driver may migrate data or move it across an interconnect. The allocation must be exportable and CPU-mappable. The complete Torch and JAX round trip is runnable with:

python examples/dlpack_roundtrip.py --framework torch --memory gpu
python examples/dlpack_roundtrip.py --framework jax --memory gpu

Sampler

Sampler is a borrowed Python buffer over a stateless Cython batch function. Loader receives a sampler object rather than duplicating its batch size, mode, or seed. A run resumes in O(1) with seek(step), without replaying draws or tracking an epoch.

from loaderx.sampler import Sampler

sampler = Sampler(1_000_000, 256, Sampler.Mode.IID, seed=42)
indices = sampler.next()       # borrowed until this sampler's next draw
saved = indices.copy()         # retain across draws only when needed

next() and iteration return a view of one reusable uint64 batch buffer. The contents stay unchanged until the next explicit draw from that Sampler; copy only plans that must outlive it. The loaders consume each view before drawing again. They borrow the sampler and never close it. Sampler exposes only next(), seek(), and iteration; it has no public indices property, close(), or context-manager protocol. Any user object can be injected instead; its next() should return a borrowed one-dimensional contiguous NumPy index array. This is a trusted hot-path contract rather than a normalized protocol, so custom policies can be ordinary NumPy code without Loaderx adapters or lifecycle methods.

1. Sequential traverses records in order and wraps at the end.
2. IID samples uniformly with replacement.
3. Cyclic traverses a fresh permutation without replacement. It does not allocate a dataset-sized permutation. Full batches are returned; a different remainder is omitted on each cycle when the length is not divisible by the batch size.

Benchmarks

Dense and Ragged are measured separately because they expose different contracts. scripts/bench_dense.py measures fixed-shape random gather; scripts/bench_ragged.py measures variable-shape records as compute-ready values + cu_seqlens + shapes; and scripts/bench.py covers machine, sampler, and the end-to-end loader comparison. Every path runs through the public Python binding, including output and metadata allocation. Every Ragged write backend starts from the same Arrow List payload and fixed-size shape column. Zrecord constructs its RaggedBatch from those Arrow buffers inside the timed write; there are no Python-record benchmark variants. List-returning competitors are packed into the common representation inside the timed gather call.

Dependencies, prepared-corpus requirements, and direct commands are documented in docs/benchmarks.md.

Methodology

The Dense and Ragged store tables below were measured on 2026-09-20 from commit 8735409 (2.8.7), after the shared allocator pool and explicit Store/View path were enabled. Each is one complete pass on a warm page cache with the default data_shards=4, not a three-run median. The benchmark verifies exact dtype, shape, order, and values before timing. CPU uses the default 256 MiB allocator pool; AMDGPU uses the default 1 GiB pool. All views in one process share their allocator's aggregate pool.

The logical write and logical gather columns report uncompressed NumPy payload bytes per elapsed second, not physical storage bandwidth. Write timing includes public-API adaptation and logical finalization, but excludes source preparation, dictionary training, writer setup, cleanup, and stable-media sync. Zrecord writers receive the complete prepared source in one append and bound execution internally. Gather timing includes allocation, reads, decompression, and output construction; Ragged packing and metadata production are included. It excludes open, close, sampler time, and destruction after return. Fresh uniform IID gathers accumulate at least two timed seconds. krecords/s is record throughput, p95 is one random batch's 95th-percentile latency, and disk is allocated blocks. Compare results within a workload table, not across different record geometries. Every listed backend and codec is required for the published matrix.

The vision source is the Oxford-IIIT Pet train split prepared by scripts/prepare_vision.py. Dense uses RGB photographs resized on the short side to 256 and center-cropped to (3, 224, 224); Ragged uses the same ordered source at native RGB resolution. Both are mmap-loaded uint8 CHW records.

Machine — one local workstation (AMD Ryzen AI 9 HX PRO 370, 12 cores / 24 threads):

The process sees all 24 threads, is not memory-limited by cgroup, and uses the ordinary page cache. The 16 GiB shared-memory mount accommodates the Torch run.

Large Vision Records

Fixed-Shape Dense

Zrecord against array-store alternatives: random batch gather over the first 2,500 Oxford-IIIT Pet train images, batch 256. Every prepared image is uint8[3,224,224]; metadata.json records the source revision, transform, decoder versions and logical SHA-256.

Fixed-resolution vision records — 147 KiB per record, 36.8 MiB per batch:

At 147 KiB per record, Zrecord-raw reaches 17.0 GiB/s and is 3.7x npy-mmap-raw in logical gather. Decoded photographs have little remaining redundancy: plain zstd reduces them only 1.15x, while dictionary mode falls to 1.12x. This is why large real images should use plain zstd or raw. LMDB and Arrow IPC are competitive raw record stores, while codecs tied to whole IPC batches or Parquet row groups pay read amplification on random gathers. Dense demonstrates that typed record ownership and per-record compression do not turn fixed tensors into an object-store slow path.

Native-Resolution Ragged

This workload contains the first 2,500 native-resolution Oxford-IIIT Pet CHW RGB images. Height is 108..2606 (median 375) and width is 117..3264 (median 500), totaling 1252.3 MiB of logical uint8 payload. Each of 50 random batches contains 256 records. Every backend persists payload plus exact shape and produces an ordered packed batch. Zrecord returns its native RaggedBatch; list-returning competitors allocate and fill equivalent packed values, int32 offsets, and uint64 shapes inside the timed call. Throughput counts payload bytes only, while metadata work remains timed.

Zrecord-raw is 5.3x LMDB and 7.2x Arrow IPC in logical gather because it already produces the contiguous compute representation. Plain and dictionary zstd both reach 1.21x storage reduction and about 2.2 GiB/s, confirming that dictionary mode is not useful for these large photographs. The Ragged matrix is intentionally asymmetric: HDF5 gzip, Arrow IPC zstd and Parquet zstd adapters are not implemented because the real 256-record batch took 1.1–1.7 seconds; TileDB variable queries took 8.6–9.6 seconds and the backend was removed entirely. ArrayRecord zstd remains as the compressed record-store comparison.

Small Token Records

Both token workloads come from the same real corpus: WikiText-103 raw train, tokenized with GPT-2 and stored as int32 IDs. Preparation is outside every measurement. scripts/prepare_tokens.py preserves nonempty text boundaries, combines fragments shorter than 16 tokens, and splits records at 512 tokens into tokens.npy plus offsets.npy; both benchmark scripts mmap those files.

Fixed Token Blocks

The Dense workload ignores text boundaries and packs the stream into 200,000 fixed int32[512] records: 2 KiB per record and 390.6 MiB logical payload. Each IID batch gathers 256 records; fresh draws continue until timed gathers accumulate at least two seconds.

The contiguous NumPy baseline is strongest for gather when the whole corpus is one fixed typed matrix. Zrecord-raw reaches 2.11 Mrecords/s while retaining independent record semantics; dictionary zstd writes 6% slower than plain zstd, uses 13% less disk and gathers 8% faster in this pass. LMDB's B-tree/page overhead is visible in both throughput and disk.

Variable Token Sequences

The Ragged workload keeps 200,000 real text records of 16..512 tokens: p10 28, median 129, mean 138.8, p90 254, totaling 105.9 MiB. It uses 100 independent correctness batches followed by fresh timed IID batches of 256 records and the same RaggedBatch output contract as native-resolution vision.

Here the record contract, not bulk byte bandwidth, is the useful scale. Zrecord-raw returns 1.20 Mrecords/s and is 7.2x LMDB and 16.1x Arrow IPC in this logical-gather comparison. Dictionary zstd writes 36% faster than plain zstd, uses 21% less disk, and gathers 8% faster in this pass.

Sampler

Index generation on its own, IID (with replacement), 1M index space, against NumPy's modern API. The counter-based Cython function is stateless at every step. The µs-scale figures fluctuate with box load; this 10K-draw-per-size run shows a 1.9–10.6x margin across batch sizes.

End-to-End Loader

The current loader benchmark measures CPU and AMDGPU output allocations separately. CPU rows use Loader(..., workers=4, prefetch=4); AMDGPU rows use Loader(..., workers=0) to limit simultaneously live allocations. Benchmark row names are loaderx-cpu, loaderx-amdgpu, and torch; Ragged row names are loaderx-cpu-ragged, loaderx-amdgpu-ragged, and torch-packed-ragged. The Torch ragged path is a packed equivalent with no padding. The current one-pass results use the backend defaults: pool_size=256 MiB for CPU and pool_size=1 GiB for AMDGPU; no explicit pool-size override is used. The snapshot was refreshed after enabling the bounded allocator pool with view-local size-class tracking, so the older Loader snapshot is not directly comparable:

• Vision (Dense) — 2,500 (3,224,224) uint8 records, batch 256 (36.8 MiB), one 200-batch pass.
• Tokens (Ragged) — 100,000 variable-length WikiText token sequences (p10/p50/p90 ≈ 28/129/255), batch 256, no padding.

Backends, per workload:

• loaderx-cpu / loaderx-cpu-ragged: zrecord raw → worker Loader → DLPack → CPU Torch tensors.
• loaderx-amdgpu / loaderx-amdgpu-ragged: zrecord raw → allocator-owned dma-buf → HIP import → DLPack → torch.from_dlpack; ragged yields values + cu_seqlens + shapes (the packed record form).
• torch: four-worker torch.utils.data.DataLoader over one .npy mmap, then .to("cuda").
• torch-packed-ragged: in-process DataLoader over per-record .npy files, packed collation, then .to("cuda").

Each backend runs in a fresh spawned process. Correctness completes first, then the source corpus mmap is evicted because every backend reads its derived Zrecord or .npy artifact during delivery. Ten warmup batches and initial GPU synchronization precede timed-pass memory sampling. peak RAM combines process-tree PSS with per-client device allocations backed by system RAM; peak VRAM comes from the same workers' DRM client accounting. This avoids both counting the inactive preparation corpus and omitting CPU-invisible device mappings. Measured on the AMD Ryzen AI 9 HX PRO 370 / Radeon 890M (gfx1150) under ROCm 10 / torch 2.12.

Vision (Dense):

Tokens (Ragged, no pad):

In these passes, loaderx AMDGPU is 5.5x Torch on Dense and 20.1x on ragged tokens. CPU Loader is 4.7x Torch on Dense and 20.7x on ragged tokens. Dense Torch also peaks at 2.9 GiB RAM versus 0.8 GiB or less for both loaderx paths. The throughput margins include storage layout, collation, and transfer differences; they are end-to-end loader comparisons, not isolated claims about one component.

Current contract. Loaderx batches sampling, gather, and AMDGPU device delivery while preserving exact record semantics. Dense serves fixed-shape arrays directly; Ragged delivers values + cu_seqlens + shapes (no padding, CPU and GPU alike). The AMDGPU path lets Store fill the final device-consumable allocation and avoids an intermediate CPU batch and explicit framework H2D transfer; file I/O remains the ordinary Store path.

These numbers measure warm-cache access, not cold disk or durability. Store reads accumulate at least two timed seconds, while sampler and loader figures can fluctuate with machine load. Treat absolute values as ballpark and cross-backend margins as the main signal.

Real-data verification

Production dataset-specific preprocessing and verification remain in the DataPipe repository. The built-in converter covers materialized Arrow-backed Hugging Face datasets; DataPipe handles sources without a common remote protocol and implements derived modalities as loader transforms without NumPy dump intermediates. Oxford-IIIT Pet is only the shared benchmark fixture.

Current Limitations

• Single-host only; multi-host training is not supported.
• A single sample must be at most 2 GiB. Store size is practically bounded by disk capacity and platform file limits.
• Stores are not portable between machines with different byte orders. All published platforms are little-endian.
• External append accepts CPU and ROCm DLPack devices. NVIDIA CUDA and other DLPack devices require an explicit transfer to a supported source.
• ROCm append and AMDGPU allocation are Linux-specific and require a working HIP runtime, an amdgpu render node, dma-buf export, and a CPU-mappable allocation. Compatibility depends on the kernel, driver, runtime, allocation type, and device topology; the published benchmark hardware is not a universal device certification matrix.
• Dense and all Ragged fields must be C-contiguous and exactly match the store schema. Ragged fields cannot mix devices, records must be nonempty, and one packed append is limited to INT32_MAX values by its offsets.
• ROCm append is synchronous and consumes one framework Allocation at a time. It avoids a full intermediate CPU tensor but may still incur CPU access, migration or interconnect traffic, cache synchronization, compression scratch, and ordinary file I/O.

设计文档

• 性能优化草案
• 分布式架构草案

Build

python3 setup.py build_ext --inplace
zig build test
python3 scripts/test_loaderx.py

The regular suite uses CPU execution plus mocks or memfd regions for device contracts. On a machine configured with an AMD GPU and ROCm Torch, run the separate hardware certification for Torch read/write and framework-owned Dense and Ragged append. Equivalent JAX checks run when a ROCm JAX backend is installed. Passing this script certifies the current machine, not every AMD architecture:

python3 scripts/test_rocm.py

Build the release wheel matrix with:

python3 scripts/build_release.py

Source distributions are intentionally not published. The wheel matrix covers the supported platforms. Source builds require Zig.