论文

COREY:选择性扫描内核的熵引导运行时块调度

COREY: Entropy-Guided Runtime Chunk Scheduling for Selective Scan Kernels

模型推理推理加速

摘要

Mamba 选择性状态空间模型 (SSM) 提供线性时间序列建模,但对选择性扫描块调度仍然敏感。我们提出了 COREY,一个\emph{概念和可行性}运行时调度程序,它将固定箱激活熵映射到块大小。我们从三个层面评估 COREY:原型成本模型、真实检查点内核时序以及现代 GPU 上的路由端到端消融。在内核级别,校准规则 \(H_{\mathrm{ref}}=\log K\) 恢复局部最优块并匹配一次性静态预言机,在消费者 GPU 上产生比未优化基准低 \(4.41\times\) 的延迟,在数据中心加速器上产生 \(3.90\times\)--\(4.04\times\) 较低的延迟。将此选择路由到修补的实时扫描内核中会关闭工程循环,而不会提高端到端速度:在统一路由消融中,最佳静态块优于所有熵引导和代理调度程序。采样直方图 COREY 增加了 \(+4.6\%\) 开销;对 Static-512 的保护后备将其减少到 \(+1.3\%\);轻量级序列长度键控表进一步将其减少到\(+0.7\%\)。然而,两者仍然比静态预言机慢,因为它们保留了调度成本。在 80 个提示的 LongBench 子集上,被动推理和路由推理在输出上完全等效,具有 \(100\%\) 贪婪词元协议和零度量增量。混合机制研究表明,单个序列长度规则与每个机制的块预言相匹配,以实现平衡服务。因此,COREY 被验证为保证质量的调度原型,但当前的熵统计数据并不能在测量的 SSM 检查点工作负载上优于静态块调整。源代码:https://github.com/mabo1215/COREY_Transformer/。