[Arxiv 2024] TriForce: Lossless Acceleration of Long Sequence Generation with Hierarchical Spec Dec

Introduction

• 作者提出 TriForce，针对长序列生成场景，draft model 仅使用部分 KV cache 进行推理，同时采用多级投机推理策略进一步降低 draft model 的推理开销

Method

• Retrieval-based Drafting. 作者针对长序列生成场景设计了 draft model，draft model 为原有的 LLM 通过检索部分重要 KV cache 进行推理；具体来说，作者将 KV cache 划分为若干 chunks，通过计算 query 和每个 chunk 的 average K cache 的 attn score 完成检索；这种方法能保存所有 tokens 的 KV cache，在精度上比 H2O 更有优势
  
• Hierarchical Speculation. 作者采用多级投机解码策略，retrieval-based draft model 进一步使用更小的 LLM + StreamingLLM cache 组成的 draft model 加速
  

Experiments

• End-to-end Results.
  
• Ablation Results.
  

References

• Sun, Hanshi, et al. “Triforce: Lossless acceleration of long sequence generation with hierarchical speculative decoding.” arXiv preprint arXiv:2404.11912 (2024).
• code: https://github.com/Infini-AI-Lab/TriForce