💬 NLP

WIDE: Boosting Adaptive LLM Inference via Token-level Dynamic Width Pruning

该论文介绍了 WIDE，这是首个端到端可微框架，通过允许对注意力头和 FFN 通道进行细粒度的选择，并结合专门的算子协同设计，实现了大语言模型（LLM）的 Token 级动态宽度剪枝，从而在保持高准确度的同时显著提升了端到端的推理加速。

原作者： Haozhe Hu, Hao Wu, Peiran Yin, Chao Han, Yunpu Ma, Xiaoyu Shen

发布于 2026-07-31

📖 1 分钟阅读☕ 轻松阅读

原作者： Haozhe Hu, Hao Wu, Peiran Yin, Chao Han, Yunpu Ma, Xiaoyu Shen

原始论文采用 CC BY 4.0 许可（http://creativecommons.org/licenses/by/4.0/）。 ✨ 这是对下方论文的AI生成解释。它不是由作者撰写或认可的。如需技术准确性，请参阅原始论文。阅读完整免责声明

想象一下，你正试图在一个微小的、电池供电的设备上运行一个庞大的、极其聪明的机器人大脑。这个机器人大脑是一个大语言模型（LLM），一种可以写故事、解数学题并像人类一样聊天的 AI 类型。但问题在于，这些大脑太大了。它们如此庞大且渴望能量，以至于通常需要超级计算机才能运行，这使得在你的手机或笔记本电脑上使用它们时既慢又昂oly。

为了解决这个问题，科学家们一直尝试对这些大脑进行“剪枝”（pruning）。把剪枝想象成修剪一株巨大的、过度生长的灌木丛。你剪掉那些没做什么工作的枝条，使植物变得更小、更快。长期以来，最好的方法是永久地剪掉灌木的整个部分，无论这株植物当天正在做什么。但这有点笨拙；有时你会剪掉一根实际上在特定任务中仍有用的枝条，导致机器人变得健忘。最近，一些更聪明的方法被发明了出来，让机器人可以在“思考时”决定使用哪些部分。然而，这些聪明的方法仍然有点粗糙——它们要么使用整个区域，要么完全跳过，就像决定是使用整个房间还是让整个房子空着一样，而不是仅仅点亮几盏特定的灯。

这篇论文介绍了一个名为 WIDE（基于宽度的推理与动态执行）的新系统。它就像是为机器人的大脑中的每一盏灯泡都配备了一个超精密的调光开关。与其关闭整个房间，WIDE 允许机器人针对它处理的每一个单词，精确地决定哪些微小的神经元组（大脑细胞）应该亮起，哪些应该保持黑暗。研究人员构建了一个特殊的“交通控制器”（路由器），它学会了做出这些瞬间的决策。他们还设计了一种新的方式，让计算机硬件能够处理这些决策而不会感到困惑或变慢。结果是，机器人在保持原有智能水平的同时，运行得更快，消耗的能量也更少，即使我们剪掉了它一半的大脑容量。

问题所在：“全或无”的陷阱

想象你是一名经营着大型厨房的厨师。你有数百名厨师（神经元）共同协作烹饪美食。在过去，如果你想节省时间，你可能会永久解雇一半的厨房员工。如果你总是烹饪同样的简单菜肴，这没问题；但如果你突然需要烤一个复杂的蛋糕，你可能已经解雇了唯一需要的面包师。这就是静态剪枝发生的情况：模型的修剪是一次性的，无论它在回答什么具体问题。

后来出现了动态剪枝，这就像聘请了一位经理，他可以告诉厨师们：“嘿，对于这道特定的订单，只需要烤架站即可；烘焙站可以休息了。”这虽然有所进步，但那位经理仍然有点笨拙。他们会说：“跳过整个烘焙站，”即便烘焙站只需要使用十个烤箱中的两个。对于大块的厨房区域来说，这是一个“全或无”的决策。

问题在于，现代 AI 模型如此复杂，跳过一大块区域往往会丢弃有用的信息，从而损害答案的质量。此外，如果经理每秒钟都在不断更改排班表，厨师们就会感到困惑，由于所有的切换开销，实际的烹饪速度并不会变得快多少。

解决方案：WIDE 的“调光开关”

这篇论文的作者（来自宁波数字孪生研究院和慕尼黑大学 LMU）提出了 WIDE。与其解雇整个部门或跳过整个房间，WIDE 让模型可以针对处理的每一个单词（token），精确地决定激活哪些小组神经元。

把模型的脑部想象成一个巨大的灯光开关网格。

旧的动态方法： “关掉房子的整个左翼。”
WIDE： “对于这个特定的单词，点亮厨房里的灯，但仅限于炉灶和冰箱上方的那些。让厨房的其他地方保持黑暗。”

WIDE 通过在模型的每一层都连接一个轻量级的“路由器”（一个微型决策者）来实现这一点。这个路由器观察当前的单词并询问：“我需要这一组注意力头（观察其他单词的部分）吗？”以及“我需要这一组 FFN 通道（处理含义的部分）吗？”它做出一个二元选择：使用这一组或跳过它。

