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

Bo Ma; Jinsong Wu; Weiqi Yan

cs.CVApr 12, 2026

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

Authors: Bo Ma, Jinsong Wu, Weiqi Yan

Organizations: Auckland University of Technology; Resideo Technologies, Inc. · Guilin University of Electronic Technology · Auckland University of Technology

Abstract

Mamba selective state space models (SSMs) provide linear-time sequence modeling but remain sensitive to selective-scan chunk scheduling. We present COREY, a 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/.

Figures & tables

Distribution	$K$	$H$ (nats)	$r{=}H/\log K$	Chunk ( $H_{\text{ref}}{=}\log K$ )	Chunk ( $H_{\text{ref}}{=}8.0$ )
Standard normal	$32$	$2.644$	$0.763$	$512$	$256$
Standard normal	$128$	$4.027$	$0.830$	$512$	$256$
Standard normal	$512$	$5.413$	$0.868$	$512$	$256$
Standard normal	$1024$	$6.106$	$0.881$	$512$	$512$
Uniform $[0,1]$	$32$	$3.466$	$1.000$	$512$	$256$
Uniform $[0,1]$	$128$	$4.852$	$1.000$	$512$	$256$

Table 1: Bin-count invariance of the calibrated rule. Under

H_{\text{ref}}{=}\log K

, the ratio

r{=}H/\log K

is nearly constant across

K\in\{32,\ldots,1024\}

for standard-normal inputs and exactly constant for uniform inputs, so the selected chunk is independent of histogram resolution. The legacy reference

H_{\text{ref}}{=}8.0

requires

K{\geq}1024

(standard normal) or

K{\geq}256

(uniform) before chunk selection saturates at

C_{\max}{=}512

, illustrating the bin-count sensitivity that the calibrated rule eliminates.

Policy	Chunk	Latency (ms)	Speedup
Static-64	$64$	$0.101\pm 0.025$	$1.00\times$
COREY (legacy $H_{\text{ref}}{=}8.0$ )	$256$	$0.017\pm 0.000$	$3.24\times$
COREY (default $H_{\text{ref}}{=}\log K$ )	$\mathbf{512}$	$\mathbf{0.013\pm 0.006}$	$\mathbf{4.41\times}$
Static-512 (offline oracle)	$512$	$0.013\pm 0.006$	$4.41\times$

Table 2: Kernel-level selective-scan policy comparison on a consumer GPU (

d{=}1024

L{=}4096

, FP16,

n{=}30

timed repeats after

5

warmup). Speedup is relative to static chunk

=64

and is derived from the end-to-end Python-dispatch chunk-sweep harness (Table 17 in the appendix; Static-64:

3.299\,\mathrm{ms}

, Static-512:

0.748\,\mathrm{ms}

, ratio

4.41\times

); the isolated kernel latencies in the Latency column reflect direct Triton timing only and yield a higher ratio (

0.101/0.013{\approx}7.8\times

) that is not the reported speedup. Under the calibrated rule

H_{\text{ref}}{=}\log K

, COREY selects the offline-profile oracle chunk size; full multi-platform results are in Appendix A .

Distribution	$H$ (nats)	Chunk	Calls	Latency (ms)	Speedup
Uniform	5.545	512	8	$0.741\pm 0.035$	$4.34\times$
Normal (standard)	4.612	256	16	$1.219\pm 0.136$	$2.64\times$
Laplace $(0,1)$	3.892	256	16	$1.140\pm 0.043$	$2.90\times$
Sparse (10% nz)	0.789	64	64	$3.188\pm 0.105$	$1.01\times$
Sparse (2% nz)	0.192	32	128	$6.096\pm 0.213$	$0.52\times$
Static-64 reference		64	64	$3.21\pm 0.13$ ms

Table 3: Activation-distribution perturbation sweep (RTX 3070,

n{=}30

repeats). COREY chunk recommendation monotonically tracks entropy: higher entropy yields larger chunks and lower latency relative to static-64. For sparse activations (entropy

<

1 nat), COREY selects small chunks (32–64), matching or falling below static-64 latency; this conservative behavior is intentional since very sparse, low-entropy activations can concentrate numerical error in large-chunk scans. The Speedup column (relative to static-64) confirms that COREY outperforms the fixed static baseline for all high-to-medium entropy distributions, which represent typical inference regimes.

Platform	Configuration	Latency (ms)	Overhead
Consumer GPU	Passive (stock)	$1160.9\pm 12.0$	—
Consumer GPU	Active, all 48 layers	$1257.5\pm 30.3$	$+8.3\%$
Consumer GPU	Active, $k{=}4$ layer sampling	${\approx}1185$	${\approx}2.1\%$
Data-center accel.	Passive (stock)	$875.0\pm 8.7$	—
Data-center accel.	Active hook only	$895.5\pm 4.6$	$+2.3\%$
Per-call scheduler cost (isolated)		$1.10\pm 0.16$ ms	—

Table 4: Inline entropy scheduling overhead (Mamba-370M, 182-token prompt, 32 generated tokens). Consumer GPU:

n{=}5

2

warmup. Data-center accelerator:

n{=}20

2

warmup. Per-call cost is isolated with synchronization fences on a captured post-convolution tensor.

