🤖 AI

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

이 논문은 디바이스 주도의 원격 스토어(remote store)가 완료되기 전에 진행되는, DiT 추론에서 간과되었던 "X-Stage" 파이프라인 단계를 식별하고, 이러한 통찰을 Burst-Gap 모델을 통해 활용함으로써 데이터 이동과 연산을 효과적으로 중첩시키고 백프레셔(backpressure)를 방지하여 분산 추론을 크게 가속화하는 융합 통신-연산 커널을 재설계한다.

원저자: Jianwen Xian, Zhiyuan Xu, Yuchen Li, Ziliang Lai, Kang He, Zhen Huang, Aichen Feng, Jinyan Chen, Yilin Zhang, Qinqin Chen, Chengru Song

게시일 2026-07-28

📖 4 분 읽기☕ 가벼운 읽기

CC BY 4.0

원저자: Jianwen Xian, Zhiyuan Xu, Yuchen Li, Ziliang Lai, Kang He, Zhen Huang, Aichen Feng, Jinyan Chen, Yilin Zhang, Qinqin Chen, Chengru Song

원본 논문은 CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/) 라이선스로 제공됩니다. ✨ 이것은 아래 논문에 대한 AI 생성 설명입니다. 저자가 작성하거나 승인한 것이 아닙니다. 기술적 정확성을 위해서는 원본 논문을 참조하세요. 전체 면책 조항 읽기

당신이 수천 대의 로봇이 복잡한 구조물을 만들고 있는 거대하고 고속인 공장을 운영하고 있다고 상상해 보십시오. 이 공장에서 로봇들은 두 가지 주요 업무를 수행합니다: 생각하기(계산 수행)와 말하기(다른 로봇들에게 설계도를 전송). 오랫동안 공장 관리자들은 이 두 업무가 엄격한 순서대로 이루어져야 한다고 생각했습니다. 즉, 로봇이 생각을 마치면 멈추고, 다른 로봇들이 메시지를 받을 때까지 기다린 다음, 그제야 다시 생각을 시작해야 한다는 것이었습니다. 이 대기 시간은 엄청난 낭비였으며, 공장 전체의 속도를 늦추는 원인이 되었습니다.

최근 엔지니어들은 로봇이 생각하는 동안에도 말할 수 있는 방법을 발견했습니다. 로봇이 공장의 고속 네트워크로 메시지를 외치기만 하면, 메시지가 반대편 끝에 도착할 때까지 기다릴 필요 없이 즉시 다음 설계도 작업을 시작할 수 있다는 것을 발견한 것입니다. 하지만 여기에는 함정이 있습니다. 공장의 네트워크는 한 번에 처리할 수 있는 "비행 중인(in-flight)" 메시지의 수가 제한되어 있습니다. 만약 로봇이 멈추지 않고 너무 빠르게 많은 메시지를 외치면, 네트워크가 막히고, 외치는 로봇이 갇히게 되어 공장 전체가 멈춰버립니다. 과학자들의 큰 고민은 이것이었습니다: 공장이 막히지 않고 최고 속도로 계속 돌아가게 하려면 정확히 언제 외치고 언제 멈춰야 하는가?

"X-Stage: DiT 추론에서 통신-연산 중첩을 위해 간과되었던 파이프라인 단계"라는 제목의 이 논문은 바로 그 문제를 깊이 파고듭니다. 연구진은 Diffusion Transformer(DiT)라고 불리는 고급 AI 모델을 연구하던 중, 통신 시스템에서 아무도 주목하지 않았던 숨겨진 "대기실"을 발견했습니다. 그들은 이를 X-Stage라고 부릅니다.

