[NeurIPS-2025] FlashMoE: Fast Distributed MoE in a Single Kernel

FlashMoE: Fast Distributed MoE in a Single Kernel

NeurIPS 2025(표기: To appear) · Omotayo A. Yamenja, Divyansh Sharma, Ryan Prout, Yiran Chen · arXiv:2506.04667v3 · 링크: https://arxiv.org/abs/2506.04667v3 ·

 

 

 5줄 요약

1. 이 논문은 분산 Mixture-of-Experts(MoE) 추론의 핵심 병목이 expert 수학 자체보다 CPU 주도 실행 구조, 동기식 All-to-All, 잦은 커널 런치에 있다는 점을 겨냥합니다.
2. 핵심 아이디어는 gate→dispatch→expert FFN→combine→inter-GPU communication 전체를 하나의 persistent GPU kernel로 융합하고, GPU 내부에서 actor 기반 런타임이 타일 단위 작업을 자율적으로 스케줄링하게 만드는 것입니다.
3. 이를 위해 Subscriber/Scheduler/Processor로 역할 분화된 GPU actor 모델, NVSHMEM 기반 one-sided async (R)DMA, symmetric tensor layout, in-place padding 기반 payload-efficient sparse communication을 제안합니다.
4. 8×H100, 최대 128 experts, 16K tokens 조건에서 FlashMoE는 최대 6.4× latency 개선, 93.17% SM utilization, 17.7 MTokens/s throughput, 최대 4× overlap efficiency를 보고합니다.
5. 결국 이 논문은 “분산 MoE의 성능 한계는 sparse 모델링이 아니라 실행기 설계에 있다”는 점을 설득력 있게 보여주며, GPU-native runtime이라는 방향을 강하게 제시합니다.

한 줄 핵심
FlashMoE는 분산 MoE를 GPU 내부 단일 persistent kernel + actor 기반 device-side runtime으로 재구성해, 토큰 재배치·원격 expert 계산·재조합·통신을 하나의 흐름으로 겹쳐 돌리는 시스템 논문입니다.

---

1. 이 논문이 푼 정확한 문제

MoE는 전체 파라미터 수를 크게 키우면서도 각 토큰이 일부 expert만 활성화하기 때문에 계산 효율이 높다는 장점이 있습니다. 하지만 대규모 모델에서는 expert들을 여러 GPU에 분산 배치해야 하므로, 각 토큰은 자신이 배정된 expert가 위치한 GPU로 이동했다가 다시 원래 순서로 돌아와야 합니다. 즉, 분산 MoE의 실제 병목은 FFN 수학 계산보다 토큰 dispatch, inter-GPU transfer, combine, 단계 간 대기, 커널 launch gap에 있습니다.

논문은 기존 분산 MoE 런타임이 보통 다음 세 가지 병목에 묶인다고 봅니다. 첫째, CPU가 커널을 계속 launch하며 실행을 오케스트레이션합니다. 둘째, 통신이 동기식 collective 중심이라 가장 느린 GPU에 전체가 끌립니다. 셋째, gate·dispatch·expert GEMM·combine이 잘게 쪼개져 수십~수백 개의 GPU op로 실행되며, 그 사이사이 idle gap이 발생합니다. FlashMoE는 바로 이 실행 구조 전체를 뜯어고치려는 작업입니다.

---

2. 이전 접근이 부족했던 지점

기존 FasterMoE, DeepEP+Megatron-LM, Megatron-TE, COMET류 접근은 일부 커널 fusion, 일부 overlap, 더 나은 collective tuning을 시도했지만, 기본 구조는 여전히 host-orchestrated multi-kernel pipeline에 가깝습니다. 즉, 어떤 부분이 빨라졌더라도 전체 분산 MoE operator가 하나의 연속적인 device-side 실행 흐름으로 바뀐 것은 아니었습니다.

