ClusterFusion: Expanding Operator Fusion Scope for LLM Inference via Cluster-Level Collective Primitive

  • 作者: Xinhao Luo, Zihan Liu, Yangjie Zhou, Shihan Fang, Ziyu Huang, Yu Feng, Chen Zhang, Shixuan Sun, Zhenzhe Zheng, Jingwen Leng, Minyi Guo
  • 机构: Shanghai Jiao Tong University, Shanghai Qi Zhi Institute, National University of Singapore

A1 主要贡献

大型语言模型（LLM）的解码阶段由于算子执行碎片化以及严重依赖片外内存进行数据交换和归约，导致延迟较高。这种执行模式限制了算子融合的机会，并带来了显著的内存流量和内核启动开销。为了解决这些问题，本文提出了一个名为ClusterFusion的执行框架，旨在通过扩展算子融合范围来加速LLM推理。

核心问题与研究目标：
现有的LLM推理系统，特别是在解码阶段，面临以下瓶ăpadă：
1. 执行碎片化：解码过程被分解为多个独立的内核（如QKV投射、注意力、输出投射），每个内核之间需要通过全局内存进行数据交换和同步。
2. 高昂的片外内存开销：线程块间的依赖关系通过将中间结果物化到全局内存来解决，导致频繁的片外数据传输。
3. 有限的算子融合范围：由于缺乏线程块间的结构化通信机制，算子融合的机会受到限制，无法有效地在片上重用数据。

尽管NVIDIA Hopper等现代GPU架构提供了分布式共享内存（DSMEM）和低延迟的集群内互连，但它们只暴露了低级的数据移动指令，缺乏用于集体片上通信的结构化抽象。研究的目标是弥合这一软硬件差距，利用新硬件特性来扩展算子融合范围，减少延迟。

创新点与主要贡献：
本文的主要贡献如下：
1. 分析与洞察：分析了现有LLM解码工作流中的通信模式和融合范围，指出现有的碎片化内核执行和片外同步是高效解码融合的关键障碍。此外，对NVIDIA Hopper GPU上的DSMEM机制进行了性能剖析，揭示了其支持低延迟块间通信和减少片外内存依赖的潜力。
2. 提出集群级集体通信原语：为了支持结构化的线程块间集体通信，提出了两个集群级集体原语：ClusterReduce和ClusterGather。这些原语抽象了在DSMEM上的归约和聚合操作，实现了线程块之间的高效协调。
3. 开发ClusterFusion执行框架：基于所提出的原语，开发了ClusterFusion，一个通过扩展算子融合来优化推理的执行框架。ClusterFusion采用“以集群为中心的数据流”（cluster-centric dataflow），将结构化的集群内通信集成到计算流程中，成功地将QKV投射、注意力和输出投射等阶段融合成一个单一的内核。这种方法无需片外内存流量即可实现计算和通信的协调，显著优于现有技术。

图1: 典型的Transformer块，作为现代LLM的基本组件。
图1: 典型的Transformer块，作为现代LLM的基本组件。
图2: 预填充和解码256个token的延迟比较。
图2: 预填充和解码256个token的延迟比较。

A3 背景知识

2.1 LLM推理工作流和瓶颈

LLM的结构与计算。大型语言模型通常基于仅解码器的Transformer架构【索引46，Attention is all you need. NeurIPS 2017】构建。如图1所示，每个Transformer块包含QKV投射、注意力、输出投射和一个前馈网络（FFN）。对于输入序列X，每个注意力头计算投射$Q_i = XW_Q^i$, $K_i = XW_K^i$, $V_i = XW_V^i$，然后进行缩放点积注意力和输出投射，公式如下：

$$Z = \text{Concat} \left( \text{Softmax} \left( \frac{Q_1 K_1^\top}{\sqrt{d_k}} \right) V_1, \dots, \text{Softmax} \left( \frac{Q_h K_h^\top}{\sqrt{d_k}} \right) V_h \right) W_O$$


其中，$d_k$是每个头的查询（Q）和键（K）的维度。FFN模块应用三个线性层，中间带有一个非线性激活函数，公式为：

$$\text{FFN}(Z) = W_3 \left( \sigma(W_1 Z) \odot W_2 Z \right)$$
其中，$\sigma$是一个非线性激活函数（例如GELU），$\odot$表示逐元素乘法。

解码阶段是推理瓶颈。在推理过程中，模型分为两个阶段：预填充和解码。在解码阶段，模型以自回归方式一次生成一个token，重用KV缓存并附加新条目。这种顺序解码模式天然地限制了并行性，并在上下文长度较长时主导了推理延迟。如图2所示，图中百分比表示不同序列长度下解码阶段的延迟占比。使用SGLang【索引52，Sglang: Efficient execution of structured language model programs. NeurIPS 2024】测量，生成256个token序列时，解码占总延迟的95%以上。这使得解码成为推理过程中的主要计算瓶颈，也是系统级性能优化的天然目标。

2.2 现有的通信模式和融合范围

