CHARM+: Cross-Hardware Attention with Re-Merge Multistream Mechanism

AI summary

Overview

Research area: Distributed systems for large language models — specifically sequence parallelism and distributed attention (the Ring-Attention family) for long-context training and inference.

Technical level: Advanced. The paper assumes familiarity with FlashAttention-style tiling, ring collectives, NCCL/RDMA bandwidth behaviour, NUMA and node topologies, and roofline-style compute-to-communication reasoning.

Scope (one sentence): The paper redesigns the communication schedule of Ring-Attention into a logarithmic merge with parallel query/key-value streams and a cross-hardware decoupling scheme, then benchmarks forward and backward passes on L20 (PCIe), H20 and H100 (NVLink) systems.

What This Paper Is About

Long-context LLMs partition the sequence across many GPUs, and Ring-Attention is the mainstream way to do this: each device passes its key/value blocks around a ring and computes attention against them as they arrive. Two problems follow. Each communication must wait for the previous one to finish, so computation and communication only partially overlap and GPUs sit idle. And on cross-hardware links (cross-NUMA or cross-node), the bandwidth is much lower than intra-hardware bandwidth, so frequent synchronisation adds further stalls. The goal of CHARM+ is to shorten that dependency chain and hide cross-hardware communication entirely behind intra-hardware work.

Key Contributions

1. A merge-based communication mechanism. Rather than passing key/value blocks around a ring, the authors reorganise key/value exchange into logarithmic merge levels. At level i, devices exchange with a partner at stride 2^(i−1). This reduces the non-overlapping critical path from O(n) in Ring-Attention to O(log n), and a sequence parallelism of 8 maps to three merge levels.

2. Re-Merge Multistream Mechanism (RM+). Inside each merge level, query communication is run in parallel with key/value communication. Because query communication does not depend on key/value communication, the two streams overlap, and the computation that would otherwise trail the final merge level is advanced into earlier levels. The authors model this as the per-level communication cost being bounded by the maximum of the two stream times instead of their sum.

3. Cross-Hardware Query Parallelism. Query communication across NUMA domains or nodes is interleaved into the intra-node pipeline, so that only intra-hardware key/value and query traffic must synchronise. The stated effect is that after all computation finishes, only the final output communication remains unhidden, and that this advantage grows as the number of nodes increases.

4. Full forward and backward implementation with broad evaluation. Both directions are implemented (backward shares the forward logic), with experiments across L20, H20 and H100 systems, ablation of RM+ and Cross-Hardware Query Parallelism, memory-overhead measurement, numerical-correctness checks against FlashAttention, and end-to-end Megatron-LM integration.

Main Findings

• Speedups over Ring-Attention (abstract averages): 1.3× forward and 1.7× backward on a dual-NUMA PCIe architecture, and 2.8× forward and 3.9× backward on a dual-node NVLink architecture.

• Speedups over Ring-Attention (maximum, as stated in the contributions): 1.4× forward and 1.9× backward on PCIe; 3.2× forward and 4.3× backward on NVLink.

• RM+ alone already beats Ring-Attention on NVLink: on a single 8-GPU H20 node, RM+ reaches a maximum 1.9× speedup over Ring-Attention. Example forward times in the ablation: at sequence length 7168 on 4× H20, Ring is 38.62 ms versus RM+ at 33.20 ms; on 8× H20, Ring is 26.79 ms versus RM+ at 13.84 ms.

• RM+ degrades on PCIe and in cross-hardware settings: on 8× L20 at sequence length 7168, Ring takes 69.85 ms while RM+ takes 98.49 ms. The cause given is PCIe link contention when multiple streams communicate concurrently, plus bandwidth imbalance between cross- and intra-hardware links.

• CHARM+ recovers and wins in cross-hardware configurations: on the same 8× L20 cross-NUMA setup at sequence length 7168, CHARM+ takes 52.40 ms versus Ring's 69.85 ms; on 16× H20 cross-node at 7168, CHARM+ takes 28.91 ms versus Ring's 84.58 ms and RM+'s 98.98 ms. The ablation reports up to 1.9× speedup in cross-NUMA and 3.5× in cross-node configurations.

