HeteroMosaic: Heterogeneous Scheduling for Energy-Efficient Edge LLM Inference on SoCs

Overview

HeteroMosaic is a heterogeneity-first scheduling framework designed for energy-efficient Large Language Model (LLM) inference on edge system-on-chips (SoCs). This framework targets the underutilization of heterogeneous resources found in modern edge SoCs, specifically combinations of CPUs, integrated GPUs (iGPUs), and neural processing units (NPUs). It addresses the challenge of optimizing performance on unified-memory platforms, where both device placement and task-graph coordination are critical.

Research Context

Existing LLM runtimes on edge SoCs commonly make device-level decisions at a coarse granularity or optimize individual operators in isolation. This approach can lead to inefficient use of the diverse computational resources available. The integration of CPUs, iGPUs, and NPUs within single SoCs presents an opportunity for more efficient LLM inference, particularly on platforms featuring unified memory architectures. The performance on these platforms is influenced by decisions regarding where computations are executed (device placement) and how the overall task graph is managed across different accelerators (task-graph coordination).

Approach

HeteroMosaic employs a multi-step approach for scheduling edge LLM inference:

1. Heterogeneous Roofline Model: The framework first utilizes a heterogeneous roofline model to determine when the combined execution on iGPUs and NPUs offers performance advantages.
2. Micro-batch Decomposition: Inference tasks are then decomposed into micro-batches. This decomposition is designed to preserve data dependencies and expose opportunities for overlapping execution across different accelerators.
3. Trace-guided Co-optimization: HeteroMosaic applies a trace-guided co-optimization strategy. This strategy simultaneously optimizes both the scheduling of tasks and their allocation to specific devices. This co-optimization accounts for practical system effects, including memory contention, Dynamic Voltage and Frequency Scaling (DVFS), variations in device performance, and NPU runtime overheads.

The implementation of HeteroMosaic was carried out in PyTorch C++.

Findings

The evaluation of HeteroMosaic was conducted on three AMD Ryzen AI platforms, specifically chosen for their varying distributions of NPU, balanced, and iGPU-heavy designs:

• On a balanced AMD Ryzen AI platform, HeteroMosaic delivered a speedup of up to 1.73X compared to an iGPU-only baseline.
• On the same balanced platform, it achieved a speedup of up to 1.78X over an NPU-only baseline.
• Furthermore, HeteroMosaic showed a performance improvement of up to 2.05X when compared to established frameworks such as llama.cpp on the balanced platform.
• In addition to performance gains, the framework demonstrated a reduction in energy consumption by up to 45.3% on the balanced platform.
• When compared against previous heterogeneous edge AI solutions, HeteroMosaic achieved a performance improvement of up to 2.35X.

Why This Matters

The findings indicate that HeteroMosaic can significantly improve the speed and energy efficiency of LLM inference on edge devices equipped with heterogeneous computing resources. This has implications for deploying LLMs in environments where power consumption and processing speed are critical constraints.