Introduction
The rapid advancement of Deep Learning (DL) accelerators has led to a significant disparity between computational throughput and memory bandwidth, a phenomenon commonly referred to as the "memory wall." While on-chip compute capabilities continue to scale, DRAM access speeds have not kept pace, often leaving high-performance processing units idle while waiting for data.
To mitigate this, researchers have turned to layer fusion, a technique that keeps intermediate data on-chip to minimize DRAM traffic. However, existing layer fusion strategies often suffer from two primary limitations: a restricted exploration of the vast fusion and tiling design space, and a failure to account for the precise timing of data movement. Current schedulers typically rely on rigid double-buffering patterns that result in "spiky" bandwidth demand—periods of extreme congestion followed by DRAM inactivity.
This paper introduces SoMa, an end-to-end scheduling framework designed to optimize DRAM communication by rethinking both the spatial and temporal aspects of data movement.
The SoMa Framework
SoMa is built upon a novel Tensor-centric Notation that characterizes schedules using six distinct attributes. These attributes define the execution order and the "living duration" of tensors within on-chip buffers, providing a comprehensive mathematical foundation for scheduling exploration.
The framework employs a two-stage exploration engine powered by Simulated Annealing (SA) to navigate the complex optimization space:
- Layer-Fusion Attribute (LFA) Optimization: This stage focuses on the spatial aspects of the schedule, determining optimal layer groupings and tiling sizes. Its primary goal is to minimize the total volume of DRAM traffic.
- DRAM-Link Scheduling Attribute (DLSA) Optimization: This stage performs temporal "surgical strikes," adjusting the timing of DRAM loads and stores. By filling idle bandwidth gaps—such as prefetching weights for future layers while the current layer is computing—DLSA smooths out the spiky demand typical of fused layers.
A centralized Buffer Allocator manages the trade-off between memory dedicated to layer fusion and memory reserved for prefetching, ensuring the hardware's on-chip resources are utilized at peak efficiency.
Experimental Methodology
The researchers evaluated SoMa across a wide spectrum of workloads and hardware configurations:
- Workloads: The suite included classical CNNs (ResNet-50, ResNet-101, Inception-ResNet-v1), irregular networks (RandWire), and Large Language Models (GPT-2 Small and XL).
- Hardware Templates: Simulations were conducted for both "Edge" (16 TOPS, 8MB buffer) and "Cloud" (128 TOPS, 32MB buffer) accelerator profiles.
- Baseline: SoMa was compared against Cocco, a previous state-of-the-art framework for layer fusion.
Key Insights
Addressing Spiky Demand
Layer fusion, while reducing total traffic, often creates highly uneven DRAM demand. The initial tile of a fused group typically requires a massive burst of weights, while subsequent tiles require minimal DRAM interaction. SoMa identifies these "quiet" periods and utilizes them to prefetch data for future operations, effectively hiding memory latency.
Buffer and Bandwidth Synergy
A significant finding of this research is that sophisticated scheduling can allow a larger on-chip buffer to compensate for lower DRAM bandwidth. This suggests a potential shift in hardware design: by utilizing a compiler like SoMa, designers might achieve high performance using more cost-effective memory solutions (e.g., LPDDR) instead of expensive High Bandwidth Memory (HBM).
LLM Bottlenecks
The study highlights a critical distinction in LLM performance. While SoMa provides substantial gains during the "Prefill" stage, the "Decode" stage (token-by-token generation) remains a challenge. Because the Decode stage is compute-sparse, it is almost entirely bound by the speed of loading weights and the KV cache, leaving little room for scheduling-based optimization.
Results and Impact
SoMa demonstrates significant improvements over existing methods:
- Performance: The framework achieved an average 2.11× speedup over Cocco, with gains as high as 2.61× on irregular models like RandWire.
- Energy Efficiency: By eliminating unnecessary DRAM accesses and reducing processor idle time, SoMa reduced energy consumption by an average of 37.3%.
- Hardware Utilization: The generated schedules reached within 3.1% of the theoretical maximum computational utilization for the target hardware.
The SoMa framework has been integrated into a commercial AI accelerator compiler and is available as an open-source project on GitHub.
Limitations and Future Directions
Despite its success, the authors note that the LLM decode stage requires different architectural approaches, such as "phase splitting" between prefill and decode hardware. Additionally, the search time for the 6D optimization space remains high, requiring significant computational resources for large-scale network suites. Future work will also involve extending the evaluator to support more diverse architectures, such as asynchronous dataflow engines.