[flashinfer] FlashInfer MoE 성능 최적화: CuTe DSL을 활용한 Deterministic Finalize 개선기

2026년 10월 2일수정: 2026년 10월 2일

PR 링크: flashinfer-ai/flashinfer#5319 상태: Merged | 변경: +307 / -8

들어가며

LLM 서빙에서 MoE(Mixture of Experts) 아키텍처는 효율적인 파라미터 활용을 가능케 하지만, 여러 전문가(Expert)의 출력을 하나로 합치는 Finalize 단계에서 병목이 발생하곤 합니다. 특히 수치적 결정론(Determinism)을 보장해야 하는 use_fused_finalize=False 설정(unfused mode)에서는 성능 저하가 더 두드러집니다.

이번 포스팅에서는 flashinfer-ai/flashinfer 레포지토리에 반영된 W4A4 및 W4A16 MoE의 unfused finalize 최적화 PR을 분석합니다. 이 PR은 기존의 standalone finalizer를 CuTe DSL 기반의 공유 커널로 교체하고, PDL(Persistent Data Loading) 스케줄링을 개선하여 Blackwell 아키텍처 등 최신 GPU에서 유의미한 성능 향상을 이끌어냈습니다.

코드 분석: 무엇이 어떻게 바뀌었나?

1. CuTe DSL 기반의 새로운 Finalize 커널 도입 (`moe_finalize.py` 신설)

가장 큰 변화는 기존의 네이티브 standalone finalizer를 대체하는 _MoeFinalizeKernel의 도입입니다. 이 커널은 CuTe DSL을 사용하여 하드웨어 친화적인 데이터 레이아웃과 연산을 정의합니다.

[Before: 기존 방식 (개념적)] 기존에는 별도의 네이티브 커널이 독립적으로 실행되며 전문가들의 출력을 BF16으로 처리했습니다. 이 과정에서 FP32 정밀도 제어나 스케줄링 최적화가 제한적이었습니다.

[After: CuTe DSL 기반 최적화 커널]

@dsl_user_op
def _fma_rn(a, b, c, *, loc=None, ip=None):
    # Match the native finalizer's --use_fast_math FP32 FMA, including FTZ.
    return cutlass.Float32(
        llvm.inline_asm(
            T.f32(),
            [cutlass.Float32(x).ir_value(loc=loc, ip=ip) for x in (a, b, c)],
            "fma.rn.ftz.f32 $0, $1, $2, $3;",
            "=f,f,f,f",
            has_side_effects=False,
            asm_dialect=llvm.AsmDialect.AD_ATT,
            loc=loc,
            ip=ip,
        )
    )

위 코드처럼 fma.rn.ftz.f32 인라인 어셈블리를 직접 사용하여, 네이티브 --use_fast_math와 동일한 Flush-To-Zero(FTZ) 동작을 수행하면서도 FP32 정밀도로 연산한 뒤 마지막에 BF16으로 변환합니다. 이는 수치적 정확도와 성능 사이의 균형을 완벽하게 맞춘 설계입니다.

2. PDL(Persistent Data Loading) 스케줄링 최적화

PDL은 데이터를 미리 로드하여 연산 유닛의 유휴 시간을 줄이는 기술입니다. 이번 PR에서는 FC2(두 번째 Fully Connected 레이어)의 실행과 Finalizer의 준비 과정을 겹치도록(overlap) 개선했습니다.

[Before: blockscaled_contiguous_grouped_gemm_finalize_fusion.py]

# 기존에는 FC2 연산이 완전히 끝난 후 혹은 동기화 시점에 의존성 제어가 이루어짐

[After: 의존성 힌트 조기 발행]

# The standalone finalizer can prepare its independent state while
# FC2 runs. Its dependency wait still protects every FC2 output load.
if cutlass.const_expr(self.enable_pdl and not self.use_fused_finalize):
    if warp_idx == self.mma_warp_id:
        with cute.arch.elect_one():
            griddepcontrol_launch_dependents()

FC2의 입력 의존성 대기가 끝난 직후 griddepcontrol_launch_dependents()를 호출하여, 후속 Finalizer가 메타데이터를 미리 로드할 수 있도록 허용합니다. 실제 FC2 출력 데이터를 읽기 전까지는 Finalizer가 대기하므로 데이터 안정성은 유지하면서도 준비 시간을 단축했습니다.

3. 벤치마크 플래그 정합성 수정

기존에는 --no-fused-finalize 플래그를 주어도 CUTLASS 백엔드에서 무시되는 경우가 있었습니다. 이를 명시적으로 전달하도록 수정되었습니다.

[Before: bench_moe_deepseek.py]

def bench_cutlass(
    # ...
    do_autotune=True,
    include_activation_quant=False,
    # use_fused_finalize 인자가 누락됨
):

[After: 플래그 전달 추가]

def bench_cutlass(
    # ...
    do_autotune=True,
    include_activation_quant=False,
    use_fused_finalize=True, # 인자 추가
):
    # ...
    results = bench['run_benchmark'](
        # ...
        use_fused_finalize=use_fused_finalize,
    )

왜 이게 좋은 최적화인가?

1. 성능 수치상의 이득

벤치마크 결과에 따르면, 특히 토큰 수가 적은 구간(1~2 tokens)에서 성능 향상이 두드러집니다.

W4A4 per-tensor: 1 token 기준 1.1118x 향상
W4A16: 1 token 기준 1.0783x 향상

대규모 배치(16384 tokens)에서도 W4A16의 경우 1.0506x의 성능 향상을 보여주며, 전 구간에서 regression 없이 고른 개선을 달성했습니다.

병목 지점의 재정의: 단순히 연산 속도만 줄이는 것이 아니라, 연산 간의 대기 시간(Latency)을 줄이기 위한 스케줄링 최적화가 얼마나 중요한지 알 수 있습니다.
도구의 적재적소 활용: CuTe DSL은 복잡한 인덱싱을 단순화해주지만, 성능의 핵심인 FMA 연산 등은 직접 제어하여 최적의 결과를 냈습니다.

MoE 기반 모델을 서빙하거나 고성능 CUDA 커널을 작성해야 하는 개발자라면, 이 PR의 moe_finalize.py 구현을 반드시 참고해보시기 바랍니다.

참고 자료

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

PR Analysis 의 다른글

이전글 [flashinfer] FlashInfer, Blackwell GPU를 위한 DeepSeek-V4 NVFP4 Sparse-MLA 디코드 최적화: Cake 백엔드 도입
현재글 : [flashinfer] FlashInfer MoE 성능 최적화: CuTe DSL을 활용한 Deterministic Finalize 개선기
다음글 [flashinfer] FlashInfer의 CAKE 커널 최적화: 아키텍처 및 형상별 프로그램 선택 전략