Memory stacked chip boosts generative image model speed and efficiency
M3D-FLOW: A Monolithic 3-D-Stackable Processing-Near-DRAM Architecture for Generative Normalizing Flow Applications
Hardware ArchitectureEmerging Technologies
Summary
Generating images using Normalizing Flows is slow and uses a lot of power because the process is repetitive and sequential. The authors present M3D-FLOW, a new computer architecture that stacks processing units directly on memory chips to reduce data bottlenecks. By using a mix of binary and stochastic calculations, and organizing tasks smartly across processing groups, this design speeds up image generation by over six times and uses one-third the energy of previous designs. They also show how shorter stochastic calculations can speed things up more at a slight cost in accuracy.
What this means in practice
- •For hardware designers: Build specialized chips that combine processing units with memory to accelerate generative AI models with improved energy use.
- •For ai infrastructure engineers: Design data center platforms that maximize performance and energy efficiency for running sequential generative normalizing flow tasks.
Authors
Bipin Thapa Magar, Ishan Thakkar, Sudeep Pasricha
Abstract
Normalizing Flows (NFs) have shown strong results in generative modeling, recently reaching state-of-the-art image generation performance. However, their repetitive, sequential operations limit the computational and memory efficiency achievable with GPUs and TPUs. We propose M3D-FLOW, a novel monolithic 3D-stackable (M3D) DRAM-based processing-near-memory (PNM) architecture that eliminates memory bottlenecks using a Cu-Cu hybrid-bonded logic die integrated with an M3D DRAM die. The M3D DRAM provides dense bank integration and a tiered organization with variable access latencies that M3D-FLOW exploits for the efficient processing of different NF workloads. To maximize compute density while minimizing area, energy, and thermal overhead, M3D-FLOW employs hybrid binary--stochastic floating-point multipliers, performing mantissa multiplication in the stochastic domain and exponent operations in the binary domain. A pipelined execution strategy maps sequential NF stages across dynamically partitioned PE groups, and together with the stochastic hardware, enables M3D-FLOW to achieve at least $6.6\times$ higher performance and $3\times$ better energy efficiency than prior PNM accelerators, with only a 2.2% bits-per-dimension (BPD) increase over bfloat16 and 4.1% over FP32. An analysis of stochastic hardware variants further shows that shorter stochastic bitstreams reduce execution time by up to $3.3\times$ at the cost of higher approximation error, providing a concrete design guide for latency--accuracy trade-offs.