至关重要的是，WIDE 不仅仅停留在做出决策，它还解决了硬件方面的难题。通常，如果你告诉计算机跳过计算的随机部分，计算机会在寻找数据位置时变得效率低下。WIDE 使用了一种被称为**掩码重排序（mask reordering）**的巧妙技巧。想象你有一堆乱七八糟的邮件，其中一些是用于投递的信件，而另一些是垃圾邮件。与其先扔掉垃圾再尝试分拣邮件，WIDE 会先将这堆邮件重新排列，使所有“保留”的信件整齐地排在顶部的一个块中，而所有“垃圾”信件则排在底部。这使得计算机可以非常高效地处理“保留”块，而不需要停下来检查每一件邮件。

研究发现：速度与智能

研究人员在流行的 AI 模型（如 Llama 3.1，80 亿参数和 Llama 3.2，30 亿参数）上测试了 WIDE。他们将其与现有的最优秀的静态和动态剪枝方法进行了对比。

以下是数据所反映的情况：

更聪明的剪枝： 当他们剪掉 50% 的模型容量（一种非常激进的修剪）时，WIDE 比其他任何方法都能更好地保持模型的智能。例如，在 Llama 3.1 模型上，WIDE 在剪枝后保留了 86.42% 的原始准确度，而排名第二的动态方法（SkipGPT）仅保留了 59.42%。甚至在没有额外微调的情况下，WIDE 已经超越了最好的静态方法。
现实世界的速度： 论文测量了模型实际运行的速度。对于“预填充”（prefill）阶段（读取长提示词），WIDE 比原始模型快了 1.68 倍。对于“解码”（decode）阶段（逐字生成答案），它快了 1.55 倍。
内核级魔法： 如果只看数学计算（忽略其他开销），加速效果更加显著，在预填充阶段达到了 1.98 倍，在解码阶段达到了 4.95 倍。这表明瓶颈在于以往方法处理“跳过”操作时的低效方式，而 WIDE 的新硬件设计解决了这个问题。

作者还发现模型学会了非常有策略性。它不仅仅是随机跳过部分，而是学会了跳过“无聊”的词（如“the”或“a”）并保留“重要”的词（如“running”或“track”）。在一项测试中，模型跳过了 66% 的注意力工作，但仅跳过了 28% 的处理工作，这表明它清楚地知道在哪里节省能量。

核心结论

WIDE 表明，高效 AI 的未来不仅仅在于把模型剪得更小，而是在于让模型更聪明地使用其剩余的部分。通过从“跳过整个房间”转向“调暗特定的灯光”，并通过重新设计厨房以高效处理这些调光器，作者创建了一个框架，使大型 AI 模型在不变得健忘的前提下，运行得更快、更高效。

虽然结果令人鼓舞，但论文指出，这是一个针对 GPU 硬件（计算机中的芯片）的特定解决方案，并且依赖于一个两阶段训练过程：模型首先学习如何做出决策，然后通过快速微调来恢复任何损失的准确度。这是向让强大的 AI 在日常设备上普及迈出的重要一步，证明了你不需要一个巨大的大脑也能变得聪明——你只需要知道在正确的时间使用大脑的哪些部分。

技术摘要：WIDE —— 通过 Token 级动态宽度剪枝提升自适应 LLM 推理效率

1. 问题陈述

大型语言模型（LLMs）在规模扩大的过程中面临着显著的效率挑战。虽然**静态结构化剪枝（static structured pruning）**具有硬件友好性并能提供吞吐量增益，但其输入无关（input-agnostic）的特性往往会导致在激进稀疏化时出现严重的精度下降，因为它对所有输入都应用固定的移除模式。相反，**动态剪枝（dynamic pruning）**方法虽然能根据单个输入调整计算，但目前受到以下限制：

粒度粗糙： 现有的动态方法主要在层级层面进行操作（决定是否执行或跳过整个层或子模块）。这种“全或无”的方法无法捕捉那些在层内仅需要部分计算的 Token，从而导致不必要的容量损失。
系统性低效： 将动态剪枝推向更细粒度（例如，Token 级的宽度分配）会引入不规则的执行模式。依赖显式 gather-scatter 操作的朴素实现会带来沉重的设备到主机（device-to-host）同步、过度的内存流量以及存储开销，阻碍了理论上的 FLOPs 减少转化为现代 GPU 上的实际端到端加速。

2. 方法论：WIDE 框架

作者提出了 WIDE，这是首个针对预填充（prefill）和解码（decode）场景设计的端到端可微 Token 级动态宽度剪枝框架。该方法由三个核心组件组成：

A. 精细化 Token 级路由

与以往在层级进行路由的方法不同，WIDE 使每个 Token 能够独立选择每一 Transformer 块内细粒度的**注意力头（attention heads）**和 **FFN 通道（FFN channels）**组。

