Attention-Level Speculation
Authors: Jack Cai
Ammar Vora
Randolph Zhang
Mark O'Connor
Mark C. Jeffrey
Venue: ICML 2025 poster (2025)
Keywords: Inference
Speculation
Attention
Transformer
Large Language Model
Primary area: deep_learning->attention_mechanisms
TL;DR: We show that attention-level speculation reduces LLM decode latency when tensor parallelism fails to scale.

Abstract: As Large Language Models (LLMs) grow in size and context length, efficient inference strategies are essential to maintain low-latency token generation. Unfortunately, conventional tensor and data parallelism face diminishing returns when scaling across multiple devices. We propose a novel form—attention-level speculative parallelism (ALSpec)—that predicts self-attention outputs to execute subsequent operations early on separate devices. Our approach overlaps attention and non-attention computations, reducing the attention latency overhead at 128K context length by up to 5x and improving end-to-end decode latency by up to 1.65x, all without sacrificing quality. We establish the fundamental pillars for speculative execution and provide an execution paradigm that simplifies implementation. We show that existing attention-approximation methods perform well on simple information retrieval tasks, but they fail in advanced reasoning and math. Combined with speculative execution, we can approximate up to 90% of self-attention without harming model correctness. Demonstrated on Tenstorrent's NPU devices, we scale up LLM inference beyond current techniques, paving the way for faster inference in transformer models.

Attention-Level Speculation
Jack Cai Ammar Vora 12
Randolph Zhang²
Mark O'Connor Mark C. Jeffrey'
Abstract
As model sizes and context lengths increase (Chitty-Venkate et al., 2024; Shocybi et al., 2019; Huang et al., 2019), sin- gle accelerator can no longer meet the memory and compute demands. Tensor parallelism (splitting the model weights across devices) and data parallelism (replicating the model weights on each device) have been widely used for large- scale LLM inference. Despite their adoption, these tech- niques exhibit the following problems during the causal decode phase, hurting coken-generation latency:
As Large Language Models (LLMs) grow size and context length, efficient inference strategies are essential to maintain low-latency token gen- eration. Unfortunately, conventional tensor and data parallelism face diminishing returns when scaling across multiple devices. We propose a novel form-attention attention-level speculative paral- lelism (AI-Spec)- that predicts self-attention out- puts to execute subsequent operations early on separate devices. Our approach overlaps atten- tion and non-attention computations, reducing the attention latency overhead at 128K context length by up to 5x and improving end-to-end de- code latency by up to 1.65x, all without sacrific- ing quality. We establish the fundamental pillars for speculative execution and provide an execu- tion paradigm that simplifies implementation. We show that existing iitentionnapproxinatiin meth- ods perform well on simple information retrieval tasks. but they fail in advanced reasoning and math. Combined with speculative execution, we can approximate up to 90% of self-attention with- out harming model correctness. Demonstrated on Tenstorrent's 's NPU devices.' we scale up LLM inference beyond current techniques, paving the way for faster inference in transformer models.
1. Tensor parallelism suffers from communication overhead. Latency for operations like all- all-gather grows with more devices, yielding diminishing returns with device scaling.
2. Data parallelism improves throughput with large batch sizes but does not improve latency.
3. As the key-value (KV) caches grow with context length, self-attention becomes bottleneck and increases overall latency (see Section 5.1 and Appendix D for details).
We observe in Figure I(left) that during token generation at small context length, the majority of latency comes from the non-self-atteenion operations, particularly the feed-forward networks. The self-attention latency grows linearly with context length. In contrast, the non-self-attention latency stays constant. This motivates our question: Can overlap attention with other operations to hide the increasing cost without compromising model correctness?
We propose attention-level speculation (ALSpec), replac- ing some full self- eeff-attention with cheaper approximations (e.g., windowed attention with attention sink (Xiao et al., 2024)) and verifying the approximation against the true at- tention in parallel. When the approximation is accurate, AL- Spec commits and skips expensive calculations. This form of speculation, when successful, allows non-self-attention operations to overlap in time with the self-attention oper- ations. We find that up to 90% of the self-attention op- erations can be safely speculated without harming output quality. In contrast decode-level speculation (Leviathan et al., 2023), ALSpec can be applied to emerging model architectures with recurrent hidden states, such as layer- looping (Eyuboglu et al., 2024) and continuous chain-of- thoughts (Hao et al., 2024). When computing resources are abundant, ALSpec could be coupled with tensor parallelism and decode-level speculation deliver even lower latency. Figure I shows that ALSpec delivers up to 5xx reduction in
1. Introduction
State-of-the-art Large Language Models (LLMs) for token generation often rely on decoder-only transformer archi- tectures (Touvron et al., 2023; Dubey et al., 2024; Jiang et al., 2024; Yang et al., 2024: Almazrouei et al., 2023). These models stack repetitive layers of self-attention and feed-forward operations, allowing efficient scaling across layers (Tay et al., 2022; Kim et al., 2023).
Unnivessity of Toronto 'Tenstorrent 2 University of Waterloo. Correspondence to: Mark C. Jeffrey <mej @ ece. uttooono..c.>..
Proceedings of the nnnd International Conference on Machine Learning, Vancouver, Canada. PMLR 267, 2025. Copyright 2025 by the author(s). Our code is at githuh.com/mciigrruu/llats

---

