[flashinfer] FlashInfer, Qwen3-30B 모델의 성능 향상을 위한 CUDA 커널 최적화: L2 캐시 힌트 도입

PR 링크: flashinfer-ai/flashinfer#5593 상태: Merged | 변경: +595 / -519

들어가며

최근 대규모 언어 모델(LLM)의 발전 속도는 눈부십니다. 이러한 모델들은 방대한 양의 파라미터를 가지며, 효율적인 추론(inference)은 실제 서비스 적용에 있어 매우 중요한 과제입니다. 특히 Mixture-of-Experts (MoE) 아키텍처는 모델의 크기를 늘리면서도 연산량을 효율적으로 관리할 수 있어 주목받고 있습니다.

이번 글에서는 LLM 추론 라이브러리인 FlashInfer의 Pull Request(PR)를 통해 Qwen3-30B 모델의 특정 설정(SM103, num_tokens 8-32)에서 발생했던 성능 병목 현상을 어떻게 해결했는지 심층적으로 분석해보겠습니다. 이 PR은 CUDA 커널 레벨에서의 미묘하지만 효과적인 최적화를 통해 상당한 속도 향상을 달성했습니다. 코드 변경 사항과 그 원리를 자세히 살펴보며, 이러한 최적화가 왜 효과적인지, 그리고 어떤 교훈을 얻을 수 있는지 알아보겠습니다.

코드 분석

이번 PR의 핵심 변경 사항은 csrc/fused_moe/warp_decode/generated/cake_warp_decode_generated_manifest.cuh 파일에서 주로 이루어졌습니다. 이 파일은 Qwen3-30B 모델의 MoE 레이어 추론을 위한 CUDA 커널들의 선언과 구현을 포함하고 있습니다. 변경의 주요 목적은 GPU의 L2 캐시 사용 방식을 최적화하여 데이터 로딩 지연 시간을 줄이는 것입니다.

1. 커널 함수 이름 변경 및 SHA256 해시 업데이트

가장 먼저 눈에 띄는 변경은 여러 CUDA 커널 함수의 이름이 변경되고, kGeneratedSourceSha256 값이 업데이트된 것입니다. 이는 내부적으로 커널의 구현이 변경되었음을 나타냅니다. 예를 들어, kernel_cake_warp_decode_7fc2fef14e509a2ea34d 함수가 kernel_cake_warp_decode_7173b39130de7a59f9c6으로 변경되었습니다. 이는 코드 생성 과정에서 발생한 변경사항으로, 실제 로직 변경의 시작점을 알리는 신호입니다.

Before:

__global__ __launch_bounds__(512, 1) void
kernel_cake_warp_decode_7fc2fef14e509a2ea34d(const __grid_constant__ CUtensorMap A, uint8_t* __restrict__ B, const __grid_constant__ CUtensorMap SFA, uint8_t* __restrict__ SFB, const __grid_constant__ CUtensorMap C, uint8_t* __restrict__ SFC, int* __restrict__ route_map, int* __restrict__ tile_expert, int* __restrict__ tile_mn_limit, int* __restrict__ num_non_exiting_ctas, int* __restrict__ work_counter, float* __restrict__ scale_c, float* __restrict__ scale_gate, float* __restrict__ clamp_limit, float* __restrict__ act_alpha, float* __restrict__ act_beta, int M_out, int K, int grid_m, int grid_n, int K_tiles, int* __restrict__ route_experts, int* __restrict__ route_slots, float* __restrict__ pack_ready, int* __restrict__ done_counter, int route_count, int top_k, int local_expert_offset, int num_experts, int initial_work, int launch_ctas);

After:

__global__ __launch_bounds__(512, 1) void
kernel_cake_warp_decode_7173b39130de7a59f9c6(const __grid_constant__ CUtensorMap A, uint8_t* __restrict__ B, const __grid_constant__ CUtensorMap SFA, uint8_t* __restrict__ SFB, const __grid_constant__ CUtensorMap C, uint8_t* __restrict__ SFC, int* __restrict__ route_map, int* __restrict__ tile_expert, int* __restrict__ tile_mn_limit, int* __restrict__ num_non_exiting_ctas, int* __restrict__ work_counter, float* __restrict__ scale_c, float* __restrict__ scale_gate, float* __restrict__ clamp_limit, float* __restrict__ act_alpha, float* __restrict__ act_beta, int M_out, int K, int grid_m, int grid_n, int K_tiles, int* __restrict__ route_experts, int* __restrict__ route_slots, float* __restrict__ pack_ready, int* __restrict__ done_counter, int route_count, int top_k, int local_expert_offset, int num_experts, int initial_work, int launch_ctas);

또한, kGeneratedSourceSha256 값도 변경되었습니다. 이는 코드 생성기의 출력이 변경되었음을 나타내며, 이 변경이 실제 성능에 영향을 미치는 핵심 로직임을 시사합니다.

Before:

inline constexpr char kGeneratedSourceSha256[] = "5b65a6b3577b69d9b12c49fbf849b5f93ae24aad994a289a7ffbf0fd8d2390c2";

After:

inline constexpr char kGeneratedSourceSha256[] = "021a8eb574f0b9e37da690f09d34c0f8cd9b549b61d5d861b9809d2a88543d40";

2. L2 캐시 힌트 도입 (cp.async.bulk.tensor ... .L2::cache_hint)

PR 설명에 따르면, 이번 최적화의 핵심은 GPU의 Tensor Memory Accelerator (TMA) 로드 명령에 L2 캐시 힌트를 추가한 것입니다. 특히, evict-first 전략을 사용하여 L2 캐시를 활용합니다. 이는 데이터를 로드할 때, 해당 데이터가 L2 캐시에 이미 존재한다면 eviction을 먼저 수행하고 새로운 데이터를 로드하도록 지시하는 방식입니다.

이 변경은 cp.async.bulk.tensor와 같은 TMA 로드 명령어에 적용됩니다. 비록 diff에서 직접적으로 cp.async.bulk.tensor 명령어가 보이지는 않지만, 커널 함수의 시그니처 변경과 SHA256 해시 값 변경은 이러한 내부적인 TMA 로드 방식의 수정이 있었음을 강력히 시사합니다. PR 설명에서 언급된 L2::cache_hint는 이러한 최적화의 핵심입니다.

이 최적화는 다음과 같은 이점을 가집니다:

• L2 캐시 재사용성 증대: evict-first 힌트는 L2 캐시의 기존 데이터를 먼저 내보내고 새로운 데이터를 적재함으로써, 캐시 미스(cache miss) 발생 시 불필요한 L2 쓰기 백(write-back)을 줄입니다.
• 메모리 대역폭 효율 향상: 이전 커널 실행으로 인해 L2 캐시에 남아있던 더티(dirty) L2 라인(line)을 다시 쓰기 위해 발생하는 비용을 절감합니다. 즉, 커널 실행 전에 L2 캐시를 깨끗하게 유지하는 데 드는 비용을 줄여, 순수하게 데이터 로딩에 집중할 수 있게 합니다.
• 콜드 L2(Cold L2) 상황에서의 성능 개선: 벤치마크 프로토콜에서 L2 캐시가 비어있는 콜드 L2 상황에서도 이 방식은 효과적입니다. 이전 커널이 남긴 더티 L2 라인을 처리하는 오버헤드 없이 바로 데이터를 로드할 수 있기 때문입니다.

PR 설명에 따르면, 이 변경은 Qwen3-30B 모델의 SM103, num_tokens 8-32 범위에서 특히 효과적입니다. 이 범위의 7개의 SM103 퓨즈드 라우트-팩 모듈들이 이 최적화의 대상이 되었습니다.

왜 이게 좋은가?

이번 PR은 다음과 같은 이유로 훌륭한 최적화라고 할 수 있습니다.

