High Performance LLM Inference in Pure Python with PyTorch Custom Ops
summary
Claims a pure-Python inference engine (SGLang/TokenSpeed) reaches parity with TensorRT-LLM by combining a multi-process scheduler, a zero-overhead batch scheduler that hides CPU work behind GPU execution, PyTorch custom ops exposing CUDA/Triton kernels, and captured CUDA graph replay. Benchmarks report 10–20% faster iteration times versus TensorRT and 2–3× faster decoding kernels in high-cache, long-context cases; the system is delivered as Python packages and shows GPUs kept fully utilized in production.
AI-written summary · about
chapters (13)
- 0:00Python LLM inference intro
- 1:27Python vs C++ inference
- 2:53Python as orchestration layer
- 3:36Multi-process runtime benefits
- 4:20CPU-GPU overlap with CUDA graphs
- 5:25PyTorch custom ops role
- 7:12State management challenges
- 9:41TokenSpeed engine overview
- 10:46Performance vs TensorRT-LLM
- 12:51Summary and future
- 13:12Q&A: State in heterogeneous systems
- 14:36Q&A: Performance guarantees vs TensorRT
- 16:43Q&A: Kernel optimizations
related talks
16mHow fast are LLM inference engines anyway?
24mLLM Inference Optimization From Simple To Sophisticated
14mLightning Talk: The Fastest Path to Production: PyTorch Inference in Python
29mMulti-Accelerator PyTorch Serving With NxD Inference and vLLM
20mUnlocking Performance: Harnessing LLMs To Streamline GPU Kernel Development in...
01h 02mWorkshop: LLMs Fine Tuning and Inferencing Using ON...
12mDemo: Optimizing Gemma inference on NVIDIA GPUs with TensorRT-LLM
24mOptimizing CPU LLM Inference in PyTorch: Lessons From VLLM