X-Stage: An Overlooked Pipeline Stage for Communication–Computation Overlap in DiT Inference

发表时间: 2026-07 · arXiv:2607.23264 (KlingAI/Tsinghua/NVIDIA)

原文: https://arxiv.org/abs/2607.23264

作者/机构: Jianwen Xian1, Zhiyuan Xu2,,‡, Yuchen Li2,,‡, Ziliang Lai1, Kang He1, Zhen Huang3, Aichen Feng3, Jinyan Chen3, Yilin Zhang3, Qinqin Chen3, Chengru Song1,†
1KlingAI Research, 2清华大学, 3NVIDIA

速读

一句话结论 本文针对分布式扩散变换器推理中的通信瓶颈，提出了量化“发出后”通信进展的X-Stage抽象及Burst-Gap模型，并据此设计了交错调度的MegaMoE与搭载式FlashAttention-A2A内核，在不增加硬件资源的前提下实现了最高1.62倍的内核加速。

要解决什么问题 现有的大模型细粒度通信与计算重叠技术（如FLUX、Comet等）普遍采用“完成耦合”模型，即假设通信发出后，发送方必须阻塞等待数据在远端完全可见，或者需要专门的资源来维持通信状态。这种抽象无法真实反映现代GPU的底层机制。实际上，远程存储的“发出”和“远端可见”是两个独立事件，发送方在发出写请求后即可继续执行后续计算。然而，这种“发出后”的推进能力是有限的。如果算子产生密集的远程存储突发，且突发之间的计算间隙太短，未完成的通信请求就会在底层路径中不断累积。当累积量耗尽下游的有效容量时，就会产生反压。这种反压会首先延长发送方的发出时间，随后通过本地的暂存资源向上传导，最终迫使核心的Tensor Core计算流水线停顿。现有的调度系统因为看不见这个反压临界点，导致在处理如MegaMoE这类细粒度内核时，经常因通信突发过于集中而引发严重的计算停顿。

怎么做的 核心思路是识别并量化从“请求发出”到“远端可见”之间的这段软件可见的执行区间，即X-Stage。作者通过微基准测试提取了三个平台参数：无反压发出时间$T_{\text{iss}}^0$、有效排空率$\mathcal{R}$和有效未完成容量$Q$。基于此构建了Burst-Gap模型，定义了当发送方注入体积为$\mathcal{V}$的通信突发，并伴随时长为$\mathcal{G}$的计算间隙时，发送方实际感受到的发出时间为：$$T_{\text{iss}}(\mathcal{G}; \mathcal{K}, \mathcal{V}) = \max\left(T_{\text{iss}}^0(\mathcal{K}, \mathcal{V}), \frac{\mathcal{V}}{\mathcal{R}} - \mathcal{G}\right)$$该公式揭示了计算间隙$\mathcal{G}$可以直接用来掩盖通信的排空时间。基于该模型，作者设计了两种内核优化策略。第一种是针对MegaMoE的“重塑注入”策略。原始的专家波次调度会连续执行Linear-2计算，导致通信突发极大且计算间隙极短，必然触发反压。作者设计了交错调度器，在连续的Linear-2通信突发之间，强行插入没有通信依赖的Linear-1计算瓦片。这在不增加总计算量的前提下人为拉长了计算间隙$\mathcal{G}$，让X-Stage有足够时间排空请求，从而消除了Tensor Core的停顿。第二种是针对Ulysses序列并行注意力的“搭载”策略。作者将注意力后的All-to-All通信与FlashAttention在瓦片粒度上融合。因为FlashAttention的Q-loop本身就提供了随序列长度线性增长的巨大计算间隙，完全满足Burst-Gap模型的容量和速率边界。因此，负责输出计算的角色在发出远程存储后，无需交出控制权或分配专用的通信线程束，直接进入下一个Q-loop恢复计算即可，通信请求会自动在X-Stage中被后台排空。

效果如何 实验在配备8张最新架构NVIDIA GPU、通过NVLink全连接的系统上进行，采用8路专家并行设置。对于MegaMoE，对比基线是代表传统完成耦合路线的专家波次调度器（Expert-Wave）。在涵盖不同模型形状、精度和路由方式的84种配置下，交错调度器实现了1.18倍的几何平均加速比和最高1.62倍的加速比；在偏斜路由下收益尤为明显，发送方可见的通信跨度从7.9至9.9微秒大幅收窄至3.6至3.9微秒。对于FlashAttention-A2A，对比基线是串行执行路线（先算注意力再通信）。在8,192序列长度下，基于FA3和FA4的融合内核分别达到了1.43倍和1.42倍的最大加速比。更关键的是，发送方侧的残余开销仅为10至14微秒的极小常数，且完全不随序列长度增长，证明通信被完美掩盖。方法的代价与局限在于：MegaMoE的跨专家交错调度有时会破坏专家权重的L2缓存局部性，导致极少数配置下出现约0.94倍的性能回退；此外，当前的X-Stage模型仅针对单边远程存储进行了表征，对于加载、原子操作或跨节点网络等其他通信模式，其有效性仍需独立验证。

A1 主要贡献

本文研究了在分布式扩散变换器（DiT）推理中，GPU计算吞吐量与节点内互连带宽之间日益扩大的差距，这一差距使得GPU间的数据移动成为分布式模型执行的关键路径。现有的细粒度、设备发起的通信系统虽然能够调度通信的发出时间和接收数据的可用时间，但忽略了通信发出后、在远端可见完成前的“后发出（post-issue）”进展阶段，这使得发送方反压（backpressure）难以预测。

核心问题与研究目标：
现有的细粒度通信与计算重叠技术，如FLUX、Comet等，虽然能在瓦片（tile）、线程协作阵列（CTA）或线程束（warp）级别分解任务，但它们的抽象模型不足以预测发送方在发出远程存储（remote store）后可以继续执行多远，也无法预知持续的数据注入何时会耗尽下游资源并对计算流水线产生反压。本文在分析DeepGEMM的MegaMoE内核时发现，传统的“完成耦合”（completion-coupled）模型无法解释其观察到的高性能，这表明发出远程存储后，发送方在等待其远程可见的同时，能够继续执行有用的工作。

创新点与主要贡献：
本文识别并提出了一个名为X-Stage的软件可见的后发出流水线阶段，即远程存储发出后到远程可见完成之间的间隔。在此阶段，已发出的请求可以与后续的计算并发进行。

  1. 提出软件可见的后发出抽象（X-Stage）：本文识别出X-Stage，即一个已接受的远程存储请求可以与后续计算并发进行的有限时间间隔。这个抽象解释了为什么完全阻塞模型或无限异步模型都无法准确捕捉发送方的可见行为。

  2. 表征与预测（Burst–Gap模型）：通过参数化的微基准测试，本文测量了无反压时的发出时间、有效排空速率和有效未完成容量。基于这些测量，提出了一个轻量级的Burst–Gap模型，该模型能预测稳态周期、发送方可见的发出开销以及由容量限制引起的反压，且无需针对特定应用重新拟合。

  3. 提出X-Stage感知的内核设计：

    • MegaMoE：通过跨专家波次（expert waves）交错执行Linear-1和Linear-2的工作，重塑集中的远程存储突发，将计算插入到密集的远程存储突发之间，从而缓解反压。
    • Ulysses序列并行注意力：将注意力后的All-to-All操作与FlashAttention进行瓦片粒度的融合。输出瓦片的拥有者在发出远程存储后立即恢复计算，利用后续的FlashAttention Q-loop作为排空窗口，避免了专用的通信线程束或流式多处理器（SM）。

  4. 性能与机制层面的验证：本文评估了84种MegaMoE配置和两代FlashAttention。通过内核时间线和逐瓦片（per-tile）的性能剖析，将测得的加速比与模型预测的发出、排空、容量和分阶段效应联系起来，验证了模型的有效性。


图1. 由epilogue发出的远程存储在X-Stage中继续进行，而下一个mainloop执行有用的计算。远程存储突发和其间的mainloop构成了本文中反复出现的“突发-间隙”模式。

A3 背景知识与动机

2.1 GPU通信与细粒度融合

