ReplaySSM:缓存 SSM 输入,而非状态
ReplaySSM: Cache SSM Inputs, Not State
基于摘要分析。ReplaySSM 针对状态空间模型(SSM)在混合模型解码时频繁读写循环状态的开销,将状态表示为检查点与近期 SSM 输入缓冲区,从而分离输出计算和检查点更新,改善推理中的状态访存、推测解码回滚及量化误差累积问题。 其核心机制是展开递推关系,周期性更新检查点,避免多数解码步骤写回状态。作者报告,这可将占主导的状态内存流量大致减半。近期输入缓冲区使推测解码回滚成为缓冲区指针操作,并支持不显式生成中间状态的仅输出验证。精度管理方面,多数步骤的输出计算使用低精度检查点,而递推更新使用更高精度检查点,使低精度副本的舍入误差不进入后续递推更新。 作者在 vLLM 中实现了面向 Mamba-2 和 Gated DeltaNet 的 ReplaySSM。作者报告,在服务批大小为 256 时,四个混合模型的标准解码相对调优后的 vLLM 基线获得 1.20–1.47 倍加速;在两个 120B 级模型上,吞吐量约为 vLLM 推测解码的 1.8 倍。作者还报告,与共享低精度检查点相比,分离状态精度可降低输出 KL 散度及下游准确率损失,但摘要未给出对应数值。具体模型、硬件、缓冲长度、精度配置、任务及完整实验协议、消融和统计信息尚未验证。 平台推测(待验证):若 EEG 流式解码器使用兼容的 SSM,且瓶颈确为循环状态访存,检查点与输入缓存分离可能具有复用价值;这是假设,不是已验证的 BCI 效果。迁移需适配输入分块、状态重建和在线延迟约束。低批量、小模型或短状态下,重算与缓存开销可能抵消收益;EEG 低信噪比也可能放大量化引起的输出扰动。可在固定 EEG 模型、数据与精度条件下,对比原始状态更新和 ReplaySSM 式更新的端到端延迟、内存流量及任务指标,并覆盖单流低批量设置。该机制不直接解决跨被试、通道差异或小数据泛化问题,已有 EEG 相关工作尚未检索确认。
值得核对其以输入缓存替代逐步状态写回的实现,以及回滚和混合精度状态管理的收益,但摘要中的加速证据限于所述混合模型推理场景,不能直接外推至 EEG/BCI。
来源:arXiv · 架构与算子 · arxiv.org