[flashinfer] FlashInfer에 cuTile 기반 Fused MoE 백엔드 도입: 성능과 유지보수성의 균형

PR 링크: flashinfer-ai/flashinfer#4646 상태: Merged | 변경: +6179 / -0

들어가며

Mixture-of-Experts(MoE) 모델은 대규모 언어 모델의 추론 효율성을 높이는 핵심 기술입니다. 하지만 다양한 GPU 아키텍처와 정밀도(Precision) 환경에서 최적의 성능을 유지하는 것은 매우 어려운 과제입니다. 기존의 cutlass_fused_moe는 강력하지만, 커스텀 최적화와 유지보수 측면에서 한계가 있었습니다. 이번 PR에서는 FlashInfer의 Unified MoE API에 cuTile 백엔드를 통합하여, 다양한 GPU(SM89, SM90, SM120, SM121) 환경에서 더 나은 이식성과 성능을 확보하고자 했습니다.

코드 분석

1. Unified MoE API 확장 (benchmarks/routines/unified_moe.py)

기존 CUTLASS 백엔드와 동일한 인터페이스를 사용하여 cuTile을 통합했습니다. 이를 통해 사용자는 동일한 입력 데이터로 두 백엔드의 성능을 쉽게 비교할 수 있습니다.

# Before: CUTLASS 전용 설정
# After: cuTile 설정 추가 및 런타임 필터링
_BACKEND_CONFIGS = {
    ("bf16", "cutlass"): CutlassBf16Config,
    ("bf16", "cutile"): CuTileBf16Config,
    ("nvfp4", "cutlass"): CutlassNvfp4Config,
    ("nvfp4", "cutile"): CuTileNvfp4Config,
}

2. cuTile NVFP4 양자화 지원

cuTile은 NVFP4 W4A4 포맷을 지원하기 위해 별도의 양자화 로직을 구현했습니다. fp4_quantize를 활용하여 BF16 가중치를 최적화된 레이아웃으로 변환합니다.

# cuTile NVFP4 양자화 로직
def _quantize_cutile_nvfp4_source(weight: torch.Tensor):
    # ... (생략) ...
    packed, scale = fp4_quantize(
        weight[expert],
        global_scale=global_scales[expert : expert + 1],
        sf_vec_size=16,
        is_sf_swizzled_layout=False,
        enable_pdl=False,
    )
    return packed_experts, scale_experts, global_scales

왜 이게 좋은가

이번 최적화의 핵심은 '이식성(Portability)'과 '성능 경쟁력'입니다.

1. 성능: RTX PRO 6000(SM120) 환경에서 Qwen3.6 모델 기준, cuTile은 CUTLASS 대비 최대 1.103배의 속도 향상을 보였습니다. 특히 대규모 토큰 처리 시 성능 이점이 두드러집니다.
2. 유지보수성: cuTile은 FlashInfer의 자체적인 커널 작성 방식을 따르므로, 외부 라이브러리인 CUTLASS에 대한 의존도를 낮추고 디버깅 경험을 개선합니다.
3. 자동 튜닝: staged GEMM-pair autotuning을 도입하여, 각 GPU 아키텍처와 입력 형상(Shape)에 맞는 최적의 커널 설정을 자동으로 선택합니다.

교훈: 범용 라이브러리(CUTLASS)를 사용하는 것과 특정 프레임워크에 최적화된 커널(cuTile)을 직접 구현하는 것 사이에서, Unified API를 통해 두 방식을 병행함으로써 성능과 안정성을 모두 잡을 수 있습니다.

리뷰 피드백 반영

리뷰 과정에서 _gemm_configs가 형상 필터링 시 모든 후보를 제거할 경우 발생하는 문제를 해결하기 위해, 명시적인 ValueError를 발생시켜 디버깅을 용이하게 했습니다. 또한 CI 파이프라인에서 발생하는 타임아웃 문제를 해결하여 안정적인 테스트 통과를 보장했습니다.

참고 자료

• https://pytorch.org/docs/stable/generated/torch.compile.html
• https://github.com/flashinfer-ai/flashinfer

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