Configuration	Scheduler	Latency (ms)	Slowdown vs. best static	Chunk distribution
Static chunk $=128$	constant	$892.71\pm 12.11$	$1.0013\times$	$\{128{:}2544\}$
Static chunk $=256$	constant	$914.43\pm 14.58$	$1.0257\times$	$\{256{:}2544\}$
Static chunk $=512$ (oracle)	constant	$\mathbf{891.51\pm 10.17}$	$\mathbf{1.0000\times}$	$\{512{:}2544\}$
Static chunk $=1024$	constant	$897.66\pm 16.59$	$1.0069\times$	$\{1024{:}2544\}$
Static chunk $=2048$	constant	$907.93\pm 10.60$	$1.0184\times$	$\{2048{:}2544\}$
No entropy	midpoint	$903.77\pm 29.67$	$1.0137\times$	$\{1024{:}2544\}$

Table 5: Unified routed scheduler ablation on a data-center accelerator (370M-parameter Mamba checkpoint, 976-token prompt, 32 generated tokens,

n{=}50

timed repeats after

3

warmup). Runtime chunk sizes are honored by a patched recurrence-preserving live scan kernel. “Slowdown vs. best static” uses the static-

512

row as the oracle. The chunk distribution gives the routed chunk count over

2544

scheduler invocations.

Regime	Chunk 128	Chunk 256	Chunk 512	Best
Short conversational ( $\approx$ 25 tok)	184.1 $\pm$ 5.6	185.5 $\pm$ 5.4	187.6 $\pm$ 5.5	128
Long-document QA ( $\approx$ 1500 tok)	307.8 $\pm$ 4.6	306.0 $\pm$ 5.2	305.0 $\pm$ 3.9	512
Source code ( $\approx$ 2100 tok)	349.1 $\pm$ 4.6	347.3 $\pm$ 4.9	345.9 $\pm$ 5.0	512
Structured logs ( $\approx$ 2048 tok)	347.6 $\pm$ 3.3	342.6 $\pm$ 4.8	342.7 $\pm$ 3.1	$256^{\dagger}$
Tabular data ( $\approx$ 2048 tok)	371.5 $\pm$ 9.9	349.8 $\pm$ 3.8	349.7 $\pm$ 3.9	512
Repetitive policy ( $\approx$ 1550 tok)	319.5 $\pm$ 4.5	312.9 $\pm$ 5.6	312.5 $\pm$ 2.7	512

Table 6: Per-regime static chunk sweep on a consumer GPU (

n{=}8

samples per regime,

12

timed repeats,

4

generated tokens, prompts truncated at

2048

tokens). Latency is mean

\pm

std (ms); bold marks the lowest mean per regime. The bottom row gives the equal-weight average; the per-regime oracle is

0.44

ms (

0.14\%

) below the global static-512 mean.

Data-center reference (different hardware; not directly comparable to rows above):
Scheduler	Avg Latency (ms)	$\Delta$ vs. Static-512
Static-512 (global baseline)	$317.1$	—
Learned-table (seq-len rule)	$316.7$	$-0.14\%$ ( $-0.44$ ms)
Per-regime oracle (upper bound)	$316.7$	$-0.14\%$ ( $-0.44\,ms$ )
Guarded sampled-histogram †	$903.0$	$+1.3\%$ †
$\dagger$ : Data-center accelerator, $n{=}50$ warmup $=3$ , 976-token prompt; guard fires (chunk $=512$ every call);

Table 7: Mixed-regime equal-weight average latency for three consumer-GPU scheduler strategies (Mamba-370M, 8 serving regimes,

8

samples/regime, 12 repeats). The guarded sampled-histogram row ( † ) is a single-prompt data-center accelerator measurement, included for reference only; it cannot be directly compared to the consumer-GPU averages above it.

	NarrativeQA	Qasper	MultiFieldQA-EN	GovReport	All
Exact match	$20/20$	$20/20$	$20/20$	$20/20$	$\mathbf{80/80}$
PPL ratio	$1.0000\times$ (all 80 prompts)

Table 8: Task-level routed quality preservation on an 80-prompt LongBench subset (Mamba-370M, sampled-histogram stride 8, 8 greedy tokens, 1024-token cap,

n{=}20

per task). Passive and routed outputs are bitwise identical on all prompts; all metric deltas are zero.

Appendix figures & tables33 assets

Supplementary material from the paper’s appendix.

Appendix

Device	Mean Latency (ms)	Std (ms)	Repeats
TPU v4-8 (PyTorch XLA)	$0.135$	$0.066$	$30$
RTX 3070 (static-512, Triton) †	$0.748$	$0.037$	$5$

Appendix

Table 9: TPU (PyTorch XLA) selective_scan_fn benchmark: batch

=1

, seq

=4096

, dim

=1024

d_{\text{state}}=16

, chunk

=512

, FP16, 30 repeats.

Parameter	Value
Fusion weights $(\alpha,\beta,\gamma)$	$(0.45,0.35,0.20)$
Default threshold $\tau$	0.52
Threshold sweep	$\{0.45,0.52,0.60\}$
Histogram bins $K$ and $\epsilon$	64, $10^{-12}$
EMA decay $\lambda_{\mathrm{ema}}$	0.85
Static-fusion group size	3 operators

Appendix

Table 10: Prototype hyperparameters and scheduler defaults.

Figure 1: Normalized histogram entropy before (red) and after (blue) Hadamard reparameterization, measured on synthetic heavy-tailed activations (prototype study; $K{=}64$ bins, five repeated runs per sequence length). The grouped bar layout makes the post-Hadamard advantage immediately readable: at every sequence length the blue (after) bar exceeds the red (before) bar. The $\Delta$ annotation above each pair reports the absolute entropy gain.

