[Paper Review] PREMA: A Predictive Multi-task Scheduling Algorithm For Preemptible Neural Processing Units

Yujeong Choi, Minsoo Rhu|arXiv (Cornell University)|Sep 6, 2019

IoT and Edge/Fog Computing4 citations

TL;DR

This paper proposes PREMA, a predictive multi-task scheduling algorithm for preemptible neural processing units (NPUs) that improves latency, throughput, and SLA compliance by intelligently preemption low-priority tasks using lightweight checkpointing. It achieves up to 7.8× lower average normalized time-to-prediction (ANTT), 1.4× higher SLA satisfaction, and 19.6× better fairness compared to non-preemptive baselines, with negligible hardware overhead.

ABSTRACT

To amortize cost, cloud vendors providing DNN acceleration as a service to end-users employ consolidation and virtualization to share the underlying resources among multiple DNN service requests. This paper makes a case for a "preemptible" neural processing unit (NPU) and a "predictive" multi-task scheduler to meet the latency demands of high-priority inference while maintaining high throughput. We evaluate both the mechanisms that enable NPUs to be preemptible and the policies that utilize them to meet scheduling objectives. We show that preemptive NPU multi-tasking can achieve an average 7.8x, 1.4x, and 4.8x improvement in latency, throughput, and SLA satisfaction, respectively.

Motivation & Objective

• To address the growing need for low-latency, high-throughput DNN inference in cloud-based MLaaS by enabling preemptible NPUs.
• To design lightweight preemption mechanisms tailored for DNN workloads that minimize context-switching overhead.
• To develop a predictive scheduling policy (PREMA) that balances latency, fairness, throughput, and SLA adherence in multi-tasked NPU environments.
• To evaluate the impact of preemption mechanisms (checkpoint vs. kill) and scheduling policies on system-wide performance metrics.

Proposed method

• Designs three NPU preemption mechanisms—checkpointing, kill, and hybrid—evaluating their impact on context state size and preemption latency.
• Employs a lightweight regression model to predict DNN inference times for scheduling decisions, enabling predictive task prioritization.
• Introduces a preemptive, high-priority-first scheduling policy that dynamically reorders tasks based on predicted execution times.
• Uses a 64-bit per-task context tracking structure in on-chip SRAM, requiring only 0.01 mm² area for 16 concurrent tasks in 32nm technology.
• Implements a memory management strategy that migrates overflowed checkpointed states to CPU memory via DMA to handle memory oversubscription.
• Evaluates performance using real DNNs (GoogLeNet, ResNet, etc.) under varying batch sizes and preemption points, with a synthetic workload model.

Experimental results

Research questions

• RQ1 Can lightweight preemption mechanisms be designed for NPUs to support low-overhead task preemption in multi-tasked DNN inference?
• RQ2 How does preemption impact the size of checkpointed context state and overall preemption latency in DNN workloads?
• RQ3 To what extent can a predictive scheduler improve latency, throughput, and SLA compliance compared to non-preemptive or priority-only schedulers?
• RQ4 What is the relative performance of checkpointing versus killing preempted tasks in terms of ANTT, fairness, and throughput?
• RQ5 How robust is the PREMA scheduler across varying workloads, batch sizes, and preemption points?

Key findings

• PREMA achieves a 7.8× improvement in average normalized time-to-prediction (ANTT) over non-preemptive FCFS scheduling.
• The algorithm improves fairness by 19.6× and throughput by 1.4× compared to the baseline non-preemptive scheduler.
• SLA violation rates are reduced by 4.8×, demonstrating strong adherence to service-level objectives.
• Checkpointing outperforms killing by 87% in ANTT, 24% in STP, and 77% in fairness, with minimal performance overhead.
• The prediction model achieves 98% correlation with actual inference times, enabling accurate scheduling decisions without oracle access.
• Hardware and energy overheads are negligible—only 0.01 mm² SRAM for 16 tasks in 32nm, with energy efficiency improved proportionally by higher throughput.