[flashinfer] FlashInfer FP8 GEMM 최적화: 윈도우 세그먼트 KG1 프로그램으로 GPU 성능 극대화

PR 링크: flashinfer-ai/flashinfer#7877 상태: Merged | 변경: +1560 / -19

들어가며

최근 대규모 언어 모델(LLM)의 발전과 함께, 모델 추론 과정에서의 효율적인 행렬 곱셈(GEMM) 연산은 GPU 성능을 극대화하는 핵심 요소가 되었습니다. 특히, FP8(8비트 부동소수점)과 같은 저정밀도 데이터 타입을 활용한 GEMM은 메모리 대역폭과 연산 속도 측면에서 큰 이점을 제공하며, FlashInfer와 같은 고성능 추론 라이브러리에서 활발히 연구되고 있습니다.

이번에 FlashInfer 프로젝트에 병합된 PR(perf(cake_grouped_fp8_gemm): window-segment KG1 program for validated FP32-scale layouts with expert runs inside a 128-row window)은 cake_grouped_fp8_gemm 커널의 특정 시나리오에서의 성능을 획기적으로 개선하는 것을 목표로 합니다. 이 커널은 FP8 E4M3 입력과 BF16 출력을 사용하며, m_indices에 의해 행이 그룹화되는 특징을 가집니다. 특히, FP32 per-128 block scales을 사용하는 경우 중, "expert runs"가 128-row window 내에서 시작되는 특정 레이아웃에서 성능 병목이 발생했습니다. 이 PR은 이러한 특정 조건에 최적화된 새로운 KG1(Kernel Group 1) 프로그램을 도입하여, 해당 시나리오에서 최대 1.4배의 속도 향상을 달성했습니다.

코드 분석: 무엇이 달라졌는가?

이 PR의 핵심 변경사항은 두 가지 주요 영역에서 이루어졌습니다. 첫째, 새로운 최적화된 CUDA 커널을 위한 호스트 바인딩 파일이 추가되었습니다. 둘째, FlashInfer의 Python 코드 내에서 이 새로운 커널을 선택하는 로직이 업데이트되었습니다.

1. 새로운 CUDA 커널 바인딩: csrc/cake_grouped_fp8_gemm/cake_grouped_fp8_gemm_b51037becc0c4cf72026_binding.cu

이 파일은 kernel_cake_grouped_fp8_gemm_b51037becc0c4cf72026이라는 새로운 CUDA 커널을 호스트(CPU)에서 호출할 수 있도록 하는 바인딩 코드입니다. 주목할 점은 NVIDIA의 TMA(Tensor Memory Accelerator)를 활용하기 위해 CUtensorMap을 사용하는 방식입니다. TMA는 GPU의 L2 캐시를 우회하여 전역 메모리에서 데이터를 효율적으로 로드하고 저장할 수 있게 해주는 하드웨어 기능으로, 대규모 행렬 연산에서 메모리 대역폭 활용도를 극대화하는 데 필수적입니다.

아래는 CUtensorMap을 사용하여 텐서 A의 데이터를 TMA 친화적인 방식으로 인코딩하는 EncodeTma_A 함수의 일부입니다.