Configuration	Fusion depth	Surr. lat. (ms)	DRAM B/tok	Diagnostic proxy
No-signal	1.00	100.94	122.98	0.0000
Static fusion	2.33	87.22	102.36	0.0851
Arithmetic-only ( $\beta{=}0.70$ )	2.33	83.79	—	—
Entropy-only ( $\alpha{=}0.45,\beta{=}0$ )	1.08	99.07	—	—
COREY default	3.50	77.97	90.58	0.0000

Appendix

Table 11: Consolidated Tier-1 prototype diagnostics (FP16, ultra-long bucket, 5 repeats, seed 7). Default setting:

\alpha{=}0.45

\beta{=}0.35

\gamma{=}0.20

\tau{=}0.52

. All values are prototype surrogate timings from a deterministic Python cost model, not GPU hardware measurements. Full per-ablation tables are in the anonymous repository.

Policy	Model	NarrQA	Qasper	GovRpt	WT103 PPL	PG19 PPL	Avg Lat. (ms)
off	Mamba-370M	0.0299	0.0458	0.1451	556.93	17.20	5142
off	Mamba-1.4B	0.0502	0.0827	0.1750	323.13	13.79	5552
off	Mamba-2.8B	0.0445	0.0399	0.1239	329.80	12.68	7343
static	Mamba-370M	0.0299	0.0458	0.1451	556.93	17.20	5401
static	Mamba-1.4B	0.0502	0.0827	0.1750	323.13	13.79	5788
static	Mamba-2.8B	0.0445	0.0399	0.1239	329.80	12.68	2,068 ‡

Appendix

Table 12: Scheduler policy comparison on 4-task LongBench subset (FP16, WSL2 RTX 3070). NarrQA/Qasper: token-F1; GovRpt: ROUGE-L. Avg Lat. is the mean per-sample latency across the four tasks (ms). Quality metrics do not vary by policy for the same model (identical FP16 inference weights); only latency reflects scheduling differences. policy_off rows used

n{=}5

without explicit warmup (first-call JIT cost included). ‡ Mamba-2.8B policy_static / policy_corey : RTX 3090, CUDA 12.1,

n{=}20

, warmup

=1

(first sample excluded; mean over NarrQA/Qasper/GovRpt/MultifieldQA); lower Avg Lat. vs. policy_off (7343 ms, RTX 3070) reflects JIT warm-up included in the policy_off

n{=}5

run. NLP scores copied from the matched policy_off row (policy-independent).

Policy	Status	Latency (ms)	Tokens/s	Token-F1
off	ok	2846.8980	11.2403	0.148148
static	ok	2309.8472	13.8537	0.148148
corey	ok	2376.1357	13.4672	0.148148

Appendix

Table 13: W1 triplet smoke comparison on real GPU (WSL2, RTX 3070, Mamba-370M, NarrativeQA,

n{=}1

). Lower latency is better. These are benchmark-path timings from the current scheduler-hook integration and are not yet fused-kernel method-vs-baseline evidence.

Scope	$n$	Mean	Std	p5	p95	Coarse bucket(s)
All prompts	80	4.022	0.092	3.871	4.139	100% in 256
NarrativeQA	20	3.999	0.079	3.887	4.138	100% in 256
Qasper	20	4.032	0.083	3.891	4.127	100% in 256
MultiFieldQA-EN	20	4.042	0.102	3.931	4.224	100% in 256
GovReport	20	4.017	0.096	3.869	4.130	100% in 256

Appendix

Table 14: Real-checkpoint input-entropy distribution for Mamba-370M over 80 LongBench prompts (20 each from NarrativeQA, Qasper, MultiFieldQA-EN, and GovReport). The runtime hook emits a continuous recommendation of 288 for all 80 prompts. Under the reviewer-requested coarse discretization

\{32,64,128,256,512\}

, all 80 prompts map to bucket 256. This indicates that the currently observed real-workload regime is narrow and medium-high entropy rather than strongly multi-modal.

Regime	Entropy mean $\pm$ std	Prompts	Aggregate chunk distribution
Templated / repetitive	$3.193\pm 0.030$	20	128:4, 256:952, 512:4
Code-like / structured	$3.107\pm 0.045$	20	256:960
Dense narrative	$3.100\pm 0.036$	20	256:960

Appendix

Table 15: H800 heterogeneous-corpus entropy check (Mamba-370M, NVIDIA H800 PCIe,

K{=}256

, 20 prompts per regime, 48 captured layer calls per prompt). The aggregate chunk counts are over layer calls, not prompts.

Regime	Entropy mean $\pm$ std	Prompts	Aggregate chunk distribution
Extra code	$2.965\pm 0.308$	20	256:960
Extra logs	$3.031\pm 0.255$	20	256:960
Extra repetition	$2.998\pm 0.276$	20	256:960
Extra tables	$2.990\pm 0.268$	20	256:960

Appendix

Table 16: Expanded H800 workload-diversity stress test (Mamba-370M, NVIDIA H800 PCIe,

K{=}256

). Each extra regime contains 20 prompts and 960 captured layer calls. All added regimes select chunk 256 for every call.

Figure 2: Per-prompt input entropy across 80 LongBench prompts on Mamba-370M (20 per task). Blue violins show the empirical kernel density and jittered scatter points mark individual prompts; red diamonds mark per-task mean $\pm$ standard deviation. Dashed horizontal lines mark the midpoints of the coarse chunk buckets $\{32,64,128,256,512\}$ used by COREY’s discretization; the dotted red line is the theoretical entropy ceiling $\log K{=}5.55$ nats at $K{=}256$ histogram bins. The empirical distribution sits well below the ceiling and entirely within the chunk 256 band, illustrating why COREY produces the same chunk recommendation across all four tasks in the current checkpoint regime and why detecting distribution shift would require prompts that reach into adjacent bands. This is the figure counterpart to Table 14 .