执行数据流的重要性。许多研究提出了各种技术【索引9，Flashattention: Fast and memory-efficient exact attention with io-awareness. NeurIPS 2022】【索引8，Flashattention-2: Faster attention with better parallelism and work partitioning. ICLR 2024】【索引38，Flashattention-3: Fast and accurate attention with asynchrony and low-precision. NeurIPS 2024】【索引10，Flash-decoding for long-context inference. 2023】【索引15，Flashdecoding++: Faster large language model inference with asynchronization, flat GEMM optimization, and heuristics. MLSys 2024】【索引49，Flashinfer: Efficient and customizable attention engine for LLM inference serving. CoRR 2025】【索引51，Bytetransformer: A high-performance transformer boosted for variablelength inputs. IPDPS 2023】来提升Transformer性能。其中，执行数据流的设计越来越受关注，它指的是在并行执行模型和内存层次结构上对计算阶段及其相关数据移动（包括分区、调度和通信）的结构化组织【索引28，Cuda c++ programming guide. https://docs.nvidia.com/cuda/cuda-c-programming-guide/index.html】。在解码工作负载中，数据流通过管理算子调度和中间数据交换，在决定端到端延迟方面起着核心作用 。

现有数据流的局限性。在现有的基于GPU的解码数据流中，线程块通常作为并行执行和调度的基本单位。图3展示了SGLang【索引52，Sglang: Efficient execution of structured language model programs. NeurIPS 2024】中Llama2-7B【索引45，Llama 2: Open foundation and fine-tuned chat models. CoRR 2023】解码阶段的数据流，涵盖QKV投射、注意力和输出投射三个阶段。在每个内核内部，线程块被分配给各个注意力头，并操作隐藏维度和KV序列的不相交瓦片（tile）。在QKV投射阶段，每个线程块对其分配的隐藏状态执行线性变换，以产生局部的Q、K和V向量，这些输出被写入全局内存。注意力阶段使用FlashDecoding【索引10，Flash-decoding for long-context inference. 2023】【索引15，Flashdecoding++: Faster large language model inference with asynchronization, flat GEMM optimization, and heuristics. MLSys 2024】【索引49，Flashinfer: Efficient and customizable attention engine for LLM inference serving. CoRR 2025】实现，其中每个块首先使用其对应的Q和分段的KV缓存计算部分注意力结果。然后一个单独的重缩放内核跨块聚合这些部分结果。最终的输出投射在聚合后的注意力输出上以块为单位执行。现有的解码数据流表现出两种形式的块间通信：交换多个块所需的中间数据（如Q/K/V向量），以及跨块归约部分结果（如注意力输出）。这些依赖关系通过片外内存和显式的内核边界来解决，导致全局同步障碍并阻碍了算子融合【索引54，Mononn: Enabling a new monolithic optimization space for neural network inference tasks on modern gpu-centric architectures. OSDI 2024】【索引47，Optimizing deep learning inference via global analysis and tensor expressions. ASPLOS 2024】。

块隔离执行的开销。这种块隔离的执行结构引入了启动开销、片外内存往返和内核间的全局同步障碍。由于缺乏跨线程块的结构化通信，中间结果必须物化到全局内存中，并由后续内核重新加载，这限制了有效的片上数据重用和更广泛融合的机会。为了克服这些限制，我们需要一种能够实现跨线程块的集体调度和通信的执行机制。

图3: 现有的QKV投射、注意力和输出投射数据流图。
图3: 现有的QKV投射、注意力和输出投射数据流图。

2.3 集群级的机遇与挑战

新硬件架构带来的机遇。前一节指出，主流方法通过全局内存同步来解决块间数据依赖。然而，如NVIDIA Hopper【索引27，Nvidia hopper architecture. https://www.nvidia.com/en-us/data-center/technologies/hopper-architecture/】等最新架构（如图4所示）可以将一组线程块分组到一个集群中。在每个集群内，线程块可以通过高速的SM-to-SM片上网络（NOC）（也称为DSMEM）直接共享数据，从而避免昂贵的全局内存访问并实现高效的集群内通信 。
图4: NVIDIA Hopper 架构。

DSMEM性能分析。为了理解这种机制的机遇与挑战，我们在NVIDIA H100 GPU上对DSMEM进行了性能剖析，将集群大小从1变为16（Hopper硬件支持的最大值）。如图5左侧所示，SM-to-SM的访问延迟在小集群尺寸下显著改善。当集群大小为2时，平均延迟达到190个周期，远低于全局内存延迟（超过470个周期）。这展示了DSMEM在低延迟片上通信方面的潜力。

DSMEM的权衡与挑战。然而，这种优势也带来了显著的权衡。随着集群大小的增加，可用的通信带宽会下降，当集群大小达到16时，带宽略低于全局内存带宽（2.90 TB/s vs. 2.96 TB/s），这是由于交叉开关架构【索引19，Uncovering real gpu noc characteristics: Implications on interconnect architecture. MICRO 2024】【索引3，https://patents.google.com/patent/US20230289189A1/en】导致的。此外，由于硬件限制，活动SM的数量也会减少。这些因素限制了基于集群的执行的可扩展性，并需要仔细配置以平衡通信效率和整体并行性 。

软件层面的挑战。除了硬件层面的权衡，还存在软件层面的挑战。NVIDIA目前仅通过低级的PTX指令暴露DSMEM和线程块集群功能，这些指令只支持线程块之间基本的、点对点的数据移动【索引27，Nvidia hopper architecture. https://www.nvidia.com/en-us/data-center/technologies/hopper-architecture/】。这种低级接口给在实际场景中表达可重用的同步和通信模式带来了巨大挑战，导致了陡峭的编程障碍，使开发者缺乏关于如何有效应用这些功能的明确指导 。
图5: H100 GPU上不同集群大小下的SM-to-SM访问延迟（左）、带宽（中）和活动SM数量（右）。

