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

Abstract: The increasing size and complexity of modern deep neural networks (DNNs) pose significant challenges for on-device inference on mobile GPUs, with limited memory and computational resources. Existing DNN acceleration frameworks primarily deploy a weight preloading strategy, where all model parameters are loaded into memory before execution on mobile GPUs. We posit that this approach is not adequate for modern DNN workloads that comprise very large model(s) and possibly execution of several distinct models in succession. In this work, we introduce FlashMem, a memory streaming framework designed to efficiently execute large-scale modern DNNs and multi-DNN workloads while minimizing memory consumption and reducing inference latency. Instead of fully preloading weights, FlashMem statically determines model loading schedules and dynamically streams them on demand, leveraging 2.5D texture memory to minimize data transformations and improve execution efficiency. Experimental results on 11 models demonstrate that FlashMem achieves 2.0x to 8.4x memory reduction and 1.7x to 75.0x speedup compared to existing frameworks, enabling efficient execution of large-scale models and multi-DNN support on resource-constrained mobile GPUs.

• The paper demonstrates that replacing monolithic weight preloading with fine-grained, operator-level memory streaming drastically reduces latency and memory use.
• It introduces an Overlap Plan Generation (OPG) model using CP-SAT optimization to schedule weight transfers with precise load capacity constraints.
• Empirical results show 2.0×–8.4× memory reduction and up to 75.0× speedup over state-of-the-art frameworks, enabling efficient multi-DNN execution.

FlashMem: Hierarchical GPU Memory Management for DNN Inference on Mobile Devices

Motivation and Problem Statement

The increasing parameter count and complexity of state-of-the-art DNNs have pushed the limits of on-device inference, particularly on mobile GPUs. These devices exhibit constrained memory footprints and unique hierarchical GPU memory architectures, with a unified memory (UM) subsystem and specialized 2.5D texture memory (TM). Existing mobile DNN frameworks (e.g., MNN, TVM, TFLite, SmartMem) predominantly adopt a weight preloading paradigm—loading all model weights into memory before execution, optimizing inference latency at the cost of exceptionally high peak and average memory consumption. This approach not only complicates the deployment of large or multiple DNNs on the same device but also incurs substantial weight transformation overheads due to suboptimal handling of texture memory layouts.

FlashMem targets efficient execution of both single large DNNs and workloads involving multiple DNNs (multi-model tasks), by replacing monolithic preload strategies with fine-grained, static memory streaming schedules and runtime overlap of weight loading and computation. The core insight is that controlling when and how much to load and transform, at the operator (and sub-operator) level, enables order-of-magnitude improvements in both latency and memory usage.

Mobile GPU Memory Hierarchy and Multi-DNN Requirements

Mobile GPUs, unlike desktop-class architectures, employ a heterogeneous memory layout incorporating UM and TM. Weight transport involves a three-stage pipeline: disk-to-UM, UM-to-TM (via kernel-based transformation), and computation on streaming multiprocessors (SMs). TM exhibits 2.5D spatial locality; hence, efficient DNN execution must both optimize the mapping of weights to TM and minimize costly transformation and redundant movement. Existing frameworks fail to exploit these physical properties, instead triggering excessive in-memory duplications and kernel launches (Figure 1).

Figure 1: (a) Mobile GPU memory hierarchy and multi-step weight transfers; (b/c) FIFO and preemptive multi-model execution paradigms.

Moreover, emerging mobile applications increasingly fuse multiple DNNs in sequential or dynamically triggered pipelines, necessitating support for multi-DNN FIFO execution and further compounding the cost of monolithic preloading. Naive approaches result in repeated load/transform costs and memory contention, particularly problematic for large models that already saturate device memory when isolated.

Formulation: Overlap Plan Generation with Load Capacity Constraints

FlashMem introduces the Overlap Plan Generation (OPG) model, cast as a CP-SAT combinatorial optimization problem with the following characteristics:

• The DNN graph is decomposed into operators, each with associated weight tensors. Decision variables designate when and where each weight is loaded/transformed, subject to constraints on UM and TM peaks.
• The objective penalizes both the number of persistently preloaded weights and their "loading distance" (early loading increases residency and hence memory).
• Weights are partitioned into uniform chunks; allocations map chunk transfers to execution layers, honoring per-layer load capacities.
• Constraints enforce completeness of allocation, per-operator load capacity limits (informed by profile-guided regression models), and global and local memory ceilings.

The CP-SAT solver is augmented by structured fallback mechanisms—incrementally relaxing infeasible constraints and employing greedy heuristics when necessary (Figure 2).

Figure 2: XGBoost-guided profiling for latency prediction and load capacity estimation across kernels.

This approach differentiates between operator classes (elemental, reusable, hierarchical), setting distinct load tolerance thresholds and informing solver decisions. Elemental and reusable ops can effectively overlap transfer and compute; hierarchical ops (e.g., Softmax/LayerNorm) cannot and are deprioritized. Fusion-aware adaptive splitting intervenes to break over-fused kernels when such fusions constrain available overlap points beyond threshold.

