veloxquant_mlx.quantizers.amc module

AMC-adapted quantizer — saliency-driven per-token tiered rank + precision.

Inspired by “Adaptive Model Compression (AMC): Saliency-Driven Resource Allocation for Ultra-Low-Power Transformer Inference” (Hu, Yuan, Hu, Yin, Li, Suchter — Apple; arXiv:2607.10109v1 [cs.IR], submitted 2026-07-11). Documented as “AMC-adapted (VeloxQuant-MLX implementation)” — not a faithful port.

No verified peer-reviewed venue as of 2026-07-14 — single-version preprint, no Comments/journal-ref field. This is the second of 40 methods in this repo shipped without a verified venue (the first was NestedKV-adapted, see veloxquant_mlx.quantizers.nestedkv), at the user’s explicit direction. This is a one-time exception, not a new standing precedent — the next new method reverts to requiring a verified venue. State this plainly everywhere AMC is documented.

Scope cut (read before using this module): AMC’s source paper is a hardware/software co-design — roughly half of it (Sections IV-V: 45nm CMOS RTL, Verilog clock-gating, the Precision-Gated Systolic Array, the Narrow-Width SRAM write-back buffer, all pJ/µJ energy figures, the EDAP/Pareto silicon comparisons) targets physical silicon that has no analogue in a software MLX library. None of that is implemented here. This module ports only the portable software half: Section II-A’s saliency engine (L1-norm score, three-tier percentile partitioning, the query-aware semantic-saliency blend, and the sequence-adaptive closed-loop threshold adjustment) and Section III’s adaptive resource scaling (Hadamard rank masking + linear fixed-point-style quantization), plus the offline SVD/PCA channel-order calibration from Algorithm 1 Phase I (see veloxquant_mlx.quantizers.amc_calibration). The paper’s headline 59.2%-energy / 2.24x-throughput / 3.6%-accuracy numbers are the paper’s own hardware-measured figures under its own synthetic 3-layer setup — not reproduced by this software port. See benchmark_scripts/benchmark_amc.py for this repo’s own (unrelated, software-only) synthetic benchmark.

Compression-only, never eviction: unlike every eviction-family method in this repo (H2O, SnapKV, PyramidKV, CurDKV, NestedKV, …), AMC never drops a token. Every token is retained; only its rank (how many hidden-dim channels survive) and bit-width (quantization precision) are reduced for low-saliency tokens. This is a structurally different family: “adaptive rank+precision,” not “eviction.”

Mechanism (per token, every prefill and decode step — the paper’s own design applies tiering continuously, not once at prefill like NestedKV):

1. Saliency score S_i — mean absolute activation magnitude (L1-norm, Eq. 1-2), optionally blended with query-aware cosine similarity (Eq. 3).

2. Percentile-threshold tier assignment: top k_high → High, next k_mid → Mid, remainder → Low (Algorithm 1 Phase II).

3. Rank masking (Eq. 6): zero out channels beyond the tier’s rank, on channels already reordered by the offline calibration permutation so the surviving prefix is the highest-variance subspace.

4. Linear quantization (Eq. 7) to the tier’s bit-width.

Adaptation notes (stated plainly, mirroring every other “-adapted” method’s honest-deviation convention):

• Query-aware saliency (Eq. 3) and sequence-adaptive closed-loop thresholds (Eq. 4-5) are opt-in and off by default — the default path is pure magnitude-only scoring (Eq. 1-2), matching the paper’s primary reported configuration.

• Requires the offline calibration permutation (see veloxquant_mlx.quantizers.amc_calibration) to be meaningful; using AMC without it truncates arbitrary, not lowest-variance, channels — the same category of footgun as Palu/SVDq/RaBitQ’s calibration requirement.

• cs.IR is an unusual arXiv category for what is fundamentally a hardware architecture paper — noted as a minor oddity, not a disqualifier.

Byte accounting:

amc_kept_bytes — sum of actual per-tier bit-width bytes for K + V full_seq_bytes — hypothetical fp16 cost if every token were 16-bit,

full rank

compression_ratio — full_seq_bytes / amc_kept_bytes (> 1 = savings)

class veloxquant_mlx.quantizers.amc.AMCTierConfig(tier, rank, bits)

Bases: object

One tier’s (rank, bit-width) pair — Algorithm 1’s Golden Model values.

Parameters:

• tier (int)

• rank (int)

• bits (int)

tier: int

rank: int

bits: int

