#  Bifurcated Attention for Single-Context Large-Batch Sampling

Ben Athiwaratkun  Sujan Kumar Gonugondla  Sanjay Krishna Gouda  Haifeng Qian
Hantian Ding  Qing Sun  Jun Wang  Jiacheng Guo  Liangfu Chen  Parminder Bhatia
Ramesh Nallapati  Sudipta Sengupta  Bing Xiang

###### Abstract

In our study, we present _bifurcated attention_ , a method developed for
language model inference in single-context batch sampling contexts. This
approach aims to reduce redundant memory IO costs, a significant factor in
latency for high batch sizes and long context lengths. Bifurcated attention
achieves this by dividing the attention mechanism during incremental decoding
into two distinct GEMM operations, focusing on the KV cache from prefill and
the decoding process. This method ensures precise computation and maintains
the usual computational load (FLOPs) of standard attention mechanisms, but
with reduced memory IO. Bifurcated attention is also compatible with multi-
query attention mechanism known for reduced memory IO for KV cache, further
enabling higher batch size and context length. The resulting efficiency leads
to lower latency, improving suitability for real-time applications, e.g.,
enabling massively-parallel answer generation without substantially increasing
latency, enhancing performance when integrated with post-processing techniques
such as reranking.

Machine Learning, ICML

##  1 Introduction

The advent of large language models (LLMs) has ushered in a new era of machine
learning, exhibiting remarkable performance on a wide array of tasks (Brown et
al., 2020; OpenAI, 2023; Chowdhery et al., 2022; Touvron et al., 2023; Chen et
al., 2021; Hoffmann et al., 2022; Li et al., 2022; Microsoft, ; Amazon, 2022;
Nijkamp et al., 2023). Despite their impressive capabilities, the deployment
of these large-scale models in practical applications poses significant
challenges, particularly in terms of inference latency and efficiency.
Enhancing these aspects is critical, as they directly influence the
computational resources required to generate predictions and enable the
practical implementation of these advanced models across various industries.

A particularly demanding inference scenario is single-context batch sampling,
where the goal is to generate multiple completions from a single context. This
task is commonly encountered in numerous applications such as code-editing IDE
tools that provide multiple recommendations, or in cases where ranking among
many generations is needed for optimal performance (via ranking metrics like
mean log probability, majority voting, etc). The incremental decoding of such
sampling scenario is memory IO intensive, which becomes a latency bottleneck
for high batches and context lengths.

In this study, we investigate two compatible strategies to address the memory
IO challenges in tranformers inference: (1) an investigation of multi-query
and its trade-offs, and (2) a novel technique called context-aware bifurcated
attention.

Our investigation begins with an analysis of the generalized multi-query
attention (Ainslie et al., 2023), which includes multi-query (Shazeer, 2019),
as well as the established multi-head attention mechanism (Vaswani et al.,
2017) for performance and latency trade-off. Our findings show smooth
performance scaling with increasing model size for a fixed value of the number
of groups $g$ for generalized multi-query111Lower values of attention groups
$g$ lead to higher compression of the key-value tensors, as in the multi-query
case where $g=1$, hence improving inference efficiency and latency due to
reduced KV cache compared to the multi-head case where $g=h$, the number of
query attention heads.. Lowering $g$ results in an upward shift of the
validation loss vs model size scaling curves. The consistent relationship
between the cache compression, model size and validation loss allows us to
trade-off inference efficiency with model size, i.e., enables us to select
higher compression for use cases requiring high efficiency, while still
matching the performance of multi-head attention by compensating with a larger
model size.

Secondly, we introduce context-aware bifurcated attention, a technique that
bifurcates any attention in the generalized multi-query family into context
and decoding components during incremental decoding. Such bifurcation involves
the same number of FLOPs and yields identical results compared to the original
attention, but can significantly reduces memory IO cost and thus latency in
high batch and context length scenarios. This approach allows the generation
of multiple real-time completions without incurring much additional latency
costs, or enables much higher batch sizes leading to improved ranking
performance. For instance, for CodeGen 16B multi-head model (Nijkamp et al.,
2022) with 2k context length, we are able to increase the batch size to $128$
with bifurcated attention, compared to batch size of only $5$ without,
resulting in the pass@k (Chen et al., 2021) increasing from $59.0\%$ to
$84.6\%$, or pass@top3 via mean log-p increasing from $55.2\%$ to $58.1\%$.

##  2 Related Work

In the literature, there are multiple avenues to improve inference latency
and/or latency. Quantization reduces memory usage by using low-bitwidth
representations such as int8, int4, and fp8 (Wei et al., 2023; Yao et al.,
2022; Dettmers et al., 2022; Frantar et al., 2022; Kuzmin et al., 2022; Xiao
et al., 2022). Quantization when applied only to model parameters offer
diminishing results as with longer sequence lengths and large batch sizes
where memory access and compute associated with dot-product attention
dominates the overall inference latency.

Sparse attention (Beltagy et al., 2020; Child et al., 2019; Zaheer et al.,
2020) has been extensively studied as a way to reduce the complexity of
attention for longer contexts and faster inference. Pope et al. (2022)
investigates generative inference efficiency of large language models by using
multi-dimensional partitioning techniques optimized for TPUs (collective
einsum) to achieve a Pareto frontier on latency and model FLOPs utilization.
The paper also shows that multi-query attention allows scaling up to 32x
larger context length with an emphasis on the efficiency under high batch
size. Paged attention (Kwon et al., 2023) enhances memory management of the KV
cache by dividing it into blocks and employing a block table for mapping
purposes. This approach effectively accommodates dynamic workload shifts and
reduces memory storage requirements through the sharing of the prompt’s KV
cache across multiple output sequences. However, this does not reduce the
memory reads of KV cache.

Speculative decoding, and its variants uses a smaller draft model to propose
multiple sequential tokens, which are processed in parallel by the main model
to accept or reject such tokens (Chen et al., 2023; Leviathan et al., 2022; Li
et al., 2024; Cai et al., 2024; Fu et al., 2023). The key idea is to enable
decoding of multiple tokens at every step, thereby amortizing the memory IO
usages of the main model. However, the latency of decoding will be still
dominated by KV cache I/O bandwidth at large context sizes, where bifurcated
attention can enhance the decoding speed further. In short, incremental
decoding focuses on lowering the amortized memory IO of model loading while
multi-query and bifurcated attention lowers the memory IO of KV cache.

##  3 Background

###  3.1 Notation

We use the following notation throughout the paper.

  * •

$K$: key tensor, $V$: value tensor, $q$: query tensor, $P_{x}$: projection
tensor associated with key, value or query tensor.

  * •

We denote $\langle A,B\rangle$ as a tensor operation between $A$ and $B$. The
actual operation can be specified in Einstein sum notation. We use $\oplus$ to
denote concatenation.

  * •

$N$ the number of model parameters, $d$: hidden dimension, $h$: number of
attention heads, $k$: $\frac{d}{h}$, or head dimension, $\ell$: number of
layers, $m$: context length (or key/value tensor length), $n$: query tensor
length where $n=m$ during context encoding and $n=1$ for incremental decoding,
$g$: number of attention groups (to be explained). We also use $v$ to
represent the head dimension for the value tensor where practically $k=v$.

###  3.2 Language Model Inference

There are many inference scenarios for language model, including batch
inference and single-context batch sampling (Figure 1). Batch inference refers
to the case where we process multiple inputs together in a batch, and generate
subsequent tokens for each batch index independently. In the case where the
batch size is 1, this reduces to the single-context inference. Another
scenario is the single-context batch sampling where we generates multiple
sequences based on a single context, where difference between the batch
inference case is that the prefill only needs to be done for a single context
to obtain the KV cache, then broadcasted to other batch indices.

Figure 1 also illustrates the two phases of language model inference: (a) the
context encoding or prefilling and (b) the incremental decoding. The context
encoding refers to a single forward pass that computes the key and value
tensors for all token positions in the context. Once the key and value tensors
are computed, we cache these key and value tensors to be used for the
attention mechanism during the incremental decoding phase, which sequentially
generates one token at a time222Or $k$ tokens at a time, in case of
speculative decoding (Chen et al., 2023; Leviathan et al., 2022).

Figure 1: Illustration of the two phases of language model inference: context
encoding and incremental decoding, as well as different inference scenarios.
In batch inference scenario, we process multiple inputs at once and perform
incremental decoding steps. In batch inference, we group multiple inputs in
batch to perform both context encoding and the subsequent incremental
decoding. In the single-context batch sampling scenario, we perform context
encoding on a single input to obtain the context KV cache, then perform
incremental decoding (with temperature sampling) to obtain potentially
different generations.

During the context encoding phase, the number of floating point operations
relative to the memory input/output (IO) operations is high, corresponding to
the compute-bound regime where the latency is influenced by the FLOPs.
However, during incremental decoding where we perform attention on a single
query token, this falls into a memory-bound regime where the number of
computation per memory access is roughly 1-to-1 (see Appendix D.1 for
details). The memory IO refers to the read and write operations from the high
bandwidth memory (HBM) (Jia et al., 2018) to the fast on-chip SRAM where the
actual computation happens. The memory IO of the incremental decoding itself
consists of two components: (1) the model parameter loading and (2) KV cache
loading. Component (1) is constant regardless of the context length $m$ or
batch size $b$ where component (2) depends on both $m$ and $b$ and dominate
the overall memory IO if $m$ or $b$ are high, which can become a significant
bottleneck for inference. Our work primarily focuses on reducing component
(2).

###  3.3 Multi-Query, Multi-Head and the Generalized Multi-Query Attention

Multi-query attention, proposed by Shazeer (2019), is an attention mechanism
for transformers models that uses a single head for the key and value tensors,
compared to $h$ heads in the traditional multi-head attention (Vaswani et al.,
2017). This technique effectively reduces the KV memory IO by $h$ times, which
leads to higher inference efficiency during incremental decoding. In effect,
the single-head key or value tensor is shared and used to attend to all the
multi-head query, hence the name multi-query. This corresponds to a
compression in representation power of the key and value tensor, which we will
see in the scaling laws study (Section 5.1) that it results in a reduced
expressiveness in terms of model parameter efficiency. Such reduced
expressiveness can be compensated by scaling the model bigger than the multi-
head counterpart to match the representation power.

We can also extrapolate these insights to a generalized multi-query attention
mechanism (Ainslie et al., 2023), which provides a framework to understand
both multi-query and multi-head attention, and everything in between. Here,
the degree of KV compression is dictated by the number of attention groups
$g$, where we alternatively refer to the generalized multi-query as multi-
group. Each attention group can be interpreted as the broadcasted attention
between a single head of key or value tensor, and multiple heads of query.

In this paradigm, multi-query attention is a special case where the number of
groups $g=1$; that is, there is exactly one such group. Conversely, multi-head
attention is another special case where the number of attention groups matches
the number of heads ($g=h$), in which case each head in the key or value
tensor attends to one head in the query. More generally, the number of groups
$g$ can lie anywhere between $1$ and $h$, indicating various degrees of
compression. For practical purposes, it is most convenient when $g$ divides
$h$. The attention mechanism in this setting can be expressed in terms of
Einstein summation as:

$\displaystyle\mathrm{logits}=\langle q,K\rangle$
$\displaystyle:\text{einsum}(bgpnk,bgmk)\to bgpnm$ (1) $\displaystyle
o=\langle w,V\rangle$ $\displaystyle:\text{einsum}(bgpmn,bgmv)\to bgpnv$ (2)

where $p=\frac{h}{g}$ represents the attention group size. Other operations in
the attention mechanism are analogous, as detailed in Appendix D.1. The memory
IO complexity for the multi-query attention becomes $bgmk$ compared to $bhmk$
in the multi-head setting, a reduction by a factor of $\frac{h}{g}$ times. The
FLOPs, however, are $bgpnmk=bdnm$, independent of the compression $g$,
implying that in the compute-bound scenario of context encoding, the latency
would be quite similar among multi-group models of different $g$’s, including
between $g=1$ and $g=h$.

