💻 computer science

MemSpec: Memory-Aware Runtime for Adaptive Draft Scheduling in Speculative Decoding on Edge Devices

MemSpec is a memory-aware runtime that enhances speculative decoding on edge devices by decoupling draft selection from execution through proactive working-set management, thereby improving throughput by 40.7% over existing adaptive methods while minimizing model switching overhead.

Original authors: Eunjeong Kim, Yeong Jun Jeon, Myeonggyun Han

Published 2026-08-12

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: Eunjeong Kim, Yeong Jun Jeon, Myeonggyun Han

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you are trying to solve a giant, complex puzzle, but you have a very strict rule: you can only look at one piece at a time. This is how many powerful computer brains, called Large Language Models (LLMs), currently work. They build sentences word by word, checking each new word against the last one before moving forward. It's accurate, but it's also incredibly slow, like trying to cross a river by hopping from stone to stone one by one. To speed this up, scientists invented a trick called "speculative decoding." Think of it like having a fast, energetic junior assistant who guesses the next few pieces of the puzzle for you. The big, slow brain then quickly checks these guesses. If the junior assistant is right, the big brain skips the hard work and moves forward faster. If the assistant is wrong, the big brain fixes the mistake and tries again. This works great, but only if the assistant is good at guessing.

The problem is that no single assistant is perfect at everything. One might be a genius at writing code but terrible at legal documents, while another is a whiz at medical advice but clumsy with poetry. On powerful computers in giant data centers, you can keep a whole team of these assistants ready to go, switching between them instantly. But on smaller devices, like the ones in your phone or a small robot, there isn't enough memory (brain space) to keep all the assistants loaded at once. If you try to swap assistants, you have to stop everything, go to the storage closet to fetch the new one, and wait. That waiting time is so long it cancels out all the speed you gained from guessing. This paper, titled "MemSpec," tackles the tricky question of how to keep a small, fast team of assistants ready on a memory-hungry edge device without getting stuck waiting in the closet.

The researchers behind MemSpec discovered that the biggest bottleneck isn't just picking the right assistant; it's making sure the right assistant is actually sitting in the chair when you need them. They found that existing methods, which try to figure out the best assistant by constantly testing different ones (a process called "exploration"), often fail on small devices. These methods might pick a great assistant, but if that assistant is still in the storage closet, the system has to stop and wait for them to arrive, wasting precious time. In fact, the time it takes to load a new assistant from storage can be 2.7 times longer than a single step of the puzzle-solving process. This means that even if you pick a better assistant, the time spent waiting for them to arrive can make the whole process slower than just using a mediocre one that was already sitting there.

To solve this, the team built a new system called MemSpec, which acts like a smart, proactive manager. Instead of waiting to see which assistant is best and then frantically running to the closet to get them, MemSpec uses a lightweight "prediction engine" to guess which assistants will be needed next based on what the user is currently talking about. If the conversation shifts from coding to math, the manager predicts this change and quietly fetches the math expert in the background before the current assistant runs out of ideas. Crucially, the system never stops the current work to wait for a new assistant. It always keeps the best currently available assistant working, while the new, better one is being loaded in the background. This way, the device is always running at full speed, and the "switching" happens without any pauses.

The team tested MemSpec on a small, powerful device called a Jetson Orin Nano, which is typical for edge computing. They compared it against the best existing methods that try to adapt by exploring different assistants. The results showed that MemSpec was a clear winner. It improved the speed of generating text by an average of 40.7% compared to the best adaptive methods currently available. It even came very close to a theoretical "perfect" scenario where the system knows exactly which assistant to use at every moment without any memory limits. The study also showed that simply predicting the right assistant wasn't enough; the system only worked well because it combined that prediction with a smart memory manager that kept the right assistants ready to go. Without this memory management, the predictions were useless because the right assistants weren't there when needed.

The researchers also looked at how different settings affected the system. They found that the system works best when it checks for new assistants every 4 steps of the puzzle, a sweet spot that balances the need for fresh predictions with the time needed to load new models. They also discovered that the longer the story or code being generated, the more MemSpec helped, because longer tasks have more changes in topic that require different types of assistants. Interestingly, adding more memory to the device didn't help MemSpec as much as it helped the older methods, because MemSpec was already so good at keeping the right two assistants ready that having a third or fourth didn't add much value. This suggests that the key to speed isn't just having more memory, but being smarter about how you use the memory you have.

In short, MemSpec proves that on small devices, the secret to speed isn't just finding the smartest assistant, but making sure the smartest available assistant is ready to work the moment you need them. By predicting future needs and managing the "closet" of assistants carefully, the system avoids the costly delays that have held back fast AI on edge devices. The paper suggests that this approach of separating the decision of "who is best" from the reality of "who is available" is the key to unlocking faster, more efficient AI on the devices we carry in our pockets every day.

Technical Summary: MemSpec

Problem Statement

The paper addresses the challenge of accelerating Large Language Model (LLM) inference on memory-constrained edge devices (e.g., mobile devices, embedded systems) using speculative decoding. While speculative decoding improves throughput by using a lightweight draft model to propose tokens for verification by a larger target model, its effectiveness on edge devices is severely limited by memory constraints.

