mSeek: learned KV-cache compression in a 3B model

mSeek is a research implementation of Lightweight Block-Compressed Attention (LBCA). It tests whether the learned block-summary principle used by DeepSeek-V4's Heavily Compressed Attention can transfer to an already pretrained, dense 3B model under a modest single-GPU adaptation budget.

LBCA replaces old token-level key/value states with one learned pooled entry per block while retaining an exact local window. The evaluated model grafts LBCA into 14 of Llama-3.2-3B's 28 layers; the remaining layers stay dense.

Headline result: the transfer was useful for storage and selected direct retrieval, but not for execution speed or general compositional reasoning. LBCA nearly halved the analytic KV cache and retained 107/108 primary needle-retrieval cases, yet it was 2.0–2.4× slower than dense attention and underperformed a trained equal-memory sliding-window baseline on multi-hop retrieval.

This is an HCA-inspired retrofit, not a reproduction of DeepSeek-V4 HCA or its complete hybrid attention system.

Results

The primary quality matrix covers 4K, 16K, 32K and 64K contexts. A smaller nine-trial NIAH extension was run at 128K and is kept separate from the primary matrix.

Result	Dense	LBCA	SWA zero-shot	SWA trained
Primary NIAH successes	108/108	107/108	4/108	101/108
Primary multi-hop successes	72/72	58/72	0/72	69/72
NIAH accuracy at 32K	1.000	0.963	0.000	0.963
Multi-hop accuracy at 32K	1.000	0.500	0.000	0.917
Analytic KV cache at 128K	14.000 GiB	7.116 GiB	7.116 GiB	7.116 GiB

Additional system-level findings:

• The evaluated half-layer LBCA retrofit reduced analytic KV memory by 45% at 4K and approximately 49% from 16K onward.
• Within an approximately 16.5 GiB KV budget, analytic capacity increased from 154,624 dense tokens to 304,343 LBCA tokens (1.97×).
• Semantic quality was evaluated only through the model's native 128K range; the 304K figure is a storage estimate, not a quality claim.
• Portable CUDA LBCA decoding took 0.073–0.076 seconds per token, compared with 0.031–0.036 seconds for dense attention in the preserved measurements.
• Trained SWA used the same analytic KV budget and achieved the stronger multi-hop result. This prevents a general claim that learned summaries beat an adapted recency-only cache.

What the result supports

The experiment supports a bounded conclusion: a small-model retrofit can keep one globally accessible learned representation of each old region, nearly halve persistent KV state, and preserve the selected salient-fact retrieval probe.

It does not establish that:

• LBCA reproduces DeepSeek-V4 HCA or isolates HCA's contribution to DeepSeek-V4;
• pooled summaries are lossless long-context memory;
• LBCA is faster than dense attention in its current portable implementation;
• LBCA is superior to a trained equal-memory eviction policy;
• performance at 128K implies useful operation at the analytic 304K capacity; or
• one training trajectory generalises across seeds, models or workloads.

Experiment at a glance

Component	Evaluated setting
Base model	meta-llama/Llama-3.2-3B
Modified layers	14 even-numbered layers; 14 layers remain dense
Local raw window	128 tokens
Block schedule	16 tokens initially; 64 tokens from the 16K stage
LBCA adaptation	4,000 continued-pretraining steps
Trained SWA control	3,500 continued-pretraining steps
Other controls	Dense attention and zero-shot SWA
Quality probes	Synthetic NIAH and multi-hop exact-match evaluation
Context range	4K–64K primary matrix; separate 128K NIAH extension
Execution environment	Single-GPU CUDA campaign

The trained arms have unequal nominal exposure and different trainable structures. Their comparison is therefore descriptive, not an architecture-only causal estimate.

How LBCA works

For each modified attention layer:

1. Project hidden states to keys and values as usual.
2. Score the tokens in each completed block.
3. Softmax the scores within the block and pool its keys and values into one learned entry.
4. Attend over every completed historical summary plus the uncompressed local window.
5. Keep incomplete blocks raw so a query can never see a summary containing future tokens.

