OmniMalloc

State-of-the-art static memory allocation for neural networks.

OmniMalloc is a Python library for static memory allocation: given buffers with known sizes and lifetimes, assign offsets so that peak memory is minimized. This is the memory-planning step at the heart of ML compilers, embedded runtimes, and accelerator toolchains.

It ships a collection of allocators and allocation algorithms behind one API, implemented with an efficient C++ backend. This includes SuperMalloc, a new allocator that outperforms the best open-source alternatives (see benchmarks below). OmniMalloc also provides a rich benchmark harness and visualization tools to develop and evaluate new allocation strategies.

Installation

Install the latest release from PyPI:

pip install omnimalloc

Or install the development version directly from GitHub:

pip install git+https://github.com/fpedd/omnimalloc.git

Usage

from omnimalloc import Allocation, Pool, run_allocation

pool = Pool(id="pool", allocations=(
    Allocation(id=0, size=64, start=0, end=10),
    Allocation(id=1, size=64, start=12, end=20),
    Allocation(id=2, size=32, start=5, end=15),
))

pool = run_allocation(pool, allocator="supermalloc_allocator", validate=True)

print(pool.size)                                     # 96
print([alloc.offset for alloc in pool.allocations])  # [0, 0, 64]

On a real problem, the result looks like this: 308 buffers of an ML workload packed with no wasted memory.

See examples for allocator selection, visualization, custom allocation sources, and benchmarking.

Benchmarks

All figures on this page are generated from a deterministic benchmark run by scripts/generate_readme_assets.py. Run your own campaigns with the benchmark harness.

Development

# Initial setup
git clone git@github.com:fpedd/omnimalloc.git
cd omnimalloc
uv sync --all-extras --group dev

# Run tests, linting, type checking
uv run pytest
uv run ruff check --fix && uv run ruff format && uv run ty check

# Setup pre-commit hooks (run once)
uv run pre-commit install

# Run pre-commit checks manually
uv run pre-commit run --all-files

License

Copyright 2025 Fabian Peddinghaus. Licensed under Apache 2.0 License. See LICENSE for details.