FlashMem: Supporting Modern DNN Workloads on Mobile with GPU Memory Hierarchy Optimizations

Zhihao Shu, Md. Musfiqur Rahman Sanim, Hangyu Zheng, Kunxiong Zhu, Miao Yin, Gagan Agrawal and 1 other

Running a modern neural network on a phone's GPU normally means loading the entire model's weights into memory before it can run, but today's big models (and workloads that run several models in a row) need more memory than a phone has, so they either fail or crawl; this paper (University of Georgia and University of Texas at Arlington — Zhihao Shu, Wei Niu, Gagan Agrawal, and colleagues) presents FlashMem, which streams weights in on demand instead of preloading them, plans exactly when each weight should arrive so peak memory stays small, and hides the loading behind computation using the mobile GPU's texture memory.

The argument in one glance

Problem

Modern models are too big to load into a phone GPU all at once

People increasingly want to run neural networks directly on their phones — for privacy, for real-time response, and to avoid depending on the cloud. But modern models are large. A speech model like Whisper or an image generator like Stable Diffusion's UNet has hundreds of millions of parameters, and its weights alone can need four to five gigabytes of GPU memory, which is more than the graphics memory many phones make available. The way today's mobile frameworks run a model makes this worse: they use a "preloading" strategy, loading the entire model — all of its weights — into GPU memory before running a single layer. That was fine for small models, but for a big one it pushes peak memory past what the device has, and for the increasingly common case where an app runs several distinct models in succession (say, transcribe audio, then summarize, then generate an image), keeping them all resident is simply impossible. On top of the memory problem, mobile GPUs have an unusual memory hierarchy, and current frameworks pay repeatedly to reshape weights into the layout the GPU wants, launching many kernels and doing redundant data transformations. That initialization and transformation cost can come to dominate the actual inference time. So the very models people most want to run on-device are the ones that either do not fit or run painfully slowly.
What people did

Preload the whole model, and optimize its memory layout

The standard approach is weight preloading: load every parameter of the model into GPU memory up front, so that when execution starts each layer's weights are already resident and inference runs without waiting on storage. To make the mobile GPU faster, recent work also optimizes how those weights are laid out — for instance, SmartMem reorganizes tensors to fit the GPU's specialized "2.5D" texture memory and eliminates some of the costly layout transformations (like Reshape and Transpose) that frameworks would otherwise repeat.
Why it fell short

Preloading blows the memory budget and cannot overlap loading with compute

Preloading is fundamentally at odds with big and multi-model workloads. Loading every weight before execution means peak memory equals the whole model — four to five gigabytes for a single large model on a phone — so a large model barely fits and running several models at once is out of the question. Layout optimizations like SmartMem help with the transformation cost but still assume the whole model is loaded, so they do not solve the memory-budget problem: they are memory-bound in exactly the cases that matter. And because loading happens entirely before execution, none of that loading time is hidden behind useful computation — the GPU sits idle while weights are transferred and transformed, and that overhead can dominate latency. What was missing was a way to bring weights in only as they are needed, keep peak memory small, and overlap the loading with the model's own computation so the transfer is nearly free.
This paper's idea

Stream weights on demand, on a plan that hides loading behind compute

FlashMem is a memory streaming framework for mobile DNN inference. Instead of preloading the whole model, it statically decides a loading schedule and then dynamically streams weights in on demand during inference, so peak memory stays within the device's limits and large or multi-model workloads become possible. It rests on four ideas. First, Optimized Overlap Plan Generation (OPG): FlashMem formalizes the question of when and where each weight should be loaded as an optimization problem, and solves it with LC-OPG, a load-aware solver that minimizes peak memory while keeping loads overlapped with computation. Second, load-capacity profiling and adaptive fusion: it classifies each operator by whether it is memory-bound or compute-bound to learn its "load capacity" — how much weight loading it can absorb while still hiding the transfer behind its own compute — and fuses operators adaptively to balance memory against speed. Third, hierarchical GPU memory optimization: it restructures weight layouts into the mobile GPU's 2.5D texture memory in small tiles, so weights are cached well and the costly repeated layout transformations are largely eliminated. Fourth, pipeline-aware kernel execution: a branch-free, pipelined kernel design interleaves computation with weight loading so all threads follow the same path and the GPU stays busy. On eleven models spanning six categories, FlashMem reports 2.0 times to 8.4 times less memory (up to 10.1 times on some models) and 1.7 times to 75.0 times speedups over state-of-the-art mobile frameworks, enabling large-scale and multi-model execution on resource-constrained phones.

Why this problem is worth solving

On-device inference is where privacy, responsiveness, and offline use come from, but it is exactly where memory is scarcest, so fitting big models onto phones has broad value.

