A-QJL：自适应分层KV缓存压缩技术助力长上下文大模型推理

A-QJL通过自适应分层分配压缩强度，在固定内存预算下显著优化长上下文LLM推理的质量与内存权衡，实现3比特每浮点数的KV缓存压缩。

随着大语言模型（LLM）应用场景的不断扩展，处理长上下文已成为刚需。从文档分析到代码理解，从多轮对话到知识检索，模型需要处理的序列长度从几千token迅速增长到数万甚至数十万token。然而，长上下文推理面临一个核心瓶颈：KV缓存（Key-Value Cache）的内存占用随序列长度线性增长。

在标准的Transformer解码过程中，模型需要缓存每一层每个注意力头的Key和Value向量，以避免重复计算。对于32层、4096维隐藏层、16个注意力头的模型，处理128K上下文时，KV缓存可能占用数十GB显存。这严重限制了长上下文能力的实际部署，尤其是在消费级GPU上。

2024年，Zandieh等人提出了QJL（Quantized Johnson-Lindenstrauss）方法，为KV缓存压缩开辟了新路径。QJL的核心思想是将Johnson-Lindenstrauss变换应用于Key嵌入，然后将结果量化为1比特（仅保留符号）。

Johnson-Lindenstrauss引理是计算几何中的经典结果，它表明高维空间中的点集可以被映射到低维空间，同时保持点间距离关系。QJL巧妙地将这一数学工具应用于注意力机制：通过对Key向量进行JL变换并量化，显著降低存储需求，同时使用非对称内积估计器来恢复注意力分数。

与传统量化方法不同，QJL采用"零开销"设计，避免了逐块量化常数的存储需求。这使得QJL在保持注意力分数相对失真最小化的同时，将KV缓存从16比特每浮点数压缩到仅3比特，实现了约5倍的内存节省。

虽然QJL已取得显著成果，但原方法对所有层采用固定的投影维度k。直观上，不同层对Key信息的敏感度可能存在差异：浅层可能更关注局部语法和词法特征，而深层则更关注语义和推理特征。

基于此观察，Bhargav Chirumamilla和Xinkai Shen提出了A-QJL（Adaptive QJL），引入自适应分层压缩策略。核心创新在于：不再对所有层使用统一的k值，而是将层分组并为每组分配不同的投影维度。

具体实现上，A-QJL将32层模型划分为多个组（如4组：0-7层、8-15层、16-23层、24-31层）。早期层获得较高的k值（较少压缩），后期层获得较低的k值（更多压缩），在固定内存预算下优化质量与压缩率的权衡。

A-QJL提供了两种分配策略：手工调优和自动敏感度分析。敏感度分析器通过评估每层Key向量的方差来确定各层对压缩的敏感度。

分析器在校准数据集（如qasper）上运行，计算每层的统计特征，输出层组边界和各组的量化比特配置。这种数据驱动的方法确保压缩资源的分配与模型的实际行为相匹配，而非依赖启发式规则。

使用示例：

python scripts/sensitivity_profiler.py --model_name "lmsys/longchat-7b-v1.5-32k" \
    --dataset_name qasper --n_calib 10 --num_groups 4 \
    --output config/aqjl_profiled.json

A-QJL在LongBench基准上进行了全面评估，测试模型包括Llama-2和Llama-3系列。实验结果表明：

内存效率：A-QJL成功将KV缓存压缩至3比特每浮点数，相比原始16比特表示节省约81%内存。这使得在相同硬件上可处理的上下文长度扩展5倍。

精度保持：在问答、摘要、代码理解等任务上，A-QJL在显著降低内存占用的同时，保持了与基线相当甚至略优的精度。自适应分层策略相比固定压缩策略展现出更优的质量-内存权衡。

推理速度：由于减少了内存带宽压力和缓存占用，A-QJL在长序列场景下展现出更快的推理速度。实验显示，随着序列长度增加，速度优势愈发明显。

兼容性：A-QJL支持多种精度格式（float16、bfloat16等），并兼容分组查询注意力（GQA）等现代架构特性，可无缝集成到Llama 2/3系列模型。

项目提供了完整的实验驱动框架，支持2组固定模式和3组以上自适应模式。用户可通过配置文件灵活调整层组边界和压缩强度：

{
  "layer_group_boundaries": "8,16,24",
  "key_quantization_bits_per_group": "512,384,256,192"
}

运行实验只需：

python scripts/aqjl_experiments.py --config config/aqjl_experiments.json

项目还包含可视化工具，自动生成性能对比图表，包括平均得分、峰值内存、每秒token数等关键指标。

A-QJL的实现充分考虑了GPU效率。核心操作通过轻量级CUDA内核加速，包括JL变换、量化和反量化。代码结构模块化，便于集成到现有推理框架。

安装流程简洁：

git clone https://github.com/Chirumamilla1522/AQJl-Memory-Efficient-Long-Context-LLM-Inference-via-QJL-and-Adaptive-Layer-Wise-Compression.git
cd AQJl-Memory-Efficient-Long-Context-LLM-Inference-via-QJL-and-Adaptive-Layer-Wise-Compression
pip install -r requirements.txt
cd qjl_kernel
python setup.py build_ext --inplace