关于此职位
Anyone AI 正在为一个专注于评估和改进 AI 工作负载内核开发任务的专项项目招募有经验的 AWS Trainium / Neuron Kernel Interface (NKI) 工程师。
我们正在寻找具备在 AWS Trainium 或 Inferentia2 硬件上构建或优化 NKI 内核的实践经验、并了解 Trainium 架构与传统 GPU 编程有何不同的工程师。
你将负责的工作
你将审查和评估涉及以下方面的技术任务：
- NKI 内核正确性和 Trainium 专用开发模式
- CUDA → NKI 内核迁移
- Trainium 性能优化和基准测试
- 跨 SBUF、PSUM 和 HBM 的内存管理
- 基于 tile 的计算和 DMA 调度
- CUDA/Triton 与 NKI 之间的跨平台数值正确性
- Trainium 专用性能瓶颈和优化机会
- 对内核实现的技术反馈和质量评估
这项工作涉及判断实现不仅在技术上正确，而且是符合 Trainium 硬件惯用方式并针对其优化的，而不是简单地从基于 GPU 的方法翻译而来。
我们在寻找
- 具备 2 年以上使用 Neuron Kernel Interface (NKI) 开发或优化内核的实践经验
- 具备使用 AWS Trainium 和/或 Inferentia2 的经验
- 深入理解：
- 基于 tile 的计算
- 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 项目的工程师来说，这是一个非常合适的机会。