veloxquant_mlx.transfer.apply module

Apply a fitted mapper to a source KV cache (paper §3, inference path).

This is the path that replaces the target’s prefill. Per target layer it does

K̂_t = (strip_rope(K_s) · W_K + b_K) · target_rope V̂_t = V_s · W_V + b_V

as one batched matmul per layer over all heads, then hands the result back as arrays the caller can install into the target model’s cache.

The paper measures 2.7–25× versus re-prefill. That ratio comes from replacing the target’s whole transformer forward pass with a matmul, so it widens as the target grows and as the sequence lengthens. It is a compute saving only — the mapped cache is exactly as large as the target’s own would have been.

veloxquant_mlx.transfer.apply.load_mapper(path)

Load a mapper saved by CrossModelMapper.save().

Weights are read lazily per layer on first use, so this is cheap even for a multi-GB artifact.

Return type:

CrossModelMapper

Parameters:

path (str | Path)

veloxquant_mlx.transfer.apply.transfer_layer(mapper, target_layer, source_kv, positions)

Map the source cache into one target layer’s keys and values.

Parameters:

• mapper (CrossModelMapper) – Fitted mapper.

• target_layer (int) – Which target layer to produce.

• source_kv (dict[int, tuple[array, array]]) – {layer: (keys, values)}, [n_kv_heads, N, head_dim], keys post-RoPE as stored by the source model.

• positions (array) – [N] absolute positions of the cached tokens.

Return type:

tuple[array, array]

Returns:

(keys, values) for the target layer, keys carrying the target’s RoPE, ready to install into the target’s cache.

veloxquant_mlx.transfer.apply.transfer_cache(source_kv, mapper, positions=None, unload_after=True)

Map a whole source KV cache into the target model’s layout.

Parameters:

• source_kv (dict[int, tuple[mx.array, mx.array]]) – {layer: (keys, values)} from the source model, each [n_kv_heads, N, head_dim], keys post-RoPE.

• mapper (CrossModelMapper) – Fitted mapper for this exact model pair.

• positions (mx.array | None) – [N] absolute positions. Defaults to 0..N-1, correct for a cache prefilled from the start of a sequence.

• unload_after (bool) – Drop each layer’s weights once it has been applied. Keeps peak memory at one layer’s weights instead of the whole multi-GB mapper — the reason this subsystem is usable on Apple Silicon at all. Set False when mapping many caches in a row.

Return type:

dict[int, tuple[mx.array, mx.array]]

Returns:

{target_layer: (keys, values)} ready for the target model.