MAS-ATTENTION: Memory-Aware Stream Processing for Attention Acceleration on Resource-Constrained Edge Devices

MAS-ATTENTION: Memory-Aware Stream Processing for Attention Acceleration on Resource-Constrained Edge Devices

Speakers: Mohammadali Shakerdargah, Shan Lu, Chao Gao, Di Niu

Conference: MLSys 2025

YouTube: https://www.youtube.com/watch?v=None

Overview

The proliferation of large language models and other foundation models has cemented the transformer architecture and its core component, attention mechanisms, as indispensable elements across diverse AI applications, from natural language processing to computer vision, diffusion models, and image restoration. However, deploying these computationally intensive models on resource-constrained edge devices – such as cell phones, IoT devices, or custom neural processing units (NPUs) and tensor processing units (TPUs) – presents significant challenges. Existing acceleration techniques, often tailored for high-end GPUs, fail to adequately address the unique memory and computational limitations of edge hardware.

This talk introduces MAS-ATTENTION, a novel framework designed to accelerate attention computations specifically for these resource-constrained edge devices. Presented by Professor Di Niu from the University of Alberta, the work tackles the critical issues of maximizing compute utilization, minimizing memory footprint, and balancing workloads between heterogeneous processing units common in edge accelerators. By adopting a memory-aware stream processing approach, MAS-ATTENTION aims to bridge the performance gap for AI workloads on the edge, enabling more efficient and widespread deployment of advanced AI capabilities beyond data centers.

The significance of MAS-ATTENTION lies in its targeted approach to a largely underserved segment of AI deployment. While solutions like FlashAttention have revolutionized attention performance on powerful GPUs, their reliance on massive parallelism and specific hardware architectures makes them unsuitable for the modest capabilities of typical edge NPUs or TPUs. MAS-ATTENTION offers a practical and effective solution for optimizing the complex attention pipeline, demonstrating substantial latency reductions and energy savings on real-world edge hardware, thereby advancing the feasibility of sophisticated on-device AI.

Background

▶ Watch: Introduction: Attention acceleration for resource-constrained edge devices (0:00)

The transformer architecture, with its self-attention mechanism, has become the de facto standard for state-of-the-art models across various AI domains. The attention operation itself, typically involving two matrix multiplications (Q times K-transpose and P times V) and a softmax activation, is a computational bottleneck. Accelerating this operation is paramount for improving the performance and efficiency of transformer-based models.

Much of the prior work in attention acceleration has focused on high-end general-purpose GPUs (GPGPUs). FlashAttention, for instance, has been a groundbreaking series of techniques designed to optimize attention on GPUs like NVIDIA's H100 and A100. Its key innovation lies in breaking down the attention computation into smaller tiles and scheduling them efficiently to leverage the massive parallelism of CUDA cores, ensuring data locality by keeping intermediate results on-chip (SRAM). FlashAttention 3, for example, further enhances performance by overlapping heterogeneous operations like matrix multiplication (matmul) and softmax across warp groups. However, these methods are inherently tied to the GPU's architecture and its extensive parallelism, making them unsuitable for direct transfer to the specialized, often more constrained, NPUs and TPUs found in edge devices.

Recognizing this gap, some efforts have attempted more general acceleration schemes. Flat, an earlier approach, aimed at a broader range of hardware including TPUs and NPUs. Its core idea was to keep all intermediate results of the matmul-softmax-matmul sequence on-chip to avoid costly off-chip memory access. While conceptually simple and effective for data locality, Flat's sequential execution of these operations significantly underutilized the available compute units, leading to suboptimal performance.

A more recent and sophisticated development is FuseMax, a collaboration between Berkeley, Tsinghua, MIT, and Nvidia. FuseMax pushed the boundaries of utilization by representing the attention workload as an "extended sum" and the softmax execution as a "one-pass sum cascade." This representation allowed for manual scheduling of operations onto heterogeneous compute units, typically a 2D array for matrix operations (Mac units) and a 1D array for vector operations (Vec units), thereby maximizing the utilization of both. FuseMax demonstrated excellent performance compared to Flat by addressing the compute utilization challenge.

