Modeling multi-chiplet architectures for TPU co-design
Abstract
We present an experience report on modeling multi-chiplet hardware accelerators for TPU co-design, and argue that detailed chiplet modeling is necessary to identify bottlenecks and optimizations for future hardware design. Our chiplet modeling tool (CMT) shows that at lower target serving latencies, serving queries/second/chip (QPS/chip) can be up to $2.2\times$ different with chiplet modeling compared to modeling with aggregated compute and memory resources. With this tool, we identify interdependencies between chiplet modeling and the rest of the system: we highlight an example where cost savings from prefetching weights from HBM to local SRAM need to be incorporated in the selection heuristic for chiplet sharding strategies. Using CMT, we explore performance sensitivity to chiplet properties in a design space exploration that highlights considerations for high-performance serving. We use an open-source MoE model (OSS-MoE) served on Ironwood-like TPU architectures as a case study. We conclude by discussing the need to establish best sharding practices to narrow this large design search space.
Research Areas