class veloxquant_mlx.quantizers.amc.AMCThresholdState(window=<factory>, calib_variance=1.0)

Bases: object

Trailing-window saliency variance tracker for closed-loop thresholding.

Parameters:

• window (RingBuffer)

• calib_variance (float)

window: RingBuffer

calib_variance: float = 1.0

veloxquant_mlx.quantizers.amc.init_amc_threshold_state(window_size, calib_variance)

Create a fresh threshold-adaptation state.

Parameters:

• window_size (int) – Trailing window length (RingBuffer capacity).

• calib_variance (float) – Nominal activation variance from offline calibration (sigma^2_calib in Eq. 4-5). Must be > 0.

Return type:

AMCThresholdState

veloxquant_mlx.quantizers.amc.amc_saliency(x)

Magnitude-based saliency score (Eq. 1-2): mean(|x|) clamped to [0, 1].

Parameters:

x (array) – [N, D] token activations.

Return type:

array

Returns:

[N] saliency scores in [0, 1].

veloxquant_mlx.quantizers.amc.amc_query_aware_saliency(x, keys, query, alpha=0.5)

Query-aware semantic saliency (Eq. 3): blend of magnitude and cosine similarity to a query/prompt vector.

S_i = alpha * mean(|x_i|) + (1 - alpha) * cosine_similarity(query, k_i)

Parameters:

• x (array) – [N, D] token activations (for the magnitude term).

• keys (array) – [N, D] key projection vectors (for the semantic term).

• query (array) – [D] embedded query/prompt vector.

• alpha (float) – Balance coefficient in [0, 1]; 1.0 == pure magnitude (amc_saliency()), 0.0 == pure semantic relevance.

Return type:

array

Returns:

[N] saliency scores in [0, 1].

veloxquant_mlx.quantizers.amc.amc_query_aware_saliency_batched(x, keys, query, alpha=0.5)

Batched-leading-axis equivalent of amc_query_aware_saliency().

Parameters:

• x (array) – [G, N, D] token activations (for the magnitude term).

• keys (array) – [G, N, D] key projection vectors (for the semantic term).

• query (array) – [G, D] per-group embedded query/prompt vector.

• alpha (float) – Balance coefficient in [0, 1].

Return type:

array

Returns:

[G, N] saliency scores in [0, 1].

veloxquant_mlx.quantizers.amc.amc_assign_tiers(saliency, k_high=0.2, k_mid=0.3)

Assign each token a tier id via percentile-threshold partitioning.

Uses veloxquant_mlx.dsa.heap.MaxHeap to select the top ceil(k_high * N) and next ceil(k_mid * N) tokens by saliency in O(N log N) via heap pops — bounded selection rather than a full sorted() call, consistent with this repo’s DSA-first convention (the same heap already backs SortedChannelIndex elsewhere in the codebase).

Parameters:

• saliency (array) – [N] saliency scores.

• k_high (float) – Fraction of tokens routed to the High tier.

• k_mid (float) – Fraction of tokens routed to the Mid tier (after High).

Return type:

list[int]

Returns:

List of length N with tier ids (HIGH/MID/LOW).

veloxquant_mlx.quantizers.amc.amc_assign_tiers_batched(saliency, k_high=0.2, k_mid=0.3)

Batched-leading-axis equivalent of amc_assign_tiers().

Replaces the per-row .tolist() + MaxHeap top-k selection with one host-sync-free mx.argsort call across all G rows (e.g. flattened B*H) at once, using the same double-argsort-for-rank trick as veloxquant_mlx.quantizers.kitty.hi_mask_from_variance_batched().

On tie-free saliency (the realistic case — saliency is a continuous function of float activations, so exact ties are measure-zero) this is verified equivalent to looping amc_assign_tiers() per row. Exact ties can select a different index than the heap does: heap pop order on equal priorities is an unspecified artifact of MaxHeap’s array layout, not a documented tie-break guarantee, so this is not considered a behavior change.

Parameters:

• saliency (array) – [G, N] saliency scores.

• k_high (float) – Fraction of tokens routed to the High tier.

• k_mid (float) – Fraction of tokens routed to the Mid tier (after High).

Return type:

array

Returns:

[G, N] int32 tier ids (HIGH/MID/LOW).

veloxquant_mlx.quantizers.amc.amc_adaptive_thresholds(tau_high_base, tau_low_base, state, new_saliency_values, gamma=0.1)

