Authors:
(1) Ben Athiwaratkun, AWS AI Labs;
(2) Sujan Kumar Gonugondla, AWS AI Labs;
(3) Sanjay Krishna Gouda, AWS AI Labs;
(4) Haifeng Qian, AWS AI Labs;
(5) Sanjay Krishna Gouda, AWS AI Labs;
(6) Hantian Ding, AWS AI Labs;
(7) Qing Sun, AWS AI Labs;
(8) Jun Wang, AWS AI Labs;
(9) Jiacheng Guo, AWS AI Labs;
(10 Liangfu Chen, AWS AI Labs;
(11) Parminder Bhatia, GE HealthCare (work done at AWS);
(12) Ramesh Nallapati, Amazon AGI (work done at AWS);
(13) Sudipta Sengupta, AWS AI Labs;
(14) Bing Xiang, Goldman Sachs (work done at AWS). Table of Links Abstract and 1 Introduction 2. Related Work 3. Background 3.1. Notation and 3.2. Language Model Inference 3.3. Multi-Query, Multi-Head and the Generalized Multi-Query Attention 4. Context-Aware Bifurcated Attention and 4.1. Motivation 4.2. Formulation and 4.3. Memory IO Complexity 5. Experiments 5.1. Comparing Capabilities of Multi-Head, Multi-Query, and Multi-Group Attention 5.2. Latencies of Capabilities-Equivalent Models 5.3. Applications 6. Conclusion and References A. FAQs B. Related Work C. Setup D. Multi-Group Attention Family E. Context-Aware Bifurcated Attention F. Applications: Additional Results G. Compatibility with Speculative Decoding and Fast Decoding techniques A. FAQs Q: If we already have an MQ model that seems to be quite efficient at large batch sampling, is bifurcated attention necessary?
A: The proposed context-aware bifurcated attention is an exact computation that provides a different way to perform attention, so one can use it "for free" without a performance tradeoff. Due to the reduced memory I/O, it enables more extreme cases of batch sampling, such as a larger batch, even for long contexts.

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

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

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

Q: How does model quantization (or lower precision arithmetic) affect the findings?
A: There are two regimes for quantization: model weight quantization and attention quantization. To date, most quantization only focuses on the weight since the attention computation is precision-sensitive and quantization has not proved to be viable. Model quantization can make incremental decoding faster due to lower memory I/O of the model itself, since the effective model size in memory is smaller. This shifts the latency curve downward for all context lengths or batch sizes. The overall conclusion for the bifurcated and multi-query attention remains the same, however, since the improvement proposed in the paper is on the attention component, which is orthogonal to the model weight. If attention quantization is viable in the future, the lower memory on the attention tensor will effectively reduce the memory I/O for KV cache by a factor of 2 in the case of int8 quantization (compared to fp16 or bf16) or a factor of 4 in the case of int4. Overall, this will flatten the latency growth with respect to batch size or context length. The overall comparative complexity (a) with or without bifurcated attention or (b) multi-head vs. multi-query remains the same.

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

Q: How does bifurcated attention differ from using attention mask for sampling as in done in SpecInfer (Miao et al., 2023)? This paper is available on arxiv under CC BY 4.0 DEED license. Authors: (1) Ben Athiwaratkun, AWS AI Labs; (2) Sujan Kumar Gonugondla, AWS AI Labs; (3) Sanjay Krishna Gouda, AWS AI Labs; (4) Haifeng Qian, AWS AI Labs; (5) Sanjay Krishna Gouda, AWS AI Labs; (6) Hantian Ding, AWS AI Labs; (7) Qing Sun, AWS AI Labs; (8) Jun Wang, AWS AI Labs; (9) Jiacheng Guo, AWS AI Labs; (10 Liangfu Chen, AWS AI Labs; (11) Parminder Bhatia, GE HealthCare (work done at AWS); (12) Ramesh Nallapati, Amazon AGI (work done at AWS); (13) Sudipta Sengupta, AWS AI Labs; (14) Bing Xiang, Goldman Sachs (work done at AWS). Authors: Authors: (1) Ben Athiwaratkun, AWS AI Labs; (2) Sujan Kumar Gonugondla, AWS AI Labs; (3) Sanjay Krishna Gouda, AWS AI Labs; (4) Haifeng Qian, AWS AI Labs; (5) Sanjay Krishna Gouda, AWS AI Labs; (6) Hantian Ding, AWS AI Labs; (7) Qing Sun, AWS AI Labs; (8) Jun Wang, AWS AI Labs; (9) Jiacheng Guo, AWS AI Labs; (10 Liangfu Chen, AWS AI Labs; (11) Parminder Bhatia, GE HealthCare (work done at AWS); (12) Ramesh Nallapati, Amazon AGI (work done at AWS); (13) Sudipta Sengupta, AWS AI Labs; (14) Bing Xiang, Goldman Sachs (work done at AWS). Table of Links Abstract and 1 Introduction Abstract and 1 Introduction 2. Related Work 2. Related Work 3. Background 3. Background 3.1. Notation and 3.2. Language Model Inference 3.1. Notation and 3.2. Language Model Inference 3.3. Multi-Query, Multi-Head and the Generalized Multi-Query Attention 3.3. Multi-Query, Multi-Head and the Generalized Multi-Query Attention 4. Context-Aware Bifurcated Attention and 4.1. Motivation 4. Context-Aware Bifurcated Attention and 4.1. Motivation 4.2. Formulation and 4.3. Memory IO Complexity 4.2. Formulation and 4.3. Memory IO Complexity 5. Experiments 5. Experiments 5.1. Comparing Capabilities of Multi-Head, Multi-Query, and Multi-Group Attention 5.1. Comparing Capabilities of Multi-Head, Multi-Query, and Multi-Group Attention 5.2. Latencies of Capabilities-Equivalent Models 5.2. Latencies of Capabilities-Equivalent Models 5.3. Applications 5.3. Applications 6. Conclusion and References 6. Conclusion and References A. FAQs A. FAQs B. Related Work B. Related Work C. Setup C. Setup D. Multi-Group Attention Family D. Multi-Group Attention Family E. Context-Aware Bifurcated Attention E. Context-Aware Bifurcated Attention F. Applications: Additional Results F. Applications: Additional Results G. Compatibility with Speculative Decoding and Fast Decoding techniques G. Compatibility with Speculative Decoding and Fast Decoding techniques A. FAQs Q: If we already have an MQ model that seems to be quite efficient at large batch sampling, is bifurcated attention necessary?
A: The proposed context-aware bifurcated attention is an exact computation that provides a different way to perform attention, so one can use it "for free" without a performance tradeoff. Due to the reduced memory I/O, it enables more extreme cases of batch sampling, such as a larger batch, even for long contexts. Q: How applicable is multi-query for single-batch inference without high batch sampling?
A: If the context is long and the number of generated tokens is high, then the benefits of multi-query are clear. Please see Section 5.2.1. Q: Is bifurcated attention applicable for the case where we process different inputs in a batch?
A: No. In that case, if we need a solution to reduce memory I/O during incremental decoding, then multi-query attention can be appealing, especially in scenarios with a high number of generated tokens where the incremental decoding phase dominates the overall latency. This is because there is an overhead to multi-query due to the context encoding phase, as outlined in the main paper. Q: Any caveats to using bifurcated attention?
A: For small workloads (low context length and batch size), due to the fact that we split the attention into two parts, there can be less parallelization of the GEMM kernels, which could lead to higher latency, especially for MQ models. However, one can get the best of both worlds given any model by triggering bifurcated attention under high workload scenarios and using normal attention otherwise. With such a workload-based switch, bifurcated attention is guaranteed to provide better latency and efficiency. Q: How does model quantization (or lower precision arithmetic) affect the findings?
A: There are two regimes for quantization: model weight quantization and attention quantization. To date, most quantization only focuses on the weight since the attention computation is precision-sensitive and quantization has not proved to be viable. Model quantization can make incremental decoding faster due to lower memory I/O of the model itself, since the effective model size in memory is smaller. This shifts the latency curve downward for all context lengths or batch sizes. The overall conclusion for the bifurcated and multi-query attention remains the same, however, since the improvement proposed in the paper is on the attention component, which is orthogonal to the model weight. If attention quantization is viable in the future, the lower memory on the attention tensor will effectively reduce the memory I/O for KV cache by a factor of 2 in the case of int8 quantization (compared to fp16 or bf16) or a factor of 4 in the case of int4. Overall, this will flatten the latency growth with respect to batch size or context length. The overall comparative complexity (a) with or without bifurcated attention or (b) multi-head vs. multi-query remains the same. Q: Does the conclusion depend on the inference implementation or different hardware?
A: Different inference platforms, such as FasterTransformers (GPUs) or PaLM inference (TPUs), can yield different latency numbers. However, the relative I/O complexity among different attention mechanisms does not change, resulting in similar relative trends among different attention mechanisms. That being said, it is possible that more efficient implementations or more performant chip/system configurations, including different tensor parallelism degrees, can result in different slopes for the latency growth with respect to context length and batch size. In that case, the trade-off points in terms of context length or batch size can be different. The comparative complexity remains the same based on the analysis. Q: How does bifurcated attention differ from using attention mask for sampling as in done in SpecInfer (Miao et al., 2023)? Q: If we already have an MQ model that seems to be quite efficient at large batch sampling, is bifurcated attention necessary? A: The proposed context-aware bifurcated attention is an exact computation that provides a different way to perform attention, so one can use it "for free" without a performance tradeoff. Due to the reduced memory I/O, it enables more extreme cases of batch sampling, such as a larger batch, even for long contexts. Q: If we already have an MQ model that seems to be quite efficient at large batch sampling, is bifurcated attention necessary? Q: A: The proposed context-aware bifurcated attention is an exact computation that provides a different way to perform attention, so one can use it "for free" without a performance tradeoff. Due to the reduced memory I/O, it enables more extreme cases of batch sampling, such as a larger batch, even for long contexts. A: Q: How applicable is multi-query for single-batch inference without high batch sampling? A: If the context is long and the number of generated tokens is high, then the benefits of multi-query are clear. Please see Section 5.2.1. Q: How applicable is multi-query for single-batch inference without high batch sampling? Q: A: If the context is long and the number of generated tokens is high, then the benefits of multi-query are clear. Please see Section 5.2.1. A: Q: Is bifurcated attention applicable for the case where we process different inputs in a batch? A: No. In that case, if we need a solution to reduce memory I/O during incremental decoding, then multi-query attention can be appealing, especially in scenarios with a high number of generated tokens where the incremental decoding phase dominates the overall latency. This is because there is an overhead to multi-query due to the context encoding phase, as outlined in the main paper. Q: Is bifurcated attention applicable for the case where we process different inputs in a batch? Q: A: No. In that case, if we need a solution to reduce memory I/O during incremental decoding, then multi-query attention can be appealing, especially in scenarios with a high number of generated tokens where the incremental decoding phase dominates the overall latency. This is because there is an overhead to multi-query due to the context encoding phase, as outlined in the main paper. A: Q: Any caveats to using bifurcated attention? A: For small workloads (low context length and batch size), due to the fact that we split the attention into two parts, there can be less parallelization of the GEMM kernels, which could lead to higher latency, especially for MQ models. However, one can get the best of both worlds given any model by triggering bifurcated attention under high workload scenarios and using normal attention otherwise. With such a workload-based switch, bifurcated attention is guaranteed to provide better latency and efficiency. Q : Any caveats to using bifurcated attention? Q A: For small workloads (low context length and batch size), due to the fact that we split the attention into two parts, there can be less parallelization of the GEMM kernels, which could lead to higher latency, especially for MQ models. However, one can get the best of both worlds given any model by triggering bifurcated attention under high workload scenarios and using normal attention otherwise. With such a workload-based switch, bifurcated attention is guaranteed to provide better latency and efficiency. A: Q: How does model quantization (or lower precision arithmetic) affect the findings? A: There are two regimes for quantization: model weight quantization and attention quantization. To date, most quantization only focuses on the weight since the attention computation is precision-sensitive and quantization has not proved to be viable. Model quantization can make incremental decoding faster due to lower memory I/O of the model itself, since the effective model size in memory is smaller. This shifts the latency curve downward for all context lengths or batch sizes. The overall conclusion for the bifurcated and multi-query attention remains the same, however, since the improvement proposed in the paper is on the attention component, which is orthogonal to the model weight. If attention quantization is viable in the future, the lower memory on the attention tensor will effectively reduce the memory I/O for KV cache by a factor of 2 in the case of int8 quantization (compared to fp16 or bf16) or a factor of 4 in the case of int4. Overall, this will flatten the latency growth with respect to batch size or context length. The overall comparative complexity (a) with or without bifurcated attention or (b) multi-head vs. multi-query remains the same. Q: How does model quantization (or lower precision arithmetic) affect the findings? Q: A: There are two regimes for quantization: model weight quantization and attention quantization. To date, most quantization only focuses on the weight since the attention computation is precision-sensitive and quantization has not proved to be viable. Model quantization can make incremental decoding faster due to lower memory I/O of the model itself, since the effective model size in memory is smaller. This shifts the latency curve downward for all context lengths or batch sizes. The overall conclusion for the bifurcated and multi-query attention remains the same, however, since the improvement proposed in the paper is on the attention component, which is orthogonal to the model weight. If attention quantization is viable in the future, the lower memory on the attention tensor will effectively reduce the memory I/O for KV cache by a factor of 2 in the case of int8 quantization (compared to fp16 or bf16) or a factor of 4 in the case of int4. Overall, this will flatten the latency growth with respect to batch size or context length. The overall comparative complexity (a) with or without bifurcated attention or (b) multi-head vs. multi-query remains the same. A: Q: Does the conclusion depend on the inference implementation or different hardware? A: Different inference platforms, such as FasterTransformers (GPUs) or PaLM inference (TPUs), can yield different latency numbers. However, the relative I/O complexity among different attention mechanisms does not change, resulting in similar relative trends among different attention mechanisms. That being said, it is possible that more efficient implementations or more performant chip/system configurations, including different tensor parallelism degrees, can result in different slopes for the latency growth with respect to context length and batch size. In that case, the trade-off points in terms of context length or batch size can be different. The comparative complexity remains the same based on the analysis. Q: Does the conclusion depend on the inference implementation or different hardware? A: Different inference platforms, such as FasterTransformers (GPUs) or PaLM inference (TPUs), can yield different latency numbers. However, the relative I/O complexity among different attention mechanisms does not change, resulting in similar relative trends among different attention mechanisms. That being said, it is possible that more efficient implementations or more performant chip/system configurations, including different tensor parallelism degrees, can result in different slopes for the latency growth with respect to context length and batch size. In that case, the trade-off points in terms of context length or batch size can be different. The comparative complexity remains the same based on the analysis. Q: How does bifurcated attention differ from using attention mask for sampling as in done in SpecInfer (Miao et al., 2023)? Q: How does bifurcated attention differ from using attention mask for sampling as in done in SpecInfer (Miao et al., 2023)? Q: This paper is available on arxiv under CC BY 4.0 DEED license. This paper is available on arxiv under CC BY 4.0 DEED license. available on arxiv

Part of HackerNoon's growing list of open-source research papers, promoting free access to academic material.

Abstract

Amazon

Why Memory I/O Efficiency Matters for AI Model Performance

Riemannian Normalization on Lie Groups

Read My Stories

Too Long; Didn't Read

Make resilience your competitive advantage

FAQs: How Bifurcated Attention Improves AI Model Efficiency

About Author

Comments

TOPICS

THIS ARTICLE WAS FEATURED IN

Related Stories

Untitled Story

A Brief Review of the Lie Group and the Geometries of SPD Manifolds

Complexity NIMBYs and the Illusion of Transcending Software Complexity Bounds

5 Practical Ways AI Can Boost Productivity for Web Developers

On the Concerns of Developers When Using GitHub Copilot

GitHub Copilot in Practice: Empirical Insights into User Experiences and Practical Challenges

A Brief Review of the Lie Group and the Geometries of SPD Manifolds

Complexity NIMBYs and the Illusion of Transcending Software Complexity Bounds

5 Practical Ways AI Can Boost Productivity for Web Developers

On the Concerns of Developers When Using GitHub Copilot

GitHub Copilot in Practice: Empirical Insights into User Experiences and Practical Challenges

Light-Mode

Classic

Newspaper

Dark-Mode

Neon Noir

Minty

HN StartUps