Despite these advancements, a critical challenge remained for truly resource-constrained edge devices. Even FuseMax, while effective, implicitly assumes a certain level of parallel processing capability and on-chip buffer size. Real-world edge NPUs or custom TPUs often feature extremely limited cores, minimal on-chip buffer memory, and highly specialized, often heterogeneous, compute units (e.g., a single Mac unit and a single Vec unit). In such environments, achieving high efficiency requires not just general data locality or improved utilization, but a finely-tuned strategy for workload balancing, pipelining, and proactive memory management that can operate effectively under severe memory and compute constraints. The problem, therefore, is to develop an attention acceleration framework specifically tailored to these ultra-resource-constrained edge scenarios, going beyond what current GPU-centric or even general NPU solutions offer.

Key Findings

▶ Watch: Review of FuseMax: enhancing utilization with extended sum (6:00)

MAS-ATTENTION presents a novel and highly effective solution for accelerating attention mechanisms on resource-constrained edge devices, addressing the limitations of prior GPU-centric and more general NPU-focused approaches. The key findings and contributions of this work can be summarized as follows:

1. Novel Stream Processing for Heterogeneous Edge Units: MAS-ATTENTION introduces a sophisticated stream processing (pipelining) paradigm that can efficiently pipeline the three core attention operations (two matrix multiplications and one softmax) even on highly constrained hardware configurations, such as those featuring only one Mac unit and one Vec unit. This stands in contrast to methods that rely on massive parallelism or sequential execution.

1. Multi-tiered Tiling and Workload Balancing: The framework employs a multi-tiered tiling strategy, applying fine-grain sub-matrix tiling for the matrix multiplications (QxKᵀ and PxV) and row-wise tiling for the softmax operation. Crucially, it integrates a search strategy, leveraging Monte Carlo Tree Search (MCTS) for tiling factor optimization and a Genetic Algorithm (GA) for compute ordering (in simulation), to perfectly balance the workload between the heterogeneous Mac (multiply-accumulate) and Vec (vector operations like softmax) units, maximizing their concurrent utilization.

1. Proactive Buffer Management: To combat the severe memory limitations of edge devices, MAS-ATTENTION implements a proactive buffer management mechanism. This strategy intelligently prioritizes the storage of intermediate softmax results over potentially reloadable input tensors (K and V). In scenarios of buffer scarcity, it allows input tensors to be overwritten, with the understanding that they can be reloaded from off-chip memory if needed, thereby preventing costly buffer spills and ensuring continuous on-chip processing.

1. Superior Performance on Edge Hardware: Through extensive simulations and real-world experiments on a Huawei MatePad Pro 13, MAS-ATTENTION demonstrates significant performance gains. In simulation, it achieves a 1.27x speedup over FuseMax, a state-of-the-art baseline for general NPUs, while incurring only a marginal 10-11% higher energy cost. On the real device, it delivers impressive 1.3x to 3x speedups specifically for the attention computation and an overall end-to-end latency saving of 6% for a 4K token length, with gains increasing for longer token sequences.

1. Energy Efficiency for On-Chip Memory: Beyond latency, MAS-ATTENTION also demonstrates notable energy savings, particularly in L1 memory access. This indicates that its memory-aware design effectively reduces the energy consumption associated with frequent data movement on the most critical and energy-sensitive on-chip caches.

These findings collectively establish MAS-ATTENTION as a robust and highly optimized solution for deploying transformer-based models on the rapidly expanding ecosystem of resource-constrained edge AI devices, overcoming critical hardware bottlenecks through intelligent algorithmic and systems co-design.

Technical Deep Dive

▶ Watch: Remaining challenges: limited resources and workload balancing on edge (8:00)

The attention mechanism, core to transformers, involves a sequence of operations: $C = QK^T$, $P = \text{softmax}(C)$, and $O = PV$. This essentially translates to two matrix multiplications (matmuls) and one softmax operation. MAS-ATTENTION's technical innovation lies in its ability to effectively pipeline and manage these operations on edge devices characterized by limited compute units (e.g., one Mac unit for matmuls, one Vec unit for softmax/activations) and minimal on-chip buffer.

