TokenWeave：粗粒度 Token 切分与融合 AllReduce-RMSNorm，张量并行推理延迟最高降 29%（MLSys 2026）

背景：张量并行推理里，通信是一笔”每层必缴”的税#

本文的主角 TokenWeave 来自 Microsoft Research India，作者是 Raja Gond、Nipun Kwatra 和 Ramachandran Ramjee。论文于 2025 年 5 月提交到 arXiv（arXiv:2505.11329），2026 年 5 月被 MLSys 2026 接收，代码已在 GitHub 上开源。它解决的问题非常具体：张量并行（Tensor Parallelism，TP）推理中，GPU 之间的通信开销即使在 NVLink 这种超高速互联上也降不下来，而现有”计算-通信重叠”方案在小 batch 下反而更慢。

先铺垫一下张量并行是怎么工作的，否则后面所有讨论都没有支点。

为什么要分布式推理#

大模型推理要跨多张 GPU，原因通常有三个：

模型放不下：Llama-3.3-70B 的权重以 bf16 存储约 140 GB，单张 80 GB 的 H100 根本装不下；
延迟 SLO：即使模型勉强塞进单卡，交互式应用的延迟要求（如 TTFT 小于数百毫秒）也可能要求多卡并行分摊计算；
KV cache 内存：分布式推理把权重分摊到多卡后，每卡省下的显存可以用来放更大的 KV cache，从而支持更大的 batch，吞吐反而更高。

最常见的并行策略就是张量并行（TP）：把单层内的权重矩阵切开，每张 GPU 只算自己那一份，然后通过集合通信把结果拼起来。具体到 Transformer 的两个子层：

FFN 层：第一个线性层按列切分（每个 GPU 持有部分列），第二个线性层按行切分（每个 GPU 持有部分行）。两个线性层之间是逐元素的 GELU 激活，不需要通信；但第一个线性层的输出必须经过一次 AllReduce 把各卡的局部求和结果汇总，才能喂给第二个线性层。最终输出还要再做一次 AllReduce。
Attention 层：沿 head 维度切分。每张 GPU 独立计算自己那部分 head 的 QKV 投影和注意力，最后的输出投影（post-projection）结果通过 AllReduce 合并。

所以每个 Transformer block 的完整前向里，有两次 AllReduce 落在关键路径上（FFN 中间一次、FFN 输出一次——Attention 的 post-projection 合并算在 attention 侧）。模型层数越多，这笔通信税交得越多。

通信开销到底有多大#

论文在 8×H100 DGX 上测了三个模型（Llama-3.3-70B、Qwen2.5-72B、Mixtral-8x22B），把纯通信时间（AllReduce 实际执行时间）与端到端推理延迟做对比：

TokenWeave 论文 Figure 1：三种模型在不同序列长度下 AllReduce 通信开销占比，误差棒为 5 次运行的方差（图来自论文 arXiv:2505.11329）

注意这是单次 AllReduce 的时间占比，是理想化的下界——现实中通信与计算不重叠，比例只会更高。结果：稠密模型在 8×H100 上，通信可以占到端到端延迟的 9%～23%，即便互联是 NVLink 4（900 GB/s 双向）且带 NVSHARP（NVSwitch 内置归约引擎）。Mixtral-8x22B 因为是 MoE，每层只有部分专家被激活，通信量相对较小，占比低一些，但也在 10% 上下。

更直观的对比看论文 Figure 2：把 Llama-3.3-70B 在 8×H100 上、不同序列长度下的前向延迟拆开。vllm-multimem 是”用上了 multimem 指令和 NVSHARP 的优化版 vLLM”，已经算很强的基线；vllm-nocomm 是论文构造的反事实基线——把所有通信从模型里删掉（结果不正确，只用来回答”通信到底拖了多少后腿”）：

TokenWeave 论文 Figure 2：Llama-3.3-70B 在 8×H100 上不同序列长度的推理延迟，虚线为相对 vllm-multimem 基线的归一化性能（图来自论文）

可以看到，序列长度越长，通信占比越大：4K 序列时 vllm-multimem 比 vllm-nocomm 慢 20% 以上。也就是说，即使把现代 GPU 上最先进的通信原语全用上，每算一层仍然有五分之一的 GPU 时间花在等数据上。