Running models directly on the phone, rather than in the cloud, is what keeps a user’s data on their device, gives instant responses without a network round trip, and works with no connection at all. Those benefits are why on-device AI keeps growing. But the phone is also the most memory-constrained place to run a model: its GPU has a fraction of a desktop’s memory, shared with everything else the device is doing. As the models people want — speech recognition, image generation, on-device assistants — get larger and start being chained together, the gap between what users want to run and what the device can hold widens. Closing that gap means those capabilities can ship to ordinary phones instead of being reserved for the cloud or high-end hardware. Because the benefit applies to every large or multi-model workload on a memory-limited device, making them fit is a broad, practical win for on-device AI.

The mobile GPU’s unusual memory hierarchy is both the source of the overhead and the lever for fixing it.

A mobile GPU does not have dedicated video memory like a desktop card; it shares memory with the CPU and relies on a layered organization — unified memory, then a specialized texture memory with its own cache, then the streaming multiprocessors that compute. Getting a weight from storage to where it can be used means moving it through these layers and reshaping it into the layout the texture memory wants, a “2.5D” tiled form. Current frameworks do this reshaping repeatedly and at the wrong granularity, launching many kernels and redoing transformations, which is a large part of why initialization and data-transformation time can dominate inference. But this same hierarchy is the opportunity: mobile GPUs can overlap memory transfers with computation using independent command queues, and the texture memory can be exploited to cache weights efficiently. So the design that hurts a naive framework is exactly what a careful one can turn into an advantage — which is why FlashMem builds its streaming, its layout, and its pipelining specifically around this hierarchy rather than treating memory as a flat pool.

What people did before, and where it broke

Preload the entire model before running. Load all weights into GPU memory up front so inference never waits on storage. Where it broke. Peak memory equals the whole model (four to five gigabytes for one large model on a phone), so large models barely fit and running several models in succession is impossible.
Optimize the weight layout for texture memory (SmartMem). Reorganize tensors into the 2.5D texture layout and eliminate repeated Reshape and Transpose transformations. Where it broke. It still assumes the whole model is preloaded, so it stays memory-bound and does not solve the memory-budget problem for big or multi-model workloads.
Swap models in and out for multi-model workloads. When another model is needed, evict the current one and load the next. Where it broke. Constant swapping re-pays the full load-and-transform cost each time, and that GPU initialization overhead comes to dominate inference latency.
Load first, then compute. Do all weight transfer and transformation before execution begins. Where it broke. None of the loading time is hidden behind computation, so the GPU sits idle during transfer and the overhead is fully exposed.
The gap. No framework brought weights in only as needed, kept peak memory small with a planned schedule, and overlapped the loading with the model’s own computation on the mobile GPU.

How FlashMem works

Streaming weights on demand and hiding the loading behind computation

↓ →
new
↓ →
new
↓ →
new
↓ →
new
↓ →
new

The idea

Bring the data in only as you need it, instead of holding it all at once. FlashMem’s founding move is to replace preloading with streaming: keep only the weights the current computation needs resident, and pull the rest in on demand, so a model far larger than memory can still run. Running a model that does not fit in memory by migrating its tensors in and out on demand, rather than requiring the whole thing be resident, is the same idea the work on tensor migration and prefetching in unified memory (DeepUM) covered elsewhere on this site develops, letting capacity be the thing you stream rather than the thing you must have. FlashMem applies it to the mobile GPU, streaming weights through the memory hierarchy so peak memory stays small.

Hide the cost of moving data behind the work that is already running. Streaming risks paying a transfer stall for every weight, so FlashMem overlaps loading with computation — the GPU computes one thing while the next weights arrive on an independent queue, making the transfer nearly free when there is compute to hide it behind. Overlapping data movement with the computation that depends on it, so the movement is masked rather than paid for serially, is the same principle the work on overlapping communication with dependent computation covered elsewhere on this site is built on, turning a serial transfer-then-compute into a pipeline. FlashMem does this per operator, using each layer’s spare compute to absorb the loading of the next.

Turn a messy scheduling decision into a problem you can actually solve. Deciding when and where every weight should be loaded, across a whole model, to minimize peak memory without stalling, is far too tangled to guess at, so FlashMem writes it down as a formal optimization — the OPG problem — and solves it with a purpose-built solver. Casting a hard resource-and-timing decision as a constrained optimization and solving it, rather than relying on hand-tuned heuristics, is the same approach the work on constraint-aware resource scheduling for inference (ExeGPT) covered elsewhere on this site takes, letting a solver find the schedule a human could not. FlashMem’s LC-OPG solver finds a loading plan that minimizes peak memory while keeping loads overlapped with compute.

Following a big model through FlashMem

Running a model that would not fit if it were preloaded

A trace, drawn from the paper, of how FlashMem executes a large model within a tight memory budget.

Step 1 of 5

A large model would exceed memory if preloaded