A2 方法细节

3.1 集群级集体通信原语

设计思路与原语定义。为了探索硬件层面的权衡并克服软件层面的挑战，我们引入了两个抽象了常见通信模式的集群级集体原语。核心思想是将一个线程块集群视为一个全连接的逻辑网络，其中所有块参与结构化的集体操作。受分布式系统中通信原语【索引42，Megatron-lm: Training multi-billion parameter language models using model parallelism. CoRR 2019】【索引4，Deepspeed-inference: Enabling efficient inference of transformer models at unprecedented scale. SC22 2022】的启发，我们设计了两个集群级原语：ClusterReduce，它使用诸如求和或求最大值之类的关联运算符跨块归约数据；以及ClusterGather，它将每个块的本地数据复制到所有其他块以进行数据共享，如图6所示。

原语实现细节。如算法1和算法2所示，ClusterReduce和ClusterGather都采用了跨越$log_2N$轮的二叉树模式，其中N是集群大小。在每一轮中，通信步长（stride）加倍，每个块与一个索引偏移当前步长的对等块交换数据。ClusterReduce在保持消息大小不变的情况下执行逐元素归约，而ClusterGather则通过在每一轮中将消息大小加倍来逐步累积远程数据。这种共享的结构有助于统一实现和硬件调优，从而实现高效的集群级同步和数据共享。

算法 1 在DSMEM上的ClusterReduce - 线程块视角

需要: 一个包含N = 2^k个线程块（k ≤ 4）的集群，每个块的秩为b ∈ [0, N - 1]，并有一个共享内存缓冲区Db，其中包含需要与同一集群中其他线程块的数据一起归约的本地数据。一个归约操作符⊕（例如，求和或求最大值）。
1: 分配与Db大小相同的共享内存缓冲区Bb ▷ 用作临时缓冲区。
2: stride ← 1.
3: while stride < N do
4:     块秩 send_to ← (b + stride) mod N.
5:     块秩 recv_from ← (b - stride + N) mod N.
6:     通过DSMEM将Db发送到块send_to的Bsend_to。
7:     通过DSMEM从块recv_from接收Drecv_from到Bb。
8:     等待Drecv_from的到达。
9:     Db ← Db ⊕ Bb ▷ 使用归约操作符聚合部分结果。
10:    stride ← stride × 2 ▷ 指数级步长增长。
11: end while
12: 返回 Db。
算法 2 在DSMEM上的ClusterGather - 线程块视角

需要: 一个包含N = 2^k个线程块（k ≤ 4）的集群，每个块的秩为b ∈ [0, N - 1]，并有一个大小为N × size的共享内存缓冲区Db。第一个段Db[0 : size]包含需要与同一集群中其他线程块的数据一起收集的本地数据。
1: stride ← 1.
2: while stride < N do
3:     块秩 send_to ← (b + stride) mod N.
4:     块秩 recv_from ← (b - stride + N) mod N.
5:     通过DSMEM将Db[0 : size × stride]发送到块send_to的Dsend_to[stride × size : 2 × stride × size]。
6:     通过DSMEM从块recv_from接收Drecv_from[0 : size × stride]到Db[stride × size : 2 × stride × size]。
7:     等待Drecv_from[0 : size × stride]的到达。
8:     stride ← stride × 2 ▷ 指数级步长增长。
9: end while
10: 返回 Db。
图6: 集群级集体通信原语示意图：ClusterReduce和ClusterGather。
图6: 集群级集体通信原语示意图：ClusterReduce和ClusterGather。

3.2 采用原语的以集群为中心的数据流设计

核心思想。基于所提出的集群级原语，我们现在阐述如何利用它们来构建一个具有扩展算子融合范围的以集群为中心的数据流。核心思想是将线程块集群视为一个协作执行和调度的单元。数据相关的维度保留在每个集群内部，以使用集群级集体原语解决块间数据依赖，从而避免片外数据交换；而数据独立的维度则分布在不同集群之间。这种组织方式将计算与内存局部性对齐，并通过允许中间结果完全在片上内存中重用，实现了无缝的内核融合。

解码阶段的数据流并行化策略。在解码阶段，投射（Projection）和注意力（Attention）模块的数据流在多个维度上进行并行化。具体来说，投射沿头数和头维度并行，而注意力沿头数和KV缓存序列长度并行。在这些维度中，计算头维度和KV缓存序列长度不同分区的线程块之间存在块间数据依赖，因为每个块要么计算一个需要归约的部分结果，要么计算一个需要收集以形成最终输出的段结果。这些线程块可以被分组到一个集群中，通过使用集群级集体原语在片上解决数据依赖。由于注意力头在所有三个模块中都是独立的，因此每个集群相应地映射到一个单独的头。如图7所示，在这种设计下，QKV投射产生的中间结果自然地保留在片上，并被注意力模块直接重用。同样，注意力模块的输出也保留在片上，并立即被输出投射消耗，从而实现了跨三个模块的无缝数据重用。
图7: 以集群为中心的融合QKV投射、注意力和输出投射的数据流图。

