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

1. Il Problema

I Modelli di Spazio di Stato (SSM), in particolare la famiglia Mamba, offrono una complessità lineare $O(L)$ per la modellazione di sequenze, rendendoli ideali per contesti lunghi e carichi di lavoro in streaming rispetto ai Transformer (che hanno complessità quadratica $O(L^2)$). Tuttavia, le implementazioni pratiche su GPU soffrono di due limiti critici:

1. Colli di bottiglia nella memoria (Memory-Bound): Le pipeline di scansione selettiva sono spesso decomposte in kernel frammentati che richiedono frequenti accessi alla memoria globale per materializzare stati intermedi. Questo trasforma il percorso critico in un'operazione limitata dalla larghezza di banda della memoria DRAM, non dalla potenza di calcolo (FLOPs).
2. Outlier nelle attivazioni: Gli SSM presentano spesso statistiche "a coda pesante" (heavy-tailed) nelle attivazioni, specialmente nei canali di output. Questi outlier destabilizzano la quantizzazione a basso bit (es. W4A8, W8A8) e aumentano la sensibilità numerica, limitando la profondità sicura della fusione dei kernel (operator fusion).

2. Metodologia: Il Framework COREY

COREY è un framework prototipale che unisce la fusione di operatori guidata dalla memoria con una riparametrizzazione basata sulla trasformata di Hadamard. L'approccio si articola in tre componenti principali:

A. Fusione Guidata dall'Entropia (Entropy-Guided Fusion)

Invece di utilizzare regole statiche o euristiche fisse per decidere quali operatori fondere, COREY utilizza l'entropia di Shannon delle attivazioni come segnale di scheduling a runtime.

• Stima dell'Entropia: Viene calcolata l'entropia empirica utilizzando istogrammi a larghezza fissa.
• Criterio di Fusione: Un punteggio di fusione $S(R)$ combina entropia normalizzata, intensità aritmetica e stime del traffico di memoria.
  $$S(R) = \alpha \hat{H}(R) + \beta \widehat{AI}(R) - \gamma \hat{M}(R)$$
  Se il punteggio supera una soglia $\tau$ e i vincoli di risorse (registri, memoria condivisa) sono soddisfatti, gli operatori vengono fusi.
• Vantaggio: L'entropia funge da ponte tra la geometria delle attivazioni e la fattibilità della fusione. Regioni ad alta entropia (distribuzioni più uniformi) permettono tile più grandi e fusione più aggressiva, mentre regioni a bassa entropia richiedono tile più piccoli per evitare instabilità.

B. Riparametrizzazione con Trasformata di Hadamard

Per mitigare gli outlier e abilitare una fusione più profonda, il paper introduce un layer di Hadamard fuso:

• Meccanismo: Si inserisce una coppia di trasformate (diretta e inversa) all'interno dei blocchi di proiezione lineare: $y = W H^\top H x = W' \hat{x}$.
• Effetto: La matrice di Hadamard normalizzata ($H$) ridistribuisce l'energia delle code pesanti tra i canali, "smussando" la distribuzione delle attivazioni.
• Teorema di Stabilità: Dimostrato teoricamente che la trasformata preserva la norma $L_2$ ma riduce la magnitudine del picco della coordinata ($L_\infty$) di un fattore $\sqrt{d}$. Questo riduce la necessità di clipping aggressivo durante la quantizzazione a basso bit.
• Integrazione: La trasformata viene fusa direttamente con le proiezioni lineari e gli operatori elementari per evitare passaggi aggiuntivi in memoria.

C. Co-Design Entropia-Tiling

L'entropia non guida solo i confini di fusione, ma anche la dimensione dei tile (tiling).

• Alta entropia: Distribuzione uniforme delle attivazioni $\rightarrow$ Tile più grandi (migliore riutilizzo dei dati, maggiore intensità aritmetica).
• Bassa entropia: Presenza di outlier o concentrazione $\rightarrow$ Tile più piccoli (riduzione della pressione sui registri, prevenzione della fusione eccessiva).

3. Contributi Chiave

