Leveraging Compute-in-Memory for Efficient Generative Model Inference in TPUs
Fuente:
arXiv
Salvato in:
| Autori principali: | , , , , , , |
|---|---|
| Natura: | Preprint |
| Pubblicazione: |
2025
|
| Soggetti: | |
| Accesso online: | |
| Tags: |
Aggiungi Tag
Nessun Tag, puoi essere il primo ad aggiungerne!!
|
| _version_ | 1866916638431903744 |
|---|---|
| author | Zhu, Zhantong Li, Hongou Ren, Wenjie Wu, Meng Ye, Le Huang, Ru Jia, Tianyu |
| author_facet | Zhu, Zhantong Li, Hongou Ren, Wenjie Wu, Meng Ye, Le Huang, Ru Jia, Tianyu |
| contents | 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. |
| format | Preprint |
| id |
arxiv_https___arxiv_org_abs_2503_00461 |
| institution | arXiv |
| publishDate | 2025 |
| record_format | arxiv |
| spellingShingle | Leveraging Compute-in-Memory for Efficient Generative Model Inference in TPUs Zhu, Zhantong Li, Hongou Ren, Wenjie Wu, Meng Ye, Le Huang, Ru Jia, Tianyu Hardware Architecture Artificial Intelligence 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. |
| title | Leveraging Compute-in-Memory for Efficient Generative Model Inference in TPUs |
| topic | Hardware Architecture Artificial Intelligence |
| url | https://arxiv.org/abs/2503.00461 |