← All roles

AI Framework Engineer

AMD· Engineering
Apply on company site
ShanghaiFull-timePosted 2mo ago

About the role

Role focused on optimizing and developing deep learning frameworks for AMD GPUs, enhancing GPU kernels, DL models, and training/inference performance across multi-GPU/multi-node systems. Collaborative with internal teams and open-source maintainers, leveraging compiler technologies to drive production-quality optimizations.

Responsibilities

  • Deep Learning & LLM Framework Optimization: Optimize major DL/LLM frameworks (TensorFlow, PyTorch, vLLM, SGLang) for AMD GPUs and contribute improvements upstream.
  • GPU Kernel & Operator Optimization: Develop and tune GPU kernels and performance-critical operators to maximize throughput and minimize latency.
  • Model & Architecture Optimization: Adapt and optimize LLM architectures (e.g., Llama, Qwen, DeepSeek) and apply advanced techniques like FlashAttention, PagedAttention, and quantization.
  • End-to-End Performance Engineering: Perform comprehensive profiling to identify bottlenecks and implement system, memory, and communication optimizations across multi-GPU and multi-node setups.
  • Compiler & Pipeline Acceleration: Leverage advanced compiler technologies and graph compilers to enhance the full deep learning and inference pipeline.
  • Research & Advanced Techniques: Prototype and integrate emerging optimization methods such as speculative decoding and weight-only quantization into production systems.
  • Cross-Team & Open-Source Collaboration: Collaborate with internal GPU library teams and open-source maintainers to align improvements and ensure seamless upstream integration.
  • Software Engineering Excellence: Apply robust engineering practices to deliver maintainable, reliable, and production-quality performance optimizations.

Qualifications

  • Inference Frameworks, Model Architectures & Optimization Expertise: Strong practical experience with vLLM or SGLang, mastery of modern LLMs (e.g., DeepSeek, Qwen), strong theoretical grounding in Transformer/Attention/MoE/KV Cache, and hands-on application of advanced inference optimizations such as FlashAttention, PagedAttention, continuous batching, and quantization (INT8/INT4/GPTQ/AWQ).
  • End-to-End LLM Performance Engineering: Demonstrated ability to profile, diagnose, and optimize compute, memory, and communication bottlenecks across multi-GPU and multi-node environments.
  • High-Performance Computing: Expert experience running and optimizing large-scale workloads on heterogeneous clusters with a focus on efficiency, reliability, and scalability.
  • Deep Learning Framework Integration: Proven ability to integrate optimized GPU kernels into TensorFlow/PyTorch to accelerate large-scale training and inference with strong scaling and throughput.
  • Software Engineering Excellence & Community Contribution: Expert in Python/C++ coding skills, effective debugging and testing practices, proven ability to deliver maintainable performance-critical software, and a track record of open-source contributions with strong self-motivation.
  • GPU Kernel Development & Optimization is a plus: Hands-on solid experience designing and tuning high-performance GPU kernels for AMD GPUs using HIP, CUDA, ASM, and tools like CK, CUTLASS, and Triton, with strong knowledge of GCN/RDNA architectures.
  • Compiler & System-Level Optimization is a plus: Expert foundational knowledge of LLVM, ROCm, and compiler-driven techniques for improving kernel and system performance.

Skills

Required:C++ developmentPython

Benefits

  • Benefits offered are described in AMD benefits at a glance.
Last updated Oct 22, 2026. Interested? Apply on the company site.
AMD logo

About AMD

together we advance_AI

AMD delivers leadership high-performance and adaptive computing solutions to advance data center AI, AI PCs, intelligent edge devices, gaming, & beyond.

AIData CenterHigh-Performance ComputingGamingEdge Devices
Company profileWebsite

More open roles at AMD