Skip to content
Preprint

Rethinking Unified Memory for NPU-PIM Systems: Dual-View Memory for Dynamic Inference of LLM

Aug 2026 · 0 citations · 64 references
Computer Science

TL;DR

PFM (PIM-as-Flexible-Memory), a dual-view memory system that decouples physical data layout from accessor-visible logical views, is presented, demonstrating its effectiveness and broad applicability as a unified memory management solution for NPU-PIM systems.

Abstract

Heterogeneous architectures that combine neural processing unit (NPU) and processing-in-memory (PIM) are increasingly adopted to accelerate LLM inference. Prior work focuses on building a unified memory that allows NPUs and PIM to share data without duplication. However, these designs implicitly assume that each tensor is bound to a fixed execution device, and therefore rely on static, device-biased data mappings. We observe that this assumption does not hold in modern LLM workloads. Due to phase changes (e.g., prefill vs. decode) and dynamic behaviors such as MoE routing, the optimal execution device for the same tensor can change at runtime. Under such dynamic execution, device-biased mappings become mismatched to access patterns, leading to substantial bandwidth underutilization and performance loss. This paper presents PFM (PIM-as-Flexible-Memory), a dual-view memory system that decouples physical data layout from accessor-visible logical views. PFM stores data in a jointly optimized physical layout and exposes different logical interpretations to NPUs and PIM, enabling efficient access across devices without data duplication or relayout. We further design accessor-aware address translation and runtime scheduling mechanisms to support dynamic execution when LLM workloads fluctuate and the optimal execution device dynamically changes. Our evaluation across LLMs shows that PFM improves end-to-end throughput by up to 2.32$\times$, demonstrating its effectiveness and broad applicability as a unified memory management solution for NPU-PIM systems.

View source

Similar papers

Book Open access Jul 2026

Aurora: A Disaggregated GPU-PNM-PIM System for High-Throughput Mixed-Length LLM Inference

Modern large language model (LLM) inference is increasingly dominated by memory-bound operations, making processing-near-memory with processing-in-memory (PNM-PIM) an attractive approach for accelerating the decode phase. However, recent long-context LLMs adopt interleaved local and global attention (ILGA), which introduces heterogeneous execution characteristics across transformer blocks. In realistic serving environments, this heterogeneity is further amplified by mixed-length requests, where inputs with widely different context lengths are processed concurrently. These trends break a key assumption underlying existing PNM-PIM systems that transformer blocks exhibit similar latency and resource demands and can be efficiently mapped to a uniform pipeline. Under ILGA and mixed-length workloads, this assumption no longer holds, leading to severe pipeline imbalance and low utilization in prior PNM-PIM designs. Moreover, PNM-PIM–only systems struggle to efficiently support long-context prefill, which remains compute-intensive and is better suited for GPU execution. In this paper, we propose Aurora, a GPU–PNM–PIM disaggregated system designed to efficiently serve mixed-length LLM inference under ILGA. Aurora introduces an ILGA-aware multi-path PNM-PIM pipeline that explicitly accounts for block-level heterogeneity and request-length diversity, improving pipeline utilization without overprovisioning tensor parallelism. Also, Aurora further adopts a stream-oriented Softmax design to reduce stage-level latency imbalance during decoding. To enable end-to-end inference under disaggregation, Aurora treats KV cache transfer as a pipeline operation and coordinates request issuance and scheduling to avoid transfer-induced stalls. Our experimental results show that Aurora achieves up to 8.5 × and 2.2 × higher throughput than the GPU-only baseline and prior PNM-PIM systems, respectively, in end-to-end Llama4-Scout inference, while requiring a comparable or smaller number of devices.

Hyeonu Kim, Seunghyuk Yu, Minjeoung Kim et al. · 0 citations
Preprint Jul 2026

Unified Static-Dynamic Pruning for Efficient LLM Inference

The increasing deployment of large language models (LLMs) has magnified the computational and memory bottlenecks of autoregressive decoding, where low compute intensity and bandwidth-bound kernels dominate inference cost. Weight pruning offers a promising remedy, but existing methods remain confined to either static pruning (SP), which permanently removes redundant weights but lacks adaptivity, or dynamic pruning (DP), which adapts to input sparsity but introduces runtime irregularity. This paper presents SPDP, a unified sparse-inference framework that integrates unstructured SP with input-adaptive DP for efficient LLM inference on GPUs. SPDP co-designs a new Tiled-Column-wise Bitmap Compressed (Tiled-CBC) format and two complementary GPU kernels: (1) a CUDA-core spMspV kernel featuring Hybrid Activation-aware Dynamic Shared-Memory Bitmap Decoding (HAD-SMBD) for fine-grained, runtime activation skipping, and (2) a Tensor-Core SpMM kernel optimized for prefill computation. This joint format-kernel design harmonizes static and dynamic sparsity, maintaining bandwidth-efficient memory access and high compute intensity under both phases of LLM inference. Comprehensive evaluations on inference-optimized GPUs demonstrate that SPDP achieves 1.24x-1.37x average speedup (up to 2.51x) over state-of-the-art sparse frameworks such as SpInfer, while matching perplexity with up to 25% higher sparsity. SPDP advances the inference efficiency-quality Pareto frontier, showing that unified static-dynamic pruning can deliver substantial throughput and performance-per-watt improvements in large-scale LLM serving.

Jin-hong Kim, Yejoo Lee, Jaeyoung Do · 0 citations
Preprint Aug 2026

ZeroLock: Concurrent Memory-Efficient LLM Training via Modular Update Decoupling

This work proposes a BP-free algorithm, called ZeroLock, that decouples the model updates into independent chunk updates by local objective construction and provides the first theoretical framework for such local objective construction-based approaches under general model chunk division by mapping local objectives to the global objective.

Wentao Dai, Xuanran Li, Yuxiang Zhang et al. · 0 citations
#small language model Book Open access Aug 2026

Balancing and Beyond: Communication-Centric Optimizations in Expert Parallelism

EPIC mitigates imbalance via performance-aware expert migration and runtime expert activation, and then improves communication with topology-adaptive transport kernels and fine-grained computation-communication overlap.

Jiamin Cao, Qingxu Li, Yaozhong Liu et al. · 0 citations
Book Open access Jul 2026

Taming Dynamic Diffusion LLM Inference through Virtual Static Execution

Prism abstracts the highly dynamic diffusion workload into a predictable, static execution flow transparent to the compiler, and achieves this via three techniques: spatial regularization, temporal stabilization, and specialized kernels that selectively bypass padding data.

Jianian Zhu, Hang Wu, Yinghui Li et al. · 0 citations
Book Open access Aug 2026

Turbo: Efficiently Serving Long-Context Large Language Models with In-Network Aggregation

This work proposes Turbo, a first-of-its-kind in-network aggregation system that accelerates long-context inference by offloading query broadcast and attention aggregation to switches and introduces a rolling forward scheme that propagates states to enable cross-stage updates.

Ying Wan, Yuchen Xu, Chuwen Zhang et al. · 0 citations