veloxquant_mlx.cache.xkv_coordinator module

xKV cross-layer coordinator — fan-in/fan-out shared-basis broadcast.

The repo’s contract is one cache object per layer; mlx_lm.generate iterates them independently. xKV needs a group of layers to jointly compute one shared SVD basis before any of them can compress — this is a fan-in (collect every member’s raw prefill keys) then fan-out (broadcast the resulting shared basis back to every member, including the one that triggered the computation) pattern, unlike:

• XQuantCoordinator — a single anchor publishes; reusers read (no fan-in, no joint computation).

• MiniCacheCoordinator — one primary publishes; exactly one merge partner consumes (pairwise fan-in of two, not N).

All XKVCache instances of a model hold a reference to one shared XKVCoordinator (injected at KVCacheBuilder.for_model build time).

State is keyed by (group_id, token_start) so every member of a group publishes for the same token range before the joint SVD runs. The basis is computed lazily on the Nth (last) publish for that key and memoized — every member’s get_shared_basis call for the same key, whenever it arrives, returns the identical basis. Each XKVCache caches its own copy locally after the first successful fetch and never calls the coordinator again for that generation, so the fan-in/fan-out round-trip happens exactly once per group per generation (at prefill).

Single-threaded by construction (mlx generate is sequential), so a plain dict-of-pending-publishes needs no locking.

class veloxquant_mlx.cache.xkv_coordinator.XKVCoordinator(max_ctx=8192)

Bases: object

Shared cross-layer basis store for one model instance / one generation.

Parameters:

max_ctx (int) – Per-group token budget; publishing beyond this raises.

Note: the shared rank / energy-threshold are NOT stored on the coordinator — every member of a group is built from the same KVCacheConfig (see KVCacheBuilder._build_xkv), so the calling XKVCache passes its own xkv_rank/xkv_energy_threshold into get_shared_basis on every call. This avoids a second place where those two values could silently drift out of sync with the config.

reset()

Clear all pending groups and published-token counts (start of a fresh generation).

Return type:

None

publish_member_keys(group_id, member_idx, token_start, n_tokens, keys)

A group member publishes its own raw keys for this token range.

Parameters:

• group_id (int) – The group this layer belongs to.

• member_idx (int) – This layer’s 0-indexed position within the group.

• token_start (int) – Offset of the first token in this write.

• n_tokens (int) – Number of tokens in this write.

• keys (array) – This layer’s own keys for the range, shape [S, D].

Return type:

None

get_shared_basis(group_id, token_start, expected_members, rank=None, energy_threshold=0.95)

Return the shared basis for this group/token range once every expected member has published; None if still waiting.

Computes the joint SVD lazily on the call that observes the last of expected_members publishes, and memoizes the result so every subsequent (or concurrent, in iteration order) caller for the same key gets back the identical arrays. rank/energy_threshold are only consulted on that triggering call (every group member shares the same KVCacheConfig values, so it does not matter which caller’s values happen to trigger the computation).

Return type:

tuple[array, array, array] | None

Parameters:

• group_id (int)

• token_start (int)

• expected_members (int)

• rank (int | None)

• energy_threshold (float)

published_tokens(group_id)

Total tokens the leader member of group_id has published.

Return type:

int

Parameters:

group_id (int)

property max_ctx: int

Configured per-group token budget.