剪枝单元：
- 注意力（Attention）： 剪枝与分组查询注意力（GQA）结构对齐的注意力头组，以保持与现有解码优化（如 XQA）的兼容性。
- FFN： 剪枝中间维度（ $D'$ ）中的连续通道组，与标准 GEMM 分块策略兼容。
路由器架构： 轻量级瓶颈路由器将输入激活映射为每个组的二元决策（执行 vs. 跳过）。在训练期间，使用 Gumbel-Softmax 松弛技术实现可微优化。
训练流水线： 采用两阶段过程：
1. 路由器训练： 冻结所有基础模型参数；仅在校准数据上训练路由器，以学习稀疏执行模式。
2. LoRA 恢复（可选）： 在线性层中添加低秩自适应（LoRA）模块以恢复模型质量，随后进行微调。

B. 统一掩码重排序预处理（Unified Mask-Reordering Preprocessing）

为了解决不规则 Token 级路由导致的低效问题，WIDE 引入了一个预处理步骤，将任意路由掩码转换为结构化执行布局，而无需显式的 gather-scatter 操作。

机制： 对于每个路由组，沿 Token 维度（ $M$ ）对二元掩码进行排序。这会将活跃 Token 聚类成一个连续的前缀。
优势： 这将不规则的稀疏性转化为 CTA 级的规整性。一个 Tile 内的活跃行是连续的，允许内核将其视为具有已知活跃前缀的稠密块，从而避免了实例化特定 Token 稀疏张量的需求。

C. 剪枝-内核协同设计（Pruning-Kernel Co-Design）

WIDE 在自定义 CUDA 内核（使用 CuTe, Triton 和 TileLang）中实现了多阶段跳过策略，以利用重排序后的掩码：

CTA 级早期退出 ( $\ell=0$ )： 如果特定组中的整个 Token 块都是非活跃的，则整个计算线程阵列（CTA）跳过执行。
块内跳过 ( $\ell=1, 2$ )：
- 加载阶段 ( $\ell=1$ )： 跳过加载非活跃内存包（例如，使用 cp.async 或 TMA gather）。
- MMA 执行阶段 ( $\ell=2$ )： 跳过非活跃 Tensor Core 计算分片。
  这种设计确保了细粒度动态稀疏性在消除不必要计算的同时，仍能保持高效的 GPU 执行模式（例如，合并内存访问）。

3. 核心贡献

首个端到端 Token 级动态宽度剪枝： WIDE 将动态剪枝从粗粒度的层/子层决策推向了神经元块级粒度，允许 Token 动态选择注意力头和 FFN 通道组。
统一加速框架： 本文识别出同步掩码到索引的转换以及不规则的 gather-scatter 是关键瓶颈。WIDE 通过独立的掩码重排序和多阶段块内谓词执行解决了这些问题，使内核能够接近可剪枝操作的理想加速比。
剪枝-内核协同设计： WIDE 是第一个将 Token 级动态宽度剪枝与针对预填充及解码的 GPU 内核协同设计的框架。它证明了细粒度动态剪枝必须与硬件执行路径同步设计。

4. 实验结果

实验在 Llama3.1-8B 和 Llama3.2-3B 上进行，使用 RedPajama-1T 数据集进行校准和 LoRA 恢复。

精度保持：
- 在 50% 稀疏度（仅校准）下，在 Llama3.1-8B 上，WIDE 比最强的动态深度基准（SkipGPT）在平均零样本准确率上提高了 8.80 个百分点（61.84 对比 53.04）。
- 使用 LoRA 恢复后，WIDE 在 Llama3.1-8B 上实现了 90.59% 的准确率保持率，显著优于 SkipGPT (86.09%) 和静态方法 DDP (76.36%)。
- WIDE 在不同的组大小（ $G$ ）下均保持了稳健的性能，其中 $G=32$ 或 $G=64$ 通常能产生最佳权衡。
加速与效率：
- 内核级： 与朴素的 gather-scatter 基准相比，WIDE 实现了高达 200 倍的加速和 1,000 倍的峰值内存减少。
- 端到端： 在 50% 稀疏度下，WIDE 在预填充阶段实现了 1.68x 的加速，在解码阶段实现了 1.55x 的加速。
- 理论极限： 该框架在预填充阶段达到了接近理论极限的 1.98x 内核级加速，在解码阶段达到了 4.95x。
- 开销： 在 0% 稀疏度下，最佳配置保留了 98.60% 的稠密吞吐量，表明路由和内核初始化的开销极小。

5. 重要性与主张

本文声称 WIDE 建立了 Token 级动态结构化剪枝的新前沿，通过证明细粒度动态剪枝必须与硬件执行路径协同设计，而不仅仅是作为一种模型压缩问题来处理。

质量-效率权衡： WIDE 成功推动了帕累托前沿（Pareto frontier），提供了比现有动态深度剪枝方法显著更好的质量保持能力，同时提供了实际的、现实世界的加速。
语义感知： 对路由模式的分析表明，WIDE 学会了根据语义内容分配计算预算（例如，跳过像 "the" 这样的虚词，同时保留像 "running" 这样的实词），验证了 Token 级自适应性的有效性。
实际可行性： 通过通过掩码重排序和内核协同设计解决不规则执行的系统级挑战，WIDE 证明了细粒度动态稀疏性能在现代 GPU 架构（特别是 NVIDIA RTX 5090 和 A100）上实现真实的端到端加速，弥合了理论稀疏性与实际推理加速之间的鸿沟。