The fundamental idea is stream processing or pipelining. Instead of executing the three operations sequentially for an entire attention block, MAS-ATTENTION breaks down the tensors into small tiles and pipelines their execution. The challenge is not merely to pipeline but to do so perfectly, ensuring maximum utilization of heterogeneous units and avoiding stalls or buffer spills.

The core of MAS-ATTENTION's pipelining strategy is a novel overlapping scheme for the two matmuls and one softmax. During the "regular phase" (after an initial warm-up), the Vec unit is actively processing the softmax for the current data chunk. Concurrently, the Mac unit is engaged in a dual task: it completes the second matmul ($PV$) for the previous chunk while simultaneously initiating the first matmul ($QK^T$) for the next chunk. This sophisticated interleaving allows both the Mac and Vec units to remain highly utilized, minimizing idle time and maximizing throughput. The design also accounts for "warm-up" and "final" stages to handle the initial and terminal chunks of data, ensuring a seamless flow.

To enable this intricate pipelining, a multi-tiered tiling strategy is crucial. For the matrix multiplications ($QK^T$ and $PV$), MAS-ATTENTION employs fine-grain tiles, breaking down the matrices into smaller sub-matrices. This allows for flexible scheduling and granular data movement. In contrast, for the softmax calculation, a row-wise tiling strategy is adopted. While theoretically possible to tile softmax at a finer granularity (e.g., partial softmax aggregation), the row-wise approach simplifies the semantic correctness of the softmax operation, which inherently operates across an entire row of values. This distinction in tiling granularity for different operation types is key to optimizing both compute and memory access patterns.

The heterogeneous nature of edge accelerators, typically comprising dedicated Mac units for multiply-accumulate operations and Vec units for vector operations (like softmax, activations), necessitates careful workload balancing. The goal is to ensure neither unit becomes a bottleneck. MAS-ATTENTION addresses this by employing an intelligent search strategy. In simulation, this involves a combination of Monte Carlo Tree Search (MCTS) to explore and optimize the tiling factors for different tensors, and a Genetic Algorithm (GA) to find the optimal compute ordering. MCTS helps in efficiently navigating the vast search space of tiling configurations, while GA refines the sequence of operations to minimize inter-unit dependencies and maximize parallelism. For real-world deployment on the Huawei MatePad Pro 13, a simplified grid search was used due to practical constraints, demonstrating that even a less complex search can yield significant benefits once the fundamental pipelining and tiling strategies are in place. The objective of these search algorithms is to find a configuration where the Mac and Vec units are as busy as possible, minimizing the time one unit waits for the other.

Perhaps one of the most critical technical contributions for resource-constrained environments is proactive buffer management. Edge devices typically have very limited on-chip buffer (SRAM) compared to GPUs. To prevent costly buffer spills (where data must be written back to slower off-chip DRAM), MAS-ATTENTION prioritizes which intermediate tensors remain in the on-chip buffer. Specifically, intermediate results from the softmax calculation (e.g., the matrix P, which is softmax(C)) are given higher priority. Input tensors like K and V are considered "expendable" in a buffer-constrained situation. The strategy defines two main scenarios:

1. If the buffer is insufficient, the intermediate result P can overwrite the input tensor V. When V is needed later for the $PV$ matmul, it is simply reloaded from off-chip memory.
2. Similarly, P can overwrite K. If K is needed again (e.g., for a subsequent $QK^T$ operation), it is reloaded.

This "recalculate/reload" strategy, while potentially introducing some minor recomputation or memory access overhead, is designed to avoid the much larger performance penalty of buffer spills, which can severely degrade performance on memory-limited devices. By intelligently managing the limited on-chip buffer, MAS-ATTENTION ensures that the most critical, transient intermediate data remains readily accessible, maintaining the integrity of the pipelined execution.

Experimental Setup & Results

▶ Watch: Mass-Attention proposal: stream processing and multi-tiered tiling (10:00)

To thoroughly evaluate MAS-ATTENTION, the researchers conducted a comprehensive set of experiments encompassing both detailed simulations and real-world deployment on an edge device.

Workloads: The evaluation focused on representative AI workloads that heavily utilize attention mechanisms. These included various language models and diffusion models, reflecting the broad applicability of transformers in modern AI. Specific model sizes or configurations were not detailed in the transcript, but the general categories cover prominent use cases for attention.