• 동기식 All-to-All 한계: collective는 전체 참여자가 맞춰 들어와야 하므로 straggler 영향이 큽니다.
• launch overhead 지속: 계산과 통신이 쪼개져 있어 CPU-GPU orchestration 공백이 남습니다.
• sparse인데 payload는 dense에 가깝게 전송: 실제 활성 토큰보다 zero-padded payload가 더 많이 오갈 수 있습니다.
• fine-grained locality 활용 부족: 준비된 작은 작업을 즉시 처리하는 device-side scheduler가 없어 타일 단위 반응성이 낮습니다.

논문 Table 1은 이 한계를 매우 직설적으로 보여줍니다. 단일 MoE layer 기준 GPU ops 수가 FlashMoE는 1개인데, COMET은 33개, Megatron-LM CUTLASS는 85개, Megatron-LM TE는 261개, Megatron-LM + DeepEP는 432개, DeepSpeedMoE는 550개입니다. 저자들의 메시지는 분명합니다. 성능 차이의 상당 부분이 수학 계산 자체보다 orchestration overhead에서 나온다는 것입니다.

---

3. 이 논문의 정확한 novelty

• 분산 MoE operator 전체를 단일 persistent kernel로 통합했습니다. gate, dispatch, expert compute, combine, 심지어 inter-GPU communication orchestration까지 커널 내부에서 처리합니다.
• actor-based GPU runtime을 도입했습니다. Subscriber, Scheduler, Processor가 GPU 안에서 서로 다른 역할을 맡아 task를 소비합니다.
• one-sided device-initiated async (R)DMA를 사용해 collective barrier를 피하고, 준비된 타일을 바로 peer GPU 메모리로 밀어 넣습니다.
• symmetric tensor layout + in-place padding으로 write conflict를 피하면서 sparse communication의 payload 낭비를 줄입니다.

따라서 FlashMoE의 새로움은 단순 kernel fusion 하나가 아니라, GPU-resident runtime + actor scheduling + one-sided async communication + sparse-aware memory layout를 결합한 전체 시스템 재설계에 있습니다.

---

4. 입력 / 출력 / 방법

입력(input)

• 입력 활성화 A ∈ R^(S×H): S는 토큰 수(또는 시퀀스 전체 토큰 수), H는 hidden dimension입니다.
• expert weights X ∈ R^(E×H×D): 로컬 expert 수 E와 intermediate dimension D를 가지는 FFN 가중치들입니다.
• routing 관련 정보: gate가 생성하는 affinity score와 토큰→expert 배정 정보가 포함됩니다.

출력(output)

• MoE 레이어를 통과한 최종 출력 O ∈ R^(S×H)입니다.
• 기능적으로는 일반 distributed MoE layer와 같은 입출력을 갖지만, 내부 실행 구조는 완전히 다릅니다.

왜 이 입출력이 중요한가

논문의 초점은 MoE 수식 자체가 아니라, 같은 입출력을 어떻게 GPU-native하게 실행할 것인가입니다. 즉, 사용자 입장에서는 같은 MoE layer인데 시스템 내부에서는 완전히 다른 실행기입니다.

실행 흐름은 대략 다음과 같습니다. 먼저 fused gate가 routing table과 gate score를 계산합니다. 이후 대부분의 block은 Processor로 동작하며 tile 단위 FFN/Combine 작업을 수행합니다. 마지막 하나의 block은 일종의 OS block으로 쓰이며, 내부에서 1 warp는 Scheduler, 3 warps는 Subscriber 역할을 맡습니다. Subscriber는 원격 GPU에서 들어온 tile packet을 해석해 실행 가능한 task descriptor로 바꾸고, Scheduler는 준비된 task를 놀고 있는 Processor block에 배정합니다.

즉, 전통적인 “stage 1 끝 → barrier → stage 2 시작”이 아니라, 준비된 타일이 생기는 즉시 실행 가능한 event-driven pipeline입니다. 이 persistent behavior가 launch overhead 제거와 fine-grained overlap의 핵심 기반입니다.