• Single-node NVLink gains are smaller than cross-node gains: on 8× H20 the speedup is slightly lower because intra-node bandwidth is very high (297.2 GB/s reported), raising the compute-to-communication ratio so computation outweighs communication.

• H100 results: across 2×8 and 4×8 configurations, average speedup of about 2× in the forward and 3× in the backward versus Ring-Attention. Maximum reported speedups are 2.28× forward and 3.69× backward on 16 H100 GPUs, and 2.69× forward and 4.21× backward on 32 H100 GPUs.

• End-to-end training on L20 (16 GPUs) with GPT-350M and GPT-1.3B: consistent speedups of 1.72× to 2.08× across model sizes and configurations, with nearly identical memory overhead. Example: GPT-350M, batch 16, sequence 2048, latency 10082.6 ms for Ring versus 4847.0 ms for CHARM+ (2.08×), memory 11,372 MB versus 11,518 MB.

• End-to-end inside Megatron-LM on 16 H20/H100 GPUs with GPT3-350M: maximum speedups of 1.68× on H20 and 2.17× on H100. The authors note Megatron-LM contains ring optimisations CHARM+ does not include and that non-attention linear layers dominate end-to-end time.

• Memory overhead is modest and constant-order: for example, batch 16, sequence 1024, 32 heads costs 2,179 MB under Ring-Attention versus 2,215 MB under CHARM+; batch 16, sequence 1024, 64 heads costs 3,701 MB versus 3,881 MB. The extra buffers are O(1), so overhead does not scale with GPU count.

• Numerical behaviour matches Ring-Attention's order of magnitude: differences versus FlashAttention in output, log-sum-exp, and Q/K/V gradients are of the same order as Ring-Attention's, indicating no substantial additional numerical error.

• Larger scale helps CHARM+: in the 2×4 and 2×8 H20 multi-node experiments, performance grows as the number of GPUs increases, which the authors attribute to cross-node communication being hidden by the intra-node pipeline.

Methodology in Plain English

The starting point is an accounting of wasted time. The authors define the "unhidden communication time" as the sum, over the critical path, of each communication step's time minus whatever computation could run alongside it, clipped at zero. Ring-Attention yields a chain of n−1 such steps. The whole design effort is aimed at shrinking that chain.

The first move is borrowed from merge sort. Instead of a single ring along which key/value blocks travel one hop at a time, devices pair up and exchange at distances of 1, then 2, then 4, and so on. Key/value data therefore aggregates in log₂(n) rounds rather than n−1 hops. Each device keeps computing against the blocks it just received while the next round of exchange proceeds.

The second move exploits an asymmetry the ring ignores: queries do not depend on key/value blocks, so query and key/value traffic can travel on separate streams at the same time. This makes each merge level's communication cost the larger of the two streams rather than their sum, and it lets the authors pull some of the trailing computation forward into earlier levels. The trade-off is that on PCIe, where concurrent transfers contend for the same links, this parallelism can make a single communication slower than in the ring — a cost that shows up in the reported measurements.

The third move addresses topology. In a multi-NUMA or multi-node system, a key/value hop crosses both a fast intra-hardware link and a slow cross-hardware link, and the fast one idles while waiting for the slow one. CHARM+ instead routes query communication across the hardware boundary into the intra-node pipeline, so each pipeline works on queries coming from other pipelines using strictly local key/value blocks. Local traffic then proceeds without waiting on the slow link, and cross-hardware traffic is hidden underneath it.

Experiments sweep batch size, sequence length and hidden dimension, with sequence lengths from 1K to 100K (up to 200K in some scenarios) and four hidden-dimension configurations formed from 12 or 20 combined with 128 or 192, keeping batch × sequence at or below 500K because of memory limits. Each configuration runs 20 warm-up iterations followed by 100 measured iterations. Communication and computation are summarised with a compute-to-communication ratio (CCR) following TASP, and timing is collected through CUDA events. Baselines are Ring-Attention and TASP plus their zigzag variants; TASP's backward is not implemented in the public code, so backward comparisons are forward-only against Ring-Attention.

