FuseMax: Leveraging Extended Einsums to Optimize Attention Accelerator Design

Authors: nandeeka  nayak
xinrui  wu
toluwanimi o odemuyiwa
michael  pellauer
joel s emer
christopher w fletcher
Venue: iscaconf_org_ISCA_2024_Workshop_MLArchSys
Published: 2024-05-04 03:48:39

Abstract: Attention for transformers is a critical workload that has recently received significant "attention" as a target for custom acceleration.
Yet, while prior work succeeds in reducing attention's memory-bandwidth requirements, it creates load imbalance between attention operators (resulting in severe compute under-utilization) and requires on-chip memory that scales with sequence length (which is expected to grow over time).

This paper ameliorates these issues, enabling attention with nearly 100\% compute utilization, no off-chip memory traffic bottlenecks, and on-chip buffer size requirements that are independent of sequence length.
The main conceptual contribution is to use a recently proposed abstraction---the cascade of Einsums---to describe, formalize and taxonomize the space of attention algorithms that appear in the literature.
In particular, we show how Einsum cascades can be used to infer non-trivial lower bounds on the number of passes a kernel must take through its input data, which has implications for either required on-chip buffer capacity or memory traffic.
We show how this notion can be used to meaningfully divide the space of attention algorithms into several categories and use these categories to inform our design process.

Based on the above characterization, we propose FuseMax---a novel mapping of attention onto a spatial array-style architecture.
On attention, in an iso-area comparison, FuseMax achieves an average $6.7\times$ speedup over the prior state-of-the-art FLAT, while using $80$\% of the energy.
Similarly, on the full end-to-end transformer inference, FuseMax achieves an average $5.3\times$ speedup over FLAT using $85$\% of the energy.

