通过熵值动态调度分块大小,提升Mamba模型推理效率
COREY: Entropy-Guided Runtime Chunk Scheduling for Selective Scan Kernels

- 用激活熵映射分块大小,实现运行时自适应调度
- 在消费级显卡上比基线低4.41倍延迟,接近静态最优性能
- 适合需要高吞吐的推理场景,但静态调参仍更优
Mamba选择性状态空间模型(SSMs)虽具备线性时间序列建模能力,但对选择性扫描分块调度敏感。本文提出COREY,一种概念与可行性验证的运行时调度器,将固定分箱激活熵映射为分块大小。评估涵盖三层次:原型成本模型、真实检查点内核计时及现代GPU上的端到端消融实验。在内核层面,校准规则 $H_{\mathrm{ref}}=\log K$ 恢复局部最优分块,匹配一次性的静态最优调度,在消费级GPU上实现比未优化基线低4.41倍延迟,在数据中心加速器上降低3.90至4.04倍。将此决策注入修补后的实时扫描内核后,虽完成工程闭环,但端到端速度未提升:统一路由消融中,最佳静态分块优于所有熵引导和代理调度器。采样直方图版本带来+4.6%开销;采用静态度量512的保护回退降至+1.3%;轻量级长度键表进一步降至+0.7%。然而两者仍慢于静态最优,因保留调度开销。在80个提示的LongBench子集上,被动与路由推理完全输出等价,贪婪词元一致率100%,指标无差异。混合模式研究显示,单一长度规则可匹配各区间最优分块策略,实现均衡服务。因此,COREY被验证为保质的调度原型,但当前熵统计无法在测量的SSM检查点工作负载上超越静态分块调优获得稳定吞吐优势。
原文摘要 · Abstract (English)
Mamba selective state space models (SSMs) provide linear-time sequence modeling but remain sensitive to selective-scan chunk scheduling. We present COREY, a \emph{concept-and-feasibility} runtime scheduler that maps fixed-bin activation entropy to chunk size. We evaluate COREY in three tiers: a prototype cost model, real-checkpoint kernel timing, and routed end-to-end ablations on modern GPUs. At the kernel level, a calibrated rule, \(H_{\mathrm{ref}}=\log K\), recovers the locally optimal chunk and matches a one-time static oracle, yielding \(4.41\times\) lower latency than an unoptimized baseline on a consumer GPU and \(3.90\times\)--\(4.04\times\) lower latency on a data-center accelerator. Routing this choice into a patched live scan kernel closes the engineering loop without improving end-to-end speed: in unified routed ablations, the best static chunk outperforms all entropy-guided and proxy schedulers. Sampled-histogram COREY adds \(+4.6\%\) overhead; a guarded fallback to Static-512 reduces this to \(+1.3\%\); and a lightweight sequence-length-keyed table further reduces it to \(+0.7\%\). However, both remain slower than the static oracle because they retain scheduling cost. On an 80-prompt LongBench subset, passive and routed inference are exactly output-equivalent, with \(100\%\) greedy-token agreement and zero metric deltas. A mixed-regime study shows that a single sequence-length rule matches the per-regime chunk oracle for balanced serving. COREY is therefore validated as a quality-preserving scheduling prototype, but current entropy statistics are not a robust throughput win over static chunk tuning on measured SSM checkpoint workloads. SourceCode: https://github.com/mabo1215/COREY_Transformer/.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。