veloxquant_mlx.memory.pooled_cache module

Block-pool memory accounting wrapper for any VeloxQuant-compressed KV cache.

Provides PooledKVCache, a KVCache decorator that delegates all compression/attend logic to an inner cache (any of the five “standalone” methods registered with KVCacheFactory) while checking blocks in and out of a shared BlockPoolAllocator at pool.config.block_size-token granularity. This lets a multi-request server observe allocation counts, reuse, and fragmentation across every active compressed cache from one pool, without the inner cache’s own storage or compression scheme having to know about pooling. Contrast with PoolBackedKVCache, which pool-tracks the plain fp16 mlx_lm cache protocol directly instead of wrapping a VeloxQuant cache.

class veloxquant_mlx.memory.pooled_cache.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