🤖 AI

COREY: A Prototype Study of Entropy-Guided Operator Fusion with Hadamard Reparameterization for Selective State Space Models

COREY is a prototype framework that enhances State Space Models by combining memory-aware operator fusion with Hadamard-based feature reparameterization and entropy-guided scheduling to reduce latency, improve throughput, and lower DRAM traffic in long-context inference.

Original authors: Bo Ma, Jinsong Wu, Hongjiang Wei, Weiqi Yan

Published 2026-04-14

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: Bo Ma, Jinsong Wu, Hongjiang Wei, Weiqi Yan

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you are running a massive, high-speed library where books (data) need to be read, processed, and summarized. In the past, the best way to do this was to have a librarian (the AI) read every single word, remember the context, and write a summary. But as the books got longer, the librarian got overwhelmed, and the process slowed down because they had to constantly walk back and forth to the shelves to grab new pages.

Enter State Space Models (SSMs), like the "Mamba" family. These are like a super-efficient librarian who can read a whole book in linear time—they don't get slower as the book gets longer. However, there's a catch: even though they read fast, they are constantly dropping off piles of notes on the floor (memory) and running back to pick them up. This "running back and forth" (moving data between memory and the processor) is the bottleneck. It's like a chef who can chop vegetables instantly but spends 90% of their time walking to the fridge to get ingredients.

This paper introduces COREY, a new way to organize the kitchen to stop the walking.

Here is the simple breakdown of how COREY works, using everyday analogies:

1. The Problem: Too Many Trips to the Fridge

In current AI systems, the "chef" (the processor) has to stop, write down a note, walk to the fridge (memory), grab a new ingredient, walk back, and continue cooking. This happens thousands of times for a single sentence. It wastes energy and time.

2. The Solution: "Operator Fusion" (Batching the Chopping)

The first trick COREY uses is Operator Fusion. Instead of chopping one onion, then walking to the fridge, then chopping a carrot, then walking back, the chef groups the tasks. They chop the onion, the carrot, and the garlic all in one go before walking to the fridge.

The Analogy: Imagine a delivery driver who usually drops off one package, drives back to the warehouse, gets another, and drives back out. Fusion is like loading the truck with 10 packages and making one efficient trip.

3. The Smart Scheduler: The "Entropy" Traffic Cop

The hard part is knowing which tasks to group together. If you group too many, the truck gets too heavy and breaks down. If you group too few, you waste gas.

The Old Way: The driver just groups 3 items every time, no matter what.
The COREY Way: COREY uses a metric called Entropy. Think of entropy as a measure of "chaos" or "surprise" in the data.
- If the data is very predictable (low entropy), the chef knows exactly what's coming. They can safely group a huge batch of tasks together because the risk of error is low.
- If the data is chaotic and full of surprises (high entropy), the chef needs to be careful. They group fewer tasks to avoid mistakes.
- The Metaphor: Imagine a traffic cop at a busy intersection. If the traffic is smooth and predictable, they let a long line of cars pass through the green light at once. If the traffic is erratic and dangerous, they let only a few cars through at a time to keep everyone safe. COREY is that traffic cop, deciding exactly how many tasks to bundle together based on how "chaotic" the data is at that moment.

4. The Secret Sauce: The "Hadamard Shuffle" (Smoothing the Outliers)

Sometimes, the data has "outliers"—weird, massive spikes in numbers that can crash the system (like a sudden, massive spike in temperature that breaks a thermometer). These spikes make it dangerous to group tasks together (fusion) because the system might get confused.

The Fix: COREY uses a mathematical trick called a Hadamard Transform.
The Analogy: Imagine you have a bag of marbles where one is a giant boulder and the rest are tiny pebbles. If you try to pour them through a narrow pipe, the boulder gets stuck.
- COREY takes the bag and shuffles the marbles. It doesn't change the total weight (the information is still there), but it breaks that giant boulder into smaller, manageable chunks and spreads them out evenly among the pebbles.
- Now, the "pipe" (the processor) can handle the flow smoothly. This "shuffling" makes the data less "spiky" and allows the chef to group more tasks together without fear of a crash.

