Trainium (NKI) Kernel Expert Job at Mercor, San Francisco, CA

  • Mercor
  • San Francisco, CA

Job Description

Evaluate the quality, correctness, and hardware-appropriateness of Neuron Kernel Interface (NKI) development tasks used to train and evaluate a frontier AI lab's models. You'll assess CUDA→NKI migration fidelity, Trainium-specific performance-optimization quality, and cross-platform numerical-correctness standards, and provide clear, rubric-based written feedback.

Basic Qualifications

  • 2+ years of hands-on experience developing or optimizing kernels using the Neuron Kernel Interface (NKI) targeting AWS Trainium/Inferentia2 hardware
  • Strong understanding of NKI-specific development patterns: tile-based computation, SBUF/PSUM/HBM memory-hierarchy management, partition-dimension constraints, and DMA orchestration
  • Demonstrated experience assessing CUDA→NKI migration quality
  • Familiarity with Trainium-specific performance profiling (NeuronCore pipeline utilization, tensor-engine throughput, memory-bandwidth bottlenecks)
  • Experience defining or evaluating cross-platform numerical-correctness standards (GPU vs Trainium accumulation order, rounding behavior, mixed-precision semantics)

Preferred Qualifications

  • Direct experience with AWS Neuron SDK, Neuron Compiler internals, or contributions to NKI kernel libraries
  • Prior CUDA or Triton kernel development
  • Familiarity with Trainium hardware specifications (NeuronCore-v2 architecture, on-chip SRAM topology, supported data types: FP32/BF16/FP8/INT8)
  • Experience benchmarking ML training workloads on Trn1/Trn2 instances

Job Tags

Contract work

Similar Jobs

Tower Legal Solutions

Legal Administrative Consultant Job at Tower Legal Solutions

 ...:Our client, a large telecommunications company is seeking temporary legal administrative staff (see shifts below) to work onsite in its Bedminster, NJ offices assisting its Court Order Compliance Group. Job function includes data entry of legal process into client'... 

Gallin Associates

Endocrinologist Job at Gallin Associates

 ...Endocrinologist Compensation: $250K-$450K per year Job Type: Flexible hours, MondayFriday, No call. A well-established, physician-led group located in Suffolk County, Long Island, is currently seeking a Board-Certified or Board-Eligible Endocrinologist to become... 

Crusoe

Material Planner (Mechanical) Job at Crusoe

 ...parameters (order policy, lot sizing, lead times, safety stock, make/buy flags, and time fences) at MBOM release. Inventory Accuracy & ABC Classification: Partner with Materials Management to reconcile ERP balances against physical stock, while maintaining ABC usage... 

Farmers Restaurant Group

Line Cook - Tysons Job at Farmers Restaurant Group

 ...portioning ingredients and preparing our recipes to standard. They also adhere to our quality and sanitation protocols to ensure the highest food safety. As a Line Cook, you will be on a dynamic culinary team and enjoy a career path that challenges you as a culinary team member... 

FedEx Office

Retail Customer Service Associate Job at FedEx Office

 ...needs, suggest alternatives and provide solutions. This customer service focused team member is knowledgeable in all areas of the Stores...  ...pricing and product information, and recommending appropriate FedEx Office products and services. The Store Consultant will operate...