--- /dev/null
+++ b/csrc/cake_grouped_fp8_gemm/cake_grouped_fp8_gemm_b51037becc0c4cf72026_binding.cu
@@ -0,0 +1,283 @@
+// ... (Copyright and License headers)
+
+#include <cuda_runtime.h>
+#include <cuda.h>
+#include <cuda_bf16.h>
+
+#include "cake_ffi_checks.h"
+
+#include <cstdint>
+
+extern "C" __global__ void kernel_cake_grouped_fp8_gemm_b51037becc0c4cf72026(const __grid_constant__ CUtensorMap A, const __grid_constant__ CUtensorMap B, const __grid_constant__ CUtensorMap C_tma, __nv_bfloat16* __restrict__ C, float* __restrict__ a_scale, float* __restrict__ b_scale, int* __restrict__ m_indices, int M, int N, int K, int G);
+
+namespace cake_binding_kernel_cake_grouped_fp8_gemm_b51037becc0c4cf72026 {
+
+// ... (CheckCuda, CheckCu, DataPtr, CheckUnitInnerStride, NumelAbove, CheckedMul, CDiv, CMod, Extent, Stride functions)
+
+CUtensorMap EncodeTma_A(const TensorView& t) {
+  TVM_FFI_CHECK(t.ndim() >= 2, ValueError)
+      << KERNEL_NAME << ": descriptor 'A' takes a source of at least 2 dims, got " << t.ndim();
+  CheckUnitInnerStride(t, "A");
+  int64_t d1 = t.size(-1);
+  int64_t d2 = t.size(-2);
+  TVM_FFI_CHECK((CMod(d1, 128LL) == 0LL), ValueError)
+      << KERNEL_NAME << ": descriptor 'A' requires (CMod(d1, 128) == 0)";
+  cuuint64_t global_dim[3] = {Extent(128LL, "A"), Extent(d2, "A"), Extent(CDiv(d1, 128LL), "A")};
+  cuuint64_t global_strides[2] = {Stride(d1, "A"), Stride(128LL, "A")};
+  cuuint32_t box_dim[3] = {128u, 128u, 1u};
+  cuuint32_t elem_strides[3] = {1u, 1u, 1u};
+  const bool allow_oob_box[3] = {false, true, false};
+  for (int axis = 0; axis < 3; ++axis) {
+    TVM_FFI_CHECK(allow_oob_box[axis] || box_dim[axis] <= global_dim[axis], ValueError)
+        << KERNEL_NAME << ": descriptor 'A' axis " << axis << " (0 is the innermost) has a TMA box of "
+        << box_dim[axis] << " elements but the source extent along it is " << global_dim[axis]
+        << ". A box may not exceed its extent unless the axis is listed in the descriptor's allow_oob_box (CK.describe_tma(..., allow_oob_box=(axis,))). List it when the box is a tile that may run past the end of the axis, a tail tile: a load fills the part past the extent with zeros (its bytes still count toward the barrier's transaction) and a store or reduce skips it, so that is safe when the kernel ignores or masks those rows and columns. If the box was meant to fit, box_dim or the extent is wrong: check box_dim, global_dim and the tensor's shape";
+  }
+  TVM_FFI_CHECK(global_strides[0] % 16 == 0 && global_strides[1] % 16 == 0, ValueError)
+      << KERNEL_NAME << ": the strides of descriptor 'A' must be multiples of 16 bytes";
+  void* address = DataPtr(t);
+  TVM_FFI_CHECK(reinterpret_cast<uintptr_t>(address) % 16 == 0, ValueError)
+      << KERNEL_NAME << ": the address of descriptor 'A' must be 16-byte aligned";
+  CUtensorMap map{};
+  CheckCu(cuTensorMapEncodeTiled(&map, CU_TENSOR_MAP_DATA_TYPE_UINT8, 3, address, global_dim, global_strides,
+                                 box_dim, elem_strides, CU_TENSOR_MAP_INTERLEAVE_NONE, CU_TENSOR_MAP_SWIZZLE_128B,
+                                 CU_TENSOR_MAP_L2_PROMOTION_NONE, CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE),
+          "cuTensorMapEncodeTiled of A");
+  return map;
+}
+// ... (EncodeTma_B, EncodeTma_C_tma, Run functions)
+}

