[flashinfer] FlashInfer, SM103a 아키텍처를 위한 cake_all_gather_matmul 최적화: CTA-pair 트윈 도입으로 성능 향상

2026년 10월 23일수정: 2026년 10월 23일

PR 링크: flashinfer-ai/flashinfer#7083 상태: Merged | 변경: +1051 / -78

들어가며

대규모 언어 모델(LLM)의 발전은 GPU 컴퓨팅 성능의 한계를 끊임없이 시험하고 있습니다. 특히 모델의 크기가 커지고 복잡해짐에 따라, 통신 및 연산 집약적인 작업의 효율성을 극대화하는 것이 중요해졌습니다. FlashInfer는 이러한 요구에 부응하여 고성능 추론 엔진을 제공하며, 이번 PR은 cake_all_gather_matmul 연산에서 특정 조건 하에 SM103a (B300) 아키텍처의 성능을 크게 향상시키는 것을 목표로 합니다.

이 PR은 cake_all_gather_matmul 연산이 특정 상황에서 더 효율적인 실행 경로를 사용하도록 라우팅을 변경합니다. 구체적으로, 4개의 GPU로 구성된 월드 사이즈(world size 4), K-major 가중치, 512개의 패딩된 행(padded rows), 그리고 2560의 출력 너비(N=2560)를 가지는 bf16 및 fp16 연산에 대해 기존의 단일 CTA(Compute Thread Array) 기반 커널 대신 CTA-pair 메인 프로그램(_cg2 트윈)을 사용하도록 변경합니다. 이는 특히 Llama-3.1-70B 모델의 TP4 QKV 연산과 같은 실제 워크로드에서 상당한 속도 향상을 가져옵니다.

코드 분석

이번 PR의 핵심 변경 사항은 기존 커널의 실행 경로를 최적화하기 위한 라우팅 규칙 수정 및 새로운 커널 프로그램의 추가입니다. 주요 변경 사항은 다음과 같습니다.

1. 라이선스 헤더 변경

여러 _kernel.cu 파일에서 라이선스 헤더가 NVIDIA에서 FlashInfer 팀으로 변경되었습니다. 이는 코드의 소유권 및 라이선스 정보를 명확히 하는 일반적인 절차입니다.

Before:

- * Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved.
+

After:

+ * Copyright (c) 2026 by FlashInfer team.
+

2. 새로운 커널 파일 추가: `cake_all_gather_matmul_72ae81d01ed72fe75fe8_kernel.cu`

가장 중요한 변경 사항은 새로운 CUDA 커널 파일이 추가된 것입니다. 이 파일은 SM103a 아키텍처에 특화된 최적화된 cake_all_gather_matmul 연산을 구현합니다. 이 커널은 _cg2 트윈으로, CTA-pair 메인 프로그램을 사용하여 이전 버전보다 효율적인 실행을 제공합니다.

+new file mode 100644
+index 0000000000..56cfa85d6a
--- /dev/null
+++ b/csrc/cake_all_gather_matmul/cake_all_gather_matmul_72ae81d01ed72fe75fe8_kernel.cu
@@ -0,0 +1,726 @@
+/*
+ * Copyright (c) 2026 by FlashInfer team.
+ * ... (라이선스 정보) ...
+ */

+// Common preamble (typedefs, tensor-map ABI, compiler helpers) shared by this export's kernels.
+#include "cake_all_gather_matmul_device_common.cuh"
+// The common preamble's fallback helper this module's helper library does not provide; the shared
+// header carries another member's library definition of the name (kept translation-unit-local).
+__device__ __forceinline__ int make_warp_uniform(int x) {
+  int result;
+  asm volatile("shfl.sync.idx.b32 %0, %1, 0, 0x1F, 0xFFFFFFFF;" : "=r"(result) : "r"(x));
+  return result;
+}

