Skip to content
pageshift.aiTalk to us (opens in a new tab)
Menu
All open roles

Engineering

ML Performance Engineer

Profile and optimize long-context JAX workloads on TPUs.

About us

Pageshift is a research lab committed to pushing the frontier of AI storytelling and creativity. We are envisioning a world in which most entertainment is personalized and AI-generated. Our goal is to build the underlying story engine that powers it all. To do this, we are not afraid to explore new ways and create novel categories of model capability.

The role

You will focus on XLA behavior and Pallas kernel development: profiling, benchmarking, implementing kernel-level optimizations and validating each improvement with data. The work includes memory behavior, sharding and end-to-end efficiency for long-context training.

What you’ll do

  • Implement and optimize JAX workloads for TPUs
  • Develop and tune Pallas kernels
  • Profile memory, sharding and execution bottlenecks
  • Benchmark changes and validate improvements with data

What we’re looking for

  • Passion for entertainment and storytelling
  • A basic understanding of TPUs, JAX and XLA
  • A willingness to optimize against a black box

Nice to have

  • A relevant project you can demonstrate
  • Experience writing Pallas kernels