Collective Algebra Above a Closed Library: A Research-Loop Agent Finds the Rewrites a Search Misses

Collective-communication libraries on closed accelerator stacks such as AWS Trainium expose a fixed set of primitives and provide little control over how they are scheduled. OverlayCCL searches above this library layer: a language model uses an on-device-calibrated simulator to generate, implement, and evaluate different collective-communication strategies, then refines the most promising ones. However, enumerating structural rewrites alone misses important optimizations. For example, OverlayCCL may miss cases where a value that appears to depend on the current computation is actually a constant determined by the routing pattern, allowing several communication operations to be replaced by a single one. It may also miss opportunities to replace a Python loop that repeatedly prepares data for an unavoidable collective operation with a vectorized operation that performs the same preparation in one step. We therefore introduce SorcarCCL, a research agent that extends OverlayCCL by proposing and evaluating such optimizations using the same simulator and hardware constraints. SorcarCCL compares ideas before implementation, retains validated improvements, records failures, and runs a separate adversarial pass. On the eight collective problems evaluated by OverlayCCL, SorcarCCL matches its results without regression. On 40 more complex production collective problems, we search every problem with nine seeds. Across all 40 problems, the median-seed rewrites run on 224 Trainium cores at 1.12×to 4.37×the performance of OverlayCCL’s solutions, with every problem improved by SorcarCCL. Integrated into end-to-end training of 9.7-billion-parameter GPT- and Llama-style transformers on 224 cores, SorcarCCL’s rewrites further reduce steady-state step time by 2.13×and 2.15× relative to OverlayCCL’s solutions.