Hybrid JIT-CUDA Graph Optimization for Low-Latency Large Language Model Inference

Technical Summary: Hybrid JIT–CUDA Graph Optimization for Low-Latency LLM Inference

1. Problem Statement

The deployment of Large Language Models (LLMs) in interactive, real-time applications (e.g., conversational agents, code assistants) is primarily bottlenecked by inference latency and latency variance (jitter). During the autoregressive decoding phase, the system must execute numerous fine-grained GPU kernels sequentially for every new token generated. This process incurs significant overhead from:

• CPU-bound dispatch: The host must constantly coordinate and launch kernels.
• Python overhead: High-level language dispatch and synchronization events add latency.
• The "Static vs. Dynamic" Trade-off:
  • CUDA Graphs can eliminate launch overhead by replaying pre-captured workloads, but they require static tensor shapes and deterministic control flow, making them incompatible with the dynamic nature of LLMs (e.g., variable sequence lengths, stochastic sampling, and KV-cache updates).
  • Just-In-Time (JIT) Compilation handles dynamic shapes and branching well but still suffers from per-launch overhead and host-side coordination.

2. Methodology

The authors propose a Hybrid JIT–CUDA Graph Runtime that partitions the transformer inference pipeline into two distinct execution domains to bridge the gap between efficiency and flexibility.

• Partitioning Strategy:

  • Static Domain (CUDA Graph): Operations with fixed shapes and deterministic flow (e.g., linear projections, layer normalization, matrix multiplications, and fused attention) are captured into CUDA Graphs. These are executed via "replay," which bypasses the CPU and Python dispatch.
  • Dynamic Domain (JIT): Operations that depend on runtime conditions (e.g., token sampling, KV-cache updates, and positional embedding extensions) are handled via PyTorch’s JIT infrastructure, allowing for data-dependent branching.

• Key Architectural Features:

  • Asynchronous Graph Generation: To prevent the system from stalling when a new sequence length is encountered, the runtime uses a background capture mechanism. It captures new subgraphs on a separate CUDA stream (Scap) while simultaneously replaying existing graphs on a replay stream (Srep).
  • Rolling Graph Buffer: The system maintains a buffer of captured graphs for various sequence lengths. When a length is no longer needed, it is evicted to manage VRAM.
  • Decoupled Execution (IPC): The architecture uses a "Context Generator" (handling JIT/dynamic logic) and a "Graph Generator" (handling static replay), coordinated via Inter-Process Communication (IPC) to maximize GPU residency and minimize host intervention.

3. Key Contributions

• Hybrid Runtime Framework: A novel architecture that combines the low-overhead replay of CUDA Graphs with the flexibility of JIT compilation.
• Asynchronous Overlap: An implementation that overlaps the "capture" of new graphs with the "replay" of current ones, amortizing the cost of graph generation.
• Optimized Pipeline Partitioning: A systematic way to classify and separate transformer operations into static and dynamic categories.
• Reduced Latency/Variance: A method to achieve near-zero CPU overhead for the most compute-intensive parts of the model without breaking the autoregressive loop.

4. Experimental Results

The framework was evaluated using LLaMA-2 7B on an NVIDIA H100 GPU with batch size 1, comparing it against HuggingFace (PyTorch Eager) and TensorRT–LLM.

• Time-to-First-Token (TTFT): The hybrid runtime significantly outperformed baselines. For prompt lengths up to 500 tokens, it achieved up to a 66.0% reduction in TTFT compared to PyTorch Eager and consistently beat TensorRT–LLM.
• Tail Latency (P99): The hybrid approach achieved much lower P99 latency, reducing it by 20.2% compared to TensorRT–LLM and 43.5% compared to PyTorch Eager. This indicates much higher execution stability (lower jitter).
• Scaling: The TTFT showed a smoother, near-linear scaling trend as prompt length increased, whereas TensorRT–LLM exhibited higher variance at longer contexts.

5. Significance

This research provides a practical optimization strategy for latency-sensitive AI applications. By recognizing that LLM workloads are a mix of highly predictable math and highly unpredictable logic, the authors demonstrate that "one-size-fits-all" compilation (either purely static or purely dynamic) is suboptimal. Their hybrid approach offers a blueprint for building next-generation inference engines that provide the high throughput of specialized compilers with the flexibility required for real-world, interactive generative AI.