+#define CAKE_INF CUDART_INF_F
+#define TMEM_NCOLS 512
+#define TMEM_ACCUM_OFFSET 0
+#define NUM_TMA_PIPE_STAGES 6
+#define NUM_MAINLOOP_PIPE_STAGES 2
+#define SMEM_SMEM_A_OFF 1024
+#define SMEM_SMEM_A_STAGE_BYTES 16384
+#define SMEM_SMEM_A_STRIDE 32768
+#define SMEM_SMEM_B_OFF 17408
+#define SMEM_SMEM_B_STAGE_BYTES 16384
+#define SMEM_SMEM_B_STRIDE 32768
+#define SMEM_TOTAL 197632
+#define THREADS 192
+
+extern "C" {
+
+__global__ __launch_bounds__(THREADS) 
+    __cluster_dims__(2, 1, 1) void kernel_cake_all_gather_matmul_72ae81d01ed72fe75fe8(
+        const __grid_constant__ CUtensorMap A_local, const __grid_constant__ CUtensorMap A_scratch, 
+        const __grid_constant__ CUtensorMap B, __half* __restrict__ C, 
+        __half* __restrict__ scratch_payload, unsigned int* __restrict__ ready, 
+        unsigned int ready_target, int rank, int M, int scratch_pitch, int signal_rows, int n_tiles, 
+        int remote_order) {
+  const int tid = threadIdx.x;
+  const int warp = make_warp_uniform(tid / 32);
+  const int lane = tid % 32;
+
+  extern __shared__ __align__(1024) char smem_raw[];
+  int smem;
+#if __CUDA_ARCH__ == 1000
+  asm volatile("{ .reg .u64 smem_ptr; cvta.to.shared.u64 smem_ptr, %1; cvt.u32.u64 %0, smem_ptr; }"
+               : "=r"(smem)
+               : "l"(smem_raw));
+  smem = make_warp_uniform(smem);
+#else
+  smem = (int)(unsigned long long)__cvta_generic_to_shared(smem_raw);
+#endif
+
+  const int mbar_base = smem;
+#define tma_full_addr (mbar_base + 0)
+#define mma_done_addr (mbar_base + 48)
+#define mainloop_done_addr (mbar_base + 96)
+#define epilogue_done_addr (mbar_base + 112)
+
+  const int bid = blockIdx.x;
+  const int num_bids = gridDim.x;
+  const unsigned int clusters_x = gridDim.x / 2;
+  const unsigned int cluster_id =
+      ((blockIdx.z * gridDim.y + blockIdx.y) * clusters_x) + blockIdx.x / 2;
+  const unsigned int num_clusters = clusters_x * gridDim.y * gridDim.z;
+
+  int cta_rank;
+  asm volatile("mov.b32 %0, %%cluster_ctarank;" : "=r"(cta_rank));
+
+  // Kernel setup ops
+  __half* smem_a = reinterpret_cast<__half*>(smem_raw + SMEM_SMEM_A_OFF);
+  const int smem_a_addr = smem + SMEM_SMEM_A_OFF;
+  __half* smem_b = reinterpret_cast<__half*>(smem_raw + SMEM_SMEM_B_OFF);
+  const int smem_b_addr = smem + SMEM_SMEM_B_OFF;
+
+  // Mbarrier init (4 pipeline groups, 0 ordered-sequence groups, 16 barriers)
+  // Mbarriers at smem_raw[0..128)
+
+  if (warp == 0) {
+    uint32_t leader = elect_sync();
+    if (leader) {
+      // --- pipeline 'tma_pipe' ---
+      // tma_full: 6 barriers, init_count=2
+      mbarrier_init(smem + 0, 2);
+      mbarrier_init(smem + 8, 2);
+      mbarrier_init(smem + 16, 2);
+      mbarrier_init(smem + 24, 2);
+      mbarrier_init(smem + 32, 2);
+      mbarrier_init(smem + 40, 2);
+      // mma_done: 6 barriers, init_count=1
+      mbarrier_init(smem + 48, 1);
+      mbarrier_init(smem + 56, 1);
+      mbarrier_init(smem + 64, 1);
+      mbarrier_init(smem + 72, 1);
+      mbarrier_init(smem + 80, 1);
+      mbarrier_init(smem + 88, 1);
+      // --- pipeline 'mainloop_pipe' ---
+      // mainloop_done: 2 barriers, init_count=1
+      mbarrier_init(smem + 96, 1);
+      mbarrier_init(smem + 104, 1);
+      // epilogue_done: 2 barriers, init_count=8
+      mbarrier_init(smem + 112, 8);
+      mbarrier_init(smem + 120, 8);
+      asm volatile("fence.mbarrier_init.release.cluster;" ::: "memory");
+    }
+  }
+
+  __syncwarp();
+
+  asm volatile("barrier.cluster.arrive.release.aligned;" ::: "memory");
+  asm volatile("barrier.cluster.wait.acquire.aligned;" ::: "memory");
+
+  // TMEM alloc (512 columns, 512 used)
+  volatile int* tmem_addr_storage = (volatile int*)(smem_raw + 128);
+  if (warp == 0) {
+    int _tmem_hold = smem + 128;
+    asm volatile(
+        "tcgen05.alloc.cta_group::2.sync.aligned.shared::cta.b32 [%0], %1;" ::"r"(_tmem_hold),
+        "r"(512

이 새로운 커널은 kernel_cake_all_gather_matmul_72ae81d01ed72fe75fe8 함수로 정의되며, SM103a 아키텍처의 특성을 활용하여 TMA (Tensor Memory Accelerator), MMA (Matrix Multiply-Accumulate), 그리고 공유 메모리(Shared Memory) 사용을 최적화합니다. 특히 make_warp_uniform, mbarrier_init, tcgen05.alloc과 같은 CUDA 저수준 기능을 사용하여 워프(warp) 간 동기화 및 메모리 접근을 효율적으로 관리합니다.

3. 기존 커널 파일의 불필요한 코드 제거 (추정)

여러 _kernel.cu 파일에서 NVIDIA 라이선스 헤더가 제거된 것을 볼 수 있습니다. 이는 위에서 언급한 라이선스 변경과 관련이 있으며, 해당 파일들이 더 이상 NVIDIA의 특정 라이선스 정책을 따르지 않거나, FlashInfer 자체 라이선스로 통합되었음을 시사합니다. 실제 코드 로직의 변경은 이 부분에서는 관찰되지 않습니다.

왜 이게 좋은가?

이 PR은 다음과 같은 이유로 좋은 최적화 및 개선이라고 할 수 있습니다.

특정 워크로드에 대한 성능 향상: PR 설명에 명시된 대로, 특정 조건(512x2560, bf16/fp16, world size 4, K-major)에서 이전 버전에 비해 1.0767 ~ 1.1143 배의 속도 향상을 달성했습니다. 이는 LLM 추론 시 병목 현상이 발생하기 쉬운 통신 및 행렬 곱셈 연산의 효율성을 직접적으로 개선한 결과입니다.

Shape per rank, M × N (K = 8192, K-major weights)	Dtype	Node	Previous version (µs)	This PR (µs)	Speedup
512 × 2560, world size 4 (Llama-3.1-70B TP4 QKV)	bf16	node 1	115.4	106.2	1.0855
512 × 2560, world size 4 (Llama-3.1-70B TP4 QKV)	bf16	node 2	114.4	105.3	1.0902
512 × 2560, world size 4	fp16	node 1	114.8	104.0	1.1143
512 × 2560, world size 4	fp16	node 2	115.3	107.3	1.0767

아키텍처 특화 최적화: SM103a (B300) 아키텍처의 새로운 기능이나 개선점을 활용하여 성능을 극대화했습니다. 특히 CTA-pair 메인 프로그램(_cg2 트윈)을 도입하여 이전의 단일 CTA 프로그램보다 더 많은 스레드와 리소스를 활용할 수 있게 된 점이 중요합니다. 이는 최신 하드웨어의 잠재력을 최대한 끌어내는 좋은 사례입니다.
점진적이고 안전한 변경: 이 PR은 기존에 배포된 커널 프로그램을 수정하거나 제거하지 않고, 특정 조건에 맞는 워크로드에 대해 새로운 최적화된 경로를 추가하는 방식으로 작동합니다. 이는 기존 기능의 안정성을 유지하면서 성능 개선을 추구하는 안전한 접근 방식입니다. SM100a (B200) 아키텍처나 다른 워크로드에 대한 변경은 없으므로, 호환성 문제의 위험을 최소화합니다.
코드 재사용 및 모듈화: 새로운 커널은 cake_all_gather_matmul_device_common.cuh와 같은 공통 헤더 파일을 포함하여 코드 재사용성을 높이고 있습니다. 이는 유지보수성을 향상시키고, 향후 유사한 최적화를 다른 커널에 적용할 때 기반이 될 수 있습니다.

일반적인 교훈

하드웨어 특화 최적화의 중요성: 최신 GPU 아키텍처는 고유한 기능과 성능 특성을 가지고 있습니다. 이러한 특성을 깊이 이해하고 활용하는 것은 최고의 성능을 달성하는 데 필수적입니다. CTA-pair와 같은 새로운 실행 모델은 이를 잘 보여줍니다.
점진적 개선 전략: 기존 시스템의 안정성을 해치지 않으면서 성능을 개선하는 가장 좋은 방법 중 하나는 특정 조건에 맞는 새로운 최적화 경로를 추가하는 것입니다. 이는 위험을 줄이고 점진적인 성능 향상을 가능하게 합니다.
측정 및 검증의 중요성: PR 설명에 포함된 상세한 성능 측정 결과는 변경 사항의 효과를 명확하게 보여줍니다. 실제 워크로드에서의 성능 측정은 최적화의 가치를 입증하는 데 매우 중요합니다.
코드 품질 및 유지보수성: 라이선스 변경, 공통 헤더 사용 등은 코드의 명확성과 유지보수성을 높이는 데 기여합니다. 잘 구조화된 코드는 장기적으로 프로젝트의 성공에 필수적입니다.

리뷰 댓글 분석

리뷰 댓글은 주로 테스트 실행 및 결과 확인에 집중되었습니다. 초기에는 일부 테스트가 실패했지만, 이는 새로운 기능의 도입으로 인한 예상 가능한 문제일 수 있습니다. 개발자가 run tests/comm/test_all_gather_matmul_cake.py tests/comm/test_all_gather_matmul_cake_e2e.py tests/comm/test_all_gather_matmul.py 명령을 통해 테스트를 재실행하도록 요청한 후, CI 파이프라인이 성공적으로 완료된 것으로 보입니다. 이는 변경 사항이 의도한 대로 작동하며, 기존 기능과의 호환성도 유지됨을 시사합니다. 특히 VR NVL72 (multi-node)에서 발생한 script failed before producing a JUnit report와 같은 인프라 관련 문제는 코드 자체의 결함이라기보다는 테스트 환경의 일시적인 문제였을 가능성이 높습니다.

References

참고 자료

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

PR Analysis 의 다른글

이전글 [Liger-Kernel] LinkedIn Liger-Kernel GRPO Loss Backward 최적화: 불필요한 연산 제거로 성능 50% 향상
현재글 : [flashinfer] FlashInfer, SM103a 아키텍처를 위한 cake_all_gather_matmul 최적화: CTA-pair 트윈 도입으로 성능 향상
다음글 [flashinfer] FlashInfer, Cake DSA 커널 최적화: Sparse Attention 성능 2배 향상

댓글

관련 포스트

PR Analysis 의 다른글