Mambalaya: Einsum-Based Fusion Optimizations on State-Space Models

Abstract: Mamba is an emerging, complex workload with various short-range and long-range dependencies, nonlinearities, and elementwise computations that are unable to run at near-peak speeds on modern hardware. Specifically, Mamba's complex dependency graph makes fusion across its full operator cascade difficult, leaving substantial inter-operator memory traffic on the table. To address these challenges, we propose Mambalaya, a novel reconfigurable accelerator that leverages fusion to overcome the limitations of Mamba. We use the recently proposed cascade-of-Einsums abstraction to characterize Mamba's full computational structure, then apply the extended Einsum framework to systematically explore inter-Einsum fusion opportunities. This principled approach yields a series of fusion mappings that reduce off-chip inter-Einsum traffic. These mappings are supported by the underlying Mambalaya architecture. Mambalaya achieves a layer performance speedup of 4.9$\times$ for prefill and 1.9$\times$ for generation over MARCA. In prefill-dominated scenarios, it achieves up to 1.5$\times$ over a recent fine-grained, memory-aware fusion accelerator for Mamba.

• The paper introduces full fusion of Mamba's tensor operations using extended Einsum notation, reducing fusion groups from 24 to one.
• It outlines a systematic taxonomy for fusion, classifying operator dependencies into RI, RSb, RSp, and RD types to maximize on-chip data reuse.
• It presents a reconfigurable accelerator architecture achieving up to 5.79× speedup and reduced memory traffic for state-space models in LLMs.

Mambalaya: Einsum-Based Fusion Optimizations on State-Space Models

Abstract and Motivation

The paper "Mambalaya: Einsum-Based Fusion Optimizations on State-Space Models" (2604.03829) introduces a novel, reconfigurable accelerator architecture—Mambalaya—for Mamba, an advanced State Space Model (SSM) widely used for LLMs. Unlike Transformers, Mamba's computational structure is characterized by intricate operator dependencies, numerous nonlinearities, and substantial elementwise computations. These properties significantly impede mapping and optimization for existing accelerator topologies, resulting in suboptimal hardware utilization and excessive memory traffic. The research leverages Einstein summation notation (Einsum) and the EDGE framework to systematically characterize and optimize these complex dependencies via fusion, realizing nearly minimal off-chip tensor traffic.

Einsum Formalism for Mamba Computation

The foundation for Mambalaya is casting all Mamba computation within the extended Einsum abstraction, enabling both a formal description of tensor algebra and a systematic, tool-based exploration of optimization strategies. This formalism supports higher-order functions, generational ranks for iterative computation, and bulk elementwise/nonlinear operations. Key features include:

• User-defined operations: Elementwise nonlinearities such as $\log$, $\exp$, $\sqrt{}$.
• Generational ranks: Recurrence along iterative token indices, foundational to SSM behavior.

The extended Einsum formalism empowers separation of concerns: algorithm specification, mapping (including fusion choices), and hardware binding. This multilevel abstraction is critical in reasoning about optimal fusion and data locality.

Figure 1: Overview of the cascade execution flow for Mamba, highlighting the rich structure of operator dependencies and tensor flows.

Memory Traffic and Roofline Analysis

Mamba workloads, unlike Transformer-based models, are heavily memory-bound due to complex dependency graphs and a high fraction of non-GEMM operations. Roofline analyses demonstrate that:

• Unfused executions alternate between compute-bound and memory-bound phases during prefill, with decode being consistently memory-bound.
• Ideal fusion—removing all inter-Einsum memory traffic—yields significant increases in computational intensity and execution speed: speedup factors of up to $5.79\times$ (prefill) and $3.8\times$ (token generation).

Figure 2: Roofline analysis showing memory-bound and compute-bound phases for unfused and fused Mamba executions.

Taxonomy and Algorithmic Recipe for Einsum Fusion

A theoretical taxonomy for fusion is introduced, classifying fusion opportunities between pairs of Einsums:

• Rank Isomorphic (RI) Fusion: Identical iteration spaces, minimal tensor footprint.
• Rank Subsetted (RSb) Fusion: Upstream iteration space is a superset, reduction.
• Rank Supersetted (RSp) Fusion: Downstream iteration space is a superset, broadcast.
• Rank Disjointed (RD) Fusion: Neither subset nor superset, both reduction and broadcast.

The greedy fusion algorithm systematically identifies fusion groups by evaluating iteration space relationships and intermediate tensor lifetimes, maximizing on-chip reuse and minimizing global memory traffic.

Figure 3: RI Fusion – example of rank isomorphic fusion mapping for an Einsum pair.

Application of Fusion Taxonomy to Mamba

By applying the fusion taxonomy and greedy stitching algorithm, the paper demonstrates that all Mamba tensor algebra can be expressed as Einsums and fused across the entire operator cascade. This drastically reduces the number of fusion groups from 24 (unfused) to 3 (fully fused), excepting cases where RD fusion necessitates careful tiling or partial recomputation.

Key findings:

• Fusion is viable for all operator phases, including complex non-GEMM regions previously thought impractical for fusion.
• Fully-fused Mamba realizes only one fusion group over the entire cascade.

  Figure 4: Fusion opportunities in the Mamba cascade, highlighting data dependencies and fusion groupings.

  

  Figure 5: Roofline utilization over time comparing various fusion strategies, illustrating latency and fusion-group reduction.

Mambalaya Accelerator Architecture

Mambalaya introduces a reconfigurable hardware architecture, featuring:

• 2D PE array: Supports both GEMM and low-intensity operations with 65,792 PEs in 2D and 1D modes.
• 1D PE subarray: Specialized for elementwise operations and efficient broadcast to the main array.
• Reconfigurable network: Dynamically adapts to fusion group requirements for optimal data locality and compute utilization.
• Support for nonlinear functions: Dedicated units for log, max, SiLU, exponential, enabling efficient fusion of bulk elementwise operations.

Mapping and binding strategies are aligned with fusion group characteristics to optimize memory locality and execution throughput.

Figure 6: Mambalaya accelerator 2D mode architecture schematic.

Numerical Results and Performance Evaluation

Simulations using Timeloop demonstrate robust numerical evidence for Mambalaya's efficacy:

• Layer speedups: $4.9\times$ for prefill and $1.9\times$ for token generation compared to MARCA, the previous SoTA Mamba accelerator.
• End-to-end scenarios: Geometric mean speedup of $3\times$ (prefill-heavy) and $1.3\times$ (decode-heavy) against MARCA-like and Geens-like baselines, respectively.
• In prefill-dominated scenarios, Mambalaya outperforms fine-grained, memory-aware fusion accelerators by up to $1.5\times$.

  Figure 7: End-to-end performance across Mamba variants for mamba-370m, visualizing ideal and achieved results.

  

  Figure 8: Comparison of best Mambalaya variant to MARCA and Geens-Like accelerators—showing superiority in compute utilization and latency.

Theoretical and Practical Implications

This work presents a decisive advance in the methodology for tensor algebra accelerator design:

• Generality of Einsum fusion: The formalism and algorithmic recipe enable systematic fusion optimization for any workload expressible as an EDGE Einsum cascade, transcending Mamba and SSMs.
• Reduced memory traffic: Robust mapping strategies drastically minimize inter-Einsum tensor traffic, directly increasing computational intensity and energy efficiency.
• Design-space exploration: The separation between algorithm specification and hardware mapping facilitates integration with existing architecture exploration tools.

On the practical front, Mambalaya sets a new standard for high-throughput, memory-efficient execution of SSM-based LLMs, with implications for broader deployment of attention-free models, hybrid architectures, and efficient LLM inference.

Conclusion

Mambalaya demonstrates for the first time that full fusion is achievable across the complex operator cascade of Mamba, offering a radically reduced memory footprint and maximized compute utilization. The combination of a systematic fusion taxonomy and reconfigurable hardware is both general and potent, enabling mechanical application of fusion strategies to any tensor workload. The bold claim—fusing any Einsum cascade by observing tensor indices—redefines the optimization landscape for LLM accelerators. Future directions include automatic design-space exploration trading inter- and intra-tensor reuse, broader application to emerging model architectures, and continued refinement of hardware support for dynamic fusion patterns.