还有一个被忽视的开销：RMSNorm#

在 TP 设置下，AllReduce 之后所有 GPU 上持有的是完全相同的 token embedding，而随后的 RMSNorm（现代模型的标准归一化层，残差加法通常也融合在它里面）是逐 token 的逐元素操作——每张 GPU 都在对同一份数据重复做一遍归一化。论文在 8×H100 上实测，单是”AllReduce 之后的 RMSNorm”这一项就贡献了 5%～9% 的推理开销（Figure 3），随序列长度增长。这和通信不同，它纯粹是冗余计算 + 冗余访存，理论上完全可以消除。

TokenWeave 论文 Figure 3：RMSNorm 在不同模型、不同序列长度下的开销占比，可达 5-9%（图来自论文）

为什么已有的”计算-通信重叠”方案救不了小 batch#

通信拖慢推理，最自然的解法是重叠（overlap）：让通信和计算同时进行，把通信延迟藏进计算时间里。问题是数据依赖——一个层的输出算完才能开始通信，通信完才能开始下一层的计算。为了制造重叠的机会，前人把大计算拆成细粒度子任务：子任务 A 算完就发 A 的通信，同时 GPU 继续算子任务 B。代表性工作有三个：

FLUX（arXiv:2406.06858，字节跳动）：把 GEMM 与通信过分解成细粒度操作后融合进一个 kernel，通过 CTA 级流式（tile 粒度）让通信与计算交错；
TileLink（arXiv:2503.20313，MLSys 2025）：把 FLUX 的思路搬进 Triton 编译器，用 tile 中心原语自动生成重叠 kernel，大幅减少手写代码；
NanoFlow（arXiv:2408.12757，2024）：从调度角度切入，把请求切成 nano-batch，在 kernel 粒度上把计算、访存、通信编排到不同 SM 分区上重叠执行。

这些方案在理论上都能把通信藏起来，但到了实际系统里，它们共同面临三个绕不开的问题，论文把它们挨个量化了：

问题一：波量化效应（wave quantization）。 GPU 以”波”为单位执行 kernel：一个 kernel 的 CTA（线程块）被分批调度到 SM 上，每批就是一”波”。如果 CTA 总数不是 SM 数的整数倍，最后一波只有部分 SM 干活。把大计算拆成小计算，会显著放大这种浪费。论文给了个具体例子：一个需要 300 个 CTA 的 GEMM，跑在 132 个 SM 的 H100 上（假设每 SM 一个 CTA），需要 2 个满波 + 1 个 36 CTA 的部分波，共 3 波。如果对半拆成两个 150 CTA 的 kernel，每个都要 132 + 18 的两波，总共 4 波——拆完反而多了 33% 的执行波次。数学上，设原计算需要 $C$ 个 CTA、GPU 有 $N_{SM}$ 个 SM，总波数为：

\text{waves}(C) = \left\lceil \frac{C}{N_{SM}} \right\rceil

拆成 $C_1 + C_2 = C$ 后总波数为 $\lceil C_1 / N_{SM} \rceil + \lceil C_2 / N_{SM} \rceil$ 。只要 $\lceil C_1/N_{SM} \rceil + \lceil C_2/N_{SM} \rceil > \lceil C/N_{SM} \rceil$ ，拆分就白白多付波数，而且拆得越碎，浪费越大——这就是为什么细粒度方案在小 batch 下特别吃亏。

问题二：拆集合通信本身更慢。 FLUX/TileLink 的做法是把一次 AllReduce 拆成 ReduceScatter（RS）再 AllGather（AG），分别与前后两个 GEMM 重叠。但论文实测（Figure 5）：RS+AG 的组合比等价的单次 AllReduce 慢最多 50% 以上，因为通信被拆小后，每次传输的报文变小，NVLink 带宽利用率大幅下降（Figure 6 显示：ReduceScatter 的带宽随 tensor 变大而显著提高，小报文跑不出带宽）。也就是说，这类方案首先要赔上”拆通信”的钱，再指望用重叠赚回来。

