[flashinfer] MoE Finalize + AllReduce 융합 커널에 shared-expert 전용 two-shot 프로그램 추가: L2 prefetch로 중간 토큰 구간 최대 3.8% 단축

PR 링크: flashinfer-ai/flashinfer#9284 상태: Merged | 변경: +15197 / -37

들어가며

FlashInfer의 flashinfer.comm.allreduce_fusion(..., moe_finalize_backend="cake")는 MoE 추론 마지막 단계를 하나의 커널로 묶습니다. 여기에는 다음 단계가 포함됩니다.

1. top-k expert 출력 row에 routing scale을 곱해 합산하고, 선택적으로 shared-expert row를 더하는 MoE finalize
2. 단일 노드 NVLink all-reduce
3. residual add
4. RMSNorm (hidden size 7168)
5. 선택적으로 NVFP4 packed 출력 양자화

이 커널은 하나의 범용 커널이 아니라, world size·dtype·출력 형식·토큰 수 구간별로 특화된 "program"들의 집합입니다. 직전 PR #9192 시점에는 아키텍처당 16종류, 63개 program이 있었습니다.

문제가 된 구간은 two-shot 경로에서 shared expert가 있는 top_k 4 배치입니다. world size 4에서는 128 토큰, world size 8에서는 512 토큰 근처가 여기에 해당합니다. 이 구간에서는 deferred shared form이 아직 쓰이지 않습니다. 그래서 shared-expert row를 일반 two-shot program이 다른 operand와 똑같이 처리했고, 매 토큰마다 shared-expert row를 읽는 메모리 지연이 그대로 critical path에 드러났습니다.

이 PR은 이 구간을 위해 3종류, 12개의 전용 program을 추가합니다. 결과적으로 아키텍처당 19종류, 75개 program이 됩니다.

코드 분석

1. 새 커널 선언 (cake_moe_finalize_allreduce_fusion_binding.cu)

커널 이름의 접미사는 다음과 같은 뜻입니다.

• ts: two-shot
• tse: two-shot early
• tss: two-shot shared (신규)
• tses: two-shot early shared (신규)
• o110 / o111: BF16·FP16 출력 / NVFP4 출력

+CAKE_MOE_FINALIZE_KERNEL(__half, kernel_cake_trtllm_moe_finalize_float16_ws4_o110_tss)
+CAKE_MOE_FINALIZE_KERNEL(__half, kernel_cake_trtllm_moe_finalize_float16_ws4_o110_tses)
+CAKE_MOE_FINALIZE_KERNEL(__half, kernel_cake_trtllm_moe_finalize_float16_ws4_o111_tss)
...
+CAKE_MOE_FINALIZE_KERNEL(__nv_bfloat16, kernel_cake_trtllm_moe_finalize_bfloat16_ws8_o111_tss)

world size 8에는 tses(early shared) 변형이 없습니다. 이 점은 아래 dispatch 로직에서 다시 등장합니다.

2. 디스패치 테이블 추가

기존 kTwoShotKernels, kTwoShotEarlyKernels와 같은 [dtype][ws][quant] 3차원 테이블이 두 개 추가됩니다.

+// Two-shot programs that also prefetch the token's shared-expert row to L2 and load it
+// read-only, same order as kKernels; world sizes 4 and 8. They run the same exchange as
+// kTwoShotKernels, so a rank with a shared-expert operand may run them while another rank
+// runs kTwoShotKernels; the binding runs them for top_k 4 only.
+const void* const kTwoShotSharedKernels[2][3][2] = {
+    {{nullptr, nullptr},
+     {reinterpret_cast<const void*>(kernel_cake_trtllm_moe_finalize_float16_ws4_o110_tss), ...},
+     {reinterpret_cast<const void*>(kernel_cake_trtllm_moe_finalize_float16_ws8_o110_tss), ...}},
+    ...
+};
+
+const void* const kTwoShotEarlySharedKernels[2][3][2] = {
+    {{nullptr, nullptr},
+     {..._float16_ws4_o110_tses, ..._float16_ws4_o111_tses},
+     {nullptr, nullptr}},   // ws8: early shared 없음
+    ...
+};

주석에 핵심 설계 두 가지가 담겨 있습니다.

• 최적화 기법: 새 program은 토큰의 shared-expert row를 L2로 미리 prefetch하고 read-only 경로로 load합니다. 기존 program은 shared row를 다른 operand와 같은 시점에 읽었습니다. 새 program은 all-reduce exchange 대기 시간과 겹치도록 L2에 미리 올려두어 지연을 숨깁니다.
• 프로토콜 호환성: 새 program은 기존 two-shot과 같은 exchange를 수행합니다. 따라서 shared-expert operand가 있는 rank는 새 program을 돌리고, 다른 rank는 기존 kTwoShotKernels를 돌려도 됩니다. 멀티 GPU collective에서 rank마다 다른 커널을 고를 수 있다는 것은 큰 장점입니다. 선택 조건이 rank-local 정보에 의존해도 deadlock이나 workspace 불일치가 생기지 않기 때문입니다.

3. 선택 로직 변경

먼저 현재 two_shot_now가 early class인지 판정합니다. 그 다음 같은 class에 해당하는 shared 변형을 고릅니다.