X-Stage를 로봇의 입과 수신자의 귀 사이를 잇는 마법 같은 컨베이어 벨트라고 생각해 보십시오. 로봇이 "원격 저장(remote store)"(다른 로봇에게 데이터를 보내는 세련된 방식)을 실행하면, 메시지는 이 X-Stage에 진입합니다. 논문에 따르면, 일단 메시지가 이 벨트에 올라타면 로봇은 즉시 새로운 계산 작업을 수행할 수 있는 자유를 얻습니다. 메시지는 로봇이 새로운 수학 계산을 하는 동안에도 스스로 벨트를 따라 계속 이동합니다. 연구진은 만약 메시지가 비워지는 속도보다 더 빠르게 메시지를 계속 외친다면 벨트가 가득 차게 되고, 결국 로봇은 멈춰야 한다는 사실을 깨달았습니다. 하지만 외치는 타이밍을 완벽하게 맞춘다면, 즉 한 번의 메시지 묶음을 외친 뒤, 이전 메시지들이 벨트에서 빠져나가는 동안 잠시 생각을 하는 시간을 갖는다면, 공장을 최대 속도로 계속 가동할 수 있습니다.

완벽한 타이밍을 찾아내기 위해 연구팀은 버스트-갭(Burst–Gap) 모델이라는 간단한 수학적 모델을 구축했습니다. 당신이 움직이는 트럭 안으로 공을 던지고 있다고 상상해 보십시오.

버스트(The Burst): 당신은 한 움큼의 공(데이터)을 매우 빠르게 던집니다.
갭(The Gap): 당신은 트럭이 공들을 실어 나르는 동안 던지는 것을 멈추고 다른 일(예: 생각하기)을 합니다.
배출율(The Drain Rate): 트럭이 공들을 얼마나 빨리 실어 나를 수 있는지의 속도입니다.
용량(The Capacity): 트ら이 움직임을 멈추기 전까지 담을 수 있는 공의 개수입니다.

연구진은 특정 유형의 강력한 컴퓨터 칩에서 "트럭"(네트워크)이 얼마나 빨리 움직이고 "공"(데이터)을 얼마나 많이 담을 수 있는지 정확히 측정했습니다. 그들은 만약 공을 너무 빠르게 많이 던지면 트럭이 막히고 기다려야 한다는 것을 발견했습니다. 하지만 한 번의 버스트를 던진 후, 트럭이 공간을 확보할 만큼 충분히 기다렸다가 다음 버스트를 던진다면, 작업 중에 한 번도 멈추지 않고 다음 버스트를 즉시 던질 수 있습니다.

이 모델을 사용하여, 팀은 두 가지 특정 AI 공정 부분을 매우 효율적으로 재설계했습니다.

첫째, 그들은 MegaMoE라고 불리는 시스템을 살펴보았습니다. 이는 서로 다른 종류의 업무를 처리하는 전문가 팀과 같습니다. 이전에는 로봇들이 한 가지 업무를 마치고 모든 메시지를 한꺼번에 외친 다음 기다렸습니다. 이는 교통 체증을 유발했습니다. 연구진은 스케줄을 변경하여, 한 그룹의 로봇이 업무를 마치고 메시지를 외치는 동안, 다른 그룹이 다른 업무를 시작하도록 했습니다. 이 "인터리빙(interleaving, 교차 실행)" 방식은 외치는 행위가 작고 관리 가능한 버스트 형태로 일어나게 하여, 네트워크가 비워질 수 있는 충분한 생각 시간을 확보해 주었습니다. 이 간단한 변화로 인해 시스템은 평균 1.18배 빨라졌으며, 최선의 경우 최대 1.62배까지 빨라졌습니다.

둘째, 그들은 긴 시퀀스의 데이터(예: 긴 이야기)를 처리하는 방법인 FlashAttention을 다루었습니다. 그들은 "생각하기" 부분과 "보내기" 부분을 결합하여, 메시지를 보내는 로봇이 멈추지 않도록 했습니다. 메시지가 비워지기를 기다리기만 하는 전담 로봇을 두는 대신, 이미 수학 계산을 하고 있는 로봇이 메시지를 보내고 즉시 다시 수학 계산으로 돌아가게 했습니다. "X-Stage" 컨베이어 벨트가 백그라운드에서 메시지 전달을 처리하게 된 것입니다. 이 접근 방식은 기존의 순차적 방식에 비해 한 버전에서는 1.43배, 다른 버전에서는 1.42배 더 빠른 속도를 보여주었습니다.

