longer context length

the best ways for extending a pretrained model to longer context length:

• ring attention
• YaRN
• RoPE scaling

• Infini-attention: https://weedge.github.io/post/paper/transformer/infini_attention/
• https://huggingface.co/blog/infini-attention
• https://github.com/huggingface/nanotron/tree/xrsrke/infini_attention_this_actually_works/examples/infinite-context-length

Vision

Cross-Attention

• 2021. CrossViT: Cross-Attention Multi-Scale Vision Transformer for Image Classification
• https://github.com/IBM/CrossViT

大语言模型内存要求 LLM memory requirement

实际上，影响 GPU 大型语言模型（LLM）内存需求的两个主要因素是模型权重和 KV 缓存。

更多见：https://weedge.github.io/post/llm/mastering-llm-techniques-inference-optimization/#kv缓存-key-value-caching

模型权重 Model weights

内存被模型参数所占用。举例来说，一个包含 70 亿个参数的模型（比如 Llama 2 7B），以 16 位精度（FP16 或 BF16）加载，大约需要 70 亿 * sizeof(FP16) ~= 14 GB 的内存。

Key-Value (KV) caching

KV（键值）缓存是decoder阶段常见的一种优化技术。在decoder阶段，每个时间步生成一个token，但每个token都依赖于之前所有token的KV张量（包括预填充时计算的输入token的KV张量，以及直到当前时间步计算的任何新KV张量）。

为了避免在每个时间步重新计算所有这些张量的值，可以将它们缓存到GPU内存中。在每次迭代中，当新的元素被计算出来时，它们被简单地添加到正在运行的缓存中，以便在下一次迭代中使用。在某些实现中，模型的每一层都有一个KV缓存。

如果我们可以只输入最近的token，然后从那里预测下一个token，那就更好了。然而，我们需要整个上下文的 K 和 V 值来完成注意力计算。幸运的是，由于句子中的每个单词只能关注句子中的前一个单词，因此 N-1 个token的 K 和 V 值在每次后续计算中都完全相同。这意味着您可以存储 K 和 V 的值，仅计算新添加的标记的新 K 和 V，然后将它们与之前存储的 K 和 V 连接起来以获得相同的完整 K 和 V 向量。

通过使用KV缓存，可以避免重复计算先前时间步骤中生成的中间状态，提高了decoder阶段的效率。这样做可以减少计算成本，加快推理速度，并有效利用GPU的内存和计算资源。这极大地加快了自回归生成的速度，尤其是在较大的模型尺寸下，但有一个问题——现在你的瓶颈是内存(memory bound)而不是计算(computer bound)。

内存计算：

在批处理中，批处理中每个请求的 KV 缓存仍然需要单独分配，并且可能具有较大的内存占用。下面的公式描述了适用于今天大多数常见LLM架构的KV缓存大小。

每个标记的KV缓存大小（以字节为单位）= 2 * (层数) * (头数 * 头尺寸) * 字节精度

Size of KV cache per token in bytes = 2 * (num_layers) * (num_heads * dim_head) * precision_in_bytes

第一个因子2代表K和V矩阵。通常，（头数 * 头尺寸）的值与transformer的隐藏大小（或模型的维度，d_model）相同。这些模型属性通常可以在模型卡片或关联的配置文件中找到。

这个内存大小适用于批量输入序列中的每个标记。假设半精度，KV缓存的总大小由下面的公式给出。

KV缓存的总大小（以字节为单位）= (批量大小) * (序列长度) * 2 * (层数) * (隐藏大小) * sizeof(FP16)

Total size of KV cache in bytes = (batch_size) * (sequence_length) * 2 * (num_layers) * (hidden_size) * sizeof(FP16)

例如，使用 16 位精度的 Llama 2 7B 模型和批量大小为 1，KV 缓存的大小将是 1 * 4096 * 2 * 32 * 4096 * 2 字节，约为 2 GB。

更好模型性能 VS KV 缓存的优化(更少的kv缓存)

回顾下原始论文中的MHA方法的作用:
这种方法是对自注意力机制（SDPA scaled dot-product attention）的一种增强，它使得注意力层能够并行执行多次，每次使用不同的、经过学习的投影矩阵作用于Q、K和V矩阵。这样的设计让模型能够同时关注来自不同表示子空间的信息，并应用于不同的位置。这些表示子空间是独立学习的，为模型提供了对输入中不同位置更丰富的理解。

有效管理这个KV缓存是一个具有挑战性的任务。随着批处理大小和序列长度线性增长，内存需求会迅速扩大，对于参数量比较大的模型，kv cache成为瓶颈， 而且训练推理成本高。因此，它限制了可以提供的吞吐量，并对长上下文输入提出了挑战。在这种背景下，为了改进MHA K/V cache， 针对原始Transformer使用 Multi-Head Attention （MHA） 多头注意力 提出了不同的优化方法， 按自己的理解归为两类：