Configuration	Chunk	Kernel calls	Latency (ms)	Speedup vs. static-64
static-32	32	128	$6.315\pm 0.385$	$0.52\times$
static-64	64	64	$3.299\pm 0.257$	$1.00\times$
static-128	128	32	$1.867\pm 0.130$	$1.77\times$
corey-256 (legacy, $H_{\text{ref}}{=}8.0$ )	256	16	$1.148\pm 0.048$	$2.87\times$
corey-512 (default, $H_{\text{ref}}{=}\log K$ )	512	8	$\mathbf{0.748\pm 0.037}$	$\mathbf{4.41\times}$
static-512	512	8	$0.748\pm 0.037$	$4.41\times$

Appendix

Table 17: Static chunk-size sweep latency curve and COREY entropy-selected point (RTX 3070, seq_len = 4096, batch = 1, FP16,

n{=}30

repeats). Latency decreases monotonically with chunk size because fewer kernel launches reduce per-call Python dispatch overhead. COREY selects chunk_size = 256 based on the measured input entropy (4.60 nats) of the standard-normal synthetic activations used in this benchmark; it is 2.87

\times

faster than the fixed static-64 baseline used in the main-text benchmark and 53.4 % slower than the exhaustive-sweep oracle ( static-512 ). Table 3 shows that the separate Uniform

[0,1]

case reaches 5.55 nats and selects 512. The static oracle requires an a-priori exhaustive sweep; COREY provides adaptive selection without it. Note: the corey-256 row was measured under the legacy

H_{\text{ref}}{=}8.0

; under the paper default

H_{\text{ref}}{=}\log K{=}5.55

nats, the same input entropy

H{=}4.60

nats gives

r{=}0.83

and selects chunk

=512

—see Table 20 .

Setting	Static-64 (ms)	COREY chunk	COREY (ms)	Speedup	Static-512 oracle (ms)
H800 FP16, $L{=}4096$	$2.299\pm 0.061$	256	$0.587\pm 0.006$	$3.915\times$	$0.300\pm 0.003$
H800 BF16, $L{=}4096$	$2.343\pm 0.110$	256	$0.581\pm 0.004$	$4.035\times$	$0.302\pm 0.002$
H800 FP16, $L{=}8192$	$4.548\pm 0.017$	256	$1.167\pm 0.005$	$3.898\times$	$0.586\pm 0.003$

Appendix

Table 18: H800 W1 kernel supplement. The COREY rows report the entropy-selected chunk-256 point from the triplet harness; Static-512 is the measured sweep oracle on the same hardware. This table strengthens the Hopper-class kernel evidence. The later runtime-chunk H800 closure in Appendix A.14 routes the selected chunk into the live checkpoint-inference scan path, but its end-to-end gain is small and scheduler-sensitive.

Hardware	Policy	Chunk	Calls	Lat. (ms)	Std (ms)	Spdup-A	Spdup-B
Consumer GPU (CUDA 12.8, WSL2)
	Static-64	64	64	$0.101$	$0.025$	$1.00\times$	$0.21\times$
	COREY (legacy, $H_{\text{ref}}{=}8.0$ )	256	16	$0.017$	$0.000$	$3.24\times$	$0.68\times$
	COREY (default, $H_{\text{ref}}{=}\log K$ )	512	8	$\mathbf{0.013}$	$\mathbf{0.006}$	$\mathbf{4.41\times}$	$\mathbf{1.00\times}$
	Static-512 (offline oracle)	512	8	$0.013$	$0.006$	$4.41\times$	$1.00\times$
TPU (torch_xla)

Appendix

Table 19: Multi-platform kernel-level selective-scan policy comparison (

d{=}1024

d_{\text{state}}{=}16

L{=}4096

, FP16,

n{=}30

timed repeats after

5

warmup). Speedup-A is relative to static chunk

=64

; Speedup-B is relative to static chunk

=512

(the offline-profile oracle). The default calibration

H_{\text{ref}}{=}\log K

selects chunk

=512

on both consumer and server-class GPUs; the legacy reference

H_{\text{ref}}{=}8.0

selects chunk

=256

and is shown only as a sensitivity point.

	W1 benchmark ( $H{=}4.60$ nats)			Real checkpoint mean ( $H{=}4.02$ nats)
$H_{\text{ref}}$	$r$	Chunk	Latency (ms)	$r$	Chunk	Latency (ms)
$\log 64\approx 4.16$	1.000	512	0.748	0.966	512	0.748
5.0	0.921	512	0.748	0.804	512	0.748
6.0	0.767	512	0.748	0.670	256	1.148
8.0 ⋆ (legacy)	0.576	256	1.148	0.503	256	1.148

Appendix

Table 20:

H_{\text{ref}}

sensitivity: chunk selected and resulting latency for two entropy scenarios. Latency values are drawn from the measured static-chunk sweep (Table 17 ) and require no additional GPU runs; the legacy default

H_{\text{ref}}{=}8.0

is marked with

\star

. The principled upper bound

H_{\text{ref}}{=}\log K{=}\log 256{\approx}5.55

(not shown but between rows 3 and 4) would select chunk = 512 for both scenarios, matching the oracle. Static-64 baseline latency is 3.30 ms.

