FuseMax: Leveraging Extended Einsums to Optimize Attention Accelerator Design

EDGE separates computations into three “actions”: map ($\bigwedge$), reduce ($\bigvee$), and populate ($=$) [38, Section 5]. Map specifies the pair-wise computation between the shared ranks of two tensors, reduce specifies the computation for the reduction step of an Einsum, and default populate ($=$) places a computed value from the right-hand side (RHS) of the Einsum to its location on the left-hand side (LHS).

Each map and reduce action contains two operations: merge and compute. Compute defines the operation to apply between two data values, and can be any user-defined function. Merge specifies which regions of the iteration space to touch; execution will not need to access the data space corresponding to culled points. Together, merge and compute precisely define the computations in an Einsum. Common merge operations include intersection ($\cap$), which touches points with non-zero values in both operands; and union ($\cup$), which touches points where at least one of the operands is non-zero.

The full EDGE specification for GEMM is then:

$\displaystyle Z_{m,n}=A_{k,m}\cdot B_{k,n}::\bigwedge_{k}\times(\cap)\bigvee_{k}+(\cup),$

(2)

where $\bigwedge_{k}$ specifies a map action between $A$ and $B$ on the $k$ rank and the intersection merge operator ($\cap$) culls $k$ points where at least one operand is zero. The compute operator ($\times$) multiplies the data values of coordinates surviving intersection. The reduce action ($\bigvee_{k}$) on the $k$ rank gathers all non-empty points in the $k$ rank and reduces them using addition ($+$).

In this work, we use three user-defined computations:

1. 1.

   Maximum ($\max(\cup)$) takes the maximum of two values. Suppose we have the following expression: $Z_{m}=A_{m}\cdot B_{m}::\bigwedge_{m}\max(\cup)$. The union merge operator ($\cup$) filters out any $m$ coordinates where both operands contain $0$ (and places 0 in the output). The $\max$ compute operator then returns the maximum of the two operands.

2. 2.

   Divide ($\div(\leftarrow)$) divides two data values. Given the following expression, $Z_{m}=A_{m}\cdot B_{m}::\bigwedge_{m}\div(\leftarrow)$, the merge operator ($\leftarrow$) only touches $m$ points where there is a non-zero value in the $B$ operand (see [38, Appendix]), and the compute operator divides the data value in $A$ with the data value in $B$.

3. 3.

   Exponentiation: we follow the example in EDGE [38, Section 7.4]. The expression $Z_{m}=e^{A_{m}}$, where $e$ is Euler’s number, applies the exponential function to every element in $A$. The exponent can also be an Einsum expression: $Z_{m}=e^{A_{m}\cdot B_{m}}::\bigwedge_{m}\times(\cap)$.

In addition to map and reduce, EDGE enables the expression of user-defined unary operations on tensors. For example, we can express the application of the non-linear, sigmoid function ($\sigma$) on each element of a tensor $A$ as $Z_{m}=\sigma(A_{m})$.

Throughout this paper, we take advantage of EDGE’s shorthand notation [38, Section 6] in the following ways:

• •

  We drop all reduce actions that consist of add and union in the compute and merge operator, respectively ($\bigvee+(\cup)$). Thus, $Z_{m}=A_{k,m}::\bigvee_{k}+(\cup)$ becomes $Z_{m}=A_{k,m}$.

• •

  We express all map actions using infix notation; that is, $A_{k,m}\cdot B_{k,n}::\bigwedge_{k}\times(\cap)$ becomes $A_{k,m}\times B_{k,n}$.

• •

  When $\max$ is part of a map action ($A_{m}\cdot B_{m}::\bigwedge_{m}\max(\cup)$), we replace it with the following shorthand: $\max(A_{m},B_{m})$

• •

  When $\div$ is part of a map action ($A_{m}\cdot B_{m}::\bigwedge_{m}\div(\leftarrow)$), we replace it with the following: $A_{m}/B_{m}$

EDGE also enables expressing Einsums that touch only a subset of the data space of their constituent tensors. For example, we may express prefix-sum of a tensor $A_{k}$ with the following Einsum:

$\displaystyle S_{i+1}=A_{k:k\leq i}$

