COMET: A Framework for Modeling Compound Operation Dataflows with Explicit Collectives

Modern machine learning accelerators are designed to efficiently execute deep neural networks, but emerging models increasingly rely on compound operations that introduce significant off-chip memory traffic challenges. As model sizes continue to grow, computation must be distributed across spatial clusters, which requires frequent and complex collective communication. Existing dataflow optimization frameworks and performance models lack the explicit modeling of these collective communication costs, limiting their applicability. To address this, we propose COMET, a framework that introduces a novel representation to explicitly model collective communication across spatial clusters, alongside latency and energy cost models for operation-level dependencies. By enabling collective-aware modeling, COMET allows for broader mapping exploration, achieving up to a $1.42 \times$ speedup for GEMM-Softmax and $3.46 \times$ for GEMM-LayerNorm and $1.82 \times$ for self-attention compared to unfused baselines.