1. Criterio Dinamico di Fusione: Introduzione di un criterio guidato dall'entropia per la selezione dinamica dei confini di fusione nei kernel SSM, superando le limitazioni delle strategie statiche.
2. Layer Hadamard Fuso: Progettazione di un layer che incorpora le trasformate di Hadamard nei blocchi di proiezione, preservando l'equivalenza funzionale ma migliorando la distribuzione delle attivazioni per la fusione e la quantizzazione.
3. Collegamento Teorico-Pratico: Dimostrazione che l'aumento dell'entropia indotto dalla rotazione di Hadamard espande la regione operativa sicura per l'esecuzione a bassa precisione, collegando la riduzione degli outlier alla fattibilità della fusione aggressiva.
4. Validazione Prototipale: Convalida del design attraverso uno studio controllato su attivazioni sintetiche a coda pesante, con metriche di latenza, throughput e traffico DRAM, distinguendo chiaramente tra risultati del prototipo e claim di produzione.

4. Risultati Sperimentali

Lo studio è stato condotto su un prototipo Python/PyTorch (senza kernel Triton finali) su workstation CPU, simulando il comportamento del scheduler, con validazioni aggiuntive su checkpoint reali (Mamba-370M, Mamba-1.4B) in ambiente WSL2.

• Prestazioni End-to-End:

  • Throughput: COREY mostra miglioramenti costanti nel throughput rispetto alle baseline "No Fusion" e "Static Fusion" (fusione fissa di gruppo 3).
  • Contesto Lungo: I guadagni aumentano con la lunghezza della sequenza. Per contesti ultra-lunghi (64k), COREY raggiunge un throughput di 840.5 k tok/s contro i 751.4 k tok/s della fusione statica e 626.9 k tok/s senza fusione.
  • Traffico DRAM: Riduzione significativa dei byte DRAM per token (es. da 123.0 a 90.6 byte/token per la bucket "Ultra-long"), confermando che il metodo riduce principalmente il traffico di memoria.
  • Latenza: Riduzione della latenza media (es. da 104.54 ms a 77.97 ms per la bucket Ultra-long).

• Quantizzazione e Stabilità:

  • In modalità W8A8 e W4A8, la riparametrizzazione Hadamard riduce il "diagnostic proxy" di degradazione della qualità rispetto alla fusione statica. Questo suggerisce che il controllo dei confini guidato dall'entropia evita la fusione eccessiva in regioni a bassa entropia dominata da outlier.
  • I risultati di qualità (perplessità su WikiText-103, PG19 e task LongBench) sono stati validati su checkpoint reali, confermando che la fusione non degrada le prestazioni del modello.

• Scalabilità Multi-GPU:

  • Test su un server remoto con 4 GPU RTX 3090 hanno mostrato uno speedup quasi lineare (circa 2x con 2 GPU) per l'inferenza data-parallel, senza degradazione della qualità per campione.

5. Significato e Impatto

Il lavoro di COREY è significativo perché:

• Supera l'approccio euristico: Trasforma la fusione degli operatori da un problema statico a uno dinamico e adattivo, basato su segnali statistici reali (entropia) delle attivazioni.
• Abilita la quantizzazione a basso bit: Dimostra che la riduzione degli outlier tramite Hadamard non è solo utile per la precisione, ma è un prerequisito per permettere una fusione di kernel più profonda senza collasso numerico.
• Ottimizzazione Hardware-Aware: Propone un metodo che si adatta alle limitazioni hardware (memoria condivisa, registri) in tempo reale, massimizzando l'intensità aritmetica laddove le attivazioni lo permettono.
• Impatto Ambientale: Riducendo il traffico di memoria DRAM, il metodo potrebbe diminuire il consumo energetico dell'inferenza SSM a lungo contesto, rendendo più sostenibile l'uso di questi modelli per applicazioni reali.

Limitazioni e Futuro:
L'autore sottolinea che i risultati attuali sono basati su un prototipo e su un subset limitato di benchmark (LongBench su 4 task). I claim di produzione su larghezza di banda reale e guadagni di qualità su dataset completi richiedono l'implementazione completa di kernel fusi in Triton e test su un ventaglio più ampio di modelli e task. Tuttavia, la validazione cross-hardware e la coerenza dei risultati di checkpoint confermano la solidità del meccanismo proposto.