← All publications

Testing DeepSeek-inspired KV sharing in Granite 4

A frozen-weight experiment inspired by DeepSeek V4.1 cuts KV-cache storage in half, but severely degrades language quality. We publish the model, method, and measurements.

Experimental model and code on Hugging Face

Can cache sharing transfer without changing the weights?

A language model’s key-value cache stores attention states from earlier tokens so they can be reused during generation. As context grows, storing a separate cache for every attention layer becomes increasingly expensive. Sharing those states across layers offers a direct way to reduce that cost.

DeepSeek V4.1’s cache architecture prompted this experiment. Its Compressed Sparse Attention 2 design shares cached states and, in some layers, sparse-attention indices. Its Causal Encoder-Decoder architecture also lets decoder layers derive global keys and values from the encoder’s final states. These mechanisms operate within a model designed and trained around them.

We tested a narrower question on IBM Granite 4.0: can adjacent layers reuse keys and values while every pretrained parameter stays frozen? Our implementation follows the simpler cross-layer attention pattern described by Brandon and colleagues. Their work trained models with sharing from the start; our experiment tests a conversion after training. It does not implement DeepSeek’s complete architecture or claim a new attention mechanism.

YOCO, by Sun and colleagues, provides related architectural background: a self-decoder produces a global KV cache that a cross-decoder reuses. That broader restructuring differs from the adjacent-layer intervention tested here. We cite it as context for sharing across a model, rather than as an architecture reproduced in this experiment.

Forty layers, twenty caches

The starting checkpoint is ibm-granite/granite-4.0-1b, the all-attention variant with approximately 1.63 billion parameters. We pair its 40 layers: layers 0 and 1 use layer 0’s keys and values, layers 2 and 3 use layer 2’s, and so on. All 40 layers still execute, but only 20 produce and store a KV cache.

Each consuming layer retains its own query projection, output projection, and feed-forward computation. Its key and value projections are bypassed during inference, with their parameters preserved in the checkpoint. The modified forward pass reuses the producer’s keys after positional encoding, along with its values.

The weight file is copied byte-for-byte from the original release. Only the configuration and inference code change. There is no training, fine-tuning, calibration, or weight quantization. The checkpoint’s size is unchanged. We switch the same loaded model between native attention and paired sharing to compare the two conditions.

Language quality collapses

The cache reduction comes with severe damage to the model’s predictions. WikiText-2 subset perplexity rises from 19.55 to 707,495.20. Perplexity measures how well the model predicts the evaluation text; lower is better. HellaSwag accuracy also falls under both scoring methods.

The released chat examples show the same failure qualitatively. The native model answers a simple arithmetic prompt correctly, while the paired model emits repetitive, incoherent text. Five examples illustrate the behavior; they do not establish an instruction-following benchmark. This release is a negative research result and is not suitable as an improved Granite model.

Quality diagnostics on fixed evaluation subsets
DiagnosticNativePaired
WikiText-2 perplexity ↓19.55707,495.20
HellaSwag accuracy, token normalized ↑62.50%31.25%
HellaSwag accuracy, raw ↑40.62%25.00%

WikiText-2: 8,176 next-token predictions. HellaSwag: 32 validation examples. These are custom subset diagnostics, not standard full-benchmark scores; the HellaSwag sample is too small for fine-grained model ranking.

KV storage falls by half

Sharing within every pair halves the number of cached layers, and measured KV storage falls by exactly 50% at each tested context length. At 4,096 tokens, the native cache occupies 320 MiB and the paired cache occupies 160 MiB.

These numbers count unique allocated key and value tensor storage immediately after prefill. They exclude model weights, other activations, and runtime overhead. A 50% reduction in KV storage therefore does not mean a 50% reduction in total memory use. We did not measure peak process or device memory.

Allocated KV tensors after prefill
Context tokensNativePaired
51240 MiB20 MiB
2,048160 MiB80 MiB
4,096320 MiB160 MiB

Batch size 1, BF16. MiB denotes 1,048,576 bytes. Contexts beyond 4,096 tokens were not tested.

Decode throughput

Median decode throughput improves in these local runs, especially at the longest tested context. The paired model reaches 5.85 tokens per second at 4,096 tokens, compared with 3.75 for native attention. Prefill changes little: the native and paired medians are 0.914 and 0.884 seconds at 512 tokens, 4.062 and 4.043 at 2,048, and 9.751 and 9.534 at 4,096.

Decoding uses the same 32 teacher-forced tokens in both conditions: each step receives a fixed evaluation token instead of feeding back the model’s generated answer. This isolates execution speed from differences in generated text or stopping behavior. It does not measure the speed of completing a useful task, and the quality collapse prevents treating these timings as an improvement in practical model performance.

Decode throughput, median of three measured runs
Context tokensNative tok/sPaired tok/sChange
51218.5819.64+5.7%
2,0489.359.67+3.4%
4,0963.755.85+55.8%

Apple M1 Pro, MPS, batch size 1. Timings are local measurements, not a prediction of CUDA or production-serving performance.

How we measured it

We ran on an Apple M1 Pro with 10 CPU cores, 14 GPU cores, and 32 GiB unified memory, using macOS 26.3.1, Python 3.12.13, PyTorch 2.8.0, and Transformers 4.57.1. Inference used MPS, BF16, scaled dot-product attention, and batch size 1.

Each condition and context length received one full warm-up before three measured repeats, with condition order alternating between repeats. We synchronized device work at timing boundaries and report medians. Prefill excludes model loading, tokenization, sampling, and network overhead, so it is not end-to-end latency or time to first token.

For WikiText-2, we used the first 16 non-overlapping 512-token windows of the raw test split, with no special tokens and a context reset for every window. This scores 511 predictions per window. For HellaSwag, we sampled 32 validation examples without replacement using seed 42 and scored continuations zero-shot, both as summed log likelihood and normalized by continuation-token count. Contexts and endings were tokenized separately, without evaluation-harness preprocessing.

The published JSON includes dataset revisions, selected HellaSwag row IDs, a hash of the WikiText token sequence, per-example scores, and all timing samples. This experiment does not evaluate multilingual quality, safety behavior, contexts beyond 4,096 tokens, or production serving frameworks.

Interpretation and next experiment

Under this conversion, adjacent layers’ learned key and value representations are not interchangeable. Replacing them with an earlier layer’s states meets the storage objective while breaking the model’s predictive behavior. Our interpretation is that architectural reuse requires more than connecting an existing query projection to a different layer’s cache. The measurements alone do not isolate which layers or representation mismatches contribute most to the failure.

The result does not establish that cache sharing is ineffective. It establishes that pairing every layer in this frozen Granite checkpoint fails to preserve quality. A useful next experiment would change one pair at a time, measure the resulting error, and test whether selective sharing can retain quality at a smaller memory saving. That work has not yet been evaluated. Adapting the complete DeepSeek V4.1 design would be a separate study.

Model and reproduction

The Hugging Face release contains the unchanged base weights, custom configuration and inference code, evaluation scripts, pinned environment, tests, integrity checks, and raw measurements. Version v0.1.0 is the artifact described here. The links in this note pin its exact commit so the evidence remains stable as later work develops.

The release validation records an exact match to native-model logits when sharing is disabled, and 20 populated cache layers when it is enabled. Loading the experimental architecture uses Transformers 4.57.1 with trust_remote_code=True; inspect the included code before running it. Compatibility with Ollama, llama.cpp, vLLM, and hosted inference has not been established. The model card includes loading and reproduction instructions, Apache-2.0 licensing, and IBM attribution.

References and artifacts