Stardust: Scalable and Transferable Workload Mapping for Large AI on Multi-Chiplet Systems

Workload partitioning and mapping are critical to optimizing performance in multi-chiplet systems. However, existing approaches struggle with scalability in large search spaces and lack transferability across different workloads. To overcome these limitations, we propose Stardust, a scalable and transferable workload mapping on multi-chiplet systems. Stardust combines learnable graph clustering to downscale computation graphs for efficient partitioning, topology-masked attention to capture structural information, and deep reinforcement learning (DRL) for optimized workload mapping. Evaluations on production-scale AI models show that (1) Stardust-generated mappings significantly outperform commonly used heuristics in throughput, and (2) fine-tuning a pre-trained Stardust model improves sample efficiency by up to 15× compared to training from scratch.