Accelerating Mobile Language Model via Speculative Decoding and NPU-Coordinated Execution

TL;DR

This work tackles the bottlenecks of on-device retrieval-augmented generation by introducing sd.npu, a system that co-designs progressive graph scheduling with NPU-optimized speculative decoding. Key innovations include overlapping graph loading with computation to hide switch latency, in-context distribution calibration to reduce draft rejection, and a draft-extension mechanism that reuses useful rejected tokens to boost NPU throughput. Empirical results on multiple smartphones, datasets, and models show significant end-to-end speedups (1.06×–3.81×) and energy savings (1.07×–4.71×), with modest memory overhead. The approach advances practical, private, responsive LLM-based mobile applications by maximizing hardware utilization without sacrificing accuracy.

Abstract

Performing Retrieval-Augmented Generation (RAG) directly on mobile devices is promising for data privacy and responsiveness but is hindered by the architectural constraints of mobile NPUs. Specifically, current hardware struggles with the variable workloads intrinsic to RAG: the transition between processing extensive contexts and generating tokens incurs significant overhead due to static graph constraints, while the memory-bound generation phase leaves computational resources underutilized. In this work, we propose a holistic acceleration framework sd.npu, designed to maximize NPU efficiency for on-device RAG ecosystem. To address the latency caused by NPU graph switching during phase transitions, we introduce a pipelined execution strategy. This approach masks the overhead of model reconfiguration by parallelizing the loading of decoding graphs with the computation of partitioned context chunks (chunked prefill), thereby ensuring continuous execution flow. Furthermore, to mitigate low hardware utilization during the decoding phase, we develop an NPU-centric speculative decoding mechanism. By calibrating generation distributions and extending draft sequences, our method effectively converts idle NPU cycles into valid token throughput. Experiments on commercial smartphones show that our framework significantly outperforms existing baselines, delivering 1.06$\times$--3.81$\times$ speedups and 1.07$\times$--4.71$\times$ energy savings across various RAG tasks.

Figures (14)

• Figure 1: RAG on NPU and NPU's architecture.
• Figure 2: Comparison on consumed energy per request across different execution backends work:mllmwork:heterollm.
• Figure 3: Trade-off between prefill and decoding with fixed-size compute graphs, highlighting the importance to deploy specialized graphs for each stage.
• Figure 4: Latency breakdown of RAG decoding phase.
• Figure 5: Comparison of verification speed with different draft lengths, and the length distribution of drafts generated by existing SD method work:pld (histogram).