4-1. Task abstraction

저자들은 FFN 계산과 combine을 모두 공통 task abstraction으로 표현합니다. 각 task descriptor는 어떤 device/tile에서 어떤 연산자(행렬곱 또는 원소곱 등)를 어떤 메타데이터와 함께 수행할지를 담습니다. 이 추상화를 통해 Subscriber, Scheduler, Processor가 동일한 작업 단위를 주고받을 수 있고, GEMM0·GEMM1·Combine 같은 단계들이 타일 granularity에서 독립적으로 준비되고 소비됩니다.

4-2. Actor model과 역할 분화

FlashMoE는 모든 block을 동일하게 취급하지 않습니다. 대부분의 block은 실제 수치 계산을 담당하는 Processor이고, 마지막 block은 OS block으로서 행정/제어를 맡습니다. 이 block 안에서 Scheduler가 work-conserving하게 task를 할당하고, Subscriber가 도착한 원격 패킷을 해석해 새로운 task를 생성합니다. 저자들의 의도는 명확합니다. 제어 로직은 아주 적은 SM 자원만 먹게 하고, 대부분의 자원은 실제 계산에 남겨두는 것입니다.

4-3. 통신 설계: one-sided async (R)DMA

기존 All-to-All은 API 수준에서 집단 동기화가 필요한 경우가 많습니다. FlashMoE는 NVSHMEM 기반 global address space를 이용해, GPU가 직접 peer GPU 메모리로 tile 또는 tile packet을 write합니다. 그리고 데이터 write와 도착 signal이 함께 전달되므로, 수신 측 Subscriber가 이를 곧바로 task로 해석할 수 있습니다. 이 구조는 collective barrier를 피하고, 준비된 작은 일감이 생기는 즉시 실행을 이어갈 수 있게 만듭니다.

4-4. Symmetric tensor layout

one-sided write를 안전하게 하려면 write-write conflict가 없어야 합니다. 이를 위해 논문은 대칭적인 텐서 레이아웃 L ∈ R^(P×R×B×E×C×H)를 정의합니다. 여기서 P는 expert parallel world size, R은 dispatch와 combine 등 통신 round, B는 staging buffer, E는 local experts, C는 expert capacity, H는 hidden dimension입니다. 핵심은 temporal buffering입니다. 여러 통신 시점과 방향이 겹치더라도 서로 다른 버퍼 좌표를 사용해 충돌 없이 비동기 접근을 하도록 설계합니다.

4-5. In-place padding과 sparse payload 효율화

기존 구현은 capacity를 맞추기 위해 실제로는 비어 있는 슬롯도 전송하는 경우가 많습니다. FlashMoE는 네트워크로는 활성 토큰만 보내고, 필요한 패딩은 로컬 symmetric tensor buffer에서 in-place로 처리합니다. 이는 sparse 모델에서 매우 중요합니다. 이 논문은 단지 계산 sparsity만 쓰는 것이 아니라, 통신 sparsity까지 회수하려고 합니다.

---

5. 핵심 Figure / Table 해설

Figure 2 — Transformer FFN, MoE, Distributed MoE의 차이

출처: arXiv HTML Figure 2 · 원문 캡션: “Transformer blocks (a) without MoE, (b) with MoE, and (c) with distributed MoE and expert parallelism.”

무엇을 보여주나: (a)는 일반 Transformer FFN, (b)는 여러 expert 중 일부만 활성화되는 MoE, (c)는 expert가 여러 GPU에 흩어진 distributed MoE를 보여줍니다.

왜 중요한가: 분산 MoE에서 진짜 병목이 단순 FFN 계산이 아니라 토큰 이동, 원격 expert 실행, 원래 순서로의 재조합이라는 점을 가장 직관적으로 설명해 줍니다.

논문의 core contribution과 어떻게 연결되나: FlashMoE는 바로 이 (c)의 실행 경로를 통째로 GPU 내부로 끌어들여, token dispatch와 remote compute를 별도 단계를 거치지 않는 흐름으로 바꾸려는 시도입니다.

Figure 3 — 기존 SOTA 방식과 FlashMoE의 구조적 차이

출처: arXiv HTML Figure 3 · 원문 캡션 일부: “Comparing FlashMoE with state-of-the-art techniques ... FlashMoE is a persistent kernel that fuses all computation and communication of the MoE operator ...”

무엇을 보여주나: 기존 방식은 계산과 통신이 분리된 단계이거나 부분 overlap 정도에 머무는 반면, FlashMoE는 전체 MoE operator가 하나의 persistent kernel 안에서 실행되는 구조를 보여줍니다.

왜 중요한가: 이 그림은 “FlashMoE가 왜 단지 더 빠른 구현이 아니라 실행 철학이 다른가”를 한 장으로 설명합니다.

논문의 core contribution과 어떻게 연결되나: CPU-orchestrated multi-kernel에서 GPU-resident event-driven kernel로의 전환이 바로 논문의 핵심 novelty입니다.

Figure 5 — FlashMoE fused kernel 내부 구조

출처: arXiv HTML Figure 5 · 원문 캡션: “FlashMoE Fused Kernel”

무엇을 보여주나: Processor block들이 다수의 tile task를 처리하고, 하나의 OS block 안에서 Scheduler와 Subscriber가 제어 기능을 맡는 구조를 시각화합니다.

왜 중요한가: GPU 안에 작은 런타임/운영체제를 넣는다는 이 논문의 감각을 가장 잘 보여줍니다. FlashMoE는 “계산 한 번 하고 끝나는 커널”이 아니라, 살아 있으면서 일을 받아 계속 수행하는 커널입니다.

논문의 core contribution과 어떻게 연결되나: persistent behavior, actor specialization, task-driven scheduling이라는 설계가 launch overhead 제거와 계산-통신 overlap의 직접 원인입니다.

Figure 9 — SM utilization 비교

출처: arXiv HTML Figure 9 · 원문 캡션 일부: “SM utilization ... average SM utilization over 100 iterations.”

무엇을 보여주나: FlashMoE의 평균 SM utilization은 93.17%이며, FasterMoE 9.67%, DeepEP+Megatron-LM 13.55%, Megatron-TE 59.11%, Comet 42.31%보다 크게 높습니다.

왜 중요한가: 이 수치는 기존 분산 MoE 구현에서 GPU가 얼마나 오래 놀고 있었는지를 보여줍니다. FlashMoE의 성과는 “한 연산을 약간 더 빠르게 했다”보다 “GPU를 거의 계속 일하게 만들었다”는 데 있습니다.

논문의 core contribution과 어떻게 연결되나: persistent single-kernel과 actor runtime이 실제로 idle gap을 줄였다는 가장 직접적인 시스템 증거입니다.

Figure 10 — GPU 수 증가 시 throughput scaling

출처: arXiv HTML Figure 10 · 원문 캡션 일부: “Throughput when scaling the number of GPUs ...”

무엇을 보여주나: FlashMoE는 8 GPUs에서 17.7 MTokens/s throughput에 도달하며 거의 선형에 가까운 scaling을 보입니다.

왜 중요한가: 논문은 FasterMoE 대비 5.7×, Megatron-TE / Megatron-CUTLASS 대비 4.9× 높은 처리량을 보고합니다. 특히 FlashMoE는 FP32인데 baseline은 FP16이라는 점이 인상적입니다.

논문의 core contribution과 어떻게 연결되나: 이 결과는 실행 구조 개선만으로도 더 높은 실효 throughput을 만들 수 있다는 논문의 핵심 주장을 뒷받침합니다.

Figure 11 — 계산-통신 overlap efficiency

출처: arXiv HTML Figure 11 · 원문 캡션 일부: “Weak scaling efficiency ...”

무엇을 보여주나: 논문은 overlap efficiency를 O_e = T(2)/T(N_G)로 정의하고, GPU 수가 늘어나도 FlashMoE의 효율 저하가 작음을 보입니다.

왜 중요한가: 4 GPUs와 8 GPUs에서 각각 최대 3.88×, 4× 더 높은 효율을 보였다고 보고합니다. 즉, 분산 규모가 커질수록 보통 통신 때문에 무너지는 효율을 상당 부분 붙들어 둔다는 뜻입니다.

논문의 core contribution과 어떻게 연결되나: async communication과 scheduler-driven task dispatch가 분산 확장성에 실질적인 영향을 준다는 증거입니다.

Table 1 — 커널 런치 수 비교

무엇을 보여주나: FlashMoE는 단일 MoE layer를 GPU ops 1개로 처리하는 반면, 기존 방식은 수십~수백 개 op를 사용합니다.

왜 중요한가: 성능 격차의 중요한 부분이 math peak 차이보다 orchestration 차이에서 나온다는 논문 메시지를 가장 압축적으로 보여주는 보조 증거입니다.

논문의 core contribution과 어떻게 연결되나: single-kernel 설계가 단순 미학이 아니라 실질적인 실행 공백 제거 장치임을 드러냅니다.

---

6. 결과를 어떻게 읽어야 하는가

실험은 8×H100 80GB, NVLink 환경에서 수행됐고, attention heads 16, embedding dimension 2048, FFN intermediate 2048, top-2 routing, capacity factor 1.0 조건의 forward-only single MoE layer를 측정합니다. 비교 대상은 Comet, FasterMoE, Megatron-CUTLASS, Megatron-TE, 일부 구간에서는 DeepEP+Megatron-LM입니다.

가장 중요한 결과는 latency, utilization, throughput, overlap efficiency 네 축입니다. Figure 8 기준으로 4 GPUs 16K tokens에서 Megatron-TE 대비 최대 4.6×, FasterMoE 대비 2.6× speedup을 보였고, 8 GPUs에서는 최대 6.4× speedup을 보고합니다. Figure 9는 평균 93.17% SM utilization을 보여줍니다. Figure 10은 8 GPUs에서 17.7 MTokens/s throughput, FasterMoE 대비 5.7×, Megatron-TE / CUTLASS 대비 4.9× 향상을 제시합니다. Figure 11은 4 GPUs 및 8 GPUs에서 각각 최대 3.88×, 4× 더 높은 overlap efficiency를 보여줍니다.

이 결과들은 “GPU arithmetic capability가 부족해서 느린 것”이 아니라, 분산 MoE가 실행 공백과 동기 통신 때문에 느렸다는 저자들의 해석을 지지합니다. FlashMoE는 sparse 모델의 이론적 효율을 실제 시스템 효율로 끌어내리려는 시도이고, 적어도 이 실험 범위에서는 그 효과가 매우 큽니다.

또 하나 흥미로운 부분은 expert scaling입니다. 논문은 experts 수가 8→128로 늘어날 때도 FlashMoE가 낮고 비교적 안정적인 latency를 유지하며, 4 H100에서 최대 4×, 8 H100에서 128 experts 기준 최대 6.6× 우위를 보인다고 설명합니다. 즉, expert가 많아질수록 orchestration 복잡도가 커지는 ultra-sparse 구간에서 특히 강합니다.

---

7. 실무적 / 학문적 의미

실무적으로 이 논문이 중요한 이유는, 앞으로 MoE 최적화의 초점을 “더 빠른 collective”나 “더 빠른 expert GEMM”만이 아니라 GPU-native runtime 설계로 옮기게 만들 가능성이 있기 때문입니다. 배치 규모가 커지고, expert가 더 많아지고, multi-node까지 가면 통신과 스케줄링의 중요도는 더 커집니다.

