# \#autotuning **URL:** https://discuss.pytorch.kr/tag/autotuning/5553.md [최신](https://discuss.pytorch.kr/latest.md) · [카테고리](https://discuss.pytorch.kr/categories.md) · [태그](https://discuss.pytorch.kr/tags.md) --- ## [Tokamax: JAX와 Pallas 기반의 커스텀 가속 커널 라이브러리](https://discuss.pytorch.kr/t/tokamax-jax-pallas/7801)
**Author:** [@9bow](https://discuss.pytorch.kr/u/9bow)\ **댓글:** 0\ **마지막 업데이트:** [9월 20, 2025, 3:30오전 UTC](https://discuss.pytorch.kr/t/tokamax-jax-pallas/7801 "2025-09-20T03:30:13Z")
Tokamax 소개 Tokamax는 NVIDIA GPU와 Google TPU에서 동작하는 커스텀 가속 커널(custom accelerator kernels)을 제공하는 Python 라이브러리입니다. 이 프로젝트는 최신 JAX와 Pallas를 기반으로 만들어졌으며, 딥러닝 모델에서 중요한 연산을 최적화된 커널로 교체해 성능을 크게 끌어올릴 수 있도록 설계되었습니다. 단순히 제공된 커널만 사용하는 것이 아니라, 사용자가 직접 자신만의 커널을 작성하고 자동 튜닝(autotuning) 기능을 통해 최적화할 수 있다는 점이 핵심적인 특징입니다. 현재 Tokamax는 개발 초기 단계에 있으며, API 변경 가능성이 크고 일부 기능은 완전하지 않습니다. 하지만 제공되는 커널은 최신 연구 성과를 반영하고 있어, 연구자와 엔지니어가 GPU 및 TPU 환경에서 고성능 실험을 진행할 때 유용합니다. 예를 들어, Transformer 모델의 핵심 연산인 tokamax.dot\_product\_attention(FlashAttention 기반), tokamax.gated\_linear\_unit(Gated linear units (SwiGLU 등)), tokamax.layer\_norm(Layer Normalization과 RMSNorm) 같은 고성능 커널을 지원합니다. Tokamax의 또 다른 장점은 GPU와 TPU 모두에서 동작하는 ragged\_dot 연산을 지원한다는 점입니다. 이는 Mixture of Experts(MoE) 아키텍처에서 중요한 연산으로, 대규모 모델의 효율적 학습과 추론을 가능하게 합니다. 따라서 Tokamax는 최신 LLM과 딥러닝 연구에 적합한 라이브러리라 할 수 있습니다. Tokamax는 단순히 커널을 제공하는 라이브러리가 아니라, 여러 구현체를 선택하거나 자동으로 최적화된 커널을 고르는 기능을 제공합니다. 예를 들어 implementation=None 옵션을 주면 Tokamax가 커널의 입력 크기와 실행 환경에 따라 최적의 구현체를 자동 선택합니다. …