+  const bool two_shot_early_class =
+      two_shot_early != nullptr &&
+      tokens <= (world_size == 8 ? kTwoShotEarlyMaxTokensWorldSize8 : kTwoShotEarlyMaxTokensWorldSize4);
+  const void* two_shot_now_shared =
+      two_shot_early_class ? kTwoShotEarlySharedKernels[dtype_index][ws_index][quant_index]
+                           : kTwoShotSharedKernels[dtype_index][ws_index][quant_index];

실제 커널 선택식은 다음과 같이 바뀝니다.

Before

kernel = two_shot_deferred ? (... ? defer_shared : two_shot_defer) : two_shot_now;

After

kernel = two_shot_deferred ? (... ? defer_shared : two_shot_defer)
  : ((v_has_shared_expert != 0 && top_k == 4 &&
      !(world_size == 8 && two_shot_early_class) &&
      two_shot_now_shared != nullptr) ? two_shot_now_shared : two_shot_now);

조건을 하나씩 보면 다음과 같습니다.

• v_has_shared_expert != 0 && top_k == 4: 벤치마크로 이득이 확인된 클래스만 대상으로 합니다.
• !(world_size == 8 && two_shot_early_class): ws8의 early 구간은 기존 kTwoShotEarlyKernels를 유지합니다. 테이블에 nullptr이 있지만, 명시적으로 한 번 더 막아 의도를 분명히 합니다.
• two_shot_now_shared != nullptr: 테이블에 구멍이 있는 조합은 안전하게 fallback합니다.

같은 식이 world size 8 large-batch 분기에도 동일하게 적용되었습니다. 덕분에 어느 진입 경로로 들어오든 일관된 선택이 보장됩니다.

4. 생성된 커널 파일

cake_trtllm_moe_finalize_bfloat16_ws4_o110_tses_kernel.cu 같은 신규 파일은 각각 1000줄이 넘는 생성 코드입니다. 다음 요소들을 직접 다룹니다.

• mbarrier.try_wait.parity.acquire.cluster
• mbarrier.arrive.expect_tx
• all-gather용 escape 인코딩 cake_ag_esc_enc: 0x8000을 0xFFFF로 치환해 sentinel 값과 충돌하지 않게 합니다.

사람이 손으로 쓰기보다 탐색·생성된 program을 바이너리처럼 추가하는 방식입니다.

왜 이게 좋은가

성능 수치

변경된 16개 owner row(ws4 @ 128 tokens, ws8 @ 512 tokens, top_k 4 + shared expert)에서 직전 #9192 대비 결과는 다음과 같습니다.

GPU	geomean	min–max
B300 (sm_103a)	1.0243	1.0084–1.0378
B200 (sm_100a)	1.0270	1.0113–1.0385

예를 들어 ws8 bf16, 512 tokens, PDL=1, NVFP4 출력 조합은 B300에서 38.23us에서 36.86us로 줄었습니다. 초기 Cake 커널(B0) 대비로는 약 1.95배 빠릅니다.

변경되지 않은 704개 row의 #9192 대비 geomean은 1.0002(B300)와 0.9996(B200)입니다. 이는 측정 노이즈 범위이며, PR 작성자도 이 구간에서는 이득을 주장하지 않는다고 명시했습니다. 측정 방식도 신뢰할 만합니다. CUPTI kernel time의 median을 쓰고, 패스별 paired ratio의 median을 취했습니다.

일반적 교훈

1. 특화 커널 + 테이블 디스패치: 범용 커널에 분기를 늘리는 대신 조건별 program을 따로 두고 lookup 테이블로 고릅니다. 다른 조합의 성능에 전혀 영향을 주지 않고 특정 구간만 정밀하게 튜닝할 수 있습니다.
2. Collective에서는 통신 프로토콜을 불변으로 유지: 새 변형이 기존과 같은 exchange를 따르기 때문에 rank별로 서로 다른 커널을 골라도 안전합니다. 분산 커널 최적화에서 가장 중요한 안전장치입니다.
3. 지연 숨기기: 통신 대기 중에 다음에 필요한 operand(shared-expert row)를 L2로 prefetch합니다. 메모리 지연과 NVLink 지연을 겹치는 전형적인 latency hiding 기법입니다.
4. 이득이 측정된 곳에만 적용: top_k == 4로 범위를 제한했습니다. 측정되지 않은 조합까지 확대하면 오히려 퇴행할 수 있다는 보수적 판단입니다.

리뷰에서는 메인테이너(yyihuang)가 tests/comm/test_cake_moe_finalize_allreduce.py와 CPU 디스패치 테스트(test_cake_moe_finalize_dispatch.py)를 CI 봇으로 실행했습니다. 커널 정확성과 디스패치 선택 로직을 각각 검증하는 구조입니다.

참고 자료

• CUDA PTX ISA: mbarrier — 생성 커널에서 사용하는 mbarrier 명령어 문서
• CUDA C++ Programming Guide — L2 캐시 및 thread block cluster 관련 공식 가이드
• CUPTI — 벤치마크 커널 시간 측정에 사용된 프로파일링 API

⚠️ 알림: 이 분석은 AI가 실제 코드 diff를 기반으로 작성했습니다.