학문적으로도 메시지가 선명합니다. sparse 모델의 장점은 모델링 관점에서만 확보되는 것이 아니라, 시스템이 sparse execution을 제대로 흡수해 줄 때 비로소 회수됩니다. FlashMoE는 계산 sparsity뿐 아니라 통신 sparsity까지 다루려 한다는 점에서 의미가 큽니다.

더 넓게 보면, 이 논문은 CPU orchestration 중심의 분산 딥러닝 구조가 dynamic workload에서는 한계가 있을 수 있으며, GPU 안에 반응형 런타임을 심는 방식이 앞으로 더 중요해질 수 있음을 보여줍니다. MoE를 넘어 다른 irregular distributed operators에도 이런 철학이 확산될 여지가 있습니다.

---

8. 한계와 주의점

• 평가 범위가 주로 추론 forward 중심입니다. end-to-end 학습 전체, backward, optimizer state, 실제 serving stack까지 검증한 것은 아닙니다.
• 엔지니어링 복잡도가 매우 높습니다. persistent megakernel, custom scheduler, NVSHMEM, custom CUTLASS GEMM, 특수 memory layout 등 재현·이식 난도가 높습니다.
• FP16 / mixed precision path 최적화는 아직 덜 성숙합니다. 부록 H에서 shared-memory layout 비효율 때문에 FP16이 FP32보다 불리할 수 있음을 인정합니다.
• 메모리 오버헤드가 0은 아닙니다. symmetric tensor layout은 temporal buffering을 위해 추가 버퍼를 요구하며, 저자들은 전체 inference 메모리 기준 최대 2.15% 수준이라 설명하지만 환경에 따라 부담이 될 수 있습니다.
• 비교 구현 성숙도 차이가 일부 결과에 영향을 줄 가능성은 있습니다. 다만 그 점을 감안해도 GPU utilization과 kernel-op 수 차이는 구조적 메시지가 매우 강합니다.

Final Summary

FlashMoE는 분산 MoE의 병목을 ‘expert 수학’이 아니라 ‘실행 구조’에서 찾은 논문입니다. gate, dispatch, expert compute, combine, inter-GPU communication을 단일 persistent kernel 안에서 처리하고, actor-based GPU runtime으로 tile-level task를 동적으로 소비하며, one-sided async (R)DMA와 sparse-aware layout으로 통신 낭비를 줄입니다. 그 결과 8×H100 환경에서 최대 6.4× latency 개선, 평균 93.17% SM utilization, 17.7 MTokens/s throughput, 최대 4× overlap efficiency를 보고합니다.

아직 학습 전체 지원, mixed precision 최적화, 구현 복잡도 같은 과제가 남아 있지만, 이 논문의 진짜 가치는 “분산 MoE는 더 좋은 커널 몇 개가 아니라 더 좋은 실행기 설계가 필요하다”는 방향을 강하게 제시했다는 데 있습니다. 읽고 나면 FlashMoE를 단순한 fast MoE implementation이 아니라, GPU-native distributed runtime 제안으로 보는 것이 가장 정확합니다.

 

[문제정의]

MoE 구조자체는 효율적이지만, GPU와 CPU간 토큰 주고받는 시간, CPU 커널 실행에서의 Latency와 같이 계산이 아니라 데이터 이동+실행 Orchestration이 병목의 원인임을 문제제기합니다.

 

[기존 방식의 한계]

기존 SOTA 기법들도 결국 구조는, 여러 단계로 쪼개서 실행(gate -> dispatch -> compute -> combine)하고, CPU의 Orchestration에서의 병목이 발생하며 이때문에 커널 실행 사이에 idle gap이 발생하고, All-to-All 통신이라 가장 느린 GPU를 기준으로 전체 Cluster가 묶인다는 한계점이 있습니다.

 

[본 논문의 핵심]

MoE 전체를 GPU 내부 하나의 커널로 돌림 (gate/dispatch/expert 계산/GPU간 통신까지)

GPU를 거의 안쉬게 만들어서 throughput과 overlap 속도를 400%이상 개선합니다.