이 코드는 CUtensorMapEncodeTiled 함수를 호출하여 텐서 A의 메모리 레이아웃을 기술합니다. global_dim은 텐서의 전체 크기를, global_strides는 각 차원의 stride를 정의합니다. box_dim은 TMA가 한 번에 로드할 타일의 크기를 나타내며, CU_TENSOR_MAP_SWIZZLE_128B는 128바이트 단위의 메모리 스위즐링을 통해 메모리 접근 패턴을 최적화합니다. 또한, TVM_FFI_CHECK와 CheckCu를 통해 텐서의 차원, stride, 메모리 정렬 등 다양한 제약 조건을 런타임에 검사하여 안정성을 확보합니다.

2. 커널 선택 로직 업데이트: flashinfer/gemm/cake_grouped_fp8_gemm.py

이 PR은 Python 레벨에서 cake_grouped_fp8_gemm 커널을 실행하기 전에, 입력 m_indices의 레이아웃과 GPU 디바이스의 특성을 분석하여 가장 적합한 CUDA 커널을 동적으로 선택하는 로직을 개선했습니다. 특히, validate_indices 함수는 새로운 "window-segment KG1 program"을 선택하는 중요한 역할을 합니다.

기존에는 특정 m_indices 레이아웃(128-row window 내에 여러 expert run이 시작되는 경우)에 대한 최적화된 프로그램이 없었지만, 이 PR을 통해 다음과 같은 조건이 충족될 때 새로운 deepk_c2_kg1_wseg 프로그램이 선택됩니다:

• FP32-scale 레이아웃을 사용하고, expert run이 128-row window 내에서 시작되는 경우.
• ceil(M / 128) x window_runs x ceil(N / 64) CTAs(Cooperative Thread Arrays)가 단일 wave에 적합한 경우.

이 새로운 프로그램은 각 128-row window에 window_runs개의 CTAs를 할당하고, 각 CTA가 하나의 expert run을 처리하도록 설계되었습니다. 또한, m_indices를 매번 읽어 유연성을 유지하며, CUDA 그래프 리플레이 시에도 올바른 동작을 보장합니다. 만약 이 최적화된 프로그램의 launch contract(예: 1 <= M <= 2**24, 64 <= N <= 2**19)를 벗어나는 경우에는 기존의 expert-aligned 프로그램으로 폴백하여 견고성을 유지합니다.

(PR diff에 Python 코드 변경사항이 직접 제공되지 않았으므로, 위 설명은 PR 설명과 리뷰 댓글을 바탕으로 한 개념적인 변경사항입니다.)

왜 이게 좋은가?

1. 압도적인 성능 향상

이 PR의 가장 큰 장점은 특정 시나리오에서 측정된 현저한 성능 향상입니다. NVIDIA GB300 (sm_103a) 및 B200 (sm_100a) GPU에서 테스트한 결과, correctness_unaligned_transition 및 random_s9_m777_g8_n256_k512와 같은 특정 shape에서 최대 1.4배에 달하는 속도 향상을 보였습니다.

#	shape	M	N	K	groups	program (this PR)	FI main (us)	this PR (us)	speedup	verdict [noise]
1	correctness_unaligned_transition	65	128	128	2	deepk_c2_kg1_wseg 4 CTAs	5.728	4.320	1.326x	FASTER
2	random_s9_m777_g8_n256_k512	777	256	512	8	deepk_c2_kg1_wseg 84 CTAs	6.720	4.896	1.373x	FASTER

이러한 성능 향상은 LLM 추론과 같이 지연 시간에 민감한 애플리케이션에서 사용자 경험을 크게 개선하고, 전체 시스템의 처리량을 증가시키는 데 기여합니다. 특히, FlashInfer의 CuTe-DSL 기반 GEMM과 비교했을 때도 WIN (Winner) 판정을 받으며 경쟁 우위를 확보했습니다.

2. 하드웨어 최적화의 정수: TMA 활용