FlashMem System Architecture

The FlashMem runtime operates as follows:

1. Model Analysis and Profiling: The model's operator-DAG is parsed and op types classified.
2. Load Capacity Profiling: Lightweight empirical measurements, regressed via XGBoost, estimate each operator's capacity for concurrent data movement.
3. OPG Solver: A CP-SAT instantiation embeds all per-layer and global constraints, solved offline to yield a deterministic, reusable overlap schedule.
4. Operator Fusion Management: The pipeline dynamically splits fused ops as solver constraints require, maximizing the number of feasible overlap windows.
5. Template-based Kernel Rewriting: GPU kernels are automatically rewritten (via Jinja templates) to integrate branch-free, pipelined loading of weight tiles, synchronized with TM's 2.5D spatial layout (Figure 3).

   Figure 3: Left—naive matrix multiplication; Right—rewritten kernel integrating tile-wise weight prefetch for pipelined compute/load overlap.

All weight transformations, tiling, and memory management are performed so as to maximally exploit physical locality and eliminate pointer arithmetic and branch divergence.

Evaluation

Latency and Memory Savings

FlashMem achieves a marked reduction in both peak and average memory use, as well as drastic speedups in total inference time, particularly on large transformer-type models (e.g., GPT-Neo, SD-UNet, DeepViT). Across 11 diverse DNNs spanning text, image, audio, and video tasks, the framework attains 2.0×–8.4× mean memory reduction and 1.7×–75.0× speedup over the best-performing baseline. Orders of magnitude savings are most pronounced in the largest models, which other frameworks fail to even accommodate (Tables in text).

Latency reduction is attributable to: (i) the avoidance of end-to-end weight preloading, (ii) fine-grained, on-demand streaming, and (iii) pipelined overlap of compute and data movement. Memory reduction emanates from the reduced residency time for weights and elimination of intermediate copies.

Detailed Optimization Contribution and Multi-Model Scenarios

Ablation analysis shows that the OPG-solver, adaptive fusion, and kernel rewriting yield stepped latency improvements (up to 8.1×) and cumulative memory savings (up to 3.8× for OPG). Multi-model scheduling experiments (Figure 4) demonstrate that peak memory never exceeds the tight constraints, even under rapid model switching, and that redundant transformations are eliminated.

Figure 4: Speedup and memory reduction breakdown against baseline (SmartMem).

Latency-Memory Tradeoff

System parameters ($M_{peak}$, $\lambda$, schedule explicitness) allow the user to trade off minimal runtime latency for lower overall memory (Figure 5). For large models, even partial preloading (≈50%) yields near-optimal inference time with significant overall memory savings, due to amortized startup costs.

Figure 5: Memory-latency tradeoff trends; as more weights are preloaded, execution latency falls but total integrated cost may rise due to increased initialization overhead.

Comparison with Alternative Overlap Policies

Baselines adopting naively early/layer-wise prefetching strategies (Figure 6) underperform, exhibiting up to 4.3× higher latency due to imbalances in data movement and kernel execution. FlashMem's load-aware, op-sensitive plan design achieves the optimal overlap configuration.

Figure 6: Comparison of overlap scheduling: integrated, per-op-type, and always-next strategies.

Power, Energy, and Portability

While average power is similar or slightly increased compared to SmartMem (due to higher GPU utilization), energy savings are drastic due to accelerated completion—up to 96%. FlashMem's model-agnostic scheduling and off-line plan generation generalize effectively across multiple mobile GPU architectures without tuning, and it enables execution of large models on devices previously unable to accommodate them.

Implications and Future Directions

FlashMem exemplifies the advantage of system-aware, graph-level static scheduling for DNN acceleration under hard physical resource limits. Its methodology demonstrates that conventional whole-model preloading and ad-hoc fusion policies are fundamentally suboptimal for next-generation mobile inference. By leveraging operator-level heterogeneity and the true architecture of hierarchical mobile GPU memory, the framework creates a path forward for handling even larger foundation models and multi-model interactive workloads.

Theoretically, the offline OPG/CP-SAT approach could generalize to datacenter-class accelerators and cloud inference, particularly as resource pooling, model concurrency, and flexible latency-memory tradeoffs become operational requirements. Practically, the template-driven kernel rewriting pipeline offers a blueprint for integrating hardware features (such as texture memory locality) into system stack optimizations, bridging the gap between model compilation and physical device topology.

Conclusion

FlashMem establishes a new standard for DNN workload management on memory-constrained mobile GPUs. By eschewing monolithic preloads in favor of schedule-driven, pipelined memory streaming and layout-aware kernel design, it supports models and workloads previously untenable on-device. Empirical results confirm substantial latency and memory savings, effective multi-DNN scheduling, and strong portability. Ongoing research directions include scaling the approach to dynamically shaped models, integrating quantized/sparse weight support, and targeting datacenter and edge-cloud accelerators.

---

References:

See (2602.15379).