2025-11-24T16:10:17.960735

Flash Inference: Near Linear Time Inference for Long Convolution Sequence Models and Beyond

Oncescu, Purandare, Idreos et al.

While transformers have been at the core of most recent advancements in sequence generative models, their computational cost remains quadratic in sequence length. Several subquadratic architectures have been proposed to address this computational issue. Some of them, including long convolution sequence models (LCSMs), such as Hyena, address this issue at training time but remain quadratic during inference. We propose a method for speeding up LCSMs' exact inference to quasilinear $O(L\log^2L)$ time, identify the key properties that make this possible, and propose a general framework that exploits these. Our approach, inspired by previous work on relaxed polynomial interpolation, is based on a tiling which helps decrease memory movement and share computation. It has the added benefit of allowing for almost complete parallelization across layers of the position-mixing part of the architecture. Empirically, we provide a proof of concept implementation for Hyena, which gets up to $7.8\times$ end-to-end improvement over standard inference by improving $110\times$ within the position-mixing part.

academic

Flash Inference: Near Linear Time Inference for Long Convolution Sequence Models and Beyond

基本信息

论文ID: 2410.12982
标题: Flash Inference: Near Linear Time Inference for Long Convolution Sequence Models and Beyond
作者: Costin-Andrei Oncescu, Sanket Purandare, Stratos Idreos, Sham Kakade (Harvard University)
分类: cs.LG, cs.AI
发表时间: arXiv预印本，2024年10月提交，2025年11月更新（v2）
论文链接: https://arxiv.org/abs/2410.12982

摘要

本文针对长卷积序列模型（LCSMs）在推理阶段的二次时间复杂度问题，提出了Flash Inference框架，将精确推理的时间复杂度降低到准线性 $O(L\log^2L)$ 。该方法受松弛多项式插值（relaxed polynomial interpolation）启发，基于分块（tiling）策略减少内存移动并共享计算。在Hyena架构上的实验表明，端到端推理获得了7.8倍加速，位置混合部分获得了110倍加速。

研究背景与动机

1. 核心问题

虽然Transformer在序列生成模型中取得了巨大成功，但其计算成本在序列长度上呈二次增长（ $O(L^2)$ ），这在训练和推理阶段都成为瓶颈。为解决这一问题，研究者提出了多种亚二次架构，包括状态空间模型（SSMs）和长卷积序列模型（LCSMs，如Hyena）。

2. 问题的重要性

训练效率已解决：LCSMs通过FFT可以在训练时达到 $O(L\log L)$ 的复杂度
推理效率未解决：在自回归推理时，由于输入序列是逐步生成的，无法直接使用FFT，导致复杂度退化为 $O(L^2)$
长上下文需求：随着大语言模型处理越来越长的上下文，推理效率问题变得更加突出

3. 现有方法的局限性

近似方法（Massaroli et al. 2024）：将卷积滤波器投影到低维LTI SSM，但这只是近似，且需要昂贵的蒸馏预计算，不支持数据依赖滤波器
递归视角：对于低维SSM可能高效，但对高维SSM（维度接近序列长度）仍然效率低下
结构利用方法：需要滤波器具有特定结构（如低秩LTI SSM），限制了模型表达能力

4. 研究动机

本文旨在提供一个精确且通用的推理加速框架，不依赖于滤波器的特定结构，同时支持数据依赖滤波器。

核心贡献

首个准线性精确推理算法：提出LCSMs的 $O(L\log^2 L)$ 时间复杂度精确推理算法，相比之前的近似方法实现了精确模拟
通用框架识别：识别了使快速推理成为可能的关键架构属性（贡献基础、查询无关），提出了适用于更广泛架构类的Flash Inference框架
跨层并行化：利用分块策略实现位置混合部分几乎完全的跨层并行计算
内存优化：通过分块方法显著减少数据移动，从 $\Omega(L^2)$ 降低到 $O(L\log L)$ ，对数据独立滤波器可节省2倍激活存储
实证验证：在Hyena架构上实现端到端7.8倍加速，卷积部分110倍加速

方法详解

任务定义

自回归序列生成：给定提示序列 $x_1, \ldots, x_p$ ，模型需要逐个生成后续token。在每个位置 $i$ ，模型通过所有层计算激活 $a^{[1,M]}_i$ ，最后从 $a^M_i$ 采样生成 $x_{i+1}$ 。

核心计算瓶颈：对于每层 $\ell$ 和每个维度，需要计算： $z_t = \sum_{i=1}^{t} y_i \cdot \rho_{t-i}$