问题三：通信 kernel 本身占 SM。 传统通信 kernel（如 NCCL 的实现）要驱动大量的 SM 才能跑满带宽，这些 SM 在重叠期间被通信占用，算不了计算，把重叠的收益吃掉一截。论文实测基于 multimem 的通信只需要 6%～8% 的 SM 就能饱和带宽（后面会细讲），而此前细粒度融合方案里通信逻辑是嵌在 GEMM kernel 内部的，会抬高 kernel 的寄存器和共享内存占用，降低 CTA 占用率。

三个问题叠加的结果是：这些方案只有 batch 内 token 数很大（8000 以上）时才有净收益。论文拿 TileLink 做了单层对比（Figure 14）：序列长度 1K 时，TileLink 的归一化延迟反而比不重叠的 vllm-multimem 基线慢 50% 以上；只有序列长度超过 4K 才开始盈利。因此直到今天，vLLM、SGLang、TensorRT-LLM 这些主流推理框架默认都不开启 TP 推理的计算-通信重叠——不是没意识到问题，是现有方案在真实负载（batch 往往只有几百到两三千 token）下不划算。

TokenWeave 全景：两路切分 + 双流流水线#

TokenWeave 的核心思想一句话概括：把 batch 里的 token 粗粒度地切成两个近似等量的子集，让一个子集的计算与另一个子集的”通信 + RMSNorm”重叠执行，并用一个融合的 AllReduce-RMSNorm kernel 把通信占用的 SM 压到 2～8 个，让重叠在 token 数只有 1024 时就开始盈利。

下图是论文的总览图，左边是标准 TP，右边是 TokenWeave：

TokenWeave 论文 Figure 8(a)：标准张量并行，计算与通信串行执行（图来自论文）

TokenWeave 论文 Figure 8(b)：TokenWeave 将 batch 切成两路，通信与 RMSNorm 融合后与另一路的计算重叠（图来自论文）

对照两图，TokenWeave 的完整结构由三个技术支柱构成：

Token-Splitting（Token 切分）：粗粒度地把输入 batch 切成两个子集，并做波感知的 Smart-splitting，保证拆分不引入额外的 kernel 波数；
RMSNorm 重排序：把 AllReduce 拆成 ReduceScatter → RMSNorm → AllGather，让 RMSNorm 的计算量除以 GPU 数 $N$ ，消除冗余；
融合 AllReduce-RMSNorm kernel：用 Hopper 的 multimem 指令把”归约 + 归一化 + 广播”融合成一个只占 2～8 个 SM 的轻量 kernel，顺带把 RMSNorm 的访存也省掉一层。

调度上，两条独立的 CUDA stream（计算流 + 通信流）配合成流水线：一个子集在计算流上跑 GEMM/Attention，另一个子集的融合通信 kernel 同时在通信流上跑。下面分别拆解这三个技术。

技术一：粗粒度 Token 切分与波感知 Smart-splitting#

为什么是”两个子集”，而不是细粒度#

前面说过，细粒度拆分（tile 级、nano-batch 级）会被波量化效应反噬。TokenWeave 的选择是只切成两半，每半承担近似等量的计算和通信。为什么两路就够了？因为重叠的本质只需要”一个在算、另一个在通信”的乒乓结构：A 算的时候 B 在通信，B 算的时候 A 在通信，两个子集互相掩护，任何时刻 GPU 都有活干。切得更多理论上能进一步压低尾波，但每一路拆分都会引入新的波量化损失与调度复杂度，两路是收益/代价的甜点。

切 batch 维度还是序列维度#

Token 级别的切分（而不是请求级别）带来一个自由：既可以把 batch 里不同请求拆开（batch 维度），也可以把单个请求的序列拆成前后两段（序列维度）。

大部分 Transformer 操作（FFN、RMSNorm、QKV 投影）是逐 token 的，拆开后两路完全独立，没有问题；
但 Attention 有依赖：序列后半段 token 的注意力要看到前半段 token 的 KV。论文的处理是采用 Sarathi 风格的 chunked attention（分块注意力，Sarathi 论文），并规定一个不变式：所有请求的前缀部分必须放进 prefix-split，后缀部分放进 suffix-split——任何时刻 suffix 的注意力计算都发生在 prefix 的 KV 已就绪之后。对于 batch 里只有部分序列被切的情况，同样保证所有被切序列的 prefix 都在 prefix-split 里。

