Performance Engineer, Kernels

Sarvam AI

Full Time5+ yearsPosted about 1 month ago

Let the right jobs find you

Get personalised suggestions from verified company career pages, matched to your role, location, level, and skills.

Overview

Position Type

Full Time

Experience

5+ years

Job Description

You will own the kernel layer. Where stock libraries - cuBLAS, cuDNN, FlashAttention, out-of-the-box Triton - leave performance on the table, you will author the custom CUDA, DSL-based, and PTX kernels that close the gap.

This is a hard, narrow, high-leverage role. We hire engineers who have shipped kernels that beat published baselines on real workloads, not engineers who have used kernels. When your code lands, production p99 moves, and you own the explanation of why.

What we're looking for

  • 5+ years in ML systems, with 2+ years authoring production CUDA kernels. You have a kernel in production that beat the prior baseline by a measurable margin.
  • CUDA at kernel-authoring level: thread-block sizing, shared-memory layout, warp primitives, async copies (cp.async, TMA), and MMA selection.
  • CUTLASS / CuTe DSL at a modify-and-extend level, with comfort in the layout algebra.
  • PTX at a debug-and-modify level - you have inserted hand-written PTX where the compiler missed.
  • Nsight Compute and Systems fluency: you read a roofline plot and propose the fix.
  • Attention kernels: you have authored or modified at least one (FlashAttention-family, paged, MLA, sliding-window, or sparse).
  • Multi-architecture awareness: what changes from Hopper to Blackwell (TMA, WGMMA, tcgen05).

Strong pluses

  • Communication kernels - NCCL / NVSHMEM authoring, custom collectives, expert-parallel dispatch (DeepEP-style), AFD bipartite comms (StepMesh-style), KV transfer (DualPath / Mooncake). Strongly desired; dedicated comms-specialist headcount is expected later.
  • Open-source kernel contributions - FlashAttention, CUTLASS examples, vLLM / SGLang kernels, DeepEP, Mooncake, or non-trivial Triton work. For this role, the GitHub filter is the highest-yield signal.
  • tcgen05, TMA, CTA-cluster launch and distributed shared memory, async pipelining, and the FP4/microscaling paths
  • Grace-side host-path optimization on GH200 / GB200.

Required Skills

CudaKernelCutlassCu TePtxNsight ComputeFlash AttentionGpu

About the Company

Sarvam AI

Bengaluru, India

Share This Job