COREY: Entropy-Guided Runtime Chunk Scheduling for Selective Scan Kernels
arXiv:2604.10597v2 Announce Type: replace
Abstract: Mamba selective state space models (SSMs) provide linear-time sequence modeling but are often limited by memory bandwidth in practice, where selective state updates are executed as fragmented kernels with repeated intermediate tensor materialization. We present COREY, a prototype scheduler that uses activation entropy estimated via fixed-width histograms as a runtime signal for chunk-size selection at the kernel-invocation level. COREY is positioned as a Concept and Feasibility contribution: a single-parameter runtime auto-tuner built on an existing Triton selective-scan kernel rather than a new fused implementation.
Evidence is organized in three tiers. Tier 1 (Python cost model) shows that entropy-guided grouping reduces surrogate latency and DRAM traffic. Tier 2a (real-checkpoint inline hook) demonstrates that entropy computation and chunk selection can run on the critical path of model.generate(); on Mamba-370M (RTX 3070, n=5), measured overhead is 8.3 percent with full instrumentation and estimated about 2 percent with sparse sampling. Tier 2b (kernel-level scan benchmark) shows that, under a principled calibration where H_ref equals log(K), COREY selects the same chunk as a one-time-profile oracle without offline sweeps and achieves up to 4.41x speedup over static chunk-64.
This work does not yet include a fully integrated end-to-end run connecting Tier 2a and Tier 2b, which remains key future work. Across 80 LongBench prompts, entropy distributions are stable, supporting COREY as a practical runtime auto-tuner within a single regime. Code and data: https://github.com/mabo1215/COREY_Transformer/.