A StorageBackend is the only thing Flatrun's runtime talks to
when it needs a tensor. Backends are registered through
flatrun.backend.registry.default_registry(); the shipped
implementations wire themselves in on first import.
class StorageBackend(abc.ABC):
name: str
supports_mmap: bool
def open(self) -> None: ...
def close(self) -> None: ...
def list_tensors(self) -> Iterator[TensorKey]: ...
def open_handle(self, name: str) -> TensorHandle: ...openis called once before any other method. It must succeed for the backend to be usable.list_tensorsreturns every tensor the backend exposes, in whatever order the backend chooses.open_handlereturns aTensorHandlefor the named tensor. The handle must be cheap to construct; reading bytes happens lazily insidehandle.view().
A TensorHandle exposes:
class TensorHandle(abc.ABC):
@property
def metadata(self) -> TensorMetadata: ... # shape, dtype, byte_size
def view(self) -> TensorView: ... # zero-copy NumPy view
def materialize(self) -> bytes: ... # eager byte copy
def close(self) -> None: ...If the handle wraps a quantised GGUF tensor, metadata.dtype
is uint8 and metadata.quantization names the GGML type
(Q4_0, Q4_K, Q8_0, ...). The forwarder resolves those
through flatrun.dequant.dequant_handle.
SafeTensorBackend parses the JSON header itself, mmap's the
file once on open, and hands out MmapTensorHandles that
are zero-copy views into the mapping. The runtime never has
to read the full file into RAM.
Native dtypes supported: F32, F16, BF16. Quantised
shards (MLX 4-bit, etc.) are exposed as uint8 bytes; the
dequant module turns them back into floats at fetch time.
When a tensor handle closes (for example after the runtime
has run the layer that uses it), the underlying mmap region
is returned to the OS via madvise(MADV_DONTNEED) for
tensors >= 1 MiB. Without this hint, the page cache would
keep the touched pages resident long after the model has
moved on to the next layer — RSS would grow with each layer
touched even though the cache cap is honoured. This is
particularly important for streamed inference on
memory-constrained hosts.
GGUFBackend reads GGUF v3 metadata and the tensor info table
in pure Python. It supports:
- Tensor name translation (Qwen3, Llama, etc.) so the same manifest builder that works for SafeTensors also works for GGUF.
bf16,f16,f32tensors.- The full set of quant types in
flatrun.dequant.gguf:Q1_0,Q4_0,Q5_0,Q5_1,Q4_K,Q5_K,Q6_K,Q8_0.
GGUF's tensor dimensions are stored fastest-axis-first
(ne[0] is the contiguous one); the backend reverses them
on read so metadata.shape matches the NumPy row-major
convention. Skipping that reversal reinterprets every flat
buffer under the wrong stride - the same buffer, shape-legal
all the way to the logits, just scrambled.
Some forks of llama.cpp add new GGML type ids. Flatrun ships
a small type table; if you hit Unknown(N) while loading a
file, add the type to _GGML_TYPES in
flatrun.backend.gguf and a matching decoder to
flatrun.dequant.gguf (and register it in
flatrun.dequant.loader._GGUF_DECODERS).
MultiBackend composes several backends under one name, so
sharded checkpoints (multiple model-00001-of-00030.safetensors
files) appear as a single backend. The runtime is
shard-agnostic: it just sees one big list of tensors and asks
for handles by name.
- Implement
StorageBackend(typically a few hundred lines). - Register it in
flatrun.backend._populate.populatewith the suffix / filename you want it to handle. - If the format uses an exotic quant, add a decoder under
flatrun.dequantand wire it into_GGUF_DECODERS. - Add a unit test that round-trips a synthetic file through the backend.
The runtime will pick up the new format automatically; no
changes to flatrun.runtime are needed.