Existing adaptive methods, which dynamically switch between multiple specialized draft models to maximize token acceptance rates, fail in this setting due to a fundamental mismatch between draft selection and draft availability. On edge devices with limited fast memory (e.g., 8GB on a Jetson Orin Nano), only a small number of draft models can reside simultaneously. Switching to a non-resident draft requires loading it from slower storage (e.g., NVMe), incurring latency that often exceeds several decoding iterations. Consequently, exploration-based adaptive methods (such as Multi-Armed Bandit approaches) frequently incur high switching overheads, causing the system to stall or execute suboptimal drafts while waiting for better ones to load. This negates the potential throughput gains from improved draft selection.

Methodology: MemSpec

To resolve this, the authors propose MemSpec, a prediction-guided, memory-aware runtime for adaptive draft scheduling. MemSpec decouples the decision of which draft to use from the execution of that draft, transforming the problem into a memory-aware scheduling task.

Core Architecture

MemSpec consists of three primary components:

Prediction Engine: A lightweight, offline-trained BERT-based model that predicts the effectiveness of candidate drafts based on the input prompt and recent generation context (the last $T$ tokens). Unlike exploration-based methods, it does not rely on online probing but instead provides a ranked list of promising drafts.
Draft Model Cache Manager: This component manages a small resident working set of draft models (capacity $K$ ) within fast memory. It uses the ranked list from the Prediction Engine to proactively prefetch high-priority non-resident drafts and evict low-priority resident ones. Crucially, it protects the currently active draft from eviction to ensure uninterrupted decoding.
Runtime Controller: Orchestrates the decoding loop. It operates on a scheduling interval (every $N$ iterations). At each interval, it selects the best currently resident draft for execution while asynchronously prefetching better candidates in the background. If the predicted optimal draft is not yet resident, the system continues with the best available resident draft, avoiding blocking.

Key Design Principles

Non-Blocking Execution: The runtime never stalls for model loading. Execution always proceeds with the best available resident draft.
Proactive Residency Management: Instead of reacting to past performance (as in bandit methods), MemSpec uses prediction to align the resident working set with future demand, ensuring high-utility drafts are loaded before they are needed.
Interval-Based Scheduling: Prediction and cache updates occur at fixed intervals to amortize overhead and allow sufficient time for asynchronous prefetching to overlap with decoding.

Key Contributions

Identification of a Fundamental Limitation: The authors demonstrate that on memory-constrained edge devices, the cost of switching to non-resident drafts often outweighs the benefits of improved selection. They show that exploration-based methods fail to translate better draft selection into higher throughput due to excessive loading latency.
MemSpec Runtime: A novel system that decouples draft selection from execution. It replaces online exploration with lightweight prediction and enables non-blocking decoding by strictly adhering to the set of currently resident models.
Memory-Aware Scheduling Framework: A coordinated prefetching and eviction strategy that maintains a small working set of drafts aligned with predicted future demand, enabling adaptive decoding without the penalty of frequent model loading.
Implementation and Evaluation: The system is implemented on a Jetson Orin Nano and evaluated across diverse workloads (code, math, law, medical, and general instruction following).

Experimental Results

The evaluation was conducted on an NVIDIA Jetson Orin Nano using LLaMA-2 7B and Qwen2.5 7B target models with four domain-specialized draft models (code, math, law, medical) and one general-purpose draft.

Throughput Improvement: MemSpec improves steady-state generation throughput by 40.7% on average over state-of-the-art bandit-based adaptive methods (MAB-Async) and by 58.8% over static baselines.
Oracle Proximity: MemSpec achieves 95–97% of the performance of an "Oracle-Dynamic" upper bound (which knows the best draft at every step but is constrained by the same memory limits), demonstrating that prediction-guided scheduling can approximate ideal behavior without exhaustive offline enumeration.
Execution Breakdown: Analysis reveals that MAB-Async spends approximately 49.4% of its time in "fallback execution" (using suboptimal drafts while waiting for better ones to load). MemSpec reduces this fallback execution to below 5.5%, significantly increasing the time spent on "desired execution."
Component Analysis: The study shows that prediction alone (without proactive prefetching) improves throughput by 21.6% over static baselines but falls short of the full MemSpec design by 30.5%. This confirms that aligning selection with residency is critical; accurate prediction alone is insufficient if the selected drafts are not available.
Sensitivity: MemSpec remains robust across different scheduling intervals ( $N=4$ to $8$) and shows increasing benefits for longer output sequences where intra-sequence heterogeneity is higher.

Significance and Claims

The paper claims that the primary bottleneck for adaptive speculative decoding on edge devices is not the quality of draft selection, but the timely availability of effective drafts under tight memory budgets.

The authors assert that MemSpec successfully bridges the gap between the theoretical benefits of dynamic adaptation and the practical constraints of edge hardware. By treating draft selection and memory management as a coupled scheduling problem, MemSpec demonstrates that it is possible to achieve near-oracle performance without the prohibitive overhead of frequent model switching. The work suggests that future adaptive inference systems on edge devices must move beyond simple selection algorithms to include proactive residency management to realize their full potential.