CHARM+가 분산 어텐션의 통신 의존성을 로그 단계 병합으로 줄인다

무엇인가

이 논문은 긴 문맥 LLM의 학습·추론에서 분산 어텐션, 특히 Ring-Attention 계열 시퀀스 병렬의 통신 병목을 다룬다. Ring-Attention은 KV 블록을 링 토폴로지로 순차 전달하므로 각 통신이 이전 통신 완료에 의존하고, 한 번의 통신이 한 번의 연산과만 겹칠 수 있어 GPU가 노는 시간이 생긴다. 또한 NUMA나 노드 경계를 넘는 교차 하드웨어 통신은 대역폭이 낮아 동기화 비용이 커진다. 저자들은 가려지지 않은 통신 시간 O를 T_comm에서 T_comp를 뺀 합으로 모델링하고, Ring-Attention의 임계 경로가 O(n)인 점을 핵심 한계로 지적한다.

어떻게 동작하나

제안 방법의 첫 축은 KV 교환을 링이 아니라 병합(merge) 단계로 재구성하는 것이다. 시퀀스 병렬 크기가 8이면 세 단계가 생기고, i번째 단계에서 장치 쌍이 2^{i-1} 간격으로 병렬 KV 집계를 수행한다. 단계 사이 의존성은 남지만 전체 임계 경로가 O(log n)으로 줄고, 각 단계에서 이전 단계에 받은 KV로 더 많은 FlashAttention 연산을 처리해 통신과 연산의 겹침 기회를 늘린다. 논문은 단계별 통신 시간 C_i^(1)과 가려지지 않은 시간 O_merge를 수식으로 제시해 이 효과를 설명한다.

무엇과 다른가

두 번째 축은 Re-Merge Multistream Mechanism(RM+)이다. q 통신은 KV 통신과 독립이라는 점을 이용해 q와 KV를 서로 다른 통신 스트림으로 병렬 전송하고, 각 병합 단계에서 통신 시간이 두 스트림 중 큰 값으로 제한되도록 만든다. 또한 각 단계의 마지막 통신을 다음 단계로 미루고 그 사이 연산을 앞당겨, 단계 i에서 2^{i-1}번의 연산이 통신을 가리도록 한다. 이로써 O_RM+가 O_merge보다 짧아질 수 있지만, PCIe에서는 다중 스트림이 링크 경합을 일으켜 단일 통신 지연이 커질 수 있다는 전제가 붙는다.

어떻게 쓰나

세 번째 축은 Cross-Hardware Query Parallelism이다. NUMA나 노드 경계에서 느린 교차 하드웨어 통신과 빠른 노드 내부 통신이 한 KV 통신에 섞이면, 내부 통신이 외부 동기화를 기다리며 놀게 된다. CHARM+는 교차 하드웨어 q 통신을 노드 내부 파이프라인에 끼워 넣고, 각 파이프라인은 로컬 KV로 연산한다. 이렇게 교차 하드웨어 통신과 하드웨어 내부 통신을 분리하면 대역폭 불균형에 따른 동기화가 줄고, 마지막 출력 통신 정도만 가려지지 않은 채 남는다. 노드 수가 늘수록 이 겹침의 이점이 커진다고 주장한다.

전제와 한계

실험은 PCIe 기반 8×L20 듀얼 NUMA, NVLink 기반 8×H20 단일 노드, H20 2노드 8·16 GPU, NVLink 32×H100(2×8, 4×8)에서 수행됐다. 시퀀스 길이는 1K~100K(일부 200K), 히든 차원은 12·20과 128·192 조합, B×S는 메모리 제약으로 500K 이하이며, 20회 워밍업 후 100회 반복해 평균을 냈다. 베이스라인은 Ring-Attention, TASP, Ring-Zigzag, TASP-Zigzag이고, TASP는 백워드 구현이 없어 포워드만 비교했다. 단일 노드에서 CHARM+는 PCIe 링크 경합이 있는 L20에서도 TASP와 Ring-Attention을 앞섰고, TASP는 L20에서 all-to-all 집합 통신이 링크를 모두 점유해 경합이 커지며 특히 성능이 나빴다. H20에서는 노드 내 대역폭이 높아 CCR이 커지면서 속도 향상 폭이 상대적으로 작았다.

멀티노드에서는 교차 노드 대역폭이 노드 내부보다 낮아 Ring-Attention과 TASP의 동기화 대기가 커지는데, CHARM+가 그 영향을 더 크게 줄였다. 논문이 보고한 평균 속도 향상은 Ring-Attention 대비 듀얼 NUMA PCIe에서 포워드 1.3배·백워드 1.7배, 듀얼 노드 NVLink에서 2.8배·2.9배이며, 최대값은 PCIe 1.4배·1.9배, NVLink 3.2배·4.3배로 제시된다. 32×H100에서는 평균 포워드 2배·백워드 3배, 최대 16 GPU에서 2.28배·3.69배, 32 GPU에서 2.69배·4.21배를 보고했다. Megatron-LM에 연산자를 통합한 종단 학습에서는 H20 최대 1.68배, H100 최대 2.17배였고, L20 16 GPU의 350M·1.3B 모델에서는 1.72~2.08배 향상과 거의 같은 메모리 사용을 보였다.

개발자에게 이 논문은 Ring-Attention을 쓰는 긴 문맥 학습·추론 스택에서 토폴로지가 성능을 좌우한다는 점을 구체화한다. NVLink처럼 링크 경합이 적고 노드 간 대역폭 차이가 큰 환경에서는 CHARM+식 로그 단계 병합과 교차 하드웨어 q 분리가 유효할 수 있다. 반면 PCIe 링크 경합이 심하거나 노드 내 대역폭이 매우 높아 연산이 병목이면 다중 스트림의 이득이 줄 수 있으므로, 도입 전에 CCR과 통신 토폴로지, 백워드 포함 여부, Megatron-LM 같은 프레임워크의 기존 링 최적화와의 중복을 확인해야 한다.

저자들이 밝힌 전제와 한계도 분명하다. RM+는 PCIe에서 링크 경합으로 단일 통신 지연이 커져 Ring-Attention보다 나빠질 수 있고, H20처럼 노드 내 대역폭이 높아 CCR이 커지면 통신-연산 겹침의 이득이 제한된다. TASP는 백워드가 없어 포워드만 비교했고, 종단 학습에서는 attention이 아닌 선형 층이 지배해 이득이 줄어든다. 메모리는 O(1) 추가 버퍼로 GPU 수에 따라 늘지 않고 증가폭이 작지만 완전히 무료는 아니며, 수치 정확도는 FlashAttention 대비 Ring-Attention과 같은 자릿수의 차이를 보이는 수준이다. 또한 제공된 원문에서는 듀얼 노드 NVLink 백워드 평균이 2.9배와 3.9배로 다르게 표기되어 있어, 인용 시 원문 표기 불일치를 확인할 필요가 있다.