现代GPU通信与细粒度融合：现代多GPU节点允许GPU通过如NVLink【27, NVLink and NVSwitch System Overview. NVIDIA Data Center Documentation】等高速互连结构直接访问对等内存。对等内存映射、统一虚拟地址空间和对称内存让GPU内核能构建远程内存地址并发起设备端的数据移动【25, CUDA C++ Programming Guide: Peer Device Memory Access. NVIDIA Developer Documentation; 28, NVSHMEM: GPU-Initiated Communication Library. NVIDIA Developer Documentation】。例如，NVSHMEM提供了一个分区全局地址空间（PGAS）模型和对称内存接口，用于细粒度的单边操作，无需返回主机或为每次传输启动单独的集合通信内核。

细粒度通信在并行计算中的重要性：在大模型并行中，细粒度通信日益重要。张量并行调用AllGather、Reduce-Scatter或All-Reduce【33, Megatron-LM: Training Multi-Billion Parameter Language Models Using Model Parallelism. 2019. arXiv】；专家并行在专家计算前后进行令牌的分发和结果的合并；Ulysses序列并行则在注意力计算前后使用两次All-to-All操作在序列分区和头分区间转换【17, DeepSpeed Ulysses: System Optimizations for Enabling Training of Extreme Long Sequence Transformer Models. 2023. arXiv】。随着低精度Tensor Core吞吐量的增加，每个计算阶段相对于GPU间数据移动的时间缩短。因此，算子或流级别的并发可能暴露出大量通信开销，这推动了在瓦片（tile）、CTA或线程束（warp）粒度上进行分解。

现有细粒度系统的局限性：近期的系统通过任务分解、内核融合、专门的执行角色和设备端同步来构建细粒度流水线【2, FLUX: Fast Softwarebased Communication Overlap on GPUs Through Kernel Fusion. 2024. arXiv; 20, MegaScale-MoE: Large-Scale Communication-Efficient Training of Mixture-of-Experts Models in Production. 2026. EuroSys; 34, ParallelKittens: Systematic and Practical Simplification of Multi-GPU AI Kernels. 2026. MLSys; 43, Comet: Fine-grained ComputationCommunication Overlapping for Mixture-of-Experts. 2025. MLSys; 44, TileLink: Generating Efficient ComputeCommunication Overlapping Kernels Using Tile-Centric Primitives. 2025. MLSys】。尽管这些设计在任务粒度、资源组织和调度机制上有所不同，但它们主要暴露了软件调度层面：哪些通信和计算瓦片可以并发执行，哪个角色发出传输，以及远程数据何时可供消费。一个瓦片级别的信号可以报告远程就绪状态，但它并未揭示发送方在此事件发生前可以前进多远，或者有限的下游资源何时会反压发出者。X-Stage通过对这种后发出进展进行建模，补充而非替代了瓦片级调度。

理想融合模型的不足：在更高层次上，一个理想的融合延迟通常被估计为 $T _ { \mathrm { i d e a l } } = \mathrm { m a x } \big ( T _ { \mathrm { c o m p u t e } } , T _ { \mathrm { c o m m } } \big )$。这个聚合模型关注总计算量是否能覆盖总通信量，但它将通信视为一个单一的时间段。实际上，远程存储的发出（issue）和远程可见的完成（remote-visible completion）是两个不同的事件。发出意味着发送方注入了一个写请求，而非数据已在目的地可见。该请求在后续计算运行时可能仍在进行中；然而，在持续注入下，未完成的请求可能累积，延长后续发出操作的时间，并反压epilogue和上游计算。仅凭聚合的计算和通信时间无法捕捉这一转变过程。

2.2 MegaMoE执行

MegaMoE概述：MegaMoE是一个代表性的细粒度、通信-计算融合的MoE内核。在一个专家并行的MoE层中，一个路由器为每个令牌选择一个或多个专家。由于专家位于不同的GPU上，Dispatch阶段首先将每个令牌发送到其目标rank。然后每个专家评估一个两层前馈网络：Linear-1将隐藏维度投影到中间维度，应用一个如SwiGLU【32, GLU Variants Improve Transformer. 2020. arXiv】的激活函数，然后Linear-2将其投影回隐藏维度。Combine阶段将专家输出返回到源rank，并使用路由器权重聚合top-k结果。

MegaMoE融合与角色分工：传统实现为Dispatch、分组GEMM、激活、第二次分组GEMM和Combine启动不同的内核。MegaMoE将这些阶段融合成一个持久化的巨核（mega-kernel），并使用线程束专门化（warp specialization）来处理Dispatch、数据移动、Tensor Core计算和epilogue工作【10, DeepGEMM: High-performance Tensor Core Kernels and Mega MoE. 2026. GitHub】。在公开实现中，Linear-2的epilogue读取累加器结果，并使用远程存储将其写入一个对称的Combine缓冲区。尽管Tensor Core mainloop和epilogue由不同角色执行，但有限的片上累加器和暂存资源将它们耦合为生产者和消费者。

MegaMoE的Expert Waves：MegaMoE将本地专家分组为专家波次（expert waves）来组织局部性和执行。原始调度通常在一个波次内先执行其Linear-1工作，然后是激活和Linear-2工作，之后再进入下一个波次。不同波次的Linear-1工作没有直接的神经网络依赖关系。原则上，一旦Dispatch数据和目标缓冲区就绪，后一个波次的Linear-1可以在前一个波次的远程Combine存储完成之前开始。

2.3 任务级建模的不匹配

任务级时间线与完成耦合解释：MegaMoE公开发布的时间线将Dispatch、Linear-1、激活、Linear-2和Combine表示为任务级阶段【11, DeepSeek-V4: Towards Highly Efficient MillionToken Context Intelligence. 2026. arXiv】；图2(a)展示了这一视图。一个自然的保守解释是，本地epilogue和暂存资源与一个Linear-2结果保持耦合，直到相应的Combine操作在远程可见。我们称之为完成耦合（completion-coupled）解释。

完成耦合模型的数学表达：令 $T _ { \mathrm { L i n 1 } }$、 $T _ { \mathrm { A c t } }$ 和 $T _ { \mathrm { L i n 2 } }$ 表示一个专家波次内Linear-1、激活和Linear-2的稳态阶段时间。令 $T _ { \mathrm { C o m b i n e } } ^ { \mathrm { c c } }$ 涵盖从本地epilogue到远程可见的整个过程。忽略流水线的填充和排空，每个波次的稳态时间近似为：$T _ { \mathrm { w a v e } } ^ { \mathrm { c c } } \approx \mathrm { m a x } ( T _ { \mathrm { L i n 1 } } , T _ { \mathrm { A c t } } ) + \mathrm { m a x } \big ( T _ { \mathrm { L i n 2 } } , T _ { \mathrm { C o m b i n e } } ^ { \mathrm { c c } } \big )$。

Combine时间的分解与模型不匹配：我们将完成耦合的Combine时间间隔分解为 $T _ { \mathrm { C o m b i n e } } ^ { \mathrm { c c } } = T _ { \mathrm { l o c a l \ e p i l o g u e } } + T _ { \mathrm { R S i s s u e } } + T _ { \mathrm { p o s t - i s s u e \ c o m p l e t i o n } }$，其中最后一项是从接受发出到远程可见的剩余时间。该模型允许通过有限的片上暂存资源在Linear-2生产者和结果发送角色之间实现临时重叠，但将远程完成视为稳态进展的约束。将此模型应用于报告的阶段时间，得出的估计加速比上限约为1.5倍，而实现达到了1.56倍【10, DeepGEMM: High-performance Tensor Core Kernels and Mega MoE. 2026. GitHub; 11, DeepSeek-V4: Towards Highly Efficient MillionToken Context Intelligence. 2026. arXiv】。虽然这个小差异本身不是机制证明，但完成耦合的估计无法解释观察到的加速，这促使我们直接测量后发出的进展。

2.4 分解Combine：从发出到远程可见

区分Issue和Remote Visibility：图2(a)将Combine视为一个单一的通信阶段。这种表示捕获了应用依赖关系和任务级时间线，但混淆了两个发送方事件。一个Linear-2 epilogue首先读取累加器，执行数据转换和地址计算，并将结果写入对称的Combine缓冲区。远程存储的发出（issue）部分在发送方接受这些写指令时结束；这并不意味着相应的数据在远程已经可见。

X-Stage的引入：图2(b)将这些事件分开。一旦远程存储被接受，发出的角色可以继续进行后续计算，而请求则继续向目的地推进。这种解耦解释了为什么完成耦合的估计会低估MegaMoE的加速。然而，仅凭应用延迟无法揭示进展速率、发出者可以领先完成多远，或者持续注入如何改变发送方可见的执行情况。因此，我们使用受控的远程存储微基准测试来直接测量有效排空、未完成容量和反压。