이 논문은 이것이 단순한 마법이 아니라 정밀한 측정임을 매우 신중하게 밝히고 있습니다. 만약 X-Stage를 무시하고 로봇이 메시지가 도착할 때까지 기다려야 한다고 가정한다면, 시스템이 얼마나 빨라질 수 있는지 과소평가하게 될 것입니다. 반대로 로봇이 멈추지 않고 영원히 달릴 수 있다고 가정한다면, 결국 네트워크를 충돌시키게 될 것입니다. "버스트-갭" 모델은 그 중간의 최적의 지점입니다.

요약하자면, 연구진은 컴퓨터들이 서로 소통하는 방식에서 숨겨진 "사이 단계"를 찾아냈습니다. 메시지가 얼마나 빨리 이동하고 얼마나 많은 공간을 차지하는지 정확히 측정함으로써, 그들은 컴퓨터가 업무와 대화를 완벽하게 조절하며 수행하도록 가르쳤습니다. 멈춰서 기다리는 대신, 컴퓨터는 메시지를 보내고, 메시지가 이동하는 동안 일부 작업을 수행한 다음, 첫 번째 메시지가 길을 비켜주는 즉시 다음 메시지를 보낼 수 있게 되었습니다. 이 작은 타이밍의 조정은 새로운 하드웨어 없이도 AI 모델을 더 빠르고 효율적으로 실행하게 만드는 엄청난 속도 향상을 가져옵니다.

기술 요약: DiT 추론에서의 X-Stage

문제 정의

분산형 확산 트랜스포머(Diffusion Transformer, DiT) 추론에서, GPU 연산 처리량과 인터커넥트 대역폭 사이의 격차가 벌어짐에 따라 GPU 간 데이터 이동이 크리티컬 패스(critical path)에 놓이게 됩니다. 최근의 시스템들(예: FLUX, Comet, MegaScale-MoE)은 타일(tile), CTA 또는 워프(warp) 단위로 작업을 분해하여 미세한 통신-연산 중첩(communication–computation overlap)을 달전성하였으나, 이들은 통신이 언제 발행(issue)되는지와 데이터가 언제 원격 가시성(remote visibility)을 갖게 되는지를 정의하는 추상화 모델에 의존하고 있습니다.

기존의 이러한 추상화 모델들은 원격 스토어(remote store)가 발행된 후부터 원격 가시성이 확보되기 전까지의 시스템 동작을 누락하고 있습니다. 구체적으로, 이들은 다음을 예측하지 못합니다:

원격 스토어가 인플라이트(in-flight) 상태인 동안, 발행 주체(issuing role)가 유용한 연산을 얼마나 더 진행할 수 있는가?
원격 스토어를 지속적으로 주입할 때, 언제 유한한 다운스트림 자원이 고갈되어 컴퓨트 파이프라인을 스톨(stall)시키는 백프레셔(backpressure)가 발생하는가?

이러한 한계는 부정확한 성능 모델링으로 이어집니다. 예를 들어, DeepGEMM의 MegaMoE 커널에서 "완료 결합형(completion-coupled)" 모델(발행자가 원격 가시성을 기다린다고 가정)은 약 1.5×의 속도 향상 상한을 예측하지만, 실제 구현은 1.56×를 달성합니다. 이러한 차이는 발행자가 포스트 이슈(post-issue) 구간 동안 유용한 작업을 재개할 수 있음을 시사하며, 이는 현재의 스케줄링 모델로는 포착되지 않는 현상입니다.

방법론: X-Stage 추상화 및 Burst–Gap 모델

저자들은 X-Stage를 원격 스토어 발행 수락과 원격 가시적 완료 사이의 유한한 간격을 나타내는 소프트웨어 가시적 파이프라인 단계로 정의합니다. X-Stage 동안, 수락된 요청은 목적지를 향해 계속 진행되는 동시에 발행 주체는 연산을 재개합니다.

이 동작을 특성화하고 예측하기 위해, 저자들은 세 가지 측정 가능한 양으로 매개변수화된 경량 Burst–Gap 모델을 제안합니다:

백프레셔가 없는 이슈 시간 ( $T^0_{iss}$ ): 백프레셔가 존재하지 않을 때 원격 스토어 버스트를 주입하는 데 걸리는 시간.
유효 드레인 속도 ( $R$ ): 수락된 요청을 처리하여 비우는 속도 (대상 시스템에서 약 717 GB/s로 측정됨).
유효 잔류 용량 ( $Q$ ): 이슈 속도가 정체되기 전까지 유지될 수 있는 유한한 데이터 부피.

이 모델은 발행자 가시적 버스트 이슈 시간 $T_{iss}$ 를 다음과 같이 정의합니다:
$T_{iss}(G; K, V) = \max\left(T^0_{iss}(K, V), \frac{V}{R} - G\right)$
여기서 $V$ 는 총 버스트 부피, $K$ 는 프로듀서의 수, $G$ 는 버스트 사이의 유용한 프로듀서 측 갭(computation time)입니다.

이 모델은 두 가지 레짐(regime)을 구분합니다:

드레인 지배형(Drain-dominated): 갭 $G$ 가 짧으면, 미결 요청이 축적되며, 이슈 시간은 드레인 속도( $V/R - G$ )에 의해 결정됩니다.
백프레셔 프리(Backpressure-free): $G$ 가 충분하여 이전 요청을 드레인할 수 있다면, 이슈 시간은 베이스라인인 $T^0_{iss}$ 로 돌아갑니다.

주요 기여

본 논문은 네 가지 주요 기여를 합니다:

X-Stage의 식별: 저자들은 X-Stage를 별도의 측정 가능한 파이프라인 단계로 정의합니다. 이 추상화는 왜 완전 블로킹(fully blocking) 모델도, 무제한 비동기(unbounded-asynchronous) 모델도 장치 주도형 통신에서의 발행자 가시적 동작을 정확하게 포착하지 못하는지를 설명합니다.
특성화 및 예측: 매개변수화된 마이크로 벤치마크를 통해 저자들은 $T^0_{iss}$ , $R$ , $Q$ 를 측정합니다. 결과적인 Burst–Gap 모델은 애플리케이션별 재적합(refitting) 없이도 정상 상태(steady-state) 기간, 발행자 가시적 오버헤드, 그리고 백프레셔의 발생 시점을 예측합니다.
X-Stage 인식 커널 설계: 이 모델은 두 가지 뚜렷한 커널 최적화를 가이드합니다:
- MegaMoE (Reshaping): DeepGEMM의 MegaMoE를 위해, 저자들은 Linear-1(연산)과 Linear-2(통신) 작업을 엑스퍼트 웨이브(expert waves) 전반에 걸쳐 혼합하는 **인터리브드 스케줄러(interleaved scheduler)**를 도입합니다. 이는 집중된 원격 스토어 버스트 사이의 간격을 넓혀 용량 고갈을 방지하도록 연산을 재분배합니다.
- FlashAttention–A2A (Piggybacking): Ulysses 시퀀스 병렬 어텐션(sequence-parallel attention)을 위해, 저자들은 포스트 어텐션 All-to-All(A2A)을 FlashAttention과 퓨즈(fuse)합니다. 출력 타일 소유자는 원격 스토어를 발행한 직지 다음 타일의 Q-루프를 재개합니다. Q-루프는 X-Stage 드레이닝을 위한 자연스러운 "갭" 역할을 수행하며, 별도의 통신 워프나 SM을 필요로 하지 않습니다.
성능 및 메커니즘 검증: 저자들은 84개의 MegaMoE 구성과 두 세대의 FlashAttention(FA3, FA4)을 통해 모델을 검증하여, 측정된 속도 향상을 이슈, 드레인 및 용량 효과와 연결합니다.

