Characterization of machine learning compilers for LLM inference on NVIDIA GPUs (link.springer.com)

Recent research has highlighted the intricate landscape of machine learning compilers (MLCs) utilized for large language model (LLM) inference on NVIDIA GPUs, addressing the critical P3 problem: balancing Performance, developer Productivity, and device Portability. By examining four key MLC tools—torch.compile, TensorRT, XLA, and ONNX Runtime—the study uncovers the trade-offs associated with deploying PyTorch-based models. Through a dual methodology that includes end-to-end benchmarks of state-of-the-art models like TinyLlama and Llama-2, the findings indicate that while Ahead-Of-Time (AOT) tools like TensorRT achieve peak performance, they are often incompatible with PyTorch, highlighting a need for specialized tools. Conversely, Just-In-Time (JIT) solutions such as torch.compile offer flexibility and portability but may not consistently accelerate LLMs. The significance of this research lies in its systematic approach to distinguishing the P3 characteristics of various MLCs and providing developers with actionable insights for model deployment. It reveals that optimizing for performance often comes at the cost of portability and productivity, necessitating a context-specific strategy when selecting an MLC. The study not only quantifies performance gains across different tools but also synthesizes guidelines to assist developers in navigating trade-offs effectively, ensuring an informed approach to leveraging the capabilities of NVIDIA GPUs for AI inference.