veloxquant_mlx.memory.block_pool module¶

Fixed-size block allocator for KV-cache storage, in the vLLM/PagedAttention style.

Provides BlockPoolAllocator, a thread-safe bookkeeping layer that pre-allocates n_blocks fixed-size blocks and hands them out per request/stream so steady-state generation performs no per-token malloc/free. This module tracks only block ownership and usage stats (PoolConfig, Block, AllocationStats); it holds no array data itself — pair it with a backing store such as MLXBlockStorage for actual mx.array buffers, or with PooledKVCache / PoolBackedKVCache for cache-level integration.

class veloxquant_mlx.memory.block_pool.PoolConfig(block_size=16, n_blocks=1024, separate_kv=True, grow_on_exhaustion=False, max_blocks=None)[source]¶

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.block_pool.Block(block_id, stream='kv', format='fp16', n_used=0, owner=None, ever_allocated=False)[source]¶

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.block_pool.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)[source]¶

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()[source]¶

Current number of allocated (non-free) blocks.

Return type:: int

fragmentation()[source]¶

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()[source]¶

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

Return type:: AllocationStats

to_prometheus(prefix='veloxquant_block_pool')[source]¶

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_pool.BlockPoolAllocator(config=None)[source]¶

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)[source]¶

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)[source]¶

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')[source]¶

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)[source]¶

Return blocks to the free list for their stream.

Parameters:: blocks (list[Block]) – Blocks previously returned by allocate().
Return type:: None

free_all(owner)[source]¶

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)[source]¶

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

Return type:: list[Block]
Parameters:: owner (int)

n_free(stream='k')[source]¶

Number of free blocks remaining for stream.

Return type:: int
Parameters:: stream (str)

shrink(stream, target_free)[source]¶

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)[source]¶

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