SPB Git

spb/wp30_uqo

Public
6.2 KB · 76 lines markdown
1# WP30 method implementation notes (one entry per named method)23Fidelity labels: **official** (upstream code/checkpoint run as released), **faithful_reimplementation** (algorithm re-implemented4from the paper/code with the same rule and defaults), **adapted** (algorithm changed to fit our setting; differences listed),5**inspired** (shares the idea, not the algorithm), **custom** (WP30 baseline with no published counterpart claimed).67## A4 abstractive summary — custom8Sender = same checkpoint as receiver (or the configured sender), greedy decoding, prompt `SUMMARY_SYSTEM`/`summary_user`9(`src/wp30/prompts.py`), hard output cap = budget tokens. Question-blind variant never sees the question; question-aware appends it.10Measured: sender input/output tokens, generation time, UTF-8 bytes of the summary; receiver tokens of the summary under its own tokenizer.1112## A5 structured evidence — custom13Same sender/budget as A4; output = JSON {facts:[{entity, attribute, value, quote}], notes?, unknowns?}. Each quote is resolved to a14character span of the source (`doc.find`); unresolved quotes are counted (`unresolved_quotes`). The receiver sees the JSON with spans15(not the quotes) — values are transmitted, spans are provenance. Truncated JSON is repaired at the last complete object.1617## A6 references + metered retrieval — custom18Message = reference index {ref, entity, attribute, span} derived from A5 with values removed. Receiver side: BM25 of the question over19the labels, fetch up to `max_lookups` spans (±40/160 chars) under `max_fetch_bytes`; every fetched byte and lookup is recorded.2021## A7 lossless transport codec — custom (zstd level 3 / gzip 6)22Compress UTF-8 text for transport, decompress at the receiver; receiver prefill unchanged. Measured codec times and bytes.2324## A8 LLMLingua-2 — official25`llmlingua` 0.2.2, `PromptCompressor(model_name="microsoft/llmlingua-2-xlm-roberta-large-meetingbank", use_llmlingua2=True, device_map="mps")`,26`compress_prompt(doc, rate=budget/source_tokens, force_tokens=["\n",".",",",":","-"], drop_consecutive=True)`. Token classification is27question-agnostic (question-blind by construction). The upstream paper evaluated MeetingBank/LongBench/ZeroScrolls/GSM8K/BBH with28GPT-3.5/Mistral-7B receivers at 2x–5x rates; our receivers are 3B models. Backend: torch 2.15 on MPS (fp32); timing reported separately29because it is a different accelerator path than MLX.3031## B2 packed KV quantization — custom (mlx_lm primitives)32`mx.quantize(K|V, group_size=g, bits=b)` along the head dimension (per token, per group of g channels), asymmetric affine (scale + bias per33group), packed into uint32 words; consumed by `quantized_scaled_dot_product_attention` (quantized execution, no dequantized dense copy).34Bits ∈ {8, 4, 2}; group 64 (32 for 2-bit). Differences vs KIVI: KIVI quantizes keys per *channel* (over tokens) and values per token, with a35full-precision residual window (R = 128) and CUDA kernels; our scheme is the mlx_lm layout (per-token for both). The KIVI motivation36(outlier key channels) predicts exactly the Qwen2.5 failure we observe at 4 bits.3738## B3 sink + recent retention — inspired (StreamingLLM)39Keep the first `sinks=4` positions plus the most recent positions up to the budget. StreamingLLM re-assigns positions *within the cache*40(keys stored un-rotated, rotated at use with arange(kv_len)); WP30 keeps original positions (keys are stored post-RoPE in mlx_lm) and41preserves the logical offset, so new tokens see a positional gap. This is a retention *pattern* test, not a StreamingLLM reproduction42(and StreamingLLM makes no claim about retrieving evicted facts).4344## B4 random retention — custom control45Uniformly random positions (seeded) at the same budget; no sinks.4647## B5 attention-informed selection — faithful_reimplementation (SnapKV rule)48Observation window = last W prompt tokens (question-aware: the fresh query tokens; question-blind: a generic 32-token instruction);49per kv-head attention of the window queries onto prefix keys (softmax, averaged over window queries and over the query heads of the50GQA group), 1-D max-pooling with kernel 7 along positions, top-k per kv-head; the window itself is always retained (`window_keep=32`).51Differences: SnapKV runs inside the prefill kernel (prompt phase) with window 32/64 and kernel 5–13 depending on config; we compute the52window queries in a separate forward pass over a *copy* of the cache (cost recorded as `compress_ms`), so the main attention kernel53remains fused. Per-head retention produces per-head index sets (the SparseKVCache gathers per head).5455## B6 layer-dependent (pyramid) budgets — faithful_reimplementation of the PyramidKV budget rule56Budget per layer from KVCache-Factory: base = cap − window; min = base // β (β = 20); max = 2·base − min; step = (max − min)/(L − 1);57budget_l = max − l·step (+ window), rescaled to preserve the total budget exactly. Selection within a layer = B5 scoring. Because retained58lengths differ across layers, multi-token query prefill is executed token by token (mask-free); recorded as `sequential_query_prefill`.5960## B9 selection + quantization — custom composition61B5 (or B3) selection first, then B2 packing of the retained tensors (order: select → quantize).6263## C2 linear KV mapper — custom (not SwaT)64See `src/wp30/mapper.py`: per receiver layer, ridge regression from the aligned sender layer's un-rotated [K, V] (512 + bias) to the65receiver's un-rotated [K, V] (512); fitted on train-split caches only; re-rotated with the receiver's RoPE. Reconstruction-only loss.66SwaT trains a ~300M-parameter mapper with task/distillation terms on Granite→SmolLM; no checkpoint or licence is public (audit_group1).6768## C0 identity transfer — negative control69Sender K/V fed as receiver cache (shapes coincide: 2 × 128). Diagnostic of "shape match ≠ compatibility".7071## Methods audited but not executed72H2O (needs accumulated attention during prefill — would disable the fused kernel for the whole prefill), Attention Matching compaction73(official code pins vLLM 0.11/CUDA; a faithful MLX re-implementation of the optimization was out of scope — see coverage table),74C2C (official fuser checkpoints exist only for Qwen3 receivers, CUDA; receiver still prefills the source), LatentMAS (CUDA/vLLM, same-architecture75agents), SwaT (no checkpoint, no licence).76
Files · d5e182fa7d7d5118aa92e61c47c40ed848da0ecf
Menu

Theme