A Fast Locality Simulator for GEMM Design-Space Exploration on Multi-Chiplet GPUs

Abstract

In multi-chiplet GPUs, memory accesses that cross the silicon interposer to a remote chiplet's high-bandwidth memory (HBM) incur extra latency and energy, making remote-traffic reduction crucial for efficiency. For general matrix multiply (GEMM), the dominant operator in LLMs, inter-chiplet traffic depends strongly on design knobs such as per-operand memory layout, cooperative thread array (CTA) traversal order, and data placement. The optimal combination is difficult to identify analytically, as locality depends strongly on CTA traversal and its interaction with the L2 cache. To this end, we present a fast, tile-level locality simulator that models data placement and CTA-to-chiplet mapping, CTA traversal, per-chiplet L2 caches, and local/remote HBM accesses. This enables rapid evaluation of locality, performance, and energy efficiency under various GEMM configurations. Using the simulator, we find that the best locality-aware configuration for each GEMM reduces remote traffic by up to 18.3x and improves energy efficiency by up to 17% over 4 KB-interleaved data with round-robin CTA-to-chiplet mapping. Moreover, using the simulator output as feedback, an agentic AI adopts a 2D block-swizzle CTA traversal that improves mean energy efficiency by 15.2% for Qwen and 6.9% for Llama relative to the best 1D traversal under 4 KB-interleaved data placement. Overall, our simulator enables fast exploration of the GEMM locality design space on multi-chiplet GPUs and is available at https://github.com/gthparch/chiplet_locality_simulator.

Poster presentation at the Workshop on Modeling & Simulation of Systems and Applications (MODSIM 2026)