새로운 커널은 CUtensorMap을 통해 NVIDIA GPU의 TMA 기능을 적극적으로 활용합니다. TMA는 GPU의 메모리 계층 구조를 최적화하여 전역 메모리 접근 지연 시간을 줄이고, 메모리 대역폭 활용률을 높입니다. 이는 특히 FP8과 같이 데이터 밀도가 높은 연산에서 데이터 로딩이 병목이 될 수 있는 상황에서 매우 중요합니다. 128바이트 스위즐링과 같은 기법은 메모리 뱅크 충돌을 줄이고 캐시 효율성을 극대화하여 실제 성능 향상으로 이어집니다.

3. 특정 패턴에 대한 정교한 CTA 스케줄링

"window-segment KG1 program"은 128-row window 내에 여러 expert run이 존재하는 특정 데이터 패턴에 맞춰 CTA 할당 및 스케줄링을 재설계했습니다. 각 CTA가 하나의 expert run을 전담하도록 함으로써, 불필요한 동기화 오버헤드를 줄이고 병렬성을 극대화합니다. 이는 범용 커널로는 달성하기 어려운 미세한 최적화이며, 특정 워크로드에 대한 깊은 이해를 바탕으로 합니다.

4. 유연성과 견고성의 균형

이 최적화는 단순히 성능만 추구하는 것이 아니라, 유연성과 견고성도 함께 고려했습니다. m_indices를 매번 읽어 다양한 입력 레이아웃에 대응할 수 있도록 했으며, 최적화된 커널의 제약 조건을 벗어나는 경우에는 기존의 안정적인 커널로 폴백(fallback)하는 메커니즘을 통해 시스템의 안정성을 보장합니다. 이는 실제 프로덕션 환경에서 다양한 입력 조건에 유연하게 대응하면서도 성능을 유지할 수 있게 합니다.

일반적 교훈

이 PR은 GPU 프로그래밍에서 성능 최적화를 위한 몇 가지 중요한 교훈을 제공합니다:

• 하드웨어 인지 최적화: GPU 아키텍처(예: TMA)에 대한 깊은 이해를 바탕으로 한 최적화는 소프트웨어적인 기법만으로는 달성하기 어려운 성능 향상을 가져올 수 있습니다.
• 워크로드 특화 커널: 모든 상황에 맞는 하나의 범용 커널보다는, 특정 데이터 패턴이나 워크로드에 특화된 커널을 개발하고 동적으로 선택하는 전략이 더 큰 성능 이득을 제공할 수 있습니다.
• 유연성과 성능의 균형: 최적화된 커널이 모든 상황에 적용될 수 없음을 인정하고, 제약 조건을 벗어나는 경우를 위한 견고한 폴백 메커니즘을 마련하는 것이 중요합니다.

결론

FlashInfer의 cake_grouped_fp8_gemm 커널에 도입된 "window-segment KG1 program"은 FP8 GEMM 연산의 특정 시나리오에서 GPU 성능을 획기적으로 개선하는 중요한 진전입니다. CUtensorMap을 통한 TMA 활용, 정교한 CTA 스케줄링, 그리고 유연한 커널 선택 로직의 조합은 고성능 딥러닝 추론 시스템을 구축하는 데 있어 FlashInfer의 기술적 리더십을 다시 한번 입증합니다. 이러한 최적화는 LLM과 같은 대규모 AI 모델의 효율적인 배포와 운영에 필수적인 요소가 될 것입니다.

참고 자료

• cuTensorMapEncodeTiled: https://docs.nvidia.com/cuda/cuda-driver-api/group__CUDA__TENSOR__MAP.html#group__CUDA__TENSOR__MAP_1g826354673105786c310c1440e210672e
• NVIDIA FP8 Data Format: https://docs.nvidia.com/deeplearning/performance/dl-performance-matrix-multiplication/index.html#fp8-format
• NVIDIA BF16 Custom Type: https://docs.nvidia.com/cuda/bfloat16-custom-type/index.html
• CUDA C++ Programming Guide - Device Functions: https://docs.nvidia.com/cuda/cuda-c-programming-guide/index.html#device-functions

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