veloxquant_mlx.memory.mlx_storage module

MLX array backing store paired with BlockPoolAllocator.

Supplies MLXBlockStorage, which allocates one mx.array buffer per block id (lazily, on first write) and reuses that buffer in place whenever the pool recycles the block id, so no new allocation happens during steady-state generation. The pool decides which block ids are free/owned; this module owns the actual memory those ids point to and reports the resident bytes of the buffers it actually holds.

class veloxquant_mlx.memory.mlx_storage.MLXBlockStorage(pool, head_dim)

Bases: object

Backing array storage for a BlockPoolAllocator.

Allocates one mx.array buffer per block id, shaped (block_size, head_dim), lazily on first write. Blocks recycled by the pool reuse their existing buffer in place — writing into a reused block overwrites its previous contents without a new mx.array allocation, which is the point of pairing this with the pool (no malloc/free churn during steady-state generation).

Because different blocks may hold different compression formats (Block.format), each block’s buffer dtype is chosen from its format at write time; a block’s buffer is only reallocated when the format assigned to that block id actually changes (rare — format is normally stable for the life of a request).

Parameters:

• pool (BlockPoolAllocator) – The allocator this storage is paired with.

• head_dim (int) – Width of each stored vector (last axis of every buffer).

buffer_for(block)

Return the mx.array buffer backing block, allocating it if needed.

Parameters:

block (Block) – A block returned by pool.allocate().

Returns:

mx.array of shape (pool.config.block_size, head_dim).

write(block, offset, values)

Write values into block starting at token-slot offset.

Parameters:

• block (Block) – Destination block.

• offset (int) – Token-slot offset within the block (0-indexed).

• values – mx.array of shape (n, head_dim) to write, n <= block_size - offset.

Return type:

None

read(block)

Return the live (used) slice of block’s buffer.

Parameters:

block (Block) – Source block.

Returns:

mx.array of shape (block.n_used, head_dim).

resident_bytes()

Total bytes currently materialized across all allocated buffers.

Only counts blocks that have had a buffer allocated (lazily, on first write/read) — this is the actual resident footprint, not the pool’s full pre-allocated capacity.

Return type:

int

release(block)

Drop the buffer for a freed block, letting MLX reclaim its memory.

Call this after pool.free([block]) if you want resident memory to shrink immediately; otherwise the buffer is kept around (still zero new-allocation cost) and silently reused/overwritten the next time this block id is handed out.

Return type:

None

Parameters:

block (Block)