论文

H-Spec:不维护独立草稿 KV 的投机解码

H-Spec: Parallel Speculative Decoding Without a Drafter-Side KV Cache

模型推理投机采样

摘要

推测性解码通过让轻量级草稿模型预测未来标记以供目标模型验证,无损地加速大型语言模型推理。最近的区块扩散起草者通过并行预测多个词元进一步减少了起草延迟。然而,现有的块绘图器将每个输入位置的目标隐藏状态投影到单独的绘图器端 KV 缓存中,从而导致每个请求的内存和 KV 写入开销随着并发性而增长;直接重用目标 KV 会删除此缓存,但无法维持整个块的草稿质量。我们提出了一种混合目标上下文注入方法,该方法仅在最后一个输入位置使用目标隐藏状态来补充直接目标 KV 重用,不需要单独的绘图者端 KV 缓存。在此设计的基础上,我们提出了 H-Spec,一种混合​​ Mamba-attention 并行绘图器,它通过互补模块消耗两个目标上下文源。 Mamba 模块使用预测的最后一个词元目标隐藏状态进行初始化,而注意力模块则在适当的位置重用目标 KV。尽管有重复的表述,Mamba 的并行扫描允许 H-Spec 保留块并行绘图。在三个目标模型和不同的任务中,H-Spec 在平均接受长度方面比最佳基线提高了 5.0--13.3%,在批量大小 1 词元间延迟加速方面提高了 5.3--12.6%。在并发服务下,H-Spec 始终能够实现更高的吞吐量,同时在评估的并发级别上保持比基线更低的 KV 缓存利用率。