融合数据流算法。通过利用集群级通信原语，我们实现了融合的QKV投射、注意力和输出投射数据流。如算法3所示，该数据流在头数上是并行的。每个头对应一个包含N = $2^k$个线程块（k ≤ 4）的集群，其中每个块被分配一个秩b ∈ [0, N-1]。在每个集群内，线程块分别划分QKV投射中的头维度、注意力中的KV缓存序列维度以及输出投射中的输出维度。对于整个数据流，每个线程块处理完整的输入隐藏状态，并在输出投射后计算相应的输出瓦片$O_b$。在此算法中，B表示批量大小，D表示输入隐藏维度，H表示总头维度，而h、s和d分别表示每个线程块划分的头维度、序列长度和输出维度的大小。

数据流变体与流量分析。根据我们以集群为中心的数据流设计原则，我们还提出了几种数据流变体，这些变体在附录B中介绍。为了评估这些变体，我们对DSMEM流量进行了定量分析【索引41，Welder: Scheduling deep learning memory access via tile-graph. OSDI 2023】【索引53，Tileflow: A framework for modeling fusion dataflow via tree-based analysis. MICRO 2023】。我们首先分析ClusterReduce和ClusterGather原语引起的DSMEM内存流量如下：
$Traffic_{Reduce}(size, N) = size \times N \times log_2N$
$Traffic_{Gather}(size, N) = size \times N \times (N-1)$
此处，$Traffic_{Reduce}$和$Traffic_{Gather}$分别表示ClusterReduce和ClusterGather的DSMEM流量。变量size表示算法1中共享内存缓冲区$D_b$的大小，以及算法2中初始段的大小。基于此分析模型，我们估计算法3中使用的数据流的总DSMEM内存流量，该数据流包括一个ClusterGather和两个ClusterReduce操作：

$$Traffic_{Total} = Traffic_{Reduce}(3h, N) + Traffic_{Gather}(H, N)$$


我们忽略了softmax统计量产生的流量，因为它只涉及两个浮点数，与张量相比可以忽略不计。根据附录B中的分析，该数据流产生的DSMEM流量最低，并且相比其他数据流变体实现了更好的性能。

通用性与框架实现。所提出的以集群为中心的数据流设计原则可以自然地推广到其他算子，包括DeepSeek MLA【索引11，Deepseek-v2: A strong, economical, and efficient mixture-of-experts language model. CoRR 2024】。相应的算法在附录B中提供。我们基于以集群为中心的数据流实现了一个端到端的执行框架ClusterFusion，该框架整合了上述融合的QKV投射、注意力和输出投射模块。对于FFN和RMSNorm等其他组件，我们采用了与现有框架如CUTLASS【索引30，Cutlass. https://github.com/NVIDIA/cutlass】和Flashinfer【索引49 ，Flashinfer: Efficient and customizable attention engine for LLM inference serving. CoRR 2025】中一致的优化实现。

A4 实验

实验环境

  • 硬件配置:

    • GPU: NVIDIA H100 SXM5 80GB

  • 软件配置:

  • 模型:

    • Llama2-7B【索引45，Llama 2: Open foundation and fine-tuned chat models. CoRR 2023】: 采用标准的多头注意力（MHA）机制。
    • DeepSeekV2-Lite【索引11，Deepseek-v2: A strong, economical, and efficient mixture-of-experts language model. CoRR 2024】: 采用多头潜在注意力（MLA）算法。

  • 实验设置:

    • 精度: 所有输入均为FP16。
    • 批量大小: 1 (多批次结果见附录C)。
    • 上下文长度: 从1K到16K不等。

  • 基线系统 (Baselines):

    • SGLang 0.4.3.post2【索引52，Sglang: Efficient execution of structured language model programs. NeurIPS 2024】
    • vLLM 0.6.4.post1【索引20，Efficient memory management for large language model serving with pagedattention. SOSP 2023】
    • TensorRT-LLM 0.18.0【索引33，Tensorrt-llm. https://github.com/NVIDIA/TensorRT-LLM 】
    • MLC-LLM 0.20.dev0【索引26，MLC-LLM, 2023-2025.】
    • 所有基线均使用官方推荐配置，并启用了CUDA Graph【索引29，Cudagraph. https://developer.nvidia.com/blog/cuda-graphs/】和Torch.compile【索引5 ，Pytorch 2: Faster machine learning through dynamic python bytecode transformation and graph compilation. ASPLOS 2024】等优化。

实验结果

端到端和核心模块评估

  • 端到端性能：使用每输出token时间（TPOT）作为指标。如图8所示，在Llama2-7B上，ClusterFusion（集群大小为4）的平均加速比分别为：SGLang 1.41倍，vLLM 1.39倍，TensorRT-LLM 1.43倍，MLC-LLM 2.03倍。在DeepSeek-V2-Lite上，平均加速比分别为1.34倍、1.37倍、1.51倍和2.39倍。
  • 核心模块性能：如图9所示，对于QKV投射、注意力和输出投射这三个核心模块，ClusterFusion在Llama2-7B上相比SGLang、vLLM、TensorRT-LLM和MLC-LLM的平均加速比分别为1.85倍、1.73倍、1.61倍和3.19倍。在DeepSeek-V2-Lite上，平均加速比分别为1.66倍、1.64倍、1.35倍和3.5倍。
  • 实际场景有效性：尽管DeepSeek MLA的优化空间相对有限，但图10显示真实数据集（ShareGPT【索引1，Sharegpt. https://sharegpt.com/】和Splitwise【索引35 ，Splitwise: Efficient generative LLM inference using phase splitting. ISCA 2024】【索引2，Splitwise. https://github.com/Azure/AzurePublicDataset/blob/master/AzureLLMInferenceDataset2023.md】）中的序列长度主要在8k以下，在此范围内ClusterFusion仍能实现显著的性能提升 。