Simulation Environment: For initial and in-depth analysis, a custom simulation tool was developed. This tool was based on established frameworks like TimeLoop and accelerator models, but critically, it was modified specifically for edge device characteristics. This involved tailoring the cost model to reflect the architectural nuances, memory hierarchies, and compute unit capabilities (e.g., Mac and Vec units) typical of edge NPUs/TPUs, which differ significantly from high-end GPUs. The simulation provided a controlled environment to explore various tiling strategies and workload balancing algorithms. It was within this simulation that the more advanced MCTS for tiling factor search and Genetic Algorithm for compute ordering were primarily employed, allowing for refined optimization.

Real-world Device Evaluation: A crucial aspect of the validation involved deploying MAS-ATTENTION on a tangible edge device. The chosen platform was the Huawei MatePad Pro 13, a tablet device, indicating a collaboration with Huawei. This device likely features a custom NPU or accelerator, representative of the target hardware for MAS-ATTENTION. This real-world experiment provided concrete evidence of the framework's practical efficacy beyond theoretical simulations.

Baselines: The primary baseline for comparison was FuseMax, a state-of-the-art attention acceleration framework developed by a consortium including Berkeley and Nvidia, known for its improved compute utilization over simpler methods like Flat. Other baselines, though not explicitly detailed, were also considered in the evaluation, ensuring a comprehensive comparison against existing techniques.

Metrics: The key performance indicators measured were:

• Latency: The time taken to complete the attention operation or the end-to-end task.
• Energy Consumption: The power consumed during execution, with a particular focus on memory energy.

Headline Results (Simulation):

• Latency Speedup: MAS-ATTENTION consistently outperformed all baselines, including FuseMax, in terms of latency on the simulated edge device architecture.
• Energy Trade-off: While achieving superior latency, MAS-ATTENTION incurred a modest 10-11% higher energy consumption than FuseMax in the simulation environment. This suggests a potential power-performance trade-off in certain configurations, though it's important to note this is a simulation result and might not perfectly reflect all real-world scenarios.
• Quantified Speedup: Specifically, MAS-ATTENTION demonstrated a 1.27x speedup over FuseMax in the simulated environment.

Headline Results (Real Device - Huawei MatePad Pro 13):

• Attention Acceleration: On the real Huawei device, MAS-ATTENTION delivered substantial speedups for the attention operation itself, ranging from 1.3x to 3x. This highlights the significant impact of the proposed techniques on the core bottleneck.
• End-to-End Latency Savings: For an end-to-end workload with a 4K token length, MAS-ATTENTION achieved a 6% latency saving. The presentation noted that this gain would be even more pronounced for larger token lengths, as attention becomes a more dominant factor in overall computation time.
• Energy Savings (L1 Memory): A critical finding in the real-device evaluation was the significant energy saving in L1 memory. The blue part of the energy consumption breakdown graph specifically highlighted MAS-ATTENTION's superiority in reducing L1 memory energy, indicating the effectiveness of its memory-aware design and proactive buffer management in minimizing energy-intensive on-chip data movement.

It's also worth noting the difference in search algorithms used: while MCTS and Genetic Algorithms were employed for refined optimization in simulation, a simplified grid search was used for the real device. This was attributed to practical constraints in real-device deployment, yet still yielded impressive results, underscoring the robustness of MAS-ATTENTION's core pipelining and tiling strategies. The experimental results unequivocally demonstrate MAS-ATTENTION's capability to deliver substantial performance and energy efficiency benefits for attention acceleration on resource-constrained edge devices.

Practical Implications

▶ Watch: Novel pipelining scheme for balancing Mac and Vac units (12:00)

MAS-ATTENTION carries significant practical implications for the burgeoning field of edge AI, particularly for the deployment of advanced machine learning models.

