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

Thank you for your thoughtful feedback and suggestions on our paper. We would like to provide more clarification.

To answer your concerns:

1. About the term "Cluster". Sorry about the confusion, we will use the term "SM Cluster" to avoid this confusion. Indeed, our primitives operate at the block/SM level. We chose the term “Cluster” to align with NVIDIA’s terminology on Hopper GPUs[1,2], where a cluster refers to a group of thread blocks with shared memory connectivity.
2. Different communication pattern. Thanks for your point. We have carefully considered both the tree pattern and the ring pattern (ScatterReduce + AllGather), and we chose the tree pattern because it requires fewer communication rounds, which leads to lower latency when the data size is relatively small. We compared the two patterns both theoretically and empirically under cluster size = $N$ and data size = $K$. In theory, the ring pattern requires $(N-1)$ rounds for ScatterReduce and another $(N-1)$ rounds for AllGather, totaling $2(N - 1)$ communication rounds, with each round transferring $\frac{K}{N}$ data per node. In contrast, the tree pattern only needs $\log_2 N$ rounds, with each round transferring $K$ data per node, as described in our paper. The following table summarizes the theoretical communication characteristics for both patterns:

For example, when $N=2$ and $N=4$:

Although the tree pattern transfers more total data, its fewer communication rounds are more favorable for latency because the data size $K$ is small (64–128 bytes) in the ClusterFusion dataflow and DSMEM bandwidth is not fully utilized. This is also reflected in our empirical results, where we report the kernel latency(us) of the core ClusterFusion kernel:

We plan to include a more detailed discussion about different communication pattern in the final version. 3. Definition of the bandwidth between SMs within cluster. The bandwidth shown in Section 2.3 and Figure 5 refers to peer-to-peer bandwidth. As cluster size increases, peer-to-peer bandwidth decreases. Based on the patent and paper referenced in that section, the interconnect topology for DSMEM in Hopper GPUs is likely a crossbar or crossbar-like architecture, which allows any SM to communicate with any other within a cluster. While a crossbar provides full connectivity, it does not provide dedicated bandwidth for all possible communication pairs simultaneously. Instead, the internal paths of the crossbar must be shared among all active transfers, and contention arises as more thread blocks in a cluster attempt concurrent communication.

4. The reason for the decreased active SMs. To the best knowledge of us based on public architectural information[1,2,3], the number of active SMs may decrease with larger cluster sizes due to hardware-level SM allocation constraints. Whenever the cluster size exceeds 2, the hardware grouping may lead to SM underutilization due to yield-related SM inactivation and cluster alignment issues. On H100 SXM GPUs, the 132 SMs are not evenly distributed: three GPCs[2] have 18 SMs, four have 16, and the last one contains 8 + 6 SMs, where the last 6 are special. For example, when using Cluster Size = 4, this grouping cannot fully utilize all SMs: in the first three GPCs, 2 SMs per GPC are left out (3 × 2 = 6), and in the last GPC, all 6 special SMs are likely excluded, totaling 12 unused SMs. This leaves only 120 active SMs which aligns with our observed drop in active SMs during profiling.
5. About kernel launch overhead in Figure 7. Thanks for pointing this out. Actually, Figure 7 shows the total kernel launch overhead of the end-to-end decoding. Our fused kernel reduces the overall kernel launch overhead across all layers. And our kernel-level launch overhead is indeed one order of magnitude lower than the CUDA Graph launch overhead used in existing frameworks. For example, in the Llama2-7B model, the end-to-end kernel launch overhead drops from 1.057 ms in the baseline to 0.169 ms with ClusterFusion.
6. The performance of multi-batch. For batch sizes >1, we observe an average 1.2× end-to-end speedup in the Appendix provided in the supplementary material. As batch size grows, the workload becomes more compute-intensive and less memory-bound, reducing the optimization space for memory traffic. Although the off-chip memory traffic increases in higher batch size, the proportion of activation memory traffic and sync decreases due to the significantly increased compute intensity. Nevertheless, ClusterFusion still achieves meaningful gains, demonstrating its robustness across diverse inference scenarios.
7. The reliance on shared memory buffer. In our current implementation, the algorithm relies on either reusing existing shared memory tiles (e.g., weight tiles) or allocating a small additional buffer. This is necessary because our reduction is single-phase: the data being transmitted and the data being reduced in each iteration reside in the same memory region, requiring intermediate buffering to avoid conflicts. However, when shared memory becomes insufficient, we can switch to buffer-free variants such as ScatterReduce. In this approach, each thread block writes its data to non-overlapping destinations, allowing communication and reduction to proceed simultaneously without intermediate buffering.
8. Comparison with other concurrent work. Thank you for pointing out two recent, concurrent work: Kernel Looping[4] and MegaKernel[5](publicly available after our submission). Kernel Looping is designed for a dataflow architecture accelerator—specifically the SambaNova SN40L Reconfigurable Dataflow Unit (RDU). It leverages the repetitive nature of decoder layers in LLM to fuse different layers, reducing multiple kernel launches to a single launch. This approach serves a role similar to CUDA Graphs in the GPU. Our method also significantly reduces kernel launch overhead on GPUs, and we empirically outperform frameworks that already leverage CUDA Graphs. The motivation of MegaKernel is similar to ours in that both aim to reduce memory bubble. Their approach is orthogonal to our work. They focus on fine-grained shared memory reuse and overlapping activation storage with weight loading. Despite this, MegaKernel still relies on global memory for intermediate result synchronization, whereas our work introduces cluster-level primitives to enable fast, on-chip synchronization. We will incorporate a detailed discussion of these recent works in the related work of the final paper.