The core validated attention paths include:

• legacy: a clear reference implementation;
• chunked: a bounded-memory implementation numerically checked against the reference in float64.

A FlexAttention CUDA path is also available for compatible environments; the chunked path remains the portable CPU/MPS fallback.

The cache, masking and attention assembly are tested for strict causality, dense-limit behaviour, cache accounting and exact training resumption.

Quick start

The default test suite is CPU-only and excludes tests requiring real 3B weights:

python3 -m venv .venv
source .venv/bin/activate
python3 -m pip install -r requirements.txt
python3 -m pytest

Run the synthetic end-to-end smoke pipeline:

python3 scripts/run_experiments.py --smoke --out runs/smoke

The smoke run uses a tiny randomly initialised model. Its accuracy is not a research result; it checks that training, evaluation, cache accounting and report generation work end to end.

Full-model training or evaluation requires a CUDA GPU, access to the upstream Llama model under its applicable licence, and the relevant checkpoint:

python3 scripts/run_experiments.py \
  --device cuda \
  --out runs/experiment \
  --run-label experiment

Store Hugging Face or GitHub credentials in the host's secret manager or environment. Never place access tokens inside a notebook, repository, result folder or shared model cache.

Repository layout

lbca/
  config.py             configuration dataclasses
  lbca_attention.py     LBCA module and cache
  graft.py              attention-module replacement and LoRA setup
  schedule.py           warm-up, block-size and context-length schedules
  train.py              continued pretraining and exact resume
  generate.py           generation through the LBCA cache
  experiments.py        experiment orchestration
  data/                 corpus loading and synthetic evaluation data
  eval/                 NIAH, multi-hop, efficiency and frontier analysis

scripts/
  run_experiments.py    smoke and full experiment entry point
  run_ablations.py      proxy-ablation runner
  merge_runs.py         merge compatible result records

tests/                  CPU, GPU and real-model acceptance tests
docs/                   mechanism guides and selected final figures
architecture.md         original build specification

Validation

The final local audit recorded 178 passed, 1 skipped and 2 deselected tests under the repository's default marker configuration.

Important numerical distinctions:

• the two dense-attention oracle tests use atol=1e-4, rtol=1e-4;
• the separate chunked-versus-reference equivalence grid uses float64 with atol=1e-12, rtol=1e-12.

The default suite omits real_model tests. Run those only where the authorised base weights and a suitable GPU are available:

python3 -m pytest -m real_model

Evidence and data boundary

The read-only run archive contains the large checkpoints and preserved evaluation outputs. It does not contain the upstream Llama weights, training-corpus caches or authentication credentials.

The repository contains selected final figures and the implementation needed to inspect the mechanism. Development diagnostics and proxy-ablation outputs are deliberately excluded because they are not the selected dissertation evidence.

The final-evaluation notebook records the revision used for the corrected LBCA evaluation and trained-SWA campaign. The earlier 4,000-step LBCA checkpoint does not itself embed the source revision used for its original training, so that revision is not assigned retrospectively.

Limitations

• One base model and one completed training trajectory per trained arm.
• Synthetic retrieval tasks rather than a broad real-document benchmark.
• Unequal training exposure between LBCA and trained SWA.
• Fourteen untouched dense layers remain in every evaluated arm.
• Single-sequence CUDA timing without a latency distribution.
• Analytic commodity-device capacity rather than a completed MPS benchmark.
• The portable implementation prioritises correctness and bounded memory over kernel-level optimisation.

Responsible use

This artefact evaluates memory representation and long-context retrieval. It does not establish that generated answers are reliable for legal, financial, insurance or other high-impact decisions. Any deployment should separately evaluate source attribution, harmful omission, abstention, privacy, access control and human review.

The shared evaluation text is synthetic. No participant data or private organisational documents are required to run the included tests.

Licence

No open-source licence has yet been granted for this repository. Until a licence file is added, normal copyright restrictions apply. The upstream Llama model and datasets remain subject to their own licences and terms.