GC-MoE가 MoE 디코드에서 통신과 GEMM을 SM 분할로 겹친다

GC-MoE: Overlapping Expert-Parallel Token Dispatch with FP8/FP4 Grouped GEMM via CUDA Green Contexts

arXiv2611.04222v1

Yong Li2026-11-03

무엇인가

이 논문은 expert parallelism으로 서빙되는 MoE 모델의 디코드 단계에서 매 레이어 반복되는 all-to-all 토큰 디스패치, grouped GEMM, combine 비용을 다룬다. 디코드에서는 랭크마다 1~256개 토큰만 처리해 전문가별 GEMM의 M 차원이 16~32행 수준으로 작아지고, 연산보다 디스패치 지연과 FP4 가중치 스트리밍이 병목이 된다. 기존 DeepGEMM Mega MoE는 하나의 persistent 융합 커널에서 디스패치와 GEMM을 겹치지만, SM마다 디스패치 warp와 스테이징 버퍼를 떼어내야 하고 모든 단계가 하나의 CTA 자원 예산을 공유한다. 또 prologue와 tail의 grid sync가 B300의 148개 SM 전체를 포함하며, 매 replay 끝에 workspace clean과 cross-rank tag-2 barrier가 붙는다. 저자들은 이 구조가 SM당 GEMM 처리량을 낮추고, 동기화 범위를 키우고, 한 단계의 최적화가 전체 커널 예산을 건드리게 만든다고 본다.

어떻게 동작하나

GC-MoE의 핵심은 CUDA Green Contexts로 디바이스를 드라이버 수준에서 작은 통신 파티션과 큰 연산 파티션으로 나누는 것이다. B300에서는 148개 SM 중 20개를 디스패치에, 128개를 연산에 배정한다. 각 역할은 별도 커널이 된다. role1은 통신 파티션에서 라우팅 계정과 all-to-all pull을 수행하고, role2는 연산 파티션에서 routed/shared grouped GEMM과 SwiGLU, requantization을 수행하며, role3은 두 역할 뒤에서 top-k와 shared expert partial을 합친다. role1과 role2는 3노드 child CUDA Graph에서 의존 엣지 없는 동시 루트로 실행되고, Hopper에서 shared expert를 융합하면 role4가 추가된다. 크로스랭크 데이터는 PyTorch distributed symmetric-memory API로 만든 대칭 메모리에 두고, 592바이트 by-value descriptor에 base 주소와 72개 상대 오프셋을 담아 NVLink peer 주소를 커널에 전달한다. 파티션은 그래프 수명 동안 정적이며, 런타임에 최소 32개 연산 SM, 짝수 연산 SM, 최소 2개 디스패치 SM, 통신과 연산 SM 합이 디바이스 SM 수와 같은지 검증한다.

무엇과 다른가

겹침은 토큰 링으로 구현된다. 디스패치는 전문가별 버퍼를 따로 잡지 않고 BM 토큰 단위 링 블록을 재사용하며, 각 블록은 l1_full_count, l1_empty_count, 64비트 l2_full_mask, l2_empty_count 신호를 갖는다. 연산 커널은 디스패치 전체가 끝나기를 기다리지 않고 첫 링 블록이 채워지는 즉시 소비를 시작한다. generation tag가 붙은 카운터와 double-buffered receive counter를 사용하고 동기화 상태를 디바이스에 영속시켜, 연속된 CUDA Graph replay 사이에 꼬리 cross-rank barrier가 필요 없게 했다. role1의 grid sync는 20개 디스패치 SM에만, role3의 grid sync는 finalize grid에만 걸린다.

어떻게 쓰나

연산 커널은 아키텍처별로 나뉜다. Blackwell에서는 operand-swapped tcgen05 block-scaled MMA를 사용해 FP4 routed 가중치와 FP8 shared 가중치를 하나의 명령 시퀀스로 처리하고, FP4 가중치를 K128 타일로 미리 패킹한다. Hopper에서는 register-A WGMMA와 in-register FP4→E4M3 디코딩, window 기반 scale repacking을 사용한다. 첫 GEMM epilogue에는 SwiGLU, 라우팅 가중치 적용, FP8 requantization이 융합되고 shared expert도 선택적으로 같은 task stream에 융합된다. Blackwell 경로의 routed activation은 K 방향 32개 값마다 UE8M0 스케일을 갖는 FP8 E4M3이고 routed weight는 같은 K32 granularity의 FP4 E2M1, shared weight는 FP8이다. Hopper 경로는 연속 FP32 K128 스케일의 E4M3FN activation과 MXFP4 라우팅 가중치를 쓴다.

전제와 한계

실험은 8×H20에서 DeepSeek-V4.1-Flash, 8×B300에서 GLM-5.2와 DeepSeek-V4.1-Flash 디코드를 DeepGEMM Mega MoE와 비교한다. 측정은 100회 replay, 10회 warmup, 30회 측정 window, L2 flush, 가장 느린 랭크 기준이다. H20 routed에서는 9개 지점 모두 빨라 평균 1.188배, 16 tokens/rank에서 최대 1.296배였고, 1 token/rank 1.04배에서 4~32 tokens/rank 1.26~1.30배 plateau를 보인 뒤 256 tokens/rank에서 1.03배로 줄었다. H20 shared expert 융합은 융합하지 않은 GC-MoE 대비 평균 1.031배(9개 중 6개 지점), unfused Mega baseline 대비 9개 지점 모두 빨라 평균 1.188배, 8 tokens/rank에서 최대 1.353배였다. 다만 32~128 tokens/rank에서는 shared expert가 직렬 연산을 더해 0.973~0.985배로 떨어지는 지점이 있었다.

