[flashinfer] FlashInfer의 실험적 NVFP4 어텐션 도입: SM103 최적화

PR 링크: flashinfer-ai/flashinfer#5283 상태: Merged | 변경: +3639 / -2

들어가며

최신 GPU 아키텍처인 NVIDIA SM103(GB300) 환경에서 어텐션 연산의 효율성을 극대화하기 위해 FlashInfer는 실험적으로 NVFP4(4-bit Floating Point) 지원을 추가했습니다. 이 PR은 flashinfer.prefill.prepare_nvfp4_attention API를 통해 BF16 입력을 NVFP4로 양자화하고, 이를 NVFP4AttentionRunner를 통해 실행함으로써 메모리 할당 없이 고성능 어텐션을 수행하는 경로를 제공합니다. 본 글에서는 이 최적화의 핵심 메커니즘을 분석합니다.

코드 분석

1. 양자화 및 데이터 준비 (cake_backend.py)

핵심은 BF16 입력을 NVFP4 형식으로 변환하는 과정입니다. _quantize_nvfp4 함수는 128차원 데이터를 16개 단위(SF_VEC)로 그룹화하여 스케일을 계산하고, 이를 E2M1/E4M3 형식으로 패킹합니다.

# Before/After: 데이터 양자화 및 패킹 로직
def _quantize_nvfp4(x):
    x_f32 = x.float()
    raw_scale = (blocks.abs().amax(dim=-1) / 6.0).clamp(
        min=2.0**-9, max=torch.finfo(torch.float8_e4m3fn).max
    )
    # ... (생략) ...
    packed = (pairs[..., 0] & 0x0F) | ((pairs[..., 1] & 0x0F) << 4)
    return packed.contiguous(), scale_fp8.view(torch.uint8).contiguous()

2. 실행 엔진 (NVFP4AttentionRunner)

NVFP4AttentionRunner는 JIT 컴파일된 모듈을 래핑하여, 입력이 변경되지 않는 한 반복적인 할당 없이 커널을 실행합니다. 이는 tvm_ffi를 통해 CUDA 스트림을 관리하며, 호출 시점에 미리 바인딩된 TMA(Tensor Memory Accelerator) 디스크립터를 사용합니다.

# 실행부
class NVFP4AttentionRunner:
    def launch(self):
        with tvm_ffi.use_torch_stream():
            self.entry(*self.arguments)
        return self.out

왜 이게 좋은가

이 최적화는 다음과 같은 이점을 제공합니다:

1. Zero-Allocation: prepare_nvfp4_attention을 통해 미리 준비된 Runner를 사용하면, 추론 시점(inference time)에 추가적인 CUDA 메모리 할당이나 해제가 발생하지 않습니다. 이는 실시간 추론 지연 시간을 최소화합니다.
2. 높은 처리량: SM103 아키텍처에서 2386~3099 TFLOP/s의 높은 성능을 달성했습니다. 이는 기존 BF16 연산 대비 메모리 대역폭 효율을 극대화한 결과입니다.
3. 정밀도 제어: E4M3 블록 스케일을 유한 범위로 제한하여, 오버플로우 발생 시 NaN이 발생하는 문제를 방지하는 견고함을 갖췄습니다.

일반적인 교훈으로, 고성능 커널 개발 시 데이터 레이아웃의 사전 패킹(Pre-packing)과 런타임 할당 제거가 지연 시간 감소에 결정적인 역할을 함을 확인할 수 있습니다.

리뷰 피드백 반영

리뷰 과정에서 yzh119는 절대적인 FLOPs 수치를 요구하였으며, 이에 따라 벤치마크 결과에 FLOPs per call 및 TFLOP/s 지표가 추가되어 성능 검증의 신뢰도를 높였습니다. 또한, 실험적 API임을 명시하여 사용자가 의도적으로 옵트인(opt-in)하도록 설계되었습니다.

참고 자료

• https://pytorch.org/docs/stable/generated/torch.nn.functional.scaled_dot_product_attention.html
• https://docs.nvidia.com/cuda/cuda-c-programming-guide/index.html#tma-descriptors

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