Scaling On-Device GPU Inference for Large Generative Models

TL;DR

ML Drift tackles the challenge of running on-device inference for very large generative models by introducing tensor virtualization and coordinate translation to decouple logical tensors from physical GPU memory, enabling flexible memory layouts and kernel optimization across OpenCL, Metal, and WebGPU. It further enhances performance through device specialization, memory management, operator fusion, and stage-aware optimizations for LLMs, along with a GPU-optimized KV cache layout. The framework achieves order-of-magnitude improvements over existing open-source GPU inference engines and supports workloads with $10$ to $100×$ more parameters across mobile, desktop, and Apple Silicon GPUs. Its results demonstrate practical impact for private, low-latency AI on edge devices, expanding the feasible scale of on-device generative models and reducing reliance on server-side computation.

Abstract

Driven by the advancements in generative AI, large machine learning models have revolutionized domains such as image processing, audio synthesis, and speech recognition. While server-based deployments remain the locus of peak performance, the imperative for on-device inference, necessitated by privacy and efficiency considerations, persists. Recognizing GPUs as the on-device ML accelerator with the widest reach, we present ML Drift--an optimized framework that extends the capabilities of state-of-the-art GPU-accelerated inference engines. ML Drift enables on-device execution of generative AI workloads which contain 10 to 100x more parameters than existing on-device generative AI models. ML Drift addresses intricate engineering challenges associated with cross-GPU API development, and ensures broad compatibility across mobile and desktop/laptop platforms, thereby facilitating the deployment of significantly more complex models on resource-constrained devices. Our GPU-accelerated ML/AI inference engine achieves an order-of-magnitude performance improvement relative to existing open-source GPU inference engines.

Figures (8)

• Figure 1: Tensor virtualization enables flexible memory layouts. A logical tensor of size $(1,2,3,5)$ can be stored as a "physical" 3D texture in $DSHWBC_4$ layout (top right), a "physical" 2D texture in $HSWBDC_4$ layout (middle right), and a "physical" 1D image buffer in $DSHWBC_4$ layout (bottom right). The squares on the right side of the illustration represent 4-channel slices.
• Figure 2: Tensor virtualization decouples logical tensor indices from physical GPU indices. An $OHWI$ weights tensor of size $(5, 2, 1, 7)$, utilized in a convolution, can be stored as a 2D texture array. This array consists of four 2D textures, each measuring $(4, 2)$, which amounts to 8 vec4 elements per texture. Dotted-line squares indicate padding used for alignment.
• Figure 3: Memory savings for Stable Diffusion 1.4 using Greedy by Size policy for offset calculation pisarchyk2020. Light squares represent naive allocation (text encoder: 62 MB, UNet: 2075 MB, VAE decoder: 2274 MB); dark squares represent optimized memory footprint (text encoder: 2 MB, UNet: 65 MB, VAE decoder: 320 MB). Each square equals 1 MB, i.e., each row equals 100 MB.
• Figure 4: Automatic operator fusion examples for large model inference. Left: fusion of element-wise operators from two branches with a fully connected operator. Right: merging residual connection and element-wise operations with a manually optimized RMS normalization operator.
• Figure 5: Single-step inference latency (milliseconds) for Stable Diffusion 1.4, by model component (text encoder, VAE decoder, UNet), on Qualcomm and Arm GPUs.