B300 routed에서는 GLM-5.2 평균 1.102배, DeepSeek-V4.1-Flash 평균 1.109배, 전체 평균 1.106배였고 18개 지점 모두 1.085~1.189배 범위에서 이겼다. B300 shared expert 융합에서는 GLM-5.2 평균 1.131배, DeepSeek-V4.1-Flash 평균 1.115배, 전체 1.123배, 1 token/rank에서 최대 1.209배였다. 두 플랫폼의 routed와 shared expert 스윕을 합친 54개 지점 전부에서 GC-MoE가 빨랐다. 유효 연산 처리량은 H20에서 8 tokens/rank 기준 73.8에서 95.6 TFLOPS로 29.5% 늘었고, 256 tokens/rank에서 1251.4 대 1287.8 TFLOPS였다. B300 GLM-5.2는 256 tokens/rank에서 7524.8 대 8490.7 TFLOPS로 12.8% 늘었고, 1 token/rank에서도 75.3에서 82.5 TFLOPS로 개선됐다. HBM working-set bandwidth 모델은 H20 32 tokens/rank에서 14,973.9 대 18,931.9 GB/s로 26.4% 높았다. 정확도는 H20 routed 47개 케이스 376개 랭크 체크, shared 47개 케이스 376개 체크, B300 94개 케이스 752개 랭크 체크가 모두 통과했고, H20에서 baseline 대비 최소 cosine은 routed 0.9999998, shared 0.9999997이었다.

개발자에게 이 논문은 MoE 디코드 서빙에서 통신 커널과 GEMM 커널을 같은 SM에 억지로 욱여넣는 대신, 드라이버 수준 SM 파티션으로 분리해 겹치는 설계가 실제로 유효하다는 근거를 준다. 다만 CUDA 12.4의 Green Contexts, CUDA Graph capture, symmetric memory, Hopper/Blackwell별 FP8/FP4 커널이 전제이므로 기존 서빙 스택에 바로 끼워 넣는 범용 최적화는 아니다. 작은 배치 디코드에서 이득이 크고, 파티션 크기와 워크큐 설정, workspace 단일 인스턴스 제약, 모델별 shape 제약을 확인해야 한다.

저자들이 밝힌 한계는 분명하다. 이 연산자는 디코드를 겨냥하며 prefill 규모의 토큰 수는 지원하지만 연산 파티션이 병목이 되어 256 tokens/rank에서 이득이 1.01~1.13배로 줄어든다. 64비트 완료 마스크, 정적 신호 크기, 72랭크 by-value peer table 같은 shape 제약은 설계상 고정된 결정이다. 이 버전에는 파티션 크기와 overlap on/off에 대한 통제된 ablation이 없고, shared expert 융합 대비 비융합 비교만 within-operator ablation으로 제시된다. workspace가 단일 인스턴스라 이를 공유하는 호출은 직렬화해야 한다.

초록 (원문)

Under expert parallelism, mixture-of-experts (MoE) layers pay all-to-all token dispatch, grouped GEMMs, and a combine step in every layer. In decode each rank holds only 1-256 tokens: dispatch latency and weight traffic dominate while tensor-core utilization is inherently low. In fused operators such as DeepGEMM's Mega MoE, dispatch and computation overlap only by sharing every SM: dispatch warps compete with TMA/MMA and epilogue warps for registers, shared memory, and issue slots; every phase shares one CTA resource budget, synchronization spans the whole SM array, and every replay ends with a workspace-clean pass and a cross-rank barrier. We present GC-MoE, an open-source FP8/FP4 MoE inference operator for NVIDIA Hopper and Blackwell GPUs that uses CUDA Green Contexts to split the device into a small communication partition and a large compute partition. Each role gets its own kernel, resource budget, and synchronization domain, so dispatch and GEMM tune independently. A three-node CUDA Graph (four on Hopper with fused shared experts) launches the dispatch and grouped-GEMM kernels as concurrent roots; a token-ring handshake starts computation on the first filled ring blocks, and double-buffered receive counters remove the trailing cross-rank barrier between replays. The compute partition runs operand-swapped tcgen05 block-scaled MMA on Blackwell and register-A WGMMA with in-register FP4 decoding on Hopper, with fused SwiGLU, requantization, and optional shared experts. On 8 H20 GPUs GC-MoE accelerates DeepSeek-V4.1-Flash decode by 1.188x on average on both the routed and shared-expert suites (peaks 1.296x and 1.353x); on 8 B300 GPUs it accelerates GLM-5.2 and DeepSeek-V4.1-Flash by 1.106x (1.123x with fused shared experts, peak 1.209x), beating Mega MoE at all 54 measured points at matched accuracy. GC-MoE reserves only 20 of 148 SMs on B300 for dispatch.