Cerebras is hiring a remote Full Stack ML Engineer!
Cerebras, the company behind a wafer-scale AI chip far larger than conventional GPUs, is hiring an ML engineer to bridge model code and accelerator hardware. The position sits at the intersection of machine learning and compilers: you take a model from a reference implementation in PyTorch and drive it down to fast, numerically faithful execution on Cerebras hardware. Day to day that means owning model bring-up, extending MLIR dialects and lowering pipelines, and tuning kernels and runtime behaviour for inference workloads. The role is fully remote and involves close collaboration with hardware, compiler, kernel, and runtime engineers. It suits someone who is equally comfortable reading a transformer diagram and a transformation pass, and who enjoys chasing correctness and performance bugs that span several layers of the stack.
Key responsibilities
- Own end-to-end bring-up of new models: interpret architectures, convert weights, build supported execution paths, and verify outputs against reference implementations.
- Extend MLIR dialects, rewrite patterns, graph transformations, and hardware-specific lowering passes.
- Enable core operators such as attention, matmul, normalization, and positional embeddings through compiler and kernel work.
- Improve execution efficiency via fusion, tensor layouts, tiling, memory allocation, data movement, and parallelism.
- Tune inference for prefill and decode, including KV cache handling, batching, and quantization, to reduce latency and memory use.
- Assess numerical differences and quantify accuracy impact from precision changes and compiler optimizations.
- Profile workloads with traces and hardware counters to locate compute, memory, communication, and runtime bottlenecks.
- Partner with hardware, compiler, kernel, and runtime teams to ship dependable model support and reproducible benchmarks.
Key requirements
- Strong C++ and Python programming ability.
- Hands-on experience bringing up and debugging models in PyTorch or similar.
- Practical MLIR experience: dialects, rewrite patterns, passes, and lowering pipelines.
- Solid compiler fundamentals, including IRs, dataflow analysis, and code generation.
- Knowledge of transformer architectures, attention, tensor operations, and numerical precision.
- Track record profiling and optimizing workloads on GPUs or other accelerators.
- Comfort debugging across model code, compiler output, kernels, and runtime.
- Nice to have: LLM inference techniques (GQA, sliding-window attention, MoE, KV caching, speculative decoding), low-precision formats such as FP16/BF16/FP8, accelerator kernel or backend development, distributed and model-parallel execution, and open-source contributions to MLIR, LLVM, or inference frameworks.
- Seniority: Not specified
- Work mode: Remote
- Key skills: C++, Python, PyTorch, MLIR, Compiler design, LLVM, Transformer architectures, Performance profiling, GPU optimization, Quantization, LLM inference, KV caching
Want more jobs?
Join our mailing list to receive updates about new data science and Full Stack ML Engineer jobs.