图2. MegaMoE专家波次时间线的任务级视图和X-Stage感知视图。任务级视图将Combine视为单一通信阶段，而X-Stage感知视图则将发送方可见的远程存储发出与后续向远程可见完成的进展分离开来。

2.5 研究问题

研究问题：以上分析引出了三个研究问题：
* RQ1：后发出进展和资源限制。 设备发起的远程存储在发出后如何进展，发出的角色可以领先远程可见完成多远，以及哪些有效资源限制对软件是可见的？
* RQ2：表征和预测。 一个基于无反压发出时间、有效排空速率和有效未完成容量的轻量级模型，能否在不同的突发量、生产者配置和计算间隙下预测执行周期、恢复和反压？
* RQ3：X-Stage感知的内核设计。 这些测量如何指导融合内核的设计？具体来说，调度器能否重塑MegaMoE的突发以减少请求累积？FlashAttention-A2A流水线能否在不预留长期通信资源的情况下覆盖后发出进展？

A2 方法细节

3. 表征与建模X-Stage

引言：第2.3节中应用层面的模型不匹配表明，远程存储的发出和远程可见的完成并非完全耦合。本节使用受控的微基准测试来表征这种后发出行为，并开发一个Burst-Gap模型，用于描述恢复、稳态突发周期和容量引起的背压。

隔离关键量的实验设计：实验旨在隔离三个量：在没有反压的情况下发送方注入一个突发的速度，已接受请求的排空速度，以及在发出操作停滞前可以有效保持未完成状态的数据量。本节中，“生产者”指发出远程存储的并发角色，除非另有说明，否则它与Tensor Core计算生产者不同。表1总结了相关符号。


图3. 远程存储微基准测试对X-Stage的表征。(a) 生产者侧间隙对发送方可见发出时间的影响；(b) 测量的突发周期与Burst-Gap模型的比较；(c) 孤立突发的有效容量拐点；(d) 不同生产者配置下的共享有效排空率。

表1. X-Stage Burst-Gap模型的符号定义。

符号 定义
$\mathcal{K}$ 并发远程存储生产者的数量。
$\mathcal{B}$ 每个生产者每次突发的字节数。
$\mathcal{V}$ 聚合突发量, $\mathcal{V} = \mathcal{KB}$; $\mathcal{V}_t$ 表示每瓦片的量。
$b_{m, n}$ 计算瓦片的M/N维度。
$\mathcal{G}$ 突发之间无远程存储的有用生产者侧时间。
$T_{\text{iss}}(\mathcal{G}; \mathcal{K}, \mathcal{V})$ 发送方可见的突发发出时间，包括反压。
$T_{\text{iss}}^0$ 无反压的突发发出时间。
$\Delta T_{\text{iss}}$ 反压开销, $\Delta T_{\text{iss}} = T_{\text{iss}} - T_{\text{iss}}^0$。
$T_{\text{period}}$ 突发开始到开始的周期, $T_{\text{period}} = T_{\text{iss}} + \mathcal{G}$。
$\mathcal{R}$ 有效聚合排空率（测量值为717 GB/s）。
$Q(\mathcal{K})$ $\mathcal{K}$个生产者的有效未完成容量。
$\mathcal{G}^*(\mathcal{K}, \mathcal{V})$ 最小恢复间隙, $\mathcal{G}^* = [\mathcal{V}/\mathcal{R} - T_{\text{iss}}^0]_+$。
$q_{\text{peak}}(\mathcal{K}, \mathcal{V})$ 无容量限制时的峰值未完成量, $[\mathcal{V} - \mathcal{R}T_{\text{iss}}^0]_+$。
$\mathcal{M}_X$ 平台参数 ($T_{\text{iss}}^0$, $\mathcal{R}$, $Q$)。

3.1 远程存储微基准测试

微基准测试设置：我们的参数化基准测试从 $\mathcal{K}$ 个并发生产者生成一个突发。每个生产者写入 $\mathcal{B}$ 字节，总聚合量为 $\mathcal{V} = \mathcal{KB}$。我们使用两种模式。

两种测试模式：
1. 周期性突发（Periodic bursts）：生产者重复发出体积为 $\mathcal{V}$ 的突发，模拟GEMM epilogue中的远程存储，然后执行一个时长为 $\mathcal{G}$ 的生产者侧间隙，模拟随后的MMA mainloop。我们记录 $T_{\text{iss}}$ 和开始到开始的周期 $T_{\text{period}}$。零间隙会驱动路径进入排空限制的稳态，并暴露有效排空率。扫描 $\mathcal{G}$ 可以揭示一个突发留下的请求如何影响下一个突发。
2. 孤立突发（Isolated bursts）：每次测量在充分恢复后开始，发出一个体积为 $\mathcal{V}$ 的突发，并记录其发送方可见的发出时间。这种模式检测单个突发是否会耗尽有效未完成容量。

本地内存控制实验：一个本地内存控制实验使用相同数量的生产者、存储宽度、地址生成和循环结构。它将普通指令开销与由远程路径上未完成请求引起的停顿分离开来。

3.2 X-Stage Burst–Gap模型

模型的三个核心量：周期性突发形成一个生产者-消费者过程：软件注入请求，下游路径将其排空。我们用三个量来描述其软件可见的行为：
* 对于一个已恢复的、参数为 $(\mathcal{K}, \mathcal{V})$ 的突发，其无反压的发出时间 $T_{\text{iss}}^0$；
* 有效聚合排空率 $\mathcal{R}$；
* 针对该生产者配置的有效未完成容量 $Q(\mathcal{K})$。

流体模型的假设与稳态分析：该模型假设在测量的范围内，存在一个工作守恒的有效排空，速率为 $\mathcal{R}$。这是一个流体、稳态的抽象，而非关于某个特定物理队列的论断。如果间隙足够长以排空前一个突发，其未完成的请求不会显著延迟下一次发出，此时 $T_{\text{iss}} = T_{\text{iss}}^0$。如果间隙较短，请求会跨周期传递。在排空主导的稳态下，流守恒要求每个周期排空该周期内注入的量：$T_{\text{period}}(\mathcal{G}; \mathcal{K}, \mathcal{V}) = \frac{\mathcal{V}}{\mathcal{R}}$。

周期定义与关键公式推导：根据定义，一个周期包括发出时间和随后的生产者侧间隙：$T_{\text{period}}(\mathcal{G}; \mathcal{K}, \mathcal{V}) = T_{\text{iss}}(\mathcal{G}; \mathcal{K}, \mathcal{V}) + \mathcal{G}$。因此，在排空主导的区域，我们得到 $T_{\text{iss}}(\mathcal{G}; \mathcal{K}, \mathcal{V}) = \frac{\mathcal{V}}{\mathcal{R}} - \mathcal{G}$。

完整的Burst-Gap关系式：发出时间不可能比无反压的基线更快，这给出了完整的Burst-Gap关系：$T_{\text{iss}}(\mathcal{G}; \mathcal{K}, \mathcal{V}) = \max\left(T_{\text{iss}}^0(\mathcal{K}, \mathcal{V}), \frac{\mathcal{V}}{\mathcal{R}} - \mathcal{G}\right)$。

开销与恢复时间的计算：由此产生的发出时间开销为 $\Delta T_{\text{iss}}(\mathcal{G}; \mathcal{K}, \mathcal{V}) = \left[\frac{\mathcal{V}}{\mathcal{R}} - \mathcal{G} - T_{\text{iss}}^0(\mathcal{K}, \mathcal{V})\right]_+$, 其中 $[x]_+ = \max(x, 0)$。调度器可以通过减少突发量或增加突发间的有用工作来减少此开销。当 $\mathcal{G} \geq \mathcal{G}^*(\mathcal{K}, \mathcal{V}) = \left[\frac{\mathcal{V}}{\mathcal{R}} - T_{\text{iss}}^0(\mathcal{K}, \mathcal{V})\right]_+$ 时，系统达到恢复的平台期。

3.3 标定有效排空率

标定有效排出率R：我们首先独立地标定 $\mathcal{R}$。通过使用零间隙的周期性突发，生产者连续注入请求并达到一个排空主导的稳态，此时 $T_{\text{period}}^{\text{ss}}(\mathcal{V}) \approx \frac{\mathcal{V}}{\mathcal{R}}$。

