Kakeya Inference Engine: A New Speculative Decoding Architecture Breaking Through KV Cache Storage Bottlenecks

Kakeya-LLM-Inference-engine uses a collaborative architecture of DLM Proposer and AR Verifier, combined with a sink+window cache strategy, to achieve a maximum KV cache compression ratio of 5500x, providing a feasible memory optimization solution for long-context inference of large models. This article will analyze it from aspects such as background, architecture, performance, and limitations.

In the inference process of Large Language Models (LLMs), the Key-Value (KV) cache is a critical component supporting autoregressive generation. As the context length increases, the memory usage of KV cache grows linearly, becoming the main bottleneck restricting long-context inference. In traditional Transformer architectures, the KV representation of each token needs to be stored; when the sequence length reaches the million level, the cache usage can reach tens of GB or even higher. This storage explosion problem not only affects the inference experience of single users but also limits batch processing scale and service concurrency. The industry has been exploring various KV cache compression techniques, including quantization, pruning, sliding window, etc., but it is often difficult to achieve an ideal balance between compression ratio and generation quality.

Core Architecture Design

Kakeya Inference Engine uses an innovative dual-model collaborative architecture, consisting of two core components:

• DLM Proposer (Diffusion Language Model Proposer)：Built based on the Qwen3-0.6B masked diffusion model, it generates candidate token blocks in a diffusion manner, does not maintain persistent KV cache, and contributes zero to KV usage.
• AR Verifier (Autoregressive Verifier)：Uses Qwen3-1.7B as the base model, implements the SinkWindowVerifier mechanism, dynamically prunes the K/V tensors of the DynamicCache layer, and only retains the KV representations of sink tokens and the most recent window.

Sink+Window Cache Strategy

This strategy divides the KV cache into two regions:

• Sink Region: Retains the first few tokens of the sequence (default 4), which contain important context information.
• Window Region: Maintains a fixed-size sliding window (default 24-64 tokens), only retains the most recent KV representations, and evicts the oldest KV when a new token is generated. This strategy is based on the locality of attention patterns, reducing memory usage while ensuring generation quality.

According to the project's benchmark test results, the Kakeya engine shows significant memory optimization effects:

Equivalence Test

When the window size covers the entire sequence (sink=4, window=64), the output is completely bit-consistent with the baseline greedy decoding, with a peak KV usage of 3.06MB, achieving a 3.86x compression compared to the full cache of 12.10MB.

Compression Mode Test

Batch Size	Sequence Length	Net Bytes per Token	Compression Ratio
8	8,192	18,582	6.17x
8	32,768	4,645.5	24.69x
8	131,072	1,161.4	98.75x
64	131,072	166.6	688.36x
64	1,048,576	20.8	5506.92x
In long-context scenarios (batch size 64, sequence length 1M), it achieves an over 5500x compression ratio, with Net Bytes per Token being only 20.8 bytes.

The project documentation discloses the following limitations:

1. Verifier Model Difference: Qwen3-1.7B is used instead of the non-existent Qwen3.6 checkpoint; the real Qwen3.5/3.6 baseline compression ratio will decrease by approximately 1.75x, but the framework code does not need modification.
2. Acceptance Rate Challenge: The current token acceptance rate is about 0.12 due to misalignment between the proposer and verifier's representation geometry. The same-family Repr-Align proposer can increase it to 0.6-0.85. A low acceptance rate does not affect correctness but reduces throughput.
3. Activation Memory Optimization Space: The proposer's activation memory is dominated by the dense logits buffer, and the 'compute logits only for masked positions' optimization has not been adopted, which needs future improvement.

Application Scenarios

• Long Document Processing: Scenarios with hundreds of thousands of token inputs such as legal contract analysis and academic paper review.
• Multi-Turn Dialogue Systems: Maintain long context while controlling memory usage.
• Batch Processing Services: Improve service throughput in high-concurrency scenarios.

Deployment Recommendations

• Adjust the sink and window sizes according to business load to balance memory and generation quality.
• Use a 4-bit quantized verifier (MLX backend) to further reduce memory; Qwen3-1.7B-4bit only requires about 1GB of resident memory.
• The product of batch size (B) and sequence length (S) should be large enough to show the proposer weight amortization effect.

The Kakeya engine represents an important exploration direction for LLM inference optimization, demonstrating the possibility of solving memory bottlenecks through architectural innovation. Its approach combining diffusion generation and autoregressive verification, along with dynamic KV management, provides new ideas for efficient inference. Although there are limitations such as low acceptance rate, the rationality and scalability of the core architecture have been verified. This open-source implementation provides a foundation for community improvements, and we look forward to more optimization work based on this architecture.