Samsung LPDDR5X-PIM: Processing-in-Memory for AI Acceleration

Samsung LPDDR5X-PIM increases internal memory bandwidth for AI compute

Samsung is implementing Processing-in-Memory (PIM) by integrating Multiply-Accumulate (MAC) units directly into LPDDR5X chips. This architecture allows compute operations to occur within the memory chip, exploiting higher internal bandwidth and reducing the latency associated with the path between DRAM and traditional compute cores.

While a standard LPDDR5X-9600 chip is limited by its external interface to a maximum bandwidth of 76.8 GB/s, the LPDDR5X-PIM utilizes PIM blocks at each of its 16 banks. By accessing these banks without the constraints of the external bus, the chip can achieve an internal bandwidth of 614 GB/s.

Hardware Architecture and Compute Throughput

Each PIM block consists of a MAC tree, register files, and control logic. The architecture is designed for low-precision AI workloads:

  • Register Files: A 1024-bit instruction register file (up to 64 16-bit instructions), a 4 kbit source register file for activation vectors, and a 2 kbit scale register for scaling model weights.
  • Data Flow: Software loads model weights into DRAM. The DRAM block provides the second operand, while the source register file provides the first.
  • Throughput: Each PIM block can sustain four INT8 or FP8 MAC operations per data clock (eight per cycle). For 4-bit input weights, throughput doubles, leading to a package-wide compute throughput of 2.4 TOPS.

To reach performance levels comparable to dedicated NPUs—such as the Intel Meteor Lake NPU—a system would require multiple LPDDR5X-PIM chips. For instance, eight 16 GB chips (128 GB total system memory) would provide approximately 9.6 INT8 TOPS.

Integration via Standard LPDDR5X Protocol

Samsung has designed LPDDR5X-PIM to remain compatible with standard memory controllers by repurposing existing protocol commands through special row addresses that act as Memory-Mapped I/O (MMIO):

  • Mode Switching: Predefined rows for mode control allow the chip to switch between "single-bank mode" (standard DRAM operation) and "multi-bank mode" (where commands are broadcast across all 16 banks to exploit internal bandwidth).
  • PIM Register Access: Special per-bank rows enable "PIM Registers Activated mode," allowing read and write commands to target PIM registers instead of standard DRAM contents.
  • Execution Flow: In a typical ML use case, software loads weights in single-bank mode, switches to multi-bank mode to write activation values and instructions into PIM registers, and then issues read commands to trigger computations. Results are accumulated in PIM vector register files (VRF) and subsequently written back to DRAM banks.
  • Address Align Mode (AAM): To prevent errors caused by memory controller reordering, AAM allows each instruction to infer its source register index from the column address being accessed.

Critical Software and Architectural Challenges

Despite the hardware efficiency, the LPDDR5X-PIM introduces severe complexities for modern operating systems and CPU architectures:

Cache Coherency and Speculative Execution

Because PIM operations modify memory contents and trigger computations via read commands, they break the fundamental assumptions of CPU memory subsystems. Samsung recommends mapping PIM memory as uncacheable, which creates several issues:

  • Performance Degradation: Disabling caches, prefetchers, and out-of-order execution to avoid incorrect PIM triggers (since a speculative read could trigger an unwanted computation) would severely cripple CPU performance.
  • Side Effects: Standard DRAM reads are idempotent; PIM reads are not. A prefetcher loading data that is later discarded by the CPU could inadvertently modify the PIM VRF state.

Multitasking and Memory Isolation

PIM mode switching is a global state change for the memory channel. This creates significant conflicts in multitasking environments:

  • Thread Interference: If one thread uses PIM while another performs a regular memory access, the non-PIM thread could trigger unintended computations or write VRF data to the wrong address.
  • OS Overhead: Managing PIM requires isolating PIM regions and potentially blocking all other threads and interrupts during PIM execution to prevent state corruption. Saving and restoring PIM state (instructions, source, scale, and VRF registers) during a context switch would be computationally expensive.

Community Perspectives and Alternatives

Technical discussion surrounding the announcement highlights skepticism regarding the practicality of the current implementation:

"The tradeoff with putting the compute in the memory is that you have to know exactly where the dependent information will be at all times. Most problems do not fit this pattern very well."

Critics note that matrix multiplication requires significant data movement to ensure inputs and outputs meet at the same multiplier, suggesting that a ring shift register or similar interconnect between banks might be necessary to avoid the host becoming the bottleneck.

To resolve these issues, the author suggests three systemic hardware changes:

  1. Expanded Interface: Adding dedicated compute commands to the DRAM interface to eliminate mode-switching complexity.
  2. Coherent Memory Controllers: Enabling the memory controller to act as a peer CPU core using Read-For-Ownership (RFO) requests to ensure PIM compute reflects the latest CPU writes.
  3. Transparent ISA Instructions: Introducing new CPU instructions (e.g., "rep macb") that allow the hardware to decide whether to execute a MAC operation in-cache or in-memory based on data size and location.

Sources

Related