其中 $y$ 是输入序列， $\rho$ 是长度为 $L$ 的卷积滤波器。朴素实现需要 $\Omega(L^2)$ 时间。

模型架构

1. 通用架构定义

模型由 $M$ 层组成，每层包含：

位置混合模块（mixer）： $\text{mixer}^\ell: \mathbb{R}^{L\times D} \to \mathbb{R}^{L\times D}$ ，使不同位置的嵌入交互
特征混合模块（block）： $\text{block}^\ell: \mathbb{R}^D \to \mathbb{R}^D$ ，包括MLP、层归一化等

激活计算： $a^\ell(x)_i = \text{block}^\ell(\text{mixer}^\ell(a^{\ell-1}(x))_i)$

2. LCSM特定定义

对于LCSMs，mixer通过卷积实现： $\text{mixer}^\ell(y)_t = \sum_{i=1}^{t} y_i \odot \rho^\ell_{t-i}$

其中 $\odot$ 是Hadamard积， $\rho^\ell \in \mathbb{R}^{L\times D}$ 是滤波器（通常由低维参数 $\theta$ 生成： $\rho = f(\theta)$ ）。

核心算法：松弛多项式插值

1. 三种计算策略

Lazy（懒惰）方法：

仅在需要时计算 $z_t = \sum_{i=1}^{t} y_i \cdot \rho_{t-i}$
每个位置需要 $O(t)$ 操作，总复杂度 $O(L^2)$

Eager（急切）方法：

当 $y_t$ 可用时，立即计算其对所有未来位置的贡献
第 $t$ 次迭代需要 $O(L-t)$ 操作，总复杂度仍为 $O(L^2)$

Relaxed（松弛）方法（本文提出）：

将贡献空间分块，使用FFT高效计算块内贡献
关键创新：平衡的矩形分块而非细长条状

2. 贡献聚合定义

定义 $\tau(y, [l,r], \rho, [l',r'])$ 为 $y_{[l,r]}$ 对 $z_{[l',r']}$ 的聚合贡献： $\tau(y, [l,r], \rho, [l',r'])_t = \sum_{i=l}^{r} y_i \cdot \rho_{t-i}, \quad \forall l' \leq t \leq r'$

Lemma 1：存在基于FFT的算法，在 $O((L_1+L_2)\log(L_1+L_2))$ 时间内计算 $\tau$ ，其中 $L_1 = r-l+1$ ， $L_2 = r'-l'+1$ 。

3. 分块策略（Algorithm 1）

for i = 1 to L-1:
    U ← 最大的能整除i的2的幂
    z_i += y_i * ρ_0  # 红色单元：直接依赖
    z[i+1:i+U] += τ(y, [i-U+1, i], ρ, [i+1, i+U])  # 灰色块：急切计算
    return z_i
    unlock y_{i+1}

关键特性：

在第 $i$ 次迭代，计算边长为 $U$ 的灰色块（ $U$ 是整除 $i$ 的最大2的幂）
红色单元处理当前位置的直接依赖
灰色块提前计算部分未来贡献

复杂度分析（Proposition 1）：

对长度为 $2^q$ 的块，有 $2^{P-1-q}$ 次调用（ $L=2^P$ ）
总时间： $\sum_{q=0}^{P-1} 2^{P-1-q} \cdot O(2^q \log 2^q) = O(L\log^2 L)$
内存： $O(L)$ （峰值由最大块决定）

LCSM推理算法（Algorithm 2）

将Algorithm 1扩展到多层多维：

for i = 1 to L-1:
    U ← 最大的能整除i的2的幂
    for ℓ = 1 to M:  # 遍历层
        b^ℓ_i += a^{ℓ-1}_i ⊙ ρ^ℓ_0  # 红色单元
        a^ℓ_i = block^ℓ(b^ℓ_i)
        b^ℓ[i+1:i+U] += τ(a^{ℓ-1}, [i-U+1, i], ρ^ℓ, [i+1, i+U])  # 灰色块
    a^0_{i+1} = sampler(a^M_i)

复杂度（Proposition 2）：

Mixer部分： $O(MDL\log^2 L)$
Block部分： $LM$ 次调用（通常 $O(MLD^2)$ ）
激活存储： $O(MLD)$

技术创新点

1. 跨层并行化（Algorithm 3）

灰色块计算可以跨所有层并行执行：

for i = 1 to L-1:
    for ℓ = 1 to M:
        处理红色单元（必须顺序）
    parallel for ℓ = 1 to M:
        处理灰色块（可并行）

优势：

小块（87.5%的块大小≤4）通常受内存延迟限制，并行化可饱和内存带宽
大块使用FFT实现，计算密集，并行化提升吞吐量