Sequence-adaptive closed-loop threshold adjustment (Eq. 4-5).

Pushes new_saliency_values into the trailing window, computes the window’s variance, and depresses/raises tau_H/tau_L by a log-ratio against the calibration-time variance.

Parameters:

• tau_high_base (float) – Baseline High-tier threshold (from offline calibration).

• tau_low_base (float) – Baseline Mid-tier threshold.

• state (AMCThresholdState) – Current AMCThresholdState (mutated in place — the returned state is the same object, appended to).

• new_saliency_values (array) – [N] saliency scores from the current step, pushed into the trailing window before computing variance.

• gamma (float) – Attenuation scaling factor.

Return type:

tuple[float, float, AMCThresholdState]

Returns:

(tau_H, tau_L, state) — updated thresholds and the (mutated) state object.

veloxquant_mlx.quantizers.amc.amc_apply_rank_mask(x, rank)

Zero channels [rank:D) of a (calibration-ordered) activation vector.

Parameters:

• x (array) – [..., D] activations, already permuted by veloxquant_mlx.quantizers.amc_calibration.amc_calibrate_channel_order() so the surviving prefix is the highest-variance subspace. Any number of leading batch axes is supported (broadcasts against the last axis only).

• rank (int) – Number of leading channels to keep.

Return type:

array

Returns:

Same shape as x with columns rank:D zeroed.

veloxquant_mlx.quantizers.amc.amc_quantize_tier(x, bits, group_size=32)

Quantize-then-dequantize a tier’s activations to bits precision.

Reuses the repo’s shared asymmetric min/max group quantizer (veloxquant_mlx.quantizers._quant_utils._group_quant_dequant()) rather than hand-rolling Eq. 7’s fixed-point rounding — same simulated quantize/dequantize round-trip, same convention as every other method here.

Parameters:

• x (array) – [N, D] activations (already rank-masked).

• bits (int) – Target bit-width (16, 8, or 4 in AMC’s tier scheme, but any value the shared group quantizer accepts works).

• group_size (int) – Token-axis group size for min/max quantization.

Return type:

array

Returns:

[N, D] quantized-then-dequantized activations, in x.dtype.

veloxquant_mlx.quantizers.amc.amc_compress_tokens_batched(x, tiers, tier_configs)

Batched equivalent of per-token rank-mask + quantize, over [..., N, D].

Replaces AMCKVCache._compress_step’s for i in range(n) loop (nested inside a for b: for h: loop over heads) with one pass per tier (there are always exactly 3: HIGH/MID/LOW) over the whole [G, N, D] array (G = flattened B*H), selected per token with mlx.core.where(). Each token is quantized independently of every other token, with min/max taken over that token’s own kept channels (rank leading channels), so MID/LOW tiers are genuinely lossy.

Parameters:

• x (array) – [G, N, D] activations (fp16 or fp32), or [N, D] for a single group.

• tiers (array) – [G, N] (or [N]) int tier ids per token (HIGH/MID/LOW), as returned by amc_assign_tiers_batched().

• tier_configs (dict[int, AMCTierConfig]) – {tier_id: AMCTierConfig} for this head_dim (see _tier_config_for_dim()).

Return type:

array

Returns:

Same leading shape as x, fp16 rank-masked, quantized-then-dequantized activations.

veloxquant_mlx.quantizers.amc.amc_pack_low_tier(quantized_codes)

Densely pack Low-tier 4-bit integer codes via BitPackBuffer.

Parameters:

quantized_codes (array) – [N] uint8 codes in [0, 15] (already quantized to 4-bit integer indices, e.g. via a group-quant codes array before dequantization).

Return type:

tuple[bytes, int]

Returns:

(packed_bytes, n) where n is the original element count (needed to unpack).

veloxquant_mlx.quantizers.amc.amc_fp16_bytes(tier_counts, head_dim)

Compute actual stored bytes given per-tier token counts (K + V).

Parameters:

• tier_counts (dict) – Mapping {HIGH: n_high, MID: n_mid, LOW: n_low}.

• head_dim (int) – Full channel dimension D (used to scale tier ranks).

Return type:

int

Returns:

Total bytes for K + V combined, across all tiers.

veloxquant_mlx.quantizers.amc.full_amc_fp16_bytes(tokens_seen, head_dim)

Hypothetical fp16 full-rank K + V byte cost if AMC were never applied.

Return type:

int

Parameters:

• tokens_seen (int)

• head_dim (int)