Bare-Metal Tensor Virtualization: Overcoming the Memory Wall in Edge-AI Inference on ARM64

AI summary

Overview

Research area: Systems-level optimization of large language model inference on ARM64 edge hardware, with a discussion of implications for computational linguistics. The paper is listed under the Natural Language Processing category (cs.CL) but its content is primarily systems programming, computer architecture, and memory-hierarchy engineering.

Technical level: Advanced. The paper assumes familiarity with the Roofline model, cache-line alignment, SIMD intrinsics, POSIX mmap, and the Llama 2 decoder architecture.

Scope: The paper presents a single-threaded, header-only C++20 inference engine for a 110M-parameter Llama 2 model that replaces standard library containers with memory-mapped tensors and hand-written NEON kernels, and benchmarks it on Apple Silicon.

What This Paper Is About

Autoregressive LLM inference is memory-bound: every generated token requires moving the entire weight set from DRAM to CPU registers, so throughput is capped by memory bandwidth rather than arithmetic speed—the "Memory Wall." Standard C++ implementations add avoidable overhead through std::vector dynamic allocation, pointer indirection, and misaligned struct layouts that cause data to straddle cache lines. The authors build a bare-metal engine that maps model weights directly into the virtual address space and processes them with hand-tuned NEON SIMD kernels, to establish a transparent, portable baseline for how fast general-purpose ARM64 silicon can run such a model without proprietary accelerators.

Key Contributions

1. Bare-metal memory management: The OS heap manager is replaced by a custom memory-mapped arena using POSIX mmap, enabling zero-copy model loading in which weight pointers point directly into the OS page cache and the kernel demand-pages weights via DMA.
2. NEON-optimized kernels: Hand-written SIMD intrinsics implement the critical GEMV (matrix-vector) path using the full 128-bit ARM NEON width, plus an optimized RMSNorm kernel, forcing use of FMLA instructions on the Firestorm microarchitecture.
3. Data-oriented Tensor Virtualization Layout (TVL): A Structure-of-Arrays layout guaranteeing all tensors are 64-byte aligned, which the authors state yields 100% cache line utilization for weight matrices and minimizes TLB misses.
4. A deterministic reference implementation: An open, portable baseline that isolates the interaction between memory bandwidth and arithmetic intensity, unlike the proprietary Apple AMX coprocessor path used by frameworks such as PyTorch.

Main Findings

• Throughput on the 110M model: The bare-metal engine reaches 61.3 tokens/second at 16.3 ms latency, versus 24 tokens/second at 41.6 ms for a scalar C++ baseline compiled with -O3 auto-vectorization. This is reported as a 2.5x speedup (Section 5.2 separately refers to a 3x speedup as the trade-off for a perplexity change).
• Comparison with PyTorch: PyTorch with the Accelerate backend reaches 298.7 tokens/second at 3.3 ms, but the paper states this is achieved by transparently offloading matrix multiplications to the undocumented Apple AMX coprocessor, which the authors characterize as an opaque black box.
• Portability claim: 61 tokens/second is presented as the effective roofline for general-purpose ARM64 cores, applicable to AWS Graviton servers and embedded Linux platforms such as Raspberry Pi 5 where AMX is unavailable.
• Roofline arithmetic: For GEMV the operational intensity is approximately 2 FLOPs/Byte. On an Apple M2 with an arithmetic peak above 3000 GFLOPS and roughly 100 GB/s bandwidth, the memory-bound roof is about 200 GFLOPS, meaning the CPU is starved for data 93% of the time.
• Memory Wall illustration: On a device with 100 GB/s bandwidth, a 70B model (140GB at FP16) would theoretically cap at under 1 token/second regardless of clock speed.
• Weight tying caused segmentation faults: The initial implementation assumed a physically distinct classifier weight region, which failed catastrophically on the 110M model that ties W_cls to the token embedding. A loader heuristic that soft-links the classifier pointer to the embedding pointer when the file is smaller than expected fixed it at zero additional memory.
• Numerical drift from SIMD summation order: Column-major accumulation in the SoA layout changed the summation order, producing a perplexity degradation of less than 0.1% compared to the reference implementation.
• Latency stability: Over N=256 generated tokens, the tight spread between P50 and P99 latency indicated deterministic execution, attributed to pre-faulted memory pages and the absence of dynamic allocation.
• Energy measurements: Using macOS powermetrics, idle power was approximately 5 mW, load power approximately 12 W, and active inference power approximately 8 W. The paper computes 25.3 mJ/token using an average throughput of 316 tokens/second in that equation—a figure that does not match the 61.3 tokens/second reported in the benchmark table, and the discrepancy is not explained in the content.
• Thermal behavior: Because the workload is memory-bound, stalled ALUs produce a "natural duty cycling" that keeps core temperatures below 65 degrees Celsius even during generation sessions longer than 10 minutes, preventing OS downclocking.
• Latency budget for dialogue: The deterministic 16 ms/token rate stays within the 200ms human turn-taking threshold discussed in psycholinguistics research, and the paper argues this makes Beam Search with B=4 feasible on consumer hardware within a 100ms-per-step interaction threshold.
• No quantized comparison is reported. The paper describes 4-bit integer quantization as the LLM standard and as future work, but all benchmarks in it are FP32/FP16.