2. 内存优化

数据移动：从 $\Omega(L^2)$ 降至 $O(L\log L)$ （平均每次迭代访问 $\log L$ 个位置）
激活复用：在位置 $i$ 用 $a^\ell_i$ 的空间存储 $b^\ell_i$ （之后不再需要 $b^\ell_i$ ）
FFT预计算：对 $\log L$ 个不同块大小预计算滤波器的DFT，节省1.5倍计算

3. 循环卷积技巧

标准FFT卷积需要4U长度的FFT（输出长度3U-1）
本文只需2U长度的循环卷积（感兴趣的输出范围 $[U, 2U-1]$ 不受循环影响）

4. 数据依赖滤波器扩展（Appendix B）

通过修改分块策略（Algorithm 5），支持 $\rho$ 依赖于数据的情况，代价是2倍计算量。

通用框架：Flash Inference

架构属性

P.1 贡献基础（Contribution-based）： Mixer通过聚合贡献工作： $\text{mixer}(y)_i = \text{read}(\text{agg}(\text{cont}(y,1,i), \ldots, \text{cont}(y,i,i)))$

其中：

$\text{cont}: \mathbb{R}^D \times \mathbb{N} \times \mathbb{N} \to \mathcal{X}$ ：贡献函数
$\text{agg}: \mathcal{X}^* \to \mathcal{X}$ ：结合律聚合函数
$\text{read}: \mathcal{X} \to \mathbb{R}^D$ ：读取函数

示例：

LCSMs： $\mathcal{X}=\mathbb{R}^D$ ， $\text{agg}=\sum$ ， $\text{cont}(y,i,j)=y_i\odot\rho_{j-i}$
Self-attention： $\mathcal{X}=\mathbb{R}^D\times\mathbb{R}$ ， $\text{cont}(y,i,j)=(v_i\cdot e^{\langle k_i,q_j\rangle}, e^{\langle k_i,q_j\rangle})$ ， $\text{read}(v,w)=v/w$

P.2 查询无关（Query-independent）： $\text{cont}(y,i,j)$ 不依赖于 $y_{[i+1,L]}$ （LCSMs满足，Transformer不满足）

通用算法（Algorithm 4）

假设存在算法 $\mathcal{A}$ 能在 $T(L_1, L_2)$ 时间内计算块贡献： $\mathcal{A}(y, [l,r], [l',r']) = \text{agg}(\text{cont}(y,l,p), \ldots, \text{cont}(y,r,p))$

Theorem 2：在P.1和P.2下，每层执行：

$L-1$ 次 $\mathcal{A}$ 调用（ $2^{P-1-q}$ 次长度 $2^q$ 的调用）
总时间： $\sum_{q=0}^{P-1} 2^{P-1-q} T(2^q, 2^q)$
跨层并行：灰色块无数据依赖，可并行

实验设置

数据集与配置

两种实验设置：

Hyena架构：真实的LCSM模型
合成设置：简化的LCSM（blocks为MLP+GELU，sampler添加噪声）

超参数扫描：

批大小 $B \in \{1,2,4,8\}$
层数 $M \in \{18, 36\}$
嵌入维度 $D \in \{256, 768, 864\}$
序列长度 $L$ ：最大能容纳内存的2的幂（ $2^{15}$ 到 $2^{18}$ ）

硬件：NVIDIA H100和A100 GPU

预热与平均：2次预热，4次运行取平均

对比方法

Baseline：

Lazy：朴素逐位置计算
Eager：提前计算所有未来贡献
Lazy NP / Eager NP：非并行版本（不使用跨层并行）

本文方法的 $\tau$ 实现（7种，4种在Pareto前沿）：

Conv1D：PyTorch默认1D卷积核（需显式填充）
Flash Conv1D：FlashFFTConv的融合核
FFT：PyTorch原生FFT卷积（DFT→逐点乘→IDFT）
FlashFFT：FlashFFTConv的融合FFT核
Hybrid：根据块大小动态选择最优实现

评价指标

端到端时间：生成全部 $L$ 个token的总时间
Mixer累积时间：仅位置混合部分的时间
每token时间：单个token的平均生成时间
加速比：相对于Lazy（并行版本）的倍数提升

实现细节

工程优化：

CUDA Graphs：将单token生成的所有核调度记录为图，后续重放以减少CPU开销（提升10-20%）
FFT预计算：为 $\log_2(L)-1$ 个块大小预计算卷积核的DFT
FlashFFT预配置：为不同块大小预初始化配置以最大化硬件性能
右填充：使用右填充而非左填充，减少一半计算时间
循环卷积：利用循环卷积性质将FFT长度减半