Distribution	$K$	$H$ (nats)	$\log K$	$H/\log K$	Chunk ⋄	Chunk (fixed $8.0$ )	Lat ⋄ (ms)
standard normal	32	2.644	3.466	0.763	512	256	0.748
standard normal	64	3.334	4.159	0.802	512	256	0.748
standard normal	128	4.027	4.852	0.830	512	256	0.748
standard normal	256	4.720	5.545	0.851	512	256	0.748
standard normal	512	5.413	6.238	0.868	512	256	0.748
standard normal	1024	6.106	6.931	0.881	512	512	0.748

Appendix

Table 21: Chunk selection as a function of histogram resolution

K

H

is the Shannon entropy (nats) estimated from

10^{6}

CPU samples per scenario. Under the principled calibration

H_{\text{ref}}{=}\log K

(columns marked ⋄ ), the normalized ratio

H/\log K

is approximately

K

-invariant, and the selected chunk locks to

C{=}512

regardless of

K

, so COREY’s output is decoupled from histogram resolution. Under the fixed default

H_{\text{ref}}{=}8.0

, small

K

systematically biases the selection toward smaller chunks because

8.0

exceeds

\log K

until

K{=}1024

. Latencies are drawn from the chunk sweep in Table 17 (RTX 3070, seq_len = 4096, FP16).

Model	GPU	Prompt tok.	Base Lat. (ms)	Hook Lat. (ms)	Delta (%)	Entropy	Tile
Mamba-370M	RTX 3070	512	2758.55	2658.39	$-3.63$	4.18	288
Mamba-370M	RTX 3090	512	2675.36	2567.73	$-4.02$	4.18	288
Mamba-1.4B	RTX 3090	512	2644.27	2609.94	$-1.30$	5.17	352

Appendix

Table 22: Real-GPU online scheduler microbenchmark (HF fallback path, moved from main text per reviewer request). Entropy hook runs on real checkpoint activations during generation. RTX 3070: WSL2 CUDA 12.8, 1-sample/0-warmup. RTX 3090: native Linux CUDA 12.1, 1-sample/1-warmup/3-repeat. Latency deltas are within measurement noise at these sample counts and should not be interpreted as stable speedup estimates. Negative delta values (

-3.6\%

-1.3\%

) reflect GPU cold-start and OS scheduling variability at

n{=}1

, zero warmup; the passive hook adds no GPU arithmetic, so any apparent latency reduction is a noise artifact, not hook-induced speedup.

Configuration	Latency (ms)	Overhead
Passive (stock fast path)	$1160.9\pm 12.0$	—
Active (all 48 layers)	$1257.5\pm 30.3$	$+8.3\%$
Active ( $k{=}4$ layer sampling, estimated)	${\approx}1185$	${\approx}2.1\%$
Per-call scheduler cost (isolated)	$1.10\pm 0.16$ ms	—

Appendix

Table 23: Active-mode inline entropy integration overhead (Mamba-370M, RTX 3070/CUDA 12.8, 32 new tokens,

n{=}5

). Per-call cost is isolated with torch.cuda.synchronize() fences on a captured post-conv tensor (

u

shape

[1,2048,182]

). Chunk distribution: 322 calls selected chunk=256 (medium-entropy layers); 14 calls selected chunk=128 (layers 0 and 15,

H{\approx}2.2

nats).

Configuration	Latency (ms)	vs. Passive
Passive (stock fast path)	$875.0\pm 8.7$	$1.000\times$
Active hook only	$895.5\pm 4.6$	$1.023\times$
Active + routed-call scaffold	$893.1\pm 5.0$	$1.021\times$

Appendix