5. The Result: Faster, Smoother, and Cooler

By combining these two ideas:

Smart Grouping: Only grouping tasks when the data is calm enough to handle it.
Smoothing: Breaking up the "boulders" so they don't cause traffic jams.

COREY achieves:

Less Walking: The processor spends less time fetching data from memory.
Faster Cooking: Tasks get done quicker, especially with very long books (long sequences).
Lower Energy: Less movement means less power consumption.

The Bottom Line

The paper admits that this is currently a "prototype" (a working model in a lab, not yet a fully commercial product). They tested it on synthetic data and found that it consistently reduces the time the AI spends waiting for memory and speeds up the processing.

In short: COREY is a smart kitchen manager that shuffles the ingredients to make them easier to handle and then tells the chef exactly how many dishes to prep at once, ensuring the kitchen runs at maximum speed without burning the food or breaking the stove.

1. Problem Statement

Selective State Space Models (SSMs), such as the Mamba family, offer linear-time sequence modeling ( $O(L)$ ), making them superior to Transformers for long-context inference. However, practical deployment on GPUs faces two critical bottlenecks:

Memory-Bound Execution: Current SSM implementations decompose selective scans into fragmented kernels (pointwise and reduction operations). This leads to frequent global memory traffic for intermediate tensor materialization, turning the critical path into a memory-bound workload rather than a compute-bound one.
Activation Outliers & Quantization Instability: SSM internals often exhibit heavy-tailed activation distributions with significant outliers. These outliers destabilize low-bit quantization (e.g., W4A8, W8A8) and limit the "safe depth" of operator fusion, as aggressive fusion in regions with high outlier concentration can amplify numeric sensitivity and degrade model quality.

Existing solutions often treat operator fusion (systems optimization) and activation smoothing (quantization/numerical stability) as separate problems.

2. Methodology: The COREY Framework

The authors propose COREY, a prototype framework that unifies memory-aware operator fusion with Hadamard-based feature reparameterization. The core innovation is using Activation Entropy as a runtime scheduling signal to bridge distribution shaping and systems optimization.

A. Entropy-Guided Operator Fusion

Instead of using static fusion rules or fixed-depth heuristics, COREY dynamically determines fusion boundaries and tile sizes based on the Shannon entropy of activation tensors.

Entropy Estimation: Activations are binned into fixed-width histograms to estimate entropy ( $H$ ).
Fusion Score: A region is fused if a composite score $S(R)$ exceeds a threshold $\tau$ :
$S(R) = \alpha \hat{H}(R) + \beta \widehat{AI}(R) - \gamma \hat{M}(R)$
Where $\hat{H}$ is normalized entropy, $\widehat{AI}$ is arithmetic intensity, and $\hat{M}$ is estimated memory traffic.
Dynamic Scheduling: High-entropy regions (where activation mass is evenly distributed) allow for larger tiles and deeper fusion. Low-entropy regions (prone to outliers) trigger smaller tiles to prevent register pressure and over-fusion.

B. Fused Hadamard Reparameterization

To address the heavy-tailed activation problem, COREY integrates a Hadamard transform directly into the linear projection layers.