实验结果

主要结果

1. Hyena架构（表1，图2）

Mixer部分加速（相对于Lazy）：

最高110×： $B=1, M=18, D=864, L=2^{17}$
平均64-110×：不同配置下持续显著加速
Eager/Lazy基线：仅0.54×（实际更慢，因为未优化）

端到端加速（表2）：

最高7.8×： $B=8, M=18, D=864, L=2^{15}$
平均3-8×：端到端提升受非mixer部分（MLP等）限制
时间分解（图2a）：mixer从主导地位降至次要部分

每token响应时间（图2c）：

低方差：93.75%的token使用块大小≤8，时间稳定
偶发尖峰：大块计算时出现（但频率低）

2. 合成设置（表3-4，图3）

Mixer加速：

Hybrid：80-124×
单一实现：Flash Conv1D（5.5-6.5×），FlashFFT（31-56×），FFT（74-119×）
Conv1D（二次复杂度）：仍有5-6×加速（验证了分块带来的算术强度提升）

端到端加速：

Hybrid：3.8-11.6×
CUDA Graphs效果：无CUDA Graphs时端到端仅1.6×，使用后达到8×

Pareto最优曲线（图3a）：

不同块大小下，不同 $\tau$ 实现最优
小块（U≤4）：Flash Conv1D最优（内存延迟限制）
中块（4<U≤64）：FlashFFT最优
大块（U>64）：FFT最优（计算密集）

消融实验

1. 跨层并行化效果

Lazy NP vs Lazy：0.76-0.91×（并行化提升10-30%）
Eager NP vs Eager：0.49-0.53×（并行化提升近2倍）
本文方法：小块占主导，并行化效果显著

2. $\tau$ 实现对比（图3b）

Hybrid始终最优或接近最优
FFT在多数情况下接近Hybrid（差距<20%）
Flash Conv1D虽为 $O(L^2)$ ，但仍比Lazy/Eager快5倍（内存友好）

3. 时间分解（图3c，图4）

非卷积部分：在所有方法中保持一致（CUDA Graphs确保）
卷积部分：Hybrid显著优于所有baseline

案例分析

累积mixer时间曲线（图2b，图3b）：

Lazy/Eager：线性增长（斜率恒定）
本文方法：对数增长（斜率递减）
交叉点：约在100-1000 token处，之后优势显著

实验发现

理论与实践一致： $O(L\log^2 L)$ 复杂度在实验中体现为显著加速
内存带宽重要性：Flash Conv1D虽为二次复杂度，但通过优化内存访问仍获得5倍加速
动态选择必要性：无单一 $\tau$ 实现在所有块大小下最优，Hybrid策略关键
CPU开销不可忽视：CUDA Graphs将端到端加速从1.6×提升至8×
并行化收益：小块占主导（87.5%），跨层并行化效果显著

结论与讨论

主要结论

理论贡献：首次为LCSMs提供 $O(L\log^2 L)$ 精确推理算法
通用框架：识别关键属性（贡献基础、查询无关），适用于更广泛架构
实证验证：Hyena上端到端7.8×加速，mixer部分110×加速
系统优化：跨层并行、内存优化、动态实现选择等工程贡献

局限性

