[flashinfer] FlashInfer의 Kimi-K3 AttnRes 연산자 최적화: SM100/SM103를 위한 성능 향상

PR 링크: flashinfer-ai/flashinfer#6980 상태: Merged | 변경: +37730 / -56984

들어가며

최근 대규모 언어 모델(LLM)의 발전은 모델의 추론 속도와 효율성을 극대화하는 기술에 대한 수요를 끊임없이 증가시키고 있습니다. 특히, 어텐션 메커니즘은 LLM의 핵심 구성 요소로서, 이 부분의 최적화는 전체 모델 성능에 지대한 영향을 미칩니다. NVIDIA의 최신 GPU 아키텍처인 Hopper (SM100) 및 Ada Lovelace (SM103)를 지원하기 위해, FlashInfer 라이브러리는 cake_kimi_k3_attn_res 연산자의 성능을 개선하는 중요한 업데이트를 진행했습니다. 이 PR은 기존의 Cake 빌드 패키지를 최신 커널 소스에 맞춰 재렌더링하고, 특정 GPU 아키텍처 및 설정에 대한 최적화된 프로그램을 추가함으로써, 특히 K=5와 같이 이전에 지원되지 않던 시나리오에 대한 성능을 향상시키는 것을 목표로 합니다.

본 글에서는 이 PR이 가져온 코드 변경 사항을 분석하고, 이러한 변경이 왜 성능 향상으로 이어지는지, 그리고 그 기술적 의미는 무엇인지 심층적으로 살펴보겠습니다.

코드 분석

이번 PR의 핵심은 flashinfer.experimental.cake_kimi_k3_attn_res 연산자의 CUDA 커널을 최신 GPU 아키텍처에 맞게 재구성하고, 다양한 K 값과 M 값 조합에 대한 최적화된 프로그램을 추가하는 것입니다. 주요 변경 사항은 README.md와 build.json 파일에서 확인할 수 있습니다.

1. README.md 변경 사항: 프로그램 설명 업데이트

README.md 파일은 cake_kimi_k3_attn_res 연산자가 지원하는 다양한 프로그램 유형과 그 특성을 설명합니다. 이번 PR에서는 새로운 native 프로그램들이 추가되고 기존 프로그램들의 설명이 업데이트되었습니다.

변경 전:

| `native_{k8,m128}` | installed native ports, one 288-thread CTA per SM (grid = the device's SM count: 148 on B200 / B300, 152 on GB300), three sources per chunk, depth 2 | the measured cells above the small-M table on the qualified SM counts `NATIVE_SM_COUNTS`: `M = 1024` (`sm_100a`) / `M = 256` (`sm_103a`) for K4 (`native_m128`), `M = 256` for K8 (`native_k8`) |

변경 후:

| `native_{k5,k7,k8,m128}` | installed native ports, one 288-thread CTA per SM (grid = the device's SM count: 148 on B200 / B300, 152 on GB300), three sources per chunk, depth 2 | the measured cells above the small-M table on the qualified SM counts `NATIVE_SM_COUNTS`: `M = 1024` (`sm_100a`) / `M = 256` (`sm_103a`) for K4 (`native_m128`), `M = 256` for K8 (`native_k8`), the round-2 native cells of the `R2_NATIVE_CELLS` table (K7 / K8 / M128), and the K5 cells of `R2_NATIVE_CELLS` (`native_k5`, whose consumed-slot release every consumer thread performs): `M = 192 / 256 / 768` on `sm_103a`, `M = 192 / 256 / 384 / 768 / 1024 / 1536 / 2048` on `sm_100a` |

분석:

native 프로그램의 설명이 확장되어 k5, k7이 추가되었고, 각 K 값에 대한 지원 M 값의 범위가 구체적으로 명시되었습니다. 특히 sm_100a (B200) 및 sm_103a (B300, GB300) 아키텍처별로 지원되는 M 값의 범위가 상세히 기술된 것은, 해당 GPU에서 최적의 성능을 내기 위한 세밀한 튜닝이 이루어졌음을 시사합니다. 또한, native_k5 프로그램의 경우 consumed-slot release every consumer thread performs라는 설명은 스레드 간의 동기화 및 자원 관리에 대한 최적화가 포함되었음을 나타냅니다.

persistent 프로그램 설명의 변경점도 주목할 만합니다.

변경 전:

| `persistent` | persistent TMEM common path, 288 threads, `nc` sources per chunk (1..5), depth 2 or 3, retraced per cell | every other dense call with `delta` and `output_norm_weight` and no snapshot write |

변경 후:

| `persistent` | persistent TMEM common path, 288 threads, `nc` sources per chunk (1..5), depth 2 or 3, retraced per cell | every other dense call with `delta` and `output_norm_weight` and no snapshot write; a key's `_crfnone` suffix names the fence-authored twin of its program (built with `consumed_release_fence_mode="none"`: the kernel states its own tcgen05 fences; schedule label `_authored_fences`), which the twelve `sm_100a` K2 / K3 cells of `R2_PERSISTENT_KEY_CELLS` at `M >= 2048` run |

분석:

_crfnone 접미사가 붙은 프로그램은 `consumed_release_fence_mode=

참고 자료

• https://github.com/flashinfer-ai/flashinfer/blob/main/flashinfer/experimental/cake_kimi_k3_attn_res/README.md
• https://github.com/flashinfer-ai/flashinfer/blob/main/flashinfer/experimental/cake_kimi_k3_attn_res/build.json

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