본문으로 건너뛰기

[flashinfer] NVIDIA R200(SM107)을 위한 SwiGLU + MXFP8 커널 최적화: 소규모 워크로드 성능 향상 기법

PR 링크: flashinfer-ai/flashinfer#8023 상태: Merged | 변경: +1282 / -14

들어가며

LLM(Large Language Model) 추론 가속화에서 가장 중요한 요소 중 하나는 연산 오버헤드를 줄이고 메모리 대역폭을 극대화하는 것입니다. 특히 SwiGLU 활성화 함수와 양자화(Quantization) 과정을 하나의 커널로 합치는 Fused Kernel 기술은 현대적인 추론 엔진의 필수 요소입니다.

이번에 분석할 FlashInfer의 PR은 NVIDIA의 최신 아키텍처 중 하나인 SM107(R200)에서 Cake 라이브러리를 활용한 SwiGLU + MXFP8 양자화 커널의 성능을 최적화한 사례입니다. 핵심은 입력 데이터의 크기($M \times K$)에 따라 가장 적합한 커널 실행 경로를 선택하도록 JIT(Just-In-Time) 라우팅 로직을 정교화한 것입니다.

코드 분석: 무엇이 변했는가?

1. cake_jit.py: 워크로드 크기에 따른 커널 라우팅 전략 수정

가장 중요한 변경 사항은 cake_jit.py에 있습니다. 기존에는 특정 아키텍처에 대해 하나의 범용 커널을 사용하거나 임계값이 높게 설정되어 있었으나, 이번 PR에서는 소규모 워크로드를 위한 전용 커널(e62d5f...)을 추가하고 임계값을 세분화했습니다.

Before (기존 로직의 일부): 기존에는 0520c1... 커널이 sm_107a의 기본 경로로 사용되거나 높은 임계값을 가졌습니다.

# 기존 cake_gated_act_mxfp8_0520c14fde53ef7bfea7 설정
"arches": [
    "sm_103a",
    "sm_107a",
],
"select_rule": {
    "elements": {
        "sm_103a": 2097152, # 2^21
    },
    "kind": "elements_at_least",
},

After (변경된 로직): 새로운 커널 e62d5f...를 도입하여 $M \times K < 2^{20}$인 구간을 담당하게 하고, 기존 커널은 그 이상의 구간을 담당하도록 변경했습니다.

# 1. 소규모 워크로드용 신규 커널 (Fallback/Small size)
"cake_gated_act_mxfp8_e62d5f6dda3083f31950": {
    "arches": ["sm_107a"],
    "role": "fwd",
    "route": "both",
    "select_rule": None, # 임계값 없음 (기본 선택)
    "launch": {
        "block": (256, 1, 1),
        "dynamic_smem_bytes": 8192,
    }
},

# 2. 기존 커널의 임계값 적용 (Medium/Large size)
"cake_gated_act_mxfp8_0520c14fde53ef7bfea7": {
    "arches": ["sm_103a", "sm_107a"],
    "select_rule": {
        "elements": {
            "sm_103a": 2097152, # 2^21
            "sm_107a": 1048576, # 2^20 (신규 추가)
        },
        "kind": "elements_at_least",
    },
}

이 변경을 통해 $M \times K$가 $2^{20}$ 미만일 때는 e62d5f... 커널이 실행되고, 그 이상일 때만 기존의 무거운 커널이 실행됩니다.

2. build.json: SM107 전용 빌드 구성 및 CUDA 13.4 요구사항

새로운 아키텍처 최적화를 위해 빌드 설정이 추가되었습니다. 특히 sm_107a 아키텍처를 명시하고 최신 CUDA 13.4 기능을 활용하도록 설정되었습니다.

{
  "arch": "sm_107a",
  "gencode": "-gencode=arch=compute_107a,code=sm_107a",
  "min_cuda": "13.4",
  "name": "cake_gated_act_mxfp8_e62d5f6dda3083f31950_sm_107a",
  "nvcc_flags": [
    "--prec-div=true",
    "--prec-sqrt=true",
    "--ftz=false"
  ]
}

왜 이게 좋은 최적화인가?

1. Occupancy(점유율) 최적화

PR 설명에 따르면, $M \times K < 2^{20}$ 구간에서 실행되는 새로운 커널은 SM당 4개의 CTA(Collaborative Thread Array)를 실행하도록 설정되었습니다. 작은 워크로드에서는 GPU의 모든 자원을 쏟아붓는 것보다, 적절한 수의 블록을 배치하여 레이턴시를 숨기고(Latency Hiding) 리소스 경합을 줄이는 것이 성능에 유리합니다.

2. 실질적인 성능 향상 (Speedup)

벤치마크 결과, 특히 작은 배치 사이즈나 시퀀스 길이에 해당하는 128 x 128 설정에서 눈에 띄는 개선이 확인되었습니다.

  • Block-scaled inputs: 2.36 µs -> 2.20 µs (1.076x 향상)
  • Unit normals: 2.10 µs -> 2.09 µs (미세 향상)

데이터의 스케일링 특성에 따라 성능 향상 폭이 다르지만, 실제 LLM 서비스 환경에서 자주 발생하는 작은 크기의 연산에서 약 7.6%의 성능 이득을 얻었다는 점은 매우 고무적입니다.

3. 정교한 커널 선택 (Kernel Selection Strategy)

모든 상황에 맞는 단 하나의 '최강 커널'은 존재하지 않습니다. 이번 PR처럼 입력 텐서의 크기를 기준으로 커널을 분기하는 전략은 다음과 같은 교훈을 줍니다.

  • Small Workloads: 오버헤드 최소화 및 적절한 Occupancy 확보가 핵심.
  • Large Workloads: 처리량(Throughput) 극대화 및 메모리 대역폭 활용이 핵심.

결론

이번 최적화는 단순히 코드를 빠르게 만드는 것을 넘어, 하드웨어(SM107)의 특성과 소프트웨어 워크로드의 크기를 결합하여 최적의 지점을 찾아낸 사례입니다. FlashInfer는 이를 통해 최신 GPU 아키텍처에서도 빈틈없는 성능을 제공하게 되었습니다. 시니어 엔지니어로서 배울 점은, 성능 최적화의 끝은 알고리즘 개선뿐만 아니라 이러한 런타임 라우팅의 정교함에 있다는 것입니다.

참고 자료

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

댓글

관련 포스트

PR Analysis 의 다른글