蚂蚁金服蚂蚁集团-算子开发工程师-上海/北京
任职要求
1、硕士及以上学历，计算机/电子/通信等专业，4年以上相关工作经验，有AI算子开发经验者优先； 1、熟悉Linux开发环境，掌握Python/C++等语言， 具备良好的编程基础以及系统设计能力； 2、熟悉GPU或S…
工作职责
1、基于自研计算芯片，研发高性能算子、对算子和应用进行调优； 2、针对业务落地过程中的性能与泛化性问题，不断迭代算子库； 3、能够编写高效的计算核和数据流，交付性能极致的算子； 4、不断迭代算子工具链，以确保其完整性和易用性，和业务一起构建生态。
1. 面向具身智能、多模态 VLA/VLM、机器人策略模型等推理场景，设计并实现高性能 CUDA/C++ 算子。 2. 针对 H系列 GPU 与 RTX 4090 等硬件平台，进行端到端推理性能优化，包括延迟、吞吐、显存占用和稳定性优化。 3. 打通具身模型推理机部署流程，支持模型导出、量化、TensorRT/TRT-LLM 编译、服务化部署与性能 profiling。 4. 有 LLM/VLM/VLA 推理链路中的关键模块的优化经验，包括 attention、KV cache、视觉编码器、动作头、后处理等。 5. 针对跨卡、跨机推理通信场景，开发和优化 NCCL/RDMA/NVLink/PCIe 相关通信算子与调度策略。
1.针对 Hopper / Blackwell 架构进行 Kernel 开发与优化； 2.基于 CUDA / PTX / Triton 设计和实现大模型训练中的核心算子 Kernel：； 3.结合 Hopper / Blackwell 特性，在线程组织、访存路径、寄存器/共享内存使用、Tensor Core 调度等层面做深度优化； 4.构建高性能训练算子库； 5.针对 WeLM 模型结构进行算子设计与融合，降低 Kernel 调度和访存开销； 6.参与 WeLM 内部高性能算子库的建设，与 PyTorch 及 cuBLAS、cuDNN、CUTLASS 等库进行集成或替换； 7.Kernel 性能分析与瓶颈优化； 8.使用 Nsight Compute / Nsight Systems / nvprof 等工具做 Kernel 级性能剖析； 9.分析 SM 利用率、Tensor Core 利用率、访存带宽、指令吞吐、分支发散等指标，持续迭代 Kernel 实现； 10.为上层训练框架提供关于 batch 设计、sequence length、并行策略等与性能相关的技术建议； 11.前沿技术跟踪与工程落地； 12.跟踪 GPU 新架构、新指令、新库（如最新 CUTLASS、NCCL 等）； 13.学习业界优秀实现（Megatron-LM、DeepSpeed、Deepseek 等）的 Kernel/算子方案，并在 WeLM 体系内做工程化落地。
1.针对大模型训练、强化学习、推理场景，负责GPU kernel开发与调优 2.深入分析GPU/主流AI芯片的硬件架构特性，设计并实现高性能算子、算法和相关组件 3.保持关注行业前沿技术，且有能力和热情开展创新研究
职位描述： 1.负责AI处理器的高性能算子方案设计，功能开发，性能优化工作； 2.负责算子自测试用例的设计和开发，保证算子的质量提升和 CI 看护； 3.探索AI高性能算子的软硬件协同优化方案，提升硬件的利用效率 4.开发算子库开发所需的功能和性能分析工具，提升算子库开发效率