• 针对MHA的改造： MQA, GQA, MLA （模型结构优化）
• 结合GPU硬件的片内内存(SRAM), 片外内存(HBM), 以及cpu, 内存(DRAM), 进行数据计算的优化（通过融合和重新排序注意力机制的不同组件而获得的相对加速度。）本质上是利用硬件的局部性原理，加速计算。（结合模型结构的工程优化），比如：
  • FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness FlashAttention是一种针对注意力机制的优化技术，旨在提高模型的计算效率和内存利用率。它专注于减少多头自注意力机制中的内存移动成本，并优化计算操作，特别是针对大规模语言模型的推理过程。(PyTorch >= 2.0 引入了flash attention, torch.nn.functional.scaled_dot_product_attention, impl from: https://github.com/Dao-AILab/flash-attention)
  • Efficient Memory Management for Large Language Model Serving with PagedAttention PagedAttention算法受操作系统中分页技术的启发，它使连续的键和值能够以非连续的方式存储在内存中。该算法将每个请求的KV缓存划分为表示固定数量标记的块，并可以非连续地存储这些块。应用于模型serving inference推理阶段， 常用的推理服务 vllm内部使用了 PagedAttention算法来管理在GPU中，attention的k/v内存。

这里主要介绍针对MHA模型结构优化。

Multi-Query Attention 多查询注意力 (MQA)

• 2019. Fast Transformer Decoding: One Write-Head is All You Need
  为了减少 MHA 中的 KV 缓存瓶颈， 原论文Transformer的二作 Shazeer 引入了多查询注意力（MQA），其中键和值在所有不同的注意力头之间共享，即它与 MHA 相同，只是不同的头共享一组键和值。这需要非常轻的 KV 缓存，从而大大加快解码器推理速度。然而，MQA 会导致质量下降和训练不稳定。

注： 过渡方案

Grouped-Query Attention 分组查询注意力 (GQA)

• 2023. Gqa: Training generalized multi-query transformer models from multi-head checkpoints
  分组查询注意力（GQA）是 MHA 和 MQA 之间的插值，通过引入多个查询头子组（少于注意力头总数），每个子组都有一个 Key 和 Value 头。与 MQA 相比，随着模型大小的增加，GQA 的内存带宽和容量保持相同比例的减少。中间数量的子组会产生比 MQA 质量更高但比 MHA 更快的插值模型。很明显，单个组的 GQA 相当于 MQA。

这种新的训练方法可以在使用少量原始训练计算资源的情况下，将最初采用多头注意力（MHA）训练的模型进行“上训练”（uptraining），转换为采用组查询注意力（GQA）。这种方法可以使模型在保持接近多头注意力（MHA）的质量的同时，保持更接近多查询注意力（MQA）的计算效率。Llama 2 70B就是一个采用组查询注意力（GQA）的模型的例子。

像多查询注意力（MQA）和组查询注意力（GQA）这样的优化方法有助于通过减少存储的键和值头的数量来减少键值（KV）缓存所需的内存。但在如何管理这种KV缓存方面可能仍存在一些低效性。
这种方法的实施可能会在计算效率和模型表现之间寻找一种平衡，因为它涉及到对查询向量进行组织和处理，有时可能会牺牲一些精确性以换取更高的计算效率。

注： 折中方案，常使用在参数量大的模型上。

Multi-Head Latent Attention 多头潜在注意力 (MLA)

• 2024. DeepSeek-V2: A Strong, Economical, and Efficient Mixture-of-Experts Language Model
  多头潜在注意力（MLA）实现了比 MHA 更优越的性能，并且显着降低了 KV 缓存提升推理效率。 MLA 不是像 MQA 和 GQA 那样减少 KV-heads，而是将 Key 和 Value 联合压缩到一个潜在向量中。

表中展示了不同注意机制之间每个令牌的 KV 缓存的比较。 MLA只需要少量的KV缓存，相当于GQA只有2.25组，但可以实现比MHA更强的性能。
PS: 在deepseek-V3版本中，scaling模型参数变大， 使用少的GPU内存，训练达到好的模型性能， 低调出圈了~ 这个是其中一个因数， 还有一个主要原因是使用了2024. DeepSeekMoE: Towards Ultimate Expert Specialization in Mixture-of-Experts Language Models

• https://planetbanatt.net/articles/mla.html
• https://github.com/madsys-dev/deepseekv2-profile/blob/main/workspace/blog/optimizing-mla.md