[flashinfer] Blackwell 아키텍처를 위한 MoE All-Reduce Fusion 최적화: FlashInfer의 'Cake' 백엔드 분석

2026년 9월 16일수정: 2026년 9월 16일

PR 링크: flashinfer-ai/flashinfer#4730 상태: Merged | 변경: +15863 / -10

들어가며

LLM(Large Language Model)의 크기가 커짐에 따라 MoE(Mixture of Experts) 구조는 효율적인 연산을 위한 필수적인 선택이 되었습니다. 하지만 MoE 모델은 전문가(Expert) 간의 데이터를 합치는 과정에서 발생하는 통신 오버헤드가 성능의 병목 점이 되곤 합니다. 특히 All-Reduce 연산은 노드 간 또는 GPU 간 대규모 데이터 교환을 필요로 하므로, 이를 어떻게 최적화하느냐가 전체 추론 속도를 결정짓습니다.

이번에 FlashInfer에 추가된 cake 백엔드는 NVIDIA의 최신 아키텍처인 Blackwell(SM100, SM103)을 타겟으로 하여, MoE의 All-Reduce 연산을 후속 작업인 Residual 및 RMSNorm과 융합(Fusion)함으로써 성능을 극대화했습니다. 이 PR은 단순한 기능 추가를 넘어, 최신 하드웨어의 특성을 어떻게 소프트웨어적으로 활용해야 하는지에 대한 정수를 보여줍니다.

코드 분석

1. Byte-Fidelity 유지를 위한 설정 변경

가장 먼저 눈에 띄는 변경 사항은 .pre-commit-config.yaml 파일입니다. 최적화된 커널 코드가 자동 포매팅 도구에 의해 변형되는 것을 방지하기 위해 특정 경로를 제외 처리했습니다.

Before:

exclude: ^(?:csrc/(?:cake_minimax_h3_mxfp8_pre_attention/|kda/flashkda_generated_.*|blackwell_msa/|concat_mla/|cake_all_gather_matmul/|cake_moe_finalize_allreduce_fusion/|cake_nvfp4_svdquant_gemm/|cake_selective_state_update/(?:cuda|host)/|cake_mamba_ssd_combined/generated/|fused_moe/warp_decode/generated/|kda/cake_flashkda_blackwell_evolution_.*\.cu)|flashinfer/(?:experimental/(?:cake_mxfp8_megamoe_ep16|sm110_gqa_decode)/csrc/|moe_ep/kernel_src/(?:cutedsl_megamoe|sm90/pull_style_cutedsl_megakernel|sm120/swapab_cutedsl_megakernel)/src/))

After:

exclude: ^(?:csrc/(?:cake_minimax_h3_mxfp8_pre_attention/|kda/flashkda_generated_.*|blackwell_msa/|concat_mla/|cake_all_gather_matmul/|cake_moe_finalize_allreduce_fusion/|cake_nvfp4_svdquant_gemm/|cake_selective_state_update/(?:cuda|host)/|cake_mamba_ssd_combined/generated/|cake_trtllm_moe_allreduce_fusion/|fused_moe/warp_decode/generated/|kda/cake_flashkda_blackwell_evolution_.*\.cu)|flashinfer/(?:experimental/(?:cake_mxfp8_megamoe_ep16|sm110_gqa_decode)/csrc/|moe_ep/kernel_src/(?:cutedsl_megamoe|sm90/pull_style_cutedsl_megkernel|sm120/swapab_cutedsl_megakernel)/src/))

왜 이게 좋은가? 이 커널들은 JIT(Just-In-Time) 컴파일 시 SHA-256 다이제스트를 통해 무결성을 검증받습니다. pre-commit 훅이 코드를 미세하게 수정(공백 추가 등)하면 체크섬이 달라져 런타임 에러가 발생할 수 있습니다. 따라서 생성된 소스 코드의 Byte-faithful(바이트 단위 일치)을 보장하는 것은 안정적인 배포를 위해 필수적입니다.

2. Blackwell 전용 All-Reduce Fusion 커널

새롭게 추가된 cake_trtllm_moe_allreduce_fusion_kernels.cu는 Blackwell 아키텍처의 핵심 기능인 Cluster-level communication을 활용합니다.

핵심 코드 인용 (PTX 기반 최적화):

__device__ __forceinline__ uint32_t mapa_to_rank(uint32_t local_addr, uint32_t rank) {
    uint32_t remote;
    asm volatile("mapa.shared::cluster.u32 %0, %1, %2;"
        : "=r"(remote) : "r"(local_addr), "r"(rank));
    return remote;
}

__global__ __launch_bounds__(224) __cluster_dims__(4,1,1) void
kernel_cake_trtllm_moe_reduction_float16_ws2_o0110(...)

분석:

__cluster_dims__(4,1,1): Blackwell 아키텍처는 여러 개의 SM(Streaming Multiprocessor)을 하나의 클러스터로 묶어 리소스를 공유할 수 있게 합니다. 이 커널은 4개의 CTA(Cooperative Thread Array)를 하나의 클러스터로 정의하여 하드웨어 가속을 활용합니다.
mapa.shared::cluster.u32: 이 PTX 명령어는 클러스터 내 다른 SM의 Shared Memory 주소를 직접 매핑합니다. 기존의 Global Memory를 거치는 통신 방식보다 훨씬 낮은 지연 시간(Latency)으로 데이터를 교환할 수 있게 해줍니다.
Fusion 전략: 커널 내부에서 MoE 전문가들의 출력을 합산(packed_acc)하는 동시에, residual을 더하고 RMSNorm까지 한 번의 커널 실행으로 처리합니다. 이는 메모리 대역폭(Memory Bandwidth) 소모를 획기적으로 줄여줍니다.

왜 이게 좋은 최적화인가?

1. 압도적인 성능 향상

PR 설명에 따르면, B200(SM100)에서 1.122930×, B300(SM103)에서 **1.121959×**의 기하평균(Geometric Mean) 속도 향상을 기록했습니다. 특히 특정 케이스에서는 최대 **1.20×**에 달하는 성능 향상을 보였습니다. 이는 통신과 연산을 융합함으로써 얻을 수 있는 이득이 얼마나 큰지를 증명합니다.

2. 하드웨어 특화 스케줄링

이 PR은 단순히 하나의 커널을 만든 것이 아니라, 하드웨어 특성에 맞춘 세부 스케줄링을 도입했습니다.

SM103: Single-token 추론 상황에 최적화된 스케줄을 적용하여 지연 시간을 최소화했습니다.
SM100: World size 8 환경에서 64/128 토큰을 처리할 때, Peer Store를 수행하기 전 'Publication Packing'을 한 번만 계산하도록 설계하여 중복 연산을 제거했습니다.

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

PR Analysis 의 다른글

이전글 [ultralytics] RT-DETR의 OpenVINO INT8 양자화 정확도 문제 해결: NNCF Transformer 모드 적용
현재글 : [flashinfer] Blackwell 아키텍처를 위한 MoE All-Reduce Fusion 최적화: FlashInfer의 'Cake' 백엔드 분석
다음글 [ultralytics] YOLO 학습 성능 25% 향상: B200 GPU 최적화 분석