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

2026-08-07

AI Summary

This work addresses the limitations of existing unified memory designs, which assume fixed tensor-device bindings and thus struggle to adapt to dynamic optimal device assignments in LLM inference—such as those caused by phase transitions (e.g., prefill/decode) or dynamic behaviors like MoE routing—leading to suboptimal bandwidth utilization and performance degradation. To overcome this, the paper proposes PFM (PIM-as-Flexible-Memory), the first dual-view unified memory architecture that decouples physical data layout from logical views, enabling NPU- and PIM-optimized access perspectives without requiring data replication or reorganization. By integrating accessor-aware address translation with a runtime scheduling mechanism, PFM efficiently supports dynamic execution demands. Experimental results demonstrate that PFM achieves up to 2.32× end-to-end throughput improvement across various LLMs, significantly enhancing bandwidth utilization and overall system performance.

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.

Problem

• unified memory
• NPU-PIM systems
• dynamic inference
• LLM
• device-biased mapping

Innovation

• dual-view memory
• unified memory
• NPU-PIM systems
• dynamic inference
• processing-in-memory