Lead Kernel Architect for AI Performance (Hybrid, Switzerland)

EPAM Systems; Zürich; Hybrid; CHF 180.000 - 240.000; Vollzeit

Benefits dieser Stelle

5 weeks of vacation

ESPP (Employee Stock Purchase Plan)

Enhanced parental leave

Extended pension plan

Sickness insurance

Employee assistance program

Travel medical and accident insurance

Learning and development opportunities

Zusammenfassung

EPAM Systems in Switzerland is seeking a Lead Kernel Engineer/Architect to drive kernel optimization for TPUs and GPUs in a hybrid work model. You will collaborate with researchers, compiler engineers and ML framework developers to deliver high-performance kernels and tools for large-scale training and inference.

The role requires deep expertise in C++/Python, kernel-level optimization, and experience with accelerator architectures, with opportunities to influence open-source ecosystems and

Qualifikationen

• Bachelor’s degree or equivalent practical experience.
• 12+ years of industry experience in software engineering or systems programming.
• 5+ years of experience in software development using C++ or Python.
• 3+ years of experience in testing, maintaining or launching software products and at least 1 year in software design or architecture.
• Hands-on experience in performance optimization at the kernel level for accelerators or high-performance systems.

Aufgaben

• Design and optimize high-performance kernels for TPU and GPU architectures using low-level programming frameworks such as Pallas, Triton or Mosaic.
• Build and maintain performance infrastructure, including benchmarking suites, autotuning systems, regression testing frameworks and tooling.
• Collaborate with ML framework developers (e.g., JAX, PyTorch) and compiler teams (XLA/MLIR) to integrate custom kernels and reduce performance bottlenecks.
• Track advancements in accelerator hardware, compiler technology and AI model design to identify opportunities for kernel-level optimization.
• Develop clear documentation, APIs and supporting OSS components that improve developer usability and adoption.
• Analyze and resolve complex performance issues impacting large-scale distributed training and inference systems.

Kenntnisse

C++

Python

Kernel optimization

Software architecture

Performance optimization

Ausbildung

Bachelor’s degree

Tools

Pallas

Triton

Mosaic

Jobbeschreibung

EPAM Systems in Switzerland is seeking a Lead Kernel Engineer/Architect to drive kernel optimization for TPUs and GPUs in a hybrid work model. You will collaborate with researchers, compiler engineers and ML framework developers to deliver high-performance kernels and tools for large-scale training and inference.

The role requires deep expertise in C++/Python, kernel-level optimization, and experience with accelerator architectures, with opportunities to influence open-source ecosystems and