결과

MegaMoE 최적화

Linear-1과 Linear-2 작업을 인터리빙함으로써, 저자들은 원격 스토어의 주입 프로파일을 재구성합니다:

속도 향상: 84개 구성에 대해 엑스퍼트 웨이브 베이스라인 대비 1.18× 기하 평균 및 1.62× 최대 커널 속도 향상을 달성했습니다.
메커니즘: 텐서 코어 타임라인 분석 결과, 인터리빙이 백프레셔로 인한 저활동 구간을 줄였음을 확인했습니다. 타일당 계측(per-tile instrumentation)을 통해 원격 스토어 스팬 분포가 광범위한 드레인 제한 레짐에서 좁은 백프레셔 프리 플로어로 이동했음을 확인했습니다.

FlashAttention–A2A 최적화

출력 소유 역할에 원격 스토어를 피기백킹(piggybacking)함으로써:

속도 향상: FlashAttention-3와 FlashAttention-4는 순차 실행 대비 각각 최대 1.43× 및 **1.42×**의 발행자 가시적 속도 향상을 달성했습니다.
정상 상태 동작: 시퀀스 길이가 증가함에 따라 Q-루프 갭이 커져, 결국 포스트 이슈 드레인 윈도우를 커버하게 됩니다. 긴 시퀀스에서 퓨즈된 커널의 정상 상태 시간은 FlashAttention 단독 실행 시간과 거의 일치하며, 발행자 가시적 잔여 오버헤드는 측정 오차 범위 내에서 거의 0에 가깝게 유지되었습니다.
검증: 결과는 "완료 결합형" 반사실적 모델(counterfactual model)이 시퀀스 길이에 따라 오버헤드가 선형적으로 증가할 것이라고 예측했던 것과 상반되었습니다. 대신 오버헤드는 평탄하게 유지되었으며, 이는 X-Stage 드레이닝이 통신 비용을 성공적으로 숨겼음을 입증합니다.

의의 및 주장

본 논문은 **포스트 이슈 진행(post-issue progress)**이 현재의 분산형 DiT 추론 시스템에서 간과되어 온 측정 가능한 스케줄링 레버임을 주장합니다. 발행과 원격 가시성 사이의 유한한 간격을 명시적으로 모델링함으로써, 저자들은 다음을 입증합니다:

백프레셔는 예측 가능하다: 이는 단순히 통신 볼륨만의 속성이 아니라, 주입 속도, 드레인 속도, 그리고 잔류 용량 간의 관계에 따른 함수입니다.
스케줄링은 하드웨어 변경 없이 최적화 가능하다: 독립적인 작업의 단순한 재정렬(MegaMoE)이나 기존의 연산 갭 활용(FlashAttention)을 통해 발행 측 백프레셔를 제거할 수 있습니다.
자원 효율성: 피기백킹 방식은 전용 통신 자원(워프/SM)을 예약할 필요가 없으며, 이는 높은 점유율(occupancy)과 연산 밀도를 유지하는 데 매우 중요합니다.

저자들은 X-Stage를 기존의 타일 수준 의존성 및 원격 준비 프로토콜에 대한 보완적인 추상화로 위치시킵니다. 이들의 특성화는 최신 NVIDIA 아키텍처의 원사이드(one-sided) 원격 스토어에 특화되어 있으며, 다른 통신 프리미티브(예: atomics, collectives)는 독립적인 검증이 필요할 수 있다고 언급합니다. 그러나 장치 주도형 원격 스토어의 범위 내에서, Burst–Gap 모델은 통신 버스트를 형성하고 발행 측 오버헤드를 숨기기 위한 엄격한 프레임워크를 제공합니다.