实验结果与R的标定值：图3(d)报告了不同生产者数量 $\mathcal{K}$ 和聚合体积 $\mathcal{V}$ 下的稳态周期。相同体积的不同 $\mathcal{K}/\mathcal{B}$ 分解产生了相似的周期，并且较大体积的数据点收敛到一条线性趋势线上。因此，是下游的排空，而不是单个生产者的本地发出吞吐量，控制着这个区域。将排空主导的数据点拟合到 $T_{\text{period}}^{\text{ss}} = \mathcal{V}/\mathcal{R}$，在被测系统上得到 $\mathcal{R} \approx 717 \text{ GB/s}$。所有后续的间隙和容量实验都使用此值，不进行逐曲线的重新拟合。

3.4 验证Burst–Gap行为

模型验证方法：在从图3(d)固定 $\mathcal{R}$ 后，我们扫描 $\mathcal{G}$ 并从大间隙的平台期获得 $T_{\text{iss}}^0$。然后，公式1和3可以在没有额外拟合参数的情况下预测整个扫描过程。

发送方可见的Issue时间验证：图3(a)展示了发送方可见的视图。在短间隙时，前一个突发留下的请求延长了下一次的发出时间。增加间隙大约以 $\mathcal{V}/\mathcal{R} - \mathcal{G}$ 的速率减少了这种额外的停顿，直到发出时间达到 $T_{\text{iss}}^0$。本地内存控制实验几乎对间隙不敏感，表明恢复来自远程路径上的后发出进展，而不是来自间隙指令或存储循环本身。

稳态周期的验证：图3(b)展示了相应的稳态周期。在短间隙时，测量值停留在排空限制的平台 $\mathcal{V}/\mathcal{R}$ 上。在预测的 $\mathcal{G}^*$ 附近，执行切换到 $T_{\text{iss}}^0 + \mathcal{G}$ 分支，并随间隙线性增加。测量结果遵循 $T_{\text{period}} = \max(T_{\text{iss}}^0 + \mathcal{G}, \mathcal{V}/\mathcal{R})$，而不是完成耦合的和 $T_{\text{period}}^{\text{cc}} = \mathcal{G} + \mathcal{V}/\mathcal{R}$。因此，生产者侧的间隙与排空重叠；它不仅仅是加在通信时间上。

3.5 测量有效未完成容量

有效未完成容量的测量方法：周期性突发揭示了跨周期的稳态。孤立突发则探究一个足够大的突发，即使在X-Stage初始为空的情况下，是否会自我反压。在一个无反压的发出时间间隔内，生产者注入 $\mathcal{V}$ 字节，同时大约 $\mathcal{R}T_{\text{iss}}^0(\mathcal{K}, \mathcal{V})$ 字节被排空。因此，在发出结束时，无约束的未完成体积为 $q_{\text{peak}}^0(\mathcal{K}, \mathcal{V}) = [\mathcal{V} - \mathcal{R}T_{\text{iss}}^0(\mathcal{K}, \mathcal{V})]_+$.

容量限制下的Issue时间：如果 $q_{\text{peak}}^0 \le Q(\mathcal{K})$，孤立的发出时间保持在 $T_{\text{iss}}^0$。否则，超出容量的注入必须等待下游排空，从而得到 $T_{\text{iss}}^{\text{iso}}(\mathcal{K}, \mathcal{V}) = \max\left(T_{\text{iss}}^0(\mathcal{K}, \mathcal{V}), \frac{\mathcal{V} - Q(\mathcal{K})}{\mathcal{R}}\right)$。

实验结果与Q的标定值：图3(c)显示了孤立发出时间与聚合体积的关系。小突发保持在无反压的基线上。超过一个依赖于生产者配置的拐点后，发出时间向排空限制的分支上升，表明领先下游完成的距离是有限的。对于 $\mathcal{K}=148$，孤立发出时间的拐点出现在每个生产者约33 KB处，对应聚合突发体积约4.77 MiB。在 $\mathcal{B}=32$ KB时，测得的无反压发出时间下限为 $T_{\text{iss}}^0 \approx 0.76 \mu s$。在拐点处应用公式6并考虑发出期间排空的数据，得到 $Q \approx 4.25$ MiB。

3.6 模型总结

平台模型总结：这三个可独立测量的量构成了平台模型 $\mathcal{M}_X = \big(T_{\text{iss}}^0(\mathcal{K}, \mathcal{V}), \mathcal{R}, Q(\mathcal{K})\big)$。$T_{\text{iss}}^0$ 捕获了特定生产者配置和突发体积下的无反压注入；$\mathcal{R}$ 捕获了聚合的后发出排空；$Q$ 捕获了相对于下游完成的有限领先距离。我们分别从恢复的发出时间平台期、零间隙稳态周期和孤立突发拐点中提取它们。

模型解释的三个观察：它们共同解释了三个观察结果：
1. 发送方可见的发出可以在远程可见的完成之前结束，允许已接受的请求与后续工作并发进行；
2. 生产者侧的间隙与后发出排空重叠，因此稳态周期遵循最大值法则，而不是完成耦合的和；
3. 当未完成体积超过有效容量时，反压首先延长远程存储发出时间，然后可能通过本地暂存传播到计算生产者。

下一节将在两种不同场景下应用此模型。MegaMoE通过跨专家波次交错工作来扩大集中的Linear-2突发之间的间隙。FlashAttention-A2A则利用紧随其后的Q-loop来覆盖前一个输出瓦片的后发出进展。

4. X-Stage感知的内核设计

引言：一旦发送方接受了一个远程存储，该请求就可以通过X-Stage继续向远程完成推进。这种后发出进展改变了融合内核的调度边界：软件不仅必须控制通信何时发出，还必须控制已接受的请求累积的速度以及在它们排空期间有多少可用工作。

两层解耦：传统的融合GEMM通常被描述为Load-Compute-Epilogue流水线。当epilogue包含远程存储时，在发出点结束描述会隐藏一个额外的阶段，如图1所示。这里存在两个不同层次的解耦。本地暂存提供了计算-发出解耦：一个Tensor Core生产者将结果存入累加器或暂存缓冲区，一个epilogue角色发出存储。X-Stage提供了发出-完成解耦：在发出被接受后，请求可以继续推进，而无需持续占用发出的线程束。

调度目标：对于一个聚合突发体积为 $\mathcal{V}$、后跟一个生产者侧间隙 $\mathcal{G}$ 的情况，Burst-Gap模型给出 $T_{\text{iss}}(\mathcal{G}; \mathcal{K}, \mathcal{V}) = \max\left(T_{\text{iss}}^0(\mathcal{K}, \mathcal{V}), \frac{\mathcal{V}}{\mathcal{R}} - \mathcal{G}\right)$ 和 $\Delta T_{\text{iss}}(\mathcal{G}; \mathcal{K}, \mathcal{V}) = \left[\frac{\mathcal{V}}{\mathcal{R}} - \mathcal{G} - T_{\text{iss}}^0(\mathcal{K}, \mathcal{V})\right]_+$。直接的调度目标是在不增加工作量或违反依赖关系的情况下减少 $\Delta T_{\text{iss}}$。

4.1 设计测试与两种动作

设计测试准则：该模型为周期性突发提供了一个紧凑的测试。在第3.2节的流体模型和稳态假设下，当以下两个条件同时满足时，X-Stage不会增加发送方可见的反压：

$$ \underbrace{\frac{\mathcal{V}}{T_{\text{iss}}^0 + \mathcal{G}} \le \mathcal{R}}_{\text{长期注入率界限}}, \qquad \underbrace{\left[\mathcal{V} - \mathcal{R}T_{\text{iss}}^0\right]_+ \le Q}_{\text{单次突发容量界限}} $$


第一个界限防止跨周期的累积；第二个防止单个孤立突发耗尽有效容量。满足这些界限可以消除模型预测的反压开销，但这本身并不证明所有发出指令、同步或完成要求都已从应用关键路径中移除。

X-Stage感知设计的步骤：一个X-Stage感知的设计遵循三个步骤。首先，使用第3节的微基准测试校准 $\mathcal{M}_X = (T_{\text{iss}}^0, \mathcal{R}, Q)$。其次，审查目标内核以获取其聚合突发体积 $\mathcal{V}$ 和自然计算间隙 $\mathcal{G}$。第三，选择以下两种动作之一：

两种调度动作：
* 违反准则：重塑注入。 如果自然间隙太短或突发太大，则在突发之间重新分配有用的工作或减少突发聚合。MegaMoE遵循此路径（第4.2节）：连续的Linear-2 epilogue注入速度超过了路径的排空速度，因此将来自后续波次的就绪Linear-1工作移到它们之间。
* 满足准则：搭载于现有计算。 如果现有的计算流水线已经提供了足够的间隙，拥有输出的角色可以发出一个短突发并返回到有用的工作中，同时由X-Stage进行排空。一个专门的通信角色在发出后等待并不能提高测得的下游排空率，尽管它在其他设计中可能对地址生成、发出吞吐量或同步有用。FlashAttention-A2A遵循此路径（第4.3节）：一个完整的Q-loop分隔了输出突发，且间隙随序列长度增长。

两种动作都使用相同的校准参数。MegaMoE改变独立的Linear-1和Linear-2瓦片的顺序以平滑注入并控制X-Stage的占用。FlashAttention-A2A让现有的拥有输出的角色发出每个突发并恢复计算，避免了仅用于等待后发出进展的长期通信线程束或SM。

4.2 MegaMoE：跨专家波次的容量感知调度

MegaMoE的问题：细粒度的专家缩短了Linear-2的MMA mainloop，这是连续epilogue突发之间的自然间隙。因此，原始的专家波次调度经常在排空主导的区域运行，并违反了公式10中的速率界限。

用有用功交换Issue Stall：图4(a)显示了原始调度。令 $\mathcal{V}_t$ 为一个Linear-2瓦片在一个生产者上发出的Combine体积，$\mathcal{K}_{\text{act}}$ 为瓦片边界对齐在一个突发中的生产者数量。X-Stage模型使用的聚合体积是 $\mathcal{V} = \mathcal{K}_{\text{act}}\mathcal{V}_t$。令 $T_{\text{L1}}^{\text{mma}}$ 和 $T_{\text{L2}}^{\text{mma}}$ 表示Linear-1和Linear-2的mainloop时间。连续的Linear-2瓦片仅提供 $\mathcal{G}_{\text{wave}} \approx T_{\text{L2}}^{\text{mma}}$ 的间隙。对于细粒度专家，$\frac{\mathcal{V}}{T_{\text{iss}}^0 + T_{\text{L2}}^{\text{mma}}} > \mathcal{R}$，因此未完成的请求会上升到有效容量。在稳态下它们不能再增长；相反，反压会延长发出时间并通过本地暂存传播到Tensor Core生产者。在排空主导的区域，公式1和2给出 $T_{\text{period}}^{\text{cc}} = \mathcal{G} + \mathcal{V}/\mathcal{R}$，因此在突发之间移动计算只会延迟下一个突发，而不是减少停顿。对于固定的通信量和排空率，调度无法缩短这个周期。然而，它可以决定周期中有多少是有效计算，多少是发送方可见的发出停顿。在专家波次调度下， $T_{\text{iss}} = \frac{\mathcal{V}}{\mathcal{R}} - T_{\text{L2}}^{\text{mma}} > T_{\text{iss}}^0$。将有用间隙增加一个单位，会使发出停顿减少一个单位，直到 $T_{\text{iss}}$ 达到 $T_{\text{iss}}^0$。这种交换依赖于后发出进展。在完成耦合模型下，移动计算只会延迟下一个突发。

交错调度器算法：算法1合并了就绪的Linear-1和Linear-2瓦片流。它贪婪地发出Linear-2，仅当Linear-1流领先至少 $D$ 个调度器行时；否则，它发出一个就绪的Linear-1瓦片，这增加了有用的工作而不产生Combine存储。在稳态下，该模式每个周期包含 $n_2$ 个Linear-2瓦片和 $n_1$ 个Linear-1瓦片。它们的比例由瓦片形状固定。领先距离 $D$ 改变了流的相位和就绪状态，而不是它们的长期比例，这区分了两个条件。

Algorithm 1: Interleaved Scheduler
g: instance id, stride G; r ∈ {0, 1}: role in pair; P: rows/stream; Pi:
pairs/row; ki: pair cursor; pi = [ki/Pi]: row; e(p): expert of row p;
Oe: its first row; D: minimum row lead of L1 over L2.
1 procedure ScHEDULER(g, r)
2   k1 ← g; k2 ← g
3   while k1 < P· P1 or k2 < P · P2 do
4     p1 ← |k1 /P1] (+∞ if L1 done)
5     p2 ← |k2/P2]
6     if k2 < P · P2 and p2 + D ≤ p1 then
7       emit L2ce(p2), m=p2−e(p2), n=2(k2 m0d P2)+r(ci:
8       k2 ← k2 + G
9     else
10      emit L1ce(p1), m=p1−0e(p1), n=2(k1 m0d P1)+r(ci:
11    end

条件1：交错是否足够？ 在一个稳态周期内平均，自然间隙为 $\mathcal{G}_{\text{wave}} \approx T_{\text{L2}}^{\text{mma}}$ 和 $\mathcal{G}_{\text{int}} \approx T_{\text{L2}}^{\text{mma}} + \frac{n_1}{n_2}T_{\text{L1}}^{\text{mma}}$。交错既不增加计算也不增加通信，它重新分配了本应在波次边界分组的Linear-1工作。将公式12代入公式9得到 $\Delta T_{\text{iss}}^{\text{int}} = \left[\Delta T_{\text{iss}}^{\text{wave}} - \frac{n_1}{n_2}T_{\text{L1}}^{\text{mma}}\right]_+$。因此，有用的计算一对一地抵消了发出开销，直到开销达到零。完全恢复需要 $G_{\text{int}} \ge G^*$, 等价于 $\frac{n_2\mathcal{V}}{n_2(T_{\text{iss}}^0 + T_{\text{L2}}^{\text{mma}}) + n_1T_{\text{L1}}^{\text{mma}}} \le \mathcal{R}$。这个测试依赖于聚合突发体积、瓦片混合和mainloop时间，以及校准的 $(T_{\text{iss}}^0, \mathcal{R})$，但与 $D$ 无关。因此，它可以在实现前预测重新排序现有计算是否足够。如果公式14不成立，调度器还必须减少聚合突发体积。对于评估的配置，该公式预测交错应将每个瓦片的发送方可见远程存储范围从排空主导区域移动到无反压的下限。我们将在第5.2.3节中测试这一预测。

从Epilogue Stall到计算关键路径：Burst-Gap模型从可用的后发出排空间隙预测epilogue的发送方可见跨度：$T_{\text{RS}} \approx T_{\text{iss}}^0 + \Delta T_{\text{iss}}$ 加上固定的暂存开销。Epilogue停顿不一定是计算停顿。双缓冲暂存允许MMA线程束组在必须重用被占用的槽位之前，推进到下一个mainloop。令 $T_{\text{cover}}^{\text{mma}}$ 表示本地暂存提供的这个特定于配置的覆盖窗口。在波次调度中，它是随后的Linear-2 mainloop；在交错调度中，当一个插入的Linear-1紧随其后时，它是Linear-1 mainloop。混合周期对每个采样的瓦片使用其实际的后继者。一个瓦片的关键路径停顿是 $\Delta t = [T_{\text{RS}} - T_{\text{cover}}^{\text{mma}}]_+$。因此，交错可以在两个独立的层面上提供帮助。它扩大了排空间隙并减少了 $T_{\text{RS}}$；它还可以用更长的Linear-1 mainloop替换直接的mainloop覆盖窗口。反之，即使公式14不能将 $\Delta T_{\text{iss}}^{\text{int}}$ 降至零，只要 $T_{\text{RS}} \le T_{\text{cover}}^{\text{mma}}$，剩余的epilogue停顿仍不在Tensor Core关键路径上。第5.2.3节使用暂存边界两侧的插桩来测试公式15。

条件2：防止饥饿的lead D：只有当Linear-2不因等待其Linear-1依赖而空转时，稳态间隙才能实现。一个Linear-2瓦片消耗跨多个SM生产并通过 l2_full 计数器同步的Linear-1输出。如果 $D$ 太小，Linear-2流会追上，因未就绪的依赖而空转，并让突发重新集中。实现中，N-tiles被映射到成对的CTA中，并以步长方式在 $\mathcal{K}_{\text{SM}}$ 个参与的SM上推进流。由此产生的跨SM相位偏差给出了近似的抗饥饿拐点 $D_{\text{knee}} = \lceil\frac{b_n\mathcal{K}_{\text{SM}}}{2h_{\text{inter}}}\rceil$，其中 $b_n$ 是N-tile的宽度， $h_{\text{inter}}$ 是专家的中间维度。这里的 $D$ 是调度器行中的领先距离；行布局决定了其对应的令牌范围。这个下限来自跨SM数据就绪性，而不是X-Stage排空，因此不包含 $\mathcal{R}$、$T_{\text{iss}}^0$ 或 $Q$。选择 $D < D_{\text{knee}}$ 会让Linear-2追上并空转，允许远程存储突发重新集中；将 $D$ 增加到拐点附近可以消除这个抗饥饿瓶颈。将 $D$ 远超拐点增加并不会改变长期的Linear-1/Linear-2比例或X-Stage速率条件，但会增加仅Linear-2的尾部和并发专家权重工作集，可能损害缓存局部性。因此 $D_{\text{knee}}$ 对模型来说是足够的；任何额外的余量都是实现特定的调优。算法1保留了所有专家内部的依赖关系，只重新排序独立的瓦片。一个未来波次的Linear-1瓦片只有在其Dispatch输入和目标环形槽就绪后才有资格执行。有限的软件流水线深度分别限制了在途的epilogue，并且必须与校准的 $Q$ 兼容。


图4. 采用独立epilogue角色的X-Stage感知MegaMoE调度。专家波次调度器连续放置Linear-2瓦片，导致它们的epilogue发出集中的远程存储突发。交错调度器在这些突发之间插入就绪的Linear-1工作，允许早期的请求在有用的计算期间排空。

4.3 FlashAttention–A2A：无需专用通信角色的搭载

Ulysses序列并行与融合目标：我们接下来考虑Ulysses序列并行，如图5所示，这是扩散变换器（DiT）工作负载的常用策略。评估的DiT设置使用完整的、非因果的注意力，而不是典型的自回归LLM使用的因果注意力【30, Scalable Diffusion Models with Transformers. 2023. ICCV】。Ulysses在注意力之前执行一次All-to-All，将序列分区转换为头分区，之后再执行一次All-to-All（A2A）以恢复原始布局【17, DeepSpeed Ulysses: System Optimizations for Enabling Training of Extreme Long Sequence Transformer Models. 2023. arXiv】。我们将FlashAttention与注意力后的All-to-All在瓦片粒度上融合：当一个输出瓦片就绪时，拥有该瓦片的角色在边界处发出该瓦片的远程存储，然后继续处理下一个瓦片。图6展示了分隔连续输出瓦片边界的Q-loop。


图5. Ulysses序列并行注意力。QKV投影和注意力前的All-to-All将序列分区的令牌转换为头分区(a-c)。FlashAttention和注意力后的All-to-All产生输出并恢复原始序列分区(d-e)；我们的FA+A2A设计融合了后两个操作。


图6. FlashAttention Q-loop和输出分块。每个CTA拥有一块查询行，并在内循环中遍历K/V瓦片。完成遍历会产生一个输出瓦片；连续的输出瓦片构成外循环。

设计审计：我们首先使用公式10审计该设计。一个FlashAttention输出瓦片的形状为 $b_m \times d$，其中 $d$ 是头维度。每个持久化CTA在瓦片边界写入 $\mathcal{B}_t = b_m d s_}}$ 字节，其中 $s_{\text{out}}$ 是输出元素大小。在保守情况下，$\mathcal{K}$ 个生产者对齐它们的边界，聚合突发为 $\mathcal{V} = \mathcal{K}\mathcal{B_t$；实际的持久化CTA通常是相位偏移的，注入更平滑。与MegaMoE不同，连续的输出突发被下一个输出瓦片的完整Q-loop分隔开，该Q-loop会访问序列长度为 $M$ 的所有 $\lceil M/b_n \rceil$ 个KV瓦片【31, FlashAttention-3: Fast and Accurate Attention with Asynchrony and Low-precision. 2024. NeurIPS】。由此产生的间隙 $\mathcal{G}_(M)$ 大致随 $M$ 线性增长。}

