Lead/Staff AI-ML Software Engineer
amd
Job Description
- Kernel Development:
- Design and implement highly optimized C++ kernel library for NPU/GPU.
- Collaborate with the research and software teams to integrate these kernels into the existing software stack.
- Vector Processor Optimization:
- Work closely with hardware engineers to understand the architecture of VLIW vector core units such as MAC, GeMM, and non-linear functions.
- Develop vectorized code that leverages SIMD (Single Instruction, Multiple Data) and ILP (instruction level parallelism) for maximum performance.
- Performance Profiling and Tuning:
- Profile and analyze the performance of existing kernels.
- Identify bottlenecks and optimize critical sections for better throughput.
- Testing and Validation:
- Develop CPU models for the ML operators in C++/ Python to validate accuracy.
- Write unit tests and integration tests to ensure correctness and reliability.
- Validate kernel performance across different hardware platforms.
- Documentation and Collaboration:
- Document design specs for new kernels and the performance improvements.
- Follow coding guidelines, use tools like git to maintain code and create pull-requests, and documentation.
- Collaborate with cross-functional teams, including machine learning researchers and software engineers.