Lightning Talk: Pluggable PyTorch LLM Inference Architecture With VLL...
summary
Modular plugin code lets PyTorch models run on Trainium by registering a torch.neuron device backend so devices appear alongside CUDA and MPU, with automatic platform detection via Python entry_points. vLLM is extended with NeuronWorker and NeuronxDistributedModelRunner classes, NKI kernels implement attention and MoE, and execution flows through torch.compile → MLIR → HLO → neuron IR. Continuous batching with prefill/decode separation is used, and HuggingFace models load and execute on Trainium through standard vLLM APIs without hardware-specific code.
AI-written summary · about
chapters (9)
- 0:00Introduction
- 1:06LLM inference hardware challenges
- 2:09torch neuron as PyTorch backend
- 3:32Eager vs torch.compile modes
- 4:36Nikki compiler tradeoffs
- 6:25Graph scoping and MLIR pipeline
- 7:52PyTorch native VLLM plugin
- 8:54Scheduler customization
- 9:37Architecture benefits summary
related talks
29mMulti-Accelerator PyTorch Serving With NxD Inference and vLLM
27mServing PyTorch LLMs at Scale: Disaggregated Inference With Kubernetes and Llm-d
35mThe State of vLLM
39mAccelerating Open-Source RL and Agentic Inference with vLLM
25mvLLM: Easy, Fast, and Cheap LLM Serving for Everyone
19mNo GPU Left Behind: Scaling Online LLM Training With...
33mAccelerating the LLM Lifecycle on the Cloud
31mThe Evolution of Multi-GPU Inference in vLLM