SPB Git

spb/wp30_uqo

Public
6.2 KB

# WP30 method implementation notes (one entry per named method)

Fidelity labels: official (upstream code/checkpoint run as released), faithful_reimplementation (algorithm re-implemented from the paper/code with the same rule and defaults), adapted (algorithm changed to fit our setting; differences listed), inspired (shares the idea, not the algorithm), custom (WP30 baseline with no published counterpart claimed).

# A4 abstractive summary — custom

Sender = same checkpoint as receiver (or the configured sender), greedy decoding, prompt SUMMARY_SYSTEM/summary_user (src/wp30/prompts.py), hard output cap = budget tokens. Question-blind variant never sees the question; question-aware appends it. Measured: sender input/output tokens, generation time, UTF-8 bytes of the summary; receiver tokens of the summary under its own tokenizer.

# A5 structured evidence — custom

Same sender/budget as A4; output = JSON {facts:[{entity, attribute, value, quote}], notes?, unknowns?}. Each quote is resolved to a character span of the source (doc.find); unresolved quotes are counted (unresolved_quotes). The receiver sees the JSON with spans (not the quotes) — values are transmitted, spans are provenance. Truncated JSON is repaired at the last complete object.

# A6 references + metered retrieval — custom

Message = reference index {ref, entity, attribute, span} derived from A5 with values removed. Receiver side: BM25 of the question over the labels, fetch up to max_lookups spans (±40/160 chars) under max_fetch_bytes; every fetched byte and lookup is recorded.

# A7 lossless transport codec — custom (zstd level 3 / gzip 6)

Compress UTF-8 text for transport, decompress at the receiver; receiver prefill unchanged. Measured codec times and bytes.

# A8 LLMLingua-2 — official

llmlingua 0.2.2, PromptCompressor(model_name="microsoft/llmlingua-2-xlm-roberta-large-meetingbank", use_llmlingua2=True, device_map="mps"), compress_prompt(doc, rate=budget/source_tokens, force_tokens=["\n",".",",",":","-"], drop_consecutive=True). Token classification is question-agnostic (question-blind by construction). The upstream paper evaluated MeetingBank/LongBench/ZeroScrolls/GSM8K/BBH with GPT-3.5/Mistral-7B receivers at 2x–5x rates; our receivers are 3B models. Backend: torch 2.15 on MPS (fp32); timing reported separately because it is a different accelerator path than MLX.

# B2 packed KV quantization — custom (mlx_lm primitives)

mx.quantize(K|V, group_size=g, bits=b) along the head dimension (per token, per group of g channels), asymmetric affine (scale + bias per group), packed into uint32 words; consumed by quantized_scaled_dot_product_attention (quantized execution, no dequantized dense copy). Bits ∈ {8, 4, 2}; group 64 (32 for 2-bit). Differences vs KIVI: KIVI quantizes keys per channel (over tokens) and values per token, with a full-precision residual window (R = 128) and CUDA kernels; our scheme is the mlx_lm layout (per-token for both). The KIVI motivation (outlier key channels) predicts exactly the Qwen2.5 failure we observe at 4 bits.

# B3 sink + recent retention — inspired (StreamingLLM)

Keep the first sinks=4 positions plus the most recent positions up to the budget. StreamingLLM re-assigns positions within the cache (keys stored un-rotated, rotated at use with arange(kv_len)); WP30 keeps original positions (keys are stored post-RoPE in mlx_lm) and preserves the logical offset, so new tokens see a positional gap. This is a retention pattern test, not a StreamingLLM reproduction (and StreamingLLM makes no claim about retrieving evicted facts).

# B4 random retention — custom control

Uniformly random positions (seeded) at the same budget; no sinks.

# B5 attention-informed selection — faithful_reimplementation (SnapKV rule)

Observation window = last W prompt tokens (question-aware: the fresh query tokens; question-blind: a generic 32-token instruction); per kv-head attention of the window queries onto prefix keys (softmax, averaged over window queries and over the query heads of the GQA group), 1-D max-pooling with kernel 7 along positions, top-k per kv-head; the window itself is always retained (window_keep=32). Differences: SnapKV runs inside the prefill kernel (prompt phase) with window 32/64 and kernel 5–13 depending on config; we compute the window queries in a separate forward pass over a copy of the cache (cost recorded as compress_ms), so the main attention kernel remains fused. Per-head retention produces per-head index sets (the SparseKVCache gathers per head).

# B6 layer-dependent (pyramid) budgets — faithful_reimplementation of the PyramidKV budget rule

Budget per layer from KVCache-Factory: base = cap − window; min = base // β (β = 20); max = 2·base − min; step = (max − min)/(L − 1); budget_l = max − l·step (+ window), rescaled to preserve the total budget exactly. Selection within a layer = B5 scoring. Because retained lengths differ across layers, multi-token query prefill is executed token by token (mask-free); recorded as sequential_query_prefill.

# B9 selection + quantization — custom composition

B5 (or B3) selection first, then B2 packing of the retained tensors (order: select → quantize).

# C2 linear KV mapper — custom (not SwaT)

See src/wp30/mapper.py: per receiver layer, ridge regression from the aligned sender layer's un-rotated [K, V] (512 + bias) to the receiver's un-rotated [K, V] (512); fitted on train-split caches only; re-rotated with the receiver's RoPE. Reconstruction-only loss. SwaT trains a ~300M-parameter mapper with task/distillation terms on Granite→SmolLM; no checkpoint or licence is public (audit_group1).

# C0 identity transfer — negative control

Sender K/V fed as receiver cache (shapes coincide: 2 × 128). Diagnostic of "shape match ≠ compatibility".

# Methods audited but not executed

H2O (needs accumulated attention during prefill — would disable the fused kernel for the whole prefill), Attention Matching compaction (official code pins vLLM 0.11/CUDA; a faithful MLX re-implementation of the optimization was out of scope — see coverage table), C2C (official fuser checkpoints exist only for Qwen3 receivers, CUDA; receiver still prefills the source), LatentMAS (CUDA/vLLM, same-architecture agents), SwaT (no checkpoint, no licence).

Files · 5c60942332e88759c6021427f923264ff96c57e3
Menu

Theme