This generalized multi-group attention mechanism thus provides a unified
perspective on the design space of attention architectures. By adjusting the
number of attention groups $g$, one can flexibly tune these trade-offs,
potentially yielding new regimes of performance for transformer models. In
Section 5.1, we will look into such capability vs latency trade-off.

##  4 Context-Aware Bifurcated Attention

In this section, we present a novel context-aware bifurcated attention method
that aims to reduce the memory IO cost during incremental decoding by
efficiently handling the computation of attention for shared context across
samples, as shown in Figure 2.

###  4.1 Motivation

We observe that the memory IO during the incremental decoding phase can be
significantly improved due to the fact that the KV corresponding to the
context are shared and can be loaded only once. During incremental decoding,
the accumulated key tensor ($K$) for a multi-head model is of size
$bhmk=bh(m_{c}+m_{d})k$. The two parts of $K$ correspond to $K_{c}$ of size
$bhm_{c}k$ and $K_{d}$ of size $bhm_{d}k$ where $m_{c}$ is length of the
original input and $m_{d}$ is the length due to previous incremental decoding
steps. Since tensor $K_{c}$ is the same across all indices in the $b$ axis, we
can also represent $K_{c}$ with a more compact shape $1hm_{c}k$ or simply
$hm_{c}k$. The query-key attention (Equation 1) is typically performed by
accessing different batch indices of $K=K_{c}\oplus K_{d}$ separately, even
though all batch indices in $K_{c}$ correspond to the same attention values.
That is, if we “naively” pass the entire tensor to the GEMM/BLAS operators,
the incurred memory I/O cost = $bhmk$, meaning that $K_{c}$ tensor is loaded
$b$ times (Figure 2). Since memory loading of $KV$ is the bottleneck for
incremental decoding, reducing such IO can bring significant reductions in
latency saving.

Figure 2: Context-aware bifurcated attention for single-context batch
sampling. The figure depicts the incremental decoding step where the batched
query $q$ attends with the cached key tensor $K$ where different colors in the
$q$ tensor correspond to different batch indices. The key tensor consists of
two parts: key cache corresponding to the single context $K_{c}$ (which was
computed during context encoding, as in Figure 1), and the key cache
corresponding to previous incremental decoding steps $K_{d}$. The query-key
attention is bifurcated into two parts, $\langle q,K_{c}\rangle$ and $\langle
q,K_{d}\rangle$, and joined back via concatenation, resulting in an identical
results using the same FLOPs but with lower memory IO (Eq. 3). The weight-
value attention is bifurcated similarly, as outlined in Eq. 4.

###  4.2 Formulation

Below outlines the proposed context-aware bifurcated attention for single-
context batch sampling. This operation splits any attention in the multi-group
family during incremental decoding into two parts: (1) attention associated
with KV cache from the single context $\langle q,K_{c}\rangle$ and (2)
attention associated with KV cache from prior incremental decoding steps
$\langle q,K_{d}\rangle$. That is,

$\displaystyle\langle q,K\rangle$ $\displaystyle=\langle
q,K_{c}\rangle\oplus\langle q,K_{d}\rangle$ (3) $\displaystyle\langle
q,K_{c}\rangle$ $\displaystyle:\text{einsum}(bgpnk,gm_{c}k)\to bgpnm_{c}$
$\displaystyle\langle q,K_{d}\rangle$
$\displaystyle:\text{einsum}(bgpnk,bgm_{d}k)\to bgpnm_{d}$

The context part computes attention with $K_{c}$ that corresponds to any batch
index, since they are all identical. Hence, the axis $b$ does not appears in
the einsum for $\langle q,K_{c}\rangle$. The result $\langle q,K_{c}\rangle$
and $\langle q,K_{d}\rangle$ are then joined together via concatenation. The
weight-value attention $\langle w,V\rangle$ is bifurcated similarly, where the
weight and value tensors are split along length $m$, and the results are
joined back via summation (Eq. 4). We also demonstrate the code for bifurcated
attention in Appendix E.3.

$\displaystyle\langle w,V\rangle$ $\displaystyle=\langle
w_{c},V_{c}\rangle+\langle w_{d},V_{d}\rangle$ (4) $\displaystyle\langle
w_{c},V_{c}\rangle$ $\displaystyle:\text{einsum}(bgpnm_{c},gm_{c}k)\to
bgpnk=bnd$ $\displaystyle\langle w_{d},V_{d}\rangle$
$\displaystyle:\text{einsum}(bgpnm_{d},bgm_{d}k)\to bgpnk=bnd$

The proposed operations yield the exact same results $\langle w,V\rangle$ as
the original attention in Equation 1 and 2, but can significantly reduce
memory I/O during incremental decoding (proof in Appendix E.1).

###  4.3 Memory IO Complexity

The memory IO complexity corresponding to loading KV changes from

memory IO w/o bifurcated attention $\displaystyle=gk\cdot bm$ (5)
$\displaystyle=gk\cdot b(m_{c}+m_{d})$ memory IO w. bifurcated attention
$\displaystyle=gk\cdot(m_{c}+bm_{d})$ (6)

The new memory IO is more efficient since $m_{c}+bm_{d}<b(m_{c}+m_{d})=bm.$
This resulting efficiency gain is applicable for all values of $g$ and can be
as high as $b$-fold in the case where $m_{c}>>m_{d}$ (high context length
compared to the number of generated tokens). The absolute efficiency gain,
however, is more substantially for high $g$ such as in the multi-head
attention case with $g=h$. For multi-query ($g=1$), the gain can be
substantial as well in the case of high $m_{c}$ or $b$.

##  5 Experiments

We first conduct experiments to see how capabilities scale with respect to
model size for each attention type in Section 5.1. We find that attention
types with higher compression (lower number of attention groups $g$) require
model size compensation, $\approx 10\%$ for multi-query ($g=1)$. We use such
findings to compare the latency between the multi-head and the larger multi-
query models of equal capabilities in Section 5.2. In Section 5.2.2, we focus
on the single-context batch sampling scenario where we demonstrate the
significant latency reduction of bifurcated attention and revisit the
comparison between multi-head and multi-query in light of bifurcated
attention. We outline inference details in Appendix C.5.

###  5.1 Comparing Capabilities of Multi-Head, Multi-Query, and Multi-Group
Attention

For a given model configuration, a multi-group model with $g<h$ has fewer
parameters in comparison to its multi-head counterpart. This reduction is a
result of the decreased size of the key and value projection matrices $P_{K}$
and $P_{V}$. Specifically, each tensor in this case has a size of
$P_{K}:d\times gk$, where $k$ is the head dimension. For instance, a 13B
multi-head model will correspond to a 11B multi-query model, with all other
model configurations fixed (see Appendix D.1 for more details).

To compare the capabilities of different attention mechanisms, one can either
scale other model configurations such as the number of layers $\ell$, the
number of heads $h$ in order to make match the total model sizes between
different attentions. However, it is often difficult to match the number of
parameters exactly. In this work, we compare different attention mechanisms
via the loss-vs-size scaling laws. For the setup, we use the model
hyperparameters similar to that of GPT-3, where the size ranges from $125$M to
$13$B, with hyperparameters such as $\ell,h,k$ increasing in tandem. Then, we
consider three cases where $g=1$ (multi-query), $g=h$ (multi-head) and $1<g<h$
(multi-group) where Appendix C.1 and C.2 shows the training and model
configuration details. We train all three attention models of each size and
plot the validation loss versus model size, shown in Figure 3. Our findings
are summarized below.

##### Higher number of attention groups $g$ leads to higher expressiveness

The results in Figure 3 shows the validation loss versus model size (log
scale). The results indicate that, for the same model size (vertical slice
across the plot), multi-head attention $g=h$ achieves the lowest validation
loss compared to $1<g<h$ (multi-group) and $g=1$ (multi-query). This trend
holds consistently over three orders of magnitude of model sizes, where the
curves corresponding to multi-head, multi-group and multi-query do not cross,
implying that the rank of model expressiveness, or relative capabilities per
number of parameters, is quite stable. An intuitive explanation is that the
lower $g$ corresponds to a lower rank representation of the key and value
tensors, which encodes lower representation power of the past context and
therefore yields lower capabilities than higher $g$, given the same model
size.

##### Scaling laws via downstream performance

We use the average scores from two code generation benchmarks, multi-lingual
HumanEval and MBXP (Athiwaratkun et al., 2022), as a proxy for model
capabilities in addition to the validation loss. This approach is similar to
that of the GPT-4 technical report (OpenAI, 2023) where HumanEval (Python)
(Chen et al., 2021) is used to track the performance across multiple
magnitudes of compute. In our case, we average across all 13 evaluation
languages and two benchmarks to obtain a more stable proxy for capabilities.
The result in Figure 3 demonstrates similar trend compared to the validation
loss where the pass rate curves indicate the same relative expressiveness for
multi-head, multi-group and multi-query attention.

##### Matching capabilities by model size compensation

Given the same capabilities (horizontal slice of the plot in Figure 3), the
distance between two curves indicates the model size difference that the
lower-rank attention needs to compensate in order to match the multi-head
model performance. Empirically, we average the distance along the interpolated
lines (log scale) and find this to correspond to $1.104$ times; that is, a
multi-query model can have the same capabilities as the multi-head model if
the size is increased by $\approx 10\%$ of the multi-head model size.
Similarly, the gap is $<10\%$ for multi-group attention. Alternatively, one
can argue that a multi-query model of the same size could match a multi-head
if the multi-query model is given more compute. However, in the regime where
we train language models until or close to convergence and the performance
saturates with respect to compute, the difference in capabilities will likely
remain. Therefore, the size compensation is likely the most fair approach for
comparison.

Figure 3: (Left) The plots of validation loss versus model size demonstrate
that the scaling laws curves of different attention mechanisms have different
expressiveness or performance efficiency. That is, the capabilities given the
same model size depends on $g$ where higher $g$ yields the best capabilities.
(Right) We demonstrate a similar trend where we use code generation abilities
as a proxy for general capabilities. Here, we average the execution pass rates
evaluated on Multi-lingual HumanEval and MBXP benchmarks under 13 programming
languages.

###  5.2 Latencies of Capabilities-Equivalent Models

As detailed in Section 5.1, we’ve observed that an increase in the multi-query
model’s size is required for it to match the performance of a multi-head
model. In this section, we focus on examining the latency trade-offs across
diverse scenarios with both multi-query and multi-head models of similar
performance capabilities. For these latency experiments, we utilize two
models, each with an approximate size of 1 billion: a multi-head model and a
multi-query model (detailed information can be found in C.3). The multi-query
model chosen for these studies is larger by a multiplicative factor $F$, where
$F=1.1$.

Overall, there is some overhead cost of using multi-query attention due to the
larger size (see Figure 4 and Appendix D.3.1 and D.3.2 for analysis). That is,
context encoding latency of the multi-query model will be slightly larger, as
well as the low-context and low-batch incremental decoding scenario. However,
multi-query can have significantly lower latency compared to multi-head in the
scenario with high number of decoding steps which makes the incremental
decoding phase being latency-dominating, and high context or batch size which
heavily impacts the memory IO of incremental decoding. We outline three
different inference scenarios below.

Figure 4: High-level latency comparison between an MH model and a larger MQ
model with comparable capabilities. Overall, there’s an overhead cost for the
initial context encoding latency due the additional compute with the larger MQ
model size. For low context and batch size, the per step latency of MQ is also
slightly higher to start due to the memory IO required for larger model size,
but does not change much as context length $m$ or batch size $b$ grow, as
supposed to the multi-head case where the per step latency can grow more
rapidly with respect to $m$ and $b$.

####  5.2.1 Single Context Scenario

