AI Framework Engineer

AMD
Shanghai
Workplace: OnsiteFull timeCNY 411,670 - 588,100 annuallyFunction: Data Science & Machine LearningExperience: 3+ yearsEducation: mastersSkills: ["Problem-solving","Collaboration","Independent work","Analytical thinking","Software engineering best practices"]

Optimize and develop deep learning frameworks for AMD GPUs, improving GPU kernels, training/inference performance, and end-to-end pipelines across multi-GPU and multi-node systems. Collaborate with internal GPU library teams and open-source maintainers to integrate performance optimizations upstream. Use advanced compiler technologies (including graph compilers) and follow software engineering best practices to deliver robust, maintainable improvements in open-source repositories like TensorFlow and PyTorch.

Loading

Loading job details...

Preparing the role view and application actions.

FursaFursa
AMD
AMD
1 day ago

AI Framework Engineer

✓ Verified Job

Canonical indexed version, validated from employer's careers page.

Source: Company careers pageValidated by: Fursa AI
Last checked: 12 hours agoStatus: Live

Job Summary

Optimize and develop deep learning frameworks for AMD GPUs, improving GPU kernels, training/inference performance, and end-to-end pipelines across multi-GPU and multi-node systems. Collaborate with internal GPU library teams and open-source maintainers to integrate performance optimizations upstream. Use advanced compiler technologies (including graph compilers) and follow software engineering best practices to deliver robust, maintainable improvements in open-source repositories like TensorFlow and PyTorch.
Location: Shanghai
Workplace: Onsite
Employment Type: Full time
Job Function: Data Science & Machine Learning
Seniority: Mid level

Key Responsibilities

  • •Optimize deep learning frameworks (e.g., TensorFlow and PyTorch) for AMD GPUs in open-source repositories.
  • •Develop and optimize GPU kernels for high performance on specific AI operations.
  • •Design and optimize deep learning models for AMD GPU performance.
  • •Collaborate with internal GPU library teams to analyze and improve training and inference performance.
  • •Work in distributed computing environments to optimize performance across multi-GPU (scale-up) and multi-node (scale-out) systems, leveraging compiler technologies and end-to-end pipeline optimizations.

Pay and Benefits

Salary: CNY 411,670 - 588,100 annually

Key Requirements

  • •3+ years of professional experience in technical software development focused on GPU optimization, performance engineering, and framework development.
  • •Strong C++ development skills in Linux environments with debugging, performance tuning, and test design.
  • •Experience optimizing GPU kernels for deep learning on AMD GPUs, including using HIP, CUDA, and/or assembly (ASM).
  • •Experience integrating optimized GPU performance into ML frameworks such as TensorFlow and PyTorch for accelerated training and inference.
  • •Bachelor’s and/or Master’s degree in Computer Science, Computer Engineering, Electrical Engineering, or a related field.
Experience:3+ yearsDeep learningHigh-performance computingOpen sourceGPU optimization
Education:Master's in Computer Science, Computer Engineering, Electrical Engineering, or related field
Skills:Problem-solvingCollaborationIndependent workAnalytical thinkingSoftware engineering best practices
Languages:English
Tech Stack:C++PythonLinuxTensorFlowPyTorchAMD GPUsGPU kernelsHIPCUDAAssembly (ASM)GCNRDNACompute Kernel (CK)CUTLASSTritonLLVMROCmGraph compilersMulti-GPUMulti-node

Company Brief

AMD
Designs and produces semiconductor products including CPUs, GPUs, and adaptive SoCs for consumer, enterprise, and embedded markets, competing across PCs, data centers, and gaming industries.
Industry: Electronics Manufacturing
Company Size: Enterprise (1,001+ employees)
Revenue: USD 1B+
Growth: Public Company
Valuation: Public Company (Market Cap in USD)
Funding: IPO / Publicly Listed
Headquarters: Santa Clara, United States
Founded: 1969
Glassdoor
Glassdoor: 3.9
WebsiteLinkedIn