图8: 在H100上Llama2-7B（左）和DeepSeek-V2-Lite（右）的TPOT。
图8: 在H100上Llama2-7B（左）和DeepSeek-V2-Lite（右）的TPOT。
图9: 在H100上Llama2-7B（左）和DeepSeek-V2-Lite（右）核心模块的延迟。
图9: 在H100上Llama2-7B（左）和DeepSeek-V2-Lite（右）核心模块的延迟。

图10: ShareGPT和Splitwise数据集中序列长度的分布。

图10: ShareGPT和Splitwise数据集中序列长度的分布。

不同配置下的核心模块性能
- 如图11所示，评估了不同集群大小和注意力头数下核心模块的性能。由于每个注意力头映射到一个集群，集群大小决定了头内部的并行度。
- 结论：当注意力头数为32和64时，集群大小为4的性能最佳。当头数增加到128时，较小的集群大小2成为最优选择。集群大小为8和16时性能下降，原因是互连延迟增加、带宽争用和活动SM数量减少，限制了核心利用率，这与图5的分析一致。这表明最优集群大小随工作负载变化，需要相应调整以最大化性能。

图11: ClusterFusion核心模块在不同设置下的延迟。
图11: ClusterFusion核心模块在不同设置下的延迟。

加速分析
- 性能优势来源：ClusterFusion的性能优势主要来自两个方面：最小化的全局内存传输量和减少的GPU内核启动开销。
- 量化分析：使用NVIDIA Nsight Systems【索引32，Nvidia nsight system. https://developer.nvidia.com/nsight-systems】 和Nsight Compute【索引31，Nvidia nsight compute. https://developer.nvidia.com/nsight-compute】进行性能剖析。如图12所示，ClusterFusion通过将QKV投射、注意力和输出投射完全在片上执行，显著减少了中间内存流量。同时，即使与经 过CUDA Graph优化的基线相比，ClusterFusion也将内核启动开销减少了近一个数量级。

图12: 全局内存数据传输大小（左）和GPU内核启动开销（右）的比较。
图12: 全局内存数据传输大小（左）和GPU内核启动开销（右）的比较。

附加分析
- 原语性能微基准测试：为了验证片上互连的重要性，对第3.1节中介绍的集群级集体通信原语进行了微基准测试。如表1所示，与片外实现相比，通过DSMEM实现的片上ClusterReduce和ClusterGather在不同数据传输大小下都表现出显著更低的延迟。
- DSMEM消融研究：对启用和禁用DSMEM的ClusterFusion进行了比较。如图13所示，禁用DSMEM后，不同序列长度下的TPOT最多增加了33%。这些结果突显了DSMEM和集群级集体原语在实现高效片上归约和聚合方面的有效性，从而提升了端到端推理性能。

表1: 使用DSMEM的片上ClusterReduce和ClusterGather与不使用DSMEM的片外实现的延迟比较。

表1: 使用DSMEM的片上ClusterReduce和ClusterGather与不使用DSMEM的片外实现的延迟比较。

图13: ClusterFusion在Llama2-7B上启用和禁用DSMEM的TPOT比较。
图13: ClusterFusion在Llama2-7B上启用和禁用DSMEM的TPOT比较。

A7 补充细节

5 关于融合范围和架构展望的讨论

当前方法的局限性。ClusterFusion建立在集群内DSMEM通信之上，其中每个融合范围受限于一个固定的集群大小（最多16个线程块）【索引27，Nvidia hopper architecture. https://www.nvidia.com/en-us/data-center/technologies/hopper-architecture/】。这给融合粒度和调度灵活性带来了限制。尽管当今主流LLM【索引45 ，Llama 2: Open foundation and fine-tuned chat models. CoRR 2023】【索引11，Deepseek-v2: A strong, economical, and efficient mixture-of-experts language model. CoRR 2024】中的大多数解码算子（如投射和注意力）都能很好地适应这个限制，但未来具有更大隐藏维度或特殊算子变体的模型可能会挑战这个边界。当融合的算子超出集群范围时，系统必须回退到全局内存通信，这会引入额外的延迟和运行时碎片化。这激发了对支持更广泛片内集体通信的硬件【索引16，Dissecting the graphcore IPU architecture via microbenchmarking. CoRR 2019】【索引6，Cerebras wse-3 chip. https://www.cerebras.ai/chip】【索引23 ，Scaling deep learning computation over the inter-core connected intelligence processor with T10. SOSP 2024】【索引14，Waferllm: A wafer-scale LLM inference system. CoRR 2025】的反思。

对未来硬件设计的启示。我们的研究结果表明，实现跨更广泛SM集合的低延迟、拓扑感知的通信，将解锁更统一和可扩展的融合策略。这样的架构支持可以将结构化协调扩展到当前集群边界之外。ClusterFusion突显了这种软硬件协同设计的机会，是支持日益复杂的架构的实用一步。

A5 结论