Why This Matters

Distributed attention is a bottleneck for long-context training and inference, and it is a bottleneck that gets worse as context windows and cluster sizes grow. CHARM+ shows that the ring schedule itself, not just the kernel arithmetic, is a design variable worth rewriting: an O(log n) communication path with better overlap can outperform a well-tuned ring even under an end-to-end training loop that includes many non-attention layers.

Real-world applications that depend on these workloads:

• Long-document training and continued pretraining, where processing long documents is stated as a requirement for model capability.
• Cross-document understanding and retrieval-augmented inference, where many documents must be attended to jointly.
• Long program comprehension, where source files exceed typical context windows.
• Multi-turn interactive dialogue, where conversation history accumulates over many turns.

Industry relevance: the gains are largest exactly where clusters are physically distributed — cross-node and cross-NUMA — which is the regime of multi-node training jobs. The method is integrated into Megatron-LM's context parallelism and evaluated on rented H100 nodes, and it works on both PCIe-only (L20) and NVLink (H20, H100) systems, so it is deployable across mixed generations of hardware rather than requiring a uniform high-bandwidth fabric.

Future Directions

• Mitigating PCIe link contention. RM+'s parallel streams hurt on PCIe because concurrent transfers contend for shared links; an adaptive scheduler that throttles stream parallelism based on measured link contention (the paper observes contention is negligible only when both batch size and sequence length are small) is a natural extension the results point to.

• Closing the gap between operator and end-to-end gains. Reported operator-level speedups (up to 4.21× backward on 32 H100) far exceed end-to-end speedups (1.68×–2.17× in Megatron-LM), because non-attention layers dominate. Combining CHARM+ with attention-layer tuning or with optimisations for the remaining layers is an open direction.

• Beyond the tested grid. Sequence lengths up to 200K appear "in some scenarios," but the reported tables stop at 53,760 in the memory study and 31,744 on H100. Extending validation to the full 200K regime and to clusters larger than 32 GPUs would test whether the single-exposed-communication claim holds at higher node counts.

• Portability to other parallelisms and hardware. The cross-hardware decoupling is described in terms of NUMA domains and nodes, but it is not evaluated against other context-parallel layouts (for example, combining with the 2D or tiled assignments from Attention-2D and Mesh-Attention) or on non-NVIDIA interconnects, which the paper does not report.

Target Audience

Systems and performance researchers working on distributed training and inference for LLMs; engineers responsible for context parallelism and long-context serving in production clusters; and practitioners integrating attention kernels into frameworks such as Megatron-LM who need to understand where operator-level speedups survive into end-to-end iteration time and where topology constrains them. Readers need a working knowledge of collective communication and FlashAttention internals to get full value; the paper reports no ML datasets or accuracy benchmarks beyond numerical-equivalence checks.

Authors’ abstract

As sequence lengths continue to grow, distributed attention based on sequence parallel (e.g., Ring-Attention) has become a mainstream solution for large language models (LLMs). Nevertheless, inherent communication dependencies limit computation-communication overlap, reducing GPU utilization. Moreover, in cross-hardware (NUMA or node) architectures, due to imbalanced communication bandwidth, frequent communication synchronization will limit overall performance. Based on these insights, we propose CHARM+, Cross-Hardware Attention with Re-Merge Multistream Mechanism. We reorganize sequential communications into several levels and achieve intra-level parallelism and computation overlap. Furthermore, we decouple cross-hardware communication from intra-hardware communication, eliminating synchronization overhead caused by bandwidth imbalance. Forward and backward algorithms were evaluated on multiple hardware architectures and consistently outperformed state-of-the-art methods. Compared to Ring-Attention, our forward and backward algorithms achieve average speedups of 1.3x and 1.7x on a dual-NUMA PCIe architecture, and 2.8x and 2.9x on a dual-node NVLink architecture, respectively.