[flashinfer] FlashInfer NVFP4 KV 타일 리팩(Repack)을 통한 성능 최적화

PR 링크: flashinfer-ai/flashinfer#4769 상태: Merged | 변경: +353 / -81

들어가며

LLM 추론 성능을 극대화하기 위해 FP8 및 NVFP4와 같은 저정밀 데이터 포맷이 널리 사용됩니다. 하지만 기존 FlashInfer 구현에서는 NVFP4 KV 타일을 처리할 때 매번 MMA(Matrix Multiply-Accumulate) 루프 내부에서 dequantization을 수행했습니다. 이 과정에는 magnitude table 조회, sign spread, scale-factor gather, 그리고 4-bit fragment swizzle을 위한 __shfl_sync 연산이 포함되어 있어, 연산 성능이 아닌 명령어 패치(instruction-fetch) 병목을 유발했습니다. 본 PR은 이 dequantization 과정을 루프 밖으로 분리하여 16-bit staging buffer로 미리 변환(repack)함으로써 성능을 최적화합니다.

코드 분석

include/flashinfer/attention/prefill.cuh

핵심 변경 사항은 NVFP4 dequantization 로직을 MMA 루프에서 분리하는 것입니다. 기존에는 루프 내부에서 매번 수행되던 복잡한 변환 과정을 repack_fp4_tile_to_16b 함수를 통해 타일 단위로 한 번만 수행하도록 변경했습니다.

Before (In-loop dequantization):

// 기존에는 MMA 루프 내부에서 매번 dequant 수행
// magnitude table, sign spread, scale-factor gather 등이 반복됨

After (Repack to 16-bit):

// 16-bit staging buffer로 미리 변환하여 MMA 연산 시에는 native ldmatrix 사용
__device__ __forceinline__ void repack_fp4_tile_to_16b(...) {
  // scale-factor gather 및 E4M3 변환을 타일 단위로 최적화
}

또한, SM100 이상에서는 E2M1 변환이 하드웨어 명령어로 지원되므로, 불필요한 staging buffer 할당을 방지하기 위해 아키텍처별 분기 로직을 추가했습니다.

// 아키텍처에 따른 최적화 정책 적용
constexpr bool kTargetsSoftwareE2M1 = fp4_repack_targets::any_below(1000);
constexpr bool kTargetsNativeE2M1 = fp4_repack_targets::any_at_least(1000);

왜 이게 좋은가

이번 최적화는 특히 소프트웨어 기반 E2M1 변환이 필요한 환경에서 극적인 성능 향상을 보여줍니다. SM80 환경에서 2048x2048 prefill 작업 시, 커널 실행 시간이 1.403ms에서 0.552ms로 약 2.5배 단축되었습니다. 이는 no_instruction stall 비율을 24.9%에서 1.3%로 낮춘 결과입니다.

일반적 교훈:

1. Instruction-bound vs Math-bound: 복잡한 dequantization 로직을 루프 내부에 두면 명령어 패치 병목이 발생합니다. 이를 미리 계산(pre-compute)하거나 별도의 staging buffer로 옮기는 것이 중요합니다.
2. Hardware-aware Optimization: 최신 아키텍처(SM100+)에서는 하드웨어 명령어로 처리 가능한 연산을 구형 아키텍처를 위해 소프트웨어로 구현할 때, 아키텍처별로 분기하여 불필요한 메모리(shared memory) 사용을 방지해야 합니다.

리뷰어 피드백 반영

리뷰어들은 아키텍처별로 use_kv_repack 정책을 일관되게 적용하는 것에 주목했습니다. 특히 KernelTraits와 SharedStorage가 동일한 정책을 공유하도록 하여, 런타임에 staging buffer가 할당되지 않음에도 불구하고 occupancy budget에서 이를 차감하는 오류를 방지했습니다.

참고 자료

• https://docs.nvidia.com/cuda/cuda-runtime-api/group__CUDART__TYPES.html
• https://docs.nvidia.com/cuda/parallel-thread-execution/index.html#ldmatrix

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