veloxquant_mlx.cache.squeeze_coordinator module

SqueezeAttention cross-layer coordinator — one-shot data-driven re-budgeting.

The repo’s contract is one cache object per layer; mlx_lm.generate iterates them independently. SqueezeAttention needs a global view to reallocate a fixed total budget across layers by their measured attention concentration. Rather than modify the model forward pass, all SqueezeAttentionCache instances of a model hold a reference to a single shared SqueezeCoordinator (injected at KVCacheBuilder.for_model build time).

The re-budget is one-shot at the prefill boundary: during the first update_and_fetch (the prompt), each layer reports its concentration score. Once every attention layer has reported, the coordinator computes the per-layer budget schedule with squeeze_budgets and publishes it; each layer then pulls its resolved budget and applies it to every head’s eviction state. Decode steps run against the frozen schedule — no further re-budgeting.

Unlike the XQuant / MiniCache coordinators (which exchange tensors every step), this coordinator exchanges only per-layer scalars and runs its allocation exactly once. Single-threaded by construction (mlx generate is sequential), so plain dicts need no locking.

class veloxquant_mlx.cache.squeeze_coordinator.SqueezeCoordinator(n_layers, avg_budget, n_sink, strength)

Bases: object

Shared per-model re-budgeting state for one generation.

Parameters:

• n_layers (int) – Number of attention-bearing layers that will report.

• avg_budget (int) – Target mean per-layer budget (the uniform baseline).

• n_sink (int) – Sink tokens each layer protects (min-budget floor).

• strength (float) – Reallocation strength in [0, 1]; 0.0 == uniform H2O.

reset()

Clear all reported/resolved state (e.g. between generations).

Return type:

None

report_concentration(layer_id, concentration)

Record a layer’s prefill concentration score.

Idempotent per layer: the first report for a layer wins (prefill only); later reports (decode steps) are ignored so the schedule stays frozen.

Once all n_layers have reported, the budget schedule is computed and cached so subsequent resolved_budget calls return real values.

Parameters:

• layer_id (int) – The reporting layer’s index.

• concentration (float) – Its concentration_score over the prefill keys.

Return type:

None

resolved_budget(layer_id)

Return the layer’s resolved budget, or None if not yet finalised.

Parameters:

layer_id (int) – The querying layer’s index.

Return type:

int | None

Returns:

The resolved budget once all layers have reported; None before that (the caller keeps using its average fallback until then).

property is_finalized: bool

True once every layer has reported and the schedule is computed.

property n_layers: int

Number of layers expected to report before the schedule is finalized.

property avg_budget: int

Average per-layer eviction budget the resolved schedule targets.

property strength: float

Configured re-budgeting strength (how aggressively budgets skew from the average).