[flashinfer] Blackwell 아키텍처를 위한 MXFP8/MXFP4 기반 고성능 MoE 추론 최적화

PR 링크: flashinfer-ai/flashinfer#4529 상태: Merged | 변경: +1332 / -16

들어가며

최신 대규모 언어 모델(LLM)의 Mixture-of-Experts(MoE) 구조는 추론 시 엄청난 메모리 대역폭을 요구합니다. 특히 가중치(Weight)가 메모리 병목의 주범이 되는 상황에서, FlashInfer는 Blackwell(SM100) 아키텍처를 활용하여 MXFP8 활성화(Activation)와 MXFP4 가중치를 결합한 새로운 W4A8 백엔드를 도입했습니다. 이 PR은 통신 비용을 2배 절감하는 MXFP8 패킹 디스패치와 CuTeDSL 기반의 커널을 통해 추론 효율을 극대화합니다.

코드 분석

1. 새로운 커널 백엔드: sm100_mxfp8_mxfp4_bf16_cutedsl

flashinfer/moe_ep/backends/split/kernel/sm100/mxfp8_mxfp4_bf16_cutedsl/ 디렉토리에 추가된 이 백엔드는 MXFP4 가중치와 MXFP8 활성화를 사용하여 BF16 출력을 생성합니다.

# Before: BF16 기반의 일반적인 split kernel 호출
# After: MXFP8/MXFP4 전용 CuTeDSL 커널 로직 적용
class Sm100_Mxfp8_Mxfp4_Bf16_Cutedsl_SplitConfig(SplitConfig):
    # ... (tactic 기반의 커널 최적화 및 pre-quantized weight 처리)

2. MXFP8 패킹 디스패치 (NCCL-EP)

nccl_ep 통신 계층에서 MXFP8 데이터를 패킹하여 전송함으로써, 기존 BF16 전송 대비 데이터 전송량을 획기적으로 줄였습니다.

# Before: BF16 텐서 전송 후 수신측에서 양자화
# After: 송신측에서 미리 패킹된 MXFP8 payload 전송
+ # [H] fp8 payload와 [H/32] UE8M0 scale bytes를 하나의 행으로 패킹
+ dispatch_packed_mxfp8(tokens, scales, nccl_ep_handle)

왜 이게 좋은가

1. 통신 비용 절감: 기존 BF16 디스패치 방식은 각 랭크에서 양자화가 필요했으나, 이제는 소스 랭크에서 한 번만 양자화하고 패킹된 상태로 전송합니다. 이는 네트워크 와이어 비용을 약 2배 절감합니다.
2. 메모리 효율성: MXFP4 가중치 사용으로 가중치 메모리 점유율을 낮추면서도, Blackwell의 하드웨어 가속을 활용해 BF16 수준의 출력 정밀도를 유지합니다.
3. 검증된 정확도: Torch Oracle을 통해 수치적 허용 오차 범위 내에서 결과가 일치함을 보장하며, 4-rank 테스트를 통해 분산 환경에서의 안정성을 확보했습니다.

이 최적화는 '통신과 연산의 분리'라는 Split-path 전략을 극대화한 사례로, 향후 대규모 MoE 모델 서빙에서 메모리 대역폭 병목을 해결하는 표준적인 접근 방식이 될 것입니다.

리뷰 피드백 반영

리뷰 과정에서 ht_recv_bufs의 shape 검증 로직이 강화되었으며, require_nccl_ep 호출 시 실패 전파(propagation)가 누락되었던 버그가 수정되었습니다. 또한, 커널 팀의 vendored 코드 스타일을 유지하기 위해 불필요한 f-string 리팩토링을 지양하는 등 유지보수성을 고려한 논의가 이루어졌습니다.

참고 자료

• FlashInfer Documentation
• CUTLASS CuTeDSL

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