For each coordinate $i$, $S_{i+1}$ is built by reducing together the subset of $A$ whose coordinates are $\leq i$. Note that this definition of prefix-sum computes the entire sum for a given $i$ without iteratively reusing the previous sum.

EDGE expresses recursion and iteration through generative/iterative ranks. We use the term standard ranks to differentiate non-iterative ranks from iterative ranks. We can express the iterative prefix-sum as follows:

	$\displaystyle S_{i+1}$	$\displaystyle=S_{i}+A_{i}$		(3)
		$\displaystyle\diamond:i\equiv K$		(4)

Here, $S$ is the iterative tensor that changes on each iteration, with the iterative rank, $i$, ranging from $0$ to $K$. Equation 4 indicates the stopping condition for the iterative expression (when $i$ is equal to $K$).

TeAAL [35] introduces the concept of cascades of Einsums, which expresses directed acyclic graphs (DAG) of Einsum expressions as a sequence of sub-Einsums. One can view the unrolled iterative expression in Equation 3 as a cascade:

	$\displaystyle S_{1}$	$\displaystyle=S_{0}+A_{0}$
	$\displaystyle S_{2}$	$\displaystyle=S_{1}+A_{1}$
	$\displaystyle...$
	$\displaystyle S_{K}$	$\displaystyle=S_{K-1}+A_{K}$

Finally, we use the EDGE Initialization label to specify computations that initialize tensors, which occur once. We use the EDGE Extended Einsum(s) label to specify the computation that occurs on each iteration of a cascade of Einsums [38] (see Einsum Cascade 5).

First, we recognize that, by the distributive property, Equation 6 can be factored to perform the reduction of $A_{k}$ first, before multiplying the result by $Y$. Doing so, we get the following cascade:

Now, because there is no read-after-write dependency between Equations 7 and 8, both Einsums can be included in the same pass. In fact, because Equation 8 reduces away the $K$ rank, Cascade 2 is a 1-pass cascade with respect to this rank. This reassociation actually provides a second benefit over Cascade 1: Equation 9 now only requires one multiplication (as opposed to $K$ multiplications in Equation 6).

Alternatively, we can iteratively construct $Y$ and $Z$ as we perform the pass through $A_{k}$. To do so, we will take a similar approach to the prefix-sum (see Sections II-C3-II-C4) and build intermediate $Y$s and $Zs$.

	$\displaystyle RY_{i+1}$	$\displaystyle=A_{k:k\leq i}\times B_{k:k\leq i}$		(16)
	$\displaystyle RZ_{i+1}$	$\displaystyle=RY_{i+1}\times A_{k:k\leq i}$		(17)

Just like with the prefix-sum, this version requires a lot of extra compute, but, because $Y=RY_{K}$ and therefore $Z=RZ_{K}$, the final result is the same.

We remove this extra work by making the $I$ ranks of $RY_{i+1}$ and $RZ_{i+1}$ iterative. This is shown in Cascade 3. Iterative $RY_{i+1}$ (Equation 12) looks very similar to the iterative prefix-sum. However, computing $RZ_{i+1}$ is a little more complicated. We start by introducing one more intermediate $S_{i}$, which is the prefix-sum for $A_{k}$:

$\displaystyle S_{i}=A_{k:k\leq i-1}$

(18)

Now, we can combine Equations 17 and 18 to write $RZ_{i}$ in terms of this prefix-sum:

$\displaystyle RZ_{i}=RY_{i}\times S_{i}$

(19)

Dividing both sides by $RY_{i}$, we derive an alternate definition for $S_{i}$:

$\displaystyle S_{i}=\frac{RZ_{i}}{RY_{i}}$

$S_{i+1}$ can also be written using this alternative definition:

$\displaystyle S_{i+1}=\frac{RZ_{i}}{RY_{i}}+A_{i}$

(20)

We can combine Equations 19 and 20 to compute $RZ_{i+1}$ in terms of $RZ_{i}$ (i.e., iteratively):

$\displaystyle RZ_{i+1}=RY_{i+1}\times\left(\frac{RZ_{i}}{RY_{i}}+A_{i}\right)$

Distributing $RY_{i+1}$ and performing some reassociation, we get Equation 13.

