论文
CAST:来自一次性块起草者的具有成本意识的推测树
CAST: Cost-Aware Speculative Trees from One-Pass Block Drafters
摘要
推测性解码通过廉价地起草未来标记并使用目标模型并行验证它们来加速大语言模型推理。区块起草者在一次前向传递中对整个未来词元区块进行评分,但标准解码仅验证得分最高的链并丢弃其他候选者。由于这些候选者已经评分,验证更多候选者会增加目标计算,但不会增加额外的起草工作。我们引入了 CAST(成本感知推测树),它将这些候选者打包到树中,并在单个目标传递中对其进行验证,而目标模型、起草者权重和解码规则保持不变。为了决定树的宽度,CAST 添加了候选树,而下一个候选树的预期收益超过了它所增加的验证时间。因此,宽度会根据延迟测量来适应每次部署,而无需扫描宽度。我们在三个 GPU 代和两个模型系列上跨五个领域评估 CAST。在其预测宽度下,CAST 在所有八种设置中都比标准链快 43%。我们还发现最佳宽度很大程度上取决于部署。当验证成本在内核边界跳跃时,128 个词元的树仅比标准链快 2%,而预测宽度的树则快 20%。此外,我们证明 CAST 在贪婪解码和采样解码下都保持目标输出分布不变。代码可在 https://github.com/js-lee-AI/CAST. 获取