yoinka

AI Software Engineer — High-Performance

AMD

Shanghai, ChinaFull TimeMid
Sign in to applyVerified 1h ago
Location
Shanghai, China
Employment
Full Time
Work model
On-Site
Level
Mid

Skills

GitLLMPyTorchPython

About this role

WHAT YOU DO AT AMD CHANGES EVERYTHING   At AMD, our mission is to build great products that accelerate next-generation computing experiences—from AI and data centers, to PCs, gaming and embedded systems. Grounded in a culture of innovation and collaboration, we believe real progress comes from bold ideas, human ingenuity and a shared passion to create something extraordinary. When you join AMD, you’ll discover the real differentiator is our culture. We push the limits of innovation to solve the world’s most important challenges—striving for execution excellence, while being direct, humble, collaborative, and inclusive of diverse perspectives. Join us as we shape the future of AI and beyond.   Together, we advance your career.    Software Engineer — High-Performance GPU Communication (MORI) Role:  Software Engineer / Systems Engineer, GPU Networking & Inference Infrastructure  Team:  ROCm System Software — Communication Primitives About the Project MORI ( Mo dular  R DMA  I nterface) is an open-source framework powering AMD GPU communication in large-scale LLM inference. It provides the RDMA + GPU-direct networking layer for MoE expert parallelism and prefill/decode disaggregation in  SGLang  and  vLLM , and owns the KVCache management and storage layer via MORI-UMBP (Unified Memory & Bandwidth Pool). MORI achieved state-of-the-art results in the SemiAnalysis InferenceX v2 evaluation.

What You'll Work On

Inference framework integration:  Own end-to-end integration of MORI primitives into SGLang and vLLM — Python operator APIs, MORI-EP dispatch/combine in MoE forward passes, and MORI-IO in KVCache transfer pipelines. MORI-UMBP:  Integrate tiered KVCache storage and distributed key-value access into inference serving stacks. PD disaggregation:  Integrate MORI-IO into the prefill/decode path, enabling high-throughput KVCache transfers over GPU-direct RDMA. Expert Parallelism (EP):  Land and maintain MORI-EP in SGLang and vLLM, covering scheduling, routing, and EPLB for MoE models like DeepSeek V3 across 8–64 GPUs. MORI-SHMEM:  Integrate and maintain the symmetric memory runtime that underpins all MORI components — managing symmetric GPU memory allocation, RDMA transport initialization (IB, AINIC, Thor2), P2P/XGMI address translation, and device-side state for GPU kernels via MORI-IR bitcode. Performance benchmarking:  Design and run end-to-end benchmarks (throughput, TTFT, ITL) across EP and PD disagg configurations; drive optimization from profiling data.

Qualifications

Required:  - Deep familiarity with at least one major LLM inference framework (SGLang, vLLM, TensorRT-LLM, or equivalent) — scheduler, attention backend, KVCache manager, and distributed execution engine. - Strong understanding of LLM serving: MoE expert parallelism, prefill/decode disaggregation, KVCache reuse, tensor/pipeline/sequence parallelism. - Solid C++ and Python; comfortable in a mixed C++/HIP/Python codebase with PyTorch custom operator extensions. - Experience contributing to large open-source projects: upstream PRs, code review, cross-team coordination. Nice to Have:  - RDMA concepts: verbs API, queue pairs, completion queues, memory registration, GPUDirect Async (IBGDA). - Collective communication libraries (NCCL, RCCL, MPI) and their integration into distributed stacks. - GPU cluster network topologies: XGMI/NVLink (intra-node), InfiniBand/RoCE (inter-node), and their impact on MoE all-to-all patterns. - NIC vendor ecosystems (Mellanox ConnectX, AMD Pollara/AINIC, Broadcom Thor2) and userspace driver libraries. - Profiling network-bound workloads with rocprofv3, Perfetto, ibstat/perfquery. - ROCm, hipcc, or AMD GPU architecture experience. What Makes This Role Unique You'll own the bridge between MORI's low-level GPU networking layer and inference frameworks running trillion-parameter models at scale. MORI-EP and MORI-IO are already merged into SGLang and vLLM — your contributions ship directly to