Stream-CQSA: Avoiding Out-of-Memory in Attention Computation via Flexible Workload Scheduling

📄 arXiv: 2604.20819v1 📥 PDF

作者: Yiming Bian, Joshua M. Akey

分类: cs.LG, cs.DC

发布日期: 2026-04-22

💡 一句话要点

Stream-CQSA：通过灵活的工作负载调度避免Attention计算中的内存溢出

🎯 匹配领域: 支柱九：具身大模型 (Embodied Foundation Models)

关键词: 长文本处理 自注意力机制 内存优化 循环仲裁集 流式计算

📋 核心要点

  1. 现有长文本处理模型受限于自注意力机制的二次方内存消耗，易导致内存溢出，限制了模型处理长序列的能力。
  2. 论文提出CQS Divide运算，将注意力分解为独立的子序列计算，保证重组后结果与完整注意力一致。
  3. Stream-CQSA框架通过调度子问题，使注意力计算适应任意内存预算，实现单GPU上数十亿token序列的精确注意力计算。

📝 摘要（中文）

长上下文大型语言模型的可扩展性受到精确自注意力二次方内存成本的根本限制，这通常导致现代硬件上的内存溢出（OOM）故障。现有方法将内存效率提高到接近线性复杂度，但假设完整的查询、键和值张量可以放入设备内存。本文通过引入CQS Divide来消除此假设，CQS Divide是一种源自循环仲裁集（CQS）理论的运算，它将注意力分解为一组独立的子序列计算，这些子序列计算的重组产生与完整序列注意力完全相同的结果。利用这种分解，我们引入了Stream-CQSA，这是一种内存自适应调度框架，可将注意力划分为适合任意内存预算的子问题。这会将注意力从逻辑上的单片操作转变为可调度的任务集合，从而实现跨设备的灵活执行，而无需设备间通信。实验表明，可预测的内存扩展，并表明可以通过流式传输在单个GPU上执行数十亿token序列的精确注意力，而无需更改注意力的底层数学定义或引入近似误差。

🔬 方法详解

问题定义：现有长文本处理模型在计算自注意力时，内存消耗呈二次方增长，当处理超长序列时，容易超出GPU内存限制，导致OOM错误。即使是一些近似方法，也仍然假设query, key, value张量可以完整放入设备内存。因此，如何降低自注意力计算的内存需求，使其能够处理更长的序列，是一个亟待解决的问题。

核心思路：论文的核心思路是将完整的自注意力计算分解为多个独立的子序列计算，每个子序列计算的内存需求较小，可以放入GPU内存中。通过循环仲裁集（CQS）理论，保证这些子序列计算的结果可以无损地重组，得到与完整自注意力计算完全相同的结果。

技术框架：Stream-CQSA框架主要包含两个阶段：CQS Divide和内存自适应调度。CQS Divide阶段将输入序列分解为多个子序列，每个子序列对应一个独立的自注意力计算任务。内存自适应调度阶段根据可用的GPU内存，动态地调度这些任务的执行。由于每个任务都是独立的，因此可以并行执行，也可以在不同的设备上执行，而无需设备间的通信。

关键创新：论文的关键创新在于提出了CQS Divide运算，它基于循环仲裁集理论，可以将自注意力计算分解为一组独立的子序列计算，并且保证重组后的结果与完整自注意力计算完全相同。这种分解方法是无损的，不会引入任何近似误差。此外，Stream-CQSA框架可以根据可用的GPU内存，动态地调度这些子序列计算任务的执行，从而实现内存自适应。

关键设计：CQS Divide运算的关键在于选择合适的循环仲裁集。论文中并没有详细说明如何选择循环仲裁集，这可能需要根据具体的应用场景进行调整。此外，内存自适应调度算法也需要根据可用的GPU内存和任务的计算量进行优化，以实现最佳的性能。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

实验结果表明，Stream-CQSA框架可以有效地降低自注意力计算的内存需求，使其能够在单个GPU上处理数十亿token的序列。与传统的自注意力计算方法相比，Stream-CQSA框架在内存使用方面具有显著的优势，并且不会引入任何近似误差。通过流式传输，可以在不改变底层数学定义的情况下，实现精确的注意力计算。

🎯 应用场景

该研究成果可应用于需要处理超长文本序列的各种场景，例如长篇文档摘要、基因组序列分析、长时间序列预测等。通过降低自注意力计算的内存需求，可以使大型语言模型能够处理更长的上下文信息，从而提高模型的性能和泛化能力。此外，该方法还可以应用于资源受限的设备上，例如移动设备和嵌入式系统。

📄 摘要（原文）

The scalability of long-context large language models is fundamentally limited by the quadratic memory cost of exact self-attention, which often leads to out-of-memory (OOM) failures on modern hardware. Existing methods improve memory efficiency to near-linear complexity, while assuming that the full query, key, and value tensors fit in device memory. In this work, we remove this assumption by introducing CQS Divide, an operation derived from cyclic quorum sets (CQS) theory that decomposes attention into a set of independent subsequence computations whose recomposition yields exactly the same result as full-sequence attention. Exploiting this decomposition, we introduce Stream-CQSA, a memory-adaptive scheduling framework that partitions attention into subproblems that fit within arbitrary memory budgets. This recasts attention from a logically monolithic operation into a collection of schedulable tasks, enabling flexible execution across devices without inter-device communication. Experiments demonstrate predictable memory scaling and show that exact attention over billion-token sequences can be executed on a single GPU via streaming, without changing the underlying mathematical definition of attention or introducing approximation error.