本文介绍了ClusterFusion，一个通过联合调度通信和计算来扩展算子融合范围的执行框架，它将解码阶段的QKV投射、注意力和输出投射等算子组合成单一的融合内核。通过引入集群级的集体通信原语，ClusterFusion有效地减少了全局内存流量和内核启动开销，实现了关键LLM解码模块的高效片上执行。我们在NVIDIA H100 GPU上对代表性模型Llama2-7B和DeepSeek-V2-Lite进行的全面评估表明，ClusterFusion在不同模型和配置下均优于最先进的LLM推理框架。

A6 附录

A 相关工作

算子级优化。许多研究探索了LLM推理的算子级优化。FlashAttention【索引9，Flashattention: Fast and memory-efficient exact attention with io-awareness. NeurIPS 2022】【索引8，Flashattention-2: Faster attention with better parallelism and work partitioning. ICLR 2024】将整个注意力操作融合成一个内存高效的单一内核。在此基础上，FlashDecoding【索引10，Flash-decoding for long-context inference. 2023】在解码期间将并行性扩展到KV缓存序列维度。FlashDecoding++【索引15，Flashdecoding++: Faster large language model inference with asynchronization, flat GEMM optimization, and heuristics. MLSys 2024】提出预先根据统计数据确定缩放因子，并引入FlatGEMM来表示解码中维度高度缩减的GEMM。FlashAttention-3【索引38，Flashattention-3: Fast and accurate attention with asynchrony and low-precision. NeurIPS 2024】证明了warp特化和异步等新硬件特性【索引27，Nvidia hopper architecture. https://www.nvidia.com/en-us/data-center/technologies/hopper-architecture/】可以对注意力产生重大影响。FlashMLA【索引18 ，Flashmla: Efficient mla decoding kernels. https://github.com/deepseek-ai/FlashMLA】受FlashAttention-3启发， 为DeepSeek MLA架构设计了一个高效的解码注意力内核。然而，这些优化采用的是块隔离的执行模式。由于缺乏跨线程块的结构化通信，中间结果被反复写入和读取全局内存，限制了更广泛的算子融合和片上重用的机会。ClusterFusion通过高效的集群范围集体原语探索了一个更通用的融合空间。

算法级优化。一些研究通过算法层面的改变来提高LLM推理效率【索引22，AWQ: activation-aware weight quantization for on-device LLM compression and acceleration. MLSys 2024】【索引24，VQLLM: high-performance code generation for vector quantization augmented LLM inference. HPCA 2025】【索引50，Native sparse attention: Hardware-aligned and natively trainable sparse attention. CoRR 2025】【索引25，Moba: Mixture of block attention for long-context llms. CoRR 2025】。量化和稀疏化等技术旨在减少计算和内存开销。量化【索引22，AWQ: activation-aware weight quantization for on-device LLM compression and acceleration. MLSys 2024】【索引24，VQLLM: high-performance code generation for vector quantization augmented LLM inference. HPCA 2025】通过将高位宽表示转换为低位宽表示来压缩模型权重和激活，从而降低算术和内存成本。剪枝【索引50，Native sparse attention: Hardware-aligned and natively trainable sparse attention. CoRR 2025】【索引25，Moba: Mixture of block attention for long-context llms. CoRR 2025】增加了权重或激活中零元素的比例，使硬件能够跳过冗余计算。这些技术与我们的工作是正交的。ClusterFusion不修改模型工作负载，但可以通过结构化和可重用的原语优化底层数据流和内核执行，从而补充这些方法。

基于核间互联硬件的系统。片上核间互连也已在Graphcore IPU【索引16，Dissecting the graphcore IPU architecture via microbenchmarking. CoRR 2019】和Cerebras WSE【索引6，Cerebras wse-3 chip. https://www.cerebras.ai/chip】等替代硬件平台上采用。一些先前的工作探索了利用这些能力来优化深度学习工作负载。T10【索引23 ，Scaling deep learning computation over the inter-core connected intelligence processor with T10. SOSP 2024】是一个针对核间互联智能处理器的端到端深度学习编译器，重点是对核间数据移动进行细粒度建模。而WaferLLM【索引14，Waferllm: A wafer-scale LLM inference system. CoRR 2025】是第一个提出专为晶圆级加速器量身定制的LLM并行解决方案的系统。此外，这些系统与提供全核间连接或将核间通信机制集成到特定算子（如GEMM和GEMV）中的定制硬件架构紧密相连。因此，它们的设计假设不能很好地转换到像NVIDIA Hopper这样的现代GPU架构，而后者仍然是LLM推理部署的主流平台。ClusterFusion引入了集群范围的通信原语，这些原语抽象了常见的聚合和归约模式，并可以在不同工作负载中灵活重用。

B 额外的数据流

本节首先介绍DeepSeek MLA的计算过程，并展示我们融合的MLA数据流，该数据流利用了集群级集体通信原语，正如在主论文中评估的那样。接着，我们描述了另一种数据流，其中线程块在注意力模块内划分头维度，与主论文中沿token维度划分KV缓存的数据流形成对比，展示了我们的原语在为相同计算启用不同数据流方面的能力。最后，我们通过分析它们相应的DSMEM流量来评估这些数据流变体。

B.1 DeepSeek MLA

