岗位职责
- 负责AI4S模型(如MLFF、Cofolding等)的训练与推理性能优化,提升模型吞吐、资源利用率和整体计算效率;
- 基于PyTorch、CUDA、Triton等技术栈,分析和解决模型训练、推理过程中的性能瓶颈;
- 利用Torch.Compile、编译优化及算子融合等技术,持续提升模型执行效率和GPU利用率;
- 建立性能分析与Benchmark体系,使用Profiling工具定位系统热点并推动优化方案落地;
- 与算法、平台及科学计算团队紧密协作,共同推动AI4S模型和基础设施的性能演进。
职位要求
- 硕士学位及以上,计算机、软件工程、数学、物理等相关专业,3年以上AI系统性能优化工作经验;
- 熟练掌握Python和C++,具备扎实的软件工程能力和良好的编码习惯;
- 熟悉PyTorch训练与推理流程,具备模型性能优化的实际项目经验;
- 熟悉CUDA编程、GPU体系结构及性能分析工具(Nsight Systems、Nsight Compute、PyTorch Profiler等),能够独立完成性能瓶颈定位与优化;
- 具备Triton、Torch.Compile、分布式训练(FSDP、DeepSpeed等)相关经验,有训练加速、推理加速或GPU Kernel优化项目经验者优先。 加分项
- 具备CUDA Kernel开发或编译优化相关经验;
- 熟悉TorchInductor、TensorRT、XLA、TVM等技术;
- 具有大模型、科学计算或AI4S相关工作负载的性能优化经验。