由于这个顺序约束，两条流的流水线总是 prefix-split 先行，suffix-split 跟随，天然形成稳定的乒乓节奏。

Smart-splitting：让波数不增加#

切分方式直接决定波量化损失。还是用论文的例子：300 CTA、132 SM。对半切（150 + 150）是 4 波，损失 33%；但如果切成 132 + 168，第一路恰好一个满波，第二路 168 = 132 + 36 是两波，总波数 3 波——与原计算完全持平。这就是 Smart-splitting：按 GPU 的 SM 数感知当前 wave 的填充情况，选择切分点使得：

\left\lceil \frac{C_1}{N_{SM}} \right\rceil + \left\lceil \frac{C_2}{N_{SM}} \right\rceil \le \left\lceil \frac{C}{N_{SM}} \right\rceil

即两路子 kernel 的波数之和不超过原 kernel 的波数。论文 Figure 9 给出了 FFN 层的实测对比（无拆分 vs 等分 vs smart-split）：

TokenWeave 论文 Figure 9：Smart-splitting 消除波量化开销，横轴为 batch 内 token 数，柱状为延迟、折线为相对无拆分归一化时间（图来自论文）

关键观察：token 数越少，等分拆分的相对损失越大（细柱越高），而 smart-splitting 几乎贴着”无拆分”的基线走。这正是 TokenWeave 敢在小 batch（1024 token）下做重叠的原因。

技术二：RMSNorm 重排序——把冗余计算除以 N#

前面提到，TP 中 AllReduce 之后每张 GPU 都持有完整的、相同的 token embedding，随后每张 GPU 独立做 RMSNorm，等于把逐 token 的归一化在 $N$ 张卡上重复了 $N$ 遍。RMSNorm 的计算式是：

\text{RMSNorm}(\mathbf{x}) = \frac{\mathbf{x}}{\sqrt{\frac{1}{d}\sum_{i=1}^{d} x_i^2 + \epsilon}} \odot \boldsymbol{\gamma}

其中 $d$ 是隐藏维度， $\boldsymbol{\gamma}$ 是逐维缩放权重， $\epsilon$ 是数值稳定项。注意它有两个阶段：先遍历一遍求均方根（ $\frac{1}{d}\sum x_i^2$ ），再遍历一遍做缩放。也就是说，传统实现里每个 token 的 embedding 要被读两遍、写一遍（一次读算方差，再一次读+写做缩放）。

TokenWeave 的重排序思路：把一次 AllReduce 拆成 ReduceScatter（RS）+ AllGather（AG）。RS 结束时，每张 GPU 拥有完整张量的 $1/N$ ——只要 RS 沿着 token 边界切分，每张 GPU 手里的 $1/N$ 就是若干条完整的 token embedding。既然 RMSNorm 是逐 token 操作，每张 GPU 就可以立即对自己这 $1/N$ 做归一化，然后再 AG 广播出去：

\text{AllReduce}(\mathbf{X}) \rightarrow \text{ReduceScatter}(\mathbf{X}) \rightarrow \text{RMSNorm}(\cdot) \rightarrow \text{AllGather}(\cdot)

这样 RMSNorm 的计算量直接除以 $N$ ：每张卡只归一化自己那份，不再重复算别人的。代价是原本一次 AllReduce 被拆成 RS + AG 两次通信——而前面 Figure 5 已经说明，拆通信本身要赔 50% 的通信效率。论文实测（Figure 4）确实如此：单纯重排序（不融合）比”AllReduce 之后再 RMSNorm”更慢，重排序省下的计算完全被拆通信的开销抵消。所以重排序必须搭配第三个技术——融合 kernel——才能真正落地。

技术三：融合 AllReduce-RMSNorm kernel#

前置知识：NVSHARP 与 multimem 指令#

Hopper 架构的第四代 NVSwitch（NVLink 4 时代）内置了 NVLink SHARP（NVLS） 引擎：交换机里的 ASIC 可以在数据经过交换矩阵时顺路完成归约运算。GPU 侧通过 multimem 系列 PTX 指令使用这一能力：

multimem_ld_reduce_add：从多播地址加载数据，交换机在转发前先把各卡的数据按位相加，GPU 拿到的是归约后的结果；
multimem_st：把数据写进多播地址，交换机自动复制一份发给订阅了该地址的每张 GPU。