Table 24: H800 active integration scaffold measurement (Mamba-370M, NVIDIA H800 PCIe, CUDA 12.8, PyTorch 2.8.0, 122-token prompt, 32 new tokens,

n{=}20

, warmup

=2

Scheduler	Prompt	Repeats	Passive (ms)	Active+routed (ms)	vs. Passive
Full histogram	122 tokens	50	$898.1\pm 31.0$	$929.2\pm 26.6$	$1.035\times$
Sampled histogram, stride 8	976 tokens	50	$902.1\pm 29.4$	$893.6\pm 18.6$	$0.9905\times$
Constant chunk 512	976 tokens	50	$910.0\pm 31.3$	$911.7\pm 32.2$	$1.0018\times$

Appendix

Table 25: H800 runtime-chunk routed live-scan measurements (Mamba-370M, NVIDIA H800 PCIe, CUDA 12.8, PyTorch 2.8.0, 32 new tokens, runtime CUDA dispatch via src.corey_selective_scan_dispatch ). The full-histogram scheduler routes into the live scan kernel but is negative; sampled histogram with stride 8 lowers scheduler cost in this passive-baseline comparison, but the unified static-oracle ablation in Table 27 remains negative.

Run	Configuration	Chunk	Active+routed (ms)	vs. best static	Chunk distribution
Full sweep	Static oracle	128	$876.6\pm 26.6$	$1.000\times$	$\{128{:}2544\}$
Full sweep	Adaptive sampled hist., stride 8	–	$905.6\pm 18.2$	$1.033\times$	$\{1024{:}2544\}$
Confirm subset	Static oracle	1024	$902.0\pm 33.1$	$1.000\times$	$\{1024{:}2544\}$
Confirm subset	Adaptive sampled hist., stride 8	–	$887.0\pm 15.8$	$0.983\times$	$\{512{:}2544\}$

Appendix

Table 26: H800 static-oracle versus adaptive routed comparison (Mamba-370M, prompt length 976, 32 new tokens,

n{=}50

, warmup

=3

). The confirm run repeats only chunk 128, chunk 1024, and adaptive sampled histogram.

Configuration	Scheduler	Active+routed (ms)	vs. passive	vs. best static	Chunk distribution
Static chunk 128	constant	$892.71\pm 12.11$	$1.0032\times$	$1.0013\times$	$\{128{:}2544\}$
Static chunk 256	constant	$914.43\pm 14.58$	$1.0045\times$	$1.0257\times$	$\{256{:}2544\}$
Static chunk 512	constant	$\mathbf{891.51\pm 10.17}$	$1.0048\times$	$\mathbf{1.0000\times}$	$\{512{:}2544\}$
Static chunk 1024	constant	$897.66\pm 16.59$	$1.0061\times$	$1.0069\times$	$\{1024{:}2544\}$
Static chunk 2048	constant	$907.93\pm 10.60$	$1.0153\times$	$1.0184\times$	$\{2048{:}2544\}$
No entropy	midpoint	$903.77\pm 29.67$	$1.0116\times$	$1.0137\times$	$\{1024{:}2544\}$

Appendix

Table 27: Unified H800 routed scheduler ablation (Mamba-370M, NVIDIA H800 PCIe, CUDA 12.8, PyTorch 2.8.0, prompt length 976, 32 new tokens, warmup

=3

n{=}50

). Runtime chunks are honoured by the patched recurrence-preserving live scan kernel. “vs. best static” uses Static-512 as the oracle row.

Task	$n$	Exact match	Metric	Passive mean	Routed mean	Metric $\Delta$	PPL ratio
NarrativeQA	20	$20/20$	token-F1	$0.0592$	$0.0592$	$+0.0000$	$1.0000\times$
Qasper	20	$20/20$	token-F1	$0.0454$	$0.0454$	$+0.0000$	$1.0000\times$
MultifieldQA-EN	20	$20/20$	exact-match	$0.0000$	$0.0000$	$+0.0000$	$1.0000\times$
GovReport	20	$20/20$	ROUGE-L	$0.0150$	$0.0150$	$+0.0000$	$1.0000\times$
Aggregate	80	80/80 ( $100\%$ )	—	—	—	$0.0000$ all tasks	$1.0000\times$ all

Appendix

Table 28: H800 task-level routed quality check on a LongBench subset (Mamba-370M, sampled-histogram scheduler with stride 8, 8 greedy new tokens, max prompt length 1024). Four LongBench tasks are sampled at

20

prompts each (80 total). For each task we report exact-match rate (greedy tokens identical between passive and routed paths), task-level metric mean for passive and routed, and the per-sample teacher-forcing perplexity-ratio mean. The routed path uses the patched runtime-chunk live selective-scan kernel. All 80 greedy generations are bitwise identical between passive and routed paths; all per-sample perplexity ratios are exactly

1.0000\times

Layer	$H(u)$ (nats)	Chunk	vs. synthetic	W1 ref. speedup
0	2.27	128	lower H; chunk differs	1.92 $\times$
8	2.91	256	consistent	3.24 $\times$
16	3.22	256	consistent	3.24 $\times$
24	3.29	256	consistent	3.24 $\times$
32	3.61	256	consistent	3.24 $\times$
40	3.54	256	consistent	3.24 $\times$

Appendix

Table 29: Cross-layer entropy and chunk selection on real Mamba-370M activations (42-token prompt,

K{=}256

bins,

H_{\text{ref}}{=}\log K{=}5.55

). The legacy setting

H_{\text{ref}}{=}8.0

produces identical chunk selections for all seven sampled layers because the real-checkpoint entropy values fall in a range where both settings map to the same discretized chunks—see Table 20 . Entropy overhead is the steady-state cost of applying the histogram estimator to the captured

u

tensor (

B{=}1

d_{\text{inner}}{=}2048

L{=}42

). The W1 reference speedup is the standalone chunk-latency ratio from Table 1 of the main paper (W1 chunked-scan benchmark) for the selected chunk vs. static-64; it is not an end-to-end checkpoint speedup.

Hardware	Policy	Latency (ms)	Launches	CUDA Kernels	CUDA Time ( $\mu$ s)
RTX 3090 / CUDA 12.1 / Linux (4 $\times$ GPU server), PyTorch 2.4.0, Triton 3.0
	policy_off	14.964	4096	16	19.1
	policy_static	2.216	64	16	18.1
	policy_corey	0.308	9	9	11.5
	Speedup vs. policy_off ‡			static: $6.75\times$	corey: $48.6\times$
RTX 4050 Laptop / CUDA 12.8 / WSL2, PyTorch 2.11.0, Triton 3.6

Appendix

Table 30: Kernel-level CUDA profile for the three scheduling policies applied to a Triton chunked scan (seq_len = 4096, batch = 1, FP16). Launches is the number of kernel calls per sequence. CUDA Kernels is the count from torch.profiler . CUDA Time is the profiler-reported total kernel device time. All results use a lightweight synthetic Triton scan kernel (not mamba_ssm ). Lower latency and fewer launches are better.

Figure 3: Overview of entropy-guided SSM operator fusion with fused Hadamard reparameterization.

Platform	Model	NarrQA	Qasper	GovRpt	PPL ${}_{\text{WT103}}$	PPL ${}_{\text{PG19}}$	Avg Lat. (ms)
RTX 3070 / CUDA 12.8	Mamba-370M	0.0135	0.0350	0.1370	809.36	14.68	1384.25
RTX 3070 / CUDA 12.8	Mamba-1.4B	0.0191	0.0460	0.1498	1128.40	11.66	2868.74
RTX 3070 / CUDA 12.8	Pythia-410M	0.0190	0.0392	0.1550	5901.47	–	624.43
RTX 3090 / CUDA 12.1 (1-GPU)	Mamba-370M	0.0135	0.0443	0.1459	–	–	27205.05
RTX 3090 / CUDA 12.1 (1-GPU)	Mamba-1.4B	0.0191	0.0569	0.1582	–	–	– ⋆
RTX 3090 / CUDA 12.1 (2-GPU)	Mamba-1.4B	0.0191	0.0569	0.1561	–	–	27971.14

Appendix

Table 31: Checkpoint-level external baseline comparison on LongBench (20 samples/task, FP16). WikiText-103 and PG19 are teacher-forcing side perplexities; “–” means not collected. MF-EN exact-match is zero across all models because 32-token generation does not match gold strings under strict exact-match; this is expected and is not a harness defect. RTX 3090 per-sample latency is dominated by multi-process I/O overhead, not GPU compute.

# GPUs	Wall-clock (s)	Agg. tok/s	Speedup
1 (sequential)	2154.3	5.04	$1.00\times$
2 (data-parallel)	1019.5	10.56	$2.11\times$
4 (data-parallel)	625.3	19.03	$3.45\times$

Appendix

Table 32: Data-parallel multi-GPU inference throughput for Mamba-370M (four-task LongBench, 20 total samples, FP16, RTX 3090). Wall-clock covers all four tasks end-to-end. Aggregate tok/s sums per-shard throughputs. Speedup relative to 1-GPU sequential wall-clock. Scores at all GPU counts agree with Table 31 within sampling noise.

$n$	Regime	Policy	Launches	Surr. lat. (ms)	Red. (%)	Solver ( $\mu$ s)
4	low	no-fusion	4	0.206	—	—
4	low	static-3	2	0.103	50.0	—
4	low	COREY	1	0.052	75.0	211.7
4	mixed	no-fusion	4	0.206	—	—
4	mixed	static-3	2	0.103	50.0	—
4	mixed	COREY	1	0.052	75.0	114.4

Appendix

Table 33: Entropy-guided fusion algorithm efficiency sweep (4

\times

RTX 3090 server / CUDA 12.1, 1000 solver-timing repeats per configuration). “Launches” is the number of kernel-launch groups produced; “Surr. lat.” is the surrogate dispatch latency at

\approx 0.052

ms per call; “Red. (%)” is reduction vs. the no-fusion baseline; “Solver” is the wall-clock time of the Python fusion solver. The surrogate latency reduction arises purely from reducing launch count, not from actual Triton kernel fusion (which is prospective; see Appendix B ).

Model	Arch	Platform	NarrQA	Qasper	GovRpt	PPL ${}_{\text{WT103}}$	Mean Lat. (ms)
Mamba-370M	Mamba-1	RTX 3070 FP16	0.0135	0.0350	0.1370	809.4	1384
Mamba-1.4B	Mamba-1	RTX 3070 FP16	0.0191	0.0460	0.1498	1128.4	2869
Pythia-410M	GPT-Neo	RTX 3070 FP16	0.0190	0.0392	0.1550	5901.5	624
RWKV-4-430M	RWKV	CPU FP32 (ref.)	0.028	0.031	0.121	13.7	134 †
GPT-2+FA2	Attn.	RTX 3090 (ref.)	0.019	0.022	0.108	29.4	58 †
Mamba2-2.7B	Mamba-2	RTX 3070 FP16 ‡	0.098	0.106	0.044 §	—	1,820

Appendix

Table 34: Extended external baseline comparison. Mamba-1.x results are from Table 31 (in-repository, FP16, WSL2 RTX 3070). RWKV-4 and FlashAttention-2 values are representative literature references. NarrQA / Qasper: token-F1; GovRpt: ROUGE-L; MF-EN: exact-match; PPL: WikiText-103 teacher-forcing perplexity. The two H800 rows are real full-model LongBench baselines collected in the 50-sample closure repeat: Pythia-410M uses FlashAttention-3; Mamba2-2.7B uses the H800 SSD path after a Transformers-generation compatibility shim. The COREY scheduler is not applied to Mamba-2 in this submission (consistent with the main-text scope statement in Section 1).

Seq. len	Latency (ms)	Tokens/s	Nominal TFLOP/s
1024	$0.103\pm 0.009$	$9.90{\times}10^{6}$	$20.8$
2048	$0.099\pm 0.003$	$20.6{\times}10^{6}$	$86.3$
4096	$0.171\pm 0.002$	$23.9{\times}10^{6}$	$200.5$
8192	$0.421\pm 0.010$	$19.4{\times}10^{6}$	$326.3$

Appendix

Table 35: FlashAttention-3 matched raw-kernel benchmark on NVIDIA H800 PCIe (causal BF16, batch

=1

, heads

=16

, head-dim

=64

, warmup

=10

n{=}100

repeats). TFLOP/s is the script’s nominal forward-attention estimate and should be read as a kernel-normalized reference rather than end-to-end model throughput.

Task	Mean Lat. (ms)	Std (ms)	Tok/s
NarrativeQA	1583	776	21.96
Qasper	1524	537	22.14
GovReport	2615	530	12.74
MultifieldQA-EN	1516	500	22.15
Mean (short-ctx)	1541	—	22.08

Appendix

Table 36: FP16 baseline for Quamba benchmark ( state-spaces/mamba-1.4b-hf , 4

\times

RTX 3090 / CUDA 12.1,

n{=}20

per task). Quamba INT4 is blocked pending CUDA 12.8/sm_89 hardware; this table establishes the FP16 reference that INT4 would be compared against.

$\dagger$ : $n{=}200$ ; all other cells $n{=}20$ .
Policy	NarrQA Lat. (ms)	Qasper Lat. (ms)	GovRpt Lat. (ms)	MField Lat. (ms)	Mean Tok/s
corey	1,361 †	1,363 †	1,543	1,437	22.9
static	1,430	1,440	1,546	1,439	22.6

Appendix

Table 37: Policy COREY ablation on state-spaces/mamba-1.4b-hf , four-task LongBench subset (FP16, 4

\times

RTX 3090 / CUDA 12.1, input truncated to 2,048 tokens, 32 output tokens). “Lat.” is mean per-sample generation latency in ms. “Tok/s” is mean generation throughput. NarrativeQA and Qasper corey used

n{=}200

; all other cells used

n{=}20

. For the balanced pair (GovReport, MultifieldQA-EN,

n{=}20

each), corey overhead vs. static is

2

–

4

ms (

<0.3\%

$\dagger$ : H800 PCIe, $n{=}50$ , warmup $=3$ , 976-token prompt; both variants route to chunk $=512$ ;
Scheduler	Avg Lat. (ms)	$\Delta$ vs. Static-512	Source
Static-512 (global baseline)	317.1	—	RTX 3070 sweep
Learned-table (seq-len rule)	316.7	$-0.14\%$ ( $-0.44$ ms)	derived from sweep
Per-regime oracle (upper bound)	316.7	$-0.14\%$ ( $-0.44$ ms)	RTX 3070 sweep
Guarded sampled-hist (fallback=512)	903.0 ms †	$+1.29\%$ ( $+11.5$ ms) †	H800 routed †
Learned-table (seq-len rule, H800)	897.6 ms †	$+0.69\%$ ( $+6.1$ ms) †	H800 routed †

Appendix

Table 38: Mixed-regime equal-weight average latency for four scheduler strategies (RTX 3070, Mamba-370M,

8

serving regimes,

8

samples/regime,

12

repeats, unless noted). Learned-table latency is derived by applying the seq-len rule (chunk

=128

for

\leq 50

tokens; chunk

=512

otherwise) to measured per-regime sweep values. Per-regime oracle selects the best static chunk independently per regime. The guarded and learned (H800) rows use a single 976-token prompt on H800 PCIe (

n{=}50

, warmup

=3

); see text.

LLM Inference Scheduling

Large language model inference scheduling allocates requests, batches, compute, and KV-cache memory to balance latency, throughput, and service-level objectives.

See all 160 papers

Scaling State-Space Models on Multiple GPUs with Tensor Parallelism
Feb 24, 2026Anurag Dutt, Nimit Shah, Hazem Masarani +1LLM Inference Acceleration Tensor Parallelism
Stony Brook University
KernelSight-LM: A Kernel-Level LLM Inference Simulator
Jun 26, 2026Xiteng Yao, Taeho Kim, Hengzhi Pei +7LLM Inference LLM Inference Serving
Amazon Web Services Santa Clara, CA, USA · Boston University Boston, MA, USA
How Much of a Real Workload Can LLM-Generated GPU Kernels Actually Reach?
Sep 17, 2026Gaurav Agarwal, Ashish Garg, Isha SinghalGPU Kernel Generation GPU Kernel Optimization

Feb 24, 2026J/K move · Enter open · S save

Scaling State-Space Models on Multiple GPUs with Tensor Parallelism

Anurag Dutt, Nimit Shah, Hazem Masarani, Anshul Gandhi

Stony Brook University

LLM Inference Acceleration Tensor Parallelism Selective SSMs

Open paper AskPDF arXiv

Selective state space models (SSMs) have rapidly become a compelling backbone for large language models, especially for long-context workloads. Yet in deployment, their inference performance is often bounded by the memory capacity, bandwidth, and latency limits of a single GPU, making multi-GPU execution increasingly necessary. Although tensor parallelism (TP) is widely used to scale Transformer inference, applying it to selective SSM blocks is non-trivial because the SSM mixer couples large projections with a sequence-wise recurrent state update and local mixing whose efficiency depends on preserving locality and avoiding synchronization in the critical path. This paper presents a communication-efficient TP design for selective SSM inference that addresses three practical engineering challenges: enabling TTFT improvements via an SSM state cache across prefill and decode, partitioning the mixer's packed parameter tensor so that recurrent updates remain local while minimizing communication, and reducing TP aggregation overhead with quantized AllReduce. We evaluate on three representative SSM-based LLMs spanning pure-SSM and hybrid architectures - Mamba, Falcon-Mamba, and Zamba - on NVIDIA A6000 and A100 clusters. Our experiments show substantial throughput gains from tensor-parallel SSM inference, improving batch-request throughput by ~1.6-2.1x on 2 GPUs and ~2.6-4.0x on 4 GPUs for Mamba, with the largest benefits at long context lengths, and achieving a further ~10-18% throughput improvement from quantized all-reduce by lowering synchronization bandwidth overhead.

Abstract

LLM Inference Scheduling

Explore similar work

Scaling State-Space Models on Multiple GPUs with Tensor Parallelism

KernelSight-LM: A Kernel-Level LLM Inference Simulator

How Much of a Real Workload Can LLM-Generated GPU Kernels Actually Reach?