[flashinfer] FlashInfer의 MXFP4 SiTU MoE 성능 최적화: Cluster Split-K와 CuTe DSL 개선

PR 링크: flashinfer-ai/flashinfer#6577 상태: Merged | 변경: +80137 / -48373

들어가며

최근 대규모 언어 모델(LLM)의 성능 향상을 위해 Mixture of Experts(MoE) 구조가 널리 사용되고 있습니다. 특히 MXFP4와 같은 저정밀도 연산은 추론 속도를 높이는 핵심 요소입니다. 이번 FlashInfer 업데이트는 Kimi-K3 MXFP4 SiTU(Single-Instruction-Tensor-Unit) MoE 백엔드에 대한 대규모 최적화를 포함하고 있습니다. 주요 목표는 GEMM 연산의 병렬성을 극대화하고, 스케줄링 효율을 높여 전체적인 처리량(throughput)을 개선하는 것입니다.

코드 분석

이번 PR의 핵심은 cake와 cake_cute 백엔드에서 GEMM1 연산의 병렬 처리 방식을 개선한 것입니다.

1. Cluster Split-K 도입

기존에는 단일 CTA(Cooperative Thread Array)에서 GEMM1을 수행했으나, 이번 변경을 통해 2개의 CTA가 하나의 클러스터를 형성하여 K 차원을 분할(Split-K) 처리하도록 변경되었습니다.

Before (기존 단일 CTA 방식):

"gemm1_swapab_situ_n16_nol2": {
  "cluster_split": false,
  "cta_v": 1,
  "template": "gemm1_swapab_situ_n16_nol2"
}

After (Cluster Split-K 적용):

"gemm1_swapab_situ_n16_nol2_cs": {
  "cluster_split": true,
  "cta_v": 1,
  "template": "gemm1_swapab_situ_n16_nol2_cs"
}

이 방식은 피어(peer) CTA가 계산한 중간 결과(128 x n fp32 accumulator)를 리더(leader) CTA의 공유 메모리로 st.async를 통해 전달하고, 리더가 최종 합산 후 SiTU 에필로그를 수행하게 함으로써 병렬성을 극대화합니다.

2. CuTe DSL 클러스터 랑데부(Rendezvous) 수정

cake_cute 백엔드에서는 클러스터 런칭 시 TMEM 할당 경합 문제를 해결하기 위해 랑데부 로직을 강화했습니다.

// CuTe DSL manifest update
"cluster_split": true,
"footer_dep_launch": true,
"gate": "e2e_registry"

이 수정은 prologue 단계에서 발생할 수 있는 race condition을 방지하여 커널의 안정성을 확보합니다.

왜 이게 좋은가

이번 최적화는 B300 GPU 환경에서 측정 결과, cake 백엔드에서 약 0.71.0%, cake_cute 백엔드에서 약 1.82.1%의 성능 향상을 보였습니다.

핵심 교훈:

1. Workload-aware Scheduling: 모든 행(row)에 대해 Split-K를 적용하는 것이 아니라, plan gating을 통해 특정 행 개수(T ≤ 16 등)에서만 최적의 성능을 내도록 선택적으로 적용했습니다.
2. Memory Locality: GEMM2 단계에서 late item에 대한 L2 prefetch를 도입하여 메모리 지연 시간을 줄였습니다.
3. Correctness: 성능 최적화와 동시에 compute-sanitizer를 통한 검증으로, 병렬 연산 증가에 따른 데이터 정합성 문제를 완벽히 해결했습니다.

이러한 방식은 고성능 커널 개발 시 단순히 연산 속도만 높이는 것이 아니라, 하드웨어의 클러스터 기능을 얼마나 효율적으로 활용하느냐가 핵심임을 보여줍니다.

참고 자료

• NVIDIA CuTe Documentation
• NVIDIA CUDA Programming Guide

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