容量边界检查：对于评估的FlashAttention-4配置，$b_m=128, d=128$，输出为bf16，且 $\mathcal{K} \le 148$。每个生产者写入32KB，低于第3.5节中测量的每个生产者约33KB的有效容量拐点。在最差的对齐情况下，$\mathcal{V} \approx 4.6$ MiB，并且 $[\mathcal{V} - \mathcal{R}T_{\text{iss}}^0]_+ \approx 4.1 \text{ MiB} < Q \approx 4.25 \text{ MiB}$。因此，容量界限成立。当 $G_(M) \ge G^}_$ 时，速率界限成立。第4.3节推导出一个保守的序列长度上界 $M_{\text{ub}}^_$；所有 $M \ge M_{\text{ub}}^*$ 都满足搭载设计测试的两个界限。

后发进展扩展了排空窗口：一个完成耦合的反事实论证明确了X-Stage的作用。如果发出的角色必须保持阻塞直到远程可见，那么每个对齐的输出突发的暴露时间将近似为 $\Delta t^{\text{cc}} = [\frac{\mathcal{V}}{\mathcal{R}} + T_{\text{lat}} - T_{\text{cover}}]_+$, 其中 $T_ \propto M$。}} \lesssim T_{\text{QK}}^{\text{mma}}$。其中 $T_{\text{lat}}$ 是一个固定的完成延迟分量，而 $T_{\text{cover}}$ 是在发出者被阻塞时，由有限的本地暂存提供的可用进展。在FlashAttention-4中，一个线程束组驱动CTA的MMA流水线，因此这个本地覆盖窗口最多为一个 $QK^T$ 瓦片mainloop的量级。其持续时间和计算量 $2b_m b_n d$ 不随 $M$ 增长。如果 $\Delta t^{\text{cc}} > 0$，那么在 $N_t \propto M$ 个输出突发波次上的总完成耦合暴露时间将增长为 $E^{\text{cc}} = N_t \Delta t^{\text{cc}

X-Stage下的反压分析：有了X-Stage，发出的角色在发送方接受存储后返回。下一个完整的Q-loop，而不是一个本地MMA步骤，成为后发出的排空窗口。如果该窗口不足，公式4给出发出侧的反压 $\Delta t(M) = [\frac{\mathcal{V}}{\mathcal{R}} - \mathcal{G}_{\text{Qloop}}(M) - T_^0]}_+$。与MegaMoE一样，这个epilogue侧的停顿只有在超过额外的双缓冲覆盖窗口 $\mathcal{G}_}}$ 后才会到达计算关键路径：$\Delta t'(M) = [\Delta t(M) - \mathcal{G_{\text{QK}}]_+$。因为 $\mathcal{G}_{\text{Qloop}}(M)$ 随 $M$ 增长而对齐的突发体积是固定的，所以 $\Delta t$ 会减小并最终达到零。更弱的关键路径条件 $\Delta t \le \mathcal{G}_$ 可能更早达到。公式17和18之间的对比源于一个特性：已接受的存储是否可以离开发出的角色并通过X-Stage进行。}


图7. 用于FlashAttention-4的搭载式FlashAttention-A2A流水线。每个CTA的单个线程束组发出 $QK^T$ 和 $PV$ 操作，而专门的角色执行softmax等操作。在输出瓦片边界，epilogue发出远程存储(RS)并立即开始下一个瓦片。前一个瓦片在下一个完整的Q-loop期间通过X-Stage排空，无需专用的通信线程束或SM。

角色分配的指导：这种区别也指导了角色分配。预留通信线程束或SM可以解耦MMA角色与地址生成和发出【20, MegaScale-MoE: Large-Scale Communication-Efficient Training of Mixture-of-Experts Models in Production. 2026. EuroSys】，但会消耗寄存器、共享内存、占用率或Tensor Core容量。在审查的区域内，发出后等待并不能改善校准的下游排空率。因此，我们将短的发出操作附加到已经拥有输出瓦片的角色上。如图7所示，瓦片 $i$ 在瓦片 $i+1$ 执行其完整Q-loop时排空。通信瓦片保留了计算瓦片的形状，因此该设计对FlashAttention数据路径的改变最小。与MegaMoE的重塑动作不同，搭载不重新排序计算；它移除了一个长期的通信角色，并利用工作负载的自然排空窗口。

