deepseek-ai

TileKernels

AI#LLM#GPU#CUDA#PyTorch#Quantization

593

// 概要

TileKernels は TileLang フレームワークを使用して LLM 操作向けに設計された高性能な GPU kernel のコレクションを提供します。本プロジェクトには Mixture of Experts のルーティング、高度な量子化技術、および manifold hyper-connection 操作のための専門的な実装が含まれています。これらの kernel はハードウェア性能を最大化するように構築されており、現在内部のトレーニングや推論ワークフローで活用されています。

// 技術解説

Tile Kernels は TileLang ドメイン固有言語を活用し、 LLM 操作に特化して最適化された高性能な GPU カーネルを提供することで、ハードウェアの演算強度とメモリ帯域幅を理論上の限界まで引き出すことを目指しています。低レベルな GPU プログラミングを Python に抽象化することで、このプロジェクトはアジャイルな開発と、 Mixture of Experts (MoE) ルーティングや高度な量子化といった複雑な操作の容易な移行を実現します。現時点では完成されたドキュメントよりもパフォーマンスを優先していますが、低レベルカーネルと高レベルな PyTorch autograd ラッパーの両方を提供することで、本番環境グレードの学習と推論のための強固な基盤を提供します。

// 主要ハイライト

Mixture of Experts アーキテクチャにおける Top-k エキスパート選択のための効率的な Gating メカニズムを実装しています。

トークンからエキスパートへのマッピングや、融合された expansion/reduction 操作を含む、包括的な MoE ルーティングサポートを提供します。

FP8、 FP4、 E5M6 キャスティングなどの高度な量子化技術を、融合された SwiGLU 操作とともにサポートします。

最適化された学習のために、 RMSNorm と重み勾配の削減を融合した特殊な Engram gating カーネルを含みます。

複雑なモデルアーキテクチャをサポートするために、 Sinkhorn 正規化のような Manifold HyperConnection カーネルを特徴としています。

開発者が低レベルカーネルを学習可能な PyTorch レイヤーに直接統合できる、高レベルな torch.autograd.Function ラッパーを提供します。

// ユースケース

Mixture of Experts (MoE) のルーティングおよび gating 操作

fused SwiGLU をサポートした FP8、FP4、および E5M6 量子化

学習可能なモデリングレイヤーのための高レベルな PyTorch autograd ラッパー

// クイックスタート

開始するには、 Python 3.10+、 PyTorch 2.10+、および NVIDIA SM90 または SM100 GPU を含む要件を環境が満たしていることを確認してください。リリースバージョンをインストールするには 'pip install tile-kernels' を、ローカル開発環境をセットアップするには 'pip install -e ".[dev]"' を使用してください。その後、プロジェクト構造を探索して特定のカーネルを利用したり、提供されている pytest スイートを実行して正確性の検証やパフォーマンスのベンチマークを行うことができます。