[1] NVIDIA CUDA C++ Programming Guide.

[2] NVIDIA H100 Tensor Core GPU Architecture.

[3] Benchmarking Thread Block Cluster, HPEC 2024.

[4] Kernel Looping: Eliminating Synchronization Boundaries for Peak Inference Performance, arxiv.

[5] Look Ma, No Bubbles! Designing a Low-Latency Megakernel for Llama-1B, Hazy Research.

Thanks for your further feedback, and we appreciate the opportunity to clarify the kernel launch details.

1. Regarding the exact number of kernels. Actually, the four modules shown in Figure 7 of our paper do not correspond to exactly four kernels in the baseline. Although existing SOTA frameworks already leverage torch.compile optimizations, there remain some fragmented kernels, as detailed in the tables below:

For Llama:

For DeepSeek:

This detailed kernel count information is also available in our supplementary materials. We will consider clarifying this in the experimental section.

2. Regarding the precise speedup. Thank you for pointing this out. Actually, our kernel launch latency reduction from 1.057ms to 0.169ms corresponds to about a 6.25× improvement, which is not a full order of magnitude. We originally used "nearly one order of magnitude lower" to emphasize the substantial decrease, but we apologize for the inaccurate phrasing that could be interpreted as a 10× improvement. We will revise this to a more precise description. And regarding Figure 12, the exact latency and corresponding speedup are as follows:

For Llama:

For DeepSeek:

Thank you for your thoughtful feedback and suggestions on our paper. We would like to provide more clarification.

To answer your concerns:

1. About novelty and the low-level PTX. While reduce and gather are common primitives in distributed system, they are rarely considered in chip-scale with interconnect. This enable new optimization scope and opportunities, and our novelty mainly lies in the exploration of on-chip dataflow via interconnect. On Hopper GPUs, blocks within a cluster are fully connected via DSMEM, as illustrated in Figure 4 and Figure 6. And as shown in Section 2.3 and Figure 5, DSMEM performance is non-trivial and not automatically efficient—naïve use often underperforms. Our primitives apply structured, tree-pattern, and collective scheduling to fully exploit DSMEM, enabling practical performance gains in real LLM workloads. And we have provided our codes in the supplementary material. We use the following set of PTX instructions in our implementation of ClusterReduce and ClusterGather:

All of these instructions are necessary for implementing our ClusterReduce and ClusterGather primitives. While these instructions are documented in PTX ISA references [1,2], they are low-level, semantically fragmented, and designed only for fine-grained behaviors (e.g., address mapping or signaling). Their correct composition into usable patterns is non-trivial. Many of them take complicated parameters (e.g., barrier state tokens, data sizes, memory space specifiers) and require strict ordering to avoid deadlocks or incorrect synchronization. Our contribution is to encapsulate these instructions into robust, reusable cluster-level primitives. This abstraction significantly reduces development cost and makes it practical to integrate DSMEM communication into LLM inference frameworks.

2. The necessity of the proposed primitives for fusion. Thanks for your insightful suggestion. Our cluster-level primitives are necessary to scale fusion across realistic LLM workloads with high parallelism. Indeed, it is possible to let a single thread block handle the tasks of both TB0 and TB1 in Figure 7 (equivalent to setting cluster size = 1), which eliminates the need for inter-block communication. We have implemented and evaluated this variant. The evaluation result of ClusterFusion core kernel latency(us) under different cluster size are:

As above results show, using only one thread block significantly degrades performance. This is because a single block often lacks sufficient parallelism. In practice, we need to split the workload across more thread blocks (e.g., along the head or token dimension to fully utilize the hardware). However, such splits introduce inter-block data dependencies that require reduction or aggregation. This is where our on-chip ClusterReduce and ClusterGather primitives become critical: they enable efficient intra-cluster communication without falling back to global memory.

[1] Parallel Thread Execution ISA Version 8.8, NVIDIA.

[2] NVIDIA H100 Tensor Core GPU Architecture.

Raising my score from 4 to 5 since the response addressed majority of my concerns.

Thank you for your thoughtful feedback and suggestions on our paper. We would like to provide more clarification.

To answer your concerns:

1. Generality on other hardwares. The techniques and rationale implemented in ClusterFusion are not limited to NVIDIA Hopper GPUs. Our fundamental goal is to enable holistic dataflow design at the larger scale. As discussed in the background section, existing fragmented design makes it difficult to achieve globally optimal. We construct ClusterReduce and ClusterGather that enable on-chip cooperation between thread blocks. On A100 GPUs, where DSMEM is not available, our on-chip collective primitives are not supported. However, as shown in Section 4.3 and Figure 13, we implement ClusterReduce and ClusterGather using global memory as a fallback, and still achieve competitive performance due to reduced kernel launches and improved holistic dataflow design. Below are the TPOT (ms) of Llama2-7B on A100:

Moreover, DSMEM-like inter-core communication is becoming a mainstream direction in modern architectures[1,2,3,4]. These trends indicate that ClusterFusion’s primitives and dataflow strategies are compatible with emerging hardware. Also, the primitives and holistic dataflow can be mapped to other platforms. On AMD GPUs[1], SMs, shared memory, and DSMEM used in our algorithm can be naturally mapped to compute units, local data share (LDS), and LDS-to-LDS interconnects (e.g., in RDNA 4). On dataflow architectures such as Graphcore IPUs[3] and Cerebras Wafer-Scale Engines[4], our fine-grained task partitioning strategy remains applicable across the many-core compute fabric of these systems. The holistic dataflow design philosophy as well as our inter-core tree-pattern ClusterReduce and ClusterGather can be further explored on these platforms to fully exploit their on-chip interconnects and distributed memory hierarchy. We are actively collaborating with engineers to implement our primitives on other GPU like MI300X/MI308X GPUs, enabling the same holistic, end-to-end dataflow optimization.

2. Integration into frameworks/compilers. We have already integrated our fused kernels into the PyTorch stack via pybind. This is foundation for future integration into real production frameworks such as SGLang, vLLM, and TensorRT-LLM. For the evaluation shown in our paper, we have already use our wrapped PyTorch-compatible interfaces. These interfaces are aligned with existing kernel stack and support drop-in replacement. For framework integration, we need to replace the original kernels used in other frameworks with our ClusterFusion kernel but no CUDA runtime modifications. For CUTLASS integration, we can append our templates of ClusterReduce/ClusterGather primitives into CUTLASS template library. For ML compilers integration, we need to add a cluster pass which automatically group some thread blocks into a thread block cluster according to the data dependency for reduce/gather.
3. About FFN, RMSNorm and RoPE. We have extended our primitives to fuse other modules like FFN, and the corresponding implementation(codes) is included in our supplementary material. In short, FFN fusion is feasible, but the performance gain is limited. Firstly, the FFN kernel must be separated from the Output projection, because the output projection in Attention introduces a all-to-one data dependency: all thread blocks must reduce their output to the same address. This is not feasible to resolve fully on-chip via DSMEM without falling back to global memory, as the limited cluster size. Secondly, we designed a fused FFN kernel where we partition the intermediate size across blocks within a cluster, and split the output dimension across different clusters. ClusterReduce is then used to compute partial sums across blocks on chip, followed by local element-wise SiLU and residual addition. However, due to the large intermediate dimensions in most mainstream LLMs (e.g., 11008 or higher), each thread block must handle a huge reduction payload. As shown in our profiling results (Section 2.3 and Figure 5), the limited DSMEM bandwidth in this case prevents us from gaining significant on-chip communication benefits. However, FFN fusion still helps reduce kernel launch overhead and global memory traffic and achieves comparable latency to existing frameworks:

Regarding RMSNorm and RoPE, these components are omitted from the paper for clarity, but are indeed included in our fused kernels. Since they can be naturally fused using shared memory or registers without requiring ClusterReduce or ClusterGather primitives. For simplicity, we did not illustrate them in the figures, but their full implementations are provided in the supplementary material.

4. The performance of multi-batch and variable length. ClusterFusion delivers speedup across both batch size = 1 and batch size > 1 scenarios. For batch size > 1, we provide empirical results in the Appendix of the supplementary material, where we observe an average end-to-end speedup of 1.2×. As for variable-length prompts, supporting fully dynamic batching requires extending the attention module to handle runtime index computation and layout alignment. This extension is technically straightforward and will not compromise the performance benefits of our fused design, since the core communication and computation patterns remain unchanged. We are actively working on this enhancement extending our implementation to better support these use cases.
5. Tuner usage. We expose cluster_size as a tunable parameter for users and offline auto-tuning. In our current implementation, cluster_size is a template parameter at the kernel level. Based on our profiling (Section 2.3), larger values like 16 are generally discouraged. So, we provide a simple tuner that performs offline benchmarking over a small candidate set (e.g., 2, 4, 8). The tuner selects the best-performing configuration for a given input models.
6. Energy optimizations. Thanks for your suggention about discussion of energy-performance tradeoff. In section 4.2 and Figure 12, we report the reduced off-chip memory traffic dut to ClusterFusion optimizaions. These changes naturally reduce the frequency of global memory accesses and kernel launches, both of which are known to contribute to energy consumption on modern GPUs. We will include a discussion of these implications in the final version.

[1] AMD "RDNA4" Instruction Set Architecture.

[2] NVIDIA Blackwell Architecture Technical Brief.

[2] Scaling Deep Learning Computation over the Inter-core Connected Intelligence Processor with T10, SOSP 2024.

[3] WaferLLM: Large Language Model Inference at Wafer Scale, OSDI 2025.

Thank you for your thoughtful feedback on our paper. We would like to provide more clarification.

To answer your concerns:

1. Cluster Size (k) Selection. Based on our profiling results in Section 2.3 and Figure 5, we observe that using an excessively large cluster size can degrade performance due to reduced effective SM bandwidth and lower active SM. To mitigate this, we recommend choosing a cluster size less than or equal to 8 for most workloads. Additionally, to facilitate data tiling and ensure alignment between data layout and thread-block mapping, we suggest using even cluster sizes such as 2, 4, or 8—resulting in a small but practical space. We also expose cluster_size as a tunable parameter in our implementation, since its optimal value is tightly coupled with several factors such as tile size and kernel-level parallelism. In practice, this introduces a trade-off: when the cluster count is large and each cluster handles a small amount of data, a smaller cluster size leads to sufficient parallelism and more favorable tiling. On the other hand, when each cluster’s workload is relatively large, a larger cluster size is needed to efficiently amortize inter-block communication. This behavior is reflected in our results in Section 4.1 and Figure 11, where smaller cluster sizes (2 or 4) are preferred for many-head attention, while larger ones become more favorable as the head count decreases or head dimension grows.
2. Kernels used in each framework. In our experiments, we adopt the default operators used by each framework under their officially recommended configuration[4,5,6,7], with torch.compile enabled to utilize optimized Triton-based kernels for decoding. The following table summarizes the actual attention and GEMM kernels used in each baseline framework:

All of these kernels represent SOTA implementations in their respective frameworks. Our method surpasses their performance, demonstrating both the generality and effectiveness of our approach. We will include detailed information about these kernel in the final version of the paper.

3. About FFN. We have extended our primitives to fuse other modules like FFN, and the corresponding implementation(codes) is included in our supplementary material. In short, FFN fusion is feasible, but the performance gain is limited. Firstly, the FFN kernel must be separated from the Output projection, because the output projection in Attention introduces a all-to-one data dependency: all thread blocks must reduce their output to the same address. This is not feasible to resolve fully on-chip via DSMEM without falling back to global memory, as the limited cluster size. Secondly, we designed a fused FFN kernel where we partition the intermediate size across blocks within a cluster, and split the output dimension across different clusters. ClusterReduce is then used to compute partial sums across blocks on chip, followed by local element-wise SiLU and residual addition. However, due to the large intermediate dimensions in most mainstream LLMs (e.g., 11008 or higher), each thread block must handle a huge reduction payload. As shown in our profiling results (Section 2.3 and Figure 5), the limited DSMEM bandwidth in this case prevents us from gaining significant on-chip communication benefits. However, FFN fusion still helps reduce kernel launch overhead and global memory traffic and achieves comparable latency to existing frameworks:

4. Evaluation on longer context and larger models. First, regarding long-context capability, we have extended our evaluation of Llama2-7B to >100K context length. In this setting, the TPOT(ms) of our fused kernel against baselines is as follows:

Now, we have already evaluated a range of model sizes under common single-GPU inference settings. As for scalability to larger and multi-modal models with long context that require distributed inference, we have confirmed that our kernel maintain measurable per-device speedups even when integrated into a multi-GPU inference pipeline. As an initial step, we tested a distributed setup for DeepSeek-R1-671B with TP=16, where our single-GPU fusion kernels still yield measurable gains in the per-device computation stage. In this setting, the TPOT(ms) of our fused kernel against baselines is as follows:

While these results are currently limited to per-device computation time, they demonstrate the effectiveness of our approach for integration into multi-GPU/node inference. When scaling to larger models such as DeepSeek-R1, the speedup diminishes slightly (e.g., from 1.3× on LLaMA2-7B to ~1.1× on R1). This is because operator fusion has less room to eliminate overheads in larger models. Our future work focus on the extensions that co-optimize communication and memory access which will further improve performance. To address cross-device scalability, our implementation allows seamless insertion of collective communication operations (e.g., ring-based all-reduce, gather) without disrupting kernel scheduling. This enables future support for end-to-end fusion across both intra- and inter-GPU workloads with minimal integration effort.

5. Generality on other hardwares. The techniques and rationale implemented in ClusterFusion are not limited to NVIDIA Hopper GPUs. Our fundamental goal is to enable holistic dataflow design at the larger scale. As discussed in the background section, existing fragmented design makes it difficult to achieve globally optimal. We construct ClusterReduce and ClusterGather that enable on-chip cooperation between thread blocks. The primitives can be mapped to other platforms. On AMD GPUs, SMs, shared memory, and DSMEM used in our algorithm can be naturally mapped to compute units, local data share (LDS), and LDS-to-LDS interconnects (e.g., in RDNA 4[1]). On dataflow architectures such as Graphcore IPUs [2] and Cerebras Wafer-Scale Engines [3], our fine-grained task partitioning strategy remains applicable across the many-core compute fabric of these systems. The holistic dataflow design philosophy as well as our inter-core tree-pattern ClusterReduce and ClusterGather can be further explored on these platforms to fully exploit their on-chip interconnects and distributed memory hierarchy. We are actively collaborating with engineers to implement our primitives on other GPU like MI300X/MI308X GPUs, enabling the same holistic, end-to-end dataflow optimization.
6. About novelty. While reduce and gather are common primitives in distributed system, they are rarely considered in chip-scale with interconnect. This enable new optimization scope and opportunities, and our novelty mainly lies in the exploration of on-chip dataflow via interconnect. On Hopper GPUs, blocks within a cluster are fully connected via DSMEM, as illustrated in Figure 4 and Figure 6. And as shown in Section 2.3 and Figure 5, DSMEM performance is non-trivial and not automatically efficient—naïve use often underperforms. Our primitives apply structured, tree-pattern, and collective scheduling to fully exploit DSMEM, enabling practical performance gains in real LLM workloads.
7. Clarity in Presentation of Limitations. As mentioned in 3. About FFN., our on-chip fusion is "split" when there are large-scale data dependencies such as all-to-one reductions that cannot be fully handled within DSMEM due to capacity constraints. In these cases, our implementation falls back to global memory synchronization. However, such cases are rare in typical LLM inference workloads and the same data dependency patterns would also cause global memory traffic in baseline frameworks. Therefore, our method still maintains a performance advantage. We will make this behavior more explicit in the final version.
8. Depth of Algorithmic Detail. We appreciate the reviewer’s suggestion. While our current pseudocode omits certain corner cases, our implementation can readily accommodate such scenarios. For example, load balancing can be enhanced by adaptive task-to-SM mapping in the start of the kernel, and DSMEM access patterns can be tuned in our primitive template to avoid bank conflicts. These extensions can be integrated with minimal change, thanks to the modular and general nature of our collective abstraction. As for formal performance model , We will include a roofline model that guide applicability across different workload sizes in the final version to improve clarity.

[1] AMD "RDNA4" Instruction Set Architecture.

[2] Scaling Deep Learning Computation over the Inter-core Connected Intelligence Processor with T10, SOSP 2024.

[3] WaferLLM: Large Language Model Inference at Wafer Scale, OSDI 2025.

[4] SGLang Blog, torchcompile-latency-optimizations, 2024-09-04.

[5] vLLM Blog, torch.compile and Piecewise CUDA Graphs, 2025/01/27.

[6] TensorRT-LLM Doc, Github.

[7] MLC-LLM Doc, Github.