AI Software Engineer

We are seeking an experienced AI Software Engineer to lead and contribute to our ONNX model optimization and model deployment infrastructure. You will work at the cutting edge of machine learning systems, bridging the gap between high-level deep learning frameworks and high-performance hardware execution through MLIR (Multi-Level Intermediate Representation) and ONNX (Open Neural Network Exchange) technologies.

Responsibilities:

1. Extensive experience with ONNX specification, protobuf definitions, and runtime implementations (ONNX Runtime, etc.)

2. Build and maintain end-to-end compilation pipelines from high-level frameworks to target hardware

3. Implement advanced optimization passes including loop transformations, memory layout optimizations, and operator fusion

4. Collaborate with hardware teams to define and implement target-specific code generation strategies

5. Lead ONNX model import/export, conversion, and optimization initiatives

6. Develop and maintain ONNX converters for internal frameworks and external tools

7. Implement ONNX graph optimizations, constant folding, and shape inference enhancements

8. Familiarity with deep learning model architectures (Transformers, CNNs, diffusion models) and their computational patterns

9. Experience with hardware accelerators (TPUs, NPUs, AI engines) and their compiler stacks

10. Architect unified pipelines connecting ONNX frontend with MLIR-based compiler backends

11. Implement quantization-aware transformations and mixed-precision support in the MLIR pipeline

12. Optimize end-to-end latency and throughput for production inference workloads

13. Profile and analyze compiler-generated code to identify optimization opportunities

14. Develop benchmarking infrastructure to validate correctness and measure performance gains

15. Experience with CI/CD for compiler infrastructure and deterministic build systems (Bazel, CMake)

Required Qualifications:

1. Bachelor's or master's degree in Data Science, Computer Science, Statistics, or a related field.

2. Fresh Graduate or Minimum of 1 year of experience working in AI, machine learning, or data science roles.

3. Prior experience in the embedded system industry or with a focus on applying AI in edge system is highly desirable.

4. Knowledge of git, docker, design pattern, machine learning algorithms and deep learning frameworks (e.g., PyTorch) is preferrable

5. Expert-level C++, Python for tooling and testing is preferrable.

6. Knowledge of GPU programming (CUDA/ROCm), vectorization, parallelization, and memory hierarchies is preferrable.