In the single batch inference scenario, the multi-query/-group attention can
achieve lower latency when the context length and the number of generated
tokens are high, as demonstrated in Figure 5. Different implementations that
are more efficient in loading KV cache (such as lower-level kernel that can
avoid duplicated IO) can cause the overall curves of MH to be flatter.
However, the overall trend still remains where given sufficiently high context
$m$, MQ will begin to be faster than MH.

Figure 5:  Incremental decoding (per step) latency and the context encoding
latency, as a function of input context length. In this plot, we compare an
multi-head model and an multi-query model of comparable capabilities, whose
size is slightly larger. (Leftmost: Per-step incremental decoding latency) For
low context length such as $m<2500$, due to the larger size of the MQ model,
the inference latency is higher. However, the growth with respect to context
length of the MQ model is much lower (almost flat), resulting in lower per
step latency when the context length is high. (Second: Context encoding
latency) The context encoding latency depends on the FLOPs where the MH and MQ
are quite similar. Note that the MQ model is slightly larger, and therefore
corresponds to a steeper curve. (Third, Fourth): Total latency for 15 or 256
generated steps The two plots illustrates the _total_ latency, which is the
sum of context encoding and the the number of steps times incremental decoding
latency. The benefits of MQ model becomes clear in the case of high decoding
steps $(256)$ whereas in the case of $15$ generated tokens, the total latency
of MQ can still be slightly higher than MH.

(a) Multi-Head

(b) Multi-Query

Figure 6: Context-aware bifurcated attention with multi-head attention (a) and
multi-query attention (b). The bifurcated attention loads the KV cache in a
context-aware manner, resulting in significantly lower latency for sampling
under high batch sizes. For instance, in the case of multi-head attention with
batch size $128$ and context length $10,000$, bifurcated attention results in
$\approx 4\times$ lower the incremental decoding latency. Additionally, growth
with respect to context length is relatively flat with bifurcated attention.
With multi-query attention, bifurcated attention permits us to use batch sizes
as high as $256$ or $512$ with lower latency than in the multi-head scenario.

####  5.2.2 Single-Context Batch Sampling

In this scenario, we are given a single context and generates multiple
completions based on temperature sampling. In this case, the context encoding
is independent of the batch size $b$ since it is performed on the single
context and broadcasted for other batch indices (Figure 1). In contrast to the
batch inference scenario, this is a more practical online inference scenario
since we are not bottlenecked by the context encoding step. Our proposed
context-aware bifurcated attention is exactly applicable for such scenario
where in this section we demonstrate the results in conjunction with both
multi-head and multi-query.

##### Multi-head benefits significantly from bifurcated attention

Figure 6(a) demonstrates the per step latency results for a multi-head model.
For instance, with batch size $8$, the per step latency without bifurcated
attention grows rapidly with context length, from $\approx 10$ ms to $\approx
100$ ms at context length $10000$. However, with bifurcated attention, the
latency remains relatively flat with respect to context length. In practice,
bifurcated attention also reduces memory consumption at high batch size and
context lengths without encountering out-of-memory error as early as without
bifurcated attention.

(a) Without bifurcated attention

(b) With bifurcated attention

Figure 7: Latency comparison between multi-head and a larger multi-query model
of equal capabilities. Without bifurcated attention, MQ is clearly much more
inference efficient. However, with bifurcated attention, MH can have better
latency than MQ in moderate scenario (up to batch size 64 in this case) where
MQ can handle more extreme scenarios better than MH.  Figure 8: Bifurcated
attention improves accuracy by enabling more generated samples over a fixed
latency budget, applicable for both multi-head attention (CodeGen) and multi-
query attention (StarCoder). Given the $n$ samples, pass@n reflects the
execution pass rate of the best sample among $n$, shown in (a) and (c).
Filtering $n$ samples with mean log probability ranking yields a subset of
best three samples, reflected by pass@top3 in (b) and (d). The increased
number of samples within the same latency budget results in increased
performance via either pass@n or pass@top-k.

##### Bifurcated attention + multi-head rivals multi-query

Figure 7 shows the comparison between MH and MQ with and without bifurcated
attention. Without bifurcated attention, MQ is clearly much more inference
efficient. However, with bifurcated attention, MQ and MH under moderate batch
size scenarios (up to 64) seems comparable, where multi-head is even has lower
latency. The results indicate that, given an existing MH model, we can support
batch sampling scenarios using bifurcated attention without the need of a
multi-query model (which requires training a new model, or at least continuous
training) (Ainslie et al., 2023). With a more inference-intensive scenarios,
including batch inference scenario where the bifurcated attention is not
applicable, switching to multi-query can be worth the effort.

##### Bifurcated attention with multi-query enables more extreme batch size
and context lengths

Multi-query has overall $h$ times lower memory IO and can already reduce
latency for some inference scenarios. With bifurcated attention, the supported
context lengths and batch sizes can become much more extreme, as demonstrated
in Figure 6(b).

###  5.3 Applications

Efficient large-scale sampling is particularly useful for downstream
applications that require multiple generations but has latency constraints,
e.g., AI code assistants. In this case, bifurcated attention enables
generating more candidates by using larger batch size without incurring much
additional latency. To verify our point, we empirically evaluate
CodeGen-16B-mono (Nijkamp et al., 2022) and StarCoder (15.5B) (Li et al.,
2023) on MBPP dataset (Austin et al., 2021), and plot pass rates with respect
to latency in Figure 8, where we also indicate the batch size $n$. We consider
two accuracy measurements: (1) pass@$n$ corresponds to the oracle scenario,
where we evaluate all the generated samples and check if any of them is
correct; (2) pass@top3, where we are only allowed to evaluate three examples
no matter how many we generate. In the top-3 case, we deduplicate the $n$
samples, and rank by their mean log probability scores (Chen et al., 2021) to
determine three candidates. All experiments use nucleus sampling with $p=0.95$
(Holtzman et al., 2020) and temperature $0.8$. The results show much sharper
improvement in either metrics relative to additional latency. This approach
opens up avenues for performance improvement given a fixed budget of latency.

##  6 Conclusion

Bifurcated attention provides a complementary approach to the existing
inference acceleration methods, with a particular focus on minimizing the
memory IO of the incremental decoding, thereby enhancing inference efficiency.
Our work helps support demanding inference scenarios due to larger context
during incremental decoding, which are emerging from, e.g., more complex
applications that requires long context such as complex reasoning, planning,
or retrieval augmented generations.