Main paper: 
 i. introduction: 
 Over the past few years, transformers [31] have emerged as the model architecture of choice for a wide range of machine learning applications, from natural language processing [11] , [18] , [29] , [30] to computer vision [12] , [20] to speech recognition [3] , [14] . This rise has been accompanied by a corresponding wave of proposals for accelerating transformers in both software [7] , [9] , [10] and hardware [16] , [34] .Fortunately, many of the layers (projections, fully connected layers, etc.) used by transformers look very similar to prior generations of machine learning models. Its resource-intensive tensor products can be described and evaluated with existing tensor algebra accelerator modeling tools [17] , [22] , [26] , and many of the other layers (e.g., layer normalization) have negligible impact on performance and can be safely ignored.However, attention [31] -usually described as a matrix multiplication, a softmax, and then another matrix multiplication-does not fit into either of these boxes. For example, the softmax is both memory intensive (featuring low algorithmic reuse) and compute intensive (featuring exponentiation and division). Furthermore, attention's characteristics preclude many "free lunches" often used to improve efficiency for other DNN models. For example, because all tensors are a function of the model inputs, there is no opportunity to amortize memory access costs with an increased batch size. Additionally, since none of the operands can be computed before the inputs are given, compression/strength reduction techniques (e.g., quantization [13] , [33] , sparsity [21] , [28] , [32] , etc.) must be applied dynamically, leading to more complicated algorithms and hardware designs.To illustrate the difficulty in accelerating attention, consider the state-of-the-art accelerator for attention: FLAT [16] . FLAT uses fusion to reduce attention memory bandwidth bottlenecks on a spatial architecture (e.g., a TPU [15] ). Specifically, FLAT maps attention's matrix multiplications to the 2D spatial array and softmax operations to a separate 1D array. While FLAT's design does make attention compute bound, it becomes compute bottlenecked in the 1D array (the softmax), causing severe under utilization of the 2D array. While one could add additional PEs to the 1D array, this results in commensurate area costs.Making matters worse, FLAT requires that the entire vector over which the softmax is performed be buffered on chip. This vector is proportional to the sequence length, which is growing rapidly with time (e.g., Google reports 10 million length sequences in research, which would require 100s of MegaBytes to buffer [1] ). When the vector/sequence length grows beyond allowable buffer capacity, FLAT is forced to spill, which contributes significantly to attention energy consumption and can even make attention memory-bandwidth bound.This paper. We address the above challenges by proposing a novel spatial architecture -FuseMax -to accelerate attention, with particular emphasis on removing bottlenecks imposed by the softmax. Our architecture addresses all of the aforementioned issues associated with FLAT. FuseMax is compute bound, but provides almost 100% utilization of both the 2D and 1D arrays throughout the attention operation, without adding additional PEs to the 1D array. Additionally, FuseMax's on-chip memory requirements are invariant to sequence length and require no extra spills to memory regardless of sequence length.The technical core of the paper is two parts. First, Section III uses the recently proposed cascade of Einsums abstraction [22] to describe, formalize and taxonomize the space of numerically stable attention proposals that appear in the literature. In a nutshell, an Einsum defines an iteration space over tensors and what computation is done on and between tensors at each point in the iteration space. A cascade of Einsums is a sequence of dependent Einsums that can be used to describe and specify a larger kernel.While prior work [22] , [25] provides a precise definition for Einsums, a major contribution in our work is to show how this definition can be leveraged to inform accelerator design. Specifically, we recognize that Einsums make explicit precisely what compute the cascade performs. We show how this can be used to inform trade-offs in designing the complicated, compute-intensive softmax kernel. Additionally, we recognize that the Einsum cascade makes explicit precisely what dependencies there are between Einsums. We show how this can be used to make non-trivial deductions about a kernel's allowed fusion granularity and algorithmic minimum per-tensor live footprint. The relationship between the live footprint and the buffer capacity, in turn, has implications for the required data movement. Given that an Einsum cascade is mapping/scheduling agnostic, both of the above observations provide insight given any possible scheduling of the cascade onto hardware.In more detail, the latter provides insight into the number of times a given element of the an input must be revisited after visiting every other element of the input. Normally, one strives to choose a dataflow that exploits maximal reuse in a given element (or tile of elements) to avoid having to come back to it later. However, some algorithms preclude this strategy. For example, in a naïve implementation of attention, one must traverse the entire softmax input to build the softmax denominator and only after that can one revisit and scale each input (softmax numerator) by the denominator. Section III-C formalizes and generalizes this notion and shows that one can meaningfully divide the space of attention approaches based on the number of passes each cascade must perform, i.e., the number of times each element is (re)visited.We note, our lower bounds on passes hold for all mapping choices, including application of fusion. For example, despite using fusion, FLAT employs a 3-pass cascade and its reliance on large on-chip buffering is a symptom of trying to avoid 3 passes-worth of DRAM traffic. Obviously, fewer passes is preferable; although, interestingly, we find that cascades with fewer passes can increase the required compute.In the second part of the techical core (Section IV), we use the insights from Section III as a starting point to develop a novel mapping for attention that can be lowered to a spatial architecture. We call our architecture FuseMax. FuseMax adopts the attention cascade used in FlashAttention-2 [9] . However, despite using the cascade from FlashAttention-2, mapping this cascade to a spatial architecture instead of a GPU is non-trivial. We overcome the differences between the architectures and demonstrate a novel mapping for the cascade that achieves high utilization for entire transformer layers. Our architecture requires only minimal changes to a standard spatial architecture and is performance/energy robust to long sequence lengths (e.g., 1M tokens and beyond).Finally, we evaluate FuseMax on BERT [11] , TrXL [8] , T5 [30] , and XLM [18] and demonstrate a 6.7× speedup on attention with 80% of the energy and a 5.3× speedup on the full end-to-end inference with 85% of the energy relative to FLAT.In this section, we describe the concepts and terminology used in the remainder of the paper. This paper focuses on algebraic computations on tensors, where a tensor is a multidimensional array. A tensor's rank refers to a specific dimension of the tensor, while the tensor's shape is the set of valid coordinates for each of the tensor's ranks. We use the same symbol for the name and shape of the rank, i.e., the rank M has shape M . A fiber is a set of points in a tensor with the same coordinate in all ranks except one.An Einsum defines a computation on a set of tensor operands using an iteration space that specifies the set of points where the computations are performed [22] , [25] . For example, we describe matrix-matrix multiplication (GEMM) computation with the following Einsum:EQUATIONWe follow the operational definition for Einsums given in [22] . In this work, we leverage the Extended General Einsums notation (EDGE) [25] (first developed for graph algorithms) to express the complex non-linear operations required for transformers.TeAAL [22] introduces the concept of cascades of Einsums, which expresses directed acyclic graphs (DAG) of Einsum expressions as a sequence of sub-Einsums. For example, a matrix multiplication can be split into two Einsums, as follows:T k,m,n = A k,m × B k,n Z m,n = T k,m,nAn Einsum specifies the computation, while a mapping indicates how computation occurs in space and time on an accelerator [6] , [26] . Mapping specifications include aspects such as loop order, partitioning, and work scheduling (sequential vs. parallel operations) [22] . Einsums in a cascade may be mapped independently or fused together to improve performance.Transformer models generally follow the architecture defined in [31] . In this work, which addresses the impact of long sequence lengths during self-attention, we focus on the encoder architecture. Figure 1a gives an overview. As the sequence length grows, the relative importance of the different operations changes. Figure 1b shows that at shorter sequence lengths, the weight-times-activation "linear" layers are a larger fraction of the total required compute, while at long sequence lengths, the attention dominates. All additional required computation has negligible impact. In the next section, we focus on describing attention more precisely, and use our analysis to understand prior work on efficient implementations.Our first contribution is to show how cascades of Einsums can be used to describe, taxonomize, and highlight trade-offs in the space of attention algorithms. The key insight is that cascades of Einsums provide a precise description of a kernel's compute requirements, the algorithmic minimum live footprint for each tensor, and the number of passes the algorithm must make through this live footprint.In the original transformer paper [31] , the kernel was described with the following equation:EQUATIONHowever, this equation says almost nothing about what the inputs Q, K, and V look like or what iteration space needs to be traversed. We clarify these points by rewriting Equation 2as a cascade of Einsums, with the exception of the softmax, whose cascade we will explore in Section III-B:EQUATIONEQUATIONHere, Equations 31 and 5 look like matrix multiplications. Taking Equation 5as an example, for each point in the iteration space F ×M ×P , we perform a multiplication using elements from two 2-tensors (A m,p and V f,m ) to produce a 2-tensor output (AV f,p ), which requires reducing across the inputs' shared rank M .Equations 3-5 can be modified to refer to the full batched, multi-head self attention [31] by adding B and H ranks to all tensors. This changes the characteristics of the kernel. Adding the B and H ranks means that Equations 3 and 5 behave like many independent matrix multiplications instead of one monolithic matrix multiplication. The challenges with attention, described in Section I, follow clearly from this modification. Because all tensors contain a B rank, the matrix multiplications are all unique to the specific batch's inputs. Therefore, none of these tensors can be computed before the inputs are given, and there is no data sharing between the different elements in the batch. To simplify notation, we assume the presence of the B and H ranks but omit writing them throughout the rest of paper.We now apply the same precise notation to the softmax. A softmax [4] over a 1-tensor is traditionally expressed with the following equation:EQUATIONIn the context of attention, this operation becomes two dimensional and can be expressed using the following cascade with input QK m,p :EQUATIONEQUATIONFor each point in the iteration space (m, p), we exponentiate QK m,p to generate the softmax numerator (SN m,p in Equation 7), reduce SN m,p with addition to produce the softmax denominator (SD p in Equation 8), and finally, divide the numerator and denominator to produce the final result (A m,p in Equation 9). 1) Improving Numerical Stability: Because e QKm,p can easily become extremely large, the above formulation suffers from overflow. Therefore, practical implementations [2] , [27] often prefer the numerically stable variant that replaces Equation 7 with:EQUATIONSN m,p = e QKm,p-GMp (11) and drop the 1 √ E term when computing QK m,p 2 . To compute the global maximum 3 GM p , we reduce QK m,p with the operator max (instead of +) 4 . Notice that subtracting GM p from QK m,p in the exponent is equivalent to dividing by e GMp , and because the 1 e GMp term appears in both the numerator (SN m,p via Equation 11) and denominator (SD p via Equation 8), the result (A m,p ) stays the same. This construction improves numerical stability by bounding the values of the softmax numerator SN m,p (i.e., before the division) to the range (0, 1]. 2 The 1 √ E term was introduced to bound the magnitude of SNm,p [31] . Because the numerically stable softmax variant already accomplishes this, the scaling is often omitted [7] , [9] , [10] . 3 "Global" here refers to over the entire M fiber. 4 We use the reduction notation defined in [25] C. Optimizing Softmax Live Footprint and Memory TrafficEinsums make explicit the dependencies between values produced and consumed across tensors in a cascade. We show how these dependencies can be used to make nontrivial deductions about a kernel's allowed fusion granularity and algorithmic minimum per-tensor live footprint. Because these deductions are made using only the cascade, they hold regardless of the choice of mapping or buffer capacity.As a simple example, let us look at the softmax cascade Equations 7-9. If we want to minimize data traffic of SN m,p , we need to choose a dataflow for each Einsum that keeps SN m,p stationary and fuses the three Einsums together. In other words, the dataflow must finish using the first element of SN m,p before moving onto the next. However, such a dataflow does not exist for this cascade. Any implementation must visit every element of a given M fiber of SN m,p to compute SD p before it can revisit any element of that fiber to compute A m,p .We define a pass a cascade performs over a particular fiber of a particular rank and tensor to be a traversal of every element of that fiber. Each time an element must be revisited after visiting every other element of that fiber, there is an additional pass. For example, Equations 7-9 perform two passes over the M rank of SN m,p .The number of passes a cascade performs is relevant because it restricts fusion schedules possible. While Einsums within a pass can be fused at will, producing and consuming a tile of the intermediate at a time, the two passes cannot be fused on the M rank. Any implementation must visit all elements of an M fiber of SN to produce SD before it can visit any of the elements of that fiber to produce A.This analysis also provides a non-trivial lower bound on the tensors' live footprints. For example, the algorithmic minimum live footprint for tensor SN is M . In other words, an architecture must either have enough buffer space to hold an entire M fiber of SN or spill and reload that fiber, incurring memory traffic proportional to the shape of M . We note that this analysis is mapping independent. There is no dataflow for this cascade that enables a smaller live footprint. We find that existing approaches to numerically stable attention can be classified as either 3-pass, 2-pass, or 1-pass cascades, where an N -pass cascade performs N passes of a given M fiber. See Table I .Next, we describe FuseMax's cascade (Cascade 1)-first proposed for FlashAttention-2. Note, despite the evidently increased compute relative to the 3-pass cascade, we will carefully design a mapping in Section IV to hide these overheads on a spatial architecture.BQK m1,m0,p = QK m1×M 0+m0,p (12) BV f,m1,m0 = V f,m1×M 0+m0 (13) RM m1:m1=0,p = -∞ (14) RD m1:m1=0,p = 0 (15) RN V m1:m1=0,p = 0 (16) Extended Einsums:LM m1,p = BQK m1,m0,p ::EQUATIONRM m1+1,p = max(RM m1,p , LM m1,p )SLN m1,m0,p = e BQKm1,m0,p-RMm1+1,pEQUATIONP RM m1,p = e RMm1,p-RMm1+1,p (22) SP D m1,p = RD m1,p × P RM m1,pEQUATIONEQUATIONEinsum Cascade 1: fooWe will start by expressing the partitioning of both of the inputs QK m,p and V f,m into M1 chunks of M0 elements each (Equations 12-13). This allows us to perform operations like maximum on individual M 0 fibers, rather than on the whole tensor (Equation 17). The problem is, of course, that the local maximum is not necessarily the same for all M 0 fibers and so will not just cancel nicely like the global maximum.We resolve this by instead using the running maximum (RM m1,p )-the global maximum of all inputs seen so farinstead of the local maximum. We recognize that M 1 can also serve as an iterative rank, and iteratively build up RM m1,p . After initializing RM 0,p to -∞ (Equation 14), we compute a new running maximum RM m1+1,p using the running maximum computed in the previous iteration RM m1,p and the new local maximum LM m1,p (Equation 18).We can now use the running maximum to compute a local numerator SLN m1,m0,p (Equation 19), a local denominator SLD m1,p (Equation 20), and even the dot product result SLN V f,m1,p (Equation 21) using the partitioned BV f,m1,m0 (Equation 13). Now consider the softmax denominator. Eventually, we would like to reduce SLD m1,p into a 0-tensor, but because its values may have been computed with different maximums, we cannot simply use addition. Instead, by introducing a new running denominator RD m1,p with iterative rank M 1, we can correct the old denominator RD m1,p to the new running maximum RM m1+1,p and then perform the addition. Again, because M 1 acts as an iterative rank for RD m1,p , we must initialize the running denominator at the start of the computation to 0 (Equation 15). Then, at each point m1, the correction factor P RM m1,p allows us to correct the previous running denominator RD m1,p with the new maximum (Equation 23). In other words, RD m1,p is downscaled by e RMm1,p . SP D m1,p "switches" the downscaling factor on RD m1,p to e RMm1+1,p by multiplying RD m1,p by e RM m1,p e RM m1+1,p (P RM m1,p ). Once SLD m1,p and SP D m1,p have the same maximum, they can be combined to produce the new running denominator RD m1+1,p (Equation 24). We can do the same to compute the running numerator-times-V ( Equations 16, [25] [26] .Finally, AV f,p can be computed by dividing the final numerator-times-V by the final denominator. By construction, at this point, RN V f,M 1,p and RD M 1,p are both downscaled by the same maximum RM M 1,p (conveniently, also the global maximum) and can be correctly combined.Based on the framework from Section III, we now describe FuseMax, an efficient mapping of a 1-pass attention cascade to a spatial array-style architecture. In this work, we use the same cascade as FlashAttention-2 [9] .The goal when mapping a cascade onto hardware is to fully utilize all available compute units. In our evaluation of prior work (Figure 4 and Section V), we observe that at short sequence lengths, the 2D PE array is under-utilized because it must wait for the 1D PE array to compute the softmax. At longer sequence lengths, both arrays are under-utilized since the workload becomes memory-bandwidth limited.FuseMax's mapping addresses these issues to achieve full utilization on both the 1D and 2D PE arrays. We do so by (1) sharing the operations beyond multiply/accumulate (max/exp) between the 1D and 2D arrays and (2) ensuring that the workload is never memory-bandwidth limited by deeply fusing all Einsums in the cascade to restrict the live footprint to only what can be buffered on-chip. No matter the sequence length, our dataflow is never forced to spill any of its intermediates off-chip.Architecture. We assume a standard spatial array-style architecture for our mapping. See Figure 2 . We set parameters to match the cloud configuration in prior work [16] . Note, although both the 1D and 2D PE arrays perform exponentiation, we implement exponentiation with 6 sequential multiplyaccumulate operations [24] , [32] and therefore do not require a dedicated exponentiation unit.Fusion and Partitioning. Prior attention accelerators [16] , [34] explore fusing many of attention's loop nests together. However, because these accelerators all use multi-pass cascades, the algorithmic minimum live footprint of some tensors (e.g., QK m,p ) is O(M ), meaning that for long sequence lengths, intermediates cannot be buffered on chip.FuseMax leverages fusion in conjunction with the 1-pass cascade to eliminate the memory traffic of these tensors, regardless of the sequence length. Specifically, we partition on both M and P (forming M 1, M 0 and P 2, P 1, P 0), and maximally fuse all levels in the attention loopnest.Parallelization and Spatial Reduction. While prior work implementing attention in hardware [16] , [34] does utilize the 2D spatial array for the tensor products, it fails to do so for the softmax, choosing instead to use the 1D array. However, because there are far fewer total PEs in the 1D array than the 2D array, the softmax becomes a bottleneck. FuseMax improves utilization of the 2D spatial array by using it for both the tensor products and the exponentiation operator in the softmax. FuseMax parallelizes across the M 0 and P 0 ranks throughout the attention kernel. We set M 0 × P 0 = # 2D Array PEs. The large spatial reductions required when parallelizing across the M 0 rank are easily handled by the low-cost inter-PE communication network.Pipelining. The dependencies between different Einsums in our cascade necessitate fine-grain pipeline parallelism to achieve high utilization of both the 1D and 2D spatial arrays. Figure 3 shows the waterfall diagram for FuseMax in the steady state. Time is broken into epochs. Each epoch performs the same set of tile-granular operations at specific tile-relative coordinates (given by a, b, c, d in the figure). Across all epochs, the kernel evaluates all tiles and each Einsum is mapped to either the 2D or 1D array for all epochs (as shown in the figure ).A major design consideration when pipelining the mapping is how to overcome the latency of fills and drains to/from the spatial array. Without careful interleaving, the kernels suffers from low utilization due to long latency spatial reductions, whose results are required for future Einsums.We address the above issues with two levels of interleaving. First, we interleave the construction of dependent tiles across epochs. This is reminiscent of software pipelining. For example, in Figure 3 Second, we interleave the construction of certain tiles within an epoch at a fine (e.g., cycle-by-cycle) granularity. See the notation 'A|B' in Figure 3 . This is to ensure high utilization of both the 2D and 1D PE arrays at all times. In each cycle, Putting everything together, as Section V will show, the above enables high utilization of all 2D and 1D array PEs.In this section, we demonstrate how the FuseMax dataflow achieves improvements in both performance and energy relative to the state of the art, for both attention and the end-to-end transformer inference.First, we present the experimental set-up details common to all following subsections.Workloads. We evaluate all accelerators and configurations using the same transformer models used by FLAT [16] : BERT-Base [11] (BERT), TrXL-wt103 [8] (TrXL), T5-small [30] (T5), and XLM [18] . We omit FlauBERT [19] because it uses the same hyperparameters as TrXL. We also note that though T5 is an encoder-decoder model, we only evaluate the encoder in this work. Following FLAT, we use a batch size B = 64 for all evaluations.Modeling with Timeloop. We perform our evaluation using the tensor algebra accelerator modeling and design space exploration too Timeloop [26] . We use these tools to build models of the accelerator architectures and evaluate each Einsum individually. Results from individual Einsums are combined using heuristics presented in prior work for evaluating full cascades [22] .Unfused Baseline. We build the unfused baseline by combining the costs of three phases: QK (Equation 3), the 3-pass softmax (Equations 10-11 and 8-9), and AV (Equation 5). Because this baseline is unfused, each phase can be scheduled independently, but proceed sequentially and require outputs to be written to memory between phases. We use Timeloop to search for efficient mappings to perform QK and AV . Additionally, we model the softmax for the unfused baseline by allowing the accelerator to load the M fibers of the input on-chip one-by-one (spilling if there is not enough space) before performing the compute. We model the memory traffic, compute, and energy required to perform all Einsums required for attention.FLAT Baseline. Our main baseline is the state-of-the-art attention accelerator FLAT [16] (cloud architecture). Though we started with the FLAT authors' original code, we found and corrected a number of bugs with confirmation from the original authors. Beyond correcting the FLAT codebase, we created and validated a Timeloop model that reproduces the FLAT authors' (corrected) code to within < 1% error. However, the FLAT codebase does not model the cost to perform the softmax. Specifically, their model ignores the cost of data transfers (between any levels of the memory hierarchy) and uses 2 30 1D PEs. When comparing FuseMax and FLAT in this work, we augment our Timeloop model to model softmax correctly per the 3-pass cascade implicitly assumed by FLAT.Hardware parameters. Figure 2 shows the selected hardware parameters, which were chosen to match FLAT's cloud accelerator. Also following FLAT, we use a 940 MHz frequency. We use Accelergy to confirm that, despite the changes to the PEs' functional units, the area difference between FLAT and FuseMax is < 0.01%.We now evaluate FuseMax to demonstrate the benefits it provides on the attention kernel by comparing it to the two baselines.Utilization. Figure 4 shows the utilization of the 2D PE array. Because of the large amount of compute required for the softmax, both baselines achieve very poor utilization of this array. On the other hand, at long sequence lengths, FuseMax achieves almost 100% utilization. We observe that both baselines do achieve slightly higher utilization on XLM, which can be attributed to the higher intensity caused by a larger embedding dimension (E/F ).Speedup. Figure 5 shows that FuseMax achieves an average speedup of 10× over the unfused baseline and 6.7× over FLAT. We note FuseMax achieves lower speedup on XLM only because the baselines are able to achieve higher utilization of the 2D array on this transformer (Figure 4 ). An energy analysis can be found in the full paper [23] .To evaluate the benefits of FuseMax on end-to-end transformer inference, we include the other required linear layers (Section II). We use Timeloop to search for optimal mappings for these linear layers and use the same mappings for all three accelerator configurations. The attention modeling remains the same as Section V-B.Figure 6 shows the performance improvement achieved by FuseMax. Across the sequence lengths tested, FuseMax achieves an average speedup of 7.6× over the unfused baseline and 5.3× over FLAT. As discussed in Section II, as sequence length grows, attention becomes a larger fraction of the total required compute. Therefore, at 1M tokens, FuseMax achieves an average 10× speedup over the unfused baseline and 7.5× speedup over FLAT. An energy analysis can be found in the full paper [23] .In Equation3, we also substitute E for d k following the notation defined in Section II, where the shape of a rank is also its rank name. 
 
 Appendix:  


Decision: Congratulations! We are pleased to inform you that your paper has been accepted for presentation at MLArchSys 2024. We look forward to your participation at the workshop. Further details regarding the schedule and format will be provided soon. See you at the workshop!

Reviews: review: title: title: A systematic analysis of einsum to optimize attention acceleration
strength_weakness: strength_weakness: (1) The evaluation is limited to a few transformer models and a specific hardware configuration (I ack that limited number of pages was a factor here). Exploring a wider range of models and cofigurations improve the quality of the paper and strengthen the claims.
questions: questions: The paper proposes a systematic methodology to analyze and optimize operations in transformers by using Einsum cascades (an already proposed methodology). The use of Einsums to explicitly model data dependencies and derive minimum live footprints is a good contribution that can inform efficient hardware design and mapping. Based on the proposed methodology and analysis, the authors propose FuseMax architecture, showing improvements in performance and energy consumption over state-of-the-art methods, such as FLAT and FlashAttention. Here are few comments/questions:

(1) I understand that the authors had to condense the materials into four pages, which may have compromised clarity and detailed description of the method, but it would be great if the authors could add a more detailed explanation with illustrative examples, esp. for readers with less familiarity to these concepts (something that I hope the authors can fix during their potential presentation at the workshop).

(2) The models used for evaluation are limited and outdated. It would be interesting to see the effectiveness of FuseMax on a wide variety of models, including recent architectures.

(3) Same goes with the hardware evaluation, which was on a specific spatial array configuration. Performing sensitivity of FuseMax's performance to different hardware parameters would be valuable.

(4) Including more detailed breakdown of the performance gains, such as latency and energy consumption of individual components of the attention mechanism.

(5) Recent work shows sparsity/quantization can be productively used in attention mechanism. I am wondering if cascade of Einsum can be extended to cover such compression methods.
title: title: Interesting approach to the use of einsums, but the rest of the paper lacks novelty and requires more investigation
strength_weakness: strength_weakness: * The paper lacks novelty. Most of the work builds upon the FLAT architecture.
* The paper assumes one configuration for the evaluation and does not give enough details on the experimental setup.
* The writing needs to be improved.
questions: questions: **Summary:** FuseMax maps attention algorithms to a spatial array architecture by leveraging Einsum cascades to describe and optimize tensor operations. 

**Strengths:**
* The paper introduces an interesting approach to use einsum cascades in accelerator design. In particular, cascade of einsums is used to describe, formalize and taxonomize the attention algorithms in the literature.
* The experimental results show substantial improvements over the state-of-the-art in both speed and energy efficiency.

**Weaknesses:**
* The paper lacks novelty. Most of the work builds upon the FLAT architecture.
* The paper assumes one configuration for the evaluation and provides limited information on the experimental setup, including the specific hardware used, software stack, and any optimizations applied during the evaluation.
* The writing needs to be improved. Especially, after Section III, the sections give the impression that they are cut. For example, Section V.A starts by "First, we present the experimental set-up details common to all following subsections." However, there is no other subsection in the remaining part of the paper.

**Questions for the authors:**
* Could you clarify how the foundational concepts from FLAT have been extended or modified to develop FuseMax?
* Could you provide more details on your implementation of FuseMax and FLAT on Timeloop and Accelergy, including the assumptions made? 
* FLAT reports multiple configurations with various on-chip buffer sizes. Which one is used for the baseline evaluation? 
* How would the utilization and speedup comparison change for various on-chip buffer sizes?
title: title: Impressive results but lack of analysis
strength_weakness: strength_weakness: - Lack of important research and technical details
- Ambigious claims
questions: questions: The paper provides impressive improvements over prior state of the art TPU-like spatial architecture. The paper uses the same cascade as proposed in FlashAttention2, and maps it onto a spatial architecture to improve compute utilization and achieve speedup.

However, the main contribution of this paper is not clear. Does the improvement in utilization (and the corresponding speedup) come purely from the fact that it uses the same 1-pass cascade as the one proposed in FlashAttention2, compared to the previous SOTA paper using a different 3-pass cascade?

I would also like much more details in Section 4 to support the claims made in the introduction. As an example, the introduction makes the following claim, "We show how this can be used to make non-trivial deductions about a kernel’s allowed fusion granularity and algorithmic minimum per-tensor live footprint." I did not find any analysis or explanation on this. Section 4 merely states "The workload is never memory bandwidth limited by deeply fusing all Einsums in the cascade to restrict the live footprint to only what can be buffered onchip." without clearly explaining how it is done (Does it come from FlashAttention2 or some new analysis?).

Another important factor on which I would like more details on is the claim made in the introduction: "However, despite using the cascade from FlashAttention-2, mapping this cascade to a spatial architecture instead of a GPU is non-trivial." What are the main differences and how are they solved? Is it solved just by sharing the non-linear operations between the 1D and 2D arrays? Are there any scheduling challenges involved in co-ordinating between these two arrays? What are the technical challenges involved in implementing the FlashAttention2 cascade on the spatial accelerator?

Additionally, most modern LLMs are decoder only architectures, so it would be interesting to explore the effect such casacdes on decoder based models. Will there be any differences in results and/or additional challenges involved in mapping these cascades for decoder based models?

The paper would benefit from a much deeper analysis of its insights and solutions. The final results are very impressive, and useful especially in the modern era where context lengths are huge. I understand this is a workshop paper, so there is a limitation on the number of pages, but I would like to see more text on analysis supporting the claims made in the introduction rather than explaining the prior work on how tensor operations can be expressed as einsum cascades.
title: title: Paper about improving attention architecture in hardware
strength_weakness: strength_weakness: 1) Despite the fact that paper presents impressive speedup numbers for important problems like alleviating memory bottleneck and increasing efficiency of attention hardware, its actual implementation is explained only at high level and very shortly at section IV. Hence, it is not clear how those high level details are implemented in the hardware. 
2) Authors do not comment on potential drawbacks and limitations of the projects as well, which raises many questions on what are some pitfalls and tradeoffs.
questions: questions: This paper presents important architectural contributions for avoiding memory bottleneck and improving the efficiency of attention hardware. They enable these contributions via revisiting math behind the softmax operation predominantly and via reconsidering the fusion behind einsum cascades generally. As a result, they were able to generalize softmax logic and achieve a high utilization of both 2D and 1D arrays of the attention hardware architecture. 

Their primary point of comparison is FLAT paper from ASPLOS'23 which uses fusion of logit and attention operations to improve the computing efficiency. However, authors of the paper argue that FLAT's on-chip memory use grows unsustainably with tokens increasing which is prohibitive. Additionally, they also argue that with FLAT, 2D arrays remain still during softmax operations which, in turn, leads to a very low utilization rates. However, despite motivation and math explanations behind this FuseMax paper is clear, they spend very little time on actual hardware mapping and implementation of their solution, so it is very hard to evaluate the work which is the main reason for rejection for me. 

Suggestions: less time for math, and more content for actual architectural implementation, more evaluation numbers on different aspects of the solution, more comments on limitation of the solution for a fair analysis, also some details of the FLAT paper can be explained a bit more for a clear vision of the competing solutions (I had to look up for the original paper to understand the FuseMax paper). 

Nevertheless, thanks for your submission and important solutions in mind!