[Hiring] System Modeling (Computation) @Unconventional

Role Description

As a Member of Technical Staff, System Modeling (Computation), you will be part of a hands-on R&D team building simulation frameworks that enable evaluation and rapid iteration across all layers of unconventional physics-based computing systems for machine learning workloads. “Extreme co-design” is our guiding principle.

System Modeling is a multi-disciplinary effort, and the team we’re building reflects that. The role involves development of physics-based system models, GPU-accelerated ML system simulations, and cross-layer system integration. You don’t need to be an expert in all of these, but you have to be very strong in at least one, and solid in the rest.

Responsibilities

• Architecting Foundational Solvers:
  • Building large-scale, GPU-accelerated, high-fidelity numerical differential equation solvers (ODE, SDE, CDE, PDE).
  • Building tools that enable rapid iteration, multiple architectures, and rich metrics/visualization, leveraging frameworks best suited for scientific ML (e.g., JAX, PyTorch, or custom CUDA/Triton kernels).
• Bridging Physics and Machine Learning:
  • Developing physics-based surrogate models of device- and system-level behavior in unconventional compute.
  • Creating clean, composable abstractions that expose algorithm–hardware tradeoffs and enable cross-layer optimization via end-to-end autodiff.
• Extreme Co-Design & Collaboration:
  • Working closely with hardware and algorithm teams to understand their simulation needs.
  • Supporting everything from high-level algorithm development to the low-level verification of novel, analog hardware.

Qualifications

• MS/PhD in a quantitative field (AI/ML, Computer Science, Physics, Electrical Engineering, Applied Math), or BS with substantial, clear evidence of equivalent research/engineering depth.
• Deep expertise in numerical differential equation solvers (e.g., ODE, SDE, DDE) and their implementation on parallel architectures (e.g., Rosenbrock methods, Euler-Maruyama, adjoint methods, implicit solvers).
• Experience with high-performance, customized GPU kernel development for numerical methods, including GPU memory optimization and multi-GPU scaling.
• Experience building effective neural network surrogate models (e.g., Neural ODEs) for complex dynamical systems.
• Solid understanding of modern AI/ML architectures and training/inference workflows.
• Strong experience implementing and debugging ML models in PyTorch (preferred) or similar, with practical experience profiling, optimizing, and stabilizing non-trivial large-scale ML systems.

Preferred Qualifications

• Exceptional Python engineering skills with a passion for Developer Experience (DX), elegant API design, strong typing, and composability.
• Experience with compiler-friendly ML paradigms and internals (e.g., JAX vmap/pmap/jit, PyTorch autograd/torch.compile, custom XLA or Triton kernels).
• A track record of building open-source tools, scientific libraries, or serious simulation/modeling frameworks from scratch.
• Familiarity with analog dynamic systems, including transient responses, and nonidealities such as nonlinearity, quantization, random noise, and feedback/stability.
• Demonstrated ability to reason across multiple layers of the stack: algorithm, software, compiler, runtime, and hardware.
• Able to cleanly connect model architecture choices to system performance implications (memory bandwidth, communication patterns, latency, energy, and numerical stability).
• Experience applying efficiency techniques natively within modeling frameworks (quantization, sparsity, pruning, distillation, kernel fusion, etc.).

Benefits

• A comprehensive package including best-in-class health benefits.
• 401k matching.
• Truly unlimited PTO.
• Complimentary meals when working from our Palo Alto office.