## References

  * Ahmad et al. [2021] W. U. Ahmad, S. Chakraborty, B. Ray, and K.-W. Chang.  Unified pre-training for program understanding and generation.  _arXiv preprint arXiv:2103.06333_ , 2021.
  * Ainslie et al. [2023] J. Ainslie, J. Lee-Thorp, M. de Jong, Y. Zemlyanskiy, F. Lebrón, and S. Sanghai.  GQA: training generalized multi-query transformer models from multi-head checkpoints.  _CoRR_ , abs/2305.13245, 2023.  doi: 10.48550/arXiv.2305.13245.  URL https://doi.org/10.48550/arXiv.2305.13245.
  * Allal et al. [2023] L. B. Allal, R. Li, D. Kocetkov, C. Mou, C. Akiki, C. M. Ferrandis, N. Muennighoff, M. Mishra, A. Gu, M. Dey, et al.  Santacoder: don’t reach for the stars!  _arXiv preprint arXiv:2301.03988_ , 2023.
  * Amazon [2022] Amazon.  Amazon code whisperer.  https://aws.amazon.com/codewhisperer/, 2022.
  * Athiwaratkun et al. [2022] B. Athiwaratkun, S. K. Gouda, Z. Wang, X. Li, Y. Tian, M. Tan, W. U. Ahmad, S. Wang, Q. Sun, M. Shang, S. K. Gonugondla, H. Ding, V. Kumar, N. Fulton, A. Farahani, S. Jain, R. Giaquinto, H. Qian, M. K. Ramanathan, R. Nallapati, B. Ray, P. Bhatia, S. Sengupta, D. Roth, and B. Xiang.  Multi-lingual evaluation of code generation models.  _CoRR_ , abs/2210.14868, 2022.  doi: 10.48550/arXiv.2210.14868.  URL https://doi.org/10.48550/arXiv.2210.14868.
  * Austin et al. [2021] J. Austin, A. Odena, M. I. Nye, M. Bosma, H. Michalewski, D. Dohan, E. Jiang, C. J. Cai, M. Terry, Q. V. Le, and C. Sutton.  Program synthesis with large language models.  _CoRR_ , abs/2108.07732, 2021.  URL https://arxiv.org/abs/2108.07732.
  * Beltagy et al. [2020] I. Beltagy, M. E. Peters, and A. Cohan.  Longformer: The long-document transformer.  _arXiv preprint arXiv:2004.05150_ , 2020.
  * Brown et al. [2020] T. B. Brown, B. Mann, N. Ryder, M. Subbiah, J. Kaplan, P. Dhariwal, A. Neelakantan, P. Shyam, G. Sastry, A. Askell, S. Agarwal, A. Herbert-Voss, G. Krueger, T. Henighan, R. Child, A. Ramesh, D. M. Ziegler, J. Wu, C. Winter, C. Hesse, M. Chen, E. Sigler, M. Litwin, S. Gray, B. Chess, J. Clark, C. Berner, S. McCandlish, A. Radford, I. Sutskever, and D. Amodei.  Language models are few-shot learners.  _CoRR_ , abs/2005.14165, 2020.  URL https://arxiv.org/abs/2005.14165.
  * Bulatov et al. [2023] A. Bulatov, Y. Kuratov, and M. S. Burtsev.  Scaling transformer to 1m tokens and beyond with rmt.  _arXiv preprint arXiv:2304.11062_ , 2023.
  * Cai et al. [2024] T. Cai, Y. Li, Z. Geng, H. Peng, J. D. Lee, D. Chen, and T. Dao.  Medusa: Simple llm inference acceleration framework with multiple decoding heads, 2024.
  * Chen et al. [2023] C. Chen, S. Borgeaud, G. Irving, J. Lespiau, L. Sifre, and J. Jumper.  Accelerating large language model decoding with speculative sampling.  _CoRR_ , abs/2302.01318, 2023.  doi: 10.48550/arXiv.2302.01318.  URL https://doi.org/10.48550/arXiv.2302.01318.
  * Chen et al. [2021] M. Chen, J. Tworek, H. Jun, Q. Yuan, H. P. de Oliveira Pinto, J. Kaplan, H. Edwards, Y. Burda, N. Joseph, G. Brockman, A. Ray, R. Puri, G. Krueger, M. Petrov, H. Khlaaf, G. Sastry, P. Mishkin, B. Chan, S. Gray, N. Ryder, M. Pavlov, A. Power, L. Kaiser, M. Bavarian, C. Winter, P. Tillet, F. P. Such, D. Cummings, M. Plappert, F. Chantzis, E. Barnes, A. Herbert-Voss, W. H. Guss, A. Nichol, A. Paino, N. Tezak, J. Tang, I. Babuschkin, S. Balaji, S. Jain, W. Saunders, C. Hesse, A. N. Carr, J. Leike, J. Achiam, V. Misra, E. Morikawa, A. Radford, M. Knight, M. Brundage, M. Murati, K. Mayer, P. Welinder, B. McGrew, D. Amodei, S. McCandlish, I. Sutskever, and W. Zaremba.  Evaluating large language models trained on code.  _CoRR_ , abs/2107.03374, 2021.  URL https://arxiv.org/abs/2107.03374.
  * Child et al. [2019] R. Child, S. Gray, A. Radford, and I. Sutskever.  Generating long sequences with sparse transformers.  _URL https://openai.com/blog/sparse-transformers_ , 2019.
  * Choquette et al. [2021] J. Choquette, W. Gandhi, O. Giroux, N. Stam, and R. Krashinsky.  Nvidia a100 tensor core gpu: Performance and innovation.  _IEEE Micro_ , 41(2):29–35, 2021.
  * Chowdhery et al. [2022] A. Chowdhery, S. Narang, J. Devlin, M. Bosma, G. Mishra, A. Roberts, P. Barham, H. W. Chung, C. Sutton, S. Gehrmann, P. Schuh, K. Shi, S. Tsvyashchenko, J. Maynez, A. Rao, P. Barnes, Y. Tay, N. Shazeer, V. Prabhakaran, E. Reif, N. Du, B. Hutchinson, R. Pope, J. Bradbury, J. Austin, M. Isard, G. Gur-Ari, P. Yin, T. Duke, A. Levskaya, S. Ghemawat, S. Dev, H. Michalewski, X. Garcia, V. Misra, K. Robinson, L. Fedus, D. Zhou, D. Ippolito, D. Luan, H. Lim, B. Zoph, A. Spiridonov, R. Sepassi, D. Dohan, S. Agrawal, M. Omernick, A. M. Dai, T. S. Pillai, M. Pellat, A. Lewkowycz, E. Moreira, R. Child, O. Polozov, K. Lee, Z. Zhou, X. Wang, B. Saeta, M. Diaz, O. Firat, M. Catasta, J. Wei, K. Meier-Hellstern, D. Eck, J. Dean, S. Petrov, and N. Fiedel.  Palm: Scaling language modeling with pathways.  _CoRR_ , abs/2204.02311, 2022.  doi: 10.48550/arXiv.2204.02311.  URL https://doi.org/10.48550/arXiv.2204.02311.
  * Costa-jussà et al. [2022] M. R. Costa-jussà, J. Cross, O. Çelebi, M. Elbayad, K. Heafield, K. Heffernan, E. Kalbassi, J. Lam, D. Licht, J. Maillard, et al.  No language left behind: Scaling human-centered machine translation.  _arXiv preprint arXiv:2207.04672_ , 2022.
  * Dao et al. [2022] T. Dao, D. Y. Fu, S. Ermon, A. Rudra, and C. Ré.  Flashattention: Fast and memory-efficient exact attention with io-awareness.  In _NeurIPS_ , 2022.  URL http://papers.nips.cc/paper_files/paper/2022/hash/67d57c32e20fd0a7a302cb81d36e40d5-Abstract-Conference.html.
  * Dathathri et al. [2019] S. Dathathri, A. Madotto, J. Lan, J. Hung, E. Frank, P. Molino, J. Yosinski, and R. Liu.  Plug and play language models: A simple approach to controlled text generation.  _arXiv preprint arXiv:1912.02164_ , 2019.
  * Dettmers et al. [2022] T. Dettmers, M. Lewis, Y. Belkada, and L. Zettlemoyer.  Llm.int8(): 8-bit matrix multiplication for transformers at scale.  _CoRR_ , abs/2208.07339, 2022.  doi: 10.48550/arXiv.2208.07339.  URL https://doi.org/10.48550/arXiv.2208.07339.
  * Farhad et al. [2021] A. Farhad, A. Arkady, B. Magdalena, B. Ondřej, C. Rajen, C. Vishrav, M. R. Costa-jussa, E.-B. Cristina, F. Angela, F. Christian, et al.  Findings of the 2021 conference on machine translation (wmt21).  In _Proceedings of the Sixth Conference on Machine Translation_ , pages 1–88. Association for Computational Linguistics, 2021.
  * Frantar et al. [2022] E. Frantar, S. Ashkboos, T. Hoefler, and D. Alistarh.  Gptq: Accurate post-training quantization for generative pre-trained transformers.  _arXiv preprint arXiv:2210.17323_ , 2022.
  * Fried et al. [2022] D. Fried, A. Aghajanyan, J. Lin, S. Wang, E. Wallace, F. Shi, R. Zhong, W.-t. Yih, L. Zettlemoyer, and M. Lewis.  Incoder: A generative model for code infilling and synthesis.  _arXiv preprint arXiv:2204.05999_ , 2022.
  * Fu et al. [2023] Y. Fu, P. Bailis, I. Stoica, and H. Zhang.  Breaking the sequential dependency of llm inference using lookahead decoding, November 2023.  URL https://lmsys.org/blog/2023-11-21-lookahead-decoding/.
  * Gehman et al. [2020] S. Gehman, S. Gururangan, M. Sap, Y. Choi, and N. A. Smith.  Realtoxicityprompts: Evaluating neural toxic degeneration in language models.  _arXiv preprint arXiv:2009.11462_ , 2020.
  * Google [2023] Google.  Bard.  https://blog.google/technology/ai/try-bard/, 2023.
  * Guu et al. [2020] K. Guu, K. Lee, Z. Tung, P. Pasupat, and M. Chang.  Retrieval augmented language model pre-training.  In _International conference on machine learning_ , pages 3929–3938. PMLR, 2020.
  * Hoffmann et al. [2022] J. Hoffmann, S. Borgeaud, A. Mensch, E. Buchatskaya, T. Cai, E. Rutherford, D. d. L. Casas, L. A. Hendricks, J. Welbl, A. Clark, et al.  Training compute-optimal large language models.  _arXiv preprint arXiv:2203.15556_ , 2022.
  * Holtzman et al. [2020] A. Holtzman, J. Buys, L. Du, M. Forbes, and Y. Choi.  The curious case of neural text degeneration.  In _8th International Conference on Learning Representations, ICLR 2020, Addis Ababa, Ethiopia, April 26-30, 2020_. OpenReview.net, 2020.  URL https://openreview.net/forum?id=rygGQyrFvH.
  * Izacard et al. [2022] G. Izacard, P. Lewis, M. Lomeli, L. Hosseini, F. Petroni, T. Schick, J. Dwivedi-Yu, A. Joulin, S. Riedel, and E. Grave.  Atlas: Few-shot learning with retrieval augmented language models.  _arXiv preprint arXiv_ , 2208, 2022.
  * Jia et al. [2018] Z. Jia, M. Maggioni, B. Staiger, and D. P. Scarpazza.  Dissecting the NVIDIA volta GPU architecture via microbenchmarking.  _CoRR_ , abs/1804.06826, 2018.  URL http://arxiv.org/abs/1804.06826.
  * Kalamkar et al. [2019] D. Kalamkar, D. Mudigere, N. Mellempudi, D. Das, K. Banerjee, S. Avancha, D. T. Vooturi, N. Jammalamadaka, J. Huang, H. Yuen, et al.  A study of bfloat16 for deep learning training.  _arXiv preprint arXiv:1905.12322_ , 2019.
  * Kaplan et al. [2020] J. Kaplan, S. McCandlish, T. Henighan, T. B. Brown, B. Chess, R. Child, S. Gray, A. Radford, J. Wu, and D. Amodei.  Scaling laws for neural language models.  _CoRR_ , abs/2001.08361, 2020.  URL https://arxiv.org/abs/2001.08361.
  * Kingma and Ba [2014] D. P. Kingma and J. Ba.  Adam: A method for stochastic optimization.  _arXiv preprint arXiv:1412.6980_ , 2014.
  * Kuzmin et al. [2022] A. Kuzmin, M. Van Baalen, Y. Ren, M. Nagel, J. Peters, and T. Blankevoort.  Fp8 quantization: The power of the exponent.  _arXiv preprint arXiv:2208.09225_ , 2022.
  * Kwon et al. [2023] W. Kwon, Z. Li, S. Zhuang, Y. Sheng, L. Zheng, C. H. Yu, J. E. Gonzalez, H. Zhang, and I. Stoica.  Efficient memory management for large language model serving with pagedattention, 2023.
  * Le et al. [2022] H. Le, Y. Wang, A. D. Gotmare, S. Savarese, and S. C. H. Hoi.  Coderl: Mastering code generation through pretrained models and deep reinforcement learning.  _Advances in Neural Information Processing Systems_ , 35:21314–21328, 2022.
  * Leviathan et al. [2022] Y. Leviathan, M. Kalman, and Y. Matias.  Fast inference from transformers via speculative decoding.  _CoRR_ , abs/2211.17192, 2022.  doi: 10.48550/arXiv.2211.17192.  URL https://doi.org/10.48550/arXiv.2211.17192.
  * Li et al. [2023] R. Li, L. B. Allal, Y. Zi, N. Muennighoff, D. Kocetkov, C. Mou, M. Marone, C. Akiki, J. Li, J. Chim, et al.  Starcoder: may the source be with you!  _arXiv preprint arXiv:2305.06161_ , 2023.
  * Li et al. [2022] Y. Li, D. Choi, J. Chung, N. Kushman, J. Schrittwieser, R. Leblond, T. Eccles, J. Keeling, F. Gimeno, A. Dal Lago, et al.  Competition-level code generation with alphacode.  _Science_ , 378(6624):1092–1097, 2022.
  * Li et al. [2024] Y. Li, F. Wei, C. Zhang, and H. Zhang.  Eagle: Speculative sampling requires rethinking feature uncertainty, 2024\.
  * Lin and Riedl [2021] Z. Lin and M. Riedl.  Plug-and-blend: A framework for controllable story generation with blended control codes.  _arXiv preprint arXiv:2104.04039_ , 2021.
  * Loshchilov and Hutter [2017] I. Loshchilov and F. Hutter.  Decoupled weight decay regularization.  _arXiv preprint arXiv:1711.05101_ , 2017.
  * Madaan et al. [2023] A. Madaan, A. Shypula, U. Alon, M. Hashemi, P. Ranganathan, Y. Yang, G. Neubig, and A. Yazdanbakhsh.  Learning performance-improving code edits.  _arXiv preprint arXiv:2302.07867_ , 2023.
  * Menick et al. [2022] J. Menick, M. Trebacz, V. Mikulik, J. Aslanides, F. Song, M. Chadwick, M. Glaese, S. Young, L. Campbell-Gillingham, G. Irving, et al.  Teaching language models to support answers with verified quotes.  _arXiv preprint arXiv:2203.11147_ , 2022.
  * Miao et al. [2023] X. Miao, G. Oliaro, Z. Zhang, X. Cheng, Z. Wang, R. Y. Y. Wong, Z. Chen, D. Arfeen, R. Abhyankar, and Z. Jia.  Specinfer: Accelerating generative LLM serving with speculative inference and token tree verification.  _CoRR_ , abs/2305.09781, 2023.  doi: 10.48550/ARXIV.2305.09781.  URL https://doi.org/10.48550/arXiv.2305.09781.
  * [46] Microsoft.  Github copilot.  https://github.com/features/copilot.
  * Mirowski et al. [2023] P. Mirowski, K. W. Mathewson, J. Pittman, and R. Evans.  Co-writing screenplays and theatre scripts with language models: Evaluation by industry professionals.  In _Proceedings of the 2023 CHI Conference on Human Factors in Computing Systems_ , pages 1–34, 2023.
  * Nadeem et al. [2020] M. Nadeem, A. Bethke, and S. Reddy.  Stereoset: Measuring stereotypical bias in pretrained language models.  _arXiv preprint arXiv:2004.09456_ , 2020.
  * Nakano et al. [2021] R. Nakano, J. Hilton, S. Balaji, J. Wu, L. Ouyang, C. Kim, C. Hesse, S. Jain, V. Kosaraju, W. Saunders, et al.  Webgpt: Browser-assisted question-answering with human feedback.  _arXiv preprint arXiv:2112.09332_ , 2021.
  * Nijkamp et al. [2022] E. Nijkamp, B. Pang, H. Hayashi, L. Tu, H. Wang, Y. Zhou, S. Savarese, and C. Xiong.  Codegen: An open large language model for code with multi-turn program synthesis.  _arXiv preprint arXiv:2203.13474_ , 2022.
  * Nijkamp et al. [2023] E. Nijkamp, H. Hayashi, C. Xiong, S. Savarese, and Y. Zhou.  Codegen2: Lessons for training llms on programming and natural languages.  _arXiv preprint arXiv:2305.02309_ , 2023.
  * [52] NVIDIA.  Fastertransformer.  https://github.com/NVIDIA/FasterTransformer.
  * OpenAI [2023] OpenAI.  GPT-4 technical report.  _CoRR_ , abs/2303.08774, 2023.  doi: 10.48550/arXiv.2303.08774.  URL https://doi.org/10.48550/arXiv.2303.08774.
  * Paszke et al. [2019] A. Paszke, S. Gross, F. Massa, A. Lerer, J. Bradbury, G. Chanan, T. Killeen, Z. Lin, N. Gimelshein, L. Antiga, A. Desmaison, A. Köpf, E. Z. Yang, Z. DeVito, M. Raison, A. Tejani, S. Chilamkurthy, B. Steiner, L. Fang, J. Bai, and S. Chintala.  Pytorch: An imperative style, high-performance deep learning library.  In H. M. Wallach, H. Larochelle, A. Beygelzimer, F. d’Alché-Buc, E. B. Fox, and R. Garnett, editors, _Advances in Neural Information Processing Systems 32: Annual Conference on Neural Information Processing Systems 2019, NeurIPS 2019, December 8-14, 2019, Vancouver, BC, Canada_ , pages 8024–8035, 2019.  URL https://proceedings.neurips.cc/paper/2019/hash/bdbca288fee7f92f2bfa9f7012727740-Abstract.html.
  * Pearce et al. [2022] H. Pearce, B. Ahmad, B. Tan, B. Dolan-Gavitt, and R. Karri.  Asleep at the keyboard? assessing the security of github copilot’s code contributions.  In _2022 IEEE Symposium on Security and Privacy (SP)_ , pages 754–768. IEEE, 2022.
  * Pope et al. [2022] R. Pope, S. Douglas, A. Chowdhery, J. Devlin, J. Bradbury, A. Levskaya, J. Heek, K. Xiao, S. Agrawal, and J. Dean.  Efficiently scaling transformer inference.  _CoRR_ , abs/2211.05102, 2022.  doi: 10.48550/arXiv.2211.05102.  URL https://doi.org/10.48550/arXiv.2211.05102.
  * Rasley et al. [2020] J. Rasley, S. Rajbhandari, O. Ruwase, and Y. He.  Deepspeed: System optimizations enable training deep learning models with over 100 billion parameters.  In _Proceedings of the 26th ACM SIGKDD International Conference on Knowledge Discovery & Data Mining_, pages 3505–3506, 2020.
  * Roziere et al. [2020] B. Roziere, M.-A. Lachaux, L. Chanussot, and G. Lample.  Unsupervised translation of programming languages.  _Advances in Neural Information Processing Systems_ , 33:20601–20611, 2020.
  * Schick et al. [2023] T. Schick, J. Dwivedi-Yu, R. Dessì, R. Raileanu, M. Lomeli, L. Zettlemoyer, N. Cancedda, and T. Scialom.  Toolformer: Language models can teach themselves to use tools.  _arXiv preprint arXiv:2302.04761_ , 2023.
  * Shazeer [2019] N. Shazeer.  Fast transformer decoding: One write-head is all you need.  _CoRR_ , abs/1911.02150, 2019.  URL http://arxiv.org/abs/1911.02150.
  * Shoeybi et al. [2019] M. Shoeybi, M. Patwary, R. Puri, P. LeGresley, J. Casper, and B. Catanzaro.  Megatron-lm: Training multi-billion parameter language models using model parallelism.  _arXiv preprint arXiv:1909.08053_ , 2019.
  * Team [2023] M. N. Team.  Introducing mpt-7b: A new standard for open-source, commercially usable llms.  2023\.  URL www.mosaicml.com/blog/mpt-7b.
  * Touvron et al. [2023] H. Touvron, T. Lavril, G. Izacard, X. Martinet, M. Lachaux, T. Lacroix, B. Rozière, N. Goyal, E. Hambro, F. Azhar, A. Rodriguez, A. Joulin, E. Grave, and G. Lample.  Llama: Open and efficient foundation language models.  _CoRR_ , abs/2302.13971, 2023.  doi: 10.48550/arXiv.2302.13971.  URL https://doi.org/10.48550/arXiv.2302.13971.
  * Tran et al. [2021] C. Tran, S. Bhosale, J. Cross, P. Koehn, S. Edunov, and A. Fan.  Facebook ai wmt21 news translation task submission.  _arXiv preprint arXiv:2108.03265_ , 2021.
  * Vaswani et al. [2017] A. Vaswani, N. Shazeer, N. Parmar, J. Uszkoreit, L. Jones, A. N. Gomez, L. Kaiser, and I. Polosukhin.  Attention is all you need.  In I. Guyon, U. von Luxburg, S. Bengio, H. M. Wallach, R. Fergus, S. V. N. Vishwanathan, and R. Garnett, editors, _Advances in Neural Information Processing Systems 30: Annual Conference on Neural Information Processing Systems 2017, December 4-9, 2017, Long Beach, CA, USA_ , pages 5998–6008, 2017.  URL https://proceedings.neurips.cc/paper/2017/hash/3f5ee243547dee91fbd053c1c4a845aa-Abstract.html.
  * Wang et al. [2020] S. Wang, B. Z. Li, M. Khabsa, H. Fang, and H. Ma.  Linformer: Self-attention with linear complexity.  _arXiv preprint arXiv:2006.04768_ , 2020.
  * Wei et al. [2023] X. Wei, S. K. Gonugondla, W. U. Ahmad, S. Wang, B. Ray, H. Qian, X. Li, V. Kumar, Z. Wang, Y. Tian, Q. Sun, B. Athiwaratkun, M. Shang, M. K. Ramanathan, P. Bhatia, and B. Xiang.  Greener yet powerful: Taming large code generation models with quantization.  _CoRR_ , abs/2303.05378, 2023.  doi: 10.48550/arXiv.2303.05378.  URL https://doi.org/10.48550/arXiv.2303.05378.
  * Wolf et al. [2019] T. Wolf, L. Debut, V. Sanh, J. Chaumond, C. Delangue, A. Moi, P. Cistac, T. Rault, R. Louf, M. Funtowicz, et al.  Huggingface’s transformers: State-of-the-art natural language processing.  _arXiv preprint arXiv:1910.03771_ , 2019.
  * Xiao et al. [2022] G. Xiao, J. Lin, M. Seznec, J. Demouth, and S. Han.  Smoothquant: Accurate and efficient post-training quantization for large language models.  _arXiv preprint arXiv:2211.10438_ , 2022.
  * Yao et al. [2022] Z. Yao, R. Y. Aminabadi, M. Zhang, X. Wu, C. Li, and Y. He.  Zeroquant: Efficient and affordable post-training quantization for large-scale transformers.  In _NeurIPS_ , 2022.  URL http://papers.nips.cc/paper_files/paper/2022/hash/adf7fa39d65e2983d724ff7da57f00ac-Abstract-Conference.html.
  * Yee et al. [2019] K. Yee, N. Ng, Y. N. Dauphin, and M. Auli.  Simple and effective noisy channel modeling for neural machine translation.  _arXiv preprint arXiv:1908.05731_ , 2019.
  * Yuan et al. [2022] A. Yuan, A. Coenen, E. Reif, and D. Ippolito.  Wordcraft: story writing with large language models.  In _27th International Conference on Intelligent User Interfaces_ , pages 841–852, 2022.
  * Zaheer et al. [2020] M. Zaheer, G. Guruganesh, K. A. Dubey, J. Ainslie, C. Alberti, S. Ontañón, P. Pham, A. Ravula, Q. Wang, L. Yang, and A. Ahmed.  Big bird: Transformers for longer sequences.  In H. Larochelle, M. Ranzato, R. Hadsell, M. Balcan, and H. Lin, editors, _Advances in Neural Information Processing Systems 33: Annual Conference on Neural Information Processing Systems 2020, NeurIPS 2020, December 6-12, 2020, virtual_ , 2020.  URL https://proceedings.neurips.cc/paper/2020/hash/c8512d142a2d849725f31a9a7a361ab9-Abstract.html.
  * Zhen et al. [2022] C. Zhen, Y. Shang, X. Liu, Y. Li, Y. Chen, and D. Zhang.  A survey on knowledge-enhanced pre-trained language models.  _arXiv preprint arXiv:2212.13428_ , 2022.