MLA机制介绍。DeepSeek引入了多头潜在注意力（MLA）机制，该机制已在其模型系列【索引11，Deepseek-v2: A strong, economical, and efficient mixture-of-experts language model. CoRR 2024】【索引12，Deepseek-v3 technical report. CoRR 2024】中被采用。在推理期间，MLA的解码阶段应用了权重吸收优化【索引37，Sglang deepseek model optimizations. https://github.com/sgl-project/sgl-learning-materials/blob/main/slides/sglang_deepseek_model_optimizations.pdf】以降低总体计算成本。原始MLA和DeepSeek-V2-Lite【索引11 ，Deepseek-v2: A strong, economical, and efficient mixture-of-experts language model. CoRR 2024】中优化版本的详细计算过程如图14所示。在原始MLA计算中，输入隐藏状态首先经过Q投射生成多头Q，并经过下投影（Down Projection）获得压缩的KV缓存。新生成的压缩KV随后与缓存的KV连接，并通过上投影（Up Projection）生成多头KV，用于MHA模块。在权重吸收版本的MLA中，过程从计算Q和K开始，其中V通过部分重用K获得：

$$Q = \mathrm{Hidden} \times W_{Q} \times W_{\mathrm{Up}}, \quad K = \mathrm{Hidden} \times W_{K}, \quad V = K[: \text{kv\_lora\_rank}]$$


接下来，计算注意力模块：

$$Z = \text{Concat} \left( \text{Softmax} \left( \frac{Q_1 K^\top}{\sqrt{d_k}} \right) V, \dots, \text{Softmax} \left( \frac{Q_h K^\top}{\sqrt{d_k}} \right) V \right)$$
最后，结果通过下投影产生最终的注意力输出：
$$\text{Output} = Z \times W_{\text{Down}}$$

MLA与MHA的关键区别。MLA与传统MHA的关键区别在于引入了额外的投射层，特别是上投影和下投影，旨在保持数学上的一致性。此外，MLA采用了MQA机制【索引40，Fast transformer decoding: One write-head is all you need. CoRR 2019】，其中所有Q头共享一个单一的KV缓存头。这种方法增加了计算强度，同时减少了与KV缓存相关的内存访问。此外，MLA中的头维度对应于kv_lora_rank的值，该值通常比其他模型中的要大。例如，在DeepSeek-V2-Lite中，它是512，而在Llama2-7B中，头维度是128。

图14: MLA计算概述：原始（左）和带权重吸收优化（右）。黑色点状部分代表新生成的潜在表示，它将被缓存以计算注意力的键和值。
图14: MLA计算概述：原始（左）和带权重吸收优化（右）。黑色点状部分代表新生成的潜在表示，它将被缓存以计算注意力的键和值。

融合的MLA数据流。根据我们以集群为中心的数据流设计原则，融合的MLA数据流如图15所示。该数据流在片上执行整个MLA计算，消除了任何中间的片外内存流量。详细算法见算法4。此数据流跨注意力头并行化，每个头分配给一个由N = $2^k$个线程块（k ≤ 4）组成的集群。集群内的每个线程块被分配一个秩b ∈ [0, N-1]。在每个集群内，线程块协作划分QKV投射的头维度、上投影和下投影模块的kv_lora_rank维度、注意力模块的KV缓存token维度以及输出投射的输出维度。每个线程块处理完整的输入隐藏状态，并在输出投射后计算其对应的输出瓦片$O_b$。在此算法中，B表示批量大小，D表示输入隐藏维度，H表示总头维度。变量h、l、s和d分别指每个线程块的头维度、kv_lora_rank、序列长度和输出维度的划分大小。为简单起见，我们省略了图14中显示的rope_dim。

图15: 以集群为中心的融合MLA数据流图。
图15: 以集群为中心的融合MLA数据流图。

MLA数据流的DSMEM流量分析。我们还估计了算法4中使用的数据流所产生的总DSMEM流量，该数据流涉及三个ClusterGather操作和三个ClusterReduce操作。ClusterGather操作的DSMEM流量为：

$$Traffic_{Gather}(h, N) + 2 \times Traffic_{Gather}(l, N)$$


而ClusterReduce操作的流量为：

$$Traffic_{Reduce}(l, N) + Traffic_{Reduce}(H, N)$$
我们忽略了softmax统计量引入的流量，因为它只涉及两个浮点值，与张量级别的数据移动相比可以忽略不计。

B.2 SplitHead数据流

替代数据流设计。根据我们以集群为中心的数据流设计原则和利用集群级通信原语，我们实现了主论文中描述的融合QKV投射、注意力和输出投射数据流，其中中间数据存储在片上共享内存中。此外，我们设计了一种名为SplitHead数据流的替代方案，它将中间数据存储在更快的片上寄存器内存中。如算法5所示，该数据流在头数上是并行的。每个头对应一个由N = $2^k$个线程块（k ≤ 4）组成的集群，每个块被分配一个秩b ∈ [0, N-1]。在每个集群内，线程块仅划分QKV投射、注意力和输出投射中的头维度。对于整个数据流，每个线程块处理完整的输入隐藏状态，并在输出投射后计算对应的部分输出$O_b$。在此算法中，B表示批量大小，D表示输入隐藏维度，H表示总头维度，S表示KV缓存序列长度，h和d分别表示每个线程块划分的头维度和输出维度的大小。在此设计中，我们需要归约$Q \times K^T$的结果，其形状为序列长度×批量大小。每个线程块持有分段的注意力输出并计算部分输出投射，然后必须将其归约并使用atomicAdd写入全局内存。

图16: 以集群为中心的融合QKV投射、注意力和输出投射数据流图 (SplitHead变体)。
图16: 以集群为中心的融合QKV投射、注意力和输出投射数据流图 (SplitHead变体)。

SplitHead数据流的流量分析与比较。如图16所示，此数据流使得像$Q_b$、$K_b$和$V_b$这样的中间结果可以存储在更快的寄存器内存中，而不是共享内存。然而，此数据流产生的总DSMEM流量为：

$$Traffic_{Total} = Traffic_{Reduce}(S, N) + Traffic_{Reduce}(D, N)$$


这高于在注意力模块中沿token维度划分KV缓存的数据流（在主论文和B.1节中提出）。SplitToken数据流的总DSMEM流量为：

$$Traffic_{Total} = Traffic_{Reduce}(H, N) + Traffic_{Gather}(3h, N)$$
这明显更低，因为它主要依赖于头维度H或融合MLA数据流中的kv_lora_rank l，这两者都远小于在SplitHead数据流中主导DSMEM流量的序列长度S。增加的通信开销超过了使用寄存器内存的好处。正如我们稍后展示的实验结果所示，SplitHead数据流导致了更高的延迟。

C 额外的实验

C.1 多批次评估和加速分析

多批次性能评估。我们使用16的批量大小进行了额外实验，其他实验设置与主论文中描述的保持一致。TPOT和核心模块的延迟结果分别如图17和图18所示。平均而言，在Llama2-7B上，ClusterFusion（集群大小为4）相对于SGLang、vLLM、TensorRT-LLM和MLC-LLM的加速比分别为1.11倍、1.09倍、1.12倍和1.32倍。对于DeepSeekV2-Lite，在相同条件下，平均加速比分别为1.15倍、1.14倍、1.07倍和1.84倍。在核心模块（包括融合的QKV投射、注意力和输出投射）方面，ClusterFusion在Llama2-7B上相对于SGLang、vLLM、TensorRT-LLM和MLC-LLM的平均加速比分别为1.14倍、1.12倍、1.2倍和1.41倍。同样，对于DeepSeek-V2-Lite，相应的加速比为1.19倍、1.18倍、1.14倍和2.04倍。

图17: 在H100上Llama2-7B（左）和DeepSeek-V2-Lite（右）的TPOT（批量大小为16）。

图17: 在H100上Llama2-7B（左）和DeepSeek-V2-Lite（右）的TPOT（批量大小为16）。
图18: 在H100上Llama2-7B（左）和DeepSeek-V2-Lite（右）核心模块的延迟（批量大小为16）。
图18: 在H100上Llama2-7B（左）和DeepSeek-V2-Lite（右）核心模块的延迟（批量大小为16）。

多批次场景下的加速分析。我们确定了两个主要因素促成了ClusterFusion相对于现有框架（即使有CUDA Graph优化）的性能优势：减少的全局内存传输量和显著降低的GPU内核启动开销【索引54，Mononn: Enabling a new monolithic optimization space for neural network inference tasks on modern gpu-centric architectures. OSDI 2024】【索引47，Optimizing deep learning inference via global analysis and tensor expressions. ASPLOS 2024】。为了量化这些好处，我们使用NVIDIA Nsight Systems【索引32，Nvidia nsight system. https://developer.nvidia.com/nsight-systems】 和Nsight Compute【索引31，Nvidia nsight compute. https://developer.nvidia.com/nsight-compute】在批量大小为16的情况下对内存流量和内核启动开销进行了剖析。结果如图19所示。观察到的性能增益主要归因于ClusterFusion将QKV投射、注意力和输出投射完全在片上执行，从而最小化了中间内存流量。此外，即使与已经 用CUDA Graph优化的基线相比，ClusterFusion在端到端场景中也将内核启动开销减少了近一个数量级。然而，在多批次场景中，减少全局内存流量的影响有限，因为KV缓存和模型权重仍然主导内存使用，而中间内存占用仍然很小。此外，随着批量大小的增加，整体计算强度显著提高，导致与主论文中提出的单批次结果相比，加速比较小。

图19: 全局内存数据传输大小（左）和GPU内核启动开销（右）的比较（批量大小为16）。
图19: 全局内存数据传输大小（左）和GPU内核启动开销（右）的比较（批量大小为16）。

C.2 SplitHead数据流评估

SplitHead vs. SplitToken性能比较。我们还实现了B.2节中描述的SplitHead数据流，并在图20中展示了SplitToken、SplitHead数据流与两个代表性基线的性能比较。当序列长度较短时，延迟差异很小，因为中间数据可以存储在寄存器内存中，这比SplitToken提高了效率，并且与SplitToken数据流相比，DSMEM流量的差距仍然很小。然而，随着序列长度的增加，SplitHead数据流的DSMEM流量显著大于SplitToken，导致延迟增加。

设计选择的合理性。从算子融合效率的角度来看，只要数据大小足够小以完全容纳在寄存器中，SplitHead数据流就能实现将中间数据存储在高速寄存器内存中的融合。然而，尽管有此好处，SplitHead数据流产生了显著更高的DSMEM流量，尤其是在序列长度增加时。这种增加的流量成为主要瓶颈，并导致整体性能变差。因此，我们以集群为中心的数据流设计不仅考虑了融合效率，还考虑了不同数据流的内存流量，旨在实现更好的整体性能。

图20: SplitToken和SplitHead数据流的延迟比较。
图20: SplitToken和SplitHead数据流的延迟比较。