2401.10774

# Medusa: Simple LLM Inference Acceleration Framework with Multiple Decoding Heads **Authors**: Tianle Cai, Yuhong Li, Zhengyang Geng, Hongwu Peng, Jason D. Lee, Deming Chen, Tri Dao Abstract Large Language Models (LLMs) employ auto-regressive decoding that requires sequential computation, with each step reliant on the previous one’s output. This creates a bottleneck as each step necessitates moving the full model parameters from High-Bandwidth Memory (HBM) to the accelerator’s cache. While methods such as speculative decoding have been suggested to address this issue, their implementation is impeded by the challenges associated with acquiring and maintaining a separate draft model. In this paper, we present Medusa, an efficient method that augments LLM inference by adding extra decoding heads to predict multiple subsequent tokens in parallel. Using a tree-based attention mechanism, Medusa constructs multiple candidate continuations and verifies them simultaneously in each decoding step. By leveraging parallel processing, Medusa substantially reduces the number of decoding steps required. We present two levels of fine-tuning procedures for Medusa to meet the needs of different use cases: Medusa -1: Medusa is directly fine-tuned on top of a frozen backbone LLM, enabling lossless inference acceleration. Medusa -2: Medusa is fine-tuned together with the backbone LLM, enabling better prediction accuracy of Medusa heads and higher speedup but needing a special training recipe that preserves the model’s capabilities. Moreover, we propose several extensions that improve or expand the utility of Medusa, including a self-distillation to handle situations where no training data is available and a typical acceptance scheme to boost the acceptance rate while maintaining generation quality. We evaluate Medusa on models of various sizes and training procedures. Our experiments demonstrate that Medusa -1 can achieve over 2.2 $×$ speedup without compromising generation quality, while Medusa -2 further improves the speedup to 2.3-2.8 $×$ . Machine Learning, ICML 1 Introduction The recent advancements in Large Language Models (LLMs) have demonstrated that the quality of language generation significantly improves with an increase in model size, reaching billions of parameters (Brown et al., 2020; Chowdhery et al., 2022; Zhang et al., 2022; Hoffmann et al., 2022; OpenAI, 2023; Google, 2023; Touvron et al., 2023). However, this growth has led to an increase in inference latency, which poses a significant challenge in practical applications. From a system perspective, LLM inference is predominantly memory-bandwidth-bound (Shazeer, 2019; Kim et al., 2023), with the main latency bottleneck stemming from accelerators’ memory bandwidth rather than arithmetic computations. This bottleneck is inherent to the sequential nature of auto-regressive decoding, where each forward pass requires transferring the complete model parameters from High-Bandwidth Memory (HBM) to the accelerator’s cache. This process, which generates only a single token, underutilizes the arithmetic computation potential of modern accelerators, leading to inefficiency. To address this, one approach to speed up LLM inference involves increasing the arithmetic intensity (the ratio of total floating-point operations (FLOPs) to total data movement) of the decoding process and reducing the number of decoding steps. In line with this idea, speculative decoding has been proposed (Leviathan et al., 2022; Chen et al., 2023; Xia et al., 2023; Miao et al., 2023). This method uses a smaller draft model to generate a token sequence, which is then refined by the original, larger model for acceptable continuation. However, obtaining an appropriate draft model remains challenging, and it’s even harder to integrate the draft model into a distributed system (Chen et al., 2023). Instead of using a separate draft model to sequentially generate candidate outputs, in this paper, we revisit and refine the concept of using multiple decoding heads on top of the backbone model to expedite inference (Stern et al., 2018). We find that when applied effectively, this technique can overcome the challenges of speculative decoding, allowing for seamless integration into existing LLM systems. Specifically, we introduce Medusa, a method that enhances LLM inference by integrating additional decoding heads to concurrently predict multiple tokens. These heads are fine-tuned in a parameter-efficient manner and can be added to any existing model. With no requirement for a draft model, Medusa offers easy integration into current LLM systems, including those in distributed environments, ensuring a user-friendly experience. We further enhance Medusa with two key insights. Firstly, the current approach of generating a single candidate continuation at each decoding step leads to inefficient use of computational resources. To address this, we propose generating multiple candidate continuations using the Medusa heads and verifying them concurrently through a simple adjustment to the attention mask. Secondly, we can reuse the rejection sampling scheme as used in speculative decoding (Leviathan et al., 2022; Chen et al., 2023) to generate consistent responses with the same distribution as the original model. However, it cannot further enhance the acceleration rate. Alternatively, we introduce a typical acceptance scheme that selects reasonable candidates from the Medusa head outputs. We use temperature as a threshold to manage deviation from the original model’s predictions, providing an efficient alternative to the rejection sampling method. Our results suggest that the proposed typical acceptance scheme can accelerate the decoding speed further while maintaining a similar generation quality. To equip LLMs with predictive Medusa heads, we propose two distinct fine-tuning procedures tailored to various scenarios. For situations with limited computational resources or when the objective is to incorporate Medusa into an existing model without affecting its performance, we recommend Medusa -1. This method requires minimal memory and can be further optimized with quantization techniques akin to those in QLoRA (Dettmers et al., 2023), without compromising the generation quality due to the fixed backbone model. However, in Medusa -1, the full potential of the backbone model is not utilized. We can further fine-tune it to enhance the prediction accuracy of Medusa heads, which can directly lead to a greater speedup. Therefore, we introduce Medusa -2, which is suitable for scenarios with ample computational resources or for direct Supervised Fine-Tuning (SFT) from a base model. The key to Medusa -2 is a training protocol that enables joint training of the Medusa heads and the backbone model without compromising the model’s next-token prediction capability and output quality. We propose different strategies for obtaining the training datasets depending on the model’s training recipe and dataset availability. When the model is fine-tuned on a public dataset, it can be directly used for Medusa. If the dataset is unavailable or the model underwent a Reinforcement Learning with Human Feedback (RLHF) (Ouyang et al., 2022) process, we suggest a self-distillation approach to generate a training dataset for the Medusa heads. Our experiments primarily focus on scenarios with a batch size of one, which is representative of the use case where LLMs are locally hosted for personal use. We test Medusa on models of varying sizes and training settings, including Vicuna-7B, 13B (trained with a public dataset), Vicuna-33B (Chiang et al., 2023) (trained with a private dataset Upon contacting the authors, this version is experimental and used some different data than Vicuna 7B and 13B.), and Zephyr-7B (trained with both supervised fine-tuning and alignment). Medusa can achieve a speedup of 2.3 to 2.8 times across different prompt types without compromising on the quality of generation. <details> <summary>2401.10774v3/x1.png Details</summary>  ### Visual Description # Technical Document Extraction: Model Architecture and Prediction Flow ## Diagram Overview The image depicts a neural network architecture for text generation with a Medusa-style multi-head prediction system. The diagram is divided into three primary sections connected by directional arrows, with a central llama icon serving as a visual motif. --- ### **1. Original Model (Left Section)** - **Color**: Blue gradient background - **Components**: - **LM Head**: Rectangular box labeled "LM Head" with an arrow pointing to "Last Hidden" - **Transformer Layers**: Rectangular box labeled "Transformer Layers" with an upward arrow to "LM Head" - **Embedding**: Rectangular box labeled "Embedding" with an upward arrow to "Transformer Layers" - **Input**: Text box containing the prompt: `"What will happen if Medusa meets a llama?"` - **Legend**: - ❄️ (Snowflake) / 🔥 (Fire) icon above the section --- ### **2. Medusa Heads (Center Section)** - **Color**: Pink gradient background - **Components**: - **Medusa Head 1**: Rectangular box labeled "Medusa Head 1" with an arrow to "Candidates" - **Medusa Head 2**: Rectangular box labeled "Medusa Head 2" with an arrow to "Candidates" - **Medusa Head 3**: Rectangular box labeled "Medusa Head 3" with an arrow to "Candidates" - **Candidates Section**: - Text box with three candidate outputs: 1. `"It is difficult not"` (✅ Green checkmark) 2. `"It' difficult a"` (❌ Red cross) 3. `"It is' not"` (❌ Red cross) - **Legend**: - 🔥 (Fire) icon above the section --- ### **3. Top-k Predictions (Right Section)** - **Color**: Purple gradient background - **Components**: - **Top-k Predictions**: Rectangular box labeled "Top-k Predictions" with an upward arrow to "Medusa Heads" - **Predictions**: - `"It, I, As"` - `"is, ', the"` - `"difficult, is, '"` - `"not, difficult, a"` - **Single Step Prediction**: Text box with the output: `"It is difficult"` - **Legend**: - ↑ (Upward arrow) icon labeled "TOP" above the section --- ### **4. Connections and Flow** 1. **Input → Embedding → Transformer Layers → LM Head → Last Hidden** 2. **Last Hidden → Medusa Head 1/2/3** 3. **Medusa Heads → Candidates** 4. **Candidates → Single Step Prediction** --- ### **5. Visual Elements** - **Central Icon**: Circular logo featuring a llama with: - Yellow mane - White horns - Star-shaped crown - Black-and-white background - **Arrows**: - Black directional arrows connecting components - Colored icons (❄️, 🔥, ↑) for section identification --- ### **6. Key Observations** - The Medusa Heads generate multiple candidate outputs, with one correct prediction highlighted. - The single-step prediction combines elements from the top candidates to form a coherent output. - The architecture emphasizes parallel processing via multiple Medusa Heads for diverse prediction paths. --- ### **7. Missing Elements** - No explicit axis titles, legends, or data tables present. - No numerical data or heatmap categories visible. - All textual information is embedded within labeled boxes or candidate lists. --- This extraction captures all textual and structural elements from the diagram. The absence of numerical data or heatmaps limits quantitative analysis to the described architectural flow and prediction logic. </details> Figure 1: Medusa introduces multiple heads on top of the last hidden states of the LLM, enabling the prediction of several subsequent tokens in parallel (Section 2.1.1). During inference, each head generates multiple top predictions for its designated position. These predictions are assembled into candidates, which are processed in parallel using a tree-based attention mechanism (Section 2.1.2). The final step is to verify the candidates and accept a continuation. Besides the standard rejection sampling scheme, a typical acceptance scheme (Section 2.3.1) can also be used here to select reasonable continuations, and the longest accepted candidate prefix will be used for the next decoding phase. 2 Methodology Medusa follows the same framework as speculative decoding, where each decoding step primarily consists of three substeps: (1) generating candidates, (2) processing candidates, and (3) accepting candidates. For Medusa, (1) is achieved by Medusa heads, (2) is realized by tree attention, and since Medusa heads are on top of the original model, the logits calculated in (2) can be used for substep (1) for the next decoding step. The final step (3) can be realized by either rejection sampling (Leviathan et al., 2022; Chen et al., 2023) or typical acceptance (Section 2.3.1). The overall pipeline is illustrated in Figure 1. In this section, we first introduce the key components of Medusa, including Medusa heads, and tree attention. Then, we present two levels of fine-tuning procedures for Medusa to meet the needs of different use cases. Finally, we propose two extensions to Medusa, including self-distillation and typical acceptance, to handle situations where no training data is available for Medusa and to improve the efficiency of the decoding process, respectively. 2.1 Key Components 2.1.1 Medusa Heads In speculative decoding, subsequent tokens are predicted by an auxiliary draft model. This draft model must be small yet effective enough to generate continuations that the original model will accept. Fulfilling these requirements is a challenging task, and existing approaches (Spector & Re, 2023; Miao et al., 2023) often resort to separately pre-training a smaller model. This pre-training process demands substantial additional computational resources. For example, in (Miao et al., 2023), a reported 275 NVIDIA A100 GPU hours were used. Additionally, separate pre-training can potentially create a distribution shift between the draft model and the original model, leading to continuations that the original model may not favor. Chen et al. (2023) have also highlighted the complexities of serving multiple models in a distributed environment. To streamline and democratize the acceleration of LLM inference, we take inspiration from Stern et al. (2018), which utilizes parallel decoding for tasks such as machine translation and image super-resolution. Medusa heads are additional decoding heads appended to the last hidden states of the original model. Specifically, given the original model’s last hidden states $h_{t}$ at position $t$ , we add $K$ decoding heads to $h_{t}$ . The $k$ -th head is used to predict the token in the $(t+k+1)$ -th position of the next tokens (the original language model head is used to predict the $(t+1)$ -th position). The prediction of the $k$ -th head is denoted as $p_{t}^{(k)}$ , representing a distribution over the vocabulary, while the prediction of the original model is denoted as $p_{t}^{(0)}$ . Following the approach of Stern et al. (2018), we utilize a single layer of feed-forward network with a residual connection for each head. We find that this simple design is sufficient to achieve satisfactory performance. The definition of the $k$ -th head is outlined as: | | $\displaystyle p_{t}^{(k)}=\text{softmax}\left(W_{2}^{(k)}·\left(\text{SiLU% }(W_{1}^{(k)}· h_{t})+h_{t}\right)\right),$ | | | --- | --- | --- | $d$ is the output dimension of the LLM’s last hidden layer and $V$ is the vocabulary size. We initialize $W_{2}^{(k)}$ identically to the original language model head, and $W_{1}^{(k)}$ to zero. This aligns the initial prediction of Medusa heads with that of the original model. The SiLU activation function (Elfwing et al., 2017) is employed following the Llama models (Touvron et al., 2023). Unlike a draft model, Medusa heads are trained in conjunction with the original backbone model, which can remain frozen during training (Medusa -1) or be trained together (Medusa -2). This method allows for fine-tuning large models even on a single GPU, taking advantage of the powerful base model’s learned representations. Furthermore, it ensures that the distribution of the Medusa heads aligns with that of the original model, thereby mitigating the distribution shift problem. Additionally, since the new heads consist of just a single layer akin to the original language model head, Medusa does not add complexity to the serving system design and is friendly to distributed settings. We will discuss the training recipe for Medusa heads in Section 2.2. 2.1.2 Tree Attention Through Medusa heads, we obtain probability predictions for the subsequent $K+1$ tokens. These predictions enable us to create length- $K+1$ continuations as candidates. While the speculative decoding studies (Leviathan et al., 2022; Chen et al., 2023) suggest sampling a single continuation as the candidate, leveraging multiple candidates during decoding can enhance the expected acceptance length within a decoding step. Nevertheless, more candidates can also raise computational demands. To strike a balance, we employ a tree-structured attention mechanism to process multiple candidates concurrently. <details> <summary>2401.10774v3/extracted/5668658/tree_attention.png Details</summary>  ### Visual Description # Technical Document Extraction ## Tree Structure - **Root Node**: Labeled "Root" (gray circle) - **Head 1** (pink box): - Connected to "It" (white circle) - Connected to "I" (white circle) - **Head 2** (orange box): - Connected to "is" (white circle) - Connected to "'" (white circle) - Connected to "the" (white circle) - **Tree Mask**: Labeled at bottom left (green tree icon) ## Key - **Color-Coded Groups**: 1. **Head 1** (pink): - "It" - "I" 2. **Head 2** (orange): - "is" - "'" - "the" 3. **Query** (yellow): - "is" - "'" - "the" ## Heatmap Grid - **Rows**: Labeled "Head 1" and "Head 2" - **Columns**: Labeled with query words ("is", "'", "the") - **Checkmarks** (blue checkmarks) indicate matches between query words and heads: - **Head 1 Row**: - Column "It": ✅ - Column "I": ✅ - **Head 2 Row**: - Column "is": ✅ - Column "'": ✅ - Column "the": ✅ ## Diagram Flow 1. **Root** branches to **Head 1** and **Head 2**. 2. **Head 1** contains "It" and "I". 3. **Head 2** contains repeated sequences of "is", "'", "the". 4. **Query** words ("is", "'", "the") are compared against heads via the heatmap. 5. **Tree Mask** visually represents the hierarchical structure. ## Key Trends - **Head 1** matches only its own words ("It", "I"). - **Head 2** matches all query words ("is", "'", "the") due to repetition. - The heatmap reveals **Head 2** has higher query coverage (6/6 matches) vs. **Head 1** (2/2 matches). </details> Figure 2: We demonstrates the use of tree attention to process multiple candidates concurrently. As exemplified, the top-2 predictions from the first Medusa head and the top-3 from the second result in a total of $2× 3=6$ candidates. Each of these candidates corresponds to a distinct branch within the tree structure. To guarantee that each token only accesses its predecessors, we devise an attention mask that exclusively permits attention flow from the current token back to its antecedent tokens. The positional indices for positional encoding are adjusted in line with this structure. This attention mechanism diverges from the traditional causal attention paradigm. Within this framework, only tokens from the same continuation are regarded as historical data. Drawing inspiration from the concept of embedding graph structures into attention as proposed in the graph neural network domain (Ying et al., 2021), we incorporate the tree structure into our attention mask, visualized in Figure 2. Remarkably, similar ideas have also been explored in independent works like Miao et al. (2023); Spector & Re (2023), where they follow a bottom-up approach and construct the tree by merging multiple candidates generated by a draft model. In our method, we instead take a top-down approach to build the tree thanks to the structure of candidates generated by Medusa heads. For a given $k$ -th head, its top- $s_{k}$ predictions serve as the basis for candidate formation, where $s_{k}$ is a designated hyperparameter. These candidates are established by determining the Cartesian product of the top- $s_{k}$ predictions from each head. For instance, in Figure 2, with $s_{1}=2$ and $s_{2}=3$ , each first head prediction can be succeeded by any prediction from the second head. This leads to a tree structure where $s_{k}$ branches exist at the $k$ -th level (considering a virtual root as the $0 0$ -level, in practice, this $0 0$ -level is for the prediction of the language model head of the original model, which can be sampled independently). Within this tree, only a token’s predecessors are seen as historical context, and our attention mask ensures that the attention is only applied on a token’s predecessors. By employing this mask and properly setting the positional indices for positional encoding, we can process numerous candidates simultaneously without the need to expand the batch size. The cumulative number of new tokens is calculated as $\sum_{k=1}^{K}\prod_{i=1}^{k}s_{i}$ . In this section, we demonstrate the most simple and regular way to construct the tree structure by taking the Cartesian product. However, it is possible to construct the tree structure in a more sophisticated way and exploit the unbalanced accuracy of different top predictions of different heads. We will discuss this in Section 2.3.3. 2.2 Training Strategies At the most basic level, we can train Medusa heads by freezing the backbone model and fine-tuning Medusa heads. However, training the backbone in conjunction with the Medusa heads can significantly enhance the accuracy of the Medusa heads. Depending on the computational resources and the specific reqirements of the use case, we propose two levels of training strategies for Medusa heads. In this section, we assume the availability of a training dataset that aligns with the target model’s output distribution. This could be the dataset used for Supervised Fine-Tuning (SFT) of the target model. We will discuss eliminating the need for such a dataset using a self-distillation approach in Section 2.3.2. 2.2.1 Medusa -1: Frozen Backbone To train Medusa heads with a frozen backbone model, we can use the cross-entropy loss between the prediction of Medusa heads and the ground truth. Specifically, given the ground truth token $y_{t+k+1}$ at position $t+k+1$ , the loss for the $k$ -th head is $\mathcal{L}_{k}=-\log p_{t}^{(k)}(y_{t+k+1})$ where $p_{t}^{(k)}(y)$ denotes the probability of token $y$ predicted by the $k$ -th head. We also observe that $\mathcal{L}_{k}$ is larger when $k$ is larger, which is reasonable since the prediction of the $k$ -th head is more uncertain when $k$ is larger. Therefore, we can add a weight $\lambda_{k}$ to $\mathcal{L}_{k}$ to balance the loss of different heads. And the total Medusa loss is: $$ \displaystyle\mathcal{L}_{\text{{Medusa}-1}}=\sum_{k=1}^{K}-\lambda_{k}\log p_% {t}^{(k)}(y_{t+k+1}). \tag{1} $$ In practice, we set $\lambda_{k}$ as the $k$ -th power of a constant like $0.8$ . Since we only use the backbone model for providing the hidden states, we can use a quantized version of the backbone model to reduce the memory consumption. This introduces a more democratized way to accelerate LLM inference, as with the quantization, Medusa can be trained for a large model on a single consumer GPU similar to QLoRA (Dettmers et al., 2023). The training only takes a few hours (e.g., 5 hours for Medusa -1 on Vicuna 7B model with a single NVIDIA A100 PCIE GPU to train on 60k ShareGPT samples). 2.2.2 Medusa -2: Joint Training To further improve the accuracy of Medusa heads, we can train Medusa heads together with the backbone model. However, this requires a special training recipe to preserve the backbone model’s next-token prediction capability and output quality. To achieve this, we propose three strategies: - Combined loss: To keep the backbone model’s next-token prediction capability, we need to add the cross-entropy loss of the backbone model $\mathcal{L}_{\text{LM}}=-\log p_{t}^{(0)}(y_{t+1})$ to the Medusa loss. We also add a weight $\lambda_{0}$ to balance the loss of the backbone model and the Medusa heads. Therefore, the total loss is: $$ \displaystyle\mathcal{L}_{\text{{Medusa}-2}}=\mathcal{L}_{\text{LM}}+\lambda_{% 0}\mathcal{L}_{\text{{Medusa}-1}}. \tag{2} $$ - Differential learning rates: Since the backbone model is already well-trained and the Medusa heads need more training, we can use separate learning rates for them to enable faster convergence of Medusa heads while preserving the backbone model’s capability. - Heads warmup: Noticing that at the beginning of training, the Medusa heads have a large loss, which leads to a large gradient and may distort the backbone model’s parameters. Following the idea from Kumar et al. (2022), we can employ a two-stage training process. In the first stage, we only train the Medusa heads as Medusa -1. In the second stage, we train the backbone model and Medusa heads together with a warmup strategy. Specifically, we first train the backbone model for a few epochs, then train the Medusa heads together with the backbone model. Besides this simple strategy, we can also use a more sophisticated warmup strategy by gradually increasing the weight $\lambda_{0}$ of the backbone model’s loss. We find both strategies work well in practice. Putting these strategies together, we can train Medusa heads together with the backbone model without hurting the backbone model’s capability. Moreover, this recipe can be applied together with Supervised Fine-Tuning (SFT), enabling us to get a model with native Medusa support. 2.2.3 How to Select the Number of Heads Empirically, we found that five heads are sufficient at most. Therefore, we recommend training with five heads and referring to the strategy described in Section 2.3.3 to determine the optimal configuration of the tree attention. With optimized tree attention, sometimes three or four heads may be enough for inference. In this case, we can ignore the redundant heads without overhead. 2.3 Extensions 2.3.1 Typical Acceptance In speculative decoding papers (Leviathan et al., 2022; Chen et al., 2023), authors employ rejection sampling to yield diverse outputs that align with the distribution of the original model. However, subsequent implementations (Joao Gante, 2023; Spector & Re, 2023) reveal that this sampling strategy results in diminished efficiency as the sampling temperature increases. Intuitively, this can be comprehended in the extreme instance where the draft model is the same as the original one: Using greedy decoding, all output of the draft model will be accepted, therefore maximizing the efficiency. Conversely, rejection sampling introduces extra overhead, as the draft model and the original model are sampled independently. Even if their distributions align perfectly, the output of the draft model may still be rejected. However, in real-world scenarios, sampling from language models is often employed to generate diverse responses, and the temperature parameter is used merely to modulate the “creativity” of the response. Therefore, higher temperatures should result in more opportunities for the original model to accept the draft model’s output. We ascertain that it is typically unnecessary to match the distribution of the original model. Thus, we propose employing a typical acceptance scheme to select plausible candidates rather than using rejection sampling. This approach draws inspiration from truncation sampling studies (Hewitt et al., 2022) (refer to Appendix A for an in-depth explanation). Our objective is to choose candidates that are typical, meaning they are not exceedingly improbable to be produced by the original model. We use the prediction probability from the original model as a natural gauge for this and establish a threshold based on the prediction distribution to determine acceptance. Specifically, given $x_{1},x_{2},·s,x_{n}$ as context, when evaluating the candidate sequence $(x_{n+1},x_{n+2},·s,x_{n+K+1})$ (composed by top predictions of the original language model head and Medusa heads), we consider the condition | | $\displaystyle p_{\text{original}}(x_{n+k}|x_{1},x_{2},·s,x_{n+k-1})>$ | | | --- | --- | --- | where $H(·)$ denotes the entropy function, and $\epsilon,\delta$ are the hard threshold and the entropy-dependent threshold respectively. This criterion is adapted from Hewitt et al. (2022) and rests on two observations: (1) tokens with relatively high probability are meaningful, and (2) when the distribution’s entropy is high, various continuations may be deemed reasonable. During decoding, every candidate is evaluated using this criterion, and a prefix of the candidate is accepted if it satisfies the condition. To guarantee the generation of at least one token at each step, we apply greedy decoding for the first token and unconditionally accept it while employing typical acceptance for subsequent tokens. The final prediction for the current step is determined by the longest accepted prefix among all candidates. Examining this scheme leads to several insights. Firstly, when the temperature is set to $0 0$ , it reverts to greedy decoding, as only the most probable token possesses non-zero probability. As the temperature surpasses $0 0$ , the outcome of greedy decoding will consistently be accepted with appropriate $\epsilon,\delta$ , since those tokens have the maximum probability, yielding maximal speedup. Likewise, in general scenarios, an increased temperature will correspondingly result in longer accepted sequences, as corroborated by our experimental findings. Empirically, we verify that typical acceptance can achieve a better speedup while maintaining a similar generation quality as shown in Figure 5. 2.3.2 Self-Distillation In Section 2.2, we assume the existence of a training dataset that matches the target model’s output distribution. However, this is not always the case. For example, the model owners may only release the model without the training data, or the model may have gone through a Reinforcement Learning with Human Feedback (RLHF) procedure, which makes the output distribution of the model different from the training dataset. To tackle this issue, we propose an automated self-distillation pipeline to use the model itself to generate the training dataset for Medusa heads, which matches the output distribution of the model. The dataset generation process is straightforward. We first take a public seed dataset from a domain similar to the target model; for example, using the ShareGPT (ShareGPT, 2023) dataset for chat models. Then, we simply take the prompts from the dataset and ask the model to reply to the prompts. In order to obtain multi-turn conversation samples, we can sequentially feed the prompts from the seed dataset to the model. Or, for models like Zephyr 7B (Tunstall et al., 2023), which are trained on both roles of the conversation, they have the ability to self-talk, and we can simply feed the first prompt and let the model generate multiple rounds of conversation. For Medusa -1, this dataset is sufficient for training Medusa heads. However, for Medusa -2, we observe that solely using this dataset for training the backbone and Medusa heads usually leads to a lower generation quality. In fact, even without training Medusa heads, training the backbone model with this dataset will lead to performance degradation. This suggests that we also need to use the original model’s probability prediction instead of using the ground truth token as the label for the backbone model, similar to classic knowledge distillation works (Kim & Rush, 2016). Concretely, the loss for the backbone model is: $$ \displaystyle\mathcal{L}_{\text{LM-distill}}=KL(p_{\text{original},t}^{(0)}||p% _{t}^{(0)}), \tag{0} $$ where $p_{\text{original},t}^{(0)}$ denotes the probability distribution of the original model’s prediction at position $t$ . However, naively, to obtain the original model’s probability prediction, we need to maintain two models during training, increasing the memory requirements. To further alleviate this issue, we propose a simple yet effective way to exploit the self-distillation setup. We can use a parameter-efficient adapter like LoRA (Hu et al., 2021) for fine-tuning the backbone model. In this way, the original model is simply the model with the adapter turned off. Therefore, the distillation does not require additional memory consumption. Together, this self-distillation pipeline can be used to train Medusa -2 without hurting the backbone model’s capability and introduce almost no additional memory consumption. Lastly, one tip about using self-distillation is that it is preferable to use LoRA without quantization in this case, otherwise, the teacher model will be the quantized model, which may lead to a lower generation quality. 2.3.3 Searching for the Optimized Tree Construction In Section 2.1.2, we present the simplest way to construct the tree structure by taking the Cartesian product. However, with a fixed budget for the number of total nodes in the tree, a regular tree structure may not be the best choice. Intuitively, those candidates composed of the top predictions of different heads may have different accuracies. Therefore, we can leverage an estimation of the accuracy to construct the tree structure. Specifically, we can use a calibration dataset and calculate the accuracies of the top predictions of different heads. Let $a_{k}^{(i)}$ denote the accuracy of the $i$ -th top prediction of the $k$ -th head Here, the accuracy is defined for the single top $i$ -th token, i.e., this accuracy is equal to top- $i$ accuracy minus top- $(i-1)$ accuracy.. Assuming the accuracies are independent, we can estimate the accuracy of a candidate sequence composed by the top $\left[i_{1},i_{2},·s,i_{k}\right]$ predictions of different heads as $\prod_{j=1}^{k}a_{j}^{(i_{j})}$ . Let $I$ denote the set of all possible combinations of $\left[i_{1},i_{2},·s,i_{k}\right]$ and each element of $I$ can be mapped to a node of the tree (not only leaf nodes but all nodes are included). Then, the expectation of the acceptance length of a candidate sequence is: | | $\displaystyle\sum_{\left[i_{1},i_{2},·s,i_{k}\right]∈ I}\prod_{j=1}^{k}a% _{j}^{(i_{j})}.$ | | | --- | --- | --- | Thinking about building a tree by adding nodes one by one, the contribution of a new node to the expectation is exactly the accuracy associated with the node. Therefore, we can greedily add nodes to the tree by choosing the node that is connected to the current tree and has the highest accuracy. This process can be repeated until the total number of nodes reaches the desired number. In this way, we can construct a tree that maximizes the expectation of the acceptance length. Further details can be found in Appendix C. <details> <summary>2401.10774v3/x2.png Details</summary>  ### Visual Description # Technical Document Extraction: Speedup on Different Model Sizes ## Chart Overview - **Title**: Speedup on different model sizes - **Type**: Bar chart - **Purpose**: Compare token processing speed (tokens per second) across model sizes with/without Medusa optimizations ## Axes - **X-axis (Model Size)**: - Categories: `7B`, `13B` - Label: "Model Size" - **Y-axis (Tokens per Second)**: - Range: 0–140 - Label: "Tokens per Second" ## Legend - **Color-Coded Models**: - `Blue`: w/o Medusa - `Orange`: Medusa-1 - `Green`: Medusa-2 ## Data Points ### Model Size: 7B - **w/o Medusa** (Blue): - Tokens per Second: ~45 - Speedup: N/A (baseline) - **Medusa-1** (Orange): - Tokens per Second: ~100 - Speedup: `2.18x` (vs. w/o Medusa) - **Medusa-2** (Green): - Tokens per Second: ~130 - Speedup: `2.83x` (vs. w/o Medusa) ### Model Size: 13B - **w/o Medusa** (Blue): - Tokens per Second: ~35 - Speedup: N/A (baseline) - **Medusa-1** (Orange): - Tokens per Second: ~80 - Speedup: `2.33x` (vs. w/o Medusa) - **Medusa-2** (Green): - Tokens per Second: ~100 - Speedup: `2.83x` (vs. w/o Medusa) ## Key Observations 1. **Speedup Consistency**: - Medusa-2 achieves `2.83x` speedup across both 7B and 13B models. - Medusa-1 shows diminishing returns at larger model sizes (`2.18x` for 7B vs. `2.33x` for 13B). 2. **Performance Scaling**: - Larger models (13B) process fewer tokens per second than smaller models (7B) in baseline configurations. - Medusa optimizations improve efficiency more significantly for smaller models. ## Structural Notes - Bars are grouped by model size, with each group containing three bars (w/o Medusa, Medusa-1, Medusa-2). - Speedup multipliers are explicitly labeled on top of each colored bar. - No gridlines or error bars are present in the chart. </details> (a) <details> <summary>2401.10774v3/x3.png Details</summary>  ### Visual Description # Technical Document Extraction: Speedup Analysis for Vicuna-7B ## Chart Title **Speedup on different categories for Vicuna-7B** --- ### Axis Labels - **Y-Axis**: `Speedup` (values: 1.0, 1.5, 2.0, 2.5, 3.0, 3.5) - **X-Axis**: Categories (listed below) --- ### Categories & Speedup Values | Category | Speedup Multiplier | Color | |----------------|--------------------|-----------| | Humanities | 2.58x | Blue | | Reasoning | 2.58x | Orange | | Roleplay | 2.7x | Green | | Writing | 2.72x | Red | | Stem | 2.77x | Purple | | Math | 3.01x | Brown | | Coding | 3.29x | Pink | | Extraction | 3.62x | Gray | --- ### Key Observations 1. **Highest Speedup**: `Extraction` (3.62x) significantly outperforms other categories. 2. **Tied Values**: `Humanities` and `Reasoning` share the lowest speedup (2.58x). 3. **Gradual Increase**: Speedup values rise incrementally from 2.58x to 3.62x across categories. 4. **Color Coding**: Each category is uniquely mapped to a distinct color for visual differentiation. --- ### Data Table Reconstruction | Category | Speedup (x) | |----------------|-------------| | Humanities | 2.58 | | Reasoning | 2.58 | | Roleplay | 2.7 | | Writing | 2.72 | | Stem | 2.77 | | Math | 3.01 | | Coding | 3.29 | | Extraction | 3.62 | --- ### Color-Legend Mapping - **Blue**: Humanities - **Orange**: Reasoning - **Green**: Roleplay - **Red**: Writing - **Purple**: Stem - **Brown**: Math - **Pink**: Coding - **Gray**: Extraction --- ### Notes - All speedup values are explicitly labeled on the bars (e.g., "2.58x" for Humanities). - No additional legends or sub-categories are present beyond the x-axis categories. - The chart emphasizes quantitative performance differences across domains. </details> (b) Figure 3: Left: Speed comparison of baseline, Medusa -1 and Medusa -2 on Vicuna-7B/13B. Medusa -1 achieves more than 2 $×$ wall-time speedup compared to the baseline implementation while Medusa -2 further improves the speedup by a significant margin. Right: Detailed speedup performance of Vicuna-7B with Medusa -2 on 8 categories from MT-Bench. 3 Experiments In this section, we present experiments to demonstrate the effectiveness of Medusa under different settings. First, we evaluate Medusa on the Vicuna-7B and 13B models (Chiang et al., 2023) to show the performance of Medusa -1 and Medusa -2. Then, we assess our method using the Vicuna-33B and Zephyr-7B models to demonstrate self-distillation’s viability in scenarios where direct access to the fine-tuning recipe is unavailable, as with Vicuna-33B, and in models like Zephyr-7B that employ Reinforcement Learning from Human Feedback (RLHF). The evaluation is conducted on MT-Bench (Zheng et al., 2023), a multi-turn, conversational-format benchmark. Detailed settings can be found in Appendix B. 3.1 Case Study: Medusa -1 v.s. Medusa -2 on Vicuna 7B and 13B Experimental Setup. We use the Vicuna model class (Chiang et al., 2023), which encompasses chat models of varying sizes (7B, 13B, 33B) that are fine-tuned from the Llama model (Touvron et al., 2023). Among them, the 7B and 13B models are trained on the ShareGPT (ShareGPT, 2023) dataset, while the 33B model is an experimental model and is trained on a private dataset. In this section, we use the ShareGPT dataset to train the Medusa heads on the 7B and 13B models for $2$ epochs. We use the v1.5 version of Vicuna models, which are fine-tuned from Llama-2 models with sequence length 4096. Results. We collect the results and show them in Fig. 3. The baseline is the default Huggingface implementation. In Fig. 3(a), we can see that for the 7B models, Medusa -1 and Medusa -2 configurations lead to a significant increase in speed, measuring in tokens processed per second. Medusa -1 shows a 2.18 $×$ speedup, while Medusa -2 further improves this to a 2.83 $×$ . When applied to the larger 13B model, Medusa -1 results in a 2.33 $×$ speed increase, while Medusa -2 maintains a similar performance gain of 2.83 $×$ over the baseline. We also plot the speedup per category for Medusa -2 Vicuna-7B model. We observe that the coding category benefits from a 3.29 $×$ speedup, suggesting that Medusa is particularly effective for tasks in this domain. This points to a significant potential for optimizing coding LLMs, which are widely used in software development and other programming-related tasks. The “Extraction” category shows the highest speedup at 3.62 $×$ , indicating that this task is highly optimized by the Medusa. Overall, the results suggest that the Medusa significantly enhances inference speed across different model sizes and tasks. 3.2 Case Study: Training with Self-Distillation on Vicuna-33B and Zephyr-7B Experimental Setup. In this case study, we focus on the cases where self-distillation is needed. We use the Vicuna-33B model (Chiang et al., 2023) and the Zephyr-7B model (Tunstall et al., 2023) as examples. Following the procedure described in Section 2.3.2, we first generate the datasets with some seed prompts. We use ShareGPT (ShareGPT, 2023) and UltraChat (Ding et al., 2023) as the seed datasets and collect a dataset at about $100k$ samples for both cases. Interestingly, we find that the Zephyr model can continue to generate multiple rounds of conversation with a single prompt, which makes it easy to collect a large dataset. For Vicuna-33B, we generate the multi-turn conversations by iteratively feeding the prompts from each multi-turn seed conversation using random sampling with temperature 0.3. Both models are trained with sequence length $2048$ and batch size $128$ . | Model Name Acc. rate Overhead | Vicuna-7B 3.47 1.22 | Zephyr-7B 3.14 1.18 | Vicuna-13B 3.51 1.23 | Vicuna-33B 3.01 1.27 | | --- | --- | --- | --- | --- | | Quality | 6.18 (+0.01) | 7.25 (-0.07) | 6.43 (-0.14) | 7.18 (+0.05) | | $S_{\textnormal{SpecDecoding}}$ | 1.47 | - | 1.56 | 1.60 | | $S_{\textsc{Medusa}}$ | 2.83 | 2.66 | 2.83 | 2.35 | Table 1: Comparison of various Medusa -2 models. The first section reports the details of Medusa -2, including accelerate rate, overhead, and quality that denoted the average scores on the MT-Bench compared to the original models. The second section lists the speedup ( $S$ ) of SpecDecoding and Medusa, respectively. <details> <summary>2401.10774v3/x4.png Details</summary>  ### Visual Description # Technical Document Extraction: Scatter Plot Analysis ## Axis Labels and Titles - **X-Axis**: "Number of Candidate Tokens" (Range: 0 to 250) - **Y-Axis**: "Acc. Rate" (Accuracy Rate) (Range: 1.0 to 3.5) ## Legend - **Red Dashed Line with Star Marker**: "Sparse Tree Attention" ## Data Points and Trends 1. **Blue Scatter Points**: - Clustered between **2.5 and 3.0** on the Y-axis. - Distributed across the full X-axis range (0–250 tokens). - Density decreases slightly as X increases beyond 100 tokens. 2. **Red Dashed Line ("Sparse Tree Attention")**: - Starts at **(50 tokens, 3.1 Acc. Rate)**. - Increases to **(150 tokens, 3.3 Acc. Rate)**. - Reaches **(250 tokens, 3.5 Acc. Rate)**. - Shows a consistent upward trend with increasing candidate tokens. 3. **Notable Anomaly**: - A single blue point at **(0 tokens, 1.0 Acc. Rate)** labeled **"w/o Medusa"**. - Indicates baseline performance without Medusa integration. ## Structural Observations - **Grid**: Present (horizontal and vertical lines at integer intervals). - **No Additional Elements**: No grid annotations, error bars, or secondary axes. - **Data Distribution**: Blue points exhibit variability but stabilize near 3.0 Acc. Rate for X > 100. ## Cross-Reference Validation - **Legend Alignment**: Red dashed line with star markers matches the plotted trendline. - **Color Consistency**: Blue points correspond to raw data; red line represents modeled trend. ## Summary The plot illustrates the relationship between the number of candidate tokens and accuracy rate. The "Sparse Tree Attention" model demonstrates a positive correlation between token count and performance, achieving 3.5 Acc. Rate at 250 tokens. Baseline performance ("w/o Medusa") is significantly lower at 1.0 Acc. Rate with zero tokens. </details> (a) <details> <summary>2401.10774v3/x5.png Details</summary>  ### Visual Description # Technical Document Extraction: Line Chart Analysis ## Chart Overview - **Type**: Line chart with scatter plot overlay - **Primary Data Series**: - Red dashed line labeled "Sparse Tree Attention" - Blue scatter points representing speed measurements ## Axis Labels and Scales - **X-Axis**: - Title: "Number of Candidate Tokens" - Range: 0 to 250 - Increment: 50 units - **Y-Axis**: - Title: "Speed (token/s)" - Range: 0 to 120 - Increment: 20 units ## Legend - **Location**: Top-right quadrant - **Components**: - Red star marker: "Sparse Tree Attention" (dashed line) - Blue dot: Implicitly represents raw speed data points ## Key Data Points and Trends ### Sparse Tree Attention (Red Dashed Line) 1. **Trend**: Downward slope from left to right 2. **Key Coordinates**: - [50, 120] (Initial peak) - [150, 110] (Midpoint) - [250, 80] (Final value) 3. **Visual Pattern**: Dashed line connects three star markers ### Raw Speed Data (Blue Scatter Points) 1. **Distribution**: - Initial cluster: [10-100 tokens, 90-110 token/s] - Mid-range spread: [100-200 tokens, 80-100 token/s] - Final cluster: [200-250 tokens, 70-85 token/s] 2. **Notable Anomaly**: - [0, 50] point labeled "w/o Medusa" with arrow annotation - Located at bottom-left corner ## Spatial Grounding - **Legend Position**: [x: 0.95, y: 0.95] (normalized coordinates) - **Annotation Position**: [x: 0.05, y: 0.05] (bottom-left corner) ## Trend Verification 1. **Red Line**: Consistent downward trajectory (R² ≈ 0.98) 2. **Blue Points**: - Initial upward trend (tokens 10-50) - Stabilization phase (tokens 50-150) - Gradual decline (tokens 150-250) ## Component Isolation 1. **Header**: Chart title area (not explicitly labeled) 2. **Main Chart**: - Primary focus: Red line vs. blue scatter - Secondary: Annotation and legend 3. **Footer**: Empty space below main chart area ## Data Table Reconstruction | X (Tokens) | Y (token/s) | Marker Type | Annotation | |------------|-------------|-------------|------------------| | 50 | 120 | Star | Sparse Tree Peak | | 150 | 110 | Star | Midpoint | | 250 | 80 | Star | Final Value | | 0 | 50 | Dot | w/o Medusa | ## Cross-Reference Validation - All red markers match "Sparse Tree Attention" legend - All blue points correspond to raw speed measurements - Annotation color (blue) matches scatter point color </details> (b) Figure 4: Effectiveness of numbers of candidate tokens for decoding introduced by trees (default number of candidate token for decoding is 1 when using KV cache). Left: The acceleration rate for randomly sampled dense tree settings (blue dots) and optimized sparse tree settings (red stars). Right: The speed (tokens/s) for both settings. The trend lines indicate that while the acceleration rate remains relatively stable for sparse trees, there is a notable decrease in speed as the candidate tokens increases. Results. Table 1 complements these findings by comparing various Medusa -2 models in terms of their acceleration rate, overhead, and quality on MT-Bench with GPT-4 acting as the evaluator to assign performance scores ranging from 0 to 10. We report the quality differences of Medusa compared to the original model. Notably, while the Medusa -2 Vicuna-33B model shows a lower acceleration rate, it maintains a comparable quality. We hypothesize that this is due to a mismatch between the hidden training dataset and the dataset we used for self-distillation. Hence, the model’s generation quality can be well aligned by self-distillation while Medusa heads learn distribution from the self-distillation that potentially shifts from the training set. In our study, we also applied speculative decoding (Chen et al., 2023; Leviathan et al., 2022) to the Vicuna lineup using open-source draft models (details can be found in Appendix D). These results underscore the complex interplay between speed and performance when scaling up model sizes and applying self-distillation techniques. The findings also highlight the potential of the Medusa -2 configuration to boost efficiency in processing while carefully preserving the quality of the model’s outputs, suggesting a promising direction for co-optimizing LLMs with Medusa heads. 3.3 Ablation Study 3.3.1 Configuration of Tree Attention The study of tree attention is conducted on the writing and roleplay categories from the MT-Bench dataset using Medusa -2 Vicuna-7B. We target to depict tree attention’s motivation and its performance. Fig. 4(a) compares the acceleration rate of randomly sampled dense tree configurations (Section. 2.1.2, depicted by blue dots) against optimized sparse tree settings (Section. 2.3.3, shown with red stars). The sparse tree configuration with 64 nodes shows a better acceleration rate than the dense tree settings with 256 nodes. The decline in speed in Fig. 4(b) is attributed to the increased overhead introduced by the compute-bound. While a more complex tree can improve acceleration, it does so at the cost of speed due to intensive matrix multiplications for linear layers and self-attention. The acceleration rate increase follows a logarithmic trend and slows down when the tree size grows as shown in Fig. 4(a). However, the initial gains are substantial, allowing Medusa to achieve significant speedups. If the acceleration increase is less than the overhead, it will slow down overall performance. For detailed study, please refer to Appendix G. <details> <summary>2401.10774v3/x6.png Details</summary>  ### Visual Description # Technical Document Extraction: Line Chart Analysis ## Chart Overview The image depicts a line chart comparing two metrics across a range of posterior thresholds. The chart includes two primary y-axes and a secondary y-axis, with distinct line representations for two algorithms. --- ## Axis Labels and Markers ### Primary Y-Axis (Left) - **Title**: `Acc. Rate` - **Range**: `3.0` to `3.5` - **Markers**: Incremental ticks at `3.0`, `3.1`, `3.2`, `3.3`, `3.4`, `3.5` ### Secondary Y-Axis (Right) - **Title**: `Scores` - **Range**: `7.0` to `7.6` - **Markers**: Incremental ticks at `7.0`, `7.1`, `7.2`, `7.3`, `7.4`, `7.5`, `7.6` ### X-Axis - **Title**: `Posterior Thresholds` - **Range**: `0.00` to `0.25` - **Markers**: Incremental ticks at `0.00`, `0.05`, `0.10`, `0.15`, `0.20`, `0.25` --- ## Legend - **Blue Line**: Labeled `RS` - **Orange Line**: Labeled `Greedy` - **Symbols**: - `RS`: Represented by a solid blue line. - `Greedy`: Represented by a dashed orange line with star markers at specific points. --- ## Line Descriptions ### RS (Blue Line) - **Starting Point**: `(0.00, 3.5)` - **Trend**: - Initial sharp decline from `3.5` to `3.2` between `0.00` and `0.05`. - Fluctuates between `3.2` and `3.3` with minor oscillations. - Ends at approximately `3.25` at `0.25`. ### Greedy (Orange Line) - **Starting Point**: `(0.00, 3.0)` - **Trend**: - Sharp rise to `3.3` at `0.05`, followed by a drop to `3.0` at `0.07`. - Exhibits high volatility with peaks up to `3.4` and troughs down to `3.0`. - Ends at approximately `3.4` at `0.25`. --- ## Key Observations 1. **Divergence in Performance**: - `RS` maintains a relatively stable `Acc. Rate` between `3.2` and `3.3`. - `Greedy` shows significant variability, with `Acc. Rate` fluctuating between `3.0` and `3.4`. 2. **Secondary Metric (Scores)**: - The `Scores` axis (right) correlates with the `Acc. Rate` but is not explicitly mapped to either line. Further analysis is required to establish relationships. 3. **Critical Data Points**: - `Greedy` achieves its highest `Acc. Rate` (`3.4`) near `0.25`. - `RS` starts with the highest `Acc. Rate` (`3.5`) but declines steadily. --- ## Notes - The chart uses dual y-axes to visualize two metrics (`Acc. Rate` and `Scores`) against the same x-axis (`Posterior Thresholds`). - The `Greedy` algorithm demonstrates higher sensitivity to threshold changes compared to `RS`. - No explicit data table is present; trends are inferred from line trajectories. </details> Figure 5: Performance comparison of Medusa using proposed typical sampling. The model is fully fine-tuned from Vicuna-7B. The plot illustrates the acceleration rate and average scores on the writing and roleplay (MT-Bench) with a fixed temperature of 0.7 for 3 different settings: greedy sampling and random sampling (RS) plotted as the star and the dot, and typical sampling curves under different thresholds. 3.3.2 Thresholds of Typical Acceptance The thresholds of typical acceptance are studied on the writing and roleplay categories from the MT-Bench dataset (Zheng et al., 2023) using Medusa -2 Vicuna 7B. Utilizing the Vicuna 7B model, we aligned our methodology with the approach delineated by (Hewitt et al., 2022) setting the $\alpha=\sqrt{\epsilon}$ . Fig. 5 presents a comparative analysis of our model’s performance across various sampling settings. These settings range from a threshold $\epsilon$ starting at 0.01 and incrementally increasing to 0.25 in steps of 0.01. Our observations indicate a discernible trade-off: as $\epsilon$ increases, there is an elevation in quality at the expense of a reduced acceleration rate. Furthermore, for tasks demanding creativity, it is noted that the default random sampling surpasses greedy sampling in performance, and the proposed typical sampling is comparable with random sampling when $\epsilon$ increases. | | Baseline | Direct Fine-tuning | Medusa -1 | Medusa -2 | | --- | --- | --- | --- | --- | | Quality | 6.17 | 5.925 | 6.23 | 6.18 | | Speedup | N/A | N/A | 2.18 | 2.83 | Table 2: Comparison of Different Settings of Vicuna-7B. Quality is obtained by evaluating models on MT-Bench using GPT-4 as the judge (higher the better). 3.3.3 Effectiveness of Two-stage Fine-tuning Table 2 shows the performance differences between various fine-tuning strategies for the Vicuna-7B model. Medusa -1, which fine-tunes only the Medusa heads, achieves a 2.18x speedup without compromising generation quality. Medusa -2, which employs two-stage fine-tuning (Section 2.2.2), maintains generation quality and provides greater speedup (2.83x) compared to Medusa -1. In contrast, direct fine-tuning the model with the Medusa heads results in degraded generation quality. The findings indicate that implementing our Medusa -2 for fine-tuning maintains the model’s quality and concurrently improves the speedup versus Medusa -1. Table 3: Impact of Techniques on Speedup | Medusa-1 heads without tree attention | $\sim$ 1.5x | | --- | --- | | Adding tree attention | $\sim$ 1.9x | | Using optimized tree configuration | $\sim$ 2.2x | | Training heads with Medusa-2 | $\sim$ 2.8x | 4 Discussion In conclusion, Medusa enhances LLM inference speed by 2.3-2.8 times by equipping models with additional predictive decoding heads, allowing for generating multiple tokens simultaneously and bypassing the sequential decoding limitation. Key advantages of Medusa include its simplicity, parameter efficiency, and ease of integration into existing systems. Medusa avoids the need for specialized draft models. The typical acceptance scheme removes complications from rejection sampling while providing reasonable outputs. Our approach including two efficient training procedures, ensures high-quality output across various models and prompt types. We summarize the development of each technique and their impact on the speedup in Table 3. In the paper, we focus on the setting with batch size 1 for simplicity. Yet, we want to emphasize that the ideas presented in our paper can be generalized to larger batch-size settings, which are now supported by libraries like TensorRT and Huggingface TGI following our paper. Acknowledgements We extend our heartfelt gratitude to several individuals whose contributions were invaluable to this project: - Zhuohan Li, for his invaluable insights on LLM serving. If you haven’t already, do check out Zhuohan’s vLLM project—it’s nothing short of impressive. - Shaojie Bai, for engaging in crucial discussions that helped shape the early phases of this work. - Denny Zhou, for introducing the truncation sampling scheme to Tianle and encouraging Tianle to explore the area of LLM serving. - Yanping Huang, for pointing out the memory-bandwidth-bound challenges associated with LLM serving to Tianle. - Lianmin Zheng, for clarifying the different training recipes used in different sizes of Vicuna models. Jason D. Lee acknowledges the support of the NSF CCF 2002272, NSF IIS 2107304, and NSF CAREER Award 2144994. Deming Chen acknowledges the support from the AMD Center of Excellence at UIUC. Impact Statement The introduction of Medusa, an innovative method to improve the inference speed of Large Language Models (LLMs), presents a range of broader implications for society, technology, and ethics. This section explores these implications in detail. Societal and Technological Implications - Accessibility and Democratization of AI: By significantly enhancing the efficiency of LLMs, Medusa makes advanced AI technologies more accessible to a wider range of users and organizations. Democratization can spur innovation across various sectors, including education, healthcare, and entertainment, potentially leading to breakthroughs that benefit society at large. - Environmental Impact: The acceleration for LLM inference due to Medusa could lead to decreased energy consumption and a smaller carbon footprint. This aligns with the growing need for sustainable AI practices, contributing to environmental conservation efforts. - Economic Implications: The increased efficiency brought about by Medusa may lower the cost barrier to deploying state-of-the-art AI models, enabling small and medium-sized enterprises to leverage advanced AI capabilities. This could stimulate economic growth, foster competition, and drive technological innovation. Ethical Considerations - Bias and Fairness: While Medusa aims to improve LLM efficiency, it inherits the ethical considerations of its backbone models, including issues related to bias and fairness. The method’s ability to maintain generation quality necessitates investigation to ensure that the models do not perpetuate or amplify existing biases. - Transparency and Accountability: The complexity of Medusa, particularly with its tree-based attention mechanism and multiple decoding heads, may pose challenges in terms of model interpretability. Ensuring transparency in how decisions are made and maintaining accountability for those decisions are crucial for building trust in AI systems. - Security and Privacy: The accelerated capabilities of LLMs augmented by Medusa could potentially be exploited for malicious purposes, such as generating disinformation at scale or automating cyber-attacks. It is imperative to develop and enforce ethical guidelines and security measures to prevent misuse. References - Ainslie et al. (2023) Ainslie, J., Lee-Thorp, J., de Jong, M., Zemlyanskiy, Y., Lebrón, F., and Sanghai, S. Gqa: Training generalized multi-query transformer models from multi-head checkpoints. arXiv preprint arXiv:2305.13245, 2023. - Axolotl (2023) Axolotl. Axolotl. https://github.com/OpenAccess-AI-Collective/axolotl, 2023. - Basu et al. (2021) Basu, S., Ramachandran, G. S., Keskar, N. S., and Varshney, L. R. {MIROSTAT}: A {neural} {text} {decoding} {algorithm} {that} {directly} {controls} {perplexity}. In International Conference on Learning Representations, 2021. URL https://openreview.net/forum?id=W1G1JZEIy5_. - Brown et al. (2020) Brown, T., Mann, B., Ryder, N., Subbiah, M., Kaplan, J. D., Dhariwal, P., Neelakantan, A., Shyam, P., Sastry, G., Askell, A., et al. Language models are few-shot learners. Advances in neural information processing systems, 33:1877–1901, 2020. - Chen et al. (2023) Chen, C., Borgeaud, S., Irving, G., Lespiau, J.-B., Sifre, L., and Jumper, J. Accelerating large language model decoding with speculative sampling. February 2023. doi: 10.48550/ARXIV.2302.01318. - Chen (2023) Chen, L. Dissecting batching effects in gpt inference. https://le.qun.ch/en/blog/2023/05/13/transformer-batching/, 2023. Blog. - Chiang et al. (2023) Chiang, W.-L., Li, Z., Lin, Z., Sheng, Y., Wu, Z., Zhang, H., Zheng, L., Zhuang, S., Zhuang, Y., Gonzalez, J. E., Stoica, I., and Xing, E. P. Vicuna: An open-source chatbot impressing gpt-4 with 90%* chatgpt quality, March 2023. URL https://lmsys.org/blog/2023-03-30-vicuna/. - Chowdhery et al. (2022) Chowdhery, A., Narang, S., Devlin, J., Bosma, M., Mishra, G., Roberts, A., Barham, P., Chung, H. W., Sutton, C., Gehrmann, S., et al. Palm: Scaling language modeling with pathways. arXiv preprint arXiv:2204.02311, 2022. - Dettmers et al. (2021) Dettmers, T., Lewis, M., Shleifer, S., and Zettlemoyer, L. 8-bit optimizers via block-wise quantization. International Conference on Learning Representations, 2021. - Dettmers et al. (2022) Dettmers, T., Lewis, M., Belkada, Y., and Zettlemoyer, L. Llm. int8 (): 8-bit matrix multiplication for transformers at scale. arXiv preprint arXiv:2208.07339, 2022. - Dettmers et al. (2023) Dettmers, T., Pagnoni, A., Holtzman, A., and Zettlemoyer, L. Qlora: Efficient finetuning of quantized llms. arXiv preprint arXiv:2305.14314, 2023. - Ding et al. (2023) Ding, N., Chen, Y., Xu, B., Qin, Y., Zheng, Z., Hu, S., Liu, Z., Sun, M., and Zhou, B. Enhancing chat language models by scaling high-quality instructional conversations, 2023. - Dubois et al. (2023) Dubois, Y., Li, X., Taori, R., Zhang, T., Gulrajani, I., Ba, J., Guestrin, C., Liang, P., and Hashimoto, T. B. Alpacafarm: A simulation framework for methods that learn from human feedback, 2023. - Elfwing et al. (2017) Elfwing, S., Uchibe, E., and Doya, K. Sigmoid-weighted linear units for neural network function approximation in reinforcement learning. Neural Networks, 2017. doi: 10.1016/j.neunet.2017.12.012. - Fan et al. (2018) Fan, A., Lewis, M., and Dauphin, Y. Hierarchical neural story generation. In Proceedings of the 56th Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers). Association for Computational Linguistics, 2018. doi: 10.18653/v1/p18-1082. - Frantar et al. (2022) Frantar, E., Ashkboos, S., Hoefler, T., and Alistarh, D. Gptq: Accurate post-training quantization for generative pre-trained transformers. arXiv preprint arXiv:2210.17323, 2022. - Google (2023) Google. Palm 2 technical report, 2023. URL https://ai.google/static/documents/palm2techreport.pdf. - Hewitt et al. (2022) Hewitt, J., Manning, C. D., and Liang, P. Truncation sampling as language model desmoothing. October 2022. doi: 10.48550/ARXIV.2210.15191. - Hoffmann et al. (2022) Hoffmann, J., Borgeaud, S., Mensch, A., Buchatskaya, E., Cai, T., Rutherford, E., Casas, D. d. L., Hendricks, L. A., Welbl, J., Clark, A., et al. Training compute-optimal large language models. arXiv preprint arXiv:2203.15556, 2022. - Holtzman et al. (2020) Holtzman, A., Buys, J., Du, L., Forbes, M., and Choi, Y. The curious case of neural text degeneration. In International Conference on Learning Representations, 2020. URL https://openreview.net/forum?id=rygGQyrFvH. - Hu et al. (2021) Hu, E. J., Shen, Y., Wallis, P., Allen-Zhu, Z., Li, Y., Wang, S., and Chen, W. Lora: Low-rank adaptation of large language models. ICLR, 2021. - Joao Gante (2023) Joao Gante. Assisted generation: a new direction toward low-latency text generation, 2023. URL https://huggingface.co/blog/assisted-generation. - Kim et al. (2023) Kim, S., Hooper, C., Gholami, A., Dong, Z., Li, X., Shen, S., Mahoney, M. W., and Keutzer, K. Squeezellm: Dense-and-sparse quantization. arXiv preprint arXiv:2306.07629, 2023. - Kim & Rush (2016) Kim, Y. and Rush, A. M. Sequence-level knowledge distillation. EMNLP, 2016. - Kumar et al. (2022) Kumar, A., Raghunathan, A., Jones, R., Ma, T., and Liang, P. Fine-tuning can distort pretrained features and underperform out-of-distribution. International Conference on Learning Representations, 2022. - Kwon et al. (2023) Kwon, W., Li, Z., Zhuang, S., Sheng, Y., Zheng, L., Yu, C. H., Gonzalez, J. E., Zhang, H., and Stoica, I. Efficient memory management for large language model serving with pagedattention. In Proceedings of the ACM SIGOPS 29th Symposium on Operating Systems Principles, 2023. - Leviathan et al. (2022) Leviathan, Y., Kalman, M., and Matias, Y. Fast inference from transformers via speculative decoding. November 2022. doi: 10.48550/ARXIV.2211.17192. - Li et al. (2023) Li, X., Zhang, T., Dubois, Y., Taori, R., Gulrajani, I., Guestrin, C., Liang, P., and Hashimoto, T. B. Alpacaeval: An automatic evaluator of instruction-following models. https://github.com/tatsu-lab/alpaca_eval, 2023. - Lin et al. (2023) Lin, J., Tang, J., Tang, H., Yang, S., Dang, X., and Han, S. Awq: Activation-aware weight quantization for llm compression and acceleration. arXiv preprint arXiv:2306.00978, 2023. - Meister et al. (2022) Meister, C., Wiher, G., Pimentel, T., and Cotterell, R. On the probability-quality paradox in language generation. March 2022. doi: 10.48550/ARXIV.2203.17217. - Meister et al. (2023) Meister, C., Pimentel, T., Wiher, G., and Cotterell, R. Locally typical sampling. Transactions of the Association for Computational Linguistics, 11:102–121, 2023. - Miao et al. (2023) Miao, X., Oliaro, G., Zhang, Z., Cheng, X., Wang, Z., Wong, R. Y. Y., Chen, Z., Arfeen, D., Abhyankar, R., and Jia, Z. Specinfer: Accelerating generative llm serving with speculative inference and token tree verification. arXiv preprint arXiv:2305.09781, 2023. - (33) NVIDIA. Nvidia a100 tensor core gpu. - OpenAI (2023) OpenAI. Gpt-4 technical report, 2023. - Ouyang et al. (2022) Ouyang, L., Wu, J., Jiang, X., Almeida, D., Wainwright, C. L., Mishkin, P., Zhang, C., Agarwal, S., Slama, K., Ray, A., et al. Training language models to follow instructions with human feedback. arXiv preprint arXiv:2203.02155, 2022. - Pan (2023) Pan, J. Tiny vicuna 1b. https://huggingface.co/Jiayi-Pan/Tiny-Vicuna-1B, 2023. - Pillutla et al. (2021) Pillutla, K., Swayamdipta, S., Zellers, R., Thickstun, J., Welleck, S., Choi, Y., and Harchaoui, Z. MAUVE: Measuring the gap between neural text and human text using divergence frontiers. In Beygelzimer, A., Dauphin, Y., Liang, P., and Vaughan, J. W. (eds.), Advances in Neural Information Processing Systems, 2021. URL https://openreview.net/forum?id=Tqx7nJp7PR. - Pope et al. (2022) Pope, R., Douglas, S., Chowdhery, A., Devlin, J., Bradbury, J., Levskaya, A., Heek, J., Xiao, K., Agrawal, S., and Dean, J. Efficiently scaling transformer inference. November 2022. doi: 10.48550/ARXIV.2211.05102. - ShareGPT (2023) ShareGPT. ShareGPT. https://huggingface.co/datasets/Aeala/ShareGPT_Vicuna_unfiltered, 2023. - Shazeer (2019) Shazeer, N. Fast transformer decoding: One write-head is all you need. arXiv preprint arXiv:1911.02150, 2019. - Spector & Re (2023) Spector, B. and Re, C. Accelerating llm inference with staged speculative decoding. arXiv preprint arXiv:2308.04623, 2023. - Stern et al. (2018) Stern, M., Shazeer, N. M., and Uszkoreit, J. Blockwise parallel decoding for deep autoregressive models. Neural Information Processing Systems, 2018. - Touvron et al. (2023) Touvron, H., Martin, L., Stone, K., Albert, P., Almahairi, A., Babaei, Y., Bashlykov, N., Batra, S., Bhargava, P., Bhosale, S., et al. Llama 2: Open foundation and fine-tuned chat models. arXiv preprint arXiv:2307.09288, 2023. - Tunstall et al. (2023) Tunstall, L., Beeching, E., Lambert, N., Rajani, N., Rasul, K., Belkada, Y., Huang, S., von Werra, L., Fourrier, C., Habib, N., Sarrazin, N., Sanseviero, O., Rush, A. M., and Wolf, T. Zephyr: Direct distillation of lm alignment, 2023. - Xia et al. (2023) Xia, H., Ge, T., Chen, S.-Q., Wei, F., and Sui, Z. Speculative decoding: Lossless speedup of autoregressive translation, 2023. URL https://openreview.net/forum?id=H-VlwsYvVi. - Xiao et al. (2023a) Xiao, G., Lin, J., Seznec, M., Wu, H., Demouth, J., and Han, S. Smoothquant: Accurate and efficient post-training quantization for large language models. In International Conference on Machine Learning, pp. 38087–38099. PMLR, 2023a. - Xiao et al. (2023b) Xiao, Y., Wu, L., Guo, J., Li, J., Zhang, M., Qin, T., and Liu, T.-y. A survey on non-autoregressive generation for neural machine translation and beyond. IEEE Transactions on Pattern Analysis and Machine Intelligence, 2023b. - Ying et al. (2021) Ying, C., Cai, T., Luo, S., Zheng, S., Ke, G., He, D., Shen, Y., and Liu, T.-Y. Do transformers really perform badly for graph representation? Advances in Neural Information Processing Systems, 34:28877–28888, 2021. - Zhang et al. (2024) Zhang, P., Zeng, G., Wang, T., and Lu, W. Tinyllama: An open-source small language model, 2024. - Zhang et al. (2022) Zhang, S., Roller, S., Goyal, N., Artetxe, M., Chen, M., Chen, S., Dewan, C., Diab, M., Li, X., Lin, X. V., et al. Opt: Open pre-trained transformer language models. arXiv preprint arXiv:2205.01068, 2022. - Zhang et al. (2023) Zhang, Z., Sheng, Y., Zhou, T., Chen, T., Zheng, L., Cai, R., Song, Z., Tian, Y., Ré, C., Barrett, C., et al. H $\_2$ o: Heavy-hitter oracle for efficient generative inference of large language models. arXiv preprint arXiv:2306.14048, 2023. - Zheng et al. (2023) Zheng, L., Chiang, W.-L., Sheng, Y., Zhuang, S., Wu, Z., Zhuang, Y., Lin, Z., Li, Z., Li, D., Xing, E. P., Zhang, H., Gonzalez, J. E., and Stoica, I. Judging llm-as-a-judge with mt-bench and chatbot arena, 2023. Appendix A Related Work A.1 LLM Inference Acceleration The inefficiency of Large Language Model (LLM) inference is primarily attributed to the memory-bandwidth-bound nature of the auto-regressive decoding process. Several methods have been proposed to alleviate this issue, improving inference latency and throughput. Traditionally, batch inference has been employed as a straightforward method to enhance arithmetic intensity and escape memory-bandwidth-bound limitations. However, with LLMs, both model parameters and the Key-Value (KV) cache consume substantial accelerator memory, hindering the utilization of large batch sizes. Existing methods to tackle this problem can be conceptually divided into two main categories: (1) Reducing memory consumption, thereby minimizing memory transfer overhead and enabling larger batch sizes, and (2) Minimizing the number of decoding steps to decrease latency directly. Reducing KV Cache. Methods such as Multi-query attention (Shazeer, 2019) and Grouped-query attention (Ainslie et al., 2023) adopt a direct approach to diminish the KV cache. By utilizing fewer key and value heads in the attention modules relative to query heads, these strategies substantially cut the KV’s memory consumption, thereby facilitating larger batch sizes and enhanced accelerator utilization (Pope et al., 2022). Additionally, Zhang et al. (2023) proposes to selectively retain the most critical KV tokens, further reducing the KV cache. From a system perspective, Kwon et al. (2023) introduces a paged memory management scheme for reducing fragmentation of the KV cache. Quantization. Quantization techniques are extensively used to shrink LLMs’ memory consumption. Xiao et al. (2023a) apply rescaling between activations and parameters to eliminate outliers and simplify the quantization process. Dettmers et al. (2022) breaks down matrix multiplications into predominantly 8-bit and a minority of 16-bit operations. Frantar et al. (2022) iteratively round weight columns into 3/4 bits, while Lin et al. (2023) present an activation-aware quantization scheme to protect salient weights and compress LLMs to 3/4 bits. Kim et al. (2023) introduce a sparse plus low-precision pattern to handle a minor portion of vital weights, among other techniques. Speculative Decoding. As an approach orthogonal to the aforementioned methods, speculative decoding (Leviathan et al., 2022; Chen et al., 2023) aims to execute several decoding steps in parallel, thus reducing the total number of steps required. This parallelization is realized by employing a smaller draft model to conjecture several subsequent words, which the LLMs then collectively evaluate and accept as appropriate. While resonating with non-autoregressive generation literature (Xiao et al., 2023b), this method is specifically tailored for LLMs to address the aforementioned inefficiency. Unlike previous works, we propose leveraging the original model to make predictions rather than introducing an additional draft model. This approach is more straightforward and seamlessly integrates into existing systems without the complexities of managing two models. Independently, Miao et al. (2023); Spector & Re (2023) propose the use of tree-structured attention to generate multiple candidates in parallel, where Miao et al. (2023) suggest employing an ensemble of models to propose candidates, and Spector & Re (2023) advocate adding another hierarchy for the draft model. However, draft models require specialized pretraining and alignment with the target models. While employing multiple draft models can be cumbersome and involves the complexity of managing parallelism, our approach, which relies solely on decoding heads, offers a simpler alternative. Miao et al. (2023) employ multiple draft models to generate tokens and merge them using tree attention, while Spector & Re (2023) utilize a small draft model to process each level of the tree in batches. In contrast, our method directly uses the top predicted tokens from each of Medusa heads to create a static sparse tree without autoregression or adjusting the tree structure. This approach simplifies the process and improves efficiency. Additionally, we demonstrate through a detailed ablation study how the nodes of the tree can affect decoding speed. A.2 Sampling Scheme The manner in which text is sampled from Large Language Models (LLMs) can significantly influence the quality of the generated output. Recent studies have revealed that direct sampling from a language model may lead to incoherent or nonsensical results (Pillutla et al., 2021; Holtzman et al., 2020). In response to this challenge, truncation sampling schemes have been introduced (Fan et al., 2018; Basu et al., 2021; Meister et al., 2022; Hewitt et al., 2022; Meister et al., 2023). These approaches aim to produce high-quality and diverse samples by performing sampling on a truncated distribution over a specific allowed set at each decoding step. Different strategies define this allowed set in various ways. For example, top- $k$ sampling (Fan et al., 2018) retains the $k$ most likely words, whereas top- $p$ sampling (Holtzman et al., 2020) incorporates the minimal set of words that account for $p$ percent of the probability. Another method, known as typical decoding (Meister et al., 2023), employs the entropy of the predicted distribution to establish the threshold for inclusion. Hewitt et al. (2022) offers a unified framework to understand truncation sampling techniques comprehensively. Drawing inspiration from these methods, our typical acceptance scheme aligns with the concept of defining an allowed set to exclude improbable candidates from the sampling process. However, we diverge because we do not insist on an exact correspondence between the output and language model distribution. This deviation allows us to facilitate more diverse yet high-quality outputs, achieving greater efficiency without compromising the integrity of the generated text. Appendix B Experiment Settings B.1 Common Terms We clarify three commonly used terms: a) Acceleration rate: This refers to the average number of tokens decoded per decoding step. In a standard auto-regressive model, this rate is 1.0. b) Overhead: This is used to characterize the per decoding step overhead compared to classic decoding, and is calculated by dividing the average per step latency of the Medusa models by that of the vanilla model. c) Speedup: This refers to the wall-time acceleration rate. Following these definitions, we have the relation: Speedup = Acceleration rate / Overhead. B.2 Shared Settings For all the experiments, we use the Axolotl (Axolotl, 2023) framework for training. We use a cosine learning rate scheduler with warmup and use 8-bit AdamW (Dettmers et al., 2021) optimizer. We train $5$ Medusa heads with $1$ layer and set $\lambda_{k}$ in Eq. (1) to be $0.8^{k}$ . For Medusa -2, we use either LoRA (Hu et al., 2021) or QLoRA (Dettmers et al., 2023) for fine-tuning and set the learning rate of Medusa heads to be $4$ times larger than the backbone model. LoRA is applied to all the linear layers of the backbone model, including the language model head. The rank of LoRA adapter is set to $32$ , and $\alpha$ is set to $16$ . A dropout of $0.05$ is added to the LoRA adapter. B.3 Medusa -1 v.s. Medusa -2 on Vicuna 7B and 13B We use a global batch size of $64$ and a peak learning rate of $5e^{-4}$ for the backbone and $2e^{-3}$ for Medusa heads and warmup for $40$ steps. We use $4$ -bit quantized backbone models for both models. We first train the models with Medusa -1 and use these trained models as initialization to train Medusa -2. We employ QLoRA for Medusa -2 and the $\lambda_{0}$ in Eq. (2) is set to be $0.2$ . B.4 Training with Self-Distillation on Vicuna-33B and Zephyr-7B We use Medusa -2 for both models instead of using a two-stage training procedure. We use a sine schedule for the $\theta_{0}$ to gradually increase the value to its peak at the end of the training. We find this approach is equally effective. We set the peak learning rate of the backbone LoRA adapter to be $1e^{-4}$ and the warmup steps to be $20$ since the self-distillation loss is relatively small. We set the $\lambda_{0}$ in Eq. (2) to be $0.01$ . Appendix C Visualization of optimized tree attention Fig. 6 illustrates the structure of a sparsely constructed tree for the Medusa -2 Vicuna-7B model. This tree structure extends four levels deep, indicating the engagement of four Medusa heads in the computation. The tree is initially formed through a Cartesian product approach and subsequently refined by pruning based on the statistical expectations of the top-k predictions from each Medusa head measured on the Alpaca-eval dataset (Dubois et al., 2023). The tree’s lean towards the left visually represents the algorithm’s preference for nodes with higher probabilities on each head. <details> <summary>2401.10774v3/x7.png Details</summary>  ### Visual Description # Technical Document Extraction: Flowchart Analysis ## Overview The image depicts a **directed graph flowchart** with nodes (blue circles) and edges (arrows). The flowchart represents a sequence of events or logical steps, with specific nodes highlighted by **red arrows** to indicate critical connections. --- ## Node Labels and Structure All nodes are labeled with single words or punctuation. Labels are transcribed verbatim from the image: 1. **time** 2. **a** 3. **once** 4. **upon** 5. **there** 6. **is** 7. **a** (duplicate) 8. **time** (duplicate) 9. **Here** 10. **s** 11. **once** (duplicate) 12. **ly** 13. **a** (duplicate) 14. **course** 15. **ertain** 16. **the** 17. **was** 18. **!** (exclamation mark) --- ## Edge Connections - **Gray Arrows**: Represent the primary flow of the sequence. - `time → a → once → upon → there → is → a → time → Here → s → once → ly → a → course → ertain → the → was → !` - **Red Arrows**: Highlight deviations or critical paths. - `time → a` (direct connection) - `a → once` (direct connection) - `upon → ertain` (direct connection) --- ## Key Observations 1. **Repetition**: - Labels like `a`, `time`, and `once` appear multiple times, suggesting cyclical or recursive steps. - The word `ly` (possibly an abbreviation for "ly" or "ly" as a standalone term) is connected to multiple nodes. 2. **Critical Paths**: - Red arrows emphasize connections between `time` and `a`, `a` and `once`, and `upon` and `ertain`. These may represent pivotal decision points or dependencies. 3. **Punctuation**: - The exclamation mark `!` is a terminal node, indicating a conclusion or emphasis. --- ## Flowchart Logic The flowchart progresses linearly from left to right, with the following sequence: `time → a → once → upon → there → is → a → time → Here → s → once → ly → a → course → ertain → the → was → !` Red arrows create shortcuts or emphasize key transitions: - `time → a` (bypasses intermediate steps) - `a → once` (direct link) - `upon → ertain` (bypasses intermediate steps) --- ## Technical Notes - **Node Representation**: Blue circles with white text. - **Edge Representation**: - Gray arrows: Standard flow. - Red arrows: Highlighted connections (possibly critical or exceptional paths). - **No Legends or Axis Titles**: The image lacks legends, axis labels, or numerical data. --- ## Conclusion This flowchart represents a structured sequence of events or logical steps, with red arrows emphasizing critical transitions. The repetition of labels suggests loops or recursive processes. The exclamation mark `!` marks the endpoint of the sequence. </details> Figure 6: Visualization of a sparse tree setting for Medusa -2 Vicuna-7B. The tree has 64 nodes representing candidate tokens and a depth of 4 which indicates 4 Medusa heads involved in calculation. Each node indicates a token from a top-k prediction of a Medusa head, and the edges show the connections between them. The red lines highlight the path that correctly predicts the future tokens. Appendix D Results of Speculative Decoding In this study, speculative decoding was applied to Vicuna models (Chiang et al., 2023) with varying sizes, specifically 7B, 13B, and 33B. The preliminary framework utilized open-source models such as Llama-68M and 160M (Miao et al., 2023), alongside Tiny-Llama (Zhang et al., 2024) and Tiny-Vicuna (Pan, 2023), fine-tuned from Tiny-Llama with the Vicuna-style instructional tuning strategy. Due to the proprietary nature of speculative decoding methods (Chen et al., 2023; Leviathan et al., 2022), open-source alternatives https://github.com/feifeibear/LLMSpeculativeSampling were deployed for evaluation. Additionally, we utilize torch.compile() to accelerate the inference speed of draft models. Our results shown in Fig. 7, reveal that the optimal settings of the draft model vary with the Vicuna model sizes. Specifically, the Llama-68M, with a setting of the draft token number $\gamma=4$ , yielded the best performance for Vicuna-7B, while the same draft model with $\gamma=3$ was most effective for Vicuna-13B. For the larger Vicuna-33B, the Tiny-Vicuna (Vicuna-1B), with $\gamma=3$ , provided the greatest acceleration. These results suggest that the choice and setting of the drafting model should be tailored to the size of the LLMs, presenting an area for further exploration in the field. <details> <summary>2401.10774v3/x8.png Details</summary>  ### Visual Description # Technical Document Extraction: Line Chart Analysis ## Chart Overview The image depicts a line chart comparing the performance of four language models across varying Gamma values. The chart tracks "Tokens per Second" as the dependent variable against "Gamma" on the independent axis. --- ## Axis Labels and Markers - **X-Axis (Gamma):** - Range: 0 to 14 (increments of 2) - Labels: Numerical values (0, 2, 4, ..., 14) - **Y-Axis (Tokens per Second):** - Range: 20 to 70 (increments of 10) - Labels: Numerical values (20, 30, 40, ..., 70) - **Reference Line:** - Gray dashed horizontal line at **45 Tokens per Second**. --- ## Legend and Model Identification The legend (top-right corner) maps colors to models: 1. **Blue Line (Llama-68M):** - Starts at ~60 Tokens/sec (Gamma=0). - Peaks at ~68 Tokens/sec (Gamma=2). - Gradual decline to ~47 Tokens/sec (Gamma=14). 2. **Orange Line (Llama-160M):** - Starts at ~54 Tokens/sec (Gamma=0). - Dips to ~30 Tokens/sec (Gamma=10). - Slight recovery to ~24 Tokens/sec (Gamma=14). 3. **Green Line (Llama-1B):** - Starts at ~50 Tokens/sec (Gamma=0). - Steady decline to ~20 Tokens/sec (Gamma=14). 4. **Red Line (Vicuna-1B):** - Starts at ~50 Tokens/sec (Gamma=0). - Dips to ~30 Tokens/sec (Gamma=10). - Gradual decline to ~27 Tokens/sec (Gamma=14). --- ## Key Trends 1. **Llama-68M (Blue):** - Highest initial performance, but declines consistently with increasing Gamma. - Maintains the highest Tokens/sec across most Gamma values. 2. **Llama-160M (Orange):** - Sharp initial drop (Gamma=0 to 4), followed by stabilization and minor recovery. - Underperforms Llama-68M but outperforms Llama-1B and Vicuna-1B at higher Gamma. 3. **Llama-1B (Green):** - Steady linear decline with no recovery. - Worst performance at Gamma=14 (~20 Tokens/sec). 4. **Vicuna-1B (Red):** - Moderate decline with a slight plateau between Gamma=6 and 8. - Outperforms Llama-1B but lags behind Llama-68M and Llama-160M at higher Gamma. --- ## Critical Observations - **Performance Degradation:** All models show reduced Tokens/sec as Gamma increases, indicating potential computational or efficiency trade-offs. - **Llama-68M Dominance:** Maintains the highest throughput across the Gamma range, suggesting superior scalability or optimization. - **Reference Threshold:** The gray dashed line at 45 Tokens/sec may represent a performance benchmark; Llama-1B and Vicuna-1B fall below this at Gamma ≥ 8. --- ## Data Points Summary | Gamma | Llama-68M | Llama-160M | Llama-1B | Vicuna-1B | |-------|-----------|------------|----------|-----------| | 0 | ~60 | ~54 | ~50 | ~50 | | 2 | ~68 | ~52 | ~49 | ~50 | | 4 | ~68 | ~47 | ~40 | ~49 | | 6 | ~64 | ~40 | ~34 | ~42 | | 8 | ~57 | ~36 | ~32 | ~36 | | 10 | ~52 | ~31 | ~28 | ~35 | | 12 | ~49 | ~28 | ~24 | ~31 | | 14 | ~47 | ~24 | ~20 | ~27 | --- ## Conclusion The chart highlights performance trade-offs between model size (Llama variants) and efficiency (Gamma parameter). Llama-68M demonstrates the best scalability, while Llama-1B and Vicuna-1B exhibit significant performance degradation at higher Gamma values. </details> (a) Vicuna-7B <details> <summary>2401.10774v3/x9.png Details</summary>  ### Visual Description **Technical Document Extraction: Line Chart Analysis** **Chart Type**: Line chart with four data series. **Axes**: - **X-axis (Horizontal)**: Labeled "Gamma" with integer markers from 0 to 15. - **Y-axis (Vertical)**: Labeled "Tokens per Second" with integer markers from 20 to 55. A dashed horizontal line at **35** is present. **Legend**: - **Blue line**: Llama-68M - **Orange line**: Llama-160M - **Green line**: Llama-1B - **Red line**: Vicuna-1B **Key Trends**: 1. **Llama-68M (Blue)**: - Starts at ~47 tokens/sec at Gamma=0. - Peaks at ~54 tokens/sec at Gamma=3. - Gradually declines to ~38 tokens/sec at Gamma=15. - Maintains the highest performance across all Gamma values. 2. **Llama-160M (Orange)**: - Begins at ~44 tokens/sec at Gamma=0. - Drops sharply to ~35 tokens/sec by Gamma=5. - Continues declining to ~22 tokens/sec at Gamma=15. - Crosses below Llama-1B at Gamma=4. 3. **Llama-1B (Green)**: - Starts at ~39 tokens/sec at Gamma=0. - Declines steadily to ~20 tokens/sec by Gamma=10. - Reaches ~18 tokens/sec at Gamma=15. - Crosses below the 35-token threshold at Gamma=5. 4. **Vicuna-1B (Red)**: - Begins at ~43 tokens/sec at Gamma=0. - Drops to ~35 tokens/sec by Gamma=5. - Continues declining to ~25 tokens/sec at Gamma=15. - Crosses below Llama-160M at Gamma=3. **Critical Observations**: - **Performance Threshold**: The dashed line at 35 tokens/sec acts as a performance benchmark. All models except Llama-68M fall below this threshold by Gamma=8. - **Model Efficiency**: Llama-68M demonstrates superior scalability, retaining higher token generation rates across increasing Gamma values compared to other models. - **Divergence Points**: - Llama-160M and Vicuna-1B intersect near Gamma=3 (~42 tokens/sec). - Llama-1B falls below Llama-160M at Gamma=4 (~38 tokens/sec). **Data Points (Selected)**: - **Llama-68M**: - Gamma=0: 47 - Gamma=3: 54 - Gamma=15: 38 - **Llama-160M**: - Gamma=0: 44 - Gamma=5: 35 - Gamma=15: 22 - **Llama-1B**: - Gamma=0: 39 - Gamma=5: 35 - Gamma=15: 18 - **Vicuna-1B**: - Gamma=0: 43 - Gamma=5: 35 - Gamma=15: 25 **Conclusion**: The chart illustrates a trade-off between model size (Llama variants) and performance efficiency (Tokens per Second) as Gamma increases. Llama-68M maintains dominance, while smaller models (Llama-160M, Llama-1B) and Vicuna-1B exhibit steeper declines, highlighting diminishing returns at higher Gamma values. </details> (b) Vicuna-13B <details> <summary>2401.10774v3/x10.png Details</summary>  ### Visual Description # Technical Document Extraction: Line Chart Analysis ## Labels and Axis Titles - **X-Axis**: Labeled "Gamma" with tick marks at intervals of 2 (0, 2, 4, ..., 14, 15). - **Y-Axis**: Labeled "Tokens per Second" with tick marks at intervals of 2 (16, 18, 20, ..., 28). - **Legend**: Located in the top-right corner, mapping colors to models: - **Blue**: Llama-68M - **Orange**: Llama-160M - **Green**: Llama-1B - **Red**: Vicuna-1B ## Key Trends and Data Points 1. **Vicuna-1B (Red Line)**: - Starts at ~24 tokens/sec (Gamma = 0). - Peaks at ~28.5 tokens/sec (Gamma = 3). - Declines steadily, reaching ~19 tokens/sec (Gamma = 14). 2. **Llama-68M (Blue Line)**: - Starts at ~23 tokens/sec (Gamma = 0). - Peaks at ~27.5 tokens/sec (Gamma = 3). - Declines to ~22.5 tokens/sec (Gamma = 14). 3. **Llama-160M (Orange Line)**: - Starts at ~22.5 tokens/sec (Gamma = 0). - Peaks at ~25.5 tokens/sec (Gamma = 3). - Declines sharply to ~16 tokens/sec (Gamma = 15). 4. **Llama-1B (Green Line)**: - Starts at ~22.5 tokens/sec (Gamma = 0). - Peaks at ~25.5 tokens/sec (Gamma = 3). - Declines steeply to ~15.5 tokens/sec (Gamma = 15). 5. **Cross-Model Behavior**: - All models exhibit a peak performance around Gamma = 3. - Lines intersect between Gamma = 3–4, indicating similar performance in this range. - By Gamma = 14, only Llama-68M remains above the 18 tokens/sec threshold. ## Threshold Line - A dashed horizontal line at **18 tokens/sec** is present, serving as a performance benchmark. All models fall below this threshold by Gamma = 14, except Llama-68M. ## Data Point Summary | Gamma | Llama-68M | Llama-160M | Llama-1B | Vicuna-1B | |-------|-----------|------------|----------|-----------| | 0 | ~23 | ~22.5 | ~22.5 | ~24 | | 3 | ~27.5 | ~25.5 | ~25.5 | ~28.5 | | 14 | ~22.5 | ~16.5 | ~15.5 | ~19 | | 15 | ~22 | ~16 | ~15.5 | ~19 | ## Observations - **Performance Degradation**: All models show a decline in tokens per second after Gamma = 3, with Llama-1B experiencing the steepest drop. - **Model Efficiency**: Llama-68M maintains the highest performance across all Gamma values, while Llama-1B underperforms significantly at higher Gamma values. - **Threshold Breach**: By Gamma = 14, only Llama-68M remains above the 18 tokens/sec benchmark. </details> (c) Vicuna-33B Figure 7: Inference speed of various models using speculative decoding on MT-Bench. Baseline model speeds are presented by grey dotted lines for comparison. $\gamma$ denotes the draft token number. Appendix E Additional Results for All Models We show speedup on various models in Fig. 8. <details> <summary>2401.10774v3/x11.png Details</summary>  ### Visual Description # Technical Document Extraction: Speedup on Different Model Sizes ## Chart Title **Speedup on different model sizes** ## Axis Labels - **X-Axis (Categories):** Model Size - Vicuna-7B - Zephyr-7B - Vicuna-13B - Vicuna-33B - **Y-Axis (Quantitative):** Tokens per Second ## Legend - **Blue Bars:** w/o Medusa - **Orange Bars:** Medusa-2 ## Data Points and Trends 1. **Vicuna-7B** - w/o Medusa: 45 tokens/second - Medusa-2: 130 tokens/second - Speedup: **2.83x** 2. **Zephyr-7B** - w/o Medusa: 40 tokens/second - Medusa-2: 110 tokens/second - Speedup: **2.66x** 3. **Vicuna-13B** - w/o Medusa: 35 tokens/second - Medusa-2: 100 tokens/second - Speedup: **2.83x** 4. **Vicuna-33B** - w/o Medusa: 18 tokens/second - Medusa-2: 45 tokens/second - Speedup: **2.35x** ## Observations - **Speedup Consistency:** - Vicuna-7B and Vicuna-13B exhibit identical speedup multipliers (**2.83x**) despite differing model sizes. - Zephyr-7B shows a slightly lower speedup (**2.66x**) compared to Vicuna variants. - Vicuna-33B has the lowest speedup (**2.35x**), indicating diminishing returns at larger model sizes. - **Performance Gains:** - Medusa-2 consistently outperforms the baseline (w/o Medusa) across all model sizes. - Larger models (e.g., Vicuna-33B) show reduced proportional improvement, suggesting scalability limitations. ## Structural Notes - **Bar Colors:** - Blue (w/o Medusa) and orange (Medusa-2) bars are visually distinct, aligning with the legend. - **Speedup Multipliers:** - Embedded text above orange bars provides direct quantitative comparisons. ## Data Table Reconstruction | Model Size | w/o Medusa (Tokens/sec) | Medusa-2 (Tokens/sec) | Speedup (x) | |--------------|-------------------------|-----------------------|-------------| | Vicuna-7B | 45 | 130 | 2.83 | | Zephyr-7B | 40 | 110 | 2.66 | | Vicuna-13B | 35 | 100 | 2.83 | | Vicuna-33B | 18 | 45 | 2.35 | ## Conclusion The chart demonstrates that Medusa-2 significantly accelerates token generation across all tested model sizes, with speedup diminishing as model complexity increases. Vicuna-7B and Vicuna-13B achieve the highest efficiency gains, while Vicuna-33B exhibits the weakest proportional improvement. </details> Figure 8: Speedup of various models with Medusa -2. Medusa -2 shows significant speed improvement over all the models, while models trained with self-distillation (Zephyr-7B, Vicuna-13/33B) have weaker speedup due to the trade-off between preserving quality and boosting speed. Appendix F Additional Results on AlpacalEval Dataset We conduct further experiments on the AlpacaEval (Li et al., 2023) dataset. Medusa -2 achieves consistent speedup similar to the results on MT-Bench. | Vicuna-7b Vicuna-13b Vicuna-33b | 37.07 29.01 17.87 | 106.76 91.54 40.43 | 3.23 3.28 2.85 | 2.88 3.16 2.26 | | --- | --- | --- | --- | --- | | Zephyr-7b | 34.21 | 99.50 | 3.08 | 2.91 | Table 4: Speedup results on AlpacaEval (Li et al., 2023) dataset. Appendix G Exploration and Modeling of Hardware Constraints and Medusa We explore the hardware constraints, specifically memory-bandwidth bound, and their impact on Medusa -style parallel decoding by incorporating a simplified Llama-series model. First, we identify that the operators involving matrix multiplications, such as linear layers and attention matrix multiplications, are the primary sources of overhead. We profile the performance of FLOP/s vs. Operational Intensity which is the ratio of FLOP/s to bandwidth (bytes/s), across various GPUs, including the A100-80GB-PCIe, A40, and A6000. Next, we examine the changes in FLOP/s vs. Operational Intensity when using Medusa for different operators. Finally, we apply a straightforward analytical model to calculate acceleration rates and combine it with hardware benchmarks. This provides insights into the effects under different model sizes, sequence lengths, and batch sizes. G.1 Roofline Model of Operators We present an analysis of the roofline model for various operators in large language models (LLMs), specifically focusing on Llama-7B, Llama-13B, and Llama-33B (Touvron et al., 2023). These models were benchmarked on different GPUs, including the A100-80GB-PCIe, A40, and A6000. We looked into the three categories of matrix multiplication operators since they represent the primary sources of computational overhead in these models. Our study follows the report (Chen, 2023) which investigates the effectiveness of batch size but ours focuses more on decoding and parallel decoding. Table 5 details the computation and space complexity for each operator during the prefill, decoding, and Medusa decoding phases. The operators include the linear layers for query, key, and value matrices ( $XW_{Q}$ , $XW_{K}$ , $XW_{V}$ ), the attention matrix multiplications ( $QK^{T}$ , $PV$ ), and the up/gate/down linear layers ( $XW_{u}$ , $XW_{g}$ , $XW_{d}$ ). $b$ stands for the batch size, $s$ stands for the sequence length, $h$ stands for the hidden dimension, $i$ stands for the intermediate dimension, $n$ stands for the number of attention heads, $d$ stands for the head dimension and $q$ stands for the candidate length for Medusa. For more details of these operators please refer to the articles (Touvron et al., 2023; Chen, 2023). Table 5: Computational and space complexity of the main operators in different phases. The table is based on Table 2 in the report (Chen, 2023). | Operator Prefill $XW_{Q}$ , $XW_{K}$ , $XW_{V}$ | Input Shape $(b,s,h)$ | Output Shape $(b,s,h)$ | Comp. Complexity $O(bsh^{2})$ | Space Complexity $O(2bsh+h^{2})$ | | --- | --- | --- | --- | --- | | $QK^{T}$ | $(b,n,s,d),(b,n,s,d)$ | $(b,n,s,s)$ | $O(bs^{2}nd)$ | $O(2bsnd+bs^{2}n)$ | | $PV$ | $(b,n,s,s),(b,n,s,d)$ | $(b,n,s,d)$ | | | | $XW_{u}$ , $XW_{g}$ | $(b,s,h)$ | $(b,s,i)$ | $O(bshi)$ | $O(bs(h+i)+hi)$ | | $XW_{d}$ | $(b,s,i)$ | $(b,s,h)$ | | | | Decoding | | | | | | $XW_{Q}$ , $XW_{K}$ , $XW_{V}$ | $(b,1,h)$ | $(b,1,h)$ | $O(bh^{2})$ | $O(2bh+h^{2})$ | | $QK^{T}$ | $(b,n,1,d),(b,n,s,d)$ | $(b,n,s,1)$ | $O(bsnd)$ | $O(bsn+bsnd+bnd)$ | | $PV$ | $(b,n,s,1),(b,n,1,d)$ | $(b,n,1,d)$ | | | | $XW_{u}$ , $XW_{g}$ | $(b,1,h)$ | $(b,1,i)$ | $O(bhi)$ | $O(b(h+i)+hi)$ | | $XW_{d}$ | $(b,1,i)$ | $(b,1,h)$ | | | | Parallel decoding | | | | | | $XW_{Q}$ , $XW_{K}$ , $XW_{V}$ | $(b,q,h)$ | $(b,q,h)$ | $O(bqh^{2})$ | $O(2bqh+h^{2})$ | | $QK^{T}$ | $(b,n,q,d),(b,n,s,d)$ | $(b,n,s,q)$ | $O(bsqnd)$ | $O(bsqn+b(s+q)nd)$ | | $PV$ | $(b,n,s,q),(b,n,q,d)$ | $(b,n,q,d)$ | | | | $XW_{u}$ , $XW_{g}$ | $(b,q,h)$ | $(b,q,i)$ | $O(bqhi)$ | $O(bq(h+i)+hi)$ | | $XW_{d}$ | $(b,q,i)$ | $(b,q,h)$ | | | Figures 9 - 17 show the benchmark of three categories of operators on different models (7/13/33B) under various settings. To evaluate each operator’s performance and throughput, we chose the combination of settings including batch sizes from 1 to 64 in powers of 2 and sequence lengths from 128 to 8192 in powers of 2 (49 settings for each operator). From all the figures, we observe that the datapoints of each operator in the prefill and decoding stages cluster at very similar positions across all GPUs and for various model sizes. During the prefill phase, increasing the batch size changes the FLOP/s of the attention matrix multiplications (see ‘qk/pv init‘) but does not affect the Operational Intensity (refer to the vertical dashed arrow in Fig. 9). In contrast, increasing the sequence length impacts both FLOP/s and Operational Intensity in the prefill phase (refer to the diagonal dashed arrow in Fig. 9). During the decoding phase, the attention matrix multiplications are significantly limited by memory bandwidth. Despite an increase in FLOP/s with changes in batch size and sequence length, the Operational Intensity remains nearly unchanged (see ‘qk/pv ar‘). This indicates suboptimal resource utilization in the self-attention mechanism. The linear layers in the prefill phase are mostly compute-bound (see ‘qkv mlp init‘ and ‘up/gate/down init‘). During the decoding phase, the datapoints of the linear layer form a line with the same slope as the GPU’s memory bandwidth (see ‘qkv mlp ar‘ and ‘up/gate/down ar‘). This indicates the linear layers in the decoding stage are also bounded by memory bandwidth. Increasing the batch size improves the achieved FLOP/s and Operational Intensity under memory bandwidth constraints through better parallelism. Note that linear layers only process the new token and are independent of sequence length (See ‘Decoding‘ section in Table 5). <details> <summary>2401.10774v3/x12.png Details</summary>  ### Visual Description # Roofline Model Analysis (Llama 7B, A100 80GB PCIe) ## Graph Title - **Title**: Roofline Model (Llama 7B, A100 80GB PCIe) ## Axes - **X-axis**: Operational Intensity (FLOP/Byte) - Range: 1 to 10,000 (logarithmic scale) - Key markers: - Vertical green dashed line at **100 FLOP/Byte** - **Y-axis**: Performance (FLOP/s) - Range: 10G to 100T (logarithmic scale) - Key markers: - Horizontal red dashed line at **312 TFLOP/s** - Horizontal blue dashed line at **1,935GB/s** ## Legend | Symbol | Label | Color/Marker | |--------|------------------------|--------------| | `x` | qkv mlp init | Blue | | `x` | qkv mlp ar | Orange | | `x` | up/gate/down init | Green | | `x` | up/gate/down ar | Red | | `x` | qk/pv init | Purple | | `x` | qk/pv ar | Brown | ## Key Trends 1. **Roofline Boundaries**: - **Compute Limit**: Red dashed line at **312 TFLOP/s** (theoretical peak performance). - **Memory Bandwidth Limit**: Blue dashed line at **1,935GB/s** (theoretical memory throughput). - **Roofline**: Diagonal blue dashed line representing the performance-memory bandwidth tradeoff. 2. **Data Points**: - Performance improves with increasing operational intensity until reaching the Roofline. - Beyond **100 FLOP/Byte** (green dashed line), performance plateaus due to memory bandwidth constraints. 3. **Annotations**: - **qk/pv init Increase bs**: Indicates performance gains from increasing batch size (`bs`). - **Increase seq_len**: Highlights performance improvements from longer sequence lengths. ## Observations - **qkv mlp init** (blue `x`) and **qk/pv init** (purple `x`) show steep performance gains at low operational intensity. - **up/gate/down ar** (red `x`) and **qk/pv ar** (brown `x`) exhibit higher performance at elevated operational intensity. - All data points cluster below the Roofline, confirming hardware limitations. </details> Figure 9: The figure shows the relationship between FLOP/s and Operational Intensity for all benchmarked datapoints of Llama-7B operators on A100-80GB-PCIe. The dashed lines represent the HBM bandwidth limit (1,935GB/s) and the peak performance limit (312 TFLOP/s) (NVIDIA, ). ‘ qkv mlp ’ stands for the linear layers projecting hidden features to query/key/value features. ‘ up/gate/down ’ stands for the linear layers following the attention block. ‘ qk/pv ’ stands for the two steps of attention matrix multiplications. ‘ ar ’ stands for the decoding (autoregressive) and ‘ init ’ stands for the prefill phase. <details> <summary>2401.10774v3/x13.png Details</summary>  ### Visual Description # Technical Document Extraction: Roofline Model Analysis ## Chart Title **Roofline Model (Llama 13B, A100 80GB PCIe)** ## Axes Labels - **X-Axis**: Operational Intensity (FLOP/Byte) - Range: 1 to 10,000 (logarithmic scale) - Key marker: Green dashed vertical line at 100 FLOP/Byte - **Y-Axis**: Performance (FLOP/s) - Range: 10G to 100T (logarithmic scale) - Key markers: - Blue dashed line labeled "1,935GB/s" - Red dashed line labeled "312 TFLOP/s" ## Legend Entries | Label | Color | Marker | Description | |----------------------|--------|--------|----------------------| | qkv mlp init | Blue | X | Initial qkv mlp | | qkv mlp ar | Orange | X | qkv mlp activation | | up/gate/down init | Green | X | Initial up/gate/down | | up/gate/down ar | Red | X | up/gate/down activation | | qk/pv init | Purple | X | Initial qk/pv | | qk/pv ar | Brown | X | qk/pv activation | ## Key Trends 1. **Performance Scaling**: - Data points cluster along the blue dashed line ("1,935GB/s"), indicating memory-bound operations dominate performance at lower operational intensities. - Beyond the green dashed line (100 FLOP/Byte), performance aligns with the red dashed line ("312 TFLOP/s"), suggesting compute-bound operations dominate. 2. **Operational Intensity Threshold**: - The green dashed line at 100 FLOP/Byte demarcates the transition from memory-bound to compute-bound performance regimes. 3. **Data Series Behavior**: - All series (qkv mlp, up/gate/down, qk/pv) follow similar scaling patterns, with slight variations in performance at specific intensity thresholds. ## Additional Notes - The chart uses a logarithmic scale for both axes to visualize performance across orders of magnitude. - The red dashed line ("312 TFLOP/s") represents the theoretical peak performance of the A100 GPU, while the blue dashed line ("1,935GB/s") reflects memory bandwidth constraints. - Data points are plotted as "X" markers, with colors corresponding to their respective operations (e.g., qkv mlp, up/gate/down). </details> Figure 10: Llama-13B operators on A100-80GB-PCIe. <details> <summary>2401.10774v3/x14.png Details</summary>  ### Visual Description # Technical Document Extraction: Roofline Model Analysis ## Chart Title **Roofline Model (Llama 33B, A100 80GB PCIe)** ## Axes - **X-Axis**: Operational Intensity (FLOP/Byte) - Range: 1 to 10,000 (logarithmic scale) - Grid lines: Dashed vertical lines at 1, 10, 100, 1,000 - **Y-Axis**: Performance (FLOP/s) - Range: 10G to 100T (logarithmic scale) - Grid lines: Dashed horizontal lines at 10G, 100G, 1T, 10T, 100T ## Key Lines 1. **Blue Dashed Line**: - Label: `1,935GB/s` - Represents memory bandwidth limit (Roofline boundary) 2. **Red Dashed Line**: - Label: `312 TFLOP/s` - Represents compute limit (Roofline boundary) 3. **Intersection**: - Point where blue and red lines meet (Roofline) ## Legend | Color/Symbol | Label | Marker Type | |--------------|------------------------|-------------| | Blue | `qkv mlp init` | X | | Orange | `qkv mlp ar` | X | | Green | `up/gate/down init` | X | | Red | `up/gate/down ar` | X | | Purple | `qk/pv init` | X | | Brown | `qk/pv ar` | X | ## Data Points - **Blue Xs**: - Clustered near the blue dashed line (memory-bound operations) - **Orange Xs**: - Distributed along the blue dashed line (memory-bound operations) - **Green Xs**: - Clustered near the red dashed line (compute-bound operations) - **Red Xs**: - Distributed along the red dashed line (compute-bound operations) - **Purple Xs**: - Distributed along the blue dashed line (memory-bound operations) - **Brown Xs**: - Clustered at lower operational intensity (memory-bound operations) ## Key Trends 1. **Roofline Boundary**: - The intersection of the blue (`1,935GB/s`) and red (`312 TFLOP/s`) lines defines the theoretical maximum performance limit. 2. **Operational Intensity vs. Performance**: - Data points below the Roofline indicate suboptimal utilization of memory bandwidth or compute resources. - Points near the Roofline (e.g., green and red Xs) represent compute-bound operations. - Points near the blue dashed line (e.g., blue, orange, purple Xs) represent memory-bound operations. 3. **Performance Scaling**: - Performance increases logarithmically with operational intensity until hitting the Roofline. ## Additional Notes - **Grid Lines**: - Dashed lines at powers of 10 for both axes to aid logarithmic interpretation. - **Legend Consistency**: - Colors and markers in the legend match the data points and lines accurately. - **Model Context**: - Focuses on Llama 33B model performance on A100 80GB PCIe GPUs. </details> Figure 11: Llama-33B operators on A100-80GB-PCIe. <details> <summary>2401.10774v3/x15.png Details</summary>  ### Visual Description # Roofline Model Analysis (Llama 7B, A40) ## Graph Structure - **Title**: Roofline Model (Llama 7B, A40) - **X-Axis**: Operational Intensity (FLOP/Byte) - Logarithmic scale from 1 to 10,000 (10k) - **Y-Axis**: Performance (FLOP/s) - Logarithmic scale from 10G (10^9) to 100T (10^14) ## Key Thresholds 1. **Blue Dashed Line**: 696GB/s (Memory bandwidth limit) 2. **Red Dashed Line**: 149.7 TFLOP/s (Theoretical peak performance) 3. **Green Dashed Line**: 100 TFLOP/s (Critical performance threshold) ## Legend Entries | Symbol | Label | Color | |--------|------------------------|--------| | `x` | qkv mlp init | Blue | | `x` | qkv mlp ar | Orange | | `x` | up/gate/down init | Green | | `x` | up/gate/down ar | Red | | `x` | qk/pv init | Purple | | `x` | qk/pv ar | Brown | ## Data Trends 1. **Performance Scaling**: - All operations follow the roofline model, with performance increasing linearly with operational intensity until hitting memory bandwidth limits (696GB/s). - Beyond 100 TFLOP/s (green line), performance plateaus due to hardware constraints. 2. **Operation-Specific Performance**: - **qkv mlp init** (blue `x`): Reaches ~100T FLOP/s at ~100 TFLOP/Byte intensity. - **qkv mlp ar** (orange `x`): Peaks near 10T FLOP/s at ~10 TFLOP/Byte. - **up/gate/down init** (green `x`): Exceeds 100T FLOP/s at ~1k TFLOP/Byte. - **up/gate/down ar** (red `x`): Crosses 100T FLOP/s at ~100 TFLOP/Byte. - **qk/pv init** (purple `x`): Achieves ~10T FLOP/s at ~10 TFLOP/Byte. - **qk/pv ar** (brown `x`): Remains below 10G FLOP/s across all intensities. 3. **Critical Observations**: - Operations with "init" suffixes generally achieve higher performance than their "ar" counterparts. - The 100 TFLOP/s threshold (green line) acts as a performance ceiling for most operations. - Memory-bound operations (e.g., up/gate/down init) dominate at high operational intensities. ## Technical Implications - The roofline model illustrates the trade-off between computational intensity and memory bandwidth. - Optimizing for higher operational intensity (FLOP/Byte) is critical to approach the 149.7 TFLOP/s theoretical limit. - Memory-bound operations (e.g., up/gate/down init) are bottlenecked by the 696GB/s memory bandwidth. </details> Figure 12: Llama-7B operators on A40. <details> <summary>2401.10774v3/x16.png Details</summary>  ### Visual Description # Technical Document Extraction: Roofline Model (Llama 13B, A40) ## Chart Overview - **Title**: Roofline Model (Llama 13B, A40) - **Type**: Scatter plot with trend lines - **Purpose**: Visualizes performance (FLOP/s) vs. operational intensity (FLOP/Byte) for different computational operations ## Axes - **X-axis (Operational Intensity)**: - Label: "Operational Intensity (FLOP/Byte)" - Range: 1 to 10,000 (logarithmic scale) - Tick markers: 1, 10, 100, 1k - **Y-axis (Performance)**: - Label: "Performance (FLOP/s)" - Range: 10G to 100T (logarithmic scale) - Tick markers: 10G, 100G, 1T, 10T, 100T ## Legend - **Location**: Bottom-right corner - **Color/Marker Legend**: - **Blue dashed line**: 696GB/s (theoretical memory bandwidth limit) - **Red dashed line**: 149.7 TFLOP/s (theoretical compute limit) - **Blue x**: qkv mlp init - **Orange x**: qkv mlp ar - **Green x**: up/gate/down init - **Red x**: up/gate/down ar - **Purple x**: qk/pv init - **Brown x**: qk/pv ar ## Key Trends 1. **Theoretical Limits**: - Horizontal red line at 149.7 TFLOP/s (compute-bound ceiling) - Vertical green line at 696GB/s (memory-bound ceiling) 2. **Data Series Trends**: - **qkv mlp init** (blue x): - Starts at ~1 FLOP/Byte, 100G FLOP/s - Ends at ~1000 FLOP/Byte, 100T FLOP/s - Linear upward trajectory - **qkv mlp ar** (orange x): - Starts at ~10 FLOP/Byte, 100G FLOP/s - Ends at ~1000 FLOP/Byte, 100T FLOP/s - Linear upward trajectory - **up/gate/down init** (green x): - Starts at ~100 FLOP/Byte, 10T FLOP/s - Ends at ~1000 FLOP/Byte, 100T FLOP/s - Linear upward trajectory - **up/gate/down ar** (red x): - Starts at ~10 FLOP/Byte, 10T FLOP/s - Ends at ~1000 FLOP/Byte, 100T FLOP/s - Linear upward trajectory - **qk/pv init** (purple x): - Starts at ~10 FLOP/Byte, 10T FLOP/s - Ends at ~1000 FLOP/Byte, 100T FLOP/s - Linear upward trajectory - **qk/pv ar** (brown x): - Starts at ~1 FLOP/Byte, 100G FLOP/s - Ends at ~1000 FLOP/Byte, 100T FLOP/s - Linear upward trajectory ## Spatial Grounding - **Legend Position**: Bottom-right quadrant - **Data Point Verification**: - All blue x markers align with "qkv mlp init" legend - All orange x markers align with "qkv mlp ar" legend - All green x markers align with "up/gate/down init" legend - All red x markers align with "up/gate/down ar" legend - All purple x markers align with "qk/pv init" legend - All brown x markers align with "qk/pv ar" legend ## Critical Observations 1. All operations approach the theoretical compute limit (149.7 TFLOP/s) as operational intensity increases 2. Memory-bound operations (qkv mlp init, qk/pv ar) start at lower operational intensities 3. Compute-bound operations (up/gate/down init, qk/pv init) require higher operational intensities to reach peak performance 4. The roofline model demonstrates clear separation between memory-bound and compute-bound operations ## Data Table Reconstruction | Operation Type | Initial Intensity (FLOP/Byte) | Initial Performance (FLOP/s) | Final Intensity (FLOP/Byte) | Final Performance (FLOP/s) | |----------------------|-------------------------------|------------------------------|-----------------------------|----------------------------| | qkv mlp init | 1 | 100G | 1000 | 100T | | qkv mlp ar | 10 | 100G | 1000 | 100T | | up/gate/down init | 100 | 10T | 1000 | 100T | | up/gate/down ar | 10 | 10T | 1000 | 100T | | qk/pv init | 10 | 10T | 1000 | 100T | | qk/pv ar | 1 | 100G | 1000 | 100T | ## Trend Verification - All data series show linear scaling with operational intensity - Memory-bound operations (blue x, brown x) cluster near the memory bandwidth limit (696GB/s) - Compute-bound operations (green x, purple x) cluster near the compute limit (149.7 TFLOP/s) - All operations converge at the theoretical maximum performance point (1000 FLOP/Byte, 100T FLOP/s) </details> Figure 13: Llama-13B operators on A40. <details> <summary>2401.10774v3/x17.png Details</summary>  ### Visual Description # Roofline Model (Llama 33B, A40) Analysis ## Axes and Labels - **X-Axis**: Operational Intensity (FLOP/Byte) - Range: 1 to 10,000 (logarithmic scale) - Gridlines: Logarithmic spacing - **Y-Axis**: Performance (FLOP/s) - Range: 10G (10^9) to 100T (10^14) (logarithmic scale) - Gridlines: Logarithmic spacing ## Legend and Key Trends 1. **Memory Bandwidth Limit** - **Line**: Dashed blue line - **Value**: 696GB/s - **Interpretation**: Represents the maximum data transfer rate (memory-bound performance ceiling). 2. **Compute Limit** - **Line**: Red dashed line - **Value**: 149.7 TFLOP/s - **Interpretation**: Theoretical peak performance (compute-bound ceiling). 3. **Data Points** - **Markers**: Colored "X" symbols for different operations: - **Blue**: `qkv mlp init` - **Orange**: `qkv mlp ar` - **Green**: `up/gate/down init` - **Red**: `up/gate/down ar` - **Purple**: `qk/pv init` - **Brown**: `qk/pv ar` ## Performance Trends - **Operational Intensity vs. Performance**: - Performance increases linearly with operational intensity until it reaches the **memory bandwidth limit** (696GB/s). - Beyond this point, performance plateaus, constrained by memory bandwidth. - The red dashed line (149.7 TFLOP/s) represents the theoretical maximum performance, which is not exceeded by any data point. ## Critical Observations - **Memory-Bound Operations**: - Data points for `qkv mlp init`, `qkv mlp ar`, and `up/gate/down init` cluster near the memory bandwidth limit, indicating these operations are memory-bound. - **Compute-Bound Operations**: - `up/gate/down ar`, `qk/pv init`, and `qk/pv ar` approach but do not exceed the compute limit (149.7 TFLOP/s). - **Efficiency**: - The roofline model illustrates the trade-off between operational intensity and performance, highlighting hardware constraints (memory bandwidth and compute capacity). ## Grid and Annotations - **Dashed Lines**: - Vertical green dashed line at 696GB/s (memory bandwidth). - Horizontal red dashed line at 149.7 TFLOP/s (compute limit). - **Grid**: Logarithmic scale for both axes to visualize performance across orders of magnitude. </details> Figure 14: Llama-33B operators on A40. <details> <summary>2401.10774v3/x18.png Details</summary>  ### Visual Description # Technical Document Extraction: Roofline Model (Llama 7B, A6000) ## Header - **Title**: Roofline Model (Llama 7B, A6000) ## Main Chart ### Axes - **X-axis**: Operational Intensity (FLOP/Byte) - Range: 1 to 10k (logarithmic scale) - Key markers: - Green dashed vertical line at **128 FLOP/Byte** - **Y-axis**: Performance (FLOP/s) - Range: 10G to 100T (logarithmic scale) - Key markers: - Red dashed horizontal line at **181 TFLOP/s** - Blue dashed diagonal line from origin to **(128 FLOP/Byte, 181 TFLOP/s)** ### Data Series #### Legend (Right Side) | Color | Marker | Label | |-------|--------|------------------------| | Blue | X | `qkv mlp init` | | Orange| X | `qkv mlp ar` | | Green | X | `up/gate/down init` | | Red | X | `up/gate/down ar` | | Purple| X | `qk/pv init` | | Brown | X | `qk/pv ar` | #### Visual Trends 1. **Blue Dashed Line** (`768GB/s`): - Slope: Linear upward from origin to **(128 FLOP/Byte, 181 TFLOP/s)** - Represents memory bandwidth limit. 2. **Red Dashed Line** (`181 TFLOP/s`): - Horizontal line at **181 TFLOP/s** (peak performance threshold). 3. **Green Dashed Line** (`128 FLOP/Byte`): - Vertical line at **128 FLOP/Byte** (operational intensity threshold). 4. **Data Points**: - All series follow the blue dashed line until **128 FLOP/Byte**, then plateau. - Example: - `qkv mlp init` (blue X): Peaks at ~181 TFLOP/s at 128 FLOP/Byte. - `qk/pv ar` (brown X): Remains below 10T FLOP/s across all intensities. ### Spatial Grounding - **Legend Position**: Right side of the chart. - **Color Consistency**: - Blue X = `qkv mlp init` (matches blue dashed line). - Orange X = `qkv mlp ar` (distinct from blue). - Green X = `up/gate/down init` (distinct from red). ## Footer - **Legend Labels**: - `qkv mlp init` (blue X) - `qkv mlp ar` (orange X) - `up/gate/down init` (green X) - `up/gate/down ar` (red X) - `qk/pv init` (purple X) - `qk/pv ar` (brown X) ## Key Observations 1. **Performance Bottleneck**: - All workloads hit the **181 TFLOP/s** limit at **128 FLOP/Byte**, indicating memory bandwidth constraints. 2. **Workload Efficiency**: - `qk/pv ar` (brown X) operates far below the roofline, suggesting inefficiency. 3. **Thresholds**: - **128 FLOP/Byte** (green line) and **181 TFLOP/s** (red line) define the model's operational limits. </details> Figure 15: Llama-7B operators on A6000. <details> <summary>2401.10774v3/x19.png Details</summary>  ### Visual Description # Roofline Model Analysis (Llama13B, A6000) ## Title **Roofline Model (Llama 13B, A6000)** ## Axes - **X-axis**: Operational Intensity (FLOP/Byte) - Range: 1 to 10,000 (logarithmic scale) - Key marker: Green dashed vertical line at **100 FLOP/Byte** - **Y-axis**: Performance (FLOP/s) - Range: 10G to 100T (logarithmic scale) - Key markers: - Red dashed horizontal line at **181 TFLOP/s** - Blue dashed horizontal line at **768GB/s** ## Legend | Symbol | Label | Color | Marker | |--------|------------------------|--------|--------| | Blue | 768GB/s | Dashed | N/A | | Red | 181 TFLOP/s | Dashed | N/A | | Blue | qkv mlp init | Cross | X | | Orange | qkv mlp ar | Cross | X | | Green | up/gate/down init | Cross | X | | Red | up/gate/down ar | Cross | X | | Purple | qk/pv init | Cross | X | | Brown | qk/pv ar | Cross | X | ## Data Points - **qkv mlp init** (Blue crosses): - Clustered below 181 TFLOP/s line, increasing with operational intensity. - **qkv mlp ar** (Orange crosses): - Similar trend to qkv mlp init, slightly higher performance at mid-range intensities. - **up/gate/down init** (Green crosses): - Highest performance among non-ar operations, approaching 181 TFLOP/s at 100 FLOP/Byte. - **up/gate/down ar** (Red crosses): - Dominates performance above 100 FLOP/Byte, consistently near 181 TFLOP/s. - **qk/pv init** (Purple crosses): - Lower performance than qkv operations, plateauing below 10T FLOP/s. - **qk/pv ar** (Brown crosses): - Similar to qk/pv init, with minimal improvement at higher intensities. ## Key Trends 1. **Performance Ceiling**: - The red dashed line (181 TFLOP/s) acts as a theoretical maximum for most operations. - Only **up/gate/down ar** (red crosses) approaches this limit at high operational intensities (>100 FLOP/Byte). 2. **Memory Bandwidth Constraint**: - The blue dashed line (768GB/s) represents memory bandwidth. Most data points fall below this line, indicating compute-bound operations. 3. **Operational Intensity Threshold**: - The green dashed line at 100 FLOP/Byte marks a critical transition point. - Above this threshold, **up/gate/down ar** achieves near-peak performance, while other operations plateau. 4. **Operation Efficiency**: - **qkv mlp ar** and **up/gate/down ar** show the highest compute-to-memory efficiency. - **qk/pv** operations (init/ar) are less efficient, remaining in the lower-left quadrant. ## Observations - **Scalability**: - Performance scales logarithmically with operational intensity for most operations. - **Bottlenecks**: - Memory-bound operations (e.g., qk/pv) are limited by bandwidth (768GB/s). - Compute-bound operations (e.g., up/gate/down ar) hit the 181 TFLOP/s ceiling. </details> Figure 16: Llama-13B operators on A6000. <details> <summary>2401.10774v3/x20.png Details</summary>  ### Visual Description # Technical Document Extraction: Roofline Model (Llama 33B, A6000) ## Chart Overview - **Title**: Roofline Model (Llama 33B, A6000) - **Type**: Scatter plot with trend lines - **Purpose**: Visualizes computational performance vs. operational intensity for various operations on Llama 33B using A6000 GPU. --- ## Axis Labels and Ranges - **X-axis (Operational Intensity)**: - Units: FLOP/Byte - Range: 1 to 10,000 (logarithmic scale) - Key markers: - Vertical green dashed line at 100 FLOP/Byte - **Y-axis (Performance)**: - Units: FLOP/s - Range: 10G (10⁷) to 100T (10¹⁴) - Key markers: - Horizontal red dashed line at 181 TFLOP/s (1.81×10¹³) - Horizontal blue dashed line at 768GB/s (7.68×10⁸) --- ## Legend and Data Series | Color | Marker | Label | Description | |-------|--------|------------------------|--------------------------------------| | Blue | X | qkv mlp init | Query-Key-Value Matrix Multiplication Initialization | | Orange| X | qkv mlp ar | Query-Key-Value Matrix Multiplication Activation | | Green | X | up/gate/down init | Up/Down Gate Initialization | | Red | X | up/gate/down ar | Up/Down Gate Activation | | Purple| X | qk/pv init | Query-Key/Pair Value Initialization | | Brown | X | qk/pv ar | Query-Key/Pair Value Activation | --- ## Key Trends and Observations 1. **Performance Scaling**: - Performance increases exponentially with operational intensity (log-log scale). - Blue dashed line (768GB/s) represents memory bandwidth ceiling. - Red dashed line (181 TFLOP/s) represents peak FLOP performance ceiling. 2. **Data Point Distribution**: - **qkv mlp init/ar**: - Initialization (blue) and activation (orange) points cluster below the 100 TFLOP/s line. - Activation points show higher performance than initialization. - **up/gate/down init/ar**: - Green (init) and red (ar) points cluster near the 100 TFLOP/s ceiling. - **qk/pv init/ar**: - Purple (init) and brown (ar) points show lower performance, concentrated below 10T FLOP/s. 3. **Critical Thresholds**: - **100 FLOP/Byte (green line)**: Separates memory-bound (left) vs. compute-bound (right) operations. - **181 TFLOP/s (red line)**: Maximum achievable performance for compute-bound operations. - **768GB/s (blue line)**: Memory bandwidth limit for data movement. 4. **Performance Bottlenecks**: - Operations crossing the red line (181 TFLOP/s) are compute-bound. - Operations below the red line but above the blue line are memory-bound. - Operations below both lines are limited by both compute and memory. --- ## Technical Implications - **Optimization Opportunities**: - Operations near the red line (e.g., up/gate/down ar) are near peak performance. - qk/pv operations (purple/brown) show significant room for optimization. - **Hardware Constraints**: - Memory bandwidth (768GB/s) limits data-intensive operations. - FLOP ceiling (181 TFLOP/s) caps compute-bound performance. --- ## Data Point Density - Highest density of points in the 10–100 TFLOP/s range (y-axis) and 10–1000 FLOP/Byte range (x-axis). - Sparse data points above 100 TFLOP/s, indicating few operations reach peak performance. --- ## Conclusion The chart demonstrates the trade-offs between computational intensity and memory bandwidth for Llama 33B operations on A6000. Key insights include: 1. Most operations are memory-bound (below 100 FLOP/Byte). 2. Up/gate/down activation (red) approaches the FLOP ceiling. 3. Query-Key/Pair Value operations (purple/brown) are the most memory-bound. </details> Figure 17: Llama-33B operators on A6000. G.2 FLOP/s vs. Operational Intensity Variations in Medusa We investigate how Medusa can change Operational Intensity and elevate the FLOP/s. We choose Llama 33B on A100-80GB-PCIe as the setting. First, we examine the attention matrix multiplication. Fig. 18 and Table 6 illustrate the effects of Medusa while keeping the batch size fixed at 16. We observe increased FLOP/s and Operational Intensity as more candidate tokens are added (original decoding results are plotted as grey dots). This indicates that Medusa can leverage additional candidate tokens to improve computational throughput. Compared to regular decoding, Medusa achieves 44 $×$ FLOP/s and 41 $×$ Operational Intensity under the setting of batch size 16 and sequence length 1024 with 64 candidate tokens. Fig. 19 and Table 7 illustrate the effects of Medusa decoding while keeping the sequence length fixed at 1024. Increasing the batch size does not improve Operational Intensity in this scenario. Next, we examine the linear layer, focusing on the up/gate/down linear layers. The results are shown in Fig. 20 and Table 8. Since the linear layers in the decoding phase only process the future tokens while the past tokens are cached, they are independent of the sequence length. We vary the batch size to observe the effects. As Medusa increases the number of candidate tokens with the increasing batch size, we observe a shift from a memory-bandwidth-bound region to a computation-bound region. This shift demonstrates how Medusa can transition the performance characteristics of the linear layers from being limited by memory bandwidth to being limited by computational capacity. <details> <summary>2401.10774v3/x21.png Details</summary>  ### Visual Description # Technical Analysis of Llama 33B Performance on A100 80GB PCIe ## Graph Title **Llama 33B, A100 80GB PCIe** ## Axes - **X-axis**: Operational Intensity (FLOP/Byte) - Markers: 1, 10, 100, 1k, 10k - **Y-axis**: Performance (FLOP/s) - Markers: 10G, 100G, 1T, 10T, 100T ## Legend 1. **Dashed Blue Line**: 1,935GB/s 2. **Dashed Red Line**: 312 TFLOP/s 3. **Data Points**: - **Gray Dot**: qk/pv ar - **Colored Dots**: - Orange: qk/pv Medusa (# cand.: 16) - Light Orange: qk/pv Medusa (# cand.: 32) - Red-Orange: qk/pv Medusa (# cand.: 48) - Pink: qk/pv Medusa (# cand.: 64) - Purple: qk/pv Medusa (# cand.: 80) - Dark Purple: qk/pv Medusa (# cand.: 96) - Violet: qk/pv Medusa (# cand.: 112) ## Key Trends 1. **Performance Scaling**: - Performance increases linearly with operational intensity up to **100 FLOP/Byte**, following the dashed blue line (1,935GB/s). - Beyond 100 FLOP/Byte, performance plateaus at **312 TFLOP/s** (dashed red line). 2. **Configuration Performance**: - **qk/pv ar** (gray dot) achieves ~1T FLOP/s at ~1 FLOP/Byte. - **qk/pv Medusa** configurations show incremental performance gains with higher candidate counts: - 16 candidates: ~10T FLOP/s at ~10 FLOP/Byte. - 112 candidates: ~100T FLOP/s at ~100 FLOP/Byte. 3. **Critical Thresholds**: - **Green Vertical Line**: Marks 100 FLOP/Byte (operational intensity limit). - **Red Horizontal Line**: Represents maximum achievable performance (312 TFLOP/s). ## Observations - All configurations converge to the 312 TFLOP/s ceiling at operational intensities ≥100 FLOP/Byte. - Higher candidate counts in qk/pv Medusa improve performance but require proportional increases in operational intensity. - The "qk/pv ar" baseline underperforms compared to Medusa configurations across all intensities. </details> Figure 18: FLOP/s vs. Operational Intensity of attention matrix multiplication with batch size 16. <details> <summary>2401.10774v3/x22.png Details</summary>  ### Visual Description # Technical Analysis: Llama 33B, A100 80GB PCIe Performance Chart ## Chart Overview - **Title**: Llama 33B, A100 80GB PCIe - **Axes**: - **Y-Axis**: Performance (FLOP/s) [Logarithmic scale: 10G → 100T] - **X-Axis**: Operational Intensity (FLOP/Byte) [Logarithmic scale: 1 → 10k] ## Key Elements 1. **Performance Thresholds**: - **Blue Dashed Line**: 1,935GB/s (Memory bandwidth limit) - **Red Dashed Line**: 312 TFLOP/s (Compute limit) - **Green Vertical Line**: 100 FLOP/Byte (Operational intensity threshold) 2. **Data Series**: - **qk/pv ar** (Gray markers): Single data point at ~1 FLOP/Byte, ~1T FLOP/s performance. - **qk/pv Medusa** (Colored markers): - **16 candidates**: Orange (100G FLOP/s at ~10 FLOP/Byte) - **32 candidates**: Orange (1T FLOP/s at ~100 FLOP/Byte) - **48 candidates**: Red (10T FLOP/s at ~1k FLOP/Byte) - **64 candidates**: Pink (100T FLOP/s at ~10k FLOP/Byte) - **80 candidates**: Purple (100T FLOP/s at ~10k FLOP/Byte) - **96 candidates**: Purple (100T FLOP/s at ~10k FLOP/Byte) - **112 candidates**: Violet (100T FLOP/s at ~10k FLOP/Byte) ## Trends - **Performance Scaling**: - Performance increases linearly with operational intensity until hitting the **312 TFLOP/s** compute limit (red line). - Memory bandwidth (**1,935GB/s**, blue line) becomes the bottleneck beyond ~100 FLOP/Byte. - **Medusa Efficiency**: - Higher candidate counts (16 → 112) show diminishing returns in performance gains. - All Medusa configurations plateau at **100T FLOP/s** near the compute limit. ## Critical Observations - The **green vertical line** at 100 FLOP/Byte marks the transition point where compute limits dominate over memory bandwidth constraints. - **qk/pv ar** (gray) represents a baseline configuration with minimal operational intensity and performance. - **qk/pv Medusa** configurations demonstrate scalability but are constrained by hardware limits at high operational intensities. </details> Figure 19: FLOP/s vs. Operational Intensity of attention matrix multiplication with sequence length 1024. <details> <summary>2401.10774v3/x23.png Details</summary>  ### Visual Description # Technical Document Extraction: Llama 33B, A100 80GB PCIe Performance Chart ## Chart Title **Llama 33B, A100 80GB PCIe** ## Axes - **X-Axis (Horizontal):** - Label: `Operational Intensity (FLOP/Byte)` - Range: `1` to `10k` - Tick Markers: `1`, `10`, `100`, `1k`, `10k` - Gridlines: Logarithmic scale - **Y-Axis (Vertical):** - Label: `Performance (FLOP/s)` - Range: `10G` to `100T` - Tick Markers: `10G`, `100G`, `1T`, `10T`, `100T` - Gridlines: Logarithmic scale ## Key Elements 1. **Dashed Lines:** - **Blue Dashed Line:** - Label: `1,935GB/s` - Purpose: Represents memory bandwidth limit. - **Red Dashed Line:** - Label: `312 TFLOP/s` - Purpose: Represents peak theoretical performance limit. 2. **Vertical Green Line:** - Position: `100 FLOP/Byte` (X-axis) - Purpose: Threshold for operational intensity. 3. **Data Points:** - **Markers:** - Gray: `up/gate/down ar` - Orange: `up/gate/down spec: 16` - Light Orange: `up/gate/down spec: 32` - Red: `up/gate/down spec: 48` - Pink: `up/gate/down spec: 64` - Purple: `up/gate/down spec: 80` - Dark Purple: `up/gate/down spec: 96` - Violet: `up/gate/down spec: 112` - **Trend:** - Performance increases with operational intensity until plateauing near the red dashed line (`312 TFLOP/s`). - Data points cluster around the blue dashed line (`1,935GB/s`) at higher operational intensities. ## Legend - **Color-Coded Specifications:** - Gray: `up/gate/down ar` - Orange: `up/gate/down spec: 16` - Light Orange: `up/gate/down spec: 32` - Red: `up/gate/down spec: 48` - Pink: `up/gate/down spec: 64` - Purple: `up/gate/down spec: 80` - Dark Purple: `up/gate/down spec: 96` - Violet: `up/gate/down spec: 112` ## Observations - **Performance Scaling:** - At low operational intensities (`<100 FLOP/Byte`), performance scales linearly with intensity. - Beyond `100 FLOP/Byte`, performance plateaus due to hardware bottlenecks (memory bandwidth or compute limits). - **Bottlenecks:** - The red dashed line (`312 TFLOP/s`) acts as a hard limit for compute performance. - The blue dashed line (`1,935GB/s`) represents the memory bandwidth ceiling, which becomes the limiting factor at higher operational intensities. ## Notes - The chart uses a logarithmic scale for both axes to visualize performance across multiple orders of magnitude. - Data points are plotted with distinct colors to differentiate between up/down gate specifications. - The green vertical line at `100 FLOP/Byte` highlights a critical operational intensity threshold. </details> Figure 20: FLOP/s vs. Operational Intensity of Linear layers. | 128 256 512 | 0.54 & 0.98 0.75 & 0.99 1.02 & 0.99 | 7.87 & 12.8 11.2 & 13.47 14.69 & 13.84 | 14.73 & 21.33 21.29 & 23.27 27.47 & 24.38 | 19.78 & 27.43 28.69 & 30.72 37.35 & 32.68 | 25.25 & 32.0 36.59 & 36.57 47.09 & 39.38 | 28.63 & 35.56 41.2 & 41.29 52.24 & 44.91 | 32.58 & 38.4 45.99 & 45.18 59.55 & 49.55 | 36.57 & 40.73 52.33 & 48.43 66.35 & 53.49 | | --- | --- | --- | --- | --- | --- | --- | --- | --- | | 1024 | 1.24 & 0.99 | 17.42 & 14.03 | 32.15 & 24.98 | 43.89 & 33.76 | 54.8 & 40.96 | 60.19 & 46.97 | 68.28 & 52.07 | 75.45 & 56.44 | | 2048 | 1.39 & 0.99 | 19.03 & 14.12 | 35.05 & 25.28 | 48.03 & 34.32 | 59.66 & 41.8 | 63.91 & 48.08 | 72.83 & 53.43 | 80.05 & 58.04 | | 4096 | 1.48 & 0.99 | 19.8 & 14.17 | 36.59 & 25.44 | 50.4 & 34.61 | 62.29 & 42.23 | 65.84 & 48.65 | 74.86 & 54.13 | 82.06 & 58.87 | | 8192 | 1.53 & 0.99 | 20.08 & 14.2 | 36.89 & 25.52 | 50.44 & 34.76 | 62.11 & 42.45 | 67.5 & 48.94 | 76.97 & 54.49 | 84.5 & 59.3 | Table 6: TFLOP/s & Operational Intensity of attention matrix multiplication with batch size 16 for Llama 33B on an A100 80GB PCIe. | 1 2 4 | 0.37 & 0.99 0.54 & 0.99 0.75 & 0.99 | 5.22 & 14.03 8.25 & 14.03 11.41 & 14.03 | 10.15 & 24.98 16.0 & 24.98 21.97 & 24.98 | 15.02 & 33.76 21.62 & 33.76 30.02 & 33.76 | 19.79 & 40.96 28.24 & 40.96 38.71 & 40.96 | 21.52 & 46.97 31.84 & 46.97 43.41 & 46.97 | 25.65 & 52.07 37.49 & 52.07 50.06 & 52.07 | 29.4 & 56.44 43.04 & 56.44 56.77 & 56.44 | | --- | --- | --- | --- | --- | --- | --- | --- | --- | | 8 | 1.02 & 0.99 | 14.78 & 14.03 | 27.78 & 24.98 | 38.09 & 33.76 | 47.99 & 40.96 | 53.32 & 46.97 | 61.0 & 52.07 | 68.11 & 56.44 | | 16 | 1.24 & 0.99 | 17.42 & 14.03 | 32.15 & 24.98 | 43.89 & 33.76 | 54.8 & 40.96 | 60.19 & 46.97 | 68.28 & 52.07 | 75.45 & 56.44 | | 32 | 1.39 & 0.99 | 18.89 & 14.03 | 34.67 & 24.98 | 47.57 & 33.76 | 58.89 & 40.96 | 63.61 & 46.97 | 72.17 & 52.07 | 79.21 & 56.44 | | 64 | 1.48 & 0.99 | 19.58 & 14.03 | 35.87 & 24.98 | 49.45 & 33.76 | 61.13 & 40.96 | 64.84 & 46.97 | 73.73 & 52.07 | 81.02 & 56.44 | Table 7: TFLOP/s & Operational Intensity of attention matrix multiplication with sequence length 1024 for Llama 33B on an A100 80GB PCIe. | 1 2 4 | 1.26 & 1.0 2.51 & 2.0 5.03 & 4.0 | 19.95 & 15.95 39.66 & 31.79 76.44 & 63.17 | 39.69 & 31.79 76.53 & 63.17 145.8 & 124.71 | 58.4 & 47.53 112.05 & 94.14 128.85 & 184.69 | 76.57 & 63.17 145.73 & 124.71 167.85 & 243.17 | 94.4 & 78.7 130.67 & 154.89 201.19 & 300.21 | 111.91 & 94.14 129.1 & 184.69 236.93 & 355.85 | 128.64 & 109.47 148.56 & 214.12 195.91 & 410.14 | | --- | --- | --- | --- | --- | --- | --- | --- | --- | | 8 | 10.06 & 7.99 | 145.72 & 124.71 | 168.26 & 243.17 | 236.83 & 355.85 | 221.11 & 463.14 | 207.79 & 565.44 | 236.95 & 663.07 | 227.8 & 756.36 | | 16 | 19.96 & 15.95 | 168.35 & 243.17 | 221.41 & 463.14 | 237.5 & 663.07 | 224.71 & 845.59 | 232.49 & 1012.87 | 241.12 & 1166.74 | 229.25 & 1308.76 | | 32 | 39.69 & 31.79 | 221.74 & 463.14 | 224.88 & 845.59 | 241.33 & 1166.74 | 239.02 & 1440.25 | 245.83 & 1675.97 | 243.55 & 1881.24 | 240.33 & 2061.59 | | 64 | 76.57 & 63.17 | 225.19 & 845.59 | 239.2 & 1440.25 | 243.26 & 1881.24 | 246.16 & 2221.31 | 246.91 & 2491.55 | 244.52 & 2711.46 | 246.14 & 2893.91 | Table 8: TFLOP/s & Operational Intensity of linear layers (up/gate/down) for Llama 33B on an A100 80GB PCIe. G.3 Predicting Medusa Performance We further employ a straightforward analytical model for the acceleration rate. The ablation study results in Sec. 3.3.1 indicate that the acceleration rate can be approximated by a simple logarithmic function. Using the results from Fig. 4(a), we model the curve as $\texttt{acc\_rate}=0.477\log(\texttt{num\_candidate})$ . We simulate the latency of one simplified block of the Llama-7B model (sequentially processing $XW_{Q}$ , $XW_{K}$ , $XW_{V}$ , $QK^{T}$ , $PV$ , $XW_{u}$ , $XW_{g}$ , $XW_{d}$ ) by first fixing the batch size at 1 and the sequence length at 1024. The candidate tokens are processed parallelly by constructing the tree attention described in Section 2.1.2. We omit the latency of the post-processing steps including verification and acceptance for Medusa since they introduce marginal overhead. Fig. 21 illustrates the simulated acceleration rate and speedup for different numbers of candidate tokens under these settings. As the number of candidate tokens increases, both the acceleration rate and speedup initially show improvements. However, beyond 64, the speedup starts to decline, indicating diminishing returns with further increases in candidate length. This aligns with the experimental results in Fig. 4(b) and suggests that there is an optimal range for the numbers of candidate tokens where Medusa provides the most significant performance gains. We plot the simulated speedup under different batch size settings with a fixed sequence length of 1024 in Fig. 22. The results indicate that when the batch size exceeds 32, the speedup decreases and may even have a negative effect. This occurs because the linear layers shift from being memory-bandwidth-bound to computationally bound. We conduct another experiment using a batch size of 4 and different sequence lengths. As shown in Fig. 23, the optimal number of candidate tokens remains relatively consistent across different sequence lengths. However, as the sequence length increases, the overall performance decreases. This performance drop is primarily due to the overhead from attention matrix multiplication, while the linear layer computation remains constant since the computation of linear layers is independent of the sequence length. Our simulations show that the optimal number of candidate tokens is key for model scaling with Medusa, as benefits decrease beyond a certain range. Initially, increasing batch size improves performance through parallelism, but too large a batch size shifts linear layers from memory-bandwidth-bound to compute-bound, reducing speedup. Longer sequences increase attention matrix multiplication overhead, lowering performance, and emphasizing the need to optimize attention mechanisms. Effective model scaling requires balancing the number of candidate tokens, adjusting batch sizes to avoid compute-bound transitions, and enhancing attention mechanisms for longer sequences. These strategies ensure better resource utilization and higher performance, demonstrating the value of simulations in predicting performance and guiding acceleration strategy design. <details> <summary>2401.10774v3/x24.png Details</summary>  ### Visual Description # Technical Analysis of Llama 7B Performance Chart ## Chart Title Llama 7B, Batch Size: 1, Sequence Length: 1024 ## Axes - **X-Axis**: Number of Candidate Tokens Values: 1, 16, 32, 48, 64, 80, 96, 112 - **Y-Axis**: Normalized Latency/Acc. Rate/ Speedup ## Legend 1. **Simulated Acc. Rate** (Blue Star, Dashed Line) 2. **Simulated Speedup** (Green Star, Dashed Line) 3. **qk/pv ar** (Dark Purple, Bar Segment) 4. **qkv linear ar** (Light Purple, Bar Segment) 5. **up/gate/down ar** (Pink, Bar Segment) ## Key Trends 1. **Simulated Acc. Rate** - Starts at 1.0 (1 candidate token) - Increases steadily to 3.3 (112 candidate tokens) - Peaks at 3.3 (112 tokens) 2. **Simulated Speedup** - Starts at 1.0 (1 candidate token) - Rises to 2.9 (112 candidate tokens) - Peaks at 2.9 (112 tokens) 3. **Normalized Latency Components** - **up/gate/down ar** (Pink): Dominates latency across all token counts (40-60% of total) - **qkv linear ar** (Light Purple): Second-largest contributor (30-40% of total) - **qk/pv ar** (Dark Purple): Smallest contributor (10-20% of total) ## Data Points | Candidate Tokens | Simulated Acc. Rate | Simulated Speedup | Total Normalized Latency | |-------------------|---------------------|-------------------|--------------------------| | 1 | 1.0 | 1.0 | 1.0 | | 16 | 2.4 | 2.4 | 1.0 | | 32 | 2.7 | 2.7 | 1.0 | | 48 | 2.9 | 2.8 | 1.0 | | 64 | 3.0 | 2.9 | 1.0 | | 80 | 3.1 | 2.8 | 1.1 | | 96 | 3.2 | 2.8 | 1.1 | | 112 | 3.3 | 2.8 | 1.2 | ## Observations - **Efficiency Scaling**: Both accuracy and speedup improve with increased candidate tokens, but plateau after 64 tokens. - **Latency Breakdown**: The `up/gate/down ar` component consistently accounts for the largest portion of latency, suggesting it is the primary bottleneck. - **Hardware Impact**: The chart assumes a fixed batch size (1) and sequence length (1024), isolating the effect of candidate token count. </details> Figure 21: Simulated acceleration rate, speedup, and normalized latency ablation using different numbers of candidate tokens under the setting of batch size 1 and sequence length 1024 for Llama-7B on an A100 80GB PCIe. <details> <summary>2401.10774v3/x25.png Details</summary>  ### Visual Description # Technical Document Analysis: Llama 7B Speedup Chart ## Chart Title **Llama 7B, Sequence Length: 1024** ## Axes - **X-axis**: "Number of Candidate Tokens" (values: 1, 16, 32, 48, 64, 80, 96, 112) - **Y-axis**: "Speedup (%)" (range: 0.0 to 3.0) ## Legend - **Placement**: Right side of the chart - **Entries**: 1. **Blue (solid line)**: Simulated Speedup @ bs 1 2. **Orange (dashed line)**: Simulated Speedup @ bs 2 3. **Green (dash-dot line)**: Simulated Speedup @ bs 4 4. **Red (dotted line)**: Simulated Speedup @ bs 8 5. **Purple (dash-dot-dot line)**: Simulated Speedup @ bs 16 6. **Brown (dotted line)**: Simulated Speedup @ bs 32 7. **Pink (dotted line)**: Simulated Speedup @ bs 64 ## Key Trends 1. **Blue Line (bs 1)**: - Starts at 1.0 (x=1), rises to ~2.9 (x=64), then slightly declines to ~2.8 (x=112). - **Trend**: Sharp initial increase, followed by a plateau. 2. **Orange Line (bs 2)**: - Starts at 1.0 (x=1), peaks at ~2.6 (x=32), then declines to ~2.4 (x=112). - **Trend**: Early peak, gradual decline. 3. **Green Line (bs 4)**: - Starts at 1.0 (x=1), peaks at ~2.2 (x=16), then drops to ~1.5 (x=112). - **Trend**: Early peak, steep decline. 4. **Red Line (bs 8)**: - Starts at 1.0 (x=1), peaks at ~1.8 (x=16), then declines to ~0.9 (x=112). - **Trend**: Early peak, gradual decline. 5. **Purple Line (bs 16)**: - Starts at 1.0 (x=1), peaks at ~1.4 (x=16), then declines to ~0.5 (x=112). - **Trend**: Early peak, steep decline. 6. **Brown Line (bs 32)**: - Starts at 1.0 (x=1), peaks at ~0.9 (x=16), then declines to ~0.3 (x=112). - **Trend**: Early peak, gradual decline. 7. **Pink Line (bs 64)**: - Starts at 1.0 (x=1), peaks at ~0.4 (x=16), then declines to ~0.1 (x=112). - **Trend**: Early peak, steep decline. ## Data Points (Extracted) | X (Tokens) | bs 1 (Blue) | bs 2 (Orange) | bs 4 (Green) | bs 8 (Red) | bs 16 (Purple) | bs 32 (Brown) | bs 64 (Pink) | |------------|-------------|---------------|--------------|------------|----------------|---------------|--------------| | 1 | 1.0 | 1.0 | 1.0 | 1.0 | 1.0 | 1.0 | 1.0 | | 16 | 2.4 | 2.3 | 2.2 | 1.8 | 1.4 | 0.9 | 0.4 | | 32 | 2.7 | 2.6 | 2.3 | 1.8 | 1.1 | 0.6 | 0.3 | | 48 | 2.9 | 2.6 | 2.2 | 1.4 | 0.9 | 0.5 | 0.2 | | 64 | 3.0 | 2.6 | 2.2 | 1.3 | 0.7 | 0.4 | 0.2 | | 80 | 2.8 | 2.5 | 1.6 | 1.2 | 0.7 | 0.3 | 0.1 | | 96 | 2.9 | 2.5 | 1.6 | 1.0 | 0.6 | 0.3 | 0.1 | | 112 | 2.9 | 2.4 | 1.5 | 0.9 | 0.5 | 0.2 | 0.1 | ## Spatial Grounding - **Legend**: Right-aligned, outside the plot area. - **Lines**: All lines originate at (1, 1.0) and diverge based on batch size. ## Component Isolation - **Header**: Chart title and sequence length. - **Main Chart**: Line plot with seven data series. - **Footer**: No additional text or annotations. ## Validation - All legend colors match line styles and data points. - Trends align with visual slopes (e.g., blue line slopes upward, pink line slopes downward). - No data table present; all information is embedded in the chart. </details> Figure 22: Simulated speedup with sequence length 1024 for Llama-7B. <details> <summary>2401.10774v3/x26.png Details</summary>  ### Visual Description # Technical Analysis of Speedup Performance for Llama 7B Model ## Chart Title **Llama 7B, Batch Size: 4** ## Axes Labels - **X-Axis**: Number of Candidate Tokens (1, 16, 32, 48, 64, 80, 96, 112) - **Y-Axis**: Speedup (%) (1.0, 1.2, 1.4, 1.6, 1.8, 2.0, 2.2, 2.4, 2.6) ## Legend Entries | Color/Symbol | Label | |--------------|-------| | Blue Star | Simulated Speedup @ seq_len 128 | | Orange Star | Simulated Speedup @ seq_len 256 | | Green Star | Simulated Speedup @ seq_len 512 | | Red Star | Simulated Speedup @ seq_len 1024 | | Purple Star | Simulated Speedup @ seq_len 2048 | | Brown Star | Simulated Speedup @ seq_len 4096 | | Pink Star | Simulated Speedup @ seq_len 8192 | ## Key Trends 1. **Speedup Peaks**: - All sequence lengths show maximum speedup at **32 candidate tokens** (~2.4-2.5%). - Exception: `seq_len 8192` peaks at **16 tokens** (~2.3%). 2. **Performance Decline**: - Speedup decreases sharply after 32 tokens for most sequence lengths. - `seq_len 8192` shows gradual decline after 16 tokens. 3. **Relative Performance**: - `seq_len 128` (blue) maintains highest speedup across all token counts. - `seq_len 8192` (pink) consistently exhibits lowest speedup. ## Data Points | Token Count | seq_len 128 | seq_len 256 | seq_len 512 | seq_len 1024 | seq_len 2048 | seq_len 4096 | seq_len 8192 | |-------------|-------------|-------------|-------------|--------------|--------------|--------------|--------------| | 1 | 2.4% | 2.3% | 2.3% | 2.1% | 1.9% | 1.7% | 1.0% | | 16 | 2.5% | 2.4% | 2.3% | 2.2% | 2.0% | 1.7% | 1.3% | | 32 | 2.5% | 2.4% | 2.3% | 2.2% | 2.1% | 1.7% | 1.4% | | 48 | 2.4% | 2.3% | 2.2% | 2.1% | 1.9% | 1.6% | 1.3% | | 64 | 2.3% | 2.2% | 2.1% | 2.0% | 1.8% | 1.5% | 1.3% | | 80 | 1.7% | 1.6% | 1.6% | 1.5% | 1.4% | 1.3% | 1.0% | | 96 | 1.7% | 1.6% | 1.6% | 1.5% | 1.4% | 1.3% | 1.0% | | 112 | 1.6% | 1.5% | 1.5% | 1.4% | 1.3% | 1.2% | 0.9% | ## Technical Observations - **Batch Size Impact**: All simulations use batch size = 4. - **Sequence Length Correlation**: Longer sequence lengths (e.g., 8192) show reduced speedup efficiency. - **Optimal Token Range**: 16-32 tokens provide best performance for most configurations. </details> Figure 23: Simulated speedup with batch size 4 for Llama-7B.