# Leveraging Compute-in-Memory for Efficient Generative Model Inference in TPUs

arXiv: https://arxiv.org/abs/2503.00461 · HTML: https://mlsanity.com/paper/2503.00461/

## Key facts

- Cited by 3 papers in the map (data through 14 Oct 2026).
- Builds on 13 of its 31 references that are in the map.
- First posted on arXiv 1 Mar 2025.

## Builds on (most cited first, 13 of 13)

- [Language Models are Few-Shot Learners](https://mlsanity.com/paper/2005.14165/) (2020)
- [High-Resolution Image Synthesis with Latent Diffusion Models](https://mlsanity.com/paper/2112.10752/) (2021)
- [Llama 2: Open Foundation and Fine-Tuned Chat Models](https://mlsanity.com/paper/2307.09288/) (2023)
- [Scalable Diffusion Models with Transformers](https://mlsanity.com/paper/2212.09748/) (2022)
- [Scaling Rectified Flow Transformers for High-Resolution Image Synthesis](https://mlsanity.com/paper/2403.03206/) (2024)
- [A Survey of Large Language Models](https://mlsanity.com/paper/2303.18223/) (2023)
- [Megatron-LM: Training Multi-Billion Parameter Language Models Using Model Parallelism](https://mlsanity.com/paper/1909.08053/) (2019)
- [Open-Sora: Democratizing Efficient Video Production for All](https://mlsanity.com/paper/2412.20404/) (2024)
- [TPU v4: An Optically Reconfigurable Supercomputer for Machine Learning with Hardware Support for Embeddings](https://mlsanity.com/paper/2304.01433/) (2023)
- [Online normalizer calculation for softmax](https://mlsanity.com/paper/1805.02867/) (2018)
- [LLM Inference Unveiled: Survey and Roofline Model Insights](https://mlsanity.com/paper/2402.16363/) (2024)
- [Gemmini: Enabling Systematic Deep-Learning Architecture Evaluation via Full-Stack Integration](https://mlsanity.com/paper/1911.09925/) (2019)
- [CiMLoop: A Flexible, Accurate, and Fast Compute-In-Memory Modeling Tool](https://mlsanity.com/paper/2405.07259/) (2024)

## Built on by (most cited first, 3 of 3)

- [LLM Inference in a Flash!](https://mlsanity.com/paper/2609.16161/) (2026)
- [When Guidance Goes Off-Scale: Recalibrating Diffusion Transformers under Analog Compute-in-Memory Nonidealities](https://mlsanity.com/paper/2608.19644/) (2026)
- [Vision-centric generative AI models: A software-hardware perspective](https://mlsanity.com/paper/2608.27199/) (2026)

## Abstract (arXiv)

With the rapid advent of generative models, efficiently deploying these models on specialized hardware has become critical. Tensor Processing Units (TPUs) are designed to accelerate AI workloads, but their high power consumption necessitates innovations for improving efficiency. Compute-in-memory (CIM) has emerged as a promising paradigm with superior area and energy efficiency. In this work, we present a TPU architecture that integrates digital CIM to replace conventional digital systolic arrays in matrix multiply units (MXUs). We first establish a CIM-based TPU architecture model and simulator to evaluate the benefits of CIM for diverse generative model inference. Building upon the observed design insights, we further explore various CIM-based TPU architectural design choices. Up to 44.2% and 33.8% performance improvement for large language model and diffusion transformer inference, and 27.3x reduction in MXU energy consumption can be achieved with different design choices, compared to the baseline TPUv4i architecture.