Attention-Level Speculation
Attention-Lvvel Specclation Overrape Attention
Tessor Parallee Decader (Tenser Porallefim)
Dats Parallel Decoder (Data Paralilliam)
Appracimated Aftention Parallal Decoder
Attenfiss-avall Specclation
(Speculative Paraliofam)
Litatiy Beallie
-aaacy Speculaziv (Dun)
Iond Scecrto
Device
A/e
New
Device
Device
Device
Devico
30
Lawnnng For Tiken vins)
Latency
neuction
Wiest Coee Seceail
20
stardon
trch
Orerap
Device
Are
Arts
Device 2
Art
I
Device
Device
Device
Arb
otenticn from
Notroorraggaiis
Opy
Atrenio Oun (ug. Feod Forword Inyers)
Approsinoee
10
Lusoeyy ofnce-ataatio opi,
low Siegl Imput
when spenu otion ia correct V
Conmuniaaiion
Carmoliable Outour Enter
25k
50k
75k
100k
125k
To Minimize Latoneye Use avory devices os possiile for lenser paraltal unti divinidinng retura, then use recrokinng devices for attendionioral spocolation
CantestLaggh
Figure 1. (Left) Our key result: Llama 3.1 8B stcode-stage latency across context lengths devices, comparing the 8-device full tensor-parallel baseline with our 4-device tensor parallel plus 4-device speculative parallel at 87.5% speculation hit rate (see Section 5.1). (Right) Comparing tensor parallel (TP), data parallel (DP), attention approximation, and ALSpec when executing 2 independent inputs on 2 devices (inputs are tensor representations of previous output tokens). TP reduces per-input latency but suffers communication overheads. DP gives highest throughput but also high per-input latency. Approximation suffer from oulput errors, despite latency improvements. ALSpec keeps correctness and reduces latency relative to TP when speculations are correct, with lateney no worse than DP.
attention latency overhead for long-context decode at 128K. also contrasts ALSpec with tensor-pate tensor-parallel, and approximated attention, which we detail in Section 2.
to an increasingly large key-value (KV) cache. Various parallelization strategics use more computing resources for lower latency or higher throughput, including tensor par- allelism (Shoeybi al., 2019), data parallelism (Huang et al., 2019), and pipeline parallelism (DeepSeek-Al et al., 2024). Figure (right)) qualitatively compares these methods when running two inputs on two devices. Tensor parallelism splits the model weights across multiple devices to acceler- ate attention and forward layers. Although it reduces latency for a single input, it suffers from communication overhead (e.g., all all-gather of partial results), which grows as the number of devices increases. This yields diminishing returns when scaling past handful of devices. Data paral- lelism replicates the model weights on all devices and pro- cesses multiple inputs (or tokens) simultancously for higher throughput. However, with interactive LLM inference, the latency per input is paramount; data parallelism provides minimal improvement because each input runs on one de- vice. Pipeline parallelism (not shown) partitions the model layers so that each device is responsible for one portion of the forward pass. Once the pipeline fills, it can achieve the same overall throughput as data parallelism. However, each input still traverses all pipeline stages, yielding a similar latency profile to data parallelism. Hence, pipeline paral- lelism, like data parallelism, does not fundamentally reduce inter-token decode latency.
parallel,
We further contribute speculative flash decode kernel that computes approximate and exact attention in one pass, adding little to no overhead for speculation. Our new execu- tion paradigm, static graph, dynamic concurrency (SGDC) with priority gating, maintains an op-by-op host view, while running approximate and exact attention paths concurrently. This mechanism delivers up to 1.65 speedup with 2x de- vices at long context lengths, even in cases where tensor parallelism fails to scale.
Our insight builds on prior observations that sparse atten- tions are effective (Kitaev et al., 2020; Beltagy et al., 2020; Xiao et al., 2024), but static approximate methods often fail on advanced reasoning or math tasks. By only accepting certain approximations at run time, we preserve accuracy across a broad set of benchmarks including reasoning, math, and information retrieval (Figure 2).
Inspired by instruction-Iever value prediction in CPUs, AL- Spec exploits attention sparsity in a dynamic and robust manner. Specifically, we demonstrate that a subset of chal- lenges in wransforme--sfeer exploiting the sparsity of self- attention can be effectively addressed through dynamic execution with speculation. In addition accelerating teaefformrr-bssed models, it points toward new opportu- nities for applying generic speculation paradigms to deep learning inference, unlocking new avenues for scaling.
Orthogonal to scaling up device count, sub-quadratic algo- rithms (Gu & Dao, 2024; Katharopoulos et al., 2020; Poli et al., 2023) and sparsity-exploitig atlention approximation methods cut attention computation and memory demands. Reformer (Kitaev et al., 2020) replaces the standard dot- product attention with locality-sensitie hashing (LSH) to reduce the per-token computational complexity from linear to logarithmic in context length. SCFA (Pagliardini et al.,
2. Motivation
Self-attention tends to dominate the latency of LLM infer- ence larger context lengths, as every new token attends
2

---

Attention-Level Speculation
Delative Store to Boxeline
Artention Sink
Amtontion 'Sink at Food Layers
Atterion-Lovel Spec. @ A-U.10 (Our Method] Baselino mark
"WALLL PRO
1.00
0.5
0.00
PPL base/PPL
IFEVAL
GPOA COT
SWIDE
GSMBK COT
MATH COT"
ECON car
BUSINESS COT
PSYCH_COM COT
Figure 2. Correctness evaluation of attention approximation methods on various benchmarks using Llama 3.1 8B. Baseline is the model with unmodified attention. Attention sink fails for tasks requiring math and reasoning on specific subject area, such MMLU_PRO with chain-of-thougle (COT). Statically allowing attention sink only at specific layers reduces degradation, but still causes a noticeable gap. Dynamically using attention sink with AL gives consistent correctness the baseline all tasks at acceptance threshold = 0.10.
2023) further reduces the computational cost of Reformer. Longformer (Beltagy et al., 2020) relies on sliding-window window (local) attention plus small number of global tokens. These approaches typically exhibit static sparsity in the sense that the approximation pattern is fixed prior to execution. Al- large-
Avg Num Layers with Accepted Approximations (out of 32) Needles found (out of 10)
10
30
Avg # Layers with Accepled Approx.
8
Number of Noedles Found
though effective in some cases (c.g., perplexity context text) or with retraining, the static nature of these methods means the model execution graph is fixed prior to execution, inevitably causing quality loss in scenarios where the model needs full attention for correct results (Figure I).
10
Aed (Ranetin))
0.0000
A 4.0.0
A.0.00
1 0.28
0.20
1 0.38
Amx (Aito Blok)
Among the approximation techniques, attention sink (Xiao al., 2024) is simple yet surprisingly powerful. It focuses the self-attention kens (e.g., 4) and a rolling window of S recent tokens from the KV cache. This makes the effective context length of attention drastically smaller than the true context length. Attention sink achieves good performance on certain tasks (e.g.. perplexity and information retrieval), even with a small S. However, any global information beyond the S-sized window or the first few tokens is inaccessible, causing at- tention sink to fail on tasks that require distant context or advanced reasoning. Figure 2 illustrates how attention sink on Llama 3.1 8B struggles on tasks like advanced math or reasoning (e.g.. M..MM. U.MLU_PRO with thain-of-thought, where yields a noticeable quality drop compared to the baseline (see Section 5.1 for methodology).
two sets of tokens: the first few to-
(a) Using attention sink at select layers does harm correctness
0.05
A 0.15
0.25
Approx, Accepted
75%
50%
20%
0%
Layer Number
(b) Average percentage of approximation accepted per layer
Figure 3. Finding needles in a haystack with Llama 3.1 8B.
selectively use attention sink with criteria that ensure only good approximations are used: layer accepts attention sink if, for the given token, the L2 distance between the attention sink output Á and full attention output A is within X L2-norm (A). A full attention sink model corresponds to = oo and the baseline model corresponds to X = 0. Figure 3a shows that for < 0.25, up to 2/3 of the layers can be approximated with attention sink without sacrificing the score. Figure 3b further shows that approximations are layer-specific: certain layers, especially the final layers, more likely to accept approximation without harm, corrobo- rating prior observations that deeper layers can be pruned with limited effect (Elhoushi et al., 2024).
The success or failure of approximation appears to be layer- specific: some layers can tolerate an approximate self-atten- tion output, while others require exact attention. Prior work shows that certain layers can be pruned or skipped without significantly impacting model quality (Elhayad et al., 2020). Consequently, one might wonder if we can statically choose subset of layers apply attention sink. To concretize this hypothesis, design synthetic Needles in a Haystack test on Llama 3.1 8B. The model inputs 16K-token con- text containing 10 randomly inserted "secret keys." It must retrieve all keys from the entire text. Figure 3a shows that a model with pure attention sink (S = 128) recovers none of the secret keys. In contrast, the baseline model finds 8 of 10. Between these configurations, we create models that
Figure 2 shows that purely relying on a fixed layer- layer-based approximation although better than applying atten-
3

---

Attention-Level Speculation
tion sink to all layers--xhibill noticeable quality degrada- tion when extended to variety of tasks. For the middle hars, we create fixed layer-based approximated models, where attention sink is always used at a layer iff the likelihood to approximate that layer is greater than p = 0.8 in Fig- ure 3b. The static approximation models do not balance performance gains with accuracy. Some inherently harder inputs or tasks require additional global context; even for the same task, the best approximation policy may vary from token to token or layer to layer. We therefore conclude that no single static approximation suffices; the key is to make decisions dynamically at run time.
tasks, we design verification step ensure correctness. In practice, we find that over half of the layers can skip full self-attention by verifying that A, R A,, Mcanwhile, ALSpec maintains baseline accuracy across tasks (Figure 2).
3.1. Low-Cost Predictor
We use a coarser version of attention sink that inputs the first and last S tokens from the KV cache. We choose S to be small relative to the context length (e.g.. S E {128, 256, 512] for 128K context length). In Section 4.1, we choose S as the chunk size of speculative flash decode.
We need a simple yet powerful way to realize dynamic exe- cution (Barad et al., 2024). Speculative parallelism involves a guess about future outcomes, a verification of that guess, and error recovery when the guess is wrong. By speculating, systems can proceed optimistically and exploit opportuni- ties for parallelism, even in the face of run-time uncertainty. Speculation is a long-standing technique in computer ar- chitecture, spanning branch, dependence, and value predic- tion (Rau & Fisher, 1993; Lipasti & Shen, 1996; Sazeides & Smith, 1997) to task-level speculation (Sohi et al., 1995; Herlihy& & Moss, 1993; Jeffrey al., 2015). Speculation enables concurrent parallel execution of instructions or blocks of code, but discards results when the guesses fail. In LLMs, speculation thus far has appeared only at the decode (or token) level via speculative decoding (Leviathan et al., 2023; Hooper et al., 2023; Spector & Re, 2023; Li et al., 2024: Miao al., 2024) which drafts several tokens in par- allel. This includes decode-level speculation that employs attention sink (Sun et al., 2024).
3.2. Low-Cost Verification
Relaxed Verification via Lipschitz Continuity: Exact equality between A, and A, is not mandatory to preserve correctness in LLMs; we only require that their difference remains bounded so that final outputs have small deviations. Lipschitz continuity (Scaman & Virmaux, 2018) provides a theoretical tool to relate small perturbations in one layer's output to eventual changes in final logits. Although self- attention alone is not Lipschitz contimuous (Kim et al., 2021; Castin et al., 2024; Geshkovski et al., 2023), ford-forward networks, layer norms, and residual connections possess (or can be bounded by) Lipschitz constants. In addition, Geshkovski et al. show that self-attention 's Lipchitz constant is bounded by & where R is the input magnitude bound. Combining these, we derive an upper bound on the final output's deviation in terms of per-layer differences | - Asl. Concretely, if each approximate attention output is within 6, in fy-norm of the ground-truth A,, the overall output deviation, c, is bounded by
3. A Attention-Leve Speculation Overview
S E (1 (1) (1+ ((R))) N-66
(1)
Attention-leve -evell speculation (ALSpec) speculates on the self- attention output within each transformer layer. Specifically, ALSpec approximates the self-attention with a technique like attention sink and subsequently computes the exact self- attention. Once the exact result arrives, ALSpec verifies how close the approximate output is. If sufficiently close, ALSpec accepts it; if not, it reverts to the exact output. The "ideal scenario" of Figure shows that speculation can overlap attention with feed-forward layers, greatly reducing latency when the prediction is correct. The "worst case" " matches data parallelism latency; if the speculation fails, the computation must fall back the exact path. This approach harnesses parallel devices (or threads) to overlap computation, hiding some of the attention latency while preventing quality degradation when approximation fails. We show that a simple approximation-namel attention sink-server serves effectively as the "predictor." For each layer, we apply it to self- self-attention, generate a speculated attention output A, in layer i, and concurrently compute the ground- truth output As. Since attention sink can fail for certain
where a, 8, and (R) are bounded constants for the Lip- schitz behavior of residuals, feed-forward networks, and self-attention inputs, respectively. A detailed proof and discussion appear in Appendix A.
Verification Algorithm:
Motivated by the above bounds we adopt simple threshold-basee metric to accept reject Å, Concretely, we measure di d(A,, A,) in the & norm and accept if: ||A,, - Allla < A ||Ar where X is small scalar hyperparameter. Algorithm I sketches the speculative execution procedure. For each layer, the algorithm (i) com- putes A, using attention sink, (ti) launches a new thread to continue the "post-attention" operations, (iii) computes A; in the main thread, then (iv) verifies if A, is sufficiently close. If yes, it commits A;; otherwise, it discards the new thread and re-runs the next operations on the main thread's ss A,. A simple thresholding strategy effectively ensures the algorithm does not propagate large attention errors. In ad- dition, we show a high probability error bound on e in the
4

---

Attention-Level Speculation
Scenario I: more resources than work
Algorithm 1 Speculative Execution with Attention Sink
d1
Attni
Input: activation x, N layers, threshold
d2
Attin
Idecl
Start execution as the main_thread. for i = 1 to N do
d3
FF
Attin
Jeli speculation correct)
d4
FFg
T - opw_beooreeeooopoeeooepeeeeeeeep wwteesseeefaeaeeeeeeeeeeee self_att((())) // // e.g. LayerNorm
d1
First speculatton wrong- Total rustine
x spec = attention sink(æ)
d2
Artng
FFabs
Artna
eouivalant to ideel except for first emtention executed in full.
new _threed = opssatr sett, Espec)
d3
ancarded
FF:
Attnn
FFg
x = reggllrraaattet attention(:))
d4
if L22 distance(1r (Fpee)) < X L2 norm(r) then
Dovice (d)
main thread = new thread
Scenario 2: more work than resources
Idle
else
d1
Attn:
FFa
Attoy
Confitets
Kill new ttheed
d2
Atting
FF.
Attng
x. =ops after self attn(x) // e.g. Feed Forward
Aftizng
end if end for
wait
A1
Attny
A4
Attin,
Weiting free
d2
FF
A2
Attiz
55. As
Aftns
resources
order of VN (Appendix A.3). This is consistent with our ex- perimental results on deeper models (Section 5.2), achieving similar results as shallower models across benchmarks.
Scenario 3: full attertion runs faster than feed forward
d1
Attn:
ool
Ff..
Ac
Cambine parallallun
speculative paralleliun reduce Ff lntency-
d2
FF,,
cel Ag
Atting:
d3
A1
Attnr p col
FP32
sc
AS
device and 28.4.
3.3. Example Speculative Executions
d4
FF
cc Ac
Attraa
col
ALSpec batween and 38-4.
Figure 4 shows three high-level scenarios that can arise when applying Algorithm on multiple devices in parallel. Scenario 1: More resources than work. With sufficiently many devices, every layer's approximate and exact attention can run in parallel without contention. In the ideal case, all speculations are correct, so each device can immediately continue with post-attention operations (e.g., feed-forward (FF) layers), yielding maximum speedup by overlapping all FF within attention. In the mnis-peculatiinn case, when the first approximation fails, the system reverts to full attention for that step. Subsequent layers resume the normal specula- tive pattern. The total latency impact of one mis- nis-sueculation remains bounded by the time needed to execute a single at- tention pass in full. Scenario 2: More work than resources. In practice, we often have fewer devices than are needed to run every speculation fully in parallel. A device may need to wait until prior layer finishes using the hardware be- fore starting its own speculation. Scenario 3: Full attention runs faster than ford-forward. In this scenario, speculation only uses 2 devices the prior device will be free when FF finishes. In such cases, we combine tensor parallelism with speculative parallelism: Figure 4 illustrates a 4-device system using 2 devices for tensor parallelism and 2 devices for speculative parallelism with 1:1 mapping.
Attn
Attention
FF
Feed Forward
Communication
Approximnted Attertion
Figure 4. Algorithm 1 run-time patterns. (Scenario 1) More re- sources than work: Each approximate and exact path concurrently on separate devices. If all speculations are correct, total latency is minimized; a mis-speculation simply triggers the ex- act attention for one layer. (Scenario 2) More work than resources: Some computations must wait for free devices, adding stalls. Oth- erwise, the overlap pattern is similar to Scenario 1. (Scenario 3) Feed-forward dominates: When FF is slower than attention, we can pipeline or combine tensor parallelism with speculative parallelism (mapping devices 1:1) to further reduce FF latency.
2018), complicating static execution graph. In contrast, Scenario 3 naturally fits an op-by-op framework: one device is used for an approximate path and another for an exact path. No further pipelining is required. Hence, in the remainder of this paper, we focus on implementing Scenario 3 and extending it with tensor parallelism real systems.
3.4. Overheads
On the communication side, combining ALSpec and TP (Scenario 3) for n devices reduces latency vs. full TP on n devices. In each TP operation. the Sandwidth-optimal all- reduce algorithm (Thakur et al., 2005) has each device send and receive n-1 of the tensor per all-reduce. While the total bandwidth demands per device remain nearly constant as TP device count increases, the latency of all-reduce operations during decode are often exposed on the critical path due to small activation sizes (Narayanan et al., 2021). This latency is highly dependent on the underlying network topology. For
In most real-world LLM usage, FF dominates run time while attention latency grows with context length (Dao, 2024). For context lengths up to 128K tokens, FF remains the primary bottleneck, Icading to Scenario 3, whereas Scenario I or Scenario 2 only become relevant at extremely large context lengths. Moreover, Scenario 1 & require fully pipelining or even speculating on still-speculative data (Jeffrey et al.,
5

---

Attention-Level Speculation
example, in ring topology (Sergeev & Balso, 2018), all- reduce requires n. hops of exposed latency. In contrast, ALSpec+TP on devices performs all-reduce across 2 devices, with 2 -1 hops of exposed latency and adds 1 hop of exposed latency for the approximating devices to send speculated activations to the speculative FF path.
Flodh Decode feod Pottern
Kor Vokee Codes
New
Core
Care 2
Core
Specslative Flush Decode Reed Puttern
Key Volv Gods
On the computation side, the ALSpec verification step in- troduces two Le norm calculations of the ground truth and speculated activations. This L2 norm cost is similar to Lay- enNorm. For decode with small activations, the L2 norm computation is negligible compared to FF and attention.
Core
Core 2
Core
Figure 5. Flash decode and speculative flash decode KV cache read patterns. Speculative flash decode reads the first and last chunks first and computes partial result as speculation.
3.5. Summary
Speculative parallelism offers simple, dynamic exten- sion to static approximation. Rather than retraining or modifying model architecture, we only overlay a predic- tion-verification mechanism. Speculative parallelism ex- tends naturally to other sparse attention approaches, creating a versatile framework for dynamic execution. This dynamic viewpoint transforms seproximate-attention methods like attention sink from "always approximate" to "approximate only when safe," mitigating their inherent latancy-quality tradeoff. By combining run-time verification with paral- lel computation, we overlap approximate and full-attention paths to reduce the critical path latency.
mitigate this, we modify the read pattern of the flash decode kernel and implement an all-fused kernel that speculates for free while computing full attention and verification at once. Appendix B shows the full algorithm.
In short, the intuition of speculative flash decode is shown in Figure 5. The original flash decode algorithm (Dao, 2024; Dao et al., 2023) processes self-attention by chunks of KV cache from the start to the decoding position, where groups of chunks are distributed across the available cores. The results are then aggregated using the statistics. We exploit the order invariance in which the chunks are processed. In speculative flash decode, we set the chunk size to be S, as- sign the first and last chunk to be read first, and send the aggregated partial result immediately to the next device to execute the speculated path. The partial result is equiva- lently the attention sink on the first and last S' tokens, and can be used later towards calculating the full self-attention result. We then combine the verification step into the same single kemne, using it as drop-in replacement for Py PyTorch's ss sca l wtw.rreeettttttteiii........
4. Implementation
Implementing Scenario 3 poses significant challenges. Tra- ditional frameworks assume an op-by-op execution pattern over static dataflow graph that maps onto well-known scal- ing paradigms like tensor and data parallelism. By contrast, attention-leve speculation requires dynamic conditional ex- ecution over the graph. An approximate path runs concur- rently with computing the exact attention, and only upon verification do we commit to one.
To make the implementation feasible, we introduce two key optimizations. First, we design a custom Speculative Flash Decode kernel that fully hides the cost of the approxi- mate attention path incur minimal penalty if speculation is rejected. Second, introduce static graph, dynamic concurrency (SGDC) with priority gating: mechanism to preserve op-by-op host code and a static computation graph, while dynamically forking and merging the specu- lative path under the hood. Verification then determines which result to commit, controlled by a lightweight priority gating mechanism.
4.2. SGDC with Priority Gating
Algorithm 2 implements Scenario 3 of Algorithm 1 on two groups of devices by using the speculative flash decode kemel and small change to the runtime, namely SGDC with Priority Gating. It uses the following methods:
Siip_compute(p// = 0): Skips a subsequent opera- tion's execution when pik] == 0, where k indicates the device index of itself.
all gather(p): Gathers p across d1, da (d with the highest priority holds the correct vz).
4.1. Speculative Flash Decode
speculative flash decode(z, p): Executes speculative decoding operation and determines the sender and re- ceiver based on p. Sender sets ps = if verification fails and otherwise pk Receiver sets pk= = 1.
In Algorithm 1, the attention sink runs before the regular attention. In the worst case, where all speculations are rejected, the latency would be worse than pipelining. To
6

---

Attention-Level Speculation
verify succeed: skip compute post & pre
ell__ gather
speculative Satpuulattve_fllshhhh www.allllsooooaaaaaaeeee decode (sfd)
all gather
Device
Send speculction result
Send
verify folls: do post & pre
Device
Send
all gother
sfd
post_attn_ops (post)
pre_attn_ops (pre)
galggaterr
Figure 6. Timing diagram of speculative exccution on two devices, hased on Algorithm 2. We show segment of exccution timeline from the all-gather op in the first layer to the all-gather op in the second layer. When speculation succeeds, Device skips the computation in post- and pre pre-attention ops. When speculation fails, it executes them fully.
Algorithm 2 SGDC with Priority Gating
5. Experiments
Input: replicated activation æ, N layers, threshold X, replicated priority vector p. device d1, d2, device k. Initialize p- = [p= = 2. P2 = 0), indicating di holds the correct with higher priority.
We evaluate the implementation of Algorithm 2 with the speculative flash decode kernel to answer the following questions: (i) Does dynamic execution really improve per- formance on real device implementation? (ii) What rate of speculation hit rate do we see across benchmarks? (iti) Why is dynamic sometimes superior to static approximation?
for i = 1 to NN do
with skip_compute(plf/ == 0):
x = ops before self attn(z)
p = all gll_ggaher(pp)
5.1. Methodology
T = speculative flash decode(x; p. A)
with akip.compute(pp) = 0):
We answer the previous questions with a case study of parallelizing Llama 3.1 8B model onto Tenstorrent N150 chips. The model architecture and weights are taken from the official Llama codebase without modifi- cation and implemented on Tenstorrent N150 chips, using the Tenstorrent Metalium (TT-Metalium) kernel library. The model is executed with default mixed precision con- figuration," 4 with BF16 activations, BF8 KV cache, and (1),,) model weights. For correctness evaluations, we run the model in BF16 precision using an NVIDIA A100 H100 GPUs; simulate speculation by running both full and approximated attention (S 128) and choosing the approximated attention output if verification succeeds.
Model & Hardware:
x www.arreeeefffeeeeeeeeee
end for
return x from the correct device based on p
Figure 6 illustrates Algorithm 2. We maintain static host code where model and input activations are replicated across two devices, effectively data parallelism. Each device main- tains priority vector to track which of them holds the correct activations. Initially, both devices hold the correct activations; once speculations and computation skipping happen, only one device holds it. We perform an all-gather of the priority vector before the speculation point to deter- mine which device has highest priority. At the speculation point, the sender device with highest priority (indicating correct activations) will execute the compute side of the speculative flash decode op, and the receiver device waits for the speculative result, then continues subsequent ops. When the full self-attention finishes, the sender will either skip the computation of the subsequent ops if the verifica- tion succeeds, or proceed executing the ops normally the verification fails. We maintain the same ops across devices and the host dispatches all the ops asynchronously. This method can be extended to tensor + speculative parallelism: when the first group comprises > 1 tensor-paralle devices, the second group has the same number of devices with 1-1 mapping.
We evaluate the correctness and specula- tion hit rate measurement at various choices of A. We group the evaluations into categories of [Question Answering (QA), Information Retrieval (IR), Reasoning, Long Con- text, and Math). We use the LM Evaluation Harness (Gao et al., 2024), a framework for few-shot language model eval- uation, with the default prompts and few-shot templates. As we focus on how ALSpec outputs diverge from the baseline, we do not explicitly optimize for baseline model scores.
Benchmarks:
"We conducted our experiments Tenstorrent N150 devices using the pre-release version of the TT-Metalium v0.55.0-rc1 soft- ware stack. Tenstorrent has active ongoing software optimizations, therefore the performance results presented here are preliminary and do not reflect the latest optimized capabilities of Tenstorrent devices. Rather, they serve to illustrate the relative performance ad- vantage provided by ALSpec. See details about TT- MT-Metali (Ten- storrent, 2024b), speculative flash decode, and SGDC with priority gating implementation in Appendix C, D & F.
"Appendix E validates SGDC with a micro- benchmark. In our real implementation, we also sync the residual stream at attention ensure the KV cache and residual are calculated correctly. We overlap it with compute so that it does not cause extra overhead.
*See Appendix F for TT-Metalium Llama implementation and performance measurement with tensor parallelism and AL.Spec.
7

---

Attention-Level Speculation
Table 1. Evaluation of attention-level speculation at various speculation verification threshold As for Llama 3.1 8B model across benchmarks, grouped by calegories. Each entry contains (evaluated scores (0-1)) / (speculation hit rate (0-100%)). Higher À corresponds to more relaxed verification, allowing higher speculation hit rates, while lower X keeps oulputs closer to the baseline.
TASK
QA
MATHEMATIC MATH
IFEVAL
GPQA COT
SWDE
FDA
GSM8K COT
MGSM COT SW
BASRLINE
0.798
0.237
0.359
0.210
0.823
0.326
0.580
A = 0.05
0.799/ 69.5%
0.257 38.0%
0.371 57.8%
0.205 / 33.5%
0.806/44 49.4%%9.4%
0.332/21.69
0.592/53.66
0.10
0.812
0.237 65.7%
0.389 83.2%
0.205
369.3%
0.799/
0.315 49.66
0.584 77.6%
A = 0.15
0.807
/ 99.8%
0.214/81.18
0.391 1901999.99
0.190/ 85.7%
0.799 / 90.2%
0.331
69.31%
0.592 / 86.0%
0.20
0.805 98.0%
0.225 490.4%
0.384/97 0944/97.000
0.172 0122/94.49
0.795 94.6%
0.323 79.4%
0.564/ 995.4%
À 0.25
0.794 99.3%
0,208 49.44%
0.359 / 98,
0.158
99.
0.719 /9.2%
0,266
86.3%
0.516 / 96.6%
TASK
LONG CONTEXT REPORENCH.P
REASONING BUSINESS
HOTPOTQA
ECON
PSYCH
BASELINE
0.920
0.756
D.328
0.470
0,289
A 0.05
0.920 18.5%
13,11%
0.341/18.99
0.477 22.8%
0.232
23.3%
À 0.10
0.915 42.8%
0.768
59.6%
0.339/
45.8%
0.490 50.0%
0.296
52.3%
0.920 56.8%
0.812
75.5%
0.321
63.0%
0.473 68.55%
0.294
/
A 0.20
0.910/71.66
0.756/
85.0%
0.316 77.5%
0.437 / 80.7%
0.292 80.9%
A 0.25
0.705 80.1%
0.742 91.8%
0.273 87.7%
0.359 / 88.3%
0.281 / 99.3%
PMO BENCHMARK WITH CHAIN OF THOUGHTS DECODING.
Measurement: We use the TT-Metalium profiler to mea- sure the device kernel time of every operation during decode at specific context lengths and speculation hit rates. This performance measurement shows an ideal scalability pattern because it does not account for host and op dispatch latency, which are subject to the host configuration. We report la- tency per token, which is estimated by taking the measured kernel run time per layer and multiplied by the total num- ber of layers (32). The metric token/s/user is calculated by 1/(latency per token).
Table 2. Evaluation of speculative execution at various verification thresholds À for Llama 3.1 70B model across selected benchmarks.
TASK
GPQA COT
GSMSK COT
HOTPOTQA
BASELINE
0.518
0.958
0.940
À 0.05
0.529/ 55.1%
0.950/ / 64.5%
0,940 32.3%
= 0.10
0.507 / 80.0%
0.951 / 86.6%
0,945 / 59.0%
0.15
0,458 / 89.9%
0.946 99..9%
0.935 74.6%
0.20
0,446 / 95, 1%
0.936 / 96.8%
0.935 / 83.3%
0.25
0.379 /97.997.6
0.897 99..3%
0.935 89.4%
to scale. In the best case (96k context length with 87.5% speculation hit rate), ALep provides 1.65x improvement in per-user throughput compared to full tensor parallelism.
5.2. Speculation Hit Rates and Performance Results
Table shows the output quality and speculation hit rate for À E f0.05, 0.10, 0.15, 0.20, 0.25) compared to baseline. ALSpec with x {0.05, 0.10} achieves on par or better correctness on all evaluated tasks, with speculation hit rates ranging from 18% (HotpotQA) to 90% (IFEval) depending on the task, and most speculation hit rates exceeding 50% for 0.10. Relaxing À to 0.15 allows the model to speculate up to 95% of attentions successfully while doing slightly worse on GPQA and FDA. Adjusting À beyond 0.15 allows flexibly trading off model accuracy for computation latency, which could be helpful depending on the use case. Table 2 shows a subset of benchmarks with Llama 3 70B (80 decoder layers), confirming that a larger and deeper model achieves similar speculation hit rate with limited correctness degradation
5.3. Reproducibility and Estimation GPUs
We corroborate the diminishing returns of tensor parallelism previously seen GPUs and estimate ALSpec performance on GPUs. Chitty-Venkata et al. find that the Llama 3 8B model on A100s with DeepSpeed-MII achieves around 75, 90, and 115 tokens/s per-user throughput when scaling across I, 2, and 4 GPUs. This scaling is similar to what we observe on N150 devices. We measure the Llama 3 8B decoding latency on 4 vs. 8 H100 GPUs doing TP using the SGLang (Zheng et al., 2024) serving framework and the FlashInfer (Ye et al., 2025) attention backend. Although we currently lack an ALSpec CUDA implementation, we es- timate the ALSpec performance using SP@65% + TP4. Table 3 shows our estimation that ALSpec cuts latency on Llama 8B by 1.28x vs. 1.09x for full TP context length 128K. In this case, the attention latency is only 63% of non-attention latency, suggesting greater gains for con- text lengths beyond 1 128K. Although AL ALSpec introduces a new attention kernel and SGDC, it does not fundamentally change the op- and static graph execution style of a GPU runtime. We expect that AL is implementable via CUDA graphs and conditionals (Gaiser al., 2024).
{0.05, 0.103.
Leveraging the previous experiments, we estimate perfor- mance for Llama 3 8B on N150 chips based on commonly observed ranges of the speculation hit rate: {50, 62.5, 75, 87.5).% %. Our results show that ALSpec provides additional scaling when tensor parallelism scaling diminishes. Figure shows that full tensor parallelism provides almost no im- provements when scaling from to 8 devices at long context length, whereas speculative + tensor parallelism continues
8

---

Attention-Level Speculation
Context Length: 1K
Context Length: 32K
Context Length: 64K
Context Length: 96K
Context Length: 128K
Full TP siggos 1P
SPgg22 5% TP SPg/5% 1D
SPg87.5 + IP
Tokens///usar
40
40
20
20
20
20
20
0
8
2
2
4
8
2
4
8
4
Number of Devices
Figure 7. Scalability patterns Llama 3.1 8B with full tensor parallelism (TP) and speculative parallelism at (50, 62.5, 75, 87.5) % speculation hit rate with tensor parallelism (SP@XS TP) on Tenstorrent N150 chips. Each plot shows the scalability at a particular decode context length (Ik, 32k, 64k, 96k, and 128k). Attention speculative parallel using the speculative flash decode kerne with chunk size 256 and tensor parallel by head. Feed-forward layers are tensor parallel by inner dimension. Performance estimation is obtained by measuring the device kernel duration of every single kernel, agnostic of dispatch and host time, At long context lengths, SP+TP (our method) continues to scale with 8 devices while full tensor parallelism scaling diminishes.
Table Latency and throughput scaling for Llama 8B inference on H100s across various context lengths. Latencies are measured per decoder layer, and throughput is in tokens per second (tok/s). Projected ALSpec results assume 65% speculation hit rate.
4XH100 ATTN
4XH100 NON-ATTN LATENCY/LAYER (US)
4xH100 TP
8xH100 TP TOK/S
TP SCALING
PROJECTED ALSPFC @ @65% HIT RATE
SLEP SCALING
CONTEXT
LATENCY/LAVER (US)
TOK/S
Ik
13
95
244.6
249.3
1.9%
262.0
7.1%
32K
29
100
214.2
231.3
8.0%
246.0
14.8%
64K
49
100
191.4
209.7
9.6%
237.8
24.2%
96K
56
100
178.2
194.2
9.0%
224.9
26.2%
128x
169.3
184.8
217.5
28.5%
Prompt:
model is given a math problem with 5-shot examples. We vi- sualize each token, shaded by its speculation rate, generated by the speculative execution model. The static approxima- tion model is unable switch to the second question during generation, hence diverging. The diverging point is a token with a low speculation rate, where it changes the topic from Question 1 to 2. Similarly, we mark the points where topics from previous contexts were introduced and they overlap with tokens with small speculation rates. Overall, dynamic execution via ALSpec detects large approximation errors and adaptively recalibrates to maintain generation quality.
The following multiple choice questions (with answers) about math. Think step by step and then finish your answer with ''hee answer is (X)"" where X the correct letter choice...
Correct Generation from Speenlative Exceution:
I'll answer each question step by step. ** Question 1** The churacteristic of ring the smallest positive unleged njsnch that times any element in the ring is zero lo this case, we have 2Z which means all elements are multiples of So, if we multiply any element by we get zero, Therefore, the charucteritti of 2. The answer (G).
""Qoesttonn 200 3 Les analyye the transformations] [1 and S: (Tnt)))) = vola) S(p(s))
Therefore, the correct answer [B] [ Question 61 (Satemmen Iis trae. Afree abeliar copics of the intcasrs, such group exists for every positive integer n. Statement 2 also true. Afinitely generated ahelian group free abelian and only if its Betti number equalss the number of elements in some generating set, Therefore, the correct answer is
group of fank is ismmrphic to direct
spee.
X Failed Generation from Static Approximation:
accepted
100
6. Concluuion
I'll answer each question slep by stup. **uustionn 1 The characteritiic of the ring answer is (G). **uusstion 240 To determine characteristic of rine
We presented ALSpec as a new paradigm for scaling LLM inference, addressing the shortcomings of existing tensor and data parallelism methods. By leveraging predictions of attention outputs, our approach overlaps computations to reduce latency while maintaining model accuracy. The successful implementation on TT-Metalium underscores the practicality of our method. While this work focuses on speculating attention, the underlying principles of specu- lative parallelism are general and hold potential for wider applicability. Future research will explore the generalization of this technique, particularly in mechanisms to repair and eliminate error accumulation, well as generic op-level speculative parallelism, which extends these principles be- yond attention to optimize a broader range of ops in LLMs.
Diverging Points
Unable to move forward Question
toples previses
re-intreduced
Figure& 8. Example of output divergence in static approximation model. Divergence points corresponds to tokens with low specula- tion accuracy, where the static approximation model fails to adjust.
5.4. When Does Static Approximation Go Wrong?
In some cases, static approximation is unable to generate key tokens that influence the direction of the conversation, indicating that full attention to the context is required. This usually happens at a change of topic token. Figure 8 illus- trates a specific example MMLU PRO MATH where the

---

Attention-Level Speculation
Acknowledgments
2024. URL https: https///pppeeeeeeeeeee net/ fer.iieummmee forum? i id=m2n2Xhh9Ec.
We sincerely thank Utku Aydonat, Steven Hill, Brian Liu, Aster Plotnik, Gilead Posluns, Aditya Saigal, Angus Wu, and the anonymous reviewers for their helpful feedback. This work was supported in part by Tenstorrent through the use of N150 and GPU devices and funding from Fujitsu.
Dao, T., Haziza, D., Massa, F., and Sizov, G. Flash-decoding for long-context inference. Stanford University Cen- ter for Research on Foundation Models Blog, Oct 2023, URL ht.pt::/tt///tt fm. stanford......... htt.fodddddd//22////3/ 12 12/flashdecoding html.
Impact Statement
DeepSeek-AI, Liu, A., Feng, B., Xue, B., Wang, B., Wu, B., Lu, C., Zhao, C., Deng. C., Zhang, C., Ruan, C., Dai, D., Guo, D., Yang, D., Chen, D., Ji, D., Li, E., Lin, F., Dai, F. Luo, F., Hao, G., Chen, G., Li, G., Zhang, H., Bao, H., Xu, H., Wang, H., Zhang, H., Ding, H., Xin, H.,, Gao, H., Li, H., Qu, H., Cai, J. L.,, Liang, J., Guo, J., Ni, J., Li, J., Wang, J., Chen, J., Chen, J., Yuan, J., Qiu, J., Li, J., Song, J., Dong, K., Hu, K., Gao, K., Guan, K., Huang, K., Yu, K., Wang. L., Zhang, L.,, Xu, L., Xia. L., Zhao, L., Wang. L., Zhang, L.., Li, M., Wang, M., Zhang, M., Zhang, M., Tang, M., Li, M., Tian, N., Huang, P., Wang, P., Zhang, P., Wang, Q. Zhu, Q.. Chen, Q., Du, Q. Chen, R. J.. Jin, R. L., Ge, R., Zhang, R., Pan, R., Wang, R., Xu, R., Zhang, R., Chen, R., Li, S. S., Lu, S., Zhou, S., Chen, S., Wu, S., Ye, S., Ye, S., Ma, S., Wang. S., Zhou, S., Yu, S., Zhou, S., Pan, S., Wang, T., Yun, T., Pei, T. Sun, T., Xiao, W. L., Zeng. W., Zhao, W., An, W., Liu, W., Liang, W., Gao, W., Yu, W., Zhang, W., Li, X. Q.. Jin, X., Wang, X., Bi, X., Liu, X., Wang, X., Shen, X., Chen, X., Zhang, X., Chen, X., Nie, X., Sun, X., Wang, X., Cheng, X., Liu, X., Xie, X., Liu, X., Yu, X., Song, X., Shan, X., Zhou, X., Yang, X., Li, X., Su, X., Lin, X., Li, Y. K., Wang. Y. Q., Wei, Y. X., Zhu, Y. X., Zhang, Y., Xu, Y., Xu, Y., Huang, Y., Li, Y., Zhao, Y., Sun, Y., Li, Y., Wang, Y., Yu, Y., Zheng, Y., Zhang, Y., Shi, Y., Xiong, Y., He, Y., Tang, Y., Piao, Y., Wang, Y., Tan, Y., Ma. Y., Liu, Y., Guo, Y., Wu, Y., Ou, Y., Zhu, Y., Wang, Y., Gong, Y., Zou, Y., He, Y., 7ha, Y., Xiong, Y., Ma, Y., Yan, Y., Luo, Y., You, Y., Liu, Y., Zhou, Y., Wu, Z. F.,
This paper presents work whose goal to advance the field of Machine Learning. There are many potential societal consequences of our work, none which feel must be specifically highlighted here.
References
Almazrouci, E., Alobeidli, H., Alshamsi, A., Cappelli, A., Cojocaru, R., Debbah, M., Étienne Goffinet, Hess- low, D., Launay, J., Malartic, Q.. Mazzotta, D., Noune, B., Pannier, B., and Penedo, G. The Falcon se- ries of open language models. arXiv preprint, 2023. 000:100485550axi12311..666
Barad, H., Achterberg. J., Chou, T. P., and Yu, J. Acceler- ated AT inference via dynamic execution methods. arXiv preprint, 2024. doi: 00.44555/XX....2111100002
Beltagy, I,, Peters, M. E., and Cohan, A. Longformer: The long-document transformer. arXiv preprint, 2020. 1001100.45000rxii2000..000.
Cai, J. Flashdecode on tenstorrent's wormhole architecture. Tech Report, Tenstorrent, 2024.
URL https://githii: http:/hhhb.ccoeeetsstttrreee tt-mmataa//bbo//mmaat//eeeeeeete setihatttentioo/llaaaeeeoooo
Castin, V., Ablin, P., and Peyré, G. How smooth is atten-
tion? In Proc. International Conference on Machine Learning (ICML), volume 235, pp. 5817-5840. PMLR, 2024. htthttss: httptt:s:/pprooeediiiie pttp///ecceee...eeeeee/ .ml mlr.press/ html.
Ren, Z. Z., Ren, Z., Sha, Z., Fu, Z., Xu, Z., Huang. Z., Zhang. Z. Xie, Z., Zhang, Z., Hao, Z., Gou, Z., Ma, Z.. Yan, Z., Shao, Z., Xu, Z., Wu, Z., Zhang, Z., Li, Z.,, Gu, Z. Zhu, Z., Liu, Z., Li, Z., Xie, Z., Song, Z., Gao, Z., and Pan, Z. Deepseck-v3 technical report. arXiv preprint, 2024. doi: 10014485500///X1122222..444
Chitty-Vnnktta, K. T., Raskar, S., Kale, B., Ferdaus, F., Tanikanti, A., Raffenetti, K., Taylor, V., Emani, M., Inference
and Vishwanath, V. benchmarking of large language models on Al accel- erators. In Proc. Workshops of the International Con- ference for High Performunce Compuling, Networking, Storage and Analysis (SC24-W), W), pp. 1362-1379, 2024. 100SS11//9 2024.0017..
Inference-
Dubey, A., Jauhri, A., Pandey, A., Kadian, A., Al-Dahle, A., Letman, A., Mathur, A., Schelten, A., Yang, A., Fan, A., Goyal, A., Hartshorn, A., Yang, A., Mitra, A., Sravanku- mar, A., Korenev, A., Hinsvark, A., Rao, A., Zhang, A.. Rodriguez, A., Gregerson, A., Spataru, A., Roziere, B., Biron, B., Tang, B., Chern, B., Caucheteux, C., Nayak, C., Bi, C., Marra, C., McConnell, C., Keller, C., Touret, C., Wu, C., Wong, C., Ferrer, C. C., Nikolaidis, C., Al- lonsius, D.,. Song, D., Pintz, D., Livshits, D., Esiobu, D., Choudhary, D., Mahajan, D., Carcia-Qlang D., Perino,
Dao, T. FlashAttention-2: Faster attention with better parallelism and work partitioning. In Proc. Interna- tional Conference on Learning Representations (ICLR),
10

---

Attention-Level Speculation
D., Hupkes, D., Lakomkin, E., AlBadawy, E., Lobanova, E.,, Dinan, E., Smith, F. M., Radenovic, F., Zhang. F., Syn- naeve, G., Lee, G., Anderson, G. L., Nail, G., Mialon, G. Pang, G., Cucurell, G., Nguyen, H., Korevaar, H., Xu, H., Touvron, H., Zarov, I., Ibarra, I. A., Kloumann, I,, Misra, I., Evtimov, I., Copet, J.. Lec, J., Geffert, J., Vranes, J., Park, J., Mahadeokar, J.. Shah, J., van der Linde, J., Billock, J., Hong. J., Lee. J.. Fu, J.. Chi, J., Huang, J., Liu, J., Wang. J., Yu, J., Bitton, J., Spisak, J., Park, J., Rocca, J., Johnstun, J., Saxe, J.. Jia, J., Alwala, K. V., Upasani, K. Plawiak, K., Li, K. Heafield, K. Stone, K., El-Arini, K. Iyer, K., Malik, K. Chiu, K., Bhalla, K., Contala-Yeary, L., van der Maaten, L.., Chen, L., Tan, L.., Jenkins, L., Martin, L.., Madaan, L.. Malo, L.., Blecher, L.., Landzaat, L., de Oliveira, L., Muzzi, M., Pasupuleti, M., Singh, M., Paluri, M., Kardas, M., Oldham, M., Rita, M., Pavlova, M., Kambadur, M.. Lewis. M., Si, M., Singh, M. K. Hassan, M., Goyal, N., Torabi, N., Bashlykov, N., Bogoychev, N., Chatterji, N., Duchenne, O., Çelebi, O., Alrassy, P., Zhang. P., Li, P., Vasic, P., Weng, P., Bhargava, P. Dubal, P. Krishnan, P., Koura, P. S., Xu, P., He, Q., Dong, Q., Srinivasan, R., Ganapathy, R., Calderer, R., Cabral, R. S., Stojnic, R., Raileanu, R., Girdhar, R., Patel, R., Sauvestre, R., Polidoro, R., Sumbaly, R., Taylor, R., Silva, R., Hou, R., Wang, R., Hosseini, S., Chennabas- appa, S., Singh, S., Bell, S., Kim, S. S., Edunov, S., Nie, S., Narang, S., Raparthy, S., Shen, S., Wan, S., Bhosale, S., Zhang, S., Vandenhende, S., Batra, S., Whitman, S., Sootla, S., Collot, S., Gururangan, S., Borodinsky, S., Her- man, T., Fowler, T., Sheasha, T., Georgiou, T., Scialom, T., Speckbacher, T., Mihaylov, T., Xiao, T. Karn, U., Goswami, V., Gupta, V., Ramanathan, V., Kerkez, V., Gonguet, V., Do, V., Vogeti, V., Petrovic, V., Chu, W., Xiong, W., Fu, W., Meers, W., Martinet, X., Wang, X_, Tan, X. E., Xie, X., Jia, X., Wang, X., Goldschlag, Y.,
Brinkman, E., Arcaute, E., Dunbar, E., Smothers, E., Sun, F., Kreuk, F., Tian, F., Ozgenel, F., Caggioni, F. Guzmán, F. Kanayet, F., Seide, F., Florez, G. M., Schwarz, G., Badeer, G., Swee, G., Halpern, G., Thattai, G., Herman, G., Sizov, G., Guangyi, Zhang, Lakssminarayanan, G., Shojanazeri, H., Zou, H., Wang, H., Zha, H., Habeeb, H., Rudolph, H., Suk, H., Aspegren, H., Goldman, H., Damlaj, I., Molybog, I. Tufanov, I., Veliche, I.-E., Gat, I. Weissman, J., Geboski, J., Kohli, J., Asher, J., Gaya, J.-B., Marcus, J., Tang. J., Chan, J., Zhen, J., Reizenstein, J., Teboul, J., Zhong, J., Jin, J., Yang. J., Cummings, J., Carvill, J., Shepard, J., McPhie, J., Torres, J., Ginsburg, J., Wang, J., Wu, K., U, K. H., Saxena, K., Prasad, K., Khandelwal, K., Zand, K., Matosich, K. Veeraragha- van, K., Michelena, K., Li, K., Huang, K., Chawla, K., Lakhotia, K., Huang, K., Chen, L., Garg, L., A, L.. Silva, L., Bell, L., Zhang. L., Guo, L., Yu, L.,, Moshkovich, I.,, Wehrstedt, I.., Khabsa, M., Avalani, M., Bhatt, M., Tsimpoukelli, M., Mankus, M., Hasson, M., Lennie, M., Reso, M., Groshev, M., Naumov, M., Lathi, M., Keneally, M., Seltzer, M. L., Valko, M., Restrepo, M., Patel, M., Vyatskov, M., Samvelyan, M., Clark, M., Macey, M., Wang. M., Hermoso, M. J., Metanat, M., Rastegari, M., Bansal, M., Santhanam, N., Parks, N., White, N., Bawa, N., Singhal, N., Egebo, N., Usunier, N., Laptev, N. P., Dong, N., Zhang, N., Cheng, N., Chernoguz, O., Hart, O., Salpekar, O., Kalinli, O., Kent, P., Parekh, P., Saab, P., Balaji, P., Rittner, P., Bontrager, P., Roux, P., Dollar, P., Zvyagina, P., Ratanchandani, P., Yuvraj. P., Liang, Q., Alao, R., Rodriguez, R., Ayub, R., Murthy, R., Nayami, R., Mitra, R., Li, R., Hogan, R., Battey, R., Wang, R., Maheswari, R., Howes, R., Rinott, R., Bondu, J. Datta, S. Chugh, S., Hunt, S., Dhillon, S., Sidorov, S., Pan, S., Verma, S. Yamamoto, S., Ramaswamy, S., Lindsay, S., Lindsay, S., Feng, S., Lin, S., Zha, S. C., Shankar, S., Zhang, S., Zhang, S., Wang. S., Agarwal, S., Sajuyigbe, S., Chintala, S., Max, S., Chen, S., Kehoe, S., Satter- field, S., Govindaprasad, S., Gupta, S., Cho, S., Virk S., Subramanian, S., Choudhury, S., Goldman, S., Remez, T., Glaser, T., Best, T., Kohler, T., Robinson, T. Li, T., Zhang, T., Matthews, T., Chou, T., Shaked, T., Vontimitta, V., Ajayi, V., Montanez, V., Mohan, V., Kumar, V. S., Mangla, V., Albiero, V., Ionescu, V., Poenaru, V., Mi- haileseu, V. T. Ivanov, V., Li, W., Wang, W., Jiang, W., Bouaziz, W., Constable, W., Tang, X., Wang, X., Wu, X., Wang, X., Xia, X., Wu, X., Gao, X., Chen, Y., Hu, Y., Jia, Y., Qi, Y., Li, Y., Zhang, Y., Zhang, Y., Adi, Y., Nam, Y., Yu, Wang, Hao, Y., Qian, Y., He, Y., Rait, Z., De Vito, Z., Rosnbrick, Z., Wen, Z., Yang, Z., and Zhao, Z. The lama 3 herd of models. arXiv preprint, 2024. doi: 11t::00.8850//xxx.240.2222222
Gaur, Y., Babaei, Y., Wen, Y., Song, Y., Zhang, Y., Li, Y.,
Mao, Y. Coudert, Z. D., Yan, Z., Chen, Z., Papakipos. Z., Singh, A., Grattafiori, A., Jain, A., Kelsey, A., Shajnfeld, A., Gangidi, A., Victoria, A., Goldstand, A., Menon, A., Sharma, A., Boesenberg, A., Vaughan, A., Baevski, A., Feinstein, A., Kallet, A., Sangani, A., Yunus, A., Lupu, A., Alvarado, A., Caples, A., Gu, A., Ho, A., Poulton, A., Ryan, A., Ramehandani, A., Franco, A., Saraf, A., Chowdhury, A., Gabriel, A., Bharambe, A., Eisemman, A., Yazdan, A., James, B., Maurer, B., Leonhardi, B., Huang, B., Loyd, B., Paola, B. D., Paranjape, B., Liu, B., Wu, B., Ni, B., Hancock, B., Wasti, B., Spence, B., Stojkovic, B., Gamido, B., Montalvo, B., Parker, C., Burton, C., Mejia, C., Wang, C., Kim, C., Zhou, C., Hu, C., Chu, C.-H., Cai, C., Tindal, C., Feichtenhofer, C., Civin, D., Beaty, D., Kreymer, D., Li, D., Wyatt, D., Adkins, D., Xu, D., Testuggine, D., David, D., Parikh, D., Liskovich, D., Foss, D., Wang, D., Le, D.,, Holland, D., Dowling, E., Jamil, E., Montgomery, E., Presani, E., Hahn, E., Wood, E.,
Elbayad, M., Gu, J., Grave, E., and Auli, M. Depth- aaahh-ddatiiee transformer. In Proc. International Conference on Learn-

---

Attention-Level Speculation
ing Representations (ICLR), 2020. URL https:// wttpeviwwwweeeettueeeeeeem ftw.oooommmmiddggg1 -
Huang, Y., Cheng, Y., Bapna, A., Firat, O., Chen, D., Chen, M., Lee, H., Ngiam, J., Le, Q. V., Wu, Y., and Chen, GPipe: Efficient training of giant neural networks using pipeline parallelism. In Proc. Conference on Neural Information Processing Systems (NewrlPS), volume 32, 2019. URL
Elhoushi, M., Shrivastava, A., Liskovich, D., Hosmer, B., Wasti, B., Lai, I.., Mahmoud, A., Acun, B., Agarwal, S., Roman, A., Aly, A., Chen, B., and Wu, C.-J. Layer- Skip: Enabling early exit inference and self self-speculativ decoding. In Proc. Anmual Meeting of the Association for Computational Linguistics (Volume I: Long Papers), pp- 12622-12642, 2024. doi: 10. 18653/v /0/553/v1/2224..0 http://.22l./ooooo//o acl.....cl-ooo665.
CC cc/pppper_iiileeppper//////1//11 neurips. 03366500802295580761122222222222 pdf.
Jeffrey, M. C., Subramanian, S., Yan, C.,, Emer, J., and Sanchez, D. A scalable architecture for ordered paral- lelism. In Proc. International Symposium Microar- chitecture (MICRO), pp. 228-241. IEEE/ACM, 2015. doi: 11o.00.111//230772288877777
Eyuboglu, S., Zinsley, D., Saad-Falcon, J., Arora, S., Rudra, A., Zou, J., and Re, C. Towards smaller language models via layer looping, In Proc. Workshop on Effi- cient Systems for Foundation Models II @ ICML2024, 2024. URL
id=2N3CtUdoFD.
Jeffrey, M. C.,, Ying, V. A., Subramanian, S., Lee, H. R., Emer, J., and Sanchez, D. Harmonizing speculative and -pecuuluu execution in architectures for ordered par- allelism. In Proc. International Symposium on Microar- chitecture (MICRO), pp- 217-230. IEEE/ACM, 2018. doi: 11//////...00...000
Gaiser, J., Fontaine, D., Hoffman, H., Jones, S., and Oh, F. Dynamic control flow in CUDA graphs with conditional nodes. NVIDIA Developer Technical Blog, May 2024.
Gao, I.., Tow, J., Abbasi, B., Biderman, S., Black, S., DiPofi, A., Foster, C., Golding, L.., Hsu, J., Le Noac"h, hh, A., Li, H., McDonell, K., Muennighoff, N., Ociepa, C., Phang, J., Reynolds, L., Schoelkopf, H.. Skowron, A., Sutawika, L., Tang, E., Thite, A., Wang, B., Wang, K., and Zou, The language model evaluation harness, July 2024.
Jiang, A. Q. Sablayrolles, A., Roux, A., Mensch, A.. Savary, B., Bamford, C., Chaplot, D. S., de Las Casas, D., Hanna, E. B., Bressand, F., Lengyel, G., Bour, G., Lample, G., Lavaud, L. R., Saulnier, L.,, Lachaux, M.. Stock, P., Subramanian, S., Yang. S., Antoniak, S., Scao, T. L., Gervet, T., Lavril, T., Wang, T., Lacroix, T., and Sayed, W. E. Mixtral of experts. arXiv preprint, 2024. doi: 06:11048550/ARXII.... 104088.
Geshkovski, B., Letrouit, C., Polyanskiy, Y., and Rigollet, P. The emergence of clusters in self-attention dynam- ics. In Proc. Conference on Neural Information Process- ing Systems (NeurIPS), volume 36, pp. 57026-57037, 2023. URL htttps://openreeee http:///ppereeeeeeeeeeeeee. net forum? id=aMjaEkkXJJ
Katharopoulos, A., Vyas, A., Pappas. N., and Fleuret, F. Transformers are RNNs: Fast autoregressive transformers with linear attention. In Proc. International Conference on Machine Learning (ICML), volume 119, pp. 5156- 5165. PMLR, 2020. URL ttpptttppp//prrreddddds mir www.ppess/llll httpp/ee//lll//aaaaarorrooooorre htt/kkkaaaappooooso0000 htmml.
Gu, A. and Dao, T. Linear-time sequence modeling with selective state spaces. In Proc. Conference on Language Modeling (COLM), 2024. URL. https:
/ wwopnneeeee..eee.... wwtp/////rfii/eeeeeeeeeeeeeeteee -.t.rmmmm..--EEx.aW22
Kim, H., Papamakarios, G., and Mnih, A. The Lipschitz stant of self-attention. Proc. International Conference on Machine Learning (ICML), volume 139, pp. 5562- 5571. PMLR, 2021. URL fttetts:/prrcceeeeee mir www.ppressseee vv19//11////////// kimk1imm111...............mmm.m1 1.hmml
Hao, S., Sukhhaatar, S., Su, D., Li, X., Hu, Z., Weston, J., and Tian, Y. Training large language models to rea- son continuous latent space. arXiv preprint, 2024. doi: 040104485500/Xi122222.066666
Herlihy, M. and Moss, E. B. Transactional mem- ory: architectural support for lock-free data struc- tures. In Proc. International Symposium Computer Architecture (ISCA), pp. 300. ACM/IEEE, 1993. doi: 14:10..114//165123165166
Kim, S., Hooper, C., Wattanawong, T., Kang, M., Yan, R., Gene, H., Dinh, G., Huang, Q., Keutzer, K., Ma- honey, M. W., Shao, Y. S., and Gholami, A. Full stack optimization of transformer inference: a survey, In Proc. Workshop on Architecture and System Support for Transformer Models (ASSYST), 2023. URL https: / //openreview.oo net http://ooroummm.dtttootoollaua
Hooper, C., Kim, S., Mohammadzadeh, H., Gene, H., Keutzer, K., Gholami, A., and Shao, S. SPEED: Speculative pipelined execution for efficient decod- ing. In Proc. NeurIPS Workshop on Efficient Natu- ral Language and Speech Processing (ENLSP), 2023. 10. X0w.44010//./...31..1111
Kitaev, N., Kaiser, L., and Levskaya, A. Reformer: The effi- cient transformer. In Proc. International Conference on Learning Representations (ICLR), 2020. URL https: / //openrevice ew, net mettfform?? d=rkg
12

---

Attention-Level Speculation
Leviathan, Y., Kalman, M., and Matias, Y. Fast in- ference from transformers via speculative decoding. In Proc. International Conference Machine Learn- ing (ICML), volume 202, pp. 19274-19286. PMLR, 2023.
Rau, B. R. and Fisher, A. Instruction-level par- allel processing: History, overview, and perspec- tive. The Journal of Supercomputing, 7(1):9-50, 1993. doi: 10..1000/FF01205181
iigs..iiiig... .1..mmmr... .ww..prss/////// html.
#202/Leviathne
Sazeides, Y. and Smith, J. The predictabilitty of data val- ues. In Proc. International Symposium on Microar- chitecture (MICRO), pp. 248-258. IEEE/ACM, 1997.
Li, Y., Wei, F., Zhang, C., and Zhang, H. EAGLE: Spec- ulative sampling requires rethinking feature uncertainty. In Proc. International Conference on Machine Learn- ing (ICML.) volume 235, pp. 28935-28948. PMIR, 2024. URL https: pproceedings.oo mlr. press/ v235 1 html.
db::10.1100/M1CR0119977 111199
Lipschitz regularity deep neural networks: analysis and efficient estimation. In Proc. Conference on Neural Infor- mation Processing Systems (NeurIPS), volume 31, 2018. URL https::/prroceeddngg neurips. cc hct/ppppee.fiille//pper///////// /ttp////o/////////e///////////// ab439aacc7766e3366655224442244-3 pdf.
Scaman, K. and Virmaux, A.
Lipasti, M. and Shen, J. Exceeding the dataflow limit via value prediction. Proc. International Symposium on Mieroarchitecture (MICRO), pp. 226-237. IEEE/ACM, 1996. doi: 10::10.11 109.11190////RR 1996.566464.
Miao, X., Oliaro, G., Zhang, Z., Cheng, X., Wang, Z., Zhang, Z., Wong. R. Y. Y., Zhu, A., Yang, L., Shi, X. Shi, C., Chen, Z., Arfeen, D., Abhyankar, R., and Jia, 7. SpecInfer: Accelerating large language model serving with tree-based speculative inference and ver- ification. In Proc. International Conference on Archi- tectural Support for Programming Languages and Op- erating Systems (ASPLOS), pp. 932-949. ACM, 2024. doi: 011.01114//32266666666 33511355.6661133
Sergeev, A. and Balso, M. D. Horovod: fast and easy distributed deep learning in TensorFlow. arXiv preprint, 2018. doi: 10.455500//ariii. 1802.05799.
Shoeybi, M., Patwary, M., Puri, R., LeGresley, P., Casper, J., and Catanzaro, B. Megatron-LMM: Training multi-billion parameter language models using model parallelism. arXiv preprint, 2019. dbi:10.48550/ariiv 10t.0i..999900..0....
Sohi, G. S., Breach, S. E., and Vijaykumar, T. N. Multi- scalar processors. In Proc. International Symposium Computer Architecture (ISCA), pp. 414-425, ACM/IEEE, 1995. doi: 10. 111114/29999222444511.
Naor, A. On the Sanach-space-valued azuma inequality and small-set isoperimetry of alon roichmon graphs. Com- binatorics, Probability and Computing, 21 2144):23-634 2012. 11..1017//09635483111000 1000757.
Narayanan, D., Shoeybi, M., Casper, J., LeGresley, P., Pat- wary, M., Korthikanti, V., Vainbrand, D., Kashinkunti, P., Bernauer, J., Catanzaro, B., Phanishayee, A., and Za- haria, M. Efficient large-scale language model training on GPU clusters using Megatron-LM. In Proc. of the International Conference for High Performance Comput- ing, Networking, Storage and Analysis (SC). ACM, 2021.
Spector, B. F. and Re, C. Accelerating LLM inference with staged speculative decoding. In Workshop on Ef- ficient Systems for Foundation Models @ ICML2023, 2023. URL https://////c. http///ppeeeeeeeeeeeeeeeeeeeeeee nwt.:frumt?f? i id=RKHF3VYJJK.
/
Sun, H., Chen, Z., Yang, X., Tian, Y., and Chen, B. Triforce: Lossless acceleration of long sequence generation with hierarchical speculative decoding. In Proc. Conference Language Modeling (COLM), 2024. URL https: / //opreevview.... net.
17.3666099
Pagliardini, M., Paliotta, D., Jaggi, M., and Fleuret, F. Fast attention over long sequences with dynamic sparse flash attention. In Proc. Conference on Neural Information Processing Systems (NeurIPS), volume 36, pp. 59808- 59831, 2023. URL https: http:///ppeeeevieeee / forum?id Darumpid-OONNuueeWt.
Tay, Y., Dehghani, M., Bahri, D., and Metzler, D. Efficient transformers: A survey. ACM Computing Surveys, 55(6), 2022. doi: 10.11..114/333011
Poli, M., Massaroli, S., Nguyen, E., Fu, D. Y., Dao, T., Baceus, S., Bengio, Y., Ermon, S., and Re, C. Hyena hierarchy: Towards larger convolutional language mod- cls. In Proc. International Conference on Machine Learning (ICML), volume 202, pp. 28043-2808 PMLR, 2023. URL https: http:ttppppppeeeeeeeeeeeeee /pppcceediiggsssee/ee mlr. .ress// v2 0202/po1122 1o0ll0112444 12aa html.
Tenstorrent. Llama3 models. GitHub Repository, 2024a. URI. tt ittmetaaiblobbvvd 55 0 re133 3/1ode3mm//////// demos mamnos//lamaa/RBaa md.
Tenstorrent. tenstorrent.c http:/rrrtt..cccoooaattteett/tee
TT-Metallium, 2024b. URL https://
13