也就是说，AllReduce 的”归约”和”广播”两个环节都可以在交换矩阵里完成，几乎不占用 GPU 的 SM 算力和 HBM 带宽。论文实测：用 multimem 实现的 AllReduce，4～8 个 SM 就能跑出接近最优的延迟（Figure 7），远远少于传统 NCCL kernel 的 SM 占用。软件侧，PyTorch 2.6 通过 SymmetricMemory API 暴露这些能力：symm_mem.empty 分配跨卡对称缓冲区，symm_mem.rendezvous 交换句柄把远端/多播指针映射进本地地址空间，之后在 Triton 或 CUDA kernel 里直接用普通内存操作读写远端，不再需要显式 NCCL 调用。

融合 kernel 怎么省访存#

融合 kernel 的流程：RS 阶段用 multimem_ld_reduce_add 把归约结果直接读进寄存器 → 立即在寄存器上累加平方和算方差 → 用方差缩放并乘上 $\gamma$ 、加上残差 → 用 multimem_st 把归一化结果直接写进多播地址完成 AG。相比”重排序但分开执行”的版本，它省掉了两笔 HBM 流量：

省一次读：方差直接在 multimem 归约结果上算，不用先把 RS 结果写回 HBM 再读出来算；
省一次写：归一化结果直接写到多播地址完成 AG，不用先写 HBM、再由 AG 读出去。

加上残差加法也融合在内，这个 kernel 对每个 token 的 HBM 访问从”RS 写 + 2 读 + 1 写 + AG 读”压缩到几乎只剩 multimem 带来的网络流量。论文实测（Figure 4），相比”AllReduce + RMSNorm”分开做的基线，融合 kernel 在所有序列长度上都有最多 40% 的延迟提升——也就是说，TokenWeave 不仅在补通信的窟窿，还顺带把归一化本身也优化了：

TokenWeave 论文 Figure 4：隐藏维度 8192、bf16 下三种实现的单次 AllReduce+RMSNorm 延迟对比——重排序不融合反而更慢，融合 kernel 全程最优，最高提升 40%（图来自论文）

kernel 的核心代码（论文 Listing 1，节选）长这样：

1
template <typename scalar_t, int width>
2
__global__ fused_rs_ln_ag_cta_kernel(...) {
3
    const int vec_hidden_size = hidden_size / width;
4
    int tokens_per_cta = (num_tokens + gridDim.x - 1) / gridDim.x;
5

6
    // 等待 RS 数据就绪（多播归约由 NVSwitch 完成）
7
    sync_remote_blocks<MemOpSem::Relaxed>(signal_pads, rank, world_size);
8
    __syncthreads();
9

10
    for (int iter = 0; iter < tokens_per_cta; iter++) {
11
        int token_id = blockIdx.x + iter * gridDim.x;
12
        if (token_id >= num_tokens) continue;
13

14
        float variance[1] = {0.0f};
15
        __shared__ float s_variance;
16
        int offset = token_id * vec_hidden_size;
17
        int offset_scalar = token_id * hidden_size;
18
        auto input_o = input_v + offset;
19
        auto residual_o = residual_v + offset;
20

21
        // 第一遍：multimem 归约读 + 残差相加 + 累加方差
22
        for (int idx = threadIdx.x; idx < vec_hidden_size; idx += blockDim.x) {
23
            auto multimem_temp = multimem_ld_reduce_add<16>(
24
                multimem_address_ptr + offset_scalar + idx * width);
25
            vec_t temp = *(reinterpret_cast<vec_t*>(&multimem_temp));
26
            temp += residual_o[idx];              // 残差融合
27
            variance[0] += temp.sum_squares();    // 直接算方差，省一次 HBM 读
28
            residual_o[idx] = temp;
29
        }
30

31
        blockReduceSum<float, 1>(variance);
32
        if (threadIdx.x == 0)
33
            s_variance = rsqrtf(variance[0] / hidden_size + epsilon);
34
        __syncthreads();
35

36
        // 第二遍：缩放 + 乘 gamma + 直接写多播地址完成 AllGather
37
        for (int idx = threadIdx.x; idx < vec_hidden_size; idx += blockDim.x) {
38
            vec_t temp = residual_o[idx] * s_variance * weight_v[idx];
39
            multimem_st<16>(mcptr + offset + idx * width,
40
                            *(reinterpret_cast<Vec<16>*>(&temp)));
41
        }
42
    }
43
    __syncthreads();
44
    sync_remote_blocks<MemOpSem::AcqRel>(signal_pads, rank, world_size);
45
}