A model like Stable Diffusion's UNet needs several gigabytes of weights — more than the phone's GPU memory allows if loaded all at once, and impossible to keep resident alongside the other models an app runs in succession.

Preloading makes peak memory equal to the whole model.

Why it works

Streaming keeps peak memory small. Preloading makes peak memory equal to the whole model → but streaming keeps only the weights the current computation needs resident → so large and multi-model workloads fit within a phone’s tight memory budget.
Overlap makes the loading nearly free. Loading on demand risks a transfer stall per weight → but overlapping loads with computation on independent command queues → so the transfer hides behind compute and costs little where there is work to mask it.
Load-capacity profiling avoids stalls. Not every operator can absorb the same amount of loading → but classifying each as memory- or compute-bound and fusing adaptively → so compute-heavy layers stream aggressively while memory-bound ones do not stall.
The 2.5D layout cuts transformation overhead. Getting weights GPU-ready normally means costly, repeated layout transformations → but restructuring them into cached 2.5D texture-memory tiles → so the transformation cost and kernel-launch overhead that dominated latency drop sharply.
Formalizing the schedule beats guessing. Deciding when and where to load every weight by hand is intractable → but writing it as the OPG problem and solving it with LC-OPG → so peak memory is minimized while every load stays overlapped with computation.

What it costs (trade-offs and scope)

Streaming only pays off when there is compute to hide it behind. If a layer cannot mask a transfer behind its own computation, on-demand loading adds latency; FlashMem manages this by profiling load capacity and keeping added per-kernel overhead within a bound, but the benefit depends on the workload having compute to overlap.
The loading plan is computed statically ahead of time. LC-OPG solves an optimization offline to produce the schedule, so there is an up-front planning cost, and the plan assumes the profiled model and device characteristics rather than adapting fully at runtime.
The gains are largest for workloads that do not fit preloaded. Very large models and multi-model workloads benefit most; a small model that already fits comfortably gains little and could even pay a small streaming overhead, so FlashMem targets the demanding cases.
It is built for the mobile GPU memory hierarchy. The texture-memory and 2.5D-layout optimizations assume a unified-memory-plus-texture-memory hierarchy like ARM Mali or Qualcomm Adreno; hardware without that structure would not see the same layout benefits.
It is a whole framework, not a drop-in flag. FlashMem’s custom kernels, branch-free pipelined execution, and layout restructuring are framework-level machinery, so its gains come with adopting the framework rather than toggling an option in an existing runtime.

The evaluation

Running modern DNN workloads on a mobile GPU

Property	Weight preloading (MNN, TVM)	Layout optimization (SmartMem)	FlashMem
Peak memory
Very large or multi-model workloads
Data-transformation overhead
Overlaps loading with compute
Reported speedup versus state of the art

A comparison of ways to run modern DNN workloads on a mobile GPU. By streaming weights on demand instead of preloading, planning the loading schedule with the LC-OPG solver, laying weights out in 2.5D texture memory, and pipelining loading with computation, FlashMem reports 2.0 times to 8.4 times less memory (up to 10.1 times on some models) and 1.7 times to 75.0 times speedups over state-of-the-art mobile frameworks (an average of 8.6 times over SmartMem) on eleven models, enabling large-scale and multi-model inference on phones. Exact per-model figures should be read from the paper.Underlined values explain themselves when clicked.

Where this stops working

What to keep in mind

This explainer is based on the paper’s full text (the authors’ arXiv copy, 2602.15379, of the ASPLOS 2026 paper). The system described is the paper’s: FlashMem is a memory streaming framework for mobile DNN inference that, instead of preloading all weights, statically determines a loading schedule and dynamically streams weights on demand, leveraging the mobile GPU’s 2.5D texture memory. It has four components: Optimized Overlap Plan Generation (OPG), which formalizes when and where each weight should load and is solved by LC-OPG, a load-aware solver that minimizes peak memory while keeping loads overlapped with computation; load-capacity profiling and adaptive fusion, which classify operators by memory and compute intensity to decide how much loading each can hide; hierarchical GPU memory optimization, which restructures weight layouts into 2.5D texture-memory tiles to reduce data-transformation overhead; and pipeline-aware kernel execution, a branch-free pipelined design interleaving compute with weight loading. The paper evaluates FlashMem on eleven models across six categories (including Whisper, GPT-Neo, and Stable Diffusion UNet) on mobile devices such as the OnePlus 12, reporting 2.0 times to 8.4 times memory reduction (up to 10.1 times on some models like DeepViT) and 1.7 times to 75.0 times speedups over state-of-the-art frameworks, with an average 8.6 times speedup over SmartMem (up to 15.8 times for GPT-Neo 1.3B and 9.3 times for SD-UNet). Exact mechanisms, the OPG formulation, the LC-OPG solver, and all quantitative results should be read from the paper.

Background

Why on-device AI keeps rediscovering the art of not holding everything at once