##  Appendix A FAQs

  1. 1.

Q: If we already have an MQ model that seems to be quite efficient at large
batch sampling, is bifurcated attention necessary?
A: The proposed context-aware bifurcated attention is an exact computation
that provides a different way to perform attention, so one can use it "for
free" without a performance tradeoff. Due to the reduced memory I/O, it
enables more extreme cases of batch sampling, such as a larger batch, even for
long contexts.

  2. 2.

Q: How applicable is multi-query for single-batch inference without high batch
sampling?
A: If the context is long and the number of generated tokens is high, then the
benefits of multi-query are clear. Please see Section 5.2.1.

  3. 3.

Q: Is bifurcated attention applicable for the case where we process different
inputs in a batch?
A: No. In that case, if we need a solution to reduce memory I/O during
incremental decoding, then multi-query attention can be appealing, especially
in scenarios with a high number of generated tokens where the incremental
decoding phase dominates the overall latency. This is because there is an
overhead to multi-query due to the context encoding phase, as outlined in the
main paper.

  4. 4.

Q: Any caveats to using bifurcated attention?
A: For small workloads (low context length and batch size), due to the fact
that we split the attention into two parts, there can be less parallelization
of the GEMM kernels, which could lead to higher latency, especially for MQ
models. However, one can get the best of both worlds given any model by
triggering bifurcated attention under high workload scenarios and using normal
attention otherwise. With such a workload-based switch, bifurcated attention
is guaranteed to provide better latency and efficiency.

  5. 5.

Q: How does model quantization (or lower precision arithmetic) affect the
findings?
A: There are two regimes for quantization: model weight quantization and
attention quantization. To date, most quantization only focuses on the weight
since the attention computation is precision-sensitive and quantization has
not proved to be viable.

Model quantization can make incremental decoding faster due to lower memory
I/O of the model itself, since the effective model size in memory is smaller.
This shifts the latency curve downward for all context lengths or batch sizes.
The overall conclusion for the bifurcated and multi-query attention remains
the same, however, since the improvement proposed in the paper is on the
attention component, which is orthogonal to the model weight.

If attention quantization is viable in the future, the lower memory on the
attention tensor will effectively reduce the memory I/O for KV cache by a
factor of 2 in the case of `int8` quantization (compared to `fp16` or `bf16`)
or a factor of 4 in the case of `int4`. Overall, this will flatten the latency
growth with respect to batch size or context length. The overall comparative
complexity (a) with or without bifurcated attention or (b) multi-head vs.
multi-query remains the same.

  6. 6.

Q: Does the conclusion depend on the inference implementation or different
hardware?
A: Different inference platforms, such as FasterTransformers (GPUs) or PaLM
inference (TPUs), can yield different latency numbers. However, the relative
I/O complexity among different attention mechanisms does not change, resulting
in similar relative trends among different attention mechanisms. That being
said, it is possible that more efficient implementations or more performant
chip/system configurations, including different tensor parallelism degrees,
can result in different slopes for the latency growth with respect to context
length and batch size. In that case, the trade-off points in terms of context
length or batch size can be different. The comparative complexity remains the
same based on the analysis.

  7. 7.

Q: How does bifurcated attention differ from using attention mask for sampling
as in done in SpecInfer (Miao et al., 2023) ?
A: The attention mask approach can have a different FLOP usage compared to the
original attention. We can consider a scenario where the attention mask
corresponds to sampling with batch $b$ and incremental decoding length $\ell$,
with the original context of length $m$. The attention FLOPs are
$O(mb\ell+b^{2}\ell^{2})$. In contrast, the original FLOPs is $O(mb\ell)$. If
$b\ell$ is sufficiently large, then the FLOPs via attention mask can be much
higher. However, for the purpose of speculative decoding where the number of
draft tokens is small, this additional FLOPs can be negligible.