保守的收敛边界：对于形状为 $b_m \times b_n$ 且归约维度为 $k$ 的GEMM瓦片，计算工作量为 $C_{\text{GEMM}} = 2b_m b_n k$。一个全注意力的输出瓦片对所有KV瓦片执行 $QK^T$ 和 $PV$，因此 $C_{\text{FA}} = 4b_m b_n d \frac{M}{b_n} \approx 4b_m d M$。将此工作量与GEMM表达式等同，得到等效归约维度 $k_{\text{equiv}} = 2d \lceil\frac{M}{b_n}\rceil \propto M$。因子2解释了注意力中的两次矩阵乘法。测量的Q-loop间隙有一个正的固定分量和一个近似线性的稳态斜率。为了得到一个保守的阈值，我们去掉正的固定分量，并使用下界 $\mathcal{G}_{\text{Qloop}}(M) \ge \sigma \lceil\frac{M}{b_n}\rceil$，其中 $\sigma \approx 0.79 \mu s$ 每个KV瓦片。模型预测的所有突发波次上的排空引起的残余为 $E_}}(M) = N_t [\Delta t(M) - \mathcal{G_{\text{QK}}]_+$。这个量不包括执行无反压发出指令的有限成本，该成本包含在第5.3节测量的残余 $E_{\text{res}} = T_{\text{fused}} - T_{\text{FA}}$ 中。去掉正的固定间隙，得到瓦片对齐的充分上界 $M_{\text{ub}}^_ = b_n \lceil\frac{\mathcal{G}^_}{\sigma}\rceil$，其中 $\mathcal{G}^_ = [\frac{\mathcal{V}}{\mathcal{R}} - T_{\text{iss}}^0]_+$。使用一次性校准的平台参数和最差对齐情况 $\mathcal{K}=148$，得到 $\mathcal{G}^_ \approx 6.0 \mu s$ 和 $M_{\text{ub}}^_ \approx 1.0\text{K}$ 令牌。对于 $M \ge M_{\text{ub}}^_$，排空引起的分量 $E_{\text{drain}}$ 消失；测量的 $E_{\text{res}}$ 可能保留一个小的发出、启动或测量下限。真实的阈值不大于此界限，因为省略的固定间隙是正的。它还可能更小，因为计算关键路径只需要 $\Delta t \le \mathcal{G}_{\text{QK}}$，并且较短的序列可能使用少于148个并发生产者。在评估的最小序列长度 $M=8,192$ 时，测量的Q-loop间隙约为 $80 \mu s$，大约是最差情况 $\mathcal{G}^*$ 的13倍。

两种模型的预测对比：评估的范围因此区分了两个先前的预测。X-Stage预测排空引起的残余不会随序列长度增长，$E_{\text{res}}^{\text{XS}}(M) \approx E_{\text{issue}}(M)$，其中 $E_{\text{issue}}$ 是一个小的无反压发出成本。完成耦合的反事实预测了一个额外的分量 $E^{\text{cc}}(M) = N_t \Delta t^{\text{cc}} \propto M$。第5.3节测试的是斜率的差异，而不仅仅依赖于单个短序列点。

4.4 总结

两种设计模式的总结：MegaMoE和FlashAttention-A2A在不同区域实例化了相同的设计测试。MegaMoE违反了速率界限，因为集中的Linear-2突发被短的mainloop分隔。跨波次的Linear-1/Linear-2交错重塑了注入；从依赖关系推导出的 $D_{\text{knee}}$ 提供了抗饥饿阈值，而不必要的大领先距离可能损害缓存局部性。FlashAttention-A2A满足容量界限，并且对于 $M \ge M_{\text{ub}}^*$，其Q-loop间隙超过了恢复间隙。因此，拥有输出的角色可以在瓦片边界搭载发出操作并恢复计算，同时由X-Stage排空请求。两种设计都保留了通信量；它们的不同之处在于工作负载是需要注入重塑，还是已经提供了足够的排空窗口。

A4 实验环境

  • 硬件配置：实验在一台配备8个GPU的系统上运行，该系统采用最新的NVIDIA架构。每个GPU有148个流式多处理器（SM）；GPU之间通过NVLink/NVSwitch全连接。
  • 软件配置：软件栈包括CUDA 13.1、PyTorch 2.9、DeepGEMM（commit 7f2a703）和FlashAttention（commit 77aacb6）。
  • 并行配置：实验采用8路专家并行（EP=8）。
  • 测量方法：
    • 通用：所有配置在测量前都进行预热，报告中位数值。所有成对比较都使用相同的输入、进程映射和内核配置。使用CUDA事件测量内核延迟，Nsight Systems的GPU指标提供Tensor Core时间线，内核内的clock64()插桩测量发送方可见的远程存储跨度。
    • MegaMoE：每次运行前进行一次不计时的JIT编译和预热，然后收集20次屏障对齐的CUDA事件试验。对每次试验，取所有rank的最大延迟，并报告这些试验最大值的中位数。
    • FlashAttention-A2A：计时采用发送方侧的稳态边界。通过背靠背重放来暴露持续的发送方侧反压。对于FA4，使用CUDA Graph重放以消除Python启动开销。

A5 实验结果

5.2 MegaMoE评估

5.2.1 内核性能

  • 实验内容：在总共84种配置（7种模型形状、W4A8/W8A8精度、均衡/偏斜路由、多种序列长度和专家配置）上，比较了Expert-Wave调度器和本文提出的交错调度器。两种实现共享相同的计算内核、线程束专业化、对称缓冲区、epilogue和同步路径，仅改变就绪的Linear-1和Linear-2瓦片的调度顺序。
  • 实验结果：如图8所示，在所有配置中，交错调度器实现了1.18倍的几何平均加速比、1.17倍的中位数加速比和1.62倍的最大加速比。偏斜路由（skewed routing）比均衡路由（balanced routing）获益更多，因为前者在热门专家上产生更长的Linear-2 epilogue运行，使得插入的Linear-1工作能提供更有价值的排空窗口。少数配置出现性能下降（约0.94倍），分析表明跨专家交错可能降低了专家权重的L2缓存局部性，抵消了通信侧的增益。


图8. MegaMoE内核加速比。X-Stage感知的交错调度器相对于专家波次调度器的加速比，涵盖不同模型、精度和路由配置。

5.2.2 Tensor Core时间线

  • 实验内容：分析Tensor Core的活动时间线，以验证增益是否由X-Stage反压引起。
  • 实验结果：如图9所示，Expert-Wave调度器在集中的远程存储注入后，Tensor Core活动出现周期性下降或长尾。而交错调度将通信突发分散在Linear-1和Linear-2工作之间，提高了平均Tensor流水线活动并缩短了低利用率区间。这表明注入重塑与利用率提高之间存在相关性。


图9. MegaMoE Tensor Core时间线（Tensor-pipe active，峰值百分比）。交错调度器将通信突发分布在Linear-1和Linear-2工作之间，减少了与反压相关的低Tensor Core活动区间。

5.2.3 逐瓦片远程存储跨度与关键路径停顿

  • 实验内容：对每个Linear-2瓦片插桩测量两个量：epilogue中发送方可见的远程存储跨度 $T_{RS}$，以及MMA线程束组在瓦片入口等待暂存槽位的时间（tmem_empty）。
  • 实验结果（跨度分布）：如图10所示，对于偏斜的W8A8路由，Expert-Wave调度的中位跨度在7.9到9.9μs之间且有明显长尾。交错后，所有七种配置的跨度都收窄到3.6-3.9μs，接近本地HBM存储的控制组结果。这表明调度顺序的改变将跨度分布从受模型影响的宽泛区域移到了一个类似下限的窄区域。
  • 实验结果（计算停顿映射）：使用公式15（$\Delta t = [T_{RS} - T_{\text{cover}}^{\text{mma}}]_+$）预测关键路径停顿，并与测量的MMA侧停顿进行比较。如表2所示，对于有明显反压的六种偏斜配置，预测的停顿（1.5-4.9μs）与测量的中位数停顿（0.65-4.3μs）的顺序和规模一致。对于交错调度，预测和测量值均接近零。这支持了从X-Stage反压经由双缓冲暂存到Tensor Core关键路径的机制。


图10. 一个MegaMoE Linear-2瓦片的发送方可见远程存储跨度。专家波次调度下连续的Linear-2工作产生更长、重尾的跨度。交错调度将分布移向无反压范围，此时Linear-1计算可以覆盖存储操作。

表2. 逐瓦片停顿模型的测试（W8A8，偏斜路由，55,808个令牌，8个GPU；时间单位μs）。跨度条目是wave/interleaved的中位数；$T^{mma}$条目是全瓦片L2/L1 mainloop的中位数。每个$\Delta t$条目是预测值/测量值。

模型 跨度 (Wave/int.) Tmma (L2/L1) ∆t (预测/测量) (Wave) ∆t (预测/测量) (Int.)
DiT-MoE 9.93 / 3.77 5.30 / 9.77 4.63 / 4.18 0 / 0.04
Qwen3.5 9.17 / 3.82 4.24 / 14.40 4.94 / 4.31 0 / 0.04
Hy3 9.68 / 3.59 6.09 / 10.79 3.60 / 3.08 0 / 0.05
MiMo-V2.5 9.22 / 3.94 7.77 / 7.57 1.45 / 0.71 0 / 0.05
GLM-5.2 9.53 / 3.63 7.99 / 18.77 1.54 / 0.92 0 / 0.04
DSv4-Flash 9.48 / 3.83 8.01 / 7.40 1.47 / 0.65 0 / 0.05
DSv4-Pro 7.88 / 3.55 9.54 / 11.57 0 / 0.06 0 / 0.04

5.3 FlashAttention–A2A端到端性能

  • 实验内容：将Ulysses All-to-All路径与FlashAttention-3（FA3）和FlashAttention-4（FA4）融合，并扫描序列长度M。比较了四种实现：FA-only、A2A-only、Serial（串行执行FA和A2A）和X-Stage-Fused。主要评估指标是相对于Serial的发送方可见加速比，以及超出FA-only的发送方侧残余时间 $E_{res} = T_{fused} - T_{FA}$。
  • 实验结果（端到端加速比）：如表3所示，在较短的序列长度（M=8,192）下，FA4和FA3分别达到1.42倍和1.43倍的最大加速比。随着序列增长，由于FlashAttention的计算时间增长快于A2A的通信时间，加速比自然下降至约1.05-1.09倍，这符合预期。
  • 实验结果（残余暴露测试）：如图11所示，测得的残余时间 $E_{res}$ 很小（10-14μs），并且不随序列长度M增长。这与X-Stage模型的预测（残余时间接近一个小的、与发出指令相关的常数下限）一致，而与完成耦合模型的预测（残余时间随M线性增长）相悖。在长序列下，残余时间在测量噪声范围内接近于零。这表明A2A的发送方侧开销几乎完全被Q-loop的计算所覆盖。


图11. FA4+A2A的相对发送方可见残余 $E_{res}(M)/T_{FA}$。先验预测使用一次性校准的平台参数和测量的计算侧时序，无应用重拟合。虚线参考代表约16μs的发出下限，蓝色带是±2%的DVFS包络。最坏情况收敛阈值 $M_{ub}^*(K=148) \approx 1.0K$ 位于所有测量点之下。

表3. 发送方可见的X-Stage感知FlashAttention-A2A性能。除隐藏率和加速比外，时间单位均为μs。

方法 / M 8,192 16,384 32,768 49,152 65,536
FA3 + A2A
FA3 only 195.2 785.5 3,236.7 7,275.6 13,249.4
Serial 295.6 968.3 3,579.6 7,767.9 13,936.6
Fused 207.0 795.5 3,241.5 7,285.7 13,218.3
隐藏率 88% 95% 99% 98% ~100%*
加速比 1.428× 1.217× 1.104× 1.066× 1.054×
FA4 + A2A†
FA4 only 85.4 346.5 1,439.4 3,281.4 5,340.7
Serial 140.9 462.4 1,660.9 3,661.7 5,813.0
Fused 99.4 356.6 1,429.2 3,282.1 5,348.9
Eres +14.0 +10.1 -10.2 +0.7 +8.2
隐藏率 75% 91% ~100%* 100% 98%
加速比 1.417× 1.297× 1.162× 1.116× 1.087×

*值等于或略高于100%反映了当A2A发出开销被隐藏在发送方侧计时边界内时的测量噪声。†FA4使用稳态CUDA Graph重放；M=65,536列是文本中描述的时钟稳定重跑结果。FA3使用每次迭代的稳态计时。指标由未四舍五入的中位数计算得出。

A7 补充细节

6.1 X-Stage作为校准的执行抽象

X-Stage作为校准的执行抽象：X-Stage抽象了在远程存储发出和远程可见完成之间观察到的聚合后发出行为。这种行为可能反映了存储流水线、缓存层次、网络注入、流控信用和接收端资源等多种因素的综合效应。X-Stage并不要求软件单独识别或建模每个组件，而是将它们的外部可见效果总结为 $\mathcal{M}_X = (T_{\text{iss}}^0, \mathcal{R}, Q)$。

X-Stage不替代完成语义：不同的系统可能会产生不同的参数值。相同的校准和调度方法仅在其外部行为能被“发出-排空-容量”模型充分描述时适用。拓扑、路由、存储宽度、生产者数量、内存布局或GPU代数的变化可能需要重新校准，并可能暴露当前流体近似模型之外的行为。排空视界 $\mathcal{V}/\mathcal{R}$ 预测了发送方何时可以返回低成本发出状态，但它不替代远程完成或内存排序语义。X-Stage感知的内核保留了所需的栅栏、信号、缓冲区生命周期规则和消费者侧的就绪检查。

6.2 两级流水线解耦

两级流水线解耦：一个融合的GEMM epilogue通常将计算结果写入远程内存。本地缓冲和线程束专门化将Tensor Core线程束组与epilogue或发出角色解耦。这个软件层可以吸收短暂的不匹配，但如果发出持续受到反压，有限的槽位最终会填满。X-Stage在接受发出后暴露了第二个层次的解耦。当有效的未完成容量可用时，请求可以向远程可见性推进，而发出角色可以恢复计算或epilogue工作。这种解耦是有限的，而非无限的。持续以高于有效排空率的速度注入会累积请求，延长后续发出操作的时间，并通过本地暂存向后传播反压。增加软件缓冲区可以延迟这种传播，但不能消除速率或容量的限制。

6.3 自适应突发整形与范围

自适应突发整形与范围：MegaMoE的交错是一种突发整形：它保留了总工作量和通信量，但在时间上重新分配了远程存储。更一般地，调度器可以使用 $\frac{\mathcal{V}}{T_{\text{iss}}^0 + \mathcal{G}} \le \mathcal{R}$ 和 $[\mathcal{V} - \mathcal{R}T_{\text{iss}}^0]_+ \le Q$ 来联合选择瓦片大小、有用间隙、软件暂存深度和发出资源数量。跨瓦片或专家的重新排序在保持过多跨波次工作在途时也可能改变缓存局部性。一个实用的生成器或运行时应该比较预测的X-Stage停顿减少量与局部性成本，并为当前的路由和工作集选择原始或交错调度。

本文证据的局限性：本文提出的证据有两个局限性。首先，微基准测试研究的是单边远程存储；加载、原子操作、具有不同进展引擎的集合通信以及跨节点网络可能表现出不同的约束。其次，该模型预测的是发送方侧的反压，但本身并未考虑同步、接收方拥塞或所有启动和调度开销。这些局限性是我们为什么将X-Stage视为一个校准的执行抽象，而不是一个通用的硬件描述的原因。

A6 结论

本文识别了X-Stage，一个在远程存储发出和远程可见完成之间的软件可见执行阶段。已接受的存储请求可以继续进行，同时发出者恢复有用的工作，但持续的注入最终会填满有限的有效容量，并对后续的发出、epilogue乃至Tensor Core生产者产生反压。一个轻量级的Burst-Gap模型通过无反压发出时间、有效排空率和有效未完成容量来捕捉这种行为。

该模型指导了两种互补的内核转换。对于MegaMoE，跨专家波次的Linear-1/Linear-2交错重塑了Combine突发，在84种配置上取得了1.18倍的几何平均和1.62倍的最大内核加速比。对于FlashAttention-A2A，一个搭载式流水线让拥有输出的角色发出远程存储并恢复计算，无需专用的通信线程束或SM。FA3和FA4分别达到了1.43倍和1.42倍的最大加速比，并且随着Q-loop覆盖了后发出排空，它们的发送方可见稳态时间接近于单独的FlashAttention时间。

X-Stage通过揭示一个缺失的调度维度——已接受但尚未完成的远程存储的速率和容量——补充了现有的瓦片依赖和远程就绪协议。当前的表征涵盖了单边远程存储，因此其他通信操作需要独立的验证。在此范围内，该模型提供了远程存储行为与融合内核调度之间一个可测量的联系。