Machine learning engineer in LLM inference optimization

Описание

Nebius is leading a new era in cloud infrastructure for the global AI economy. We are building a full-stack AI cloud platform that supports developers and enterprises from data and model training through to production deployment, without the cost and complexity of building large in-house AI/ML infrastructure.

Задачи

• Own optimization projects for specific model families, customer endpoints, or serving backends.
• Compare inference engines and recommend serving configurations suited to individual workloads.
• Diagnose model-quality and performance regressions during production rollouts.
• Improve LLM and VLM endpoint latency, throughput, memory efficiency, GPU utilization, quality, and cost per token.
• Deploy, configure, benchmark, and extend engines such as vLLM, SGLang, TensorRT-LLM, Triton Inference Server, and NVIDIA Dynamo.
• Develop production model-compression workflows covering quantization, quantization-aware training, distillation, low-bit serving, and accuracy recovery.
• Implement or integrate speculative decoding, draft-model approaches, KV-cache optimization, prefix caching, chunked prefill, continuous batching, and disaggregated prefill/decode serving.
• For prefill–decode disaggregation (PDD), evaluate KV-cache transfer, worker placement, and capacity balancing to determine which workloads benefit from the architecture.
• Design and improve LLM request routers and scheduling policies that balance worker utilization, queueing, request characteristics, and KV-cache locality while meeting latency and reliability targets.
• Scale dense and mixture-of-experts inference across multiple GPU nodes.
• Select and tune tensor, pipeline, data, and expert parallelism—including wide expert parallelism (WideEP)—with attention to hardware topology, expert load balance, and communication overhead.
• Build reproducible benchmark harnesses measuring TTFT, TPOT, tokens per second per GPU, p95/p99 latency, GPU memory, reliability, and cost per token.
• Validate architecture choices under representative traffic and consistent GPU budgets.
• Collaborate with GPU kernel and platform engineers to trace bottlenecks across model code, kernels, runtime, scheduler, gateway, and cluster layers.
• Produce design documents, performance reports, rollout plans, and customer-facing technical explanations.

Требования

• Strong engineering skills in Python and PyTorch.
• Hands‑on experience deploying or optimizing LLM, VLM, or high‑throughput transformer inference.
• Practical knowledge of at least one modern inference stack, such as vLLM, SGLang, TensorRT-LLM, Triton Inference Server, NVIDIA Dynamo, Ray Serve, KServe, or an equivalent internal system.
• A strong understanding of transformer inference bottlenecks involving KV cache, attention, memory bandwidth, batching, parallelism, and long‑context serving.
• Ability to quantify trade‑offs among latency, throughput, quality, utilization, and cost.
• Experience designing or optimizing distributed inference system architecture, with hands‑on work in one or more areas such as request routing, distributed scheduling, PDD, multi‑node inference, or expert parallelism.
• Strong communication and collaboration skills across research, kernel, infrastructure, product, and customer teams.
• Будет плюсом: Experience with quantization‑aware training, post‑training quantization, FP8, INT8, INT4, NVFP4, MXFP4, AWQ, GPTQ, or SmoothQuant;
• Experience with distillation, speculative decoding, EAGLE, Medusa, multi‑token prediction, or related inference acceleration methods;
• Experience supporting agentic workloads involving tool calling, structured outputs, streaming APIs, high concurrency, and multi‑step orchestration;
• Familiarity with CUDA or Triton;
• The role does not require kernel engineering to be the candidate’s primary specialization;
• Contributions to open‑source projects such as vLLM, SGLang, TensorRT-LLM, FlashInfer, LMCache, PyTorch, Triton, Ray, or KServe;
• Implementation experience with cache‑aware request routing, PDD, or WideEP, including diagnosing communication bottlenecks and load imbalance;
• Familiarity with GPU communication libraries and interconnects, such as NCCL, NVLink, InfiniBand, or RoCE.

Условия

• Competitive compensation
• Career growth and learning opportunities
• Flexibility and ownership
• Collaborative and innovative culture
• Opportunity to work on impactful AI projects
• International environment and talented teams
• Fast moving
• Bold thinking
• Constant growth
• Meaningful impact
• Trust and real ownership
• Opportunity to shape the future of AI