[flashinfer] FlashInfer Kimi-K3 Fused Router 최적화: GB300 및 B200에서 최대 1.22배 성능 향상

PR 링크: flashinfer-ai/flashinfer#7748 상태: Merged | 변경: +9920 / -4828

들어가며

최신 대규모 언어 모델(LLM)의 핵심 기술인 Mixture-of-Experts(MoE)는 추론 시 효율적인 전문가 라우팅이 성능을 좌우합니다. FlashInfer의 kimi_k3_fused_router는 FP32 로짓과 바이어스를 입력받아 Top-16 전문가를 선택하고, 그룹화된 MoE GEMM을 위한 라우트 플랜을 단일 커널 실행으로 생성합니다. 이번 업데이트는 NVIDIA GB300 및 B200 아키텍처에서 특정 토큰 배치 범위에 최적화된 커널 프로그램을 도입하여, 특히 작은 배치 사이즈에서 최대 1.22배의 성능 향상을 달성했습니다.

코드 분석

1. cake_backend.py: 커널 디스패치 전략 고도화

이번 변경의 핵심은 토큰 수와 하드웨어 아키텍처에 따라 더 세분화된 커널 프로그램을 선택하도록 디스패치 로직을 개선한 것입니다.

# Before (기존)
# 4-CTA-cluster family kernels...

# After (개선)
# Added Q4S_E16, Q4SP_PG, Q4SP_LR, Q4SPW_LR variants
# Q4S_E16: plan emit gives each lane 16 tokens instead of 32
# Q4SP_PG: skips the emit word pass for batches <= 1024

특히 Q4S_E16은 기존 32토큰 단위의 emit을 16토큰으로 줄여 레지스터 압박을 완화했고, Q4SP_PG는 비트맵 워드 생성이 불필요한 경우 해당 패스를 건너뛰어 오버헤드를 줄였습니다.

2. README.md 및 데이터 정렬 검증

성능 최적화뿐만 아니라 안정성을 위해 16바이트 정렬(alignment) 검증이 추가되었습니다. 커널이 int4 및 float4 벡터 연산을 수행하므로, 입력 텐서의 메모리 주소가 16바이트로 정렬되어야 합니다.

# 16-byte alignment validation logic
if tensor.data_ptr() % 16 != 0:
    raise ValueError("tensor must be 16-byte aligned")

왜 이게 좋은가

이번 최적화는 단순히 알고리즘을 개선한 것이 아니라, GPU 아키텍처별로 최적화된 커널(Kernel Specialization)을 토큰 수에 따라 동적으로 선택하도록 설계되었습니다.

• 성능 수치: GB300에서 1토큰/block_m 16 조건 시 기존 3.232us에서 2.656us로 1.22배 단축되었습니다.
• 교훈:
  1. 메모리 정렬의 중요성: 벡터 로드/스토어(ld.global.nc.v4, cp.async)를 활용할 때 16바이트 정렬은 필수적이며, 이를 강제함으로써 메모리 대역폭 활용도를 극대화할 수 있습니다.
  2. 커널 세분화(Kernel Specialization): 모든 토큰 수에 대해 하나의 범용 커널을 사용하는 대신, 특정 토큰 수 범위에 최적화된 작은 커널들을 다수 배치함으로써 캐시 히트율과 레지스터 효율을 높였습니다.

결론

이번 FlashInfer의 업데이트는 MoE 추론 파이프라인에서 라우팅 단계가 병목이 되지 않도록 정교하게 튜닝된 사례입니다. 특히 하드웨어 특성(sm_100a vs sm_103a)을 고려한 커널 선택 로직은 고성능 CUDA 프로그래밍의 좋은 본보기입니다.

참고 자료

• https://pytorch.org/docs/stable/generated/torch.Tensor.data_ptr.html
• https://docs.nvidia.com/cuda/cuda-c-programming-guide/index.html#alignment-requirements

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