Methodology in Plain English

The researchers wrote a minimal inference engine from scratch rather than adapting an existing framework. Instead of reading the model file into heap memory, they map the file into the process's address space, so weights become pointers into the operating system's page cache and the kernel loads pages on demand. They lay out all tensors so every buffer starts on a 64-byte boundary, meaning a 16-byte SIMD load always touches exactly one cache line instead of splitting across two. Activation buffers are split into physically separate, non-overlapping arrays (a "Virtual Register File" or RunState) so the compiler can treat them as independent streams. For the dominant matrix-vector multiply, they use NEON intrinsics with four separate accumulators to hide the 4-cycle FMLA latency, and they use a reciprocal-square-root estimate path for RMSNorm. They then benchmarked generation of 256 tokens on a MacBook Pro with an M2 Pro chip and 16GB of unified memory, measuring wall-clock time from the start of the forward call to the return of logits, and compared against scalar C++ and PyTorch CPU baselines.

Why This Matters

Impact on research: The paper argues that frameworks which silently dispatch to proprietary accelerators obscure the relationship between memory bandwidth and arithmetic intensity, making it hard to study the Memory Wall on general-purpose silicon. A bare-metal baseline with no opaque dependencies provides a reproducible reference point for measuring what standard ARMv8 hardware can actually do.

Real-world applications:

• On-device assistants and conversational interfaces that must respond within the 200ms turn-taking threshold to preserve perceived intelligence.
• Deployment on ARM64 cloud and embedded platforms—AWS Graviton servers, Raspberry Pi 5, embedded Linux—where Apple AMX or similar coprocessors do not exist.
• Battery-constrained mobile inference, since the paper argues that every avoided cache miss saves energy as well as time.
• Advanced decoding research (beam search, contrastive search) on consumer hardware, where maintaining multiple parallel hypotheses multiplies bandwidth pressure.

Industry relevance: The work positions open, portable, dependency-light runtimes as an alternative to vendor-locked acceleration, and argues that for single-user local inference the M2 architecture with this runtime occupies an optimal ISO-Energy point, since server-class H100 GPUs in the 300-400W TDP range match that efficiency only at large batch sizes (B > 128).

Future Directions

• Int4 quantization kernels: Implementing a custom 4-bit decoding kernel with NEON vdotq_s32 and vqdmlal instructions to unpack weights into registers on the fly, potentially doubling throughput by halving the memory bandwidth requirement.
• Linguistics-aware mixed-precision KV-caching: Storing stop-words in Int4 while keeping rare named entities in FP16, dynamically adjusting memory pressure based on Shannon Information Content.
• Closing the gap to vendor libraries: Whether the portable NEON path can be pushed closer to the AMX coprocessor's 298.7 tokens/second without adopting opaque dependencies.
• Resolving the reported energy/throughput discrepancy: The 25.3 mJ/token calculation uses a 316 tokens/second figure, while the benchmark table reports 61.3 tokens/second; the paper does not reconcile these, and reconciling them would be needed before the energy claims can be treated as settled.

Target Audience

Systems programmers and performance engineers working on edge inference; researchers studying the memory wall, cache behavior, and SIMD kernel design on ARM64; practitioners deploying LLMs on ARM servers or embedded Linux devices without vendor accelerators; and computational linguists interested in how inference latency bounds constrain realistic dialogue and decoding strategies. Readers without background in computer architecture, SIMD intrinsics, or C++ memory management will find the implementation sections difficult, though the framing sections on the memory wall and latency thresholds are broadly accessible.

Authors’ abstract

The deployment of Large Language Models (LLMs) on edge devices is fundamentally constrained by the "Memory Wall" the bottleneck where data movement latency outstrips arithmetic throughput. Standard inference runtimes often incur significant overhead through high-level abstractions, dynamic dispatch, and unaligned memory access patterns. In this work, we present a novel "Virtual Tensor Core" architecture implemented in software, optimized specifically for ARM64 microarchitectures (Apple Silicon). By bypassing standard library containers in favor of direct memory mapping (mmap) and implementing hand-tuned NEON SIMD kernels, we achieve a form of "Software-Defined Direct Memory Access (DMA)." Our proposed Tensor Virtualization Layout (TVL) guarantees 100% cache line utilization for weight matrices, while our zero-copy loader eliminates initialization latency. Experimental results on a 110M parameter model demonstrate a stable throughput of >60 tokens/second on M2 hardware. While proprietary hardware accelerators (e.g., Apple AMX) can achieve higher peak throughput, our architecture provides a fully open, portable, and deterministic reference implementation for studying the memory bottleneck on general-purpose ARM silicon, meeting the 200ms psycholinguistic latency threshold without opaque dependencies.