veloxquant_mlx.memory package

KV-cache-aware block pool allocator (issue #249).

Fixed-size block allocation and reuse for KV-cache storage, addressing fragmentation and repeated malloc/free overhead during long-running or multi-request inference — the memory-allocation analogue of the bit-level compression the rest of VeloxQuant-MLX provides.

Core pieces:

• BlockPoolAllocator / PoolConfig — pure-Python block bookkeeping (allocate, free, reuse, per-owner tracking, fragmentation and allocation stats). No MLX dependency; usable and testable on any platform.

• MLXBlockStorage — pairs with a pool to provide the actual mx.array buffers, one per block id, reused in place across requests and supporting mixed compression formats (different blocks may hold fp16, int8, int4, … data) within the same pool.

• PooledKVCache — wraps any existing VeloxQuant KVCache so its token appends check out/return blocks from a shared pool, without changing the wrapped cache’s own compression logic.

• PoolBackedKVCache / build_pooled_caches() — a drop-in replacement for mlx_lm.models.cache.KVCache whose growth is tracked by a shared pool. Unlike PooledKVCache (which wraps VeloxQuant’s own append_key/append_value/attend interface, implemented only by the 5 “standalone” methods), this implements mlx_lm’s update_and_fetch protocol directly, so it plugs into mlx_lm.generate() for any model — this is what actually puts the pool in a real model’s decode hot path.

Typical usage:

from veloxquant_mlx.memory import BlockPoolAllocator, PoolConfig, PooledKVCache
from veloxquant_mlx import KVCacheBuilder, KVCacheConfig

pool = BlockPoolAllocator(PoolConfig(block_size=16, n_blocks=512))
inner = KVCacheBuilder().with_method("kivi").with_head_dim(128).build()
cache = PooledKVCache(inner, pool, owner=request_id, format="int2")
...
cache.release()  # return blocks to the pool for reuse by the next request

Driving real mlx_lm.generate() traffic through the pool:

import mlx_lm
from veloxquant_mlx.memory import BlockPoolAllocator, PoolConfig, build_pooled_caches

model, tokenizer = mlx_lm.load("mlx-community/Llama-3.2-3B-Instruct-4bit")
pool = BlockPoolAllocator(PoolConfig(block_size=16, n_blocks=4096))
cache = build_pooled_caches(model, pool, owner=request_id)
text = mlx_lm.generate(model, tokenizer, prompt, prompt_cache=cache)
cache[0].release()  # every layer shares one owner, so any one releases all

class veloxquant_mlx.memory.AllocationStats(n_blocks, n_allocations=0, n_frees=0, n_reused=0, n_exhausted=0, n_grown=0, n_retired=0, peak_blocks_in_use=0, _fragmented_count=0)

Bases: object

Running counters for a BlockPoolAllocator.

All counts are cumulative since pool construction unless noted.

Variables:

• n_blocks – Total blocks managed by the pool.

• n_allocations – Number of successful allocate() calls.

• n_frees – Number of free() calls.

• n_reused – Allocations satisfied by a block that had previously been freed and returned to the pool (as opposed to a block that had never been handed out before).

• n_exhausted – Allocation attempts that failed because no free block was available.

• n_grown – Number of blocks added to the pool by exhaustion-triggered growth (see PoolConfig.grow_on_exhaustion).

• n_retired – Number of blocks permanently removed from the pool by shrink() (see BlockPoolAllocator.shrink).

• peak_blocks_in_use – High-water mark of concurrently allocated blocks.

Parameters:

• n_blocks (int)

• n_allocations (int)

• n_frees (int)

• n_reused (int)

• n_exhausted (int)

• n_grown (int)

• n_retired (int)

• peak_blocks_in_use (int)

• _fragmented_count (int)

n_blocks: int

n_allocations: int = 0

n_frees: int = 0

n_reused: int = 0

n_exhausted: int = 0

n_grown: int = 0

n_retired: int = 0

peak_blocks_in_use: int = 0

blocks_in_use()

Current number of allocated (non-free) blocks.

Return type:

int

fragmentation()

Fraction of managed blocks that are neither free nor fully used.

A block counts as fragmented once it has been allocated (n_used > 0) but is not full. Returns 0.0 for an allocator with no blocks.

Return type:

float

snapshot()

Return an independent copy of this stats object for history/export.

Return type:

AllocationStats

to_prometheus(prefix='veloxquant_block_pool')

Render these stats as Prometheus text-exposition-format gauges/counters.

Parameters:

prefix (str) – Metric name prefix.

Return type:

str

Returns:

Newline-terminated text in the Prometheus exposition format, suitable for serving directly from a /metrics endpoint or appending to a scrape response.

class veloxquant_mlx.memory.Block(block_id, stream='kv', format='fp16', n_used=0, owner=None, ever_allocated=False)

Bases: object

A single fixed-size block of KV-cache storage.

Variables:

• block_id – Index into the pool’s backing storage; stable for the lifetime of the pool (blocks are reused in place, never moved).

• stream – Which logical stream (“k” or “v”) this block belongs to. Always “kv” when the owning pool was built with separate_kv=False.

• format – Compression-format tag of the data currently stored in this block (e.g. “fp16”, “int2”, “int4”). Lets a single pool host blocks written by different quantization methods concurrently.

• n_used – Number of token slots currently occupied (<= block_size).

• owner – Opaque request id the block is currently checked out to, or None if the block is free.

• ever_allocated – True once this block id has been handed out by BlockPoolAllocator.allocate() at least once. Used to distinguish a fresh allocation from a reuse for stats purposes.

Parameters:

• block_id (int)

• stream (str)

• format (str)

• n_used (int)

• owner (int | None)

• ever_allocated (bool)

block_id: int

stream: str = 'kv'

format: str = 'fp16'

n_used: int = 0

owner: int | None = None

ever_allocated: bool = False

class veloxquant_mlx.memory.BlockPoolAllocator(config=None)

Bases: object

Fixed-size block allocator for KV-cache storage.

Pre-allocates n_blocks fixed-size blocks up front and hands them out on request, avoiding per-token malloc/free during generation. Freed blocks return to a free list and are reused by later allocations (LIFO — a just-freed block tends to be cache-hot), so steady-state generation performs zero new-memory allocation once the pool is warm.

Keys and values can be tracked on independent free lists (PoolConfig.separate_kv=True, the default) so K-heavy and V-heavy workloads (e.g. asymmetric bit-widths) don’t compete for the same blocks, or share a single free list for tighter packing when K/V pressure is uneven across requests.

Thread-safe: every mutating and read method acquires an internal lock, so one pool can be shared across threads/concurrent request handlers without external synchronization. The lock is a plain (non-reentrant) threading.Lock — do not call back into the pool from inside a callback invoked while holding it.

This class only manages block bookkeeping (which block ids are free, who owns which, how full each is, fragmentation/latency stats). It does not itself own array storage — pair it with a backing store (e.g. veloxquant_mlx.memory.mlx_storage.MLXBlockStorage) that allocates the actual mx.array buffers indexed by block id.

Parameters:

config (PoolConfig | None) – Pool sizing/behavior configuration.

Example:

pool = BlockPoolAllocator(PoolConfig(block_size=16, n_blocks=256))
blocks = pool.allocate(stream="k", n_tokens=40, owner=request_id, format="int2")
...
pool.free_all(owner=request_id)

history: deque[AllocationStats]

register_owner(owner)

Mark owner as active, guarding against accidental id reuse.

allocate() calls this automatically for any owner not yet seen, but a second allocate() for an owner that is already active is treated as that same request continuing to grow, not a collision (see allocate()’s docstring) — it can’t otherwise tell the two cases apart. Call register_owner() explicitly yourself (once, before any allocate() calls for this owner) if you want a genuine collision — e.g. a request-id counter that wrapped around, or two callers racing on the same id — to raise immediately, including on a second register_owner() call for a still-active id.

Raises:

OwnerAlreadyActiveError – If owner already has blocks checked out or was registered and not yet released.

Return type:

None

Parameters:

owner (int)

release_owner(owner)

Clear owner from the active-owner set without freeing blocks.

Rarely needed directly — free_all() calls this once an owner’s last block is freed. Useful if you registered an owner via register_owner() but never actually allocated anything for it.

Return type:

None

Parameters:

owner (int)

allocate(stream, n_tokens, owner, format='fp16')

Allocate enough blocks to hold n_tokens slots.

Parameters:

• stream (str) – “k” or “v” (ignored, treated as “kv”, when the pool was built with separate_kv=False).

• n_tokens (int) – Number of token slots needed.

• owner (int) – Opaque id (e.g. request id) the blocks are checked out to. Repeated allocate() calls for an owner that is already active on this pool are treated as that same request growing (the normal case — see PoolBackedKVCache) and never raise. To catch a different caller accidentally reusing an id that is still checked out, call register_owner() yourself before the first allocate() — a second register_owner() for an id that is still active raises OwnerAlreadyActiveError.

• format (str) – Compression-format tag to record on each block.

Return type:

list[Block]

Returns:

List of newly allocated Block instances, in order.

Raises:

• BlockPoolExhaustedError – If fewer than the required number of free blocks remain and the pool cannot grow enough to cover the gap (see PoolConfig.grow_on_exhaustion/max_blocks). No blocks are allocated on failure (all-or-nothing).

• ValueError – If n_tokens <= 0 or stream is invalid.

free(blocks)

Return blocks to the free list for their stream.

Parameters:

blocks (list[Block]) – Blocks previously returned by allocate().

Return type:

None

free_all(owner)

Free every block currently checked out to owner.

Parameters:

owner (int) – Opaque id passed to a prior allocate() call.

Return type:

None

Unlike free(), which frees an arbitrary caller-supplied subset and must locate+remove each block id from the owner’s list, this already has the owner’s complete, correct block-id list – so it pops that list once (O(1)) instead of doing N membership checks + N list.remove() calls (VeloxQuant-MLX#510; was 13.4ms at 256K tokens / 16K blocks, confirmed superlinear: 2x blocks -> 2.0-3.4x time).

blocks_for(owner)

Return the blocks currently checked out to owner, in allocation order.

Return type:

list[Block]

Parameters:

owner (int)

n_free(stream='k')

Number of free blocks remaining for stream.

Return type:

int

Parameters:

stream (str)

shrink(stream, target_free)

Permanently retire fully-free blocks from stream down to target_free.

Use this to give memory back when a pool was sized for peak load that turned out much higher than steady-state usage. Only blocks currently on the free list are eligible — in-use blocks are never touched, so this never breaks a live request. Retired blocks are removed from the pool’s bookkeeping entirely (not just marked free/unfree), so any external backing storage indexed by block id (e.g. MLXBlockStorage) may drop the corresponding buffer.

Parameters:

• stream (str) – “k” or “v” (ignored, treated as “kv”, when the pool was built with separate_kv=False).

• target_free (int) – Desired number of free blocks to retain for this stream after shrinking. If the stream already has this many or fewer free blocks, this is a no-op.

Return type:

int

Returns:

Number of blocks actually retired.

Raises:

ValueError – If target_free < 0.

mark_used(block, n_used)

Update block.n_used and keep fragmentation stats consistent.

Callers that track a block’s fill level themselves (e.g. PooledKVCache, which appends token-by-token to an already-checked-out block) must go through this instead of writing block.n_used directly – AllocationStats._fragmented_count is maintained incrementally (not recomputed by a full scan every allocate()/free()/shrink() call, which would be O(n_blocks) on every single-block growth step in the decode hot path), so it only stays correct if every write to a block’s fill level is observed here.

Parameters:

• block (Block) – A block currently checked out from this pool (owner is not None). Blocks that are free are never fragmented by definition, so calling this on one is a no-op for stats (n_used is still updated).

• n_used (int)

Return type:

None

class veloxquant_mlx.memory.PoolConfig(block_size=16, n_blocks=1024, separate_kv=True, grow_on_exhaustion=False, max_blocks=None)

Bases: object

Configuration for a BlockPoolAllocator.

Variables:

• block_size – Number of token slots stored per block.

• n_blocks – Total number of blocks the pool manages initially.

• separate_kv – If True, keys and values are allocated from independent free lists (K and V blocks are never interchanged). If False, a single free list backs both keys and values, which packs tighter when K and V pressure is uneven across requests.

• grow_on_exhaustion – If True, an allocate() call that would otherwise raise BlockPoolExhaustedError instead grows the exhausted stream’s block count just enough to satisfy the request (see max_blocks). If False (default), exhaustion always raises.

• max_blocks – Upper bound on total blocks across all streams once grown. None means unbounded growth. Ignored when grow_on_exhaustion is False. Growth that would exceed this cap still raises BlockPoolExhaustedError.

Parameters:

• block_size (int)

• n_blocks (int)

• separate_kv (bool)

• grow_on_exhaustion (bool)

• max_blocks (int | None)

block_size: int = 16

n_blocks: int = 1024

separate_kv: bool = True

grow_on_exhaustion: bool = False

max_blocks: int | None = None

class veloxquant_mlx.memory.PooledKVCache(inner, pool, owner=-1, format='fp16')

Bases: KVCache

Wraps any VeloxQuant KVCache with block-pool-backed memory accounting.

The wrapped cache still owns and drives its own storage arrays (its compression format may be anything registered with KVCacheFactory); this wrapper checks out and returns fixed-size blocks from a shared BlockPoolAllocator in lock-step with token appends, so a multi-request server can track allocation counts, reuse, and fragmentation across every active cache from one pool instead of each cache doing its own ad hoc growth.

Block granularity is pool.config.block_size tokens: a new block is requested from the pool only when the previous one fills up, and every block owned by this cache is returned to the pool in one call when the request completes (release()). No block is ever allocated or freed on a per-token basis, matching the “avoid malloc/free during generation” goal in issue #249.

Parameters:

• inner (KVCache) – The underlying KVCache that performs actual compression.

• pool (BlockPoolAllocator) – Shared BlockPoolAllocator to draw block accounting from.

• owner (int) – Opaque id identifying this cache’s request/sequence. Two PooledKVCache instances must use different owners to avoid each other’s release() reclaiming the wrong blocks.

• format (str) – Compression-format tag recorded on each block this cache checks out (purely informational — passed through to pool.allocate(format=...)).

append_key(k)

Check out a K block from the pool if needed, then append to the inner cache.

Parameters:

k (Any) – Key vector, shape (d,), fp16.

Return type:

None

append_value(v)

Check out a V block from the pool if needed, then append to the inner cache.

Parameters:

v (Any) – Value vector, shape (d,), fp16.

Return type:

None

attend(q)

Delegate attention computation to the inner cache.

Parameters:

q (Any) – Query vector, shape (d,), fp16.

Return type:

Any

Returns:

Attention output, shape (d,), fp16.

memory_bytes()

Return memory usage of the inner (compressed) cache.

Return type:

int

release()

Return every block this cache holds to the pool.

Call this once the request finishes (or is evicted) so its blocks become available for reuse by later requests.

Return type:

None

n_blocks_held()

Total K + V blocks currently checked out to this cache.

Return type:

int

reset()

Return all held blocks to the pool and clear the inner cache.

Return type:

None

class veloxquant_mlx.memory.PoolBackedKVCache(pool, owner=-1, step=None)

Bases: object

Drop-in replacement for mlx_lm.models.cache.KVCache whose growth is tracked by a shared BlockPoolAllocator.

mlx_lm’s stock KVCache grows its backing mx.array in fixed step=256-token chunks, hardcoded per instance, with no visibility into how many times that growth happened across a server’s caches or how much of the last chunk is wasted. This class keeps the same contiguous-buffer growth strategy — attention needs one contiguous (B, n_kv_heads, seq_len, head_dim) tensor every step, and RoPE / masking read cache.offset directly, so growth cannot be replaced with true block-paged storage without a per-step gather cost — but routes every growth step through pool.allocate() so it shows up in the pool’s AllocationStats (allocation count, reuse, fragmentation) the same way any other pool consumer’s growth does. The growth chunk size itself defaults to stock’s 256-token cadence (VeloxQuant-MLX#510), independent of the pool’s (typically much finer) accounting block_size – growing in 16-token chunks to match a default block_size=16 pool cost 8.6-10.2% measured decode throughput vs. stock, since MLX dispatch/graph-construction overhead tracks mx.concatenate call count, not bytes moved.

This is the piece that puts BlockPoolAllocator in a real model’s decode hot path: unlike PooledKVCache, which wraps VeloxQuant’s own append_key/append_value/attend interface (only implemented by the 5 “standalone” methods — see STANDALONE_METHODS in veloxquant_mlx/cache/base.py), this class implements mlx_lm’s update_and_fetch protocol directly, so it plugs into mlx_lm.generate() for any model exactly where a stock KVCache would.

Parameters:

• pool (BlockPoolAllocator) – Shared BlockPoolAllocator to draw growth accounting from.

• owner (int) – Opaque id identifying this cache’s request/sequence. Two concurrent PoolBackedKVCache instances must use different owners so release() only reclaims their own blocks.

• step (int | None) – Token-chunk size to grow the backing buffer by. Defaults to max(pool.config.block_size, 256), matching stock KVCache’s hardcoded 256-token growth cadence (see _DEFAULT_GROWTH_STEP) regardless of the pool’s (typically much smaller) accounting block size. Pass an explicit value to override either direction.

update_and_fetch(keys, values)

Append one decode step’s keys/values and return the full history.

Parameters:

• keys (Any) – New keys, shape (B, n_kv_heads, S, head_dim).

• values (Any) – New values, shape (B, n_kv_heads, S, head_dim).

Returns:

(keys, values) accumulated so far, each shape (B, n_kv_heads, offset, head_dim).

release()

Return every block this cache’s growth checked out to the pool.

Call this once the request finishes so its accounted blocks free up pool headroom for the next request. The backing mx.array itself is reclaimed by MLX’s normal garbage collection when this object goes out of scope; this only clears the pool’s bookkeeping.

Return type:

None

size()

Number of tokens currently stored (the cache’s offset).

Return type:

int

property state

(keys, values) trimmed to the current offset.

property meta_state

No extra metadata for this cache type; always empty.

is_trimmable()

Always True: offset can be rewound via trim().

Return type:

bool

trim(n)

Rewind the offset by up to n tokens; returns the amount actually trimmed.

Return type:

int

Parameters:

n (int)

make_mask(*args, **kwargs)

Build an attention mask for the current offset via mlx_lm’s create_attention_mask.

empty()

True if no keys/values have been written yet.

Return type:

bool

property nbytes: int

Combined byte size of the backing keys and values buffers.

veloxquant_mlx.memory.build_pooled_caches(model, pool, owner, step=None)

Build one PoolBackedKVCache per language-model layer.

Drop-in replacement for mlx_lm.models.cache.make_prompt_cache / model.make_cache() for the “plain fp16, pool-tracked growth” case — every attention-bearing layer gets a PoolBackedKVCache sharing the same pool and owner, so a single pool.free_all(owner) (or calling PoolBackedKVCache.release() on any one of them) releases every layer’s growth accounting for this request at once.

Parameters:

• model – Loaded mlx_lm model instance.

• pool (BlockPoolAllocator) – Shared BlockPoolAllocator every layer’s cache draws from.

• owner (int) – Opaque id (e.g. request id) shared by every layer’s cache for this request.

• step (int | None) – Growth chunk size in tokens; defaults to max(pool.config.block_size, 256) (see PoolBackedKVCache).

Return type:

list

Returns:

List of PoolBackedKVCache instances, one per language-model layer.

Submodules

• veloxquant_mlx.memory.block_pool module
  • PoolConfig
    • PoolConfig.block_size
    • PoolConfig.n_blocks
    • PoolConfig.separate_kv
    • PoolConfig.grow_on_exhaustion
    • PoolConfig.max_blocks
  • Block
    • Block.block_id
    • Block.stream
    • Block.format
    • Block.n_used
    • Block.owner
    • Block.ever_allocated
  • AllocationStats
    • AllocationStats.n_blocks
    • AllocationStats.n_allocations
    • AllocationStats.n_frees
    • AllocationStats.n_reused
    • AllocationStats.n_exhausted
    • AllocationStats.n_grown
    • AllocationStats.n_retired
    • AllocationStats.peak_blocks_in_use
    • AllocationStats.blocks_in_use()
    • AllocationStats.fragmentation()
    • AllocationStats.snapshot()
    • AllocationStats.to_prometheus()
  • BlockPoolAllocator
    • BlockPoolAllocator.history
    • BlockPoolAllocator.register_owner()
    • BlockPoolAllocator.release_owner()
    • BlockPoolAllocator.allocate()
    • BlockPoolAllocator.free()
    • BlockPoolAllocator.free_all()
    • BlockPoolAllocator.blocks_for()
    • BlockPoolAllocator.n_free()
    • BlockPoolAllocator.shrink()
    • BlockPoolAllocator.mark_used()
• veloxquant_mlx.memory.mlx_storage module
  • MLXBlockStorage
    • MLXBlockStorage.buffer_for()
    • MLXBlockStorage.write()
    • MLXBlockStorage.read()
    • MLXBlockStorage.resident_bytes()
    • MLXBlockStorage.release()
• veloxquant_mlx.memory.pool_backed_cache module
  • PoolBackedKVCache
    • PoolBackedKVCache.update_and_fetch()
    • PoolBackedKVCache.release()
    • PoolBackedKVCache.size()
    • PoolBackedKVCache.state
    • PoolBackedKVCache.meta_state
    • PoolBackedKVCache.is_trimmable()
    • PoolBackedKVCache.trim()
    • PoolBackedKVCache.make_mask()
    • PoolBackedKVCache.empty()
    • PoolBackedKVCache.nbytes
  • build_pooled_caches()
• veloxquant_mlx.memory.pooled_cache module
  • PooledKVCache
    • PooledKVCache.append_key()
    • PooledKVCache.append_value()
    • PooledKVCache.attend()
    • PooledKVCache.memory_bytes()
    • PooledKVCache.release()
    • PooledKVCache.n_blocks_held()
    • PooledKVCache.reset()