ML Compiler Intern - Kernel Auto-Tuning

What You'll Learn

• Project Overview: Build a cost-model-guided auto-tuning framework for PyTorch and Triton kernels on Samsung’s compute-near-memory AI accelerator, and make it placement-aware, so that each kernel is tuned for where its data lives. You’ll connect it to torch.compile and measure the gains on LLM inference.
• Skills You’ll Learn:
  • Performance optimisation for a cache-less, DMA-driven accelerator: search spaces, legality rules and cost models
  • PyTorch compiler internals (torch.compile, FX, Inductor autotuning) and Triton kernel compilation
  • Benchmarking and profiling LLM inference kernels (GEMM, attention) on simulators and hardware

What You’ll Do

Our ML Compiler team builds the software that runs PyTorch models on Samsung’s compute-near-memory AI accelerator: a PyTorch device and compiler backend, a Triton-based kernel compiler and an LLVM back end. We work closely with the runtime and hardware teams and with the PyTorch and Triton open-source communities. Over the summer you’ll own one piece of making ordinary PyTorch models fast on new memory-centric hardware, from research idea to merged, measured code.

Location: Daily onsite presence at our San Jose Headquarters in alignment with our Flexible Work policy

Reports to: Director

• Design a search space for Triton kernel configurations on the accelerator: tile sizes, launch-grid mapping, on-chip buffer budgets and load style.
• Build legality rules with our compiler engineers that prune configurations that will not compile or will not fit in on-chip memory.
• Develop a kernel cost model, calibrated against our simulator and profiler, that ranks candidates so only the most promising are benchmarked.
• Build a reliable benchmarking harness and a persistent tuning cache, so that tuning cost stays practical.
• Integrate the tuner with the torch.compile Inductor autotuning path and with Triton’s autotuner.
• Extend tuning to data placement: tune kernels for a given tensor layout and export tuned costs that a placement pass can use.
• Measure kernel-level and end-to-end LLM inference gains, and share results in weekly syncs and a final presentation.
• Collaborate with PyTorch and Triton developers on upstream autotuning hooks, and hand off documented, tested code.
  
• Complete other responsibilities as assigned.

What You Bring

• Pursuing a Masters or PhD in Computer Science, Computer Engineering or Electrical Engineering; PhD preferred.
• Must have at least 1 academic quarter/semester remaining
• Strong Python and C++, with experience writing or tuning Triton kernels or working with PyTorch compiler internals (torch.compile, FX, Inductor).
• Grounding in compilers or performance engineering: IR and passes, cost models, roofline analysis and profiling.
• Preferred: auto-tuning or auto-sharding experience (e.g. TVM/Ansor, the Halide autoscheduler, Kernel Tuner, Alpa, GSPMD), or search methods such as Bayesian optimisation or ILP.
• Preferred: MLIR/LLVM back ends, scratchpad or near-memory architectures, LLM inference kernels, or publications or open-source PyTorch/Triton contributions.
  
• You’re inclusive, adapting your style to the situation and diverse global norms of our people.
• An avid learner, you approach challenges with curiosity and resilience, seeking data to help build understanding.
• You’re collaborative, building relationships, humbly offering support and openly welcoming approaches.
• Innovative and creative, you proactively explore new ideas and adapt quickly to change.

#LI-AD1

Hourly Base Pay Rate 28 - 68 USD