Cascade 3 is also a 1-pass cascade, performing one pass of the $K$ rank of $A_{k}$ (indexed with the variable $i$) and iteratively building $RY_{i+1}$ and $RZ_{i+1}$. Unfortunately, unlike Cascade 2, Cascade 3 does require extra compute over the original Cascade 1. However, memory bandwidth-limited workloads can afford to trade off extra compute for reduced memory traffic, and Cascade 3 may still provide benefit.

Because $e^{QK_{m,p}}$ can easily become extremely large, the above formulation suffers from overflow. Therefore, practical implementations [3, 41] often prefer the numerically stable variant that replaces Equation 26 with:

	$\displaystyle GM_{p}$	$\displaystyle=QK_{m,p}::\bigvee_{m}\text{max}(\cup)$		(29)
	$\displaystyle SN_{m,p}$	$\displaystyle=e^{QK_{m,p}-GM_{p}}$		(30)

and drop the $\frac{1}{\sqrt{E}}$ term when computing $QK_{m,p}$²²2The $\frac{1}{\sqrt{E}}$ term was introduced to bound the magnitude of $SN_{m,p}$ [48]. Because the numerically stable softmax variant already accomplishes this, the scaling is often omitted [16, 15, 13].. To compute the global maximum³³3“Global” here refers to over the $M$ fiber. $GM_{p}$, we reduce $QK_{m,p}$ with the operator max (instead of $+$). Notice that subtracting $GM_{p}$ from $QK_{m,p}$ in the exponent is equivalent to dividing by $e^{GM_{p}}$, and because the $\frac{1}{e^{GM_{p}}}$ term appears in both the numerator ($SN_{m,p}$ via Equation 30) and denominator ($SD_{p}$ via Equation 27), the result ($A_{m,p}$) stays the same. This construction improves numerical stability by bounding the values of the softmax numerator $SN_{m,p}$ to the range $(0,1]$.

The 3-pass cascade is the straightforward, numerically stable cascade that we already discussed in Section IV-C1, namely Equations 29-30 followed by Equations 27-28, reproduced in Cascade 4 for clarity.

In Pass 1, we compute $GM_{p}$; in Pass 2, we compute $SN_{m,p}$ and $SD_{p}$; and in Pass 3, we compute $A_{m,p}$. Notice that we must finish an entire $M$ fiber of Equation 33 (reading an entire $M$ fiber of $QK_{m,p}$) before $GM_{p}$ is ready to start Equation 34 (where we must read the same $M$ fiber of $QK_{m,p}$ again). Similarly, we must finish an entire $M$ fiber of Equation 35 (reading an entire $M$ fiber of $SN_{m,p}$) before $SD_{p}$ is ready to start Equation 36 (where we must read the same $M$ fiber of $SN_{m,p}$ again). Regardless of the mapping (including fusion), this cascade must perform three passes, since they are a consequence of the dependencies between Einsums.

We now briefly summarize the 2-pass cascade, deferring details due to space. Rather than computing the global max and then starting the softmax (as in the 3-pass cascade), the 2-pass cascade first partitions the input, computes a per-partition local max and applies it to form a variant of $SN_{m,p}$ whose elements are adjusted by the local max and likewise partitioned. Analogously, each partition gets a local denominator (also adjusted by the same local max). While this is occurring, it builds the global max from the local max values. Next, in a second pass, it uses the global max to correct the per-partition numerators and denominators and compute the softmax output.

While prior work proposes multiple different 1-pass cascades [16, 15] that take advantage of the reassociations presented in Section III-C. However, the main ideas are the same. First, modify the cascade to multiply the softmax numerator-times-$V$ and then compute the division (as described in Section IV-D). This reassociation combines the second and third passes of Cascade 4 (see Section III-C1). To ensure numerical stability, we cannot use this strategy to combine the first and second passes, so we instead use the iterative approach (see Section III-C2). Rather than using the per-partition local max to compute the local numerator and denominator, instead keep a running max that represents the max value seen so far. Each time a new running max is computed, adjust previous results (e.g., numerator-times-$V$, denominator, etc.) with this max.

Next we describe FlashAttention-2’s 1-pass cascade (Cascade 5) because we use it to build FuseMax. Note, despite the evidently increased compute relative to the 3-pass cascade, we will carefully design a mapping in Section V to hide these overheads on a spatial architecture.