几个关键实现细节：

每 CTA 处理 tokens_per_cta 个 token，kernel 只需要 2～8 个 SM 的 grid（Figure 10 显示 8 个 SM 已经接近最优，大序列时 4 个就够）——这么小的 SM 占用是它能和另一路计算重叠的前提；
方差计算和缩放两遍循环都直接基于 multimem 读到的数据，multimem_st 完成 AG 的写出，整个 kernel 对本地 HBM 的读写趋近于零；
开头和结尾各有一个 sync_remote_blocks，用信号量（signal pads）与对端 GPU 同步，保证 RS 数据就绪后才读、写完才通知别人；
当前实现只支持 bf16，向量宽度 width=16 字节。

为什么”少用 SM”这么关键#

重叠方案里，通信 kernel 占的 SM 越少，留给计算 kernel 的资源就越多。传统通信实现要占十几个甚至二十个 SM，重叠时计算 kernel 被挤到剩余 SM 上，吞吐打折，经常把重叠的收益吃光。TokenWeave 的融合 kernel 只占 2～8 个 SM，而且即使因为 SM 太少而慢一点也没关系——只要它跑得比另一路的计算快，慢的部分就被完全隐藏。这给了 TokenWeave 一个别人没有的余量：RMSNorm 这个访存密集型操作也能被重叠掉，而不只是通信本身。

系统集成：vLLM V1 里的双流流水线#

论文在 vLLM 的 V1 引擎上完成了实现与开源：基于 vLLM==v0.8.5.dev106，通信集合基于 PyTorch 2.6 的 SymmetricMemory + Triton 3.2 的 multimem 配方实现。集成方式不是侵入式的框架改造，而是在模型文件层面（llama.py 这类文件）把”每层两个 AllReduce”替换为”两路切分 + 融合通信 kernel”的流水线，因此扩展新模型只需要改模型定义。

调度上使用两条 CUDA stream：计算流跑一路的 FFN/Attention kernel，通信流跑另一路的融合 AllReduce-RMSNorm kernel。vLLM 本身提供的 CUDA Graph 捕获机制对重叠场景是个挑战（图上所有 kernel 必须静态确定），论文的实现为重叠路径做了配套处理，decode-only 的纯小 batch 则走非重叠的融合 kernel 路径。

实际部署策略是混合的：对 batch 内 token 数 ≥ 1024 的批次（Mixtral 是 ≥ 4096）启用两路重叠，对 decode-only 的小批次仍然使用融合 kernel 但不做切分重叠。这是因为 decode 批次往往只有几十个 token，切分后每路的计算时间比通信还短，重叠没有意义。测试环境是 8×H100 DGX（带 NVSHARP），TP-4 实验使用其中 4 张卡；实验关闭了 prefix caching、微基准测试前会 flush L2 cache，保证测量的公平性。

实验效果：18%～29% 延迟、15%～26% 吞吐#

单次前向延迟#

论文先测”单次迭代的前向延迟”，覆盖三种模型在 8×H100 和 4×H100 上的五种配置（Mixtral-8x22B 塞不进 4×H100，所以少一种）。结果（Figure 11）：

稠密模型在 8×H100 上，序列长度 1K～32K 全程保持 18%～29% 的延迟收益，即使在 1K 这样的小序列也有 18%——这是此前所有方案做不到的（TileLink 在 1K 是净损失）；
序列长度 ≥ 4K 时，TokenWeave 的延迟低于 vllm-nocomm——不仅把通信开销全部赚回来，还因为融合 kernel 优化了 RMSNorm 而额外倒赚一截；
4×H100 上收益略低于 8×H100（GPU 少，通信占比本来就低），但从 1K 起全程为正；
MoE 模型是例外：Mixtral-8x22B 在 1K、2K 序列时 TokenWeave 反而是净开销。原因是 MoE 的 FFN 把 token 分发到 8 个专家上，小序列时每个专家的负载很小，FFN 变成访存受限（权重加载占主导），此时再切两路，每路的访存效率进一步下降，重叠省下的通信时间补不回来。