Mechanism: For a linear projection $y = Wx$, the method inserts a Hadamard transform pair: $y = (WH^\top)(Hx) = W'\hat{x}$ .
Fusion: The Hadamard transform ($Hx$), the projection ( $W'$ ), and adjacent element-wise operators are fused into a single kernel, avoiding extra memory passes.
Effect: The Hadamard rotation redistributes heavy-tailed energy across channels, smoothing the activation distribution. This increases the empirical entropy and reduces the magnitude of peak coordinates (outliers), thereby expanding the safe operating region for aggressive fusion and low-precision execution.

C. Theoretical Guarantees

The paper provides theoretical proofs supporting the approach:

Entropy Growth: Under doubly-stochastic mixing (approximated by Hadamard rotation over fixed bins), the histogram entropy strictly increases if the pre-rotation distribution is non-uniform.
Fusion Depth Bound: A derived bound shows that fusion depth is limited by shared memory constraints, justifying the need for entropy-guided tiling to balance depth and resource usage.
Stability: Hadamard rotation preserves total variance ( $\ell_2$ norm) but reduces the peak coordinate magnitude ( $\ell_\infty$ norm) by a factor of $\sqrt{d}$ , tightening clipping bounds for low-bit quantization.

3. Key Contributions

Entropy-Driven Criterion: Introduction of a dynamic, runtime-friendly entropy signal to select fusion boundaries and tile granularity in SSM kernels, moving beyond static graphs.
Fused Hadamard Layer: Design of a reparameterization layer that embeds forward/inverse Hadamard transforms into projection blocks while preserving functional equivalence and reducing peak-coordinate concentration.
Unified Link: Explicitly connecting entropy gain (from Hadamard smoothing) to the feasibility of aggressive fusion and low-precision execution.
Prototype Validation: A reproducible prototype study demonstrating consistent reductions in proxy latency, improved throughput, and lower DRAM traffic compared to unfused and static-fusion baselines.

4. Experimental Results

The study was conducted as a controlled prototype using synthetic heavy-tailed activations and validated against Mamba checkpoints (370M, 1.4B, 2.8B) on a WSL2/CUDA environment.

Performance Gains (FP16):
- Throughput: COREY consistently outperformed "No Fusion" and "Static Fusion" baselines. For ultra-long sequences (64k), throughput increased from 751 tok/s (Static) to **840 tok/s** (COREY).
- Latency: Latency was reduced by approximately 10-15% across all sequence lengths compared to static fusion.
- Memory Traffic: DRAM bytes per token were significantly reduced (e.g., from 102.4 to 90.6 in ultra-long sequences), confirming the method acts primarily by reducing memory materialization.
Quantization Stability:
- In W8A8 and W4A8 settings, the Hadamard-reparameterized COREY showed lower "diagnostic proxy" degradation scores compared to static fusion. This indicates better stability in low-entropy, outlier-dominated regions.
- The method successfully mitigated the "over-fusion" penalty in low-entropy regions while maximizing gains in high-entropy regions.
Scalability:
- Speedups increased with sequence length, confirming the strategy targets memory-bound bottlenecks dominant in long contexts.
- Multi-GPU data-parallel sharding demonstrated near-linear speedup (2x on 2 GPUs) without accuracy loss.
Checkpoint Validation:
- Cross-hardware validation (RTX 3070 vs. RTX 3090) confirmed that task scores (NarrativeQA, Qasper, etc.) and perplexity were platform-independent.
- Note: The authors explicitly state these are prototype results. The "quality drop" metric is a synthetic diagnostic proxy, not a full perplexity benchmark, and the study does not claim deployment-grade hardware gains yet.

5. Significance and Future Work

Significance: COREY demonstrates that activation geometry (specifically entropy and outlier distribution) can be used as a direct control signal for systems optimization (fusion and tiling). It proves that smoothing activations via Hadamard rotation is not just a quantization trick but a systems enabler that allows for deeper, more efficient kernel fusion.
Limitations:
- Current results are based on a Python-level prototype; real GPU kernel timings (Triton) are not yet fully integrated into the main performance tables.
- Quantization support is limited to specific Mamba variants (Quamba) due to compatibility issues with generic AWQ/GPTQ stacks.
- The "quality drop" metric is a proxy; full perplexity and downstream task accuracy comparisons on real fused kernels are future work.
Future Directions: The authors plan to extend the benchmark matrix to broader LongBench samples, implement full Triton fused kernels, and conduct real hardware power measurements to quantify energy savings.

In summary, COREY offers a novel, theoretically grounded approach to optimizing SSMs by treating activation distribution properties as first-class citizens in the compiler scheduling process, promising significant efficiency gains for long-context AI inference.