For practitioners and model deployers, MAS-ATTENTION directly addresses a critical bottleneck: the efficient execution of transformer-based models on resource-constrained edge devices. This means that sophisticated models, including larger language models, vision transformers, and diffusion models, which are typically confined to cloud GPUs, can now be realistically deployed on-device. This opens up new possibilities for applications requiring low latency, privacy, and offline operation, such as intelligent assistants on smartphones, real-time analytics in IoT devices, or advanced computer vision in embedded systems. The demonstrated 1.3x to 3x speedups for attention and 6% end-to-end latency saving (for 4K tokens) directly translate to a more responsive user experience and expanded capabilities for edge AI products.

For infrastructure teams and custom accelerator designers, MAS-ATTENTION provides valuable insights into optimizing hardware-software co-design for NPUs and TPUs. The framework's ability to maximize utilization even with minimal, heterogeneous compute units (e.g., one Mac and one Vec unit) offers a blueprint for designing highly efficient, specialized accelerators. The emphasis on proactive buffer management and its impact on L1 memory energy savings underscores the importance of intelligent memory hierarchies and data flow strategies in hardware design. This work suggests that optimizing data movement and on-chip memory usage is as critical as maximizing raw compute power for edge efficiency.

However, MAS-ATTENTION also presents certain tradeoffs and limitations. While simulation results showed impressive latency gains, they also indicated a 10-11% increase in overall energy consumption compared to FuseMax in some scenarios. This implies that while MAS-ATTENTION is highly performant, there might be a subtle power-performance tradeoff to consider depending on the specific application's requirements. For battery-powered devices, this slight increase in total energy might need careful consideration, although the significant L1 memory energy savings on the real device could mitigate this in practice.

Another limitation is the current focus on inference (forward pass). Many advanced AI applications, particularly in federated learning or on-device personalization, require efficient on-device training (backward pass). The authors acknowledge this as a future direction, indicating that MAS-ATTENTION's principles might need adaptation for the increased complexity and memory demands of training workloads. Furthermore, while the multi-tiered tiling strategy is effective, the choice of row-wise tiling for softmax (vs. finer granularity) represents a simplification. Future work could explore more advanced partial softmax aggregation techniques to potentially extract even more parallelism or reduce memory pressure. Finally, the practical constraints of deploying on a real device led to the use of a simplified grid search for workload balancing compared to the more advanced MCTS and genetic algorithms used in simulation. This suggests that the full theoretical potential of MAS-ATTENTION might be even higher if more sophisticated optimization techniques can be integrated into real-world deployment flows.

Despite these considerations, MAS-ATTENTION represents a significant leap forward in making complex attention models viable for the vast and diverse landscape of edge devices, paving the way for more powerful and ubiquitous AI applications.

Key Takeaways

• Targeted Edge Acceleration: MAS-ATTENTION is specifically designed to accelerate transformer attention on resource-constrained edge devices (NPUs, TPUs, mobile, IoT), addressing a critical gap left by GPU-centric solutions like FlashAttention.
• Novel Stream Processing: It employs a sophisticated pipelining strategy that efficiently overlaps two matrix multiplications and one softmax operation, even with minimal heterogeneous compute units (e.g., one Mac and one Vec unit).
• Intelligent Workload Balancing: A multi-tiered tiling strategy (fine-grain for matmuls, row-wise for softmax) combined with advanced search algorithms (MCTS/GA in simulation) optimizes and balances the workload between Mac and Vec units.
• Proactive Buffer Management: To counter limited on-chip memory, MAS-ATTENTION intelligently prioritizes intermediate softmax results, allowing less critical input tensors (K, V) to be overwritten and reloaded, effectively preventing costly buffer spills.
• Significant Performance Gains: Experiments show substantial latency improvements: up to 1.27x speedup over FuseMax in simulation and 1.3x to 3x attention speedup on a real Huawei MatePad Pro 13, leading to 6% end-to-end latency savings for 4K tokens.
• Memory Energy Efficiency: The framework achieves notable L1 memory energy savings on the real device, underscoring its effectiveness in optimizing on-chip data movement and energy consumption.

About the Speaker(s)

The talk was presented by Di Niu, who is a Professor in the Department of Electrical and Computer Engineering (ECE) at the University of Alberta. He is one of the co-authors of the paper, alongside Mohammadali Shakerdargah, Shan Lu, and Chao Gao. His presentation highlighted the team's research into optimizing AI/ML workloads for resource-constrained edge environments.