通过合并多个词元,实现更高效的大模型推理。
HAMburger: Accelerating LLM Inference via Token Smashing
- 将多个词元合并为一个缓存,减少每步计算和存储开销。
- KV缓存计算减少最多2倍,每秒生成速度提升至2倍。
- 适合对推理效率要求高、需兼顾长文本质量的场景。
大语言模型推理对效率的需求日益增长,亟需在算法、系统和硬件层面进行整体优化。然而,现有方法极少改变生成模式:每个词元需一次前向传播和一次键值缓存。我们发现,大模型能精准判断单个键值缓存应存储的信息量,许多词元可无需全局上下文自信生成。基于此洞察,我们提出 HAMburger,一种分层自回归模型,通过在基础模型间嵌入组合嵌入器与微步解码器,将多个词元压缩为单一键值缓存,并每步生成多个词元。同时,它作为推测解码框架,可无条件信任自动生成的词元。结果表明,哈曼堡以亚线性方式增长键值缓存与前向浮点运算量,其推理速度随查询困惑度与输出结构动态调整。大量实验显示,该方法使键值缓存计算最多减少2倍,每秒生成速度最高提升2倍,且在短文本与长文本任务中均保持高质量。本方法探索了需兼顾计算与内存效率的极具挑战性的推理范式,具有硬件无关设计特性。
原文摘要 · Abstract (English)
The growing demand for efficient Large Language Model (LLM) inference requires a holistic optimization on algorithms, systems, and hardware. However, very few works have fundamentally changed the generation pattern: each token needs one forward pass and one KV cache. This can be sub-optimal because we found that LLMs are extremely capable of self-identifying the exact dose of information that a single KV cache can store, and many tokens can be generated confidently without global context. Based on this insight, we introduce HAMburger, a Hierarchically Auto-regressive Model that redefines resource allocation in LLMs by moving beyond uniform computation and storage per token during inference. Stacking a compositional embedder and a micro-step decoder in between a base LLM, HAMburger smashes multiple tokens into a single KV and generates several tokens per step. Additionally, HAMburger functions as a speculative decoding framework where it can blindly trust self-drafted tokens. As a result, HAMburger shifts the growth of KV cache and forward FLOPs from linear to sub-linear with respect to output length, and adjusts its inference speed based on query perplexity and output structure. Extensive evaluations show that HAMburger reduces the KV cache computation by up to 2$\times$ and achieves up to 2$\times$ TPS, while maintaining quality in both short- and long-context tasks. Our method explores an extremely challenging inference regime that requires both computation- and memory-efficiency with a hardware-agnostic design.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。