본문으로 건너뛰기

[flashinfer] FlashInfer FP8 GEMM 성능 최적화: 스케일 곱셈을 커널 내부로 통합

PR 링크: flashinfer-ai/flashinfer#6800 상태: Merged | 변경: +82 / -41

들어가며

대규모 언어 모델(LLM)의 추론 성능은 행렬 곱셈(GEMM) 연산의 효율성에 크게 좌우됩니다. 특히, FP8(8비트 부동소수점) 데이터 타입을 활용한 GEMM은 메모리 대역폭과 연산량을 줄여 추론 속도를 크게 향상시킬 수 있는 잠재력을 가지고 있습니다. 하지만 FP8 연산 후에는 종종 스케일링(scaling) 과정이 필요하며, 이 과정이 제대로 최적화되지 않으면 전체 성능에 병목 현상을 일으킬 수 있습니다.

이번 PR(Pull Request)은 FlashInfer 라이브러리에서 FP8 GEMM 연산의 후처리 단계(epilogue)에서 발생하는 스케일링 오버헤드를 제거하여 성능을 개선하는 것을 목표로 합니다. 기존에는 FP8 GEMM의 결과에 적용되는 두 개의 스케일(a_scale, b_scale)을 호스트(CPU)에서 곱한 후, 그 결과를 GPU 커널로 전달했습니다. 이 과정은 불필요한 커널 실행과 데이터 전송 오버헤드를 유발했습니다. 이 PR은 이 스케일 곱셈 연산을 GPU 커널 내부로 통합함으로써 이러한 오버헤드를 제거합니다.

코드 분석: 파일별 변경 사항

이번 PR의 핵심 변경 사항은 flashinfer/gemm/kernels/bmm_fp8_blackwell.py, flashinfer/gemm/kernels/bmm_fp8_rubin.py, 그리고 flashinfer/gemm/kernels/bmm_fp8_wrapper.py 파일에 집중되어 있습니다.

1. flashinfer/gemm/kernels/bmm_fp8_blackwell.py 및 flashinfer/gemm/kernels/bmm_fp8_rubin.py

이 파일들은 각각 Blackwell 및 Rubin 아키텍처에 특화된 FP8 GEMM 커널 구현을 포함합니다. 주요 변경 내용은 다음과 같습니다:

  • __call__ 메소드 시그니처 변경: 기존에는 output_scale_tensor라는 단일 인자를 받았지만, 이제는 a_scale_tensor와 b_scale_tensor라는 두 개의 인자를 받도록 변경되었습니다.

    -        output_scale_tensor: cute.Tensor,
+        a_scale_tensor: cute.Tensor,
+        b_scale_tensor: cute.Tensor,

    이는 스케일링 연산이 더 이상 호스트에서 미리 계산되어 전달되지 않음을 의미합니다.

  • kernel 함수 내부 로직 변경: GPU 커널 내부에서 스케일 곱셈이 수행되도록 로직이 수정되었습니다.

    -        output_scale = output_scale_tensor[0]
+        output_scale = a_scale_tensor[0] * b_scale_tensor[0]

    기존에는 output_scale_tensor에서 단일 스케일 값을 로드했지만, 이제는 a_scale_tensor와 b_scale_tensor에서 각각 값을 로드하여 곱한 후, 그 결과를 output_scale로 사용합니다. 이 연산은 각 CUDA 스레드 블록(CTA) 내에서 수행됩니다.

  • bmm 래퍼 함수 변경: bmm 함수 역시 output_scale_tensor 대신 a_scale_tensor와 b_scale_tensor를 인자로 받도록 수정되었습니다.

    -    output_scale_tensor: cute.Tensor,
+    a_scale_tensor: cute.Tensor,
+    b_scale_tensor: cute.Tensor,

    -    gemm_op(a, b, c, max_active_clusters, stream, output_scale_tensor, epilogue_op)
+    gemm_op(
+        a,
+        b,
+        c,
+        max_active_clusters,
+        stream,
+        a_scale_tensor,
+        b_scale_tensor,
+        epilogue_op,
+    )

    이 변경은 bmm 함수가 스케일 텐서를 직접 커널로 전달하며, 스케일 곱셈은 GPU에서 이루어짐을 명확히 합니다.

2. flashinfer/gemm/kernels/bmm_fp8_wrapper.py

이 파일은 커널 컴파일 및 래핑 로직을 담당합니다. 주요 변경 사항은 다음과 같습니다:

  • _create_fake_tensors 함수 수정: CUDA 그래프 호환성을 위해 생성되는 가짜(fake) 텐서 목록에 scale_fake 하나 대신 a_scale_fake와 b_scale_fake 두 개의 텐서가 생성되도록 변경되었습니다.

    -    scale_fake = cute.runtime.make_fake_compact_tensor(
+    # One fake tensor per input scale; the kernel multiplies them on device.
+    a_scale_fake = cute.runtime.make_fake_compact_tensor(
+        cutlass.Float32,
+        (1,),
+        assumed_align=4,
+    )
+    b_scale_fake = cute.runtime.make_fake_compact_tensor(
         cutlass.Float32,
         (1,),
         assumed_align=4,
     )

-    return a_fake, b_fake, c_fake, scale_fake, stream_fake
+    return a_fake, b_fake, c_fake, a_scale_fake, b_scale_fake, stream_fake

    이는 런타임 시 두 개의 스케일 텐서가 전달될 것임을 반영합니다.

  • _compile_and_create_tensor_api 함수 수정: 컴파일 시 전달되는 가짜 텐서 목록에 a_scale_fake와 b_scale_fake가 포함되도록 변경되었습니다.

    -        gemm, a_fake, b_fake, c_fake, scale_fake, stream_fake, cluster_shape_mn
+        gemm, a_fake, b_fake, c_fake, a_scale_fake, b_scale_fake, stream_fake, cluster_shape_mn

  • tensor_api 함수 시그니처 및 호출 변경: 최종적으로 사용자에게 노출되는 tensor_api 함수의 시그니처가 a_scale_tensor와 b_scale_tensor를 받도록 변경되었고, 커널 호출 시에도 이 두 텐서를 전달하도록 수정되었습니다.

    -        a_tensor: torch.Tensor,
-        b_tensor: torch.Tensor,
-        c_tensor: torch.Tensor,
-        output_scale_tensor: torch.Tensor,
+        a_tensor: torch.Tensor,
+        b_tensor: torch.Tensor,
+        c_tensor: torch.Tensor,
+        a_scale_tensor: torch.Tensor,
+        b_scale_tensor: torch.Tensor,

    -            a_tensor, b_tensor, c_tensor, output_scale_tensor,
+            a_tensor, b_tensor, c_tensor, a_scale_tensor, b_scale_tensor,

    이 변경은 CUDA 그래프 캡처 시 호스트-디바이스 동기화 없이 스케일 텐서를 GPU에서 직접 사용할 수 있도록 보장합니다.

왜 이게 좋은가: 성능 향상과 일반화된 교훈

성능 향상

이 PR의 가장 큰 장점은 불필요한 커널 실행을 제거하여 성능을 향상시킨다는 점입니다. 기존 방식에서는 FP8 GEMM 연산 후, 스케일 곱셈을 위해 별도의 vectorized_elementwise_kernel이 호출되었습니다. 이는 특히 작은 M 차원을 가지는 GEMM 연산에서 두드러지는 고정적인 커널 실행 지연 시간(launch latency)을 유발했습니다. 리뷰어 jimmyzho가 제공한 벤치마크 결과에 따르면, 이 변경으로 인해 평균 1.11x ~ 1.12x의 속도 향상이 관찰되었습니다. 이는 특히 M 차원이 작은 경우에 더 큰 성능 향상을 보이며, 가장 큰 M x N x K 형태 (9216 x 8192 x 1280)에서도 1.07x ~ 1.08x의 향상을 기록했습니다.

| m x n x k | cute-dsl untuned, pr vs base | cute-dsl tuned, pr vs base |
|---|---|---|
| 9216 x 8192 x 1280 | 1.07x | 1.08x |
| 2048 x 2560 x 8192 | 1.18x | 1.15x |
| ... |
| **geomean (15 shapes)** | **1.11x** | **1.12x** |

torch.profiler 트레이스 결과에서도 이 점이 명확히 드러납니다. 변경 전에는 cute-dsl 호출이 두 개의 CUDA 커널(vectorized_elementwise_kernel과 SM107PersistentDenseGemmKernel)로 구성되었지만, 변경 후에는 GEMM 커널(SM107PersistentDenseGemmKernel) 하나만 실행됩니다. GEMM 커널 자체의 실행 시간은 변하지 않았으므로, 성능 향상은 전적으로 불필요한 커널 실행 제거에서 비롯됩니다.

또한, 이 변경은 CUDA 그래프(CUDA graph) 캡처 동작에 영향을 주지 않습니다. 스케일 텐서가 여전히 GPU에서 커널 실행 시점에 로드되므로, 호스트-디바이스 동기화 없이 CUDA 그래프를 효율적으로 사용할 수 있습니다.

일반화된 교훈

  1. 커널 퓨전(Kernel Fusion)의 중요성: 연산의 작은 부분이라도 호스트와 디바이스 간의 통신 또는 별도의 커널 호출을 유발한다면, 이를 GPU 커널 내부로 통합하는 것이 성능 향상의 지름길이 될 수 있습니다. 특히 반복적으로 호출되는 함수나 연산에서 이러한 최적화는 큰 효과를 발휘합니다.
  2. FP8 연산의 디테일: FP8과 같은 저정밀도 연산은 성능 향상의 큰 이점을 제공하지만, 스케일링과 같은 후처리 단계의 최적화가 전체 성능을 좌우할 수 있습니다. 스케일링 파라미터의 관리 및 적용 방식을 신중하게 설계해야 합니다.
  3. CUDA 그래프 호환성: 고성능 추론 시스템에서는 CUDA 그래프를 활용한 오버헤드 감소가 필수적입니다. 변경 사항이 CUDA 그래프의 동작을 방해하지 않도록 주의해야 하며, 이 PR은 이를 성공적으로 달성했습니다.
  4. 작은 연산의 누적 효과: 비록 개별 스케일 곱셈 연산 자체는 매우 작지만, 수백만 번의 GEMM 호출에서 이러한 작은 연산들이 누적되면 상당한 성능 향상을 가져올 수 있습니다. 특히 LLM과 같이 반복적인 연산이 많은 워크로드에서는 이러한 최적화가 더욱 중요합니다.

References

  • CuTe-DSL - FlashInfer에서 사용하는 DSL (Domain Specific Language)입니다.
  • [cutlass/epilogue](https://github.com/ NVIDIA/cutlass/blob/main/include/cutlass/epilogue/epilogue.h) - CUTLASS 라이브러리의 에필로그 구현 관련 문서입니다. (정확한 함수 링크는 찾기 어려우나, 에필로그 개념 이해에 도움이 됩니다.)
  • torch.compile - PyTorch의 컴파일 기능으로, 이 라이브러리에서 사용되는 JIT 컴파일과 유사한 개념입니다.
  • FP8 Data Format - NVIDIA의 FP8 데이터 포맷 및 활용에 대한 문서입니다.

참고 자료

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

댓글

관련 포스트

PR Analysis 의 다른글