数据依赖滤波器：虽理论支持，但需2倍计算量，实验未充分验证
内存需求：仍需存储全部激活 $O(MLD)$ （vs 递归视角的 $O(MD')$ ）
适用范围：
- 不适用于Transformer（不满足查询无关）
- 对极低维SSM（ $D' \ll \log^2 L$ ），递归视角可能更优
提示阶段：长提示时，预填充（prefill）仍主导时间，本文优化的自回归生成相对收益有限
硬件依赖：加速效果依赖GPU内存带宽特性

未来方向

架构设计：设计满足Flash Inference要求且高质量的新架构
因果数据依赖滤波器：如何使滤波器数据依赖同时保持因果性（Arora et al., Karami & Ghodsi已显示潜力）
混合方法：结合递归视角（小状态维度）和卷积视角（大状态维度）
更多架构：扩展到其他满足框架属性的模型（如某些注意力变体）
分布式推理：多GPU/多节点场景下的优化

深度评价

优点

1. 理论严谨性

复杂度分析完整：从Lemma 1到Theorem 2，证明链条清晰
通用框架抽象：P.1和P.2属性抽象恰当，既包含LCSMs又排除不适用情况（如Transformer）
数学工具选择：松弛多项式插值理论应用巧妙

2. 方法创新性

分块策略：平衡的矩形分块（vs 细长条）是关键洞察
跨层并行：识别灰色块无依赖，突破传统层序执行限制
动态实现选择：Hybrid策略体现对硬件特性的深刻理解

3. 实验充分性

多维度评估：端到端、mixer、每token时间
参数扫描全面：21种配置组合（B, M, D, L）
消融实验详尽：7种 $\tau$ 实现，并行vs非并行，CUDA Graphs效果
两种设置：真实Hyena + 合成（排除无关因素）

4. 工程贡献

系统级优化：CUDA Graphs、FFT预计算、循环卷积等实用技巧
开源潜力：算法描述详细，易于复现
内存分析：Appendix D/E对内存使用的细致讨论

5. 写作清晰度

可视化优秀：图1的分块示意图直观
符号一致：全文符号系统清晰
附录完善：扩展讨论、证明、额外实验组织良好

不足

1. 实验局限

无真实模型训练：使用随机初始化权重，未验证对模型质量的影响
缺少端到端对比：未与Mamba等其他高效架构对比
提示阶段分析不足：长提示场景下的实际收益未充分探讨
数据依赖滤波器未实测：Algorithm 5仅理论讨论，无实验验证

2. 方法限制

内存开销： $O(MLD)$ 激活存储在长序列/多层时仍可能成为瓶颈
峰值内存：最大块需额外 $O(LD)$ 空间（虽可通过顺序处理缓解）
适用性受限：
- 不适用于Transformer（主流架构）
- LCSMs本身质量可能不如Transformer
- 需架构满足特定属性

3. 理论分析

常数因子： $O(L\log^2 L)$ 中的常数可能较大（实验显示小块时FFT不是最优）
最优性：未证明 $\log^2 L$ 是否为下界
内存复杂度权衡：未深入分析时间-内存Pareto前沿

4. 对比不足

与近似方法：未实验对比Massaroli et al.的质量-速度权衡
与递归视角：何时递归视角更优的定量分析不足（仅定性讨论 $D' \in O(\log^2 L)$ ）
与结构利用：未对比扩张卷积等特定结构方法

影响力

1. 学术贡献

开创性：首次为LCSMs提供准线性精确推理
理论深度：连接松弛多项式插值与序列模型推理
框架价值：通用属性识别可指导未来架构设计

2. 实用价值

立即可用：Hyena等现有模型可直接应用
工程启发：系统优化技巧（CUDA Graphs等）可迁移
局限性：LCSMs在实际应用中不如Transformer普及，限制直接影响

3. 可复现性

算法清晰：伪代码详细，易于实现
实验细节：超参数、硬件配置明确
开源潜力：虽未提及代码发布，但描述足够复现
硬件依赖：需高端GPU（H100/A100）验证全部结果

适用场景

1. 理想场景

长序列生成： $L > 10^4$ ，复杂度优势明显
自回归主导：生成token数远多于提示长度
LCSM架构：已训练的Hyena等模型
高端硬件：GPU内存带宽高，支持并行化

2. 不适用场景

短序列： $L < 1000$ ，常数开销可能抵消优势
长提示短生成：预填充主导，自回归优化收益有限
Transformer模型：不满足查询无关属性
极低维SSM： $D' \ll \log^2 L$ ，递归视角更优

3. 潜在扩展

混合架构：Transformer + LCSM层（部分层应用本文方法）
近似变体：结合本文精确方法与低秩近似
其他模态：音频、视频生成（卷积更常见）

参考文献（关键文献）

van der Hoeven, J. (1997). Lazy multiplication of formal power series. ISSAC. 理论基础
Poli, M. et al. (2023). Hyena hierarchy: Towards larger convolutional language models. 主要应用对象
Massaroli, S. et al. (2024). Laughing hyena distillery: Extracting compact recurrences from convolutions. NeurIPS. 近似方法对比
Gu, A. & Dao, T. (2023). Mamba: Linear-time sequence modeling with selective state spaces. SSM相关工作
Fu, D. Y. et al. (2023). FlashFFTConv: Efficient convolutions for long sequences with tensor cores. 实现基础
Agarwal, N. et al. (2024). FutureFill: Fast generation from convolutional sequence models. 并行工作

总体评价：这是一篇理论与实践结合紧密的优秀论文。理论上，它为LCSMs推理提供了首个准线性精确算法并抽象出通用框架；实践上，通过系统级优化实现了显著加速。主要局限在于LCSMs本身在实际应用中不如Transformer普及，以及数据依赖滤波器的实验验证不足。该工作为序列模型推理优化提供了新的视角，特别是对未来架构设计具有指导意义。推荐给关注模型效率、序列建模和系统优化的研究者。