端到端吞吐#

生产环境更关心吞吐。论文按 vLLM 的默认配置开启 chunked prefill（分块预填充） 做混合批处理（prefill chunk 与 decode 请求同批），固定输入/输出长度的负载以及 ShareGPT 真实对话轨迹都测了：

TokenWeave 论文 Figure 12(a)：8×H100 上端到端负载的吞吐提升（图来自论文）

chunked prefill 大小为 2K（vLLM 默认 max_num_batched_tokens）时，稠密模型吞吐提升约 20%，把大部分通信开销收复；
把 chunk 大小从 1024 扫到 8192，收益稳定在 15%～26%，对 chunk 大小不敏感，说明它能和现有调度器无缝协作。

这也是论文强调的实用价值：vLLM 默认的 chunked prefill 就是 2048 token，Sarathi 式调度器的 chunk 也在这个量级，TokenWeave 恰好在这个尺度上盈利；而把 prefill/decode 彻底分离的架构（如 Splitwise、DistServe）里，batch 内 token 数通常 ≥ 4096，TokenWeave 的收益在 20% 以上。

与 TileLink、NanoFlow 的对比#

vs TileLink（单层基准，batch=1）：1K 序列时 TileLink 是净开销（比基线慢 50% 以上），4K 以上才开始盈利，且收益封顶约 21%；TokenWeave 在 1K 就有约 20% 的单层收益，最高 38%。端到端收益比单层低是因为 embedding、采样等非层内开销不参与优化；
vs NanoFlow：论文把 NanoFlow 移植到 H100 后在其自有框架内做了相对对比，NanoFlow 的收益只有 5%～8%（与它论文自报的 7% 通信收益一致），而 TokenWeave 在相同负载上约 20%。

局限与未解决的问题#

TokenWeave 的边界条件论文交代得很清楚，这里逐条列出：

MoE 小序列场景失效。专家分发让 FFN 变访存受限，两路切分雪上加霜，1K/2K 序列时是净损失。MoE 正在成为主流（DeepSeek、GLM、Qwen-MoE 都是），这个缺口值得后续工作处理；
强依赖 NVSHARP/multimem 硬件。融合 kernel 的省访存魔法全部建立在 NVSwitch 4 的交换机内归约上，A100（NVLink 3 无 SHARP）跑不了，AMD 平台也没有对应原语。好在 Blackwell 继续支持 multimem，论文代码库的说明也提到在 B200 上有评估；
只支持 bf16。kernel 的向量宽度和归约路径都是按 bf16 写的，fp16/int8 量化模型要额外适配；
两路切分的粒度上限。batch 极小（几十 token）时不切分，中等 batch 切两路，但没有尝试三路、四路——更多路数可能进一步压缩尾波，但需要更复杂的流调度；
模型文件级集成。当前改动在 llama.py 这类模型定义里，开源代码覆盖 Llama-3.3-70B、Qwen2.5-72B、Mixtral-8x22B 三个模型，换架构要改模型文件；SGLang、TensorRT-LLM 尚未集成；
流与 CUDA Graph 的兼容成本。vLLM 依赖 CUDA Graph 摊薄 kernel 启动开销，双流重叠路径要小心处理图捕获，论文的实现做了配套处理，但这部分是工程上最容易出暗坑的地方。

小结#

TokenWeave 的启示不在于”又一个重叠方案”，而在于它把重叠的盈亏账算清楚了：细粒度拆分赔在波量化，拆通信赔在带宽利用率，通信占 SM 赔在计算资源——三个坑它分别用波感知的两路粗切分、RMSNorm 重排序 + 融合 kernel、multimem 低 SM 占用来填，最后还靠融合 kernel 白赚了 RMSNorm 的访存优化。结果是主流推理框架终于有了一个在 1024 token 级别就能稳定盈利、且已经开源进 vLLM 生态的 TP 通信重叠方案。对正在学 GPU 推理优化的读者，这篇论文是”如何从硬件原语出发设计系统优化”的很好范本：先量化瓶颈，再逐个击破，最后用端到端数据说话。

音乐

音乐