TileKernels: DeepSeek AI가 공개한, LLM 연산을 위한 고성능 TileLang 기반 GPU 커널 라이브러리

읽을거리&정보공유
, , , , , , ,
1

TileKernels: DeepSeek AI가 공개한, LLM 연산을 위한 고성능 TileLang 기반 GPU 커널 라이브러리
TileKernels: DeepSeek AI가 공개한, LLM 연산을 위한 고성능 TileLang 기반 GPU 커널 라이브러리1536×1024 330 KB

TileKernels 소개

현대 대규모 언어 모델(LLM)의 학습과 추론에서 GPU 커널 최적화는 연산 효율을 좌우하는 핵심 요소입니다. 특히 Mixture of Experts(MoE) 아키텍처가 널리 쓰이면서 게이팅(Gating), 전문가 라우팅(Expert Routing), 양자화(Quantization) 등 다양한 연산 커널의 하드웨어 수준 최적화 필요성이 높아졌습니다. 기존에는 커스텀 CUDA 커널을 작성하기 위해 C++와 CUDA 내부를 깊이 이해해야 했기 때문에 진입 장벽이 높았고, 새로운 아키텍처가 등장할 때마다 커널을 처음부터 다시 작성해야 하는 문제가 있었습니다.

TileKernels는 DeepSeek AI가 공개한 오픈소스 GPU 커널 라이브러리로, LLM 연산에 특화된 고성능 커널 모음입니다. 이 프로젝트는 TileLang을 기반으로 구축되었습니다. TileLang은 Python으로 고성능 GPU 커널을 표현할 수 있도록 설계된 도메인 특화 언어(Domain-Specific Language, DSL)로, 쉬운 마이그레이션, 빠른 개발, 자동 최적화를 특징으로 합니다. 기존 CUDA C++ 대신 Python 수준에서 커널을 작성할 수 있어 개발 생산성이 크게 향상됩니다.

프로젝트의 커널 대부분은 연산 집약도(Compute Intensity)와 메모리 대역폭(Memory Bandwidth) 측면에서 하드웨어 성능 한계에 근접하도록 설계되었습니다. 일부 커널은 이미 DeepSeek 내부의 학습 및 추론 시나리오에서 사용되고 있으며, MoE 기반 모델의 핵심 연산들을 담고 있습니다. 다만 개발팀은 현재 코드가 베스트 프랙티스를 대표하지는 않으며, 코드 품질과 문서화를 지속적으로 개선하는 중이라고 밝히고 있습니다.

TileKernels의 주요 커널 구성

TileKernels는 LLM 연산에 필요한 다양한 커널을 체계적으로 제공합니다. 각 커널 모음은 저장소의 tile_kernels/ 디렉토리 아래 영역별로 정리되어 있습니다.

게이팅(Gating): MoE(Mixture of Experts) 라우팅에서 각 토큰이 어떤 전문가에게 배정될지를 결정하는 Top-k 전문가 선택 및 스코어링 커널입니다. MoE 모델의 추론 효율에 직접적인 영향을 미치는 핵심 연산입니다.

MoE 라우팅(MoE Routing): 토큰-전문가 매핑, 융합된 확장/축소(Fused Expansion/Reduction), 가중치 정규화(Weight Normalization)를 처리합니다. 관련 커널은 tile_kernels/moe/에 위치하며, 토큰을 전문가에게 효율적으로 디스패치하고 결과를 집계하는 과정을 최적화합니다.

양자화(Quantization): 토큰 단위(Per-token), 블록 단위(Per-block), 채널 단위(Per-channel)의 FP8/FP4/E5M6 캐스팅을 tile_kernels/quant/에서 지원하며, SwiGLU와 양자화를 하나로 합친 융합 연산(Fused SwiGLU + Quantization)도 포함됩니다. 저정밀도 양자화는 모델 크기와 추론 속도 모두에서 중요한 최적화 기법입니다.

전치(Transpose): tile_kernels/transpose/의 배치 전치(Batched Transpose) 연산을 고성능으로 처리합니다.

Engram: tile_kernels/engram/의 Engram 게이팅 커널로 융합된 RMSNorm, 순전파/역전파 패스, 가중치 그래디언트 리덕션을 포함합니다. 메모리 접근 패턴을 최적화하여 학습 효율을 향상시킵니다.

Manifold HyperConnection: tile_kernels/mhc/의 하이퍼커넥션 커널로, Sinkhorn 정규화와 믹스 분할/적용(Mix Splitting/Application)을 포함합니다.

모델링(Modeling): tile_kernels/modeling/에서 torch.autograd.Function을 상위 수준으로 래핑하여 저수준 커널들을 학습 가능한 레이어(Engram Gate, mHC 파이프라인)로 조합합니다. PyTorch의 자동 미분 시스템과 자연스럽게 통합됩니다.

이 외에도 tile_kernels/torch/에는 PyTorch 기반 레퍼런스 구현이, tile_kernels/testing/에는 테스트 및 벤치마크 유틸리티가 포함되어 있습니다.

TileKernels 설치 및 사용법

TileKernels는 Python 3.10 이상, PyTorch 2.10 이상, TileLang 0.1.9 이상이 필요하며, NVIDIA SM90(Hopper, H100 계열) 또는 SM100(Blackwell) 아키텍처 GPU와 CUDA Toolkit 13.1 이상이 필요합니다.

릴리즈 버전 설치:

pip install tile-kernels

개발 버전 설치:

git clone https://github.com/deepseek-ai/TileKernels.git
cd TileKernels
pip install -e ".[dev]"

테스트 실행:

테스트는 pytest 로 수행합니다. 단일 테스트 파일은 정확성만 검증하거나 벤치마크를 함께 돌릴 수 있습니다.

# 정확성 검증 (워커 4개)
pytest tests/transpose/test_transpose.py -n 4

# 정확성 검증 + 벤치마크
pytest tests/transpose/test_transpose.py --run-benchmark

# 압박 테스트(pressure test)
TK_FULL_TEST=1 pytest -n 4 --count 2

저수준 커널들은 tile_kernels/modeling/의 torch.autograd.Function 래퍼를 통해 학습 가능한 레이어로 조합되므로, PyTorch 학습 파이프라인에 그대로 통합할 수 있습니다.

TileKernels의 라이선스

TileKernels는 MIT 라이선스로 공개되어 있어 개인 및 상업적 목적으로 자유롭게 사용, 수정, 배포할 수 있습니다.

:github: TileKernels 프로젝트 GitHub 저장소

더 읽어보기



이 글은 GPT 모델로 정리한 글을 바탕으로 한 것으로, 원문의 내용 또는 의도와 다르게 정리된 내용이 있을 수 있습니다. 관심있는 내용이시라면 원문도 함께 참고해주세요! 읽으시면서 어색하거나 잘못된 내용을 발견하시면 덧글로 알려주시기를 부탁드립니다. :hugs:

:pytorch:파이토치 한국 사용자 모임:south_korea:이 정리한 이 글이 유용하셨나요? 회원으로 가입하시면 주요 글들을 이메일:love_letter:로 보내드립니다!
텔레그램(Telegram)이나 Slack/Discord/Teams/Dooray/GoogleChat 등으로도 새 글 알림을 받으실 수 있습니다. :smiley:

:wrapped_gift: 아래:down_right_arrow:쪽에 좋아요:+1:를 눌러주시면 새로운 소식들을 정리하고 공유하는데 힘이 됩니다~ :star_struck: