この求人について
Anyone AI は、AIワークロード向けのカーネル開発タスクの評価と改善に特化したプロジェクトのために、AWS Trainium / Neuron Kernel Interface (NKI) の経験豊富なエンジニアを募集しています。
AWS Trainium または Inferentia2 ハードウェア上で NKI カーネルの構築や最適化を実際に行った経験があり、Trainium のアーキテクチャが従来の GPU プログラミングとどう異なるかを理解しているエンジニアを求めています。
業務内容
以下の技術的タスクのレビューと評価を担当していただきます：
- NKI カーネルの正確性と Trainium 固有の開発パターン
- CUDA → NKI カーネルの移行
- Trainium のパフォーマンス最適化とベンチマーク
- SBUF、PSUM、HBM 間のメモリ管理
- タイルベースの計算と DMA スケジューリング
- CUDA/Triton と NKI 間のクロスプラットフォームの数値的正確性
- Trainium 固有のパフォーマンスボトルネックと最適化の機会
- カーネル実装に対する技術的フィードバックと品質評価
この業務では、実装が技術的に正確であるだけでなく、GPU ベースのアプローチから単純に移植されたものではなく、Trainium ハードウェア向けに慣用的かつ最適化されているかどうかを判断することが求められます。
求める要件
- Neuron Kernel Interface (NKI) を用いたカーネルの開発または最適化の実務経験 2 年以上
- AWS Trainium および/または Inferentia2 の使用経験
- 以下に関する深い理解：
- タイルベースの計算
- SBUF / PSUM / HBM のメモリ階層
- パーティションディメンションの制約
- DMA オーケストレーション
- Trainium 固有の最適化手法
- CUDA → NKI 移行を評価する能力
- Trainium 上でのワークロードのプロファイリングと最適化の経験
- GPU と Trainium バックエンド間の数値的な差異に関する理解
- 複雑な技術的実装を分析し、明確な書面によるフィードバックを提供する能力
歓迎するスキル
- AWS Neuron SDK または Neuron Compiler の使用経験
- CUDA または Triton カーネル開発の経験
- NeuronCore-v2 アーキテクチャの知識
- FP32、BF16、FP8、INT8 ワークロードの経験
- Trn1 または Trn2 インスタンス上でのワークロードのベンチマーク経験
- nki.language、@nki.jit、または XLA カスタムコールに関する知識
- 技術評価、AI/ML データプロジェクト、RLHF、またはルーブリックベースの評価の経験
契約形態
勤務形態：リモート 契約形態：パートタイム、プロジェクトベースのコンサルティング フォーカス：AWS Trainium / NKI カーネルエンジニアリングと技術評価
AWS Trainium インフラストラクチャと低レベルの ML カーネル最適化に深く携わってきたエンジニアで、その専門知識を技術的に難易度の高い AI プロジェクトに活かしたい方に最適です。