##  Appendix B Related Work

###  B.1 Applications of Single-Context Batch Sampling

The observed latency reduction we achieve can have a profound impact on many
applications. Some of these applications include:

  * •

Code Generation: In software development, AI-assisted code generation can
benefit greatly from reduced latency, especially when generating multiple code
snippets or suggestions for a given context. This can lead to a more
responsive and efficient user experience for developers using AI-powered
Integrated Development Environments (IDEs) or code completion tools (Nijkamp
et al., 2023, 2022; Chen et al., 2021; Le et al., 2022; Fried et al., 2022; Li
et al., 2022; Allal et al., 2023; Li et al., 2023; Ahmad et al., 2021).

  * •

Machine Translation: In situations where multiple translations are needed for
a single input, such as generating translations with varying degrees of
formality or generating translations for different dialects, the context-aware
bifurcated attention can provide more efficient computation, resulting in
faster and more scalable machine translation services (Costa-jussà et al.,
2022; Farhad et al., 2021; Tran et al., 2021; Yee et al., 2019).

  * •

Chatbots and Conversational AI: Conversational agents often need to generate
multiple responses to handle different interpretations of a user’s input or to
provide multiple suggestions. The reduced latency offered by the proposed
method can significantly improve the responsiveness of chatbots, leading to a
more natural and fluid conversation with users (Google, 2023).

  * •

Creative Content Generation: In applications like poetry, story, or
advertisement generation, the ability to generate multiple variations for a
given prompt is crucial. The proposed method enables more efficient generation
of diverse content, making it more feasible for real-time or large-scale
applications (Lin and Riedl, 2021; Mirowski et al., 2023; Team, 2023; Yuan et
al., 2022).

  * •

Data Augmentation: In the context of data augmentation for machine learning,
generating multiple alternative examples for a given input can help improve
model robustness and generalization. With the reduced latency provided by
context-aware bifurcated attention, the process of generating augmented data
can be made faster, enabling more efficient use of computational resources
during training.

  * •

General Large Scale Evaluation: In addition to the aforementioned use-cases
there are many niche use-cases where LLM and other open-ended generation
models are explored for toxicity (Dathathri et al., 2019; Gehman et al., 2020;
Nadeem et al., 2020), detection of vulnerable code in generations (Pearce et
al., 2022), performance improving code edit generation (Madaan et al., 2023),
programming language translations (Roziere et al., 2020) and many others. In
all of these scenarios many generations per each prompt are gathered for a
deeper understanding of the models, bifurcated attention can drastically speed
up the generation process in such cases.

In conclusion, the proposed context-aware bifurcated attention method can
significantly reduce memory I/O cost and improve latency in various
applications, leading to increased efficiency and scalability. This method has
the potential to enable new use cases and enhance the user experience in
numerous AI-powered systems, making them more practical for real-world
deployment.

###  B.2 Supporting Long Context Requires IO-Efficient Attention

As language models are becoming general purpose and highly capable, the demand
for language models to handle longer context sequences has grown
significantly. Recently, there is an ongoing focus on models that can handle
even longer context sequences (Bulatov et al., 2023; OpenAI, 2023; Team,
2023). As of today, GPT-4 (OpenAI, 2023) supports context length of 32k
tokens, and MPT-7B (Team, 2023) extends it to 64k while Anthropic’s Claude
333https://www.anthropic.com/index/100k-context-windows supports as long as
100k input length. Most recently, Bulatov et al proposed 1M token input
context length for transformers. These models push the boundaries of context
understanding and generation capabilities, enabling more comprehensive
discourse understanding and contextually informed responses.

This trend is driven by the need for comprehensive discourse understanding in
applications like Retrieval-Augmented Generation (RAG), as well as many
complex prompting methods. Applications such as RAG (Guu et al., 2020; Izacard
et al., 2022; Menick et al., 2022; Zhen et al., 2022) retrieve extensive
passages or documents from external corpora, providing rich and grounded
context for generating responses. Additionally, models like Toolformer (Schick
et al., 2023) and WebGPT (Nakano et al., 2021) leverage external tools, such
as APIs and search engines, to expand the context and enhance generation.

Long context is disproportionately expensive for transformer family models
because for vanilla self-attention both memory and time complexity are
quadratic to the sequence length. To effectively handle longer context
sequences, optimizing memory I/O and reducing computational overhead are
critical. Currently, the dominant approaches to addressing this challenge have
been to make the attention computation less expensive. Beltagy et al. (2020)
proposed to sparsify self-attention using various attention patterns. Wang et
al. (2020) explores low-rank approximation of self-attention. In addition to
the compute bound improvements, advancements in memory-efficient attention
mechanisms and techniques for reducing memory I/O will continue to propel the
field forward, facilitating the handling of longer context sequences in
language models. FlashAttention (Dao et al., 2022) is proposed to speed up
self-attention and reduce the memory footprint without any approximation. It
leverages fused kernel for matrix multiplication and softmax operation which
greatly reduces memory IO during training.

##  Appendix C Setup

###  C.1 Model Training Details

We trained multiple models with varying sizes, ranging from 125 million
parameters to 13 billion parameters, using code data with a context size of
2048 and adjusting the per-GPU batch size and total number of steps according
to the model size. For model training we used multiple p4 instances each
equipped with 8 40GB Nvidia A100 GPUs per instance.

For our largest model family, the 13 billion parameter model, we used a global
batch size of 1024, which approximately translates to 2 million tokens per
batch. The settings for each model within each model-size family were kept
consistent. The remaining training hyperparameters are summarized in the
following table 1.

We use AdamW optimizer (Kingma and Ba, 2014) with $\beta_{1}=0.9$,
$\beta_{2}=0.95$, and $\epsilon=10^{-8}$. The warm-up steps were set to 2000,
and a cosine annealing learning rate schedule was employed after reaching the
peak learning rate. The minimum learning rate was set to 10% of the
corresponding peak learning rate. A weight decay (Loshchilov and Hutter, 2017)
of 0.01 and gradient clipping of 1.0 were applied to enhance training
stability. Following the approach in (Shoeybi et al., 2019), the standard
deviation for random weight initialization was rescaled for larger models. Our
training pipeline is based on PyTorch Lightning and we use bfloat16 (Kalamkar
et al., 2019) and DeepSpeed (Rasley et al., 2020) for training optimization.
Finally, a random split of 0.1% of the data was reserved as a validation set.

Table 1: Training Hyperparameters Model Size | Total Training Steps | Batch Size | Compute Nodes | Max Learning Rate
---|---|---|---|---
125M | 400k | 256 | 8 | $2.5\times 10^{-4}$
672M | 200k | 256 | 8 | $2.5\times 10^{-4}$
2.8B | 200k | 512 | 16 | $1.6\times 10^{-4}$
13B | 100k | 1024 | 32 | $1.0\times 10^{-4}$

###  C.2 Model Configurations

For each model size we train three models with attention variations; multi
head where $g=h$, multi group where $1<g<h$ and multi query where $g=1$.
Additionally, for 672m and 2.8b models we train a multi group model variant
where the fanout in feed forward layer is decreased from $4\times d$ to
$2\times d$. Each model variant yields different number of total parameters
therefore we group these models into family of model sizes. The detailed
architectural choices for each of the model family is found in the table 2.

Table 2: Model Specifications table presenting architecture details for the three variants: multi head (MH), multi query (MQ), and multi group (MG) including parameter count, number of attention groups, head dimensions, and number of layers. The additional fanout-based MG variant is described here as MG + $2\times d$ Model Family |  Attention Type |  $groups$ |  $d_{head}$ |  $n_{layer}$ |  $N_{params}$ (billions)
---|---|---|---|---|---
125M |  MH |  12 |  64 |  12 |  0.125
MG |  4 |  64 |  12 |  0.115
|  MQ |  1 |  64 |  12 |  0.112
672M |  MH |  20 |  72 |  24 |  0.672
MG |  4 |  72 |  24 |  0.592
|  MG + $2\times d$ |  4 |  72 |  24 |  0.393
|  MQ |  1 |  72 |  24 |  0.578
2.8B |  MH |  24 |  128 |  24 |  2.878
MG |  4 |  128 |  24 |  2.501
|  MG + $2\times d$ |  4 |  128 |  24 |  1.595
|  MQ |  1 |  128 |  24 |  2.444
13B |  MH |  40 |  128 |  40 |  12.852
MG |  8 |  128 |  40 |  11.174
|  MQ |  1 |  128 |  40 |  10.807

###  C.3 Model Details of 1B Latency Experiment

In Section 5.2.2, we use candidate models of sizes roughly 1B to study the
effect of bifurcated attention. We outline the hyperparameters of such models
below.

Table 3:  Model Specifications for Latency Experiment in Section 5.2.2.  Model Family |  Attention Type |  $groups$ |  $d_{head}$ |  $n_{layer}$ |  $N_{params}$ (billions)
---|---|---|---|---|---
1B |  MH |  20 |  128 |  12 |  1.077
MG |  4 |  128 |  15 |  1.156
|  MQ |  1 |  128 |  16 |  1.193

###  C.4 Ablation Studies: $2d$ Intermediate Feature Dimension

One can also argue that different $g$ results in different balance of the
number of parameters in the feedforward versus the attention components. We
performed an ablation study where we reduce the typical intermediate feature
size of $4d$ to $2d$ and train models for three model sizes (which we will
refer to as the $2d$ experiment). The ablation study reveals that the scaling
laws curves for the $2d$ experiment crosses the usual $4d$ curves, which
implies that the reduced size of the attention component alone compared to
feedforward does not provide a consistent explanation of model capabilities.
This can be seen from Figure 9.

Figure 9:  Capabilities versus size plots including the $2d$-intermediate-size
feedforward model. The plot shows that the balance between the number of
feedforward parameters and the attention parameters alone does not explain the
relative expressiveness of multi-head, multi-group, and multi-query
attentions. Rather, we argue that what explains relative expressiveness is the
representation power associated with the key and value tensors (Section 5.1).

###  C.5 Inference Setup

We use Nvidia A100 GPUs for inference hardware (Choquette et al., 2021). We
perform latency studies using Deepspeed inference (Rasley et al., 2020) on top
of Huggingface transformers (Wolf et al., 2019), where we wrote custom code to
handle the generalize multi-group attention as well as bifurcated attention.
Future work includes extending the implementation to FasterTransformer
(NVIDIA, ).

##  Appendix D Multi-Group Attention Family

###  D.1 Detailed Analysis on Memory Access

We show in Table 4 that the memory IO cost for $\langle q,K\rangle$ is
dominated by the loading of $K$ which costs $bmhk$ in the case of multi-head
where $g=h$. This cost is particularly high due to the coupling of batch size
$b$, context length $m$, and the entire hidden dimension $d$. Compared to the
number of computations, which has complexity $bmd$, this attention module
requires one memory IO per one tensor operation (memory-io bound). In
contrast, other operations such as feedforward has much lower ratio of memory
IO per compute (compute bound). These attention computation can be the main
bottleneck for incremental decoding and our paper aims to tackle such
problems.

Concretely, we can see that the context encoding in single-batch scenario in
Appendix 5.2.1 is $400$ ms for context length $10000$, implying that the
amortized latency per token during this phase is $0.04$ ms per token. However,
the per token latency during incremental decoding is in the order of $\approx
10$ ms per token, which is $\frac{10}{0.04}=250$ times slower. This number
clearly demonstrates that compute is not a dominating factor, but the memory
IO required to load both model and KV cache.

Table 4: Comparison of memory access and computation between Multi Head, Multi Query, and Multi Group attention mechanisms. The memory access is for incremental decoding with the query length $n=1$.  Operation | Einsum | Memory Access | Computation
---|---|---|---
Input ($x$): $bd$ |  |  |
$q=\langle x,P_{q}\rangle$ | $bd,hdk\rightarrow bhk$ | $bd+hdk=bd+d^{2}$ | $bdhk=bd^{2}$
$K=\langle x,P_{k}\rangle\ (+K_{prev})$ |  [MH] $bd,hdk\rightarrow bhk\ (+bmhk)$ | $bd+d^{2}$ | $bdhk=bd^{2}$
|  [MQ] $bd,dk\rightarrow bk\ (+bmk)$ | $bd+dk$ |
|  [MG] $bd,gdk\rightarrow bgk\ (+bgmk)$ | $bd+gdk$ |
$V=\langle x,P_{v}\rangle\ (+V_{prev})$ |  [MH] $bd,hdv\rightarrow bhv\ (+bmhv)$ | $bd+d^{2}$ | $bdhv=bd^{2}$
|  [MQ] $bd,dv\rightarrow bv\ (+bmv)$ | $bd+dv$ |
|  [MG] $bd,gdv\rightarrow bgv\ (+bgmv)$ | $bd+gdv$ |
logits $=\langle q,K\rangle$ |  [MH] $bhk,bhmk\rightarrow bhm$ | $bhk+bhmk=bd+bmd$ | $bhmk=bmd$
|  [MQ] $bhk,bmk\rightarrow bhm$ | $bd+bmk+bhm$ |
|  [MG] $bhk,bgmk\rightarrow bhm$ | $bhk+bgmk+bhm$ |
weights: softmax |  | $bhm$ | $bhm$
out($O$) $=\langle$ weights, $V\rangle$ |  [MH] $bhm,bhmv\rightarrow bhv$ | $bhm+bhmv=bhm+bmd$ | $bhmv=d$
|  [MQ] $bhm,bmv\rightarrow bhv$ | $bhm+bmv+bhv$ |
|  [MG] $bhm,bgmv\rightarrow bhv$ | $bhm+bgmv+bhv$ |
$y=\langle O,P_{O}\rangle$ | $bhv,hdv\rightarrow bd$ | $bd+d^{2}$ | $bdhv=bd^{2}$
Total: Multi Head |  | $bd+bmd+d^{2}$ | $bhm+bmd+bd^{2}\approx bd^{2}$
Total: Multi Query |  | $bd+bmk+d^{2}$ |
Total: Multi Group |  | $bd+bgmk+d^{2}$ |
$r$: Multi Head |  | $1/d+m/d+1/b$ |
$r$: Multi Query |  | $1/d+m/(dh)+1/b$ |
$r$: Multi Group |  | $1/d+g/(dh)+1/b$ |

###  D.2 Model FLOPs

The scaling laws by Kaplan et al. (2020) shows that the model-related FLOPs
during the forward pass is $2N$ where $N$ is the number of parameters (without
the embeddings). We show that it holds for a general multi-group model as
well. The only difference between the multi-group and the multi-head case is
the projection $P_{K}$ and $P_{V}$ where they are of size $dgk$ instead of
$dhk$. Since this is a linear layer, the forward pass FLOPs for any input is
still proportional such projection size. Therefore, it follows that for any
multi-group attention, including multi-head, the forward FLOPs is $2N$ where
$N$ is the respective number of parameters.

###  D.3 Comparing Capabilities-Equivalent Models

This section outlines the analysis of latency change when we switch from an MH
model to an MG model with $F$ times the size.

####  D.3.1 Context encoding

The dominating factor for latency in context encoding is the compute rather
than the memory IO. The compute can be broken down into two parts (a) tensor
projections related to model parameters and (b) KV attention involving no
model parameters. For both parts, the large multi-group model will involve
higher latency proportional to the size factor $F$. The context encoding time
is $\propto N\times bm$ where $N$ is the model size except embeddings for (a)
since the FLOPs per token is $2N$ (Kaplan et al., 2020), which holds for all
multi-group attention (Appendix D.2). For (b), the encoding time is
$\propto\ell\cdot bhm^{2}\propto Nbm^{2}$ for (b). Overall, the multi-group
model with similar capabilities as the multi-head model will incur slightly
higher context encoding time due to the larger size since $N$ to increase to
$FN$.

####  D.3.2 Incremental Decoding

The incremental decoding component can dominate the overall inference latency
compared to the context encoding, especially in the scenario where we decode
in many steps. Incremental decoding is memory bound, meaning that the latency
of this step is limited by memory I/O throughput. We can divide the memory I/O
usage into two parts: reading (a) model parameters and (b) cached key-value
pairs. With multi-group, we expect the model parameters to increase by a
factor of $F(g)$, leading to an increase in I/O usage in (a) by the same
factor. The memory IO in (b) changes by a factor of $\frac{g}{h}$ when moving
from multi-head with KV cache size $2bhmk$ to multi-group with cache size
$2bgmk$ (more precisely $\frac{g}{h}\cdot F(g)$ but $\frac{g}{h}$ is a
dominating term since $F(g)$ is close to $1$).

##  Appendix E Context-Aware Bifurcated Attention

###  E.1 Proof

Here, we outline the proof that the proposed bifurcated attention in Equation
3 and 4 recovers the same attention as the operations in 1 and 2 for the case
of single-context batch sampling. We use the fact that the KV part
corresponding to context length, all the batch indices correspond to the
tensors.

$\displaystyle\langle q,K\rangle$ $\displaystyle:\text{einsum}(bgpnk,bgmk)\to
bgpnm$ $\displaystyle=\text{einsum}(bgpnk,bg(m_{c}+m_{d})k)\to bgpnm$
$\displaystyle=\text{einsum}(bgpnk,bgm_{c}k)\to bgpnm\oplus$ $\displaystyle\ \
\ \ \ \text{einsum}(bgpnk,bgm_{d}k)\to bgpnm$
$\displaystyle=\text{einsum}(bgpnk,bgm_{c}k)\to bgpnm\oplus$ $\displaystyle\ \
\ \ \ \text{einsum}(bgpnk,gm_{d}k)\to bgpnm$ $\displaystyle=\langle
q,K_{c}\rangle\oplus\langle q,K_{d}\rangle$ $\displaystyle\langle w,V\rangle$
$\displaystyle:\text{einsum}(bgpnm,bgmk)\to bgpnk=bnd$
$\displaystyle=\text{einsum}(bgpnm_{c},bgm_{c}k)\to bgpnk+$ $\displaystyle\ \
\ \ \ \text{einsum}(bgpnm_{d},bgm_{d}k)\to bgpnk$
$\displaystyle=\text{einsum}(bgpnm_{c},gm_{c}k)\to bgpnk+$ $\displaystyle\ \ \
\ \ \text{einsum}(bgpnm_{d},bgm_{d}k)\to bgpnk$ $\displaystyle=\langle
w_{c},V_{c}\rangle+\langle w_{d},V_{d}\rangle$

###  E.2 Detailed Memory I/O Analysis

Overall, the memory I/O complexity changes from

  * •

Original memory I/O cost: $bhnk+bgmk+bhnm$ (for $\langle q,K\rangle$) +
$bhnm+bgmk+bnd$ (for $\langle w,V\rangle$)

  * •

Bifurcated attention memory I/O cost: $bhnk+gm_{c}k+bgm_{d}k+bhnm$ (for
$\langle q,K\rangle$) + $bhnm+gm_{c}k+bgm_{d}k+bnd$ (for $\langle w,V\rangle$)

There is an associated memory IO to write the $\langle w,V_{c}\rangle$ and
$\langle w,V_{d}\rangle$ output twice. However, it is typically very small
($bnd$) compared to the IO of KV cache component $bgmk$ since $m>>n=1$.

###  E.3 Implementation of Bifurcated Attention

Despite the dramatic gain in inference efficiency of the bifurcated attention,
we demonstrate the simplicity of our implementation involving 20 lines of code
using Pytorch (Paszke et al., 2019).

⬇

1def attn(query, key, value, bifurcated_attention):

2 # <q,K>

3 if bifurcated_attention and type(key) == dict:

4 # g = number of groups

5 # h = gp where p = num heads per group

6 # n = 1 for incremental decoding

7 attn_weights_context = torch.einsum(

8 "bgpnk,gmk->bgpnm", query, key["context_past_key"][0]

9 )

10 attn_weights_incremental = torch.einsum(

11 "bgpnk,bgmk->bgpnm", query, key["incremental_past_key"]

12 )

13 attn_weights = torch.cat(

14 [attn_weights_context, attn_weights_incremental], dim=-1

15 )

16 else:

17 attn_weights = torch.einsum(

18 "bgpnk,bgmk->bgpnm", query, key

19 )

20 # softmax and causal mask (omitted)

21 # <w,V>

22 if bifurcated_attention and type(value) == dict:

23 # n = 1 for incremental decoding

24 context_past_value_length = value["context_past_value"].size(-2)

25 attn_output_context = torch.einsum(

26 "bgpnm,gmv->bgpnv",

27 attn_weights[:, :, :, :, :context_past_value_length],

28 value["context_past_value"][0],

29 )

30 attn_output_incremental = torch.einsum(

31 "bgpnm,bgmv->bgpnv",

32 attn_weights[:, :, :, :, context_past_value_length:],

33 value["incremental_past_value"],

34 )

35 attn_output = attn_output_context + attn_output_incremental

36 else:

37 attn_output = torch.einsum(

38 "bgpnm,bgmv->bgpnv", attn_weights, value

39 )

40 return attn_output

(a) MBXP Java

(b) MBXP JavaScript

Figure 10: Bifurcated attention enables high batch sampling with minimal
increase in latency with respect to batch size, resulting in more diverse and
improved quality samples as indicated by pass@n and pass@top3 on MBXP-Java
(Figure 10(a)). This trend extends to other evaluation such as JavaScript
(Figure 10(b)) and Python (Figure 8).

##  Appendix F Applications: Additional Results

We demonstrate additional results to the evaluation in Section 5.3 on MBXP-
Java and MBXP-Javascript, in addition to the Python results. We replace
CodeGen-16B-mono with CodeGen-16B-multi for the evaluation on Java and
JavaScript and use the same StarCoder model. From Figure 10, we observe
similar trends as in Python (Figure 8), which furthers demonstrates the wide
applicability of of bifurcated attention in improving accuracy under latency-
constrained scenarios.

##  Appendix G Compatibility with Speculative Decoding and Fast Decoding
techniques

Unlike standard auto-regressive decoding, fast decoding techniques such as
Speculative decoding(Chen et al., 2023; Leviathan et al., 2022), Medusa (Cai
et al., 2024), Lookahead (Fu et al., 2023), and Eagle (Li et al., 2024)
attempt to decode multiple tokens at each step. This reduces I/O bandwidth
requirements because model parameters and KV cache are fetched only once per
step and can be amortized across all generated tokens.

The fundamental principle behind these techniques is to first draft (or guess)
a set of tokens (denoted as $n_{g}$) and then validate their accuracy by
parallelly decoding with the model. After each step, up to a tokens (where
$a\leq n_{g}$) may be accepted as valid, allowing for memory usage
amortization across these accepted tokens. This approach is successful because
decoding is primarily constrained by memory I/O.

The benefits of bifurcated attention are orthogonal to those of speculative
sampling, leading to further memory I/O improvements. This can be observed by
extrapolating per-step memory I/O costs from Section E.2 with $n_{g}$
raplacing $n$. Since $m>>n_{g}$ continues to hold, the advantages of
bifurcated attention persist even when combined with speculative decoding.


ILD .eps.gz .eps.gzpdf.pdfzcat #1 | epstopdf –filter –outfile=

###### Abstract