---

Attention-Level Speculation
Tenstorrent. T7-Metaliumth Guide, 2024c. URL https:
Zheng, L., Yin, L., Xie, Z.., Sun, C., Huang, J., Yu, C. H., Cao, S., Kozyrakis, C., Stoica, I., Gonzalez, J. E., Barrett, C., and Sheng, Y. SGLang: Efficient execution of struc- tured language model programs. In Proc. Conference on Neural Information Processing Systems (NeurIPS), 2024. URL https:: https://ppeneeview net/ forum? id=vgkAKiibpq
com ///ggthtb..cmm/eestttrrreee//// http:///tnsttoreettteteeeeee blob/v0. 55. rc13/ ro13/MSTAIUN_GUDDE.. md.
Thakur, R., Rabenseifner, R., and Gropp, W. Opti- mization collective communication operations in MPICH. The International Journal of High Perfor- mance Computing Applications, 19(1):49-66, 2005. 111101177700444200555552222
Touvron, H., Lavril, T., Izacard, G., Martinet, X., Lachaux, M.-A., Lacroix, T., Rozière, B., Goyal, N., Ham- bro, E., Azhar, F., Rodriguez, A., Joulin, A., Grave, E., and Lample, G. Llama: Open and efficient foundation language models. arXiv preprint, 2023. doi: (6.:104485500axXi2332..19999
Vasiljevic, J., Bajic, L., Capalija, D.,, Sokorac, S., Ignjatovic, D., Bajic, L., Trajkovic, M., Hamer, I., Matosevic, I., Cejkov, A., Aydonat, U., Zhou, T., Gilani, Z., Paiva, A., Chu, J., Maksimovic, D., Chin, S. A., Moudallal, Z.,, Rakhmati, A., Nijjar, S., Bhullar, A., Drazic, B., Lee, C.. Sun, J., Kwong, K.-M., Connolly, J., Dooley, M., Farooq, H., Chen, J. T., Walker, M., Dabiri, K., Mabee, K., Lal, R. S., Rajatheva, N., Retnamma, R., Karodi, S., Rosen, D., Munoz, E., Lewycky, A., Knezevic, A., Kim, R., Rui, A., Drouillard, A., and Thompson, D. Compute substrate for software 2.0. IEEE Micro, (12)550--5, 2021. doi: 10. 1..0.119900
Xiao, G., Tian, Y., Chen, B., Han, S., and Lewis, M. Ef- ficient streaming language models with attention sinks. In Proc. International Conference on Learning Represen- tations (ICLR), 2024. URL https: openreview.e.. net/ fett/formm?e NN=NNGG75111
Yang, A., Yang, B., Hui. B., Zheng, B., Yu, B., Zhou, C., Li, C., Li, C., Liu, D., Huang, F., Dong, G., Wei, H., Lin, H., Tang, J., Wang, J., Yang, J., Tu, J., Zhang. J., Ma, J., Yang, J.. Xu, J., Zhou, J., Bai, J., He, J., Lin, J., Dang, K., Lu, K., Chen, K., Yang, K., Li, M., Xue, M., Ni, N., Zhang, P., Wang, P., Peng, R., Men, R., Gao, R., Lin, R., Wang. S., Bai, S., Tan, S., Zhu, T., Li, T., Liu, T., Ge, W. Deng, X., Zhou, X.. Ren, X., Zhang. X., Wei,
X.. Ren, X., Liu, X., Fan, Y., Yao. Y., Zhang, Y., Wan, Y., Chu, Y., Liu, Y., Cui, Z., Zhang, Z., Guo, Z., and Fan, 7. Qwen2 technical report. arXiv preprint, 2024. 14011044850//axx..440...
Ye. Z. Chen, L., Lai, R., Lin, W., Zhang, Y., Wang, S., Chen, T., Kasikci, B., Grover, V., Krishnamurthy, A., and Ceze, L. Flashinfer: Efficient and customizable attention engine for Ilm inference serving. In Proc. Conference on Machine Learning and Systems (MLSys), 2025. URL https:/ http///pppreeieeeeeeeeeeeeeeeeee wtwp:/evvieenneeeeorroeeeee nerumm 1d=RXPofAshh
14

---

Attention-Level Speculation
A. A Proof for Verification Error Bound
Error bound on speculation steps is essential to rigorously quantify the trade-off's between efficiency and accuracy. Such bounds ensure that the simplified model maintains acceptable performance levels while providing theoretical guarantees on the maximum potential loss in accuracy.
We accept the speculated attention output when its value is close enough to the ground truth (i.e., within a threshold ô) in some metric space. To find this appropriate metric space and its corresponding 6, however, requires search in an almost infinite space if done empirically.
Instead, we frame the problem as follows:
On a model with M speculation points, given some metric d. for c >> O. for each speculation point E 1...N, find threshold d, such that for every ground truth hidden state H; and speculated hidden state Hl, if d(H,, H.) for all i, then the distance between the final output d(HN, Hx) <
To derive the relationship between 5 and c, we apply the Lipschitz regularity of neural network (Scaman & Virmaux, 2018). We consider the Llama (Touvron et al., 2023) architecture, consists of layer normalization, self-attention, feed-forward neural networks, and residual connections. When the self-attention output A, is being speculated as A,, a verifier with threshold 5, g guarantees d( (A,,, Á.) < 6,, otherwise the speculation would be rejected.
A.1. Worst Case Error Bound
Since we are interested in bounding the error in IR which is a Banach space, we denoted the distance function d(A, B) as | Bl p that the p-norm of TRd. Note that the norm induced metrics are translation invariant in Banach spaces:
| Y|l = (I + A) (Y + A)II, VX, YX R
Those metrics also follow triangular inequality:
| Y |X 2| < X 1 AX, Y, 7 e Rd
Follow the same assumption as in Section 3.2, suppose that for each layer, the self-attention output speculation error is bounded by ||A,- - -A,ll < 64. Then we can easily obtain the bound for residual connections:
[I, A,) (H, A = ||4 - - < 6;
Following the residual connection of self-attention is the LayerNorm and MLP. Since both LayerNorm and MLP are Lipschitz continuous, their composite is also Lipschizt continuous.
|LP(MM NormMLP(A.)) aallA, - A
follows
(A: + Norm.LP((.)) (A:+ NormMLP(As)
(A, -A)) (NormMLP(AJ) NormMLP(A.)
< (A. A + ( +(((NoMMMP(AA)) NormMIP((A))
<8, + ag,
-(1 + a;)6;
At the next speculation point (the self-attention in layer i 1) the non-speculative and speculative inputs are H4+1 and H4+1 respectively. We denote self-attention outputs A Å4+1. and Å as the non-speculative path output, speculative input followed by non-speculative output, and speculative input followed by speculative output respectively. We note that to obtain bound for d(A;41, Å we can once again apply triangular inequality of the metric space and the result from (Castin et al., 2024), which the Lipschitz constant of self-attention, f f(2),, is function of the input magnitude Rt:
|| Acsill
< ||AH1 - A ll + | Aml
$ f(R)]] - H 1l + 8:41
15

---

Attention-Level Speculation
In general, let H2 and Hj be the non-speculative and speculative input of layer and oz and 3 be Lipschitz constants of layer then:
|1 +1,+11l s (1 + where H2 H2ll - (1 + 01)61
-H,1())
(2)
Simplifying Equation by assuming the same upper bound a and / across all layers, we obtain the following relationship between {6,} and e: 6= [N Hvll S EN
E (1 + ao -1 + ((R)8)
A.2. Free choice of Lp and linear approximation of Equation
For any finite dimension Banach space, any two norms on the space are equivalent (#). Here we show a simplified proof for our norms Ly where S p S co. Let E IR" and < p < r < oc.
When x = 0, this obviously holds. Consider x + Then
|
ok 1////r
2 (max |x 1)) ) 1/r
ken
max ma REn
||l oo
Note that
|
Σn [x
(maxien lare))"
|ark ir
(maxten (x%]))"
|20.|]
(maNGen [
|20
(maxten (2;)) )2
|20 p
(maxien (xs))
a
Which gives
|
12
| | 1
|
||
|x
1//
|
|0
(
y 1/r
||=1||p
|
Note this holds for p. r > We can use the same argument to show that ||>||p S ||2111- Also,
n||l=||
n2
max [r
SEn
|
16

---

Attention-Level Speculation
This shows that | constant of that only depends the dimensionality but not the magnitude of the input vector.
| < | < ||||| < =||=e|ba. Using this fact, the choice of Ip would at most differ by
When a and 3 are small, we can approximate the error bound linearly. We have that
S 101 (1 + a)
N- :+1
((1 ((R)3) '
N-++1 ++1
Σ
--i+1
N-
Σ
NN- (N5 (F(R)O) 8,
j
j=0
N ++
Σ
NN -i)
a -
Σ
N-1
( (x)))
6:
j
j=0
N
(1+ (N - (1+ -N- -
Σ
i)( (N 6,
(N+ (a + 8(R)3)8
Which 2 6, is an overall bound. This bound holds for the worst case; all errors are in the same direction. However, this is highly unlike in practice and it has a dependence on N As a result, the threshold chosen using this bound for rejecting speculation is too small and not useful.
A.3. High Probability Error Bound
However, with stronger assumptions, we can obtain a better error bound. Assume that the errors in the speculation rounds are independent and follow some distribution with mean O. In other words, the direction of the error will not be biased towards any direction. By rejecting errors that exceeding some bound, we ensure that speculated results are bounded by the ball centered the ground truth ær, with radius X ||a=||- This is martingale:
E E(H++1 H.
With the causal softmax attention
H = Vsoftmax(QK3)
The approximation is obtained as
Ā Ā= QKT B)
And
f = V(B, sootmmax
Where B is a subset of columns of KT, With the error being bounded:
H H1+1l1a S ^|| 1|
With Azuma-type inequality for Banach space-valued martingales (Naor, 2012), there exists some universal constant c >>> 0 such that.
cE2
P(I - H 122 2 Ee exp
1= X2 H
Rearranging the equation, with probability & E (0,1), 1). we have
(4 log 8)12 EN A7
17

---

Attention-Level Speculation
Choose 8 0.05, with some addition assumption that |HH "on the same scale" due layer norm, we replace them with universal constant || 22. Now with probability 95%,
log N
JHN HNlle Sv
-NNll2 IHN
7N
<X1
|H|1
c
This bound has much better dependence on N which is the total number of speculations than the last bound. However, relies on unbiased assumption on speculation error that is yet to be proved. In practice, À = 0.1 worked well varies benchmarks and this bound much closer to 0.1.
18

---

Attention-Level Speculation
B. Speculative Flash Decode Algorithm
The speculative flash decode algorithm builds upon the standard Flash Attention 2 (Dao, 2024). In the scope of this work, we only consider the forward pass for inference application. During the decode stage, the activation query matrix (Q), with sequence length I, is very small and fits into the on-chip SRAM. The keys (K) and values (V) matices, particularly at large context length, are large and can only live in the slower DRAM. In flash decode, imput is replicated across parallel workers and each worker computes partial attention output (with local statistics) using a portion of K and V by running the standard flash attention algorithm with chunk size S. Next, the partial outputs are gathered and reduced on the reducer worker using the local statistics by applying the online softmax trick. In the extreme case of having only a single worker, flash decode is the same as flash attention. In the algorithm below, we skip the details of how flash attention is computed and flash attent i on (Q, K, V) to obtain the partial output O and statisties m
we simply call O, m,
Speculative flash decode builds upon the flash decode algorithm with few modifications. The kernel consists of the sender, which computes the full and speculative attention outputs, and the receiver, which waits for the speculative results to arrive. First, one (or two) of the workers compute the partial output using first and last chunk of K and V, which the result serves as the speculated attention output. Then, the worker continues to finish the rest of its own assigned K and V chunks. Finally, the reducer worker computes the full output using the partial outputs from each worker and the speculative output. We show speculative flash decode Algorithm below:
Algorithm 3 Speculative Flash Decode (forward pass for inference)
1: Require:
2: Toworkers total number of parallel workers
3: m: total number of chunks
4: Q: query of size 1 x (fits in on-chip SRAM)
5: (Ki....; } key chunks DRAM
6: (Vi...., ]: value chunks in DRAM
7: Routines:
8: Elash.attention(0. K] [V]) (0,m.))
9: send_to_recceiver() send combine())
10:
// SENDER (run on each worker)
12: chunks per worker [m/ 1m/nworkers)
13: assigned.chunks range of chunk indices for this worker, excluding 1 and m
14: if worker-id = then
15:
(O., (0,,,,,14 + flash.attentiion K,, Km,, [V,, Vnl) (Speculative output)
16: send to receiver(O.)
17:
serddtt..utpptt.wooker(...... (..
end if
19: (06, m() (0)) flash.ttteniion (Q, Kı |j assigned (V, assigned_chunkss)
20: senotto.ooupppt.woreee.... (1) (())
21: // OUTPUT WORKER (on sender)
22: Gather all partial outputs (0 (m,)))))))))))) from each worker
23: Gather speculative (O., m., l.) from worker
24: (O, m,,) +- wwwinneieeeeeeooo www.sffaaa...mmmaoae
(((0(), (10) Mlworkees U Uf(O,, mp;
25: return O
26: // RECEIVER
27: wait for speculative output O, to arrive from worker
19

---

Attention-Level Speculation
C. Tenstorrent N150 Chips Specification and Tenstorrent Metalium (TT-Metalium) Framework
C.1. Tenstorrent Overview
Tenstorrent specializes in sesh-pefrrmance computing solutions for artificial intelligence (AI) and machine learning (ML) workloads. Their primary focus on designing innovative hardware accelerators and software frameworks that enable efficient inference and training of deep learning models. By leveraging custom processor architectures optimized for Al workloads, Tenstorrent aims to push the boundaries of energy efficiency and computational throughput.
C.2. N150 Wormhole Chips
The N150 Wormhole chip an AI accelerator designed for High Performance Computing (HPC) applications such as high- throughput and low-latency machine learning inference. These chips are designed to be highly flexible and scalable. Each chip contains mesh of 72 Tensix TM compute cores connected via network-on- (NOCs), offering 262 Tera-FLOPs of FP8 performance. Moreover, with on-board ethernet, N150 chips can be interconnected in various topologies consisting of 2, 4, 8, and 32 N150 devices.
Grayskull chip and Tensix core
LPODR4
LPDOR4
LPODR4
LPODR4
Comprn CRanK
lensic Core
T T
T
T
NoC
#OUTTEn
T
T T
T
T
SPAM
T
T
T T
T
Tensix core
5 RISC-V processors control:
Grayskull: 120 Tensix cores
2 NoCs, Matrix Vector engines, SRAM
tenstorrent
Conticenaal
Figure 9. Architecture of GraySkull device, the predecessor to the Wormhole TM Image taken from (Tenstorrent, 2024c).
C.3. Tensix TM Core
The Tensix TM core is a crucial building block on the Wormhole device. It consists of 5 small RISC-V V processors (also known as Baby RISCVs), that handle compute, data movement, and dispatch of instructions. For compute, Matrix Engine (FPU) is used to perform various binary operations on small 32x32 matrices (known as tiles), such multipiiiiiii,innnn, dot-product, and other element-wise operations. Alongside, a Vector engine (SPU) is used for special operations such as GELU, exponential, and sqrt. For memory, each core has access 1.5MB of high- bandwidth SRAM (L1) as well as 12GB of GDDR6 off-chip DRAM. A Data movement engine is also available on each core, connected to 2 NOCs. By creating mesh of cores via these NOCs, the N150 provides access to a high capacity of 108 MB of SRAM, allowing for the Tensix TM cores to operate at "silicon peak" performance. To allow users to leverage the powerful capabilities of the N150 chip. Tenstorrent has developed a low-level kernel library: TT-Metalium (Tenstorrent, 2024c).
C.4. TT-Metalium
TT-Metalium, also known as TT-Metalium, is an open-source low-level kernel operations library that allows users to enable peehorerformecc machine learning applications on Tenstorrent hardware, such as the N150 chip. At a fundamental level, TT-Metalium assigns a 1:1 mapping between each Tensix TM core and a thread, discarding the need for complex thread
20

---

Attention-Level Speculation
scheduling. Using TT-Metalium, can leverage bare-metal programming write kernels in C/C++ that directly each Tensix and perform computation tilized input tensors. With the fine-gr tne-grained control enabled by the TT-Metalium framework, the mesh of Tensix cores on the N150 chip can be leveraged in robust and powerful manner, overall supporting highly intensive machine learning workloads.
C.5. Kernel programming with TT-Metalium
Programming Tensix TM core TT-Metalium mainly consists of writing three kernels reader, writer, and compute. There are other advanced types of kernels available, such as dispatch and ethernet data movement kernels, however, they are not relevant to the user when bringing up general operations, and can be accessed via an abstraction layer. To implement an operation, the reader kernel is programmed read data from L1 or DRAM memory and pass it off the compute kernel via local buffers LI. After the compute kernel has completed the operation (using the FPU/SFU), it packs out the data into another local buffer, where then accessed by the writer kernel. The writer kernel is responsible for formatting the data required and providing outputs either in I.1 DRAM. As these three kernels in parallel, there native support for overlapping data movement with computation, where the reader/writer kernel can transfer data while the compute kernel is busy. It is also important to note that the compute kernel only supports tilized inputs to maximize utilization, and as such, the reader/writer kernels must format input tensors in this manner. Moreover, set of data movement APIs, that leverage the NOC, are available inside each reader/writer kernel, allowing each core to access the LI memory of any other core. Putting all of this together, TT-Metalium can be used create efficient data patterns across the N150 chip, resulting powerful implementations of operations used in machine learning.
Data Data-movement and compute kernels running Iensix Core
Compute
Data movement
Ten x Core
RISC-V 2
RISC-V 3
RISC-V
RISC-V
Data Into
Data the
ROUTER
Matrix & Vector engines
ROUTER
NoC
NoC
for
SRAM
for.
tenstorrent
Figure 10. The workflow of a Tensix TM core. Image taken from (Tenstorrent, 2024c).
C.6. TT-NN and Execution Paradigm
Powered by TT-Metalium in the background, TT-NN user-friendly library implemented in Python that enables fast bring-up of machine learning models without sacrificing performance. It provides easy-to-use APIs to perform machine learning operations on N-dimensional tensors, similarly to PyTorch. Additionally, TT-Metalium objects such as a Multi- Device tensor are exposed in TT-NN, allowing users to easily implement algorithms and paradigms that span multiple N150 chips and leverage distributed computing. When running models implemented in TT-NN, operations are dispatched by the host and scheduled to run on Tenstorrent hardware via the PCIe link. Once results are ready on device, they are sent back to the host.
21

---

Attention-Level Speculation
D. Speculative Flash Decode Implementation on TT-Metalium
implement the speculative flash decode kernel on TT-Metalium based on the existing flash decode kernel implemen- tation (Cai, 2024). All kernels on TT-Metalium are programmeed directly on the Tensix cores take advantage of the parallelism and tile-based granularity offered by the hardware. Each N150 chips consists of 64 worker cores and math is done tiles with size 32 32. In the case of flash decode, parallelize the batch and key/value heads across the cores, tile.
and put the query heads (usually S 32) within
We illustrate this idea using Llama 3.1 8B's group query attention, which contains 32 query heads and 8 k/v heads, with head dimension 128. Assuming batch 1, during decode, query has sequence length of so the shape of Q is [S= 1, B 1, H = 32, D 128), and the shape of shape of K and are [B = 1, 8, S, = 128]. In this case, we parallelize the k/v heads across the 64 cores, with cores per head. Each of the cores then each takes chunk of K and V fractured in the sequence length dimension and run the speculative flash decode algorithm (Appendix B). We illustrate the implementation in Figure 11:
1. Speculative portion finishes
Sender Chip
Receiver Chip
Key Value Cache
s/w
F/w
CCL
CCL
First
Last chunk
chunk
s/w
F/w
w
W
Core Core/Work WWrk Assignment
s/w
r/w
Head
s/w
r/w
CCL
linter-dhip
commuticcation)
2. All worker finishes, do reduction step
Head
s/w
r/w
w
W
s/w
r/w
w
Semaphore INC
Head
s/w
r/
w
W
Write via Ethernet
s/w
r/w
w
Speculate
/worker
Reduce
worker
worker
s/w
r/w
w
w
3. Reducer do verification and update priority tensor
Figure 1.. Speculative flash decode implementation on TT-Metalium. Each k/v head is split across cores in the sequence length dimension. For each head, core computes speculative and partial attention output, core computes partial attention output and reduction, and the rest of the cores computes partial attention output only. When speculative results is finish, the worker core signals the CCL core to send result to the receiver chip via ethernet link. The rest follows the standard flash decode implementation, where partial outputs are gathered and reduced on the reduction core. In the end, we fused the verification step and update the priority vector.
Overall, speculative flash decode yields almost no overhead compared to standard flash decode in our performance measurement. The speculative portion computes towards the final output, and communication across chip done on separate core via ethernet. When parallelizing attention across multiple devices using tensor parallelism, we split the k/v heads across the chip so that each head gets more workers per head. For example, splitting onto chips gives 4 k/v heads per chip, and split across 8 chips gives k/v head per chip. In practice, parallelizing across the heads yields diminishing returns when each chip receives "too little" work. In the case of having all 64 cores computing a single head, the bottleneck is NOC rather than compute due to the NOC transactions at the reduction step. For this reason, we see diminishing returns when scaling to 8 devices doing full tensor parallelism, and a better alternative would be doing device tensor parallelism + ALSpec for scaling up to 8 devices. Figure 12 demonstrates speculative flash decode, sender & receiver side kernel durations, across context length for tensor parallel Llama 3.1 8B attention across 1,2,4 & 8 devices with chunk size (128, 256, 512).
22

---

Attention-Leve Speculation
Llama 3.1 8B Attention with Speculative Flash Decode
Chunk Size 128
Chunk Size
Chunk 512
chip sender
dhip receiver
2 chipe TP sender
chip sender
dhip receiver
2 chips TP sender
chip sender chip recelver
2 chips TP sender
chips
chips
chipa TP recaivrr
8 chipp TP sender
chipa TP receiver
8 chips TP sender
chipa TP receiver
chips TP sender
1.5
1.5
Labonoy (mn))
Latency (ms)
caboa Sonoooo (mm))
0.5
0.5
0304 www. www.://wwwwo.o.ooo..o..... www..r/waworteeeeeeeeo.......... 2o
Longhh
Longhh
Context Length
Figure 12. Speculative flash decode kernel duration Llama 3.1 shape chunk size {128, 256, 512).
23

---

Attention-Level Speculation
E. SGDC with Priority Gating Implementation on TT-Metalium
The static graph dynamic concurrency (SGDC) paradigm, refers to having a static operation graph generated by the host but having the devices execute concurrently based on function, which in our case is the priority gating mechanism. For the speculative flash decode algorithm, the dynamic concurrency comes in the form of % speculation acceplance and the resulting overlap between attention and post-attention operations, which varies layer-to-layer and token-to-token. To support this paradigm, we made change to the run-time behavior of the TT-Metalium framework.
We leverage the fact that any kernel for any operation (data movement, compute, ethernet) in TT-Metalium is launched through a small set of C++ kernel source code files. By passing a flag at the build stage of these kernel source code files, we can control the behaviour of the kernel launch stage. At the start model execution, we create multi-device priority vector 'p', where each device tensor contains the device's priority. This tensor will have the same lifetime as the entire model and used the gating mechanism. Then, pass the base memory address of this tensor to the kernel source code files using C++ define, SKIP -CMPPCTT When a kernel is launched, it first passes through the kernel source code, where the address specified by SKIP_COMPUS priority of the device is O, then the operation will be skipped. Note, the ability to use the SKIP .COMPUTE can be toggled by adding/removing the define from the build arguments of the kemne files. We make sure to remove the define before running the speculative_fuel watpaitivvellshdeeee decode kernel, so as to not skip this operation.
is used to check whether the kernel should be launched. As described in the paper, if the
Using this mechanism, we can enable the behaviour described by scenario in Figure When device is sender device and finds that the speculated result correct within certain threshold, it can completely skip the post-attention operations using the SKIP _COMPUTE flag, relying on the concurrent computation performed device 2. When viewed from the host, this results in a static operation graph, however at the device level, we have achieved dynamic concurrency.
Flash Decode + MLP vs Speculative Flash Decode + MLP
FD + MLP
SFD @ 100% + MLP
4
3
Time [ms]
0
1k
8k
16k
32k
64k
128k
Sequence Length
Figure 13. End-to-ent cod performance using Llama3.1 8b shapes for Flash decode followed by MLP, compared to Speculative Flash Decode 100% specualation followed by MLP. For FD one device is used and for SFD two devices used.
24

---

Attention-Level Speculation
F. Llama 3.1 8B Implementation on TT-Metalium and Performance Estimation
TT-Metalium supports native mixed precision computation at 4 different data formats on N150 chips: float32, Block Floating Point (bfloat) 16, bfloat8 b, and bfloat4 These data formats can be used interchangeably within the model execution to opportunistically speed up computation when full precision is not required. In Llama 3.1 8B's implementation, 4/8-bit feedforward weights and 8-bit precision KV cache has been used to reduce memory bottleneck, while operations sensitive to precision, such as rotary embedding, are performed in full or half precisions. To ensure correctness, op, module, and model level unit tests are performed to ensure output quality. Specifically, Pearson Correlation Coefficient (PCC) and top 1/5 scores has been used as common metrics in unit tests. Generally, quantization is applied to models only if the PCC score is above 0.99 and top 1/5 scores are above 0.9/0.99. The empirically verified strict thresholds ensure that the model benefits as much as can from the mixed-precision inference without hurting the output quality. Therefore, to ensure our performance estimation of ALSpec matches and benefits the existing Llama model implementations on TT-Metalium, we use the default mixed-precision models brought up on the repository (Tenstorrent, 2024a).
We estimate the pure device performance of the Llama 3.1 8B model during decode phase using the following method:
We measure the kernel latency (Device FW Duration) of the operations within single decoder layer, excluding the flash decode op. We call this per_layer_lase
We also measure the kernel latency of operations outside the decode (embedding and LM-head ops). We call this fixed latency.
We sweep the latency of flash decode and speculative flash decode op across context length, from 1K to 128K. We denote the latency (K)) and SF ((K) for the latency at context length K.
For speculative flash decode, we record both the sender (SFDs(K)) and receiver (SFDR(K)) kernel latency.
We repeat the measurement on 1, 2, 4 and 8 devices tensor parallelism.
For decode phase, the input is always the same shape, but only the KV cache gets longer as the context length increase. In the Llama 3 implementation, the only operation that changes latency with respect to context length is the flash decode operation. Therefore, with 32 layers in the Llama 3.1 8B model, we can obtain the token latency and tokens/s/user throughput (1/(token latency)) at context length K using the following formula for full tensor parallelism:
token latency = fixed latency + 32 (per Jayer fayer_latecyyy + FDK))
To estimate the token latency at a particular speculation hit rate r we use the following formula:
token latency = fixed latency + [32 r (sF(sDD(K)) +
per layer layer
+ 322 (1 r)] pper_layer.ateeccyy + SFDs(K))
In short, this formula takes the overlapped attention latency when speculation hits, and full layer latency when speculation misses. It serves as an upper bound for latency because we always round down for the number layers with speculation hits 6.
We present the full latency breakdown for per layer Jatency and fixed latency across 1, 2, 4 and 8 devices tensor parallel in Table 4, 5, 6, & 7. One can verify the estimation in Figure 7 based the table and the speculative flash decode latency in Figure 12. The latency breakdown also shows how diminishing returns happen with tensor parallelism, while compute operations such as matmul get lower latency when parallelized to more devices, CCL operation latency such as all-gather increases, making the overall latency reduction minimal.
$ 'Note that in our measurement, we take r as a multiple of 1/32 to ensure that full layers are doing speculation. 'This formula doesn't It account for the last attention speculation, which should overlap with the fixed latency (LM Head). We are aware of this limitation and we use the formula for its simplicity. In the amortized case of running multiple token generation, this approximation holds as the final attention can overlap with the pre-attention ops the next token.
25

---

Attention-Level Speculation
Table 4. Llama 3.1 8B latency break down on x N150 chip
Per Layer DEVICE FW Duration [ns]
Embedding & Lm Head
OP CODE
CODE
DEVICE FW Duration [ns]
LayerNorm
12454
Embeddings
5745
Matmul
109174
Embeddings
5639
ShartedTolnterleveed
4299
Transpose
13808
NLPCreateHeads] Decode
21048
Transpose
13512
RotaryEmbeddtingl Jama
4828
Slice
2091
RolaryEmbeddiigl..n lama
4827
Slice
2089
PagedUpdlateaccle
13166
SntrrleavedToSharreed
1694
PagedUpdateCachee
13049
InterteavedToSharee
1759
69853
Reshard
3386
Interlcaved
1701
LayerNorm
13243
3055
Matmul
546986
Matmul
76389
Shardedfollmtereaveed
9003
Binary
1823
Matmul
545955
LayerNorm
12487
ShardedTolnterleaved
8833
Matmul
141261
Matmul
545707
Matmul
141772
Shardeedollnterleavedd
8944
Binary
30204
Matmul
547068
Matmul
252556
ShardedTollntrreavve
9055
Binary
2654
Concat
29186
Total / Layer without flash decode
846747
Total LM Head
2313703
Table 5. Llama 3.1 8B latency break down on Tensor Parallel 2 x N1.50 chips.
Per Layer Context Length 1K
OP CODE
DEVICE FW Duration [ns]
AllGather
19567
LayerNorm
12471
Matmul
56690
Embedding Lm Head
ShardeeTolnterleaved
2541
OP CODE
DEVICE FW Duration [ns]
NLPCreateHeeads] Decode
10972
Embeddings
5549
RotaryEmbeddiing
4834
Embeddings
5767
Rotary Embeeddiegennaammmm
4820
Transpose
13570
PagedUpdateCache
8447
Transpose
13590
PagedU UpdateCache
8383
Slice
2094
64309
Slice
2222
mnerlaavedToShhrdde
1833
InterleavedToShardee
1748
NI W...ConaaHHaddDeeacc
2320
InterleavedToShardeed
1701
Matmul
41730
AllGather
20607
ReduceScatter
21842
LayerNorm
12931
Binary
1621
Matmul
546309
AllGather
20887
ShardeeFoolntereeavee
8779
LayerNorm
12486
Matmul
546729
Matmul
76942
ShardeddTolnterleaved
8992
Matmul
77430
Concat
17048
Binary
16543
Matmul
128695
Total LM Head
1207636
ReduceScatter
19108
Reshard
3233
Binary
1620
Total / Layer without flash decode
555015
26

---

Attention-Level Speculation
Table Llama 3.1 8B latency breakdown on Tensor Parallel 4 x N150 chips.
Per Layer
OP CODE
DEVICE FW Duration [ns]
AllGather
22378
LayerNorm
12485
Matmul
39022
Sharteedolnterleaved
2057
Embedding & Lm Head
MLPCreateHeadsDeeco
6269
CODE
DEVICE Duration [ns]
RotaryEmbeddiiggL..
4837
DotaryEmbeddiigee
4834
Embeddings
5622
Cache
6260
Embeddings
5767
PagedUpdate UpdateCache
6342
Transpose
13328
56024
Transpose
13585
InterleavedTooShareee
1694
Slice
2088
HLaDeonatteaadDeeodeee
1955
Slice
2193
Matmul
24505
AnteteevedTToShardee
1795
ReduceScatter
24753
InterlcaeddToShaceee
1703
Binary
1517
AllGather
24259
AllGather
22055
LayerNorm
12906
LayerNorm
13423
Matmul
546508
Matmul
39711
ShardedTolInerllawee
8775
Matmul
39791
Total LM Head
638529
Binary
15800
Matmul
68424
ReduceScatter
24764
Reshard
3212
Binary
1496
Total / Layer without flash decode
387584
Table 7. Llama 3.1 8B latency breakdown Tensor Parallel on 8 x N150 chip.
Per Layer
OP CODE
DEVICE FW Duration [ns]
AllGather
34069
LayerNorm
12507
Matmul
35745
Sharded
1620
NLPCreateHeadd
3976
Embedding & Lm Head
RotaryEmbedding
4839
OP CODE
DEVICE FW Duration [ns]
Rotary Embeedding
4837
Embeddings
5641
PugedUpdateCachhe
5007
Embeddings
5607
PagedU UpdateCache
5078
Transpose
13717
50096
Transpose
13680
Shhrrded
1708
Slice
2095
NL.CoonatHccadseecode
1774
Slice
2092
Reshard
2274
Interleave/ToShardeed
1707
AllGatherMatmul
88825
InterleavedTTohharde
1726
Reshard
2847
AllGather
43463
Binary
1464
LayerNorm
12909
AllGather
54386
Marmul
276107
LayerNorm
15771
ShardedTolIntereennn
5716
Matmul
22705
Matmul
22899
Total LM Head
384460
Binary
15804
Matmul
38486
ReduceScatter
59572
Reshard
2817
Binary
1441
Total Layer without flash decode
440451
27

---

Attention-Level Speculation
G. Section 2 Needle in Haystack Experiment Details
The tiny needle in a haystack experiment served as a motivational experiment rather than rigorous benchmark. conducted on custom dataset created with the following method:
First, we create large context using the text of the paper "Compute Substrate for Software 2.0" (Vasiljevic et al., 2021) and the first chapters of the book "A Tale of the Two Cities" by Charles Dickens. This makes the total context length around 16K.
Then, we generate 7 random MD5 hashes and insert at random places in the text.
To ensure the result is not biased towards hash strings, we make up 3 random facts and insert at random places in the text.
We then prompt the model with the following questions after the context:
Start with text "Answer the following questions if you can. If you don't know the answer, answer don't know."
Then, we put questions about fact 1, 2, and 3.
In the end, we ask the model to repeat hashes 1-7.
28


Reviews:
[{"review_id": "spr3mjUKMo", "reviewer": "Reviewer_LXts", "summary": "The paper \"Attention-Level Speculation\" introduces ALSpec, a novel parallelism paradigm designed to accelerate transformer-based LLM inference by overlapping self-attention computations with subsequent non-attention operations (e.g., feed-forward layers). Key contributions include the core idea of self-attention outputs using approximations (e.g., attention sink on first/last tokens) and verifies them in parallel. If the approximation error is within a threshold (controlled by hyperparameter $λ$), subsequent operations proceed early; otherwise, they fall back to exact attention results.\nAchieves 1.65× end-to-end decode latency reduction at 128K context length with 87.5% speculation hit rate, outperforming tensor/data parallelism scaling limits. Maintains baseline correctness on reasoning (GSM8K), math (MATH), and retrieval tasks by dynamically rejecting harmful approximations. Reduces attention latency overhead by 5× via overlapping computations, validated on Tenstorrent NPUs.\nKey Insight involve static approximations (e.g., fixed sparse attention) fail on tasks requiring global context (e.g., topic shifts in multi-step reasoning). ALSpec adapts layer- and token-specifically, accepting approximations in 50–90% of layers without quality loss.Combines with tensor parallelism, showing continued scaling where pure tensor parallelism plateaus (e.g., 8 devices).", "questions": "", "limitations": "", "rating": 4, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We are grateful for your detailed feedback which will greatly improve the quality of the work.\n\n## Absence of compositional reasoning and low-res. languages\n\nWe extend our eval to GSM8K in Swahili (low-res. language), HotpotQA, and RepoBench-P. Results are below.\n\n| Config\\Tasks | MGSM_CoT_Sw | HotpotQA | RepoBench-P |\n|---|---|---|---|\n| no_spec | 0.58 | 0.92 | 0.756 |\n| l@0.05  | 0.59/54% | 0.92/19%  | 0.74/33% |\n| l@0.10  | 0.58/78% | 0.92/43% | 0.77/60% |\n| l@0.15  | 0.59/86% | 0.92/57%  | 0.81/76% |\n| l@0.20  | 0.56/95% | 0.91/72%  | 0.76/85% |\n| l@0.25  | 0.52/97% | 0.91/80% | 0.74/92% |\n\nThese new results show that ALSpec maintains correctness on low-res. languages (e.g., Swahili) or long context multi-hop QA (HotpotQA) with high speculation hit rate.\n\nALSpec's attn sink applies positional bias to recent tokens to approx. full attn. The 128 window size should capture key contextual ideas regardless of language, while early tokens address the softmax off-by-one bias [Evan Miller, July 2023], not positional preference.\n\n## GPU/TPU generalization\n\nWe believe that ALSpec can reproduce latency improvements on GPUs/TPUs. We confirmed this with new experiments, serving 8B models on 4 vs. 8 H100s with TP using the SGLang framework and FlashInfer attn backend. Although we don't (yet) have an ALSpec implementation on GPU, this provides an estimation of the perf gain. The results are below:\n\n| Context Len | 4xH100 Attn Latency (us) | 4xH100 Non-Attn Latency (us) | 4xH100 TP Tok/s | 8xH100 TP Tok/s | TP Scaling | Proj. ALSpec @ 65% Hit Rate | ALSpec Scaling |\n|---|---|---|---|---|----|---|---|\n|32k| 29 | 100 | 214.2 | 231.3 | 8.0% | 246.0 | 14.8% |\n|64k| 49 | 100 | 191.4 | 209.7 | 9.6% | 237.8 | 24.2% |\n|128k| 63 | 101 | 169.3 | 184.8 | 9.2% | 217.5 | 28.5% |\n\nOur new estimation shows that ALSpec cuts latency on Llama 8B by 1.28x vs. 1.09x for full TP at context len 128K. In this case, the attn latency is only 63% of non-attn latency--- meaning more gains for context len beyond 128K. Although ALSpec introduces a new attn kernel and SGDC, it does not fundamentally change the Op by Op and static graph execution style on modern platforms. Therefore, we believe ALSpec is implementable on GPUs/TPUs by expert kernel writers.\n\n## Lipschitz continuity analysis assumptions\n\nHua et al. (2023) can be a game changer to our induction method, and we would be grateful if LXts could point us to the exact paper title, as we weren't able to find it. The $\\alpha$ and $\\beta$ are not assumed to be a universal bound---they are some upper bounds of all Lipschitz constants across all layers and used to simplify the expression.\n\nRegarding mean-0 approximation error, we have the approximation to be\n$\n\\tilde{A} = QK^T(:,B)\n$\nwhere $B$ is a subset of columns of $K^T$ and\n$\n\\tilde{H} = V(B, :)\\text{softmax}(\\tilde{A}).\n$\n\nUnder some condition of distribution of cols of $V$ and rows of $K^T$, we might be able to show\n$\n\\mathbb{E}_{\\pi_V, \\pi_A}\\left[\\tilde{H}\\right] = H\n$\n\nFull independence is not required for Azuma-style concentration result, but we agree that we are yet to show that conditional expectation holds in martingale conditions. The boundedness of the error is a direct consequence of the algorithm; we abort the thread if the error exceeds the threshold.\n\n\n## Host-device communication overhead\n\nModern computing frameworks capture (trace) operations that are going to be executed on device, ahead of time. This trace can then be executed on a device (e.g. GPU), without any interaction with the host. Since all ops in ALSpec run fully on device, we execute it using trace. Hence, we analyze only the device duration to present findings that are agnostic to the host, while maintaining realism and feasibility.\n\n## SGDC Sync Overhead\n\nThe SGDC mechanism ensures that each device can compare its priority against its pair, thereby determining its role in SFD (ie. sender/receiver). Thus, the sync is an all-gather on 2 devices, where each device collects priority from its pair. GPUs/NPUs/TPUs support direct p2p communication via a topology (e.g. mesh), where collective ops (e.g. all-gather) can be implemented without host overhead. Considering (i) the small size of the priority tensor and (ii) that the cost of syncing is constant as seq len scales, we conclude that SGDC sync is a low-cost mechanism that does not inflate real-world gains from SFD.\n\n## Token-level pipelined speculation\n\nA key advantage of ALSpec is that it leaves the model unchanged (no fine tuning). SPEED require retraining for weight sharing. ALSpec addresses long context len issue through attn overlap. SPEED's 1.8x latency reduction is at small context len while ALSpec's 1.65x latency reduction is at context len 128K. SPEED and ALSpec are optimized for different scenarios and could complement each other when used together.\n\n## Additional literature\n\nThank you for pointing us to these related work. We will cite, compare, and contrast ALSpec with them in our final version.", "claims_and_evidence": "1. The 5× attention latency reduction and 1.65× end-to-end decode latency improvement at 128K context length are validated through empirical benchmarks on Tenstorrent N150 chips (Tables 1–2, Figure 8). The scalability analysis demonstrates diminishing returns for pure tensor parallelism, while ALSpec + tensor parallelism continues scaling.\n\n2. Dynamic verification (via L2-norm thresholds) maintains baseline accuracy on tasks like GSM8K, MATH, and MMLU PRO (Figure 2). Static approximations (e.g., fixed sparse attention) fail on reasoning tasks, while ALSpec selectively accepts approximations with a hit rate of 50–90% per layer.\n\nProblematic Claims:\n\n1. The reported latency gains are tied to Tenstorrent’s NPU architecture and proprietary kernels (e.g., speculative flash decode). Without ablation studies on GPUs/TPUs or open-sourced kernels, reproducibility is unclear. Host dispatch overheads (e.g., CPU-to-NPU communication) are excluded from latency measurements, potentially inflating real-world gains. \n\n2. The Lipschitz continuity analysis (Appendix A) assumes independent, zero-mean speculation errors. Real-world error propagation may violate these assumptions, risking unbounded deviations in practice. The L2 verification threshold (λ) is empirically set without theoretical justification for its sufficiency across layers/task. \n\n3. While ALSpec preserves accuracy on retrieval and math tasks, its performance on compositional reasoning (e.g., multi-hop QA) or low-resource languages is untested. The \"needles in a haystack\" experiment (Figure 3) uses synthetic data, which may not reflect real-world long-context retrieval challenges", "methods_and_evaluation_criteria": "The methods and evaluation criteria demonstrate technical validity but exhibit notable limitations in scope and generalizability:\n\n1. The L2-norm threshold verification ($ \\| \\tilde{A}_i - A_i \\|_2 < \\lambda \\| A_i \\|_2 $) is empirically effective but lacks theoretical justification. While Lipschitz continuity analysis provides error bounds (Equation 1), it assumes:  \n1.1 Independent, zero-mean speculation errors  \n1.2 Constant layer-wise Lipschitz factors ($ \\alpha, \\beta $). These assumptions may not hold in practice, risking unbounded error propagation.\n\n2. Integrates attention sink approximation (first/last $ S $ tokens) with exact attention in a fused kernel, reducing overhead by 5×. However:  \n2.1Chunk size $ S $ is fixed (128–512) rather than adaptive to input.\n2.2 Prioritizing first/last KV cache chunks biases toward positional extremes, potentially harming mid-context retrieval.\n\n3.  Maintains static execution graphs while dynamically routing computations. While effective on Tenstorrent NPUs, host-device dispatch overhead (CPU-NPU communication) is excluded from latency metrics, inflating real-world gains.\n\n4:\n4.1. Results are confined to Tenstorrent N150 NPUs. Tensor parallelism’s diminishing returns (Figure 8) may differ on GPU/TPU architectures due to distinct communication patterns.\n4.2. Task Coverage:\n4.2.1 No evaluation on low-resource languages (e.g., Swahili, Bengali) despite claimed multilingual support\n4.2.2 Absence of compositional reasoning benchmarks (e.g., DROP, HotpotQA).", "theoretical_claims": "The authors provides directional guidance but lacks robustness guarantees for real-world deployments. Independent empirical validation of Lipschitz constants and error distributions is needed to trust the bounds.\n\n1. Lipschitz Continuity Error Bound : The derivation in Appendix A establishes an upper bound on output deviation: \n$$\n\\epsilon \\leq \\sum_{i=1}^N (1 + \\alpha)^{N-i+1} (1 + f(R)\\beta)^{N-i} \\delta_i\n$$  \nwhere $$\\delta_i = \\|\\tilde{A}_i - A_i\\|_2$$.  \n\nI find few issues with this:\n1.1. This assumes identical $\\alpha$ (feed-forward/LayerNorm Lipschitz) and $\\beta$ (attention Lipschitz) across all layers. Real transformer layers exhibit heterogeneous operations (e.g., early vs. late layers), violating this assumption.  \n1.2. For $N=32$ layers, coefficients grow as $(1 + \\alpha)^{32}$, making the bound practically vacuous unless $\\alpha \\ll 1$. The paper notes this but provides no empirical validation of $\\alpha < 0.01$.  \n1.3. The high-probability bound assumes speculation errors are independent and zero-mean. Real approximation errors (e.g., attention sink) exhibit structured biases (e.g., positional bias toward first/last tokens), violating independence.  \n\n2. The L2-norm threshold ($$\\lambda$$) is empirically set without theoretical justification. The paper claims:  \n$$\n\\|\\tilde{A}_i - A_i\\|_2 < \\lambda \\|A_i\\|_2 \\implies \\text{safe approximation}\n$$  \n\nIssues:\n1. 1. Layer-Wise Thresholding: Ignores error accumulation across layers. A per-layer threshold $\\lambda = 0.1$ could allow $\\epsilon = O(N\\lambda)$ deviation, violating final output fidelity.  \n2.2. $\\|A_i\\|_2$ varies significantly with context length and input entropy, making fixed $\\lambda$ suboptimal.  \n\n3.  The 87.5% speculation hit rate (Table 1) suggests the bounds are overly pessimistic, but no ablation study isolates the impact of violating assumptions like uniform $$\\alpha/\\beta$$.  The analysis doesn’t address the attention sink’s inherent positional bias – a systematic error source excluded from the zero-mean error assumption.", "experimental_designs_or_analyses": "The paper's experimental design demonstrates technical rigor but has critical limitations in scope and generalizability:\n\n1. Task Coverage in Correctness Evaluation  \n1.1 Benchmarks exclude multi-hop QA (DROP, HotpotQA) and low-resource languages.  \n1.2 \"Needles in haystack\" uses artificial key insertion (Fig 3), failing to capture real-world long-context QA patterns.  \n1.3 Compares only against attention sink, omitting dynamic variants of LSH (Reformer) or sliding windows (Longformer).  \n\n2. \n2.1 Attention sink prioritizes first/last tokens, violating the paper's assumption of zero-mean independent errors (Appendix A).  \n2.2 Lipschitz constants (α, β) are assumed uniform across layers, ignoring layer-specific dynamics (early vs late layers).  \n2.3 Fixed L2 threshold (λ=0.1) lacks theoretical justification for error propagation across layers.  \n2.4 Exponential error bound $$ \\epsilon \\leq \\sum_{i=1}^N (1+\\alpha)^{N-i+1}\\delta_i $$ becomes vacuous for N=32 unless α ≪ 1 (unvalidated empirically).", "relation_to_broader_scientific_literature": "ALSpec bridges cognitive theories (attentional resource allocation), neural evidence (predictive coding), and ML systems (speculative execution) to address transformer inference bottlenecks. Its dynamic, hardware-aware approach advances beyond static approximations and token-level speculation, offering a generalizable framework for adaptive computation in LLMs.", "essential_references_not_discussed": "The paper overlooks several critical areas of related work that contextualize its contributions:\n1. The paper cites Elhoushi et al. (2024) for layer pruning but misses recent advances in self-speculative decoding (Elhoushi §4.2) and adaptive computation time (ACT) transformers. For example:  \n1.1 SPEED (Hooper et al., 2024): Overlaps layer computations across devices via pipelined speculation, achieving 1.8× speedups on GPUs without attention approximation.  \n1.2 LayerLoop (Eyuboglu et al., 2024): Reuses layer outputs for computational savings, relevant to ALSpec's focus on overlapping FF/attention.  \n\n2.While ALSpec uses attention sink, it omits comparison to:  \n2.1 Blockwise Parallel Transformers (BPT) (Google, 2023): Dynamically adjusts sparse attention blocks using gradient-based importance scores.  \n2.2 FLASH (H2O, 2023): Hybrid sparse-dense attention with runtime pattern selection, achieving 2.1× speedups on 32K contexts.  \n\n3. The Lipschitz analysis lacks connection to:  \n3.1 Kim et al. (2021) Proved self-attention Lipschitz constants are unbounded without layer normalization, contradicting ALSpec's assumption of uniform $\\alpha$. [1]\n3.2 Zhu et al. (2024): Proves high probability excess risk bounds of $O(1/n^2)$ via algorithmic stability under strong convexity/smoothness. [2]\n3.3 Lei et al. (2023): Analyzes gradient stability for stochastic optimization but focuses on generalization bounds rather than error propagation rates. [3]\n\n4. ALSpec's Tenstorrent NPU focus ignores:  \n4.1 FlashDecoding++ (NVIDIA, 2023): Achieves 4.2× speedup over vanilla FlashAttention on 128K contexts via asynchronous state management.  \n4.2 vLLM (Berkeley, 2023): Paged attention for dynamic KV cache management, critical for real-world long-context deployments.  \n\n5. No discussion of subquadratic attention methods that reduce compute without approximation: \n5.1 Hyena (Poli et al., 2023): Replaces attention with implicitly parameterized convolutions.  \n5.2 Mamba (Gu & Dao, 2023): Selective state-space models for linear-time sequence modeling.  \n\nThese omissions weaken ALSpec's claims of novelty in:  \n- Dynamic execution (overshadowed by SPEED/LayerLoop)  \n- Error analysis (lacks modern Lipschitz bounds)  \n- Hardware generality (no GPU/TPU benchmarks vs FlashDecoding++/vLLM)  \n\nIncluding these would better position ALSpec within the broader landscape of efficient transformer inference.\n\n\n[1] Kim, Hyunjik, George Papamakarios, and Andriy Mnih. \"The lipschitz constant of self-attention.\" International Conference on Machine Learning. PMLR, 2021.\n[2] Zhu, Bowei, Shaojie Li, and Yong Liu. \"Stability and Sharper Risk Bounds with Convergence Rate $ O (1/n^ 2) $.\" arXiv preprint arXiv:2410.09766 (2024).\n[3] Lei, Yunwen. \"Stability and generalization of stochastic optimization with nonconvex and nonsmooth problems.\" The Thirty Sixth Annual Conference on Learning Theory. PMLR, 2023.", "other_strengths_and_weaknesses": "The paper \"Attention-Level Speculation\" demonstrates notable strengths in technical innovation and practical implementation but has limitations in theoretical robustness and generalizability. \n\n1. Strengths: \n1.1 Combines speculative execution (from CPU architecture) with transformer inference, extending token-level speculation (Leviathan et al., 2023) to attention layers. This hybrid approach is novel compared to static approximations like Reformer or Longformer.  \n1.2 Integrates attention sink (Xiao et al., 2024) with Tenstorrent NPU optimizations (speculative flash decode kernel, SGDC), advancing beyond GPU-centric methods like FlashAttention.  \n1.3 Demonstrates that deeper layers tolerate more approximation (Fig 4), aligning with pruning literature (Elhoushi et al., 2024) but adding dynamic verification.  \n\n2. Significance  \n2.1 Achieves 1.65× speedup at 128K context (Fig 8), addressing the critical bottleneck of long-context LLM inference.  \n2.2 Combines with tensor parallelism, circumventing its diminishing returns (e.g., 8 devices yield 60 tokens/s vs. 40 tokens/s for pure tensor parallelism at 128K context).  \n2.3 Validated on real hardware (Tenstorrent N150) with mixed-precision support, showing feasibility for deployment.  \n\n\n3 Theoretical Gaps:\n3.1 The error bound $\\epsilon \\leq \\sum_{i=1}^N (1+\\alpha)^{N-i+1} \\delta_i$ assumes uniform $$ \\alpha, \\beta $$ across layers, contradicting evidence of layer-wise dynamics (Kim et al., 2021).  \n3.2 Empirically set to 0.1 without theoretical justification for sufficiency across tasks/layers.  \n\n4. Fails to compare with recent dynamic methods like SPEED (Hooper et al., 2024) or Hyena (Poli et al., 2023), which offer alternative efficiency gains.   Provides a template for op-level speculation beyond attention (e.g., FFN layers), though this is not explored.", "other_comments_or_suggestions": "1. Page 3, Fig 1: Define \"ccl\" (collective communication ops) in the caption.  \n2. Page 5, Algorithm 1: Clarify \"ops before/after self attn\" (e.g., LayerNorm, residual adds).  \n3. Appendix A: Add intermediate steps between Equations 2 and 3 for readability.", "questions_for_authors": "1.Can ALSpec reproduce the reported latency gains (1.65× at 128K context) on GPUs/TPUs, particularly compared to FlashDecoding++ (NVIDIA) or vLLM? If not, what architectural features of Tenstorrent NPUs (e.g., NOC design, fused kernels) are indispensable for ALSpec’s gains?\n\n2. How do you empirically validate the assumption of uniform Lipschitz constants (α, β) across layers? Does layer-wise measurement of α (e.g., via power iteration) reveal significant variance, and if so, how does this affect the error bound in Eq. 1?\n\n3. Does ALSpec maintain correctness on low-resource languages (e.g., Swahili) or multi-hop QA (HotpotQA), where static approximations fail? If untested, could positional bias in attention sink harm non-English token distributions?\n\n4. What is the CPU↔NPU communication overhead for priority tensor synchronization in SGDC? Does excluding this from latency metrics inflate real-world gains (e.g., 5× attention reduction)?\n\n5. How does ALSpec compare to token-level pipelined speculation (Hooper et al., 2024) in terms of latency reduction per additional device? Does ALSpec outperform SPEED’s 1.8× GPU speedups when both use 8 devices?", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "yenCntFtkd", "reviewer": "Reviewer_GYW3", "summary": "The paper introduces attention-level speculative parallelism (ALSpec), a dynamic method for approximating self-attention in large language models. ALSpec computes an approximate attention output  and decides whether to accept the approximation. It uses a specialized flash decode kernel and SGDC with priority gating to overlap expensive attention and feed-forward operations. Experimental results on Llama 3.1 8B model using Tenstorrent NPU devices show significant reductions in latency (up to 5× in attention overhead at long contexts) and improved throughput (about 1.65× speedup at high speculation hit rates) compared to traditional tensor parallelism approaches.\n\n## update after rebuttal\nThe author's response addressed my concerns, leading me to raise my rating from 3 to 4.", "questions": "", "limitations": "", "rating": 4, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We are grateful for your detailed feedback which will greatly improve the quality of the work.\n\n# Impact on Larger Models\nWe performed additional experiments in terms of correctness analysis on the Llama 3.3 70B Instruct model. We showed that in terms of speculation hit rate and benchmark correctness across different lambdas, the Llama 70B model, despite deeper and larger, manifest the same behavior as the 8B model. Specifically, we benchmarked the 70B models on GPQA with CoT, GSM8K with CoT, Multi-lingual GSM with CoT in Swahili, and Hotpot QA. The results are shown below:\n\n| Config\\Tasks | GPQA_COT | GSM8K_COT | MGSM_COT_Swahili | Hotpot QA |\n|---|---|---|---|---|\n|Baseline|  0.518 | 0.958 | 0.852 | 0.940 |\n|lambda=0.05|  0.529/55.1% |  0.950/64.5% | 0.856/73.9% | 0.940/32.3% |\n|lambda=0.10|  0.507/80.0% | 0.951/86.6% | 0.852/90.1% | 0.945/59.0% |\n|lambda=0.15|  0.458/89.9% | 0.946/93.9% | 0.840/95.6% | 0.935/74.6% |\n|lambda=0.20|  0.446/95.1% | 0.936/96.8% | 0.820/98.1% | 0.935/83.3% |\n|lambda=0.25 |  0.379/97.6% | 0.897/98.3% | 0.816/99.1% | 0.935/89.4% |\n\nFor models in the sizes of 50 billion parameters such as Llama 70B, running the base model would usually require 8 devices with tensor parallelism (TP). As a result, using ALSpec would require at least 16 devices. To show ALSpec at the point where TP diminishes return (so that ALSpec provides real benefits for 70B model) would require even more devices. The resources required for these exepriments are beyond our current experiment setup. However, with the correctness analysis and our implemented kernels and methods, we believe that our approach is scalable to larger models when computing resources are available.\n\n# Additional Details in Experimental Setup for Kernel Fusion and SGDC\nWe thank the reviewer for pointing this out, and we will add more information in the camera-ready version, particularly in Appendix D and E, where we talk about the kernel fusion and SGDC in detail. For the reviewer's concern about fused vs unfused attention kernel, the speed up from the fused kernel is simply the difference of running the flash decode on the entire context length vs. running flash decode on entire context length + first and last chunk context. This is because the fused kernel changes the order of computation, allowing us to obtain the intermediate result of first and last chunk for free. \n\n# Impact of Batch Size\nWe focused on batch 1 in this paper, as usually during long context decode, the batch sizes are very small due to limited on-chip memory for the large KV cache. We will add more discussion on how batch sizes are handled in the camera ready version. In short, we have experimented with per batch threshold, where we only accept the speculation if all batches pass the lambda test. Alternatively, we also experimented with treating all batches as a single tensor and perform a single lambda test. The trade-off between these two methods are not within the scope of this paper, as the existing contents are already hard to fit within 8 pages. We plan to deep dive into this area in future work.\n\n# Why TP Fails to Scale\nAs the reviewer points out, communication overhead is one reason of TP's diminishing scalability. Another reason is kernels such as GEMM have a constant cost regardless of the shape. When a model's parameters are sharded across many devices, the constant part of the kernel latency dominates, and as a result, the overall scalability diminishes.\n\n# Comparisons to DP\nOur proposed method is to optimize for latency. The problem with DP, as depicted in Figure 1, is that despite it gives good scalability (always close to 2x throughput with 2x more devices), it only improves throughput. If the objective is to optimize throughput only, then DP is always a better choice than ALSpec and full TP. On the other hand, the latency per user is always lower for DP than ALSpec and full TP. Therefore, we believe that a fair comparison should be done between AlSpec and full TP rather than DP.", "claims_and_evidence": "The claims in this paper are supported by clear empirical observations and experiments.", "methods_and_evaluation_criteria": "Yes. The methods are well-aligned with the goal of reducing inference latency in large language models while preserving output quality.", "theoretical_claims": "I did not verify the correctness of the proofs, as the determination of appropriate speculation verification threshold values is primarily based on experimental results rather than theoretical derivations.", "experimental_designs_or_analyses": "I am interested in understanding how the two primary optimization techniques implemented in this paper—namely the Speculative Flash Decode Kernel and SGDC—impact the inference latency and throughput.", "relation_to_broader_scientific_literature": "The paper extends previous work on fixed attention approximations (Attention Sink) and token-level speculative decoding by introducing dynamic, attention-level speculative parallelism (ALSpec). It verifies approximations on the fly with hardware-efficient kernels and SGDC to achieve better performance.", "essential_references_not_discussed": "No.", "other_strengths_and_weaknesses": "Strengths:\n1. This paper introduces a dynamic speculative execution framework for self-attention, combining approximate and exact computations with on-the-fly verification, which improves latency without sacrificing accuracy.\n2. Although the experiments were conducted on a specialized device (Tenstorrent N150), I believe the proposed methods could be applicable to other computational platforms as well.\n\nWeakness:\n1. The experimental scope is relatively limited, as it does not encompass larger models, particularly those exceeding 10 billion parameters.\n2. The experimental setup is not sufficiently comprehensive. There are no experiments demonstrating how much benefit the Kernel Fusion and SGDC-related optimizations provide in terms of latency and throughput. Additionally, there is no specific comparison showing the performance difference between unfused speculative attention and the fused kernel implementation.\n3. The paper does not explain how the proposed method handles different batch sizes.", "other_comments_or_suggestions": "Clarify the baseline depicted in Figure 2.", "questions_for_authors": "1. In Figure 8, SP demonstrates scalability advantages compared to TP. Is TP's scalability disadvantage due to communication overhead? \n2. Compared to TP, SP seems more similar to DP. In the experimental setup, shouldn't the comparison include DP+TP rather than just Full TP?", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "Y67sQeUbRF", "reviewer": "Reviewer_JxHd", "summary": "LLM is resource-intensive, and serving LLMs is difficult. Model parallelism is bottlenecked by communication, when the communication bandwidth is low, and data parallelism is great at throughput but not good as inference latency. Approximate attention is robust at instruction-following and knowledge retrieval tasks but not good at context-heavy tasks and complex reasoning tasks. The paper uses approximate attention as a way to speed up self-attention in a speculative decoding way but at the per-layer level. In the setup introduced in the paper, it uses two NPU per execution/shard, where one is main, another one is speculative. The main thread will execute their in-house attention kernel that does StreamingLLM attention and the full attention at the same time, the StreamingLLM finishes first, and kick start the FFN computation on the speculative thread, which is on a different NPU. When the full attention finishes, it uses a threshold to check whether the StreamingLLM attention output is close to the full attention. If true, the main thread now becomes the speculative thread, and vice versa. If not, it will continue generate following the original path. \nThe takeaway is that compared with a 8 cards tensor parallelism, \nthe upper bound in efficiency of the method is 4 cards StreamingLLM tensor parallelism. \nThe lower bound in efficiency of the method is 4 cards full-attention tensor parallelism. \nBecause of the hit-rate and lower overhead of communication, the system is shown to beat the tensor parallelism in efficiency, while also preserve the serving efficiency.", "questions": "", "limitations": "", "rating": 4, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We are grateful for your detailed feedback which will greatly improve the quality of the work.\n\n## GPU/TPU generalization\n\nWe ran new experiments for a correctness analysis on the Llama 3.3 70B Instruct model, which are more suitable for frontier GPUs like H100 with higher communication bandwidth and speed. We find that in terms of speculation hit rate and benchmark correctness across different lambdas, the Llama 70B model, despite being deeper and larger, manifests the same behavior as the 8B model. Specifically, we benchmarked the 70B models on GPQA with CoT, GSM8K with CoT, Multi-lingual GSM with CoT in Swahili, and Hotpot QA. The results are:\n\n| Config\\Tasks | GPQA_COT | GSM8K_COT | MGSM_COT_Swahili | Hotpot QA |\n|---|---|---|---|---|\n|Baseline|  0.518 | 0.958 | 0.852 | 0.940 |\n|lambda=0.05|  0.529/55.1% |  0.950/64.5% | 0.856/73.9% | 0.940/32.3% |\n|lambda=0.10|  0.507/80.0% | 0.951/86.6% | 0.852/90.1% | 0.945/59.0% |\n|lambda=0.15|  0.458/89.9% | 0.946/93.9% | 0.840/95.6% | 0.935/74.6% |\n|lambda=0.20|  0.446/95.1% | 0.936/96.8% | 0.820/98.1% | 0.935/83.3% |\n|lambda=0.25 |  0.379/97.6% | 0.897/98.3% | 0.816/99.1% | 0.935/89.4% |\n\nAdditionally, we ran a performance analysis serving the 8B models on 4 vs. 8 H100 doing tensor parallelism using the SGLang serving framework and FlashInfer attention backend. Althoguh we don't (yet) have an ALSpec implementation, it provides us an estimation of the performance gain if the method is implemented on H100s. The results are summarized in the table below:\n\n| Context Length | 4xH100 Attn Latency (us) | 4xH100 Non-Attn Latency (us) | 4xH100 TP Tok/s | 8xH100 TP Tok/s | TP Scaling | Projected ALSpec @ 65% Hit Rate | ALSpec Scaling |\n|---|---|---|---|---|----|---|---|\n|1k| 13 | 95 | 244.6 | 249.3 | 1.9% | 262.0 | 7.1% |\n|32k| 29 | 100 | 214.2 | 231.3 | 8.0% | 246.0 | 14.8% |\n|64k| 49 | 100 | 191.4 | 209.7 | 9.6% | 237.8 | 24.2% |\n|96k| 56 | 100 | 178.2 | 194.2 | 9.0% | 224.9 | 26.2% |\n|128k| 63 | 101 | 169.3 | 184.8 | 9.2% | 217.5 | 28.5% |\n\nOur new estimation on 8xH100s shows that ALSpec would cut latency on Llama 8B by 1.28x compared to 1.09x for full TP at context length 128K. In this case, the attention latency is only 66% of non-attention latency, which means more gains can be achieved for context length beyond 128K. For implementation on GPUs/TPUs, although ALSpec introduces a new attention kernel and run time modification SGDC, it does not fundamentally change the Op by Op and static graph execution style on modern platforms. Therefore, we believe ALSpec is implementable on those platforms by expert kernel writers.\n\nWe also believe multi-node serving would be ideal for models with more parameters such as Llama 70B, as running these models on a single node with 8 devices TP would still give strong scaling. Only doing TP across multiple nodes would they start to show diminishing return, and we believe that this would be a perfect situation to apply ALSpec. Unfortunately, the resources required for these exepriments are beyond our current experiment setup. However, with the correctness analysis and our implemented kernels and methods, we believe that our approach is scalable to larger models when computing resources are available.\n\n## Scaling Context Length Beyond 128K\n\nWhen full-attention becomes more expensive than FFN, there are bubbles in NPUs. We thank the reviewer for pointing this out. This situation falls into Scenario 2 as depicted in Figure 5. We are aware of this issue and we denote this as one future extension of the work. Since this work introduces the new idea of ALSpec and the content is already very full for an 8-page manuscript, we have limited our scope to Scenario 3 in Figure 5. Fortunately, most widely used modern models have 128K context length and usually have attention running faster than the FFN.\n\n## Additional literature\n\nThank you for the pointer to this valuable reference. We will add a citation and in-text compare/contrast to our camera ready.", "claims_and_evidence": "The important per-condition of the proposed method is that doing tensor parallelism on multiple NPUs has huge communication overhead. The claim is justified in left side of Figure 1. It would be better if the paper also discuss GPUs such as H100s that have maybe higher communication bandwidth and speed to see whether it holds universally. \n\nMost claims in the paper are backed strongly by abundance of experiments.", "methods_and_evaluation_criteria": "The method relies on the approximate attention to speculate in the per-layer basis. Also, using two GPU to overlap the FFN (during speculation) and full-attention (main thread) makes intuitive sense. \n\nThe evaluation is comprehensive for 8B parameter models. It comprehensively covers important aspects of LLMs' performance.", "theoretical_claims": "The threshold used for branching decision is derived to make sure that the accumulated error at the end of the LLM is bounded. The proof in Appendix A contains no obvious error.", "experimental_designs_or_analyses": "The advantage of the method: \n1. The paper most claims are strongly backed by comprehensive experiments, from approximate attention weakness to end-to-end system speedup and performance. The paper is highly well-written. \n2. The method proposed in the paper is novel in design, and the kernel designed is inspiring, as by perturbing the sequence of block computation in Flash Decoding, the kernel can compute both StreamingLLM attention and full attention in the same time without damaging the full attention latency. \n3. The per-layer analysis of StreamingLLM replacement is also insightful and can benefit the speculative decoding community. \nThe weakness of the method: \n1. The core limitation is that the method effectiveness is doubtful for frontier GPUs with the highest speed communication bandwidth. \n2. There is some computation waste during rejection in branch that is innate by design. \n3. Another minor issue is that the method might not scale well when the context length goes beyond 128k, which makes full-attention more expensive than FFN, injecting more bubbles to the NPUs.", "relation_to_broader_scientific_literature": "The per-layer speculation and verification with partial attention is quite novel. \nThe per-layer analysis of StreamingLLM replacement (Figure 4) is also insightful and can benefit the speculative decoding community.", "essential_references_not_discussed": "There are some methods before that uses StreamingLLM to do output token-level speculative decoding, although they are very different from per-layer speculative decoding, they should be discussed and cited by the paper. \nSun, H., Chen, Z., Yang, X., Tian, Y., & Chen, B. (2024). Triforce: Lossless acceleration of long sequence generation with hierarchical speculative decoding. arXiv preprint arXiv:2404.11912.", "other_strengths_and_weaknesses": "I have no comments on other strengths and weaknesses.", "other_comments_or_suggestions": "No other comments.", "questions_for_authors": "1. Does the method works for frontier devices such as H100s? \n2. Can the method scale up potentially to multi-node serving situation?", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "PFh4WKCnfH", "reviewer": "Reviewer_xU2t", "summary": "This paper presents a novel infra method that accelerates the model forwarding (interence time) speed via attention-level speculation.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We are grateful for your detailed feedback which will greatly improve the quality of the work.\n\n## Error Propagation for Deeper Models\nTo make sure our method works for larger model with deeper layers, empirically, we have conducted experiments on the Llama 3.3 70B model for correctness analysis, which has 80 layers compared to 32 layers in Llama 3.1 8B. We showed that in terms of speculation hit rate and benchmark correctness across different lambdas, the Llama 70B model, despite deeper and larger, manifest the same behaviour as the 8B model. Specifically, we benchmarked the 70B models on GPQA with CoT, GSM8K with CoT, Multi-lingual GSM with CoT in Swahili, and Hotpot QA. The results are shown below:\n\n| Config\\Tasks | GPQA_COT | GSM8K_COT | MGSM_COT_Swahili | Hotpot QA |\n|----------|----------|----------|----------|----------|\n|Baseline|  0.518 | 0.958 | 0.852 | 0.940 |\n|lambda=0.05| 0.529/55.1% |  0.950/64.5% | 0.856/73.9% | 0.940/32.3% |\n|lambda=0.10|  0.507/80.0% | 0.951/86.6% | 0.852/90.1% | 0.945/59.0% |\n|lambda=0.15|  0.458/89.9% | 0.946/93.9% | 0.840/95.6% | 0.935/74.6% |\n|lambda=0.20|  0.446/95.1% | 0.936/96.8% | 0.820/98.1% | 0.935/83.3% |\n|lambda=0.25 |  0.379/97.6% | 0.897/98.3% | 0.816/99.1% | 0.935/89.4% |\n\nOverall, our new benchmark results empirically confirms that ALSpec works with larger models with deeper layers.\n\n## Theoretical Claims\nAs shown in Appendix A, the worst-case error is in the order of $O(N^2\\delta)$, where $N$ is the number of layers and $\\delta$ is the threshold, which is certainly a limitation for deeper models. However, with stronger assumptions, the bound only has a dependence on $\\sqrt{N}$ which is much more scalable. This is consistent with our new experiment results on the Llama 70B model, which achieved as good correctness as the 8B model despite being much deeper (80 layers) compared to the 8B model (32 layers).\n\n## Impact of Batch Size\nSpeed results under different sizes and settings are indeed important. Our work focuses on a specific case of decode, which is decoding at long context length with small batch size (e.g. batch_size=1). Usually during long context decode, the batch sizes are very small due to limited on-chip memory for KV cache. Since this work already introduces the new idea of ALSpec and our space is already very limited for this 8 pages manuscript, we therefore limited our scope to this scenario. Relevantly, this is the case for long context decode on most accelerators such as H100s or Tenstorrent's N150, where the on-chip memory limits the use of large batch size at long context length.", "claims_and_evidence": "Most claims are convincing. Though I'm not convinced that such error is actually controllable (as shown in the main figure) if the model layer is very deep.", "methods_and_evaluation_criteria": "I think the benchmark results are sufficient. Evaluation metric is correct in terms of effciency.", "theoretical_claims": "No. I'm not a theory guy and at the same time I don't think the theory introduced will ensure the error is under control with the proposed threshould acceptance method.", "experimental_designs_or_analyses": "Experiment is appropriate. But I do believe speed results under different model sizes and settings (like memory-bounded and computing-bounded) are necessary. e.g., when the batch size is large, what trend will it show as the accelerators scale up?", "relation_to_broader_scientific_literature": "inference acceleration", "essential_references_not_discussed": "No", "other_strengths_and_weaknesses": "NA", "other_comments_or_suggestions": "NA", "questions_for_authors": "Can you show the error propagation results (or benchmark scores) with larger models with deeper layers? if so, I will be convinced that it's really useful.", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Meta-review: The authors present a compelling approach with Attention-Level Speculation (ALSpec), demonstrating significant gains in inference efficiency for transformer-based LLMs. The empirical validation is robust, especially showcasing the practicality and effectiveness of ALSpec on real hardware platforms. The authors' detailed engagement during the rebuttal period clearly address reviewer concerns with thoughtful responses and clarifications. Overall, this work presents meaningful progress towards efficient inference for LLMs and reflects a rigorous commitment from the authors to strengthen their contributions. Therefore, I recommend acceptance of this work.
Decision: Accept (poster)