Kakeya推理引擎：打破KV缓存存储瓶颈的推测解码新架构

Kakeya-LLM-Inference-engine通过DLM提议器与AR验证器的协作架构，结合sink+window缓存策略，实现最高5500倍的KV缓存压缩比，为大模型长上下文推理提供可行的内存优化方案。本文将从背景、架构、性能、局限等方面展开分析。

大语言模型(LLM)的推理过程中，键值(KV)缓存是支撑自回归生成的关键组件。随着上下文长度不断增加，KV缓存的内存占用呈线性增长，成为制约长上下文推理的主要瓶颈。传统Transformer架构中，每个token的KV表示都需要被存储，当序列长度达到百万级别时，缓存占用可达数十GB甚至更高。这种存储爆炸问题不仅影响单用户的推理体验，更限制了批处理规模和服务并发能力。业界一直在探索各种KV缓存压缩技术，包括量化、剪枝、滑动窗口等方法，但往往在压缩率和生成质量之间难以取得理想平衡。

核心架构设计

Kakeya推理引擎采用创新的双模型协作架构，由两个核心组件构成：

• DLM提议器（Diffusion Language Model Proposer）：基于Qwen3-0.6B的掩码扩散模型构建，以扩散方式生成候选token块，不维护持久化KV缓存，KV贡献为零。
• AR验证器（Autoregressive Verifier）：采用Qwen3-1.7B作为基础模型，实现SinkWindowVerifier机制，动态裁剪DynamicCache层的K/V张量，仅保留sink token和最近窗口内的KV表示。

Sink+Window缓存策略

该策略将KV缓存划分为两个区域：

• Sink区域：保留序列起始的几个token（默认4个），包含重要上下文信息。
• Window区域：维护固定大小的滑动窗口（默认24-64个token），仅保留最近的KV表示，新token生成时逐出最旧KV。 这种策略基于注意力模式的局部性，在保证生成质量的同时降低内存占用。

根据项目基准测试结果，Kakeya引擎展现显著内存优化效果：

等效性测试

当窗口大小覆盖整个序列时（sink=4, window=64），输出与基线贪婪解码完全比特一致，峰值KV占用3.06MB，相比完整缓存12.10MB实现3.86倍压缩。

压缩模式测试

批大小	序列长度	Net Bytes per Token	压缩比
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
长上下文场景（批大小64，序列长度1M）下，实现超5500倍压缩比，Net Bytes per Token仅20.8字节。

项目文档披露以下局限：

1. 验证器模型差异：使用Qwen3-1.7B替代不存在的Qwen3.6检查点，真实Qwen3.5/3.6基线压缩比会降低约1.75倍，但框架代码无需修改。
2. 接受率挑战：当前token接受率约0.12，因提议器与验证器表示几何未对齐，同家族Repr-Align提议器可提升至0.6-0.85，低接受率不影响正确性但降低吞吐量。
3. 激活内存优化空间：提议器激活内存由密集logits缓冲区主导，未采用“仅掩码位置计算logits”优化，需未来改进。

应用场景

• 长文档处理：法律合同分析、学术论文综述等数十万token输入场景。
• 多轮对话系统：保持长期上下文同时控制内存占用。
• 批处理服务：高并发场景提升服务吞吐量。

部署建议

• 根据业务负载调整sink和window大小，平衡内存与生成质量。
• 使用4-bit量化验证器（MLX后端）进一步降低内存，Qwen3-1.7B-4bit仅需约1GB常驻内存。
• 批处理规模(B)和序列长度(S)乘积需足够大，显现提议器权重摊销效应。

Kakeya引擎代表LLM推理优化的重要探索方向，展示通过架构创新解决内存瓶颈的可能性。其结合扩散生成与自回归验证、动态KV管理的路线，为高效推理提供新思路。尽管存在接受率偏低等局限，核心架构的合理性和可扩展性已被验证。该开源实现为社区改进提供基础，期待更多基于此架构的优化工作。