Although the LHC experiments have searched for and excluded many proposed new
particles up to masses close to 1 TeV, there are many scenarios that are
difficult to address at a hadron collider. This talk will review a number of
these scenarios and present the expectations for searches at an electron-
positron collider such as the International Linear Collider. The cases
discussed include the light Higgsino, the stau lepton in the coannihilation
region relevant to dark matter, and heavy vector bosons coupling to the
s-channel in e+e- annihilation. The studies are based on the ILD concept at
the ILC.

##  1 ILC and BSM

The ILC - the International Linear Collider - is a power-efficient
$e^{+}e^{-}$ collider with initial $E_{CMS}$ = 250 GeV, upgradable up to 1 TeV
[1, 2]. As electrons are point-like objects, the initial state is known. Since
the production mechanism is electro-weak, the backgrounds will be low meaning
that the detectors can be thin and have close to $\sim 4\pi$ coverage, and can
be operated without trigger. In addition, at ILC, both beams will be
polarised. This combination of polarised beams, low background, known in-
state, hermetic detectors and energy upgradability makes the ILC the ideal
environment for Beyond the Standard Model (BSM) searches [3].

###  1.1 BSM at ILC: the SUSY case

SUSY is the most complete theory of BSM. It also serves as a boiler-plate for
BSM, since almost any new topology can be obtained in some flavour of SUSY. In
addition it is the most studied model with serious simulation: In most cases
presented here, used full simulation of the International Large Detector
concept at ILC (the ILD [4]), with all SM backgrounds, and all beam-induced
backgrounds included.

Although the LHC experiments have searched for and excluded many proposed new
particles up to masses in the 1 TeV range, there are many scenarios that are
difficult to address at a hadron collider [5, 6]. In SUSY, all is known for
given masses, due to SUSY-principle, relating sparticle properties to those of
the corresponding SM particles. These relations do not depend on the SUSY
breaking mechanism. Obviously, there is one Next-to-lightest SUSY particle
(NLSP), and it must have 100 % BR to it’s (on- or off-shell) SM-partner and
the lightest SUSY particle (the LSP). Therefore, one can perform model
independent evaluations of exclusion or discovery reach in $M_{NLSP}-M_{LSP}$
plane - with no loop-holes. One can do this for any possible NLSP, and
concentrate on the "worst" cases - the ones with lowest cross-section and most
difficult signatures. These are on one hand the bosinos, and on the other the
$\tilde{\tau}$. The reaches for these cases are shown in Fig. 1.

(a)

(b)

Figure 1: The expected exclusion and discovery reaches at ILC for bosinos (a)
[7] and $\tilde{\tau}$:s (b)[8], together with obtained or projected exclusion
limits from LHC/HL-LHC and LEPII (references can be found in [7]). NB: At ILC,
exclusion and discovery is almost the same, even for difficult channels.

Since the discovery and exclusion reaches at the ILC are quite similar, after
a possible discovery, one will quickly enter the realm of precision
measurements.

In Fig. 2 typical signals of bosino production are shown, for charginos (Fig.
2(a)) and neutralinos (Fig. 2(c)). Both models are higgsino-LSP ones. The one
to the left is one of three studied natural SUSY models with moderate mass
differences (15-20 GeV) [9], while the one to the right is a cosmology-
motivated model, and has a sub-GeV difference [10]. In the natural SUSY
analysis, the combination of the measured masses, BR’s and Higgs properties,
all 10 weak-scale parameters gets constrained, for all three bench-marks. In
particular, the bino and wino SUSY breaking masses $M_{1}$ and $M_{2}$ \- the
ones most directly related to the higgsino masses - can be determined at
percent level. The fitted weak-scale parameters can be evolved with the
appropriate RGE’s to higher scales. This allows to verify or discard the idea
of GUT-scale unification of $M_{1}$ and $M_{2}$, see Fig 2(b).

(a) Natural SUSY

(b) Extrapolation to GUT

(c) cosmology-motivated model

Figure 2: ILC bosino measurements: (a) chargino signal in a natural SUSY
model; (b) RGE extrapolation of the gaugino mass-parameters to the GUT scale,
with experimental inputs obtained in natural SUSY; (c) neutralino signal in
the cosmology-motivated model. The SUSY background was included in the
analysis, but was found to completely removed after all cuts were applied.

In Fig. 3, typical slepton signals are illustrated. The end-points of the
spectra of the slepton decay-products, enables to measure the slepton and LSP
masses to percent level accuracy. Fig. 3(a) and (b) are examples of
$\tilde{\tau}_{1}$ and $\tilde{\mu}_{\mathrm{R}}$ signals, in a $\tilde{\tau}$
co-annihilation model [6]. Fig. 3(c) shows the signal that could be obtained
from evaluating the decay-kinematics in a model with cascade decays of
$\tilde{\mu}$’s [11].

(a) $\tilde{\tau}_{1}$

(b) $\tilde{\mu}_{R}$

(c) $\tilde{\mu}_{R}$ from cascades

Figure 3: Slepton signals at ILC: (a) $\tau$ decay-product spectrum in
$\tilde{\tau}_{1}$ pair-production; (b) $\mu$ spectra in $\tilde{\mu}_{R}$
pair-production; (c) reconstructed $\tilde{\mu}_{R}$ mass in a model where
$\tilde{\chi}^{0}_{2}$ decays to $\mu\tilde{\mu}_{R}$. Open: signal, red: SM
background, blue: SUSY background.

After the full ILC program, and depending on model, channel, and polarisation,
we find experimentally that measured $\delta(\mathrm{masses})$ = 0.5-1 % ,
$\delta(\sigma\times\mathrm{BR})$ = 1-6%

###  1.2 BSM at ILC: not only SUSY.

Dark matter can be searched for at ILC in
$\mbox{${\,\mathrm{e}}^{+}{\mathrm{e}}^{-}\to$}(DM)(DM)+\mathrm{ISR}\gamma$,
i.e. in Mono-photon searches. Results of such searches are shown in Fig. 4(a)
for heavy mediators where a model independent EFT approach is appropriate
[12]. Also arbitrary mediators have been studied and also show potential
beyond HL-LHC [13]

(a) Mono-photon - heavy mediators

(b) New scalar

(c) Dark photon

Figure 4: Examples of non-SUSY BSM searches at ILC

New Higgs-like scalar (S), produced in
$\mbox{${\,\mathrm{e}}^{+}{\mathrm{e}}^{-}\to$}Z^{*}\rightarrow ZS$ with
unknown decays of S, can be searched for it in a decay-mode insensitive way at
ILC, using the recoil-mass, i.e. the mass of the system recoiling against the
measured $Z$ [14]. Couplings down to a few percent of the SM-Higgs equivalent
can be excluded, as shown in Fig. 4(b).

At ILC one can also search for Dark photon/Z’. Generically, the kinetic mixing
term $\frac{\epsilon}{2\cos{\theta_{W}}}F^{\prime}_{\mu\nu}B^{\mu\nu}$ in the
Lagrangian leads to a tiny, narrow resonance, but still wide enough to make
decays prompt. One can search for this as a $\mu\mu$ resonance above
background in
$\mbox{${\,\mathrm{e}}^{+}{\mathrm{e}}^{-}\to$}Z^{\prime}+\mathrm{ISR}\rightarrow\mu^{+}\mu^{-}+\mathrm{ISR}$.
Results (from EPPSU) shown in Fig. 4(c).

Another study that ILC allows for is indirect BSM. Such studies do not only
show an important discovery potential, but also can achieve model separation.
In [15] the results of an SM effective field theory (SMEFT) study, using ILC
results on Higgs properties and TGCs are given There, one selected models that
are not discoverable at HL-LHC, and it was shown that, at ILC, not only are
the models separatable at 5 $\sigma$ from the SM, but also from each other.

##  2 Conclusion

Sometimes, the capabilities for the direct discovery of new particles at the
ILC exceed those of the LHC, since ILC provides a well-defined initial state,
a clean environment without QCD backgrounds, extendability in energy and
polarised beams. Thanks to the low background levels, detectors - like ILD -
can be factors more precise than their LHC counterparts, can be hermetic, and
do not need to be triggered.

Many ILC - LHC synergies from energy-reach vs. sensitivity are expected. E.g.
if SUSY is reachable at ILC, it means 5 $\sigma$ discovery, and precision
measurements. This input might be just what is needed for LHC to transform a 3
$\sigma$ excess to a discovery of states beyond the reach of ILC. Similar
synergies can also be expected in many other models for BSM physics.

##  3 Acknowledgements

We would like to thank the LCC generator working group and the ILD software
working group for providing the simulation and reconstruction tools and
producing the Monte Carlo samples used in this study. This work has benefited
from computing services provided by the ILC Virtual Organisation, supported by
the national resource providers of the EGI Federation and the Open Science
GRID.

## References

  * [1] Philip Bambade  “The International Linear Collider: A Global Project”, 2019  arXiv:1903.01629 [hep-ex]
  * [2] “The International Linear Collider Technical Design Report - Volume 1: Executive Summary”, 2013  arXiv:1306.6327 [physics.acc-ph]
  * [3] Keisuke Fujii  “The Potential of the ILC for Discovering New Particles”, 2017  arXiv:1702.05333 [hep-ph]
  * [4] Halina Abramowicz  “International Large Detector: Interim Design Report”, 2020  arXiv:2003.01116 [physics.ins-det]
  * [5] Mikael Berggren  “What pp SUSY limits mean for future e+e- colliders”  In _International Workshop on Future Linear Colliders_ , 2020  arXiv:2003.12391 [hep-ph]
  * [6] Mikael Berggren et al.  “Non-simplified SUSY: $\widetilde{\tau}$ -coannihilation at LHC and ILC”  In _Eur. Phys. J. C_ 76.4, 2016, pp. 183  DOI: 10.1140/epjc/s10052-016-3914-2
  * [7] Ichiro Adachi  “The International Linear Collider: Report to Snowmass 2021”, 2022  arXiv:2203.07622 [physics.acc-ph]
  * [8] María Teresa Núñez Pardo de Vera, Mikael Berggren and Jenny List  “Evaluating the ILC SUSY reach in the most challenging scenario: $\widetilde{\tau}$ NLSP, low $\Delta M$ , lowest cross-section”  In _2022 Snowmass Summer Study_ , 2022  arXiv:2203.15729 [hep-ph]
  * [9] Howard Baer et al.  “ILC as a natural SUSY discovery machine and precision microscope: From light Higgsinos to tests of unification”  In _Phys. Rev. D_ 101.9, 2020, pp. 095026  DOI: 10.1103/PhysRevD.101.095026
  * [10] Mikael Berggren et al.  “Tackling light higgsinos at the ILC”  In _Eur. Phys. J. C_ 73.12, 2013, pp. 2660  DOI: 10.1140/epjc/s10052-013-2660-y
  * [11] Mikael Berggren  “Reconstructing sleptons in cascade-decays at the linear collider”  In _International Conference on Linear Colliders (LCWS 04)_ , 2005, pp. 907–910  arXiv:hep-ph/0508247
  * [12] Moritz Habermehl, Mikael Berggren and Jenny List  “WIMP Dark Matter at the International Linear Collider”  In _Phys. Rev. D_ 101.7, 2020, pp. 075053  DOI: 10.1103/PhysRevD.101.075053
  * [13] Jan Kalinowski et al.  “Sensitivity of future linear $\hbox{e}^{+}\hbox{e}^{-}$ colliders to processes of dark matter production with light mediator exchange”  In _Eur. Phys. J. C_ 81.10, 2021, pp. 955  DOI: 10.1140/epjc/s10052-021-09758-6
  * [14] Yan Wang, Mikael Berggren and Jenny List  “ILD Benchmark: Search for Extra Scalars Produced in Association with a $Z$ boson at $\sqrt{s}=500$ GeV”, 2020  arXiv:2005.06265 [hep-ex]
  * [15] Tim Barklow et al.  “Improved Formalism for Precision Higgs Coupling Fits”  In _Phys. Rev. D_ 97.5, 2018, pp. 053003  DOI: 10.1103/PhysRevD.97.053003