It helps to see FlashMem within a broad pattern: as models outgrow the devices meant to run them, systems keep converging on the same answer — do not keep everything resident, stream it in as needed, and hide the cost of moving it. This is the central tension of on-device inference, where the model wants more memory than the hardware has, and where the loading of parameters, not just the math, becomes a first-class performance concern. The response is always some blend of the same ingredients: bring data in on demand, decide carefully what to load when, and overlap that loading with computation so the transfer disappears into work the device was doing anyway. That FlashMem needs a formal scheduler, a memory-hierarchy-aware layout, and a pipelined execution engine to pull this off on a phone shows how much machinery good streaming requires. This concern for squeezing demanding models onto memory-limited edge devices by pipelining how their parameters are loaded is the same terrain the work on turbocharging NLP inference at the edge via elastic pipelining (STI) covered elsewhere on this site occupies, treating parameter loading as something to schedule and overlap rather than a fixed up-front cost. FlashMem’s contribution is to show that, with a planned schedule, a texture-memory-aware layout, and pipelined kernels, even models several times larger than a phone’s GPU memory can run, and several can run in succession, on ordinary hardware.

If you remember only three things

1. Modern DNN workloads on phones — very large models like Whisper or Stable Diffusion UNet, and workloads that run several distinct models in succession — do not fit the standard mobile approach, which preloads every weight into GPU memory before running. Preloading makes peak memory equal to the whole model (four to five gigabytes for one large model), so large models barely fit and multi-model workloads are impossible, and the mobile GPU's memory hierarchy adds repeated layout transformations and initialization overhead that come to dominate inference latency.
2. FlashMem is a memory streaming framework that brings weights in on demand instead of preloading, built on four ideas. It formalizes when and where each weight should load as the Optimized Overlap Plan Generation (OPG) problem and solves it with LC-OPG, a load-aware solver that minimizes peak memory while keeping loads overlapped with computation. It profiles each operator's load capacity (memory- versus compute-bound) and fuses adaptively so streaming hides behind compute. It restructures weights into the GPU's 2.5D texture-memory tiles to cut data-transformation overhead. And it uses branch-free, pipelined kernels that interleave compute with loading so the GPU never idles.
3. On eleven models across six categories, evaluated on mobile devices such as the OnePlus 12, the paper reports FlashMem achieves 2.0 times to 8.4 times memory reduction (up to 10.1 times on some models) and 1.7 times to 75.0 times speedups over state-of-the-art mobile frameworks, including an average 8.6 times speedup over SmartMem, enabling large-scale and multi-model execution on resource-constrained phones. The trade-offs: streaming pays off only when there is compute to hide the transfer behind, the loading plan is computed statically ahead of time, the biggest gains are for workloads that do not fit preloaded, and the layout optimizations assume a mobile GPU with a texture-memory hierarchy.

Check yourself

Check your understanding

1. Why does the standard "weight preloading" strategy fail for modern DNN workloads on phones?
2. What is the core idea of FlashMem, and how does it keep peak memory small?
3. How does FlashMem avoid paying a transfer stall every time it streams a weight on demand?
4. What does the paper report about FlashMem's results across its evaluation?

Sources

The paper is “FlashMem: Supporting Modern DNN Workloads on Mobile with GPU Memory Hierarchy Optimizations” by Zhihao Shu, Md Musfiqur Rahman Sanim, Hangyu Zheng, Kunxiong Zhu, Miao Yin, Gagan Agrawal, and Wei Niu (University of Georgia and University of Texas at Arlington), published at ASPLOS 2026, DOI 10.1145/3779212.3790164. It is available on arXiv as 2602.15379.

This explainer is written from the paper’s full text (the authors’ arXiv copy). FlashMem is a memory streaming framework for mobile DNN inference that replaces weight preloading with on-demand streaming, statically determining a loading schedule and dynamically streaming weights during execution while leveraging the mobile GPU’s 2.5D texture memory. It comprises: Optimized Overlap Plan Generation (OPG), which formalizes when and where each weight should load and is solved by the load-aware LC-OPG solver to minimize peak memory while overlapping loads with computation; load-capacity profiling and adaptive fusion, classifying operators by memory and compute intensity to decide how much loading each can hide; hierarchical GPU memory optimization, restructuring weight layouts into 2.5D texture-memory tiles to cut data-transformation overhead; and pipeline-aware kernel execution, a branch-free pipelined design interleaving compute with loading. Evaluated on eleven models across six categories on devices such as the OnePlus 12, the paper reports 2.0 times to 8.4 times memory reduction (up to 10.1 times on some models) and 1.7 times to 75.0 times speedups over state-of-the-art mobile frameworks, with an average 8.6 times speedup over SmartMem. Numbers are the paper’s own; exact mechanisms and full results should be confirmed in the paper.