1. 명확한 성능 향상: PR 설명의 테이블에서 볼 수 있듯이, num_tokens 8-32 범위에서 기존 대비 1.12x ~ 1.16x의 속도 향상을 달성했습니다. 이는 기하 평균으로 약 1.13x ~ 1.14x에 해당하며, 이전 PR의 1.00x ~ 1.03x 대비 상당한 개선입니다. 특히, num_tokens 31-32에서는 약 6.6 µs의 지연 시간을 절감했습니다.

   Speedup vs baseline = baseline latency B (FlashInfer trtllm-gen NVFP4 fused-MoE path) divided by this export's latency E_o in the same paired population (B/E_o); > 1.00x means this PR is faster than the baseline. For num_tokens 8–32 the value is the median of the two independent B300 allocations (both listed); for num_tokens 1–7 it is the median of three measurements. "Previous PR" is the merged export before this PR (#5574).

   Compared with the previous PR, this PR speeds up the 25 shapes num_tokens 8–32 (bold rows): speedup vs baseline goes from 1.00–1.03x to 1.12–1.16x, geometric mean 1.13x (session 1) / 1.14x (session 2) vs about 1.015x before, and the export latency drops by 3.5–6.7 µs per call (3.5 µs at num_tokens 8, 6.7 µs at num_tokens 31, 6.6 µs at num_tokens 32). num_tokens 1–7 use unchanged modules and are re-measured for non-regression only.

2. 기존 기능과의 호환성 유지: PR 설명에서 강조하듯이, 워프 역할(warp roles), 배리어(barriers), 파이프라인 깊이(pipeline depth), 타일 모양(tile shapes), 에필로그(epilogues), 라우트(routes), 경계(boundaries) 및 물리적 계약(physical contract)은 변경되지 않았습니다. 즉, 모든 변경된 모듈의 출력은 이전 버전과 비트 단위로 동일(bitwise identical)합니다. 이는 기능적 변경 없이 순수하게 성능만을 개선했음을 의미하며, 기존 시스템과의 호환성을 보장합니다.

3. 메모리 계층 구조 최적화: 이 PR은 GPU 메모리 계층 구조, 특히 L2 캐시의 동작 방식을 이해하고 이를 최적화하는 좋은 예시를 보여줍니다. evict-first와 같은 캐시 힌트를 적절히 사용함으로써, 데이터 로딩 성능을 크게 향상시킬 수 있음을 입증했습니다. 이는 복잡한 딥러닝 모델에서 메모리 I/O가 병목이 되는 경우가 많기 때문에 매우 중요한 최적화 기법입니다.

4. 체계적인 성능 측정 및 검증: PR에서는 다양한 시나리오(콜드 L2, 여러 세션)에 걸쳐 상세한 성능 측정 결과를 제공합니다. 또한, compute-sanitizer를 이용한 동시성(synccheck) 및 경쟁 조건(racecheck) 검사를 통과했으며, Nsight Compute를 통한 성능 분석 결과도 제시하여 변경 사항의 타당성과 안정성을 검증했습니다.

일반적인 교훈

이 PR을 통해 얻을 수 있는 몇 가지 일반적인 교훈은 다음과 같습니다:

• 하드웨어 특성 이해의 중요성: GPU 아키텍처, 특히 메모리 계층 구조(L1, L2 캐시, 레지스터 파일, TMA 등)의 동작 방식을 깊이 이해하는 것은 고성능 컴퓨팅에서 필수적입니다. 이번 PR은 L2 캐시 힌트의 효과를 잘 보여줍니다.
• 프로파일링과 병목 분석: 성능 병목 지점을 정확히 파악하는 것이 최적화의 첫걸음입니다. 이 PR은 특정 모델과 설정에서 발생한 메모리 로딩 지연 시간을 식별하고 이를 해결하는 데 집중했습니다.
• 작은 변경으로 큰 효과: 때로는 커널 내의 작은 명령어 수준 변경(예: 캐시 힌트 추가)이 전체 성능에 상당한 영향을 미칠 수 있습니다. 복잡한 시스템일수록 이러한 미세한 최적화가 중요합니다.
• 철저한 검증: 성능 개선뿐만 아니라, 기능적 정확성과 안정성을 보장하는 것이 중요합니다. 비트 단위 동일성(bitwise identical) 검증, 동시성 검사 등은 이러한 신뢰성을 높여줍니다.

References

• Tensor Memory Accelerator (TMA) - NVIDIA Documentation
• CUDA C++ Programming Guide - Cache Hints

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