COMET: Communication and Memory Co-Design for Fine-Grained AI Inference in MCM Accelerators

Chiplet-based architectures emerge as a promising approach to overcoming the physical and manufacturing constraints faced by monolithic chips, enabling the scalable integration of computing resources to meet the growing demands of AI workloads. However, efficient inter-chiplet communication still faces significant bottlenecks, especially under the fine-grained and bursty Direct Memory Access (DMA) request patterns generated by processing elements in modern AI tasks. Existing communication models and simulators fail to capture these characteristics, which limits the accuracy of performance analysis and the effectiveness of optimization strategies. These limitations hinder DMA-communication inefficiencies in chiplet-based AI systems and pose challenges for designing HPC architectures. To address these challenges, we present the first comprehensive chiplet communication model that explicitly incorporates finegrained DMA traffic observed in realistic AI workloads. Building on this model, we propose COMET, a novel framework that intelligently searches for optimal DMA request aggregation and memory address mapping strategies tailored to chiplet environments. COMET dynamically consolidates small DMA transfers to improve bandwidth utilization and reduce communication latency, while also adapting on-chip memory mapping to align with workload-specific dataflows. This mitigates synchronization overhead across diverse AI tasks. Compared with inference on conventional chiplet communication schemes, COMET achieves$1.1 \times-2.6 \times$speedup and$1.5 \times-4.4 \times$higher bandwidth utilization across different DNN and LLM workloads.