Every Kernel Is a Join: Automatic Multi-GPU Parallelism for AI Computations in Einsummable

Every Kernel Is a Join: Automatic Multi-GPU Parallelism for AI Computations in Einsummable Zhimin Ding, Chen-Kuan Liao, Chima Adiole, Brianna Barrow, Fangzhou Du, Yu Hsiao, Ge Huang, Yicheng Jin, Ismail Syed, Chris Jermaine Rice University Houston, USA {zd21,bl86,ca31,bmb15,fd30,rh120,gh31,yj67,is38,cmj4}@rice.edu

arXiv:2609.03905v1 [cs.DC] 3 Sep 2026

ABSTRACT We present Einsummable, a prototype system that accepts a PyTorch-like description of an AI computation and automatically distributes it across a multi-GPU server. Einsummable models every operation as a relational join followed by an aggregation over tensor relations, in which the tuples contain sub-tensors. Each operation exposes its possible decompositions through join-agg specs. An optimizer selects decompositions across the whole computation to minimize a communication-cost proxy. Because it searches decompositions rather than a menu of strategies, Einsummable discovers plans that mesh-based auto-parallelizers cannot. Each decomposed operation is implemented by synthesizing an exchange program, which is a topology-aware generalization of Volcano’s exchange operator. Despite being fully automatic, Einsummable can outperform custom-designed implementations. For example, on LLaMA transformer blocks on an eight-GPU A100 server, Einsummable achieves a geometric-mean runtime of 8.97 ms, versus 13.65 ms for hand-tuned PyTorch and 14.87 ms for vLLM. PVLDB Reference Format: Zhimin Ding, Chen-Kuan Liao, Chima Adiole, Brianna Barrow, Fangzhou Du, Yu Hsiao, Ge Huang, Yicheng Jin, Ismail Syed, Chris Jermaine. Every Kernel Is a Join: Automatic Multi-GPU Parallelism for AI Computations in Einsummable. PVLDB, 20(1): XXX-XXX, 2027. doi:XX.XX/XXX.XX PVLDB Artifact Availability: The source code, data, and/or other artifacts have been made available at https://anonymous.4open.science/r/Einsummable-CE45.

1

INTRODUCTION

Determining how to shard an AI computation across the GPUs of a multi-GPU server, so that adding GPUs actually reduces the time required to run the computation, is one of the central problems in systems-for-AI [20, 32, 42]. The difficulty is that distributing the operations that make up an AI computation (such as dot-product attention [36] or matrix multiplication) can incur significant communication overhead. This overhead can be so significant that it is easily possible to distribute a computation across a set of GPUs and find that adding additional GPUs increases the running time. This work is licensed under the Creative Commons BY-NC-ND 4.0 International License. Visit https://creativecommons.org/licenses/by-nc-nd/4.0/ to view a copy of this license. For any use beyond those covered by this license, obtain permission by emailing info@vldb.org. Copyright is held by the owner/author(s). Publication rights licensed to the VLDB Endowment. Proceedings of the VLDB Endowment, Vol. 20, No. 1 ISSN 2150-8097. doi:XX.XX/XXX.XX

In this paper, we describe a prototype AI system called Einsummable that accepts as input a PyTorch-like [27] description of an AI computation and automatically distributes it across the GPUs of a multi-GPU server, with no human expertise required. Einsummable automatically optimizes the intra-operator parallelism [42] of the computation. Each individual operation is decomposed so that it runs on all GPUs at once, ideally achieving a speedup close to linear in the number of GPUs. “Intra-operator parallelism” is a general class of parallel decompositions that includes data parallelism, tensor (model) parallelism [32, 39], sequence parallelism [22], sharded data parallelism [28], and so on. Crucially, because Einsummable searches a space of decompositions rather than a menu of named strategies, it is not limited to parallelization schemes that have already been invented.1 Optimizing for intra-operator parallelism is fundamentally a database problem. Intra-operator parallelism was the foundation of the shared-nothing parallel database systems of the 1980s and 1990s [10], and its automatic orchestration was the centerpiece of Volcano’s exchange operator [12]. Einsummable brings key ideas from database systems to bear on AI workloads. Einsummable leverages the fact that every operation in a modern AI computation (dot-product attention, softmax, layer normalization, element-wise operations such as the ReLU activation, matrix multiplication, and so on) can be modeled as a relational join followed by a relational aggregation over so-called tensor relations, in the spirit of the tensor-relational model of computation [41]. A “tensor relation” is a set of (sub_tensor_key, sub_tensor) pairs used to represent a tensor. For example, a distributed matrix multiply is precisely a join on the shared inner index followed by a SUM aggregation. In the degenerate, purely relational case, the sub_tensor_key identifies a single cell in the tensor and the sub_tensor is a scalar value. However, no AI system that manipulates individual scalar values can be competitive for modern, dense AI computations. A competitive system must instead invoke efficient kernels, routines (such as a cuBLAS matrix multiply or FlashAttention [7]) that run on an accelerator (GPU, TPU, etc.) over sub-tensors. Thus, the question addressed by Einsummable’s design is: How to decompose each AI operation in a larger AI computation into joins and aggregations over tensor relations such that (a) each “joined” set of tuples represents enough work that a high-performance kernel can be applied effectively, and (b) the communication overhead introduced by the decomposition is minimized? Einsummable brings

1We do not consider multi-server computations, where the slow server-to-server

interconnect means that throughput-enhancing methods such as pipeline parallelism [15] are more practical. A modern server can present 72 GPUs in a single NVLink domain (e.g., NVIDIA’s GB200 NVL72) so single machines cover many use cases.

• We introduce join-agg specs through which each operation exposes its valid decompositions (Section 3). • We present a dynamic-programming optimizer that selects decompositions for an entire computation so as to minimize total communication (Section 4). • We introduce exchange programs through which each decomposed operation is converted to a directed acyclic graph of kernel invocations and data movements, replacing canned collective libraries with communication synthesized per computation (Section 5). • Experimental results show that Einsummable automatically discovers and implements decompositions whose performance meets and sometimes exceeds that of mature systems such as JAX [4] and PyTorch, even expertly tuned. (Section 6).

two classical database ideas to the problem of multi-GPU AI computation: the separation of a logical specification from its physical implementation, and cost-based optimization. Existing auto-parallelizing systems such as FlexFlow, GSPMD, and Alpa [20, 39, 42] search over per-operator sharding annotations drawn from a fixed vocabulary, in which tensor axes are mapped to a logical device mesh. More recent systems such as nnScaler [25] let each operator declare which of its dimensions may be partitioned, but every dimension is still partitioned independent of the data. Unlike these systems, in Einsummable, each operation defines its own space of relational decompositions. The added expressivity is crucial. The classical 3D matrix-multiplication algorithm, for example, shards all three loops of a matrix multiply simultaneously, and therefore cannot be expressed over a two-dimensional device mesh. Or, consider dot-product attention over packed sequences: the sequence dimension may be split only at the boundaries between sequences, positions that depend on the data rather than on the tensor’s shape. And under grouped-query attention, the queryheads dimension cannot be sharded independently, but inherits its sharding from the key/value heads. These constraints are properties of the operator and its input, not of the mesh. Einsummable uses a three-part solution to the problem of automated decomposition. First, every operation available to an Einsummable programmer (such as dot-product attention or matrix multiplication) must export the ability to produce what we call join-agg specs. Each join-agg spec describes one way to decompose the operation’s input tensors into tensor relations so that exactly 𝑝 tuples are produced by the join that is used to implement the operation, where 𝑝 is the number of processors; by varying its requests, the system enumerates the operation’s decompositions. Second, because the communication required to move the tuples of a tensor relation from GPU to GPU and to convert one tensor-relational representation into another can be very costly, Einsummable uses a dynamic programming algorithm to automatically choose the most effective set of decompositions for the overall computation, minimizing an estimate of the total number of bytes that must be communicated. Finally, Einsummable physically realizes the decomposed computation using exchange programs. An exchange program is a topologyaware generalization of Volcano’s exchange operator [12], in which all of the aggregation, repartitioning, and replication between producing and consuming operations is expressed in a domain-specific language. Einsummable synthesizes an exchange program for each operation, guided by a simulator of the actual hardware topology, producing a directed acyclic graph of kernel invocations and datamovement operations, which is compiled into a CUDA graph and executed on the GPU server. Notably, Einsummable never invokes a canned communication pattern such as an NCCL all-reduce. Every communication and aggregation pattern is special-purpose, derived during compilation for the specific data movement required. In summary, this paper makes the following contributions:

2

EINSUMMABLE OVERVIEW

At the highest level, Einsummable allows a programmer to create highly-optimized AI computations in a few lines of code. Einsummable computations run in parallel on multiple AI accelerators without the programmer worrying about how the parallelism happens. In this section, we give a simple Einsummable example from a programmer’s view, and use it to give an overview of the compilation process. Consider a two-layer feed-forward network (FFN), Y = ReLU(XW1 )W2, where X ∈ Rtokens×d_model, W1 ∈ Rd_model×d_ff, W2 ∈ Rd_ff×d_model, and consequently Y ∈ Rtokens×d_model . A programmer would construct this computation in Einsummable as follows: import einsummable as es graph = es.createGraph() x

= graph.insertTensor( es.tensorDesc([tokens, d_model], dtype)) w1 = graph.insertTensor( es.tensorDesc([d_model, d_ff], dtype)) w2 = graph.insertTensor( es.tensorDesc([d_ff, d_model], dtype)) act = graph.matmul(x, w1) hidden = graph.unary(act, "relu", 0.0) y = graph.matmul(hidden, w2) graph.save(y) es.compileGraph(graph) The two calls to matmul and the call to unary add abstract operators to the graph (abstract operators are defined formally in Definition 3.3 of Section 3). The call to graph.save identifies an output that must remain available after execution. The call to compileGraph asks Einsummable to map the abstract computation onto hardware. By default, Einsummable queries the local CUDA runtime and compiles for every visible GPU; the programmer may instead name a subset of devices, as in es.compileGraph(graph,

• We show that the operations making up modern AI computations can be modeled as joins followed by aggregations over tensor relations, giving a single relational framework that captures data, tensor, sequence, and other intraoperator parallelism strategies as special cases (Section 3). 2

Zero-bounded tensors are those whose tensor type has a lower-bound vector of all zeros. Einsummable users manipulate zero-bounded tensors, but internally Einsummable tensor types may have non-zero lower bounds, as it allows the system to describe/manipulate sub-tensors more easily. For human interpretability, zero-bounded tensor types are often specified using named indices. A named index is a textual string (such as num_rows or num_cols) that corresponds to an upper bound (for example, we may have num_rows = 256 and num_cols = 512). Then when we specify a tensor type as [num_rows, num_cols], it corresponds to the tensor type T with T𝑟 = 2, T𝑙 = ⟨0, 0⟩, and T𝑢 = ⟨256, 512⟩. We decompose computations over tensors so that they can be run in parallel. This is done by decomposing the original computation into smaller computations whose input and output types tile the input and output types to the original computation.

device=[0, 1, 2, 3]). The program contains no device assignment for an individual operator, no sharding annotations on an individual tensor, and no explicit communication operations. The number of devices named in this call is the processor count 𝑝 over which all join-agg specs are formed. Join-agg specs are defined in Definition 3.6 of Section 3. Every abstract operator can be asked to produce join-agg specs, which are formal specifications for the tensor-relational decompositions that can be used to implement the operator. During a logical optimization phase, Einsummable considers the join-agg specs available for each abstract operator and chooses compatible decompositions across the entire FFN so as to minimize communication cost, as described in Section 4. Once a logical plan has been produced, the compilation process then produces a data-movement schedule among the chosen devices linking each abstract operator (specified as an exchange program, described in Section 5.2). The default, schedule_mode="naive", uses a naive exchange program. Setting schedule_mode="ai" asks an optimizer, guided by a hardware-topology simulator, to search for a lower-cost schedule of the same decomposed operators (Section 5.4). After compilation, the programmer loads the input tensors (as NumPy arrays), executes the graph, and offloads the result. From the programmer’s perspective, Einsummable resembles a conventional tensor graph API. However, unlike other systems, each operation denotes an abstract operator with many possible join-agg implementations, and compilation selects devices, decompositions, and a data-movement schedule without programmer involvement. The remainder of the paper describes the machinery beneath the API, such as the tensor and join-agg formalism (Section 3), the logical optimizer that chooses decompositions (Section 4), and the physical optimizer that synthesizes exchange programs (Section 5).

3

Definition 3.2 (Tiling). Given a tensor type T , we say that the set of tensor types 𝑇 = {T1 , T2, · · · } tiles T if (a) the indices of the Ð tensor types in 𝑇 cover the indices of T ; that is, I (T ) = 𝑖 I (T𝑖 ) and (b) the indices of the tensor types in 𝑇 are pairwise disjoint; that is, for all 𝑖 ≠ 𝑗, it is the case that I (T𝑖 ) ∩ I (T𝑗 ) = ∅.

3.2

Einsummable programmers build computations using abstract operators over tensors. Abstract operators are operations such as matrix multiplication, flash attention, softmax, and so on. To a programmer, the actual implementation of the abstract operator (how its inputs are sharded, how it is parallelized) is invisible. In this subsection we formally define the notion of an abstract operator, as well as how they are implemented. Definition 3.3 (Abstract operator). An abstract operator in Einsummable is a function over zero-bounded tensors. Given a list of zero or more zero-bounded input tensor types T (𝑖 1 ) , T (𝑖 2 ) , . . ., and a list of zero or more zero-bounded output tensor types T (𝑜 1 ) , T (𝑜 2 ) , . . ., an abstract operator is a function from the set T (𝑖 1 ) × T (𝑖 2 ) × · · · to the set T (𝑜 1 ) × T (𝑜 2 ) × · · · .

TENSORS AND OPERATORS

In this section, we describe the formalisms that underlie our Einsummable implementation.

3.1

Abstract Operators and Join-Agg Specs

Tensors in Einsummable

First, we define the notion of tensors and tensor types. Example 3.4. Consider matrix multiplication. If we have named indices i, j, and k, then matrix multiply is a function from [i, j] × [j, k] to [i, k].

Definition 3.1 (Tensor type; tensor). A tensor type T is a set of functions from index vectors to real numbers, and is specified by three quantities: a rank T𝑟 (a non-negative integer; “r” stands for “rank”, and matrices are rank-2 tensors), an upper-bound vector T𝑢 , and a lower-bound vector T𝑙 , where the bound vectors are integer vectors of length T𝑟 satisfying T𝑙 [ 𝑗] ≤ T𝑢 [ 𝑗] for each 𝑗. Define I (T ) to be the set

Every abstract operator can be decomposed and implemented as a join followed by an aggregation, in a way that is described using a join-agg spec. We first need three auxiliary notions. Definition 3.5 (Projection; function specification; concrete operator). We say that a tensor type T̂ is a projection of a tensor type T if T̂𝑟 = T𝑟 and, componentwise, T𝑙 ≤ T̂𝑙 and T̂𝑢 ≤ T𝑢 ; note that then I ( T̂ ) ⊆ I (T ). Given a tensor U ∈ T , the projection of U along T̂ , denoted Û, is the restriction of U to the domain I ( T̂ ). A function specification is an expression of the form F = T̂ (𝑖 1 ) × T̂ (𝑖 2 ) × · · · → T̂ (𝑜 1 ) × T̂ (𝑜 2 ) × · · · , that is, a list of input tensor types together with a list of output tensor types. Finally, a concrete operator 𝑓 is a polymorphic function over tensors; we write 𝑓 [F ] for the instantiation of 𝑓 at function specification F , which is a function from T̂ (𝑖 1 ) × T̂ (𝑖 2 ) × · · · to T̂ (𝑜 1 ) × T̂ (𝑜 2 ) × · · · .

{T𝑙 [0], . . . , T𝑢 [0] − 1} × {T𝑙 [1], . . . , T𝑢 [1] − 1} × · · · × {T𝑙 [T𝑟 − 1], . . . , T𝑢 [T𝑟 − 1] − 1}. This is the set of all indices or keys that obey the upper and lower bounds of T (lower bounds are inclusive, upper bounds are exclusive). T is then the set of all functions from I (T ) to the set of real numbers, and a tensor (denoted with an upper-case bold letter such as U) having tensor type T is simply one of the functions in T . That is, we say U has tensor type T if U ∈ T . 3

Imagine that we decompose one or more of the tensors that are input into an abstract operator, so that each such input tensor is represented as a set of sub-tensors—that is, as a set of (sub_tensor_key, sub_tensor) pairs, where each key is the tensor type of the corresponding sub-tensor (a tensor relation [41]). The abstract operator is then implemented in three phases. If the abstract operator has 𝑛 inputs, the first phase is a join where 𝑛-tuples of input tensors are formed from the decomposed inputs, and those 𝑛-tuples are sent to the abstract operator’s corresponding concrete operator 𝑓 . Each invocation of 𝑓 happens on a different processor (or AI accelerator). The second phase is an aggregation where some of the outputs of the various 𝑓 invocations may need to be combined. And the final phase is the recompose, where the various partial results are stitched together to form the abstract operator’s final output. We now formally define a join-agg spec and its implementation.

This join-agg spec implements matrix multiply. Take 𝑓 to be the polymorphic block matrix multiply, which for a function specification Â × B̂ → Ĉ multiplies the two projected blocks, contracting over the index range of Â’s second dimension, and take ⊕ to be addition. In step (1) of Definition 3.7, the processor assigned F1 computes the partial products over the first half of j, and the processor assigned F2 over the second half. In step (2), the two results have the identical output type C and are summed elementwise. Step (3) is not necessary, as C is the only output block. We can further decompose so as to require recomposition. Additionally split i into two halves. The resulting join-agg spec has four function specifications F𝑚,𝑛 for 𝑚, 𝑛 ∈ {1, 2}, where F𝑚,𝑛 multiplies the block of the first input with rows in the 𝑚th half of i and columns in the 𝑛th half of j by the block of the second input with rows in the 𝑛th half of j, and has the single output type C𝑚 , with (C1 )𝑙 = ⟨0, 0⟩, (C1 )𝑢 = ⟨128, 256⟩, (C2 )𝑙 = ⟨128, 0⟩, and (C2 )𝑢 = ⟨256, 256⟩. Now two distinct output types appear in the spec. In step (2), the outputs of F𝑚,1 and F𝑚,2 are summed to produce the aggregated block of type C𝑚 , and in step (3) the two aggregated blocks, which tile the output index set, are stitched into the final output tensor.

Definition 3.6 (Join-agg spec). Consider an abstract operator of the form T (𝑖 1 ) × T (𝑖 2 ) × · · · → T (𝑜 1 ) × T (𝑜 2 ) × · · · . A join-agg spec for this abstract operator is a set J = {F1, . . . , F𝑝 } of function specifications in which each input type T̂ (𝑖 𝑗 ) is a projection of T (𝑖 𝑗 ) and each output type T̂ (𝑜 𝑗 ) is a projection of T (𝑜 𝑗 ) , and which satisfies the following tiling condition: for each output index 𝑗, the distinct 𝑗th output types appearing in J tile T (𝑜 𝑗 ) . When a join-agg spec is executed, each of its 𝑝 function specifications is assigned to its own processor.

3.3

Shardable Indices

One or more of the named indices in the specification for an abstract operator are shardable. Shardable indices are used by Einsummable to request join-agg specs from an abstract operator. Every shardable index will either have an associated list of positions that it can be sharded on, or else it will have an annotation any indicating that it can be sharded at any position. To request a join-agg spec from an abstract operator, Einsummable decomposes each shardable index into one or more intervals and provides the intervals to the abstract operator, which returns a corresponding join-agg spec. If intervals[𝑠] is the number of intervals associated with the 𝑠th shardable index, then the returned join-agg spec must conÎ tain exactly 𝑠 intervals[𝑠] function specifications, one for each combination of intervals. Crucially, since the abstract operator is itself responsible for building the join-agg spec, there is a great deal of flexibility in how operators can be decomposed.

Definition 3.7 (Implementation). We say that a join-agg spec J implements an abstract operator if there exist a concrete operator 𝑓 and a binary, associative, and commutative aggregation operation ⊕ over the reals (applied elementwise to tensors of a common type) such that, for every input to the abstract operator, the following three-step procedure recovers the correct output: (1) Join. For each F ∈ J : project each input tensor U (𝑖 𝑗 ) along the corresponding T̂ (𝑖 𝑗 ) from F to obtain Û (𝑖 𝑗 ) , then execute 𝑓 [F ] on Û (𝑖 1 ) , Û (𝑖 2 ) , . . . to obtain Û (𝑜 1 ) , Û (𝑜 2 ) , . . ., where each Û (𝑜 𝑗 ) ∈ T̂ (𝑜 𝑗 ) . (2) Aggregate. For each output index 𝑗 and each distinct type T̂ appearing as a 𝑗th output type in J , combine using ⊕ the 𝑗th outputs of all function specifications whose 𝑗th output type is T̂ , yielding a single tensor V̂ ∈ T̂ . (3) Recompose. For each output index 𝑗, the tiling condition ensures that there is a unique tensor V (𝑜 𝑗 ) ∈ T (𝑜 𝑗 ) that agrees with each aggregated tensor from step (2) on that tensor’s index set; this V (𝑜 𝑗 ) is the 𝑗th output.

Example 3.9 (Matrix multiply). For the matrix multiply of Example 3.4, all three named indices i, j, and k are shardable with the annotation any, since a matrix multiply may be decomposed between any two rows or columns of its inputs. Now let i = k = 256 and j = 512. Suppose that during optimization, Einsummable decomposes i into the intervals [0, 128) and [128, 256), j into the intervals [0, 256) and [256, 512), and leaves k as the single interval [0, 256). Since intervals[i] = 2, intervals[j] = 2, and intervals[k] = 1, the abstract operator must return a join-agg spec containing 2 · 2 · 1 = 4 function specifications, and it returns the join-agg spec {F𝑚,𝑛 } of Example 3.8, where F𝑚,𝑛 corresponds to the 𝑚th interval of i paired with the 𝑛th interval of j. Note the differing roles of the indices. j indexes the contracted dimension, so specifications that share an i interval but differ in their j interval are combined by ⊕ = + during aggregation, whereas the decomposition of i passes through to an output tiling.

Example 3.8 (Matrix multiply, decomposed). Let i = k = 256 and j = 512, and consider the matrix-multiply abstract operator of Example 3.4, from [i, j] × [j, k] to [i, k]. (This may be read as the first matrix multiply of the FFN of Section 2, with i = tokens, j = d_model, and k = d_ff.) A simple join-agg spec for this operator splits the shared index j into two halves. Let J = {F1, F2 }, where F1 = A × B → C and F2 = A ′ × B ′ → C, where A and A ′ are the left and right halves (columns [0, 256) and [256, 512)) of the first input type, B and B ′ are the top and bottom halves of the second, and C is the entire output type [i, k]. Each of these is a projection of the corresponding input or output type of the abstract operator, and since C is the only output type appearing in J , with I (C) = I ( [i, k]), the tiling condition holds trivially.

Example 3.10 (Flash attention). Consider flash-attention with grouped-query attention, with input types [heads_q, sequence, 4

d_head] (the queries), [heads_kv, sequence, d_head] (the keys), and [heads_kv, sequence, d_head] (the values), and with output type [heads_q, sequence, d_head]. Let heads_q = 32 and heads_kv = 8, so that query heads 4𝑚, . . . , 4𝑚 + 3 are associated with the key-value head 𝑚. Of these named indices, only heads_kv and sequence are shardable. heads_q is not directly shardable, yet sharding heads_kv has the effect of sharding heads_q. If Einsummable decomposes heads_kv into its eight unit intervals, the abstract operator returns a join-agg spec with eight function specifications, where the 𝑚th specification projects the key and value inputs onto key-value head 𝑚 and the query input and output onto the query-head interval [4𝑚, 4(𝑚 + 1)). This logic is handled by the abstract operator when it produces the join-agg spec. This is why join-agg spec production is left to individual abstract operators. Decomposition logic can be operator specific and the Einsummable system does not need to directly reason about them. The index sequence is an example of a shardable index without the any annotation. Suppose sequence = 4096 holds three independent sequences packed end to end, ending at positions 1024, 2560, and 4096. Attention never crosses a sequence boundary, but positions within a single sequence will attend to one another, so sequence may be sharded only at the ends of individual sequences.

{out(J ) : J ∈ specs(Op)}. Note that this restriction does not determine the join-agg spec as several specs may produce the same output tiling. For example, in the matrix multiply of Example 3.9, specs that differ only in how the contracted index j is decomposed all induce the same tiling of the output. Consider an abstract operator Op whose inputs are produced by abstract operators Op1, . . . , Op𝑛 , and let specs(Op) denote the set of join-agg specs over 𝑝 processors that Op can produce. This set is obtained by requesting one join-agg spec from Op for each decomposition of its shardable indices into intervals whose counts multiply to 𝑝. The recurrence is Opt(Op,𝑇 ) =

J ∈specs(Op) : out( J )=𝑇 𝑇1 ,...,𝑇𝑛

𝑛 h ∑︁

Opt(Op𝑘 ,𝑇𝑘 )

𝑘=1

i + Cost(J,𝑇1, . . . ,𝑇𝑛 ) , where each𝑇𝑘 ranges over the candidate tilings of the corresponding producer Op𝑘 . That is, to produce the output of Op tiled according to 𝑇 , choose a join-agg spec J for Op that induces the tiling 𝑇 , and choose a tiling 𝑇𝑘 for each input; pay each producer Op𝑘 to deliver its output tiled according to 𝑇𝑘 , and pay to execute J on inputs that are so tiled (including the cost of repartitioning each input from 𝑇𝑘 into the blocks that J requires). The base case is a tensor that is an input to the computation. Such a tensor may be materialized across the processors in any tiling, so Opt(Op,𝑇 ) = 0 for every 𝑇 when Op is an input. The cost of the best logical plan is then min𝑇 Opt(Op∗,𝑇 ), where Op∗ is the operator producing the computation’s designated output tensor, and the plan itself (a choice of join-agg spec for every abstract operator) is recovered by tracing back the minimizing choices. To fill the table for an operator, it is enough to compute, once per spec, 𝑛 h ∑︁ i 𝐵(J ) = min Opt(Op𝑘 ,𝑇𝑘 ) + Cost(J,𝑇1, . . . ,𝑇𝑛 ) ,

4 LOGICAL OPTIMIZATION 4.1 Overview Recall from Section 2 that compilation proceeds in two phases, which roughly correspond to logical and physical optimization in classical database systems. This section describes logical optimization, where Einsummable associates a join-agg spec with each abstract operator, chosen so as to minimize the communication associated with the computation. The result is a logical plan.

4.2

min

Choosing Join-Agg Specs

A computation is a directed, acyclic graph where vertices are abstract operators, and there is an edge (A, B) if abstract operator B consumes the output of A. For the time being we assume that each output is consumed only once, and that each operator produces only a single output. Logical optimization is powered by a dynamic programming algorithm operating over tilings. Let Opt(Op,𝑇 ) denote the lowest communication cost possible for computing the output of abstract operator Op, subject to the constraint that the tiling of the output is 𝑇 . Assume we have a cost function Cost(J, 𝑇1,𝑇2, · · · ) that computes the communication cost required to implement join-agg spec J given that the first input has tiling 𝑇1 , the second 𝑇2 , and so on. This includes the cost of repartitioning each input from its tiling into the input projections required by J , as well as the cost of the join and the aggregation; we describe the details subsequently. Further, for a join-agg spec J , let out(J ) denote the tiling of the output induced by J (that is, the set of distinct output types appearing in J that tiles the output). A key observation is that in Opt(Op,𝑇 ), the tiling 𝑇 need only range over the output tilings that Op’s join-agg specs actually produce. It is never better to first repartition an output into some intermediate tiling and then repartition again at the consumer than it is to repartition directly. Consequently, nothing is lost by restricting the entries of the dynamic programming table for Op to the tilings

𝑇1 ,...,𝑇𝑛 𝑘=1

after which Opt(Op,𝑇 ) is the minimum of 𝐵(J ) over the specs with out(J ) = 𝑇 . Since the table for Op has one entry per distinct tiling, the number of entries is at most the number of specs. If each abstract operator has at most 𝑠 join-agg specs and at most 𝑛 inputs, so that each tensor has at most 𝑡 ≤ 𝑠 candidate tilings, the Í total running time is 𝑂 ( Op 𝑠 · 𝑡 𝑛 ), where the 𝑡 𝑛 factor comes from enumerating combinations of input tilings. Algorithm 1 gives the complete procedure. It processes the operators bottom-up, computing 𝐵(J ) for each spec. For each output tiling, it records the cheapest join-agg spec that induces it, together with the input tilings that achieved the minimum. A final backwards pass walks the tree from the designated output.

4.3

Cost Model

Our implementation of Cost(J,𝑇1, . . . ,𝑇𝑛 ) does not take into account computation cost, as for every abstract operator available in the standard toolkit, the computational cost of implementing each join-agg spec for the operator is the same in terms of the number of floating point operations required. As Einsummable requests decompositions into near-equal intervals, the candidate join-agg 5

Algorithm 1 Logical optimization

numbers. Note that if the 𝑘th input is an input to the computation the tensor is materialized directly into the blocks required by J , and this sub-cost is zero. (2) Join. Consider each function specification in J of the form Í T̂ (𝑖 1 ) × T̂ (𝑖 2 ) ×· · · → T̂ (𝑜 ) . Simply add 𝑘 |I ( T̂ (𝑖𝑘 ) )| to the overall cost. Intuitively: we have to form every 𝑛-tuple that is the input to some invocation of the concrete operator 𝑓 , and this counts the total number of floating point numbers in those 𝑛-tuples. This penalizes join-agg specs that require large inputs to 𝑓 . (3) Aggregation. Consider each distinct type T̂ (𝑜 ) appearing as an output type in J . Count the number of times this distinct type appears in J ; let this be 𝑚. Add (𝑚 −1)|I ( T̂ (𝑜 ) )| to the overall cost. Intuitively: these 𝑚 copies need to be aggregated down to one copy, which (naively) requires moving (𝑚 − 1) copies to the location of the last copy, and this counts the total number of floating point numbers moved.

Require: a computation tree whose leaves are input tensors and whose output is produced by Op∗ , a processor count 𝑝 Ensure: a join-agg spec spec[Op] for every abstract operator Op 1: for each input tensor X of the computation do 2: cand [X] ← {⊥}; Opt[X, ⊥] ← 0 3: end for 4: for each operator Op, in bottom-up (topological) order do 5: let Op1, . . . , Op𝑛 be the producers of the 𝑛 inputs of Op 6: cand [Op] ← ∅ 7: for each J ∈ specs(Op) do 8: choose (𝑇1★, . . . ,𝑇𝑛★) ∈ cand [Op1 ] × · · · × cand [Op𝑛 ] to Í min 𝐵 ← 𝑛𝑘=1 Opt[Op𝑘 ,𝑇𝑘 ] + Cost(J,𝑇1, . . . ,𝑇𝑛 ) 9: 𝑇 ← out(J ) 10: if 𝑇 ∉ cand [Op] or 𝐵 < Opt[Op,𝑇 ] then 11: cand [Op] ← cand [Op] ∪ {𝑇 } 12: Opt[Op,𝑇 ] ← 𝐵 13: choice[Op,𝑇 ] ← (J,𝑇1★, . . . ,𝑇𝑛★) 14: end if 15: end for 16: end for 17: 𝑇 ★ ← the 𝑇 ∈ cand [Op∗ ] minimizing Opt[Op∗ ,𝑇 ] 18: Assign(Op∗ ,𝑇 ★ )

4.4

Detailed Example: 3D Matrix Multiply

We now use this machinery end to end, and show that the optimizer recovers the classical “3D” matrix-multiplication algorithm [1, 9], which is communication-optimal for square matrices [16]. Consider a computation consisting of a single matrix multiply (Example 3.4) with i = j = k = 𝑁 , whose two input matrices are inputs to the computation, to be run on 𝑝 processors; for simplicity, assume 𝑝 = 𝑞 3 for an integer 𝑞, and that the interval counts under consideration divide 𝑁 . As in Example 3.9, every candidate join-agg spec for this operator decomposes i, j, and k into 𝑎, 𝑏, and 𝑐 equal intervals with 𝑎𝑏𝑐 = 𝑝, and contains one function specification per cell of the resulting 𝑎 × 𝑏 × 𝑐 grid. The specification for a given cell multiplies an (𝑁 /𝑎) × (𝑁 /𝑏) block of the first input by an (𝑁 /𝑏) × (𝑁 /𝑐) block of the second producing a partial result for an (𝑁 /𝑎) × (𝑁 /𝑐) block of the output. We evaluate Cost on such a spec. Sub-cost (1) is zero as both inputs are inputs to the computation. For sub-cost (2) each of the 𝑎𝑏𝑐 2 2 function specifications has input blocks totaling 𝑁𝑎𝑏 + 𝑁𝑏𝑐 numbers,  2  2 for a total of 𝑎𝑏𝑐 𝑁𝑎𝑏 + 𝑁𝑏𝑐 = (𝑎 + 𝑐)𝑁 2 . Intuitively, each element of the first input is replicated 𝑐 times, and each element of the second input 𝑎 times. For sub-cost (3), the distinct output types are the 𝑎𝑐 blocks of an 𝑎 × 𝑐 tiling of the output, and each appears 𝑚 = 𝑏 times, contributing (𝑏 − 1)𝑁 2 . Hence

19: procedure Assign(Op,𝑇 )

(J,𝑇1, . . . ,𝑇𝑛 ) ← choice[Op,𝑇 ] spec[Op] ← J 22: for 𝑘 = 1, . . . , 𝑛 do 23: if the 𝑘th input of Op is produced Op𝑘 then 24: Assign(Op𝑘 ,𝑇𝑘 ) 25: end if 26: end for 27: end procedure

20:

21:

specs are load-balanced and so the computation time is effectively the same across the candidate specs. However, different join-agg specs can have very different communication costs. Thus Cost computes a proxy for communication, just as a logical query optimizer in a relational database system costs candidate plans using a proxy such as the number of intermediate tuples produced, deferring precise, hardware-aware costing to physical optimization. Our Cost function estimates the number of floating point numbers that may need to be transferred to run the join-agg spec J , ignoring which processor each block of data actually resides on; data placement is considered during physical optimization. Assuming that each abstract operator produces a single output, the estimate includes three separate sub-costs:

Cost = (𝑎 + 𝑐) 𝑁 2 + (𝑏 − 1) 𝑁 2 = (𝑎 + 𝑏 + 𝑐 − 1) 𝑁 2 . √3 By the arithmetic/geometric mean inequality, 𝑎+𝑏+𝑐 ≥ 3 𝑎𝑏𝑐 = 3𝑞, with equality exactly when 𝑎 = 𝑏 = 𝑐 = 𝑞. The optimizer therefore selects the 𝑞 × 𝑞 × 𝑞 grid, at a cost of (3𝑞 − 1)𝑁 2 = (3𝑝 1/3 − 1)𝑁 2 (each processor multiplies an (𝑁 /𝑞) × (𝑁 /𝑞) block of the first input by an (𝑁 /𝑞) × (𝑁 /𝑞) block of the second, and the 𝑞 partial results for each output block are summed). This is precisely the 3D matrixmultiplication algorithm [1].

(1) Repartition. Consider the 𝑘th input tensor to the operator in question, with type T (𝑖𝑘 ) , and let in𝑘 (J ) denote the set of distinct 𝑘th input types appearing in the function specifications of J . If in𝑘 (J ) ≠ 𝑇𝑘 , then add |I (T (𝑖𝑘 ) )| to the overall cost. Intuitively: if we have to repartition the tensor, we may have to move every floating point number in the tensor around the network; there are |I (T (𝑖𝑘 ) )| such

4.5

Edge Cases

The discussion so far assumes that each abstract operator has one output, consumed by one operator. Three extensions cover the cases that arise in practice. If an operator has several outputs that all flow to the same consumer, Opt is indexed by a tuple of tilings, 6

5.2

one per output, and the consumer selects a single table entry of that producer rather than choosing the tilings independently; the graph remains a tree and optimal substructure is preserved. If an output has several consumers, the computation is no longer a tree, and an exact algorithm would have to track the joint demands on every shared operator, a state space that grows exponentially with the amount of sharing. Instead, the graph is cut at every edge out of a multiply-consumed tensor, and the resulting trees are optimized with Algorithm 1 one at a time in decreasing order of size, each subject to the join-agg specs already fixed (a tensor produced in an already-optimized tree enters the recurrence like a computation input, except that Cost charges any repartitioning). Finally, because the per-operator work is 𝑂 (𝑠 · 𝑡 𝑛 ) in the number of inputs 𝑛, our implementation applies the same tactic to linear paths of operators: each operator on a path draws at most one input from within the path, so the inner minimization in 𝐵(J ) ranges over a single producer’s tilings and the work drops to 𝑂 (𝑠 · 𝑡). In both cases the plan is optimal within each piece, though not necessarily globally.

The ExProg Language

State. An exchange program manipulates named atoms which are tensors, each with a size (in GB). The state of the machine at any instant is a set of pairs (aggregate, site), where an aggregate is a ⊕-combination of atoms, written A+B+C, and a site is a GPU. Because ⊕ is associative and commutative, A+B and B+A denote the same aggregate; the symbol + denotes the aggregation operation ⊕ of the operator at hand, which need not be addition. A program declares an initial state and a required final state, and its body must transform the former into the latter. • GPUS: lists the valid sites, and PARTITIONS: optionally names groups of sites; the groups must be disjoint and cover the declared GPUs. A partition may also be declared by size alone, as in G1 {size 4}, in which case the compiler chooses its membership, and the final state must hold for some choice of membership. • IN: places each atom at a site, optionally with a size, e.g., IN: A @ 0 (size: 1.0 GB) (sizes default to 1 GB). • OUT: lists the aggregates that must exist, and at which sites or partitions, e.g., OUT: A+B @ 0, 1.

5 PHYSICAL OPTIMIZATION 5.1 Overview

Actions. The body of a program is a set of actions, each naming an aggregate, the site at which it is formed, and an optional destination list. A+B @ 0 -> 1 aggregates A and B at site 0 and sends the result to site 1. A @ 0 -> 1, 2, 3 multicasts A to three sites. A+B @ 0 aggregates locally with no network traffic. A+B @ G1 aggregates at every site of partition G1 simultaneously. When convenient, we write several actions on one line, separated by semicolons. There are three key rules. First, a program is a set, not a sequence. An action becomes eligible the moment every atom of its aggregate is present at its source, and all eligible actions proceed concurrently. Intermediate states emerge from dataflow, not program order. Second, every action with an arrow is a two-stage pipeline: a local aggregation (reads and writes at HBM bandwidth), followed by a transfer across the link. Third, aggregating at the source materializes the aggregate there, satisfying the final-state requirement at that site, and other actions may reuse it.

Einsummable’s physical optimizer accepts the output of the logical optimizer (a DAG of abstract operators, each labeled with a selected join-agg spec) and produces an executable artifact. In the current version of Einsummable, this artifact is an NVIDIA CUDA graph, in which each vertex is a GPU kernel or a network operation, which can be executed directly on an NVIDIA server. Physical optimization proceeds through the graph of abstract operators in topologically sorted order, so that when an operator op is reached, every operator supplying one of its inputs has already been compiled. Each producer is compiled through its join so that when op is reached, the unaggregated partial blocks of each of op’s inputs are pinned to specific GPUs. It is op’s job to aggregate those partial blocks as it brings them where they are needed. To do so, the compiler generates the preamble of a program in a domain-specific language called ExProg (for exchange program). This includes an IN declaration recording the location and size of every partial block, and an OUT declaration, derived from op’s joinagg spec, requiring each aggregated input block at the sites that will run op’s concrete operator. The sites in the OUT declaration are expressed as partitions declared by size, so that the placement of op’s 𝑝 concrete-operator instances is itself chosen during synthesis. Exchange programs are a topology-aware generalization of Volcano’s exchange operator [12]. As with the exchange operator, an exchange program is where repartitioning, replication, and interprocessor parallelism are visible. Unlike the exchange operator, an exchange program also performs aggregation, and it is written against an explicit model of the hardware topology. Given the preamble, the body of the exchange program is synthesized so as to minimize the execution time predicted by a hardware simulator. The compiler uses the synthesized program and its simulated execution trace (fixing the dependency structure among transfers, aggregations, and compute kernels) to create the portion of the CUDA graph that links op to its inputs and invokes each of the 𝑝 instances of the concrete operator 𝑓 implementing op.

Cost. A hardware simulator assigns each program a completion time. DGX-A100 and DGX-H100 machines are modeled as fully nonblocking all-to-all NVSwitch fabrics in which the only contention is at endpoints. A multicast to 𝑘 destinations divides the source’s egress bandwidth by 𝑘. The DGX-V100 is far more restrictive as its eight GPUs form a hybrid cube-mesh with 50 GB/s double NVLinks on some pairs, 25 GB/s single links on others, and no direct link on the rest so transfers between unlinked pairs fall back to a PCIe bridge with 15.75 GB/s of total capacity shared by all such transfers. Example 5.1 (A peer-to-peer swap). Two GPUs each hold one atom, and both must end up with the aggregate: IN: A @ 0; B @ 1 OUT: A+B @ 0, 1 A @ 0 -> 1;

B @ 1 -> 0;

A+B @ 0, 1

Both transfers are eligible immediately and proceed concurrently and when complete, the aggregations run. On a DGX-H100 with 1 7

GB atoms, the simulator reports completion in roughly 3.1 milliseconds with 2.2 ms for the (overlapping) transfers, and the remainder for the local aggregations.

families of blocks must be considered. First, the distinct output types of J𝑘 tile T (𝑖𝑘 ) , and one copy of each such block per function specification of J𝑘 that lists it exists somewhere on a GPU. Second, op’s join-agg spec J requires the blocks in𝑘 (J ). These must be delivered, fully aggregated, to the GPU of every function specification of J consuming it. To ensure that no blocks need to be split during the exchange, the compiler forms the overlay of the two. In each dimension, it collects the lower and upper bounds, in that dimension, of every block in both families. The resulting grid of tensor types tiles T (𝑖𝑘 ) and its cells are the atoms for input 𝑘. Every physically existing block and every required block is a disjoint union of atoms. The compiler brackets the exchange program with local kernel calls. Before the program runs decompose kernels slice each resident copy of a producer output block into copies of the atoms it contains. After the program runs, recompose kernels concatenate the aggregated atoms delivered to each GPU into the contiguous input blocks expected by the concrete operator 𝑓 . Given this, the IN declaration lists one ExProg atom per (atom, copy) pair at the GPU holding that copy, with size equal to the cardinality of the atom’s index set times the width of the datatype. The PARTITIONS declaration creates one partition of size one per function specification of op. The OUT declaration requires, for each atom, the ⊕-aggregate of its 𝑚 copies at the partition of every function specification of op whose input blocks contain that atom.

Example 5.2 (Recursive doubling). Four GPUs each hold one atom, and all four must obtain the total aggregate: IN: A @ 0; B @ 1; C @ 2; D @ 3 OUT: A+B+C+D @ 0, 1, 2, 3 A @ 0 -> 1; B @ 1 -> 0 C @ 2 -> 3; D @ 3 -> 2 A+B @ 0 -> 2; A+B @ 1 -> 3 C+D @ 2 -> 0; C+D @ 3 -> 1 A+B+C+D @ 0, 1, 2, 3 The first four actions are eligible at time zero. The second group becomes eligible as its operands arrive. Partial aggregates cross the network in the second round, halving the traffic relative to gathering all four atoms everywhere. On a DGX-H100 with 1 GB atoms, the simulator reports completion in 6.2 milliseconds. Example 5.3 (Topology-aware placement on a DGX-V100). The following specification leaves placement to the compiler: the two partitions are declared by size alone, so the final state must hold only for some choice of two four-GPU groups. GPUS: 0, 1, 2, 3, 4, 5, 6, 7 PARTITIONS: G1 {size 4} G2 {size 4} IN: A @ 0; B @ 1; C @ 2; D @ 3; E @ 4; F @ 5; G @ 6; H @ 7 OUT: A+B+E+F @ G1; C+D+G+H @ G2 On a DGX-V100, the compiler chooses G1 = {0, 1, 4, 5} and G2 = {2, 3, 6, 7}, aligning each group with the machine’s 50 GB/s doublelinks (0↔4, 1↔5, 2↔6, 3↔7), and emits: A @ 0 -> 1; B @ 1 -> 0; C @ 2 -> 3 D @ 3 -> 2; E @ 4 -> 5; F @ 5 -> 4 G @ 6 -> 7; H @ 7 -> 6 A+B @ 0 -> 4; A+B @ 1 -> 5; C+D @ 2 -> 6 C+D @ 3 -> 7; E+F @ 4 -> 0; E+F @ 5 -> 1 G+H @ 6 -> 2; G+H @ 7 -> 3 A+B+E+F @ G1; C+D+G+H @ G2 The first round of swaps uses the 25 GB/s single links within each CPU socket. The second round crosses the socket boundary on the 50 GB/s pillars and no transfer uses the 15.75 GB/s PCIe bridge. The simulator reports completion in roughly 67 milliseconds.

5.3

Example 5.4 (Generating a preamble). Suppose the producer is the matrix multiply of Example 3.8 compiled with its fourspecification join-agg spec {F𝑚,𝑛 }, and that the compiler assigned F1,1, F1,2, F2,1, F2,2 to GPUs 0, 1, 2, 3. Then GPUs 0 and 1 each hold an unaggregated copy of the output block C1 (rows [0, 128)) and GPUs 2 and 3 each hold a copy of C2 (rows [128, 256)). The consumer op is a row-wise softmax whose join-agg spec decomposes i into four intervals of 64 rows, so that in1 (J ) consists of the four blocks with rows [0, 64), [64, 128), [128, 192), [192, 256), each consumed by exactly one function specification. The overlay has row split points {0, 64, 128, 192, 256}, giving four atoms. Each producer block is the union of two atoms (which we call X𝑗 and Y𝑗 for the two copies of atom 𝑗). The generated preamble is: GPUS: 0, 1, 2, 3 PARTITIONS: S1 {size 1} S2 {size 1} S3 {size 1} S4 {size 1} IN: X1 @ 0; X2 @ 0; Y1 @ 1; Y2 @ 1; X3 @ 2; X4 @ 2; Y3 @ 3; Y4 @ 3 OUT: X1+Y1 @ S1; X2+Y2 @ S2; X3+Y3 @ S3; X4+Y4 @ S4 Synthesis chooses S1, . . . , S4 and the program body.

From Join-Agg Specs to Exchange Programs

As Einsummable’s physical optimizer proceeds through the directed, acyclic graph of abstract operators, for a given abstract operator op, the optimizer must produce the preamble for the exchange program used to physically implement op. To do this, the compiler must determine the atoms that the program will operate over. These atoms represent the finest pieces of data that the exchange program can possibly handle so that the exchange program can do its work by aggregating and routing whole atoms. Splitting and concatenation happen only just before or just after the compiled program is executed, in kernel calls generated by the compiler. Consider the 𝑘th input of op, with tensor type T (𝑖𝑘 ) produced by an already-compiled operator whose join-agg spec is J𝑘 . Two

5.4

Synthesizing Program Bodies

Given a preamble we may synthesize the body in a number of ways. For example, we may start with a trivial, valid program, and repeatedly rewrite to a new valid program, using the simulator to cost each candidate. Regardless of the algorithm that chooses the order of transformations applied, they must be sound (they never produce an invalid program) and complete (every valid program is reachable, so in particular a cost-optimal one is). We now define a set of sound and complete transformations. 8

• Source substitution: redirect a transfer of X to draw from another site at which X is materialized. Guarded: the redirected transfer no longer materializes X at the original source.

Definition 5.5 (Closure; validity). A state is a set of pairs (𝑋, 𝑠), where 𝑋 is an aggregate and 𝑠 is a site. The initial state is 𝑆 IN = {({A}, 𝑠) : A @ 𝑠 ∈ IN}. An action 𝑋 @ 𝑠 -> 𝐷 is enabled by a state 𝑆 if 𝑆 contains aggregates 𝑌1, . . . , 𝑌𝑚 at site 𝑠 that partition 𝑋 (that is, 𝑋 = 𝑌1 ∪ · · · ∪ 𝑌𝑚 with the 𝑌𝑖 pairwise disjoint). Firing it adds (𝑋, 𝑠) and (𝑋, 𝑑) for each 𝑑 ∈ 𝐷. The closure 𝑆 ∗ (𝑃) of a program 𝑃 is the least state that contains 𝑆 IN and is closed under firing the enabled actions of 𝑃. 𝑃 is well-formed if every action of 𝑃 is enabled by 𝑆 ∗ (𝑃), and valid if it is well-formed and (𝑋, 𝑠) ∈ 𝑆 ∗ (𝑃) for every requirement 𝑋 @ 𝑠 of OUT.

In our implementation, Definition 5.5 ensures the validity of every candidate exchange program that the search procedure proposes, and Theorem 5.6 guarantees the space of valid programs being searched contains a cost-optimal program.

6

EXPERIMENTAL EVALUATION

Our evaluation asks three questions. First, can our automatic system meet the performance of mature systems expertly tuned? Second, how effective is Einsummable’s optimization: how important is cost-based selection of join-agg specs (Section 4), and how much does the synthesis of exchange programs matter (Section 5)? Third, is the communication-cost proxy of Section 4.3 reasonable? Einsummable is implemented in C++ using CUDA. Concrete operators are high-performance kernels: cuBLAS matrix multiplies, FlashAttention-style attention kernels, and Triton-generated kernels for elementwise operations, aggregation, and the decomposition/recomposition of atoms (Section 5.3). Every plan is compiled to a CUDA graph and executed directly. Exchange-program synthesis uses an off-the-shelf commercial LLM to propose candidate programs. Each candidate is checked for validity (Definition 5.5), costed by the hardware simulator, and the cheapest is kept. Experiments are run on two servers. Most experiments use an NVIDIA DGX A100 with eight 40 GB A100 GPUs connected by a fully non-blocking NVSwitch. Experiment 4 additionally uses an NVIDIA DGX V100 with eight 32 GB V100 GPUs, whose hybrid cube-mesh topology (Section 5.2) is challenging. All computations run in FP16 and are executed with two warm-up rounds followed by ten timed rounds. Figures report mean runtimes in milliseconds with 95% confidence intervals (which are vanishingly small). Two workload families are used throughout. The first is a LLaMAscale transformer block [13]: model dimension 4,096, thirty-two query heads and eight key/value heads (grouped-query attention, as in Example 3.10) of head dimension 128, and FFN dimension 14,336, with attention computed by fused FlashAttention-style kernels. The block is evaluated on five token workloads: single sequences of 1K, 8K, and 128K tokens, and batches of eight and 128 sequences of 1K tokens each. The second family is five fixed matrix chains, labeled A through E, designed so that each multiply averages roughly 81923 scalar multiplies while stressing different plan structures: chains A and B have bushy optimal parenthesizations (three and one independent sub-products, respectively, verified by the classical matrix-chain dynamic program); chain C is a threematrix chain whose optimal order is right-associated, the opposite of what left-deep data parallelism produces; and chains D and E have large end matrices multiplied across thin interior dimensions, so that each plan is forced into a single large terminal aggregation. The parenthesization is fixed and identical for every system. Writing 𝑀𝑖 for the 𝑖th matrix of a chain, the parenthesizations
are: Chain A, (𝑀1 𝑀2 ) ((𝑀3 𝑀4 )((𝑀5 𝑀6 )(𝑀7 𝑀8 ))) 𝑀9 , nine matrices with dimensions alternating 16,384 and 6,144 between ends of 10,240; Chain B, (𝑀1 𝑀2 ) ((𝑀3 𝑀4 ) 𝑀5 ), five matrices with dimensions alternating 20,480 and 5,120 between ends of 9,216; Chain C, 𝑀1 (𝑀2 𝑀3 ), with dimensions 10,240×12,288×32,768×2,048; Chain D,

For partitions declared by size alone, a program is valid if some choice of membership makes it valid; below we treat the membership as fixed. A transformation is a rule that maps a program to a program, and it is sound if it maps every valid program to a valid program for the same IN and OUT. Two primitive transformations suffice for completeness. Insert(𝑎) maps 𝑃 to 𝑃 ∪ {𝑎} and is permitted whenever the result is well-formed. Delete(𝑎) maps 𝑃 to 𝑃 \ {𝑎} and is permitted whenever the result is valid. Both are sound. Deletion is always sound due to the guard, and insertion because of monotonicity (adding an action can only grow the closure). Theorem 5.6 (Completeness). Let 𝑃 and 𝑄 be valid programs for the same IN and OUT. Then 𝑃 can be transformed into 𝑄 by |𝑄 \𝑃 | insertions followed by |𝑃 \𝑄 | deletions, with every intermediate program valid. Proof. Insert the actions of 𝑄 \ 𝑃 in a firing order of 𝑄; enabledness is monotone in the program, so each insertion is permitted, and this reaches 𝑃 ∪ 𝑄. Then delete the actions of 𝑃 \ 𝑄 latest-firing first: every intermediate contains 𝑄, so OUT stays satisfied, and every remaining action is enabled by actions that fire strictly earlier and are still present. □ Since a valid program always exists whenever OUT is satisfiable at all (gather the atoms of each required aggregate at each of its target sites and aggregate locally), the theorem implies that all valid programs are reachable from any valid starting program. Insert and Delete are used merely to argue completeness and no reasonable search strategy would use them. We give a more reasonable set of rules below. Note that these are not unconditional rewrites. Deleting an action also deletes the materializations it produced, and those materializations may be needed elsewhere, by an OUT requirement, or to enable another action. • Multicast fission and fusion: X @ s -> 𝐷 1 ∪𝐷 2 ↔ the pair X @ s -> 𝐷 1 and X @ s -> 𝐷 2 , trading egress contention against action count. Both directions leave the closure unchanged, so both are always permitted. • Relay introduction and elimination: X @ s -> d ↔ X @ s -> r -> d. This is a “bounce” move that routes around slow or contended links. Introduction only adds a materialization at r and is always permitted; elimination is permitted only when X at r is not otherwise required. • Aggregation pushdown and pullup: the pair X @ r -> s, Y @ r -> s ↔ X+Y @ r -> s. This is the network analogue of eager versus lazy partial aggregation [40]. Both directions are guarded: pushdown deletes the separate materializations of X and Y at s, while pullup deletes that of X+Y (a local aggregation X+Y @ s may be inserted to restore it). 9

Table 1: Pearson correlation between predicted communication cost and measured runtime. Transformer workload

𝑟

Matrix chain

𝑟

1 seq × 1K tokens 1 seq × 8K tokens 1 seq × 128K tokens 8 seq × 1K tokens 128 seq × 1K tokens

0.72 0.92 0.80 0.85 0.75

Chain A Chain B Chain C Chain D Chain E

0.73 0.81 0.87 0.47 0.92

Experiment 5: Validating the communication-cost proxy. The twenty random plans per workload from Experiment 3 give, for each of the ten (workload, chain) setups, twenty (predicted cost, measured runtime) pairs. Because FLOPs are invariant across decompositions, runtime should be a linear-plus-noise function of communication, so we report the Pearson correlation (Table 1).

6.2

Over all experiments, fully automatic Einsummable met or exceeded the performance of every system tested. On the transformer block at eight GPUs, Einsummable’s geometric-mean runtime over the five workloads is 8.97 ms, versus 13.65 ms for hand-tuned PyTorch and 14.87 ms for vLLM. In Experiment 1, the gap is biggest on the long-sequence workloads. On the single 128K-token sequence, Einsummable requires 143 ms where PyTorch requires 506 ms and vLLM 498 ms (roughly 3.5× longer). The plan Einsummable selects for this workload is instructive. A single sequence offers exactly one eight-way decomposition of the attention operator. Shard heads_kv, so that each GPU computes attention for one key/value head and its four query heads over the entire sequence. The token activations (131,072 × 4,096, about 1 GB in FP16) dwarf every weight matrix, so the optimizer shards the normalization, elementwise, projection, and FFN operators by tokens with replicated weights, and connects the tokensharded and head-sharded regions by repartitioning the activations at the boundaries of the attention operator rather than sharding the contraction of the output projection. This is the hand-designed DeepSpeed-Ulysses scheme [17] (sequence parallelism outside attention, head parallelism inside it, joined by all-to-all exchanges) which Einsummable derives automatically. With a batch of one, the PyTorch implementation can only shard activations by sequence, so every GPU must gather the full keys and values and absorb the causal-mask load imbalance across sequence shards. In Experiment 2, the three systems are effectively tied. At eight GPUs, every per-chain difference is within about 6%, and all three scale by roughly 6× from one to eight GPUs. We believe that Einsummable’s parity is actually a strong result. With no manual decomposition of any kind, it handles classical workloads as well as systems with many person-years of development effort do. Experiment 3 shows that random plan selection is catastrophic on the transformer block. The median random plan has a geometricmean runtime of 22.8 ms against the optimizer’s 8.97 ms (2.5× worse), and even the best of twenty draws is 17.6 ms. On the matrix chains the penalty for a poor logical plan is smaller. The median random plan is 34% worse in geometric mean, while the best of twenty draws essentially matches the optimizer (within 3%, occasionally beating it on an individual chain). The greedy selector is more interesting. In geometric mean it costs only 2.3% on the transformer block (and on a few individual workloads it is slightly faster, within the noise) but 7% on the chains. This may suggest that exchange programs are able to effectively mask the added cost associated with operator-to-operator repartitions. Experiment 4 confirms that physical optimization can be crucial. On the NVSwitch-based A100, where every pair of GPUs enjoys a dedicated, non-blocking path, optimized and naive exchange programs are statistically indistinguishable (geometric-mean ratio

𝑀1 (𝑀2 (𝑀3 (𝑀4 (𝑀5 (𝑀6 (𝑀7 𝑀8 )))))) 𝑀9 , nine matrices with ends of 24,576, necks of 1,536, and a 20,480-wide interior; and Chain E,
𝑀1 (𝑀2 (𝑀3 (· · · (𝑀11 𝑀12 ) · · · ))) 𝑀13 , thirteen matrices with ends of 49,152, necks of 1,024, and a 20,480-wide interior.

6.1

Discussion

Experiments Run

Experiment 1: Comparing Einsummable with mature, expertlytuned systems on transformer workloads. We compare against JAX [4] (XLA-compiled, with GSPMD sharding annotations [39] and cuDNN fused attention), PyTorch [27] (NCCL communication, flash scaled-dot-product attention kernels, weights replicated and activations sharded by batch or by sequence, whichever is better for the workload), and vLLM’s FlashAttention-2 kernels [24] driven directly. Each baseline was tuned by hand; Einsummable is fully automatic. All four systems run all five token workloads on one, four, and eight GPUs, except that JAX was unable to complete the 128K single-sequence workload. Results for the four larger workloads are given in Figure 1. On the 1K single-sequence workload, omitted from the figure for space, every system finishes in 0.8–2.2 ms at eight GPUs (JAX 0.83, Einsummable 1.21, PyTorch 1.72, vLLM 2.19), and it is included in the geometric means reported below. Experiment 2: Comparing Einsummable with mature systems on matrix-chain arithmetic. On chains A through E we compare against JAX (JIT-compiled, row-sharded) and PyTorch (cuBLAS multiplies with NCCL row sharding), again on one, four, and eight GPUs. Results are given in Figure 2. Experiment 3: Testing the value of coordinated, cost-based logical optimization. We run the transformer block and the matrix chains on eight GPUs under three plan selectors, all with physical optimization on: (a) the full logical optimizer of Section 4 (labeled “EinDecomp” in the figures, after the decomposition algorithm [3] on which it is based), (b) a greedy selector, which is the same dynamic program with the repartition sub-cost removed (since repartitioning is the only sub-cost that depends on the input tilings, this decouples the recurrence and each operator independently chooses its cheapest join-agg spec), and (c) random selection, which chooses uniformly among valid join-agg specs, with twenty independent draws per workload. Results are given in Figures 3 and 4. Experiment 4: Testing the value of physical optimization on a restrictive topology. We compare optimized exchange programs against naive ones (each aggregated input block is simply gathered and aggregated at its consuming sites, with no topology-aware routing, relaying, or placement) for the transformer block and the matrix chains on the DGX V100, which has a challenging network topology. Results are given in Figures 5 and 6. For calibration, we also ran the same comparison on the NVSwitch-based DGX A100. 10

JAX (1) vLLM (1)

JAX (4) vLLM (4)

PyT.

JAX (8) vLLM (8)

PyTorch (1) Einsummable (1)

PyT.

PyTorch (4) Einsummable (4)

224 66.9 33.3

JAX

vLLM Einsum.

306 77.3 38.5

16.2 4.16 2.51

PyT.

19.3 5.18 3.03

JAX

vLLM Einsum.

273 68.5 34.2

JAX

254 58.1 28.9

vLLM Einsum.

17.3 4.59 2.65

no data

0

128 seq × 1K tokens

860 269 143

969 986 506

PyT.

19.2 6.06 4.01

JAX

8 seq × 1K tokens

14.6 3.84 2.20

2

21.3 9.20 5.70

23

19.8 9.52 6.00

26

991 982 498

1 seq × 128K tokens

29 39.1 10.5 5.77

Runtime (ms, log)

1 seq × 8K tokens 212

vLLM Einsum.

PyTorch (8) Einsummable (8)

Figure 1: Transformer-block mean runtimes (95% CIs) for Einsummable vs. JAX, PyTorch, and vLLM on one, four, and eight A100 GPUs (parenthesized numbers in the legend), four LLM workloads.

PyTorch (1)

PyT.

Einsum.

PyTorch (8)

Einsum. (1)

PyT.

Einsum.

PyTorch (4)

47.6

JAX

JAX

14.2 9.13

JAX (8)

9.28 5.85

JAX (4)

34.0

31.6

JAX (1)

10.2 6.14

30.7

Einsum.

10.4 6.57

PyT.

8.58

JAX

8.90

Einsum.

2.19 1.31

2−1

2.43 1.38

5.52 2.84

PyT.

8.59

4.80 2.68

JAX

21

2.59 1.62

4.90 2.70

16.9

17.1

17.2

34.7

33.5

11.9 5.90

23

9.94 5.67

32.8

25

52.6

Chain E 48.1

Chain D

14.2 8.62

Chain C

10.1 5.86

Runtime (ms, log)

2

Chain B

14.6 9.19

Chain A 7

JAX

PyT.

Einsum.

Einsum. (4)

Einsum. (8)

Figure 2: Matrix-chain mean runtimes (95% CIs) for Einsummable vs. JAX and PyTorch on one, four, and eight A100 GPUs (parenthesized numbers in the legend), over chains A to E.

0

0

Greedy

128 32.5

33.3

50

32.7

2.51

0

Random best

104

150 100

2.54

100

2.51

144

5

0

EinDecomp

7.15

10

128×1K

8.74

8×1K

221

181

143

4.01

200

137

1×128K

300

9.39 4.01

3.83

1

5

5.82

2.67

2.16

1×8K 10

1.28

2

1.50

3 1.21

Runtime (ms)

1×1K tokens

0

Random median

Naive schedule

0

2

9.07

11.2

10.3

9.39

10

9.08

7.82

Chain E

15 5.76

6.04

5.97

5.77

1.31

1.59

1.18

1.50

1.30

Chain D 10

5

1

5

0

EinDecomp

Chain C 2 2.81

4.22

2.99

4

2.81

6.08

8.32

6.68

Chain B 2.66

5

Chain A 6.36

10 5.85

Runtime (ms)

Figure 3: Logical-optimization ablation on the transformer block (eight A100 GPUs): the full optimizer vs. greedy, random, and naive-schedule plan selection.

0

Greedy

0

Random best

Random median

0

Naive schedule

Figure 4: Logical-optimization ablation on chains A-E (eight A100 GPUs). 1.006 on the transformer block). With no topology to exploit and the logical optimizer having already removed unnecessary repartitions, a naive solution does well and there is little improvement possible beyind the naive solution. However, on the V100’s hybrid cubemesh, optimized programs are up to 11% faster on the transformer

block (geometric mean 5.6%), and 1.41× and 1.32× faster on chains A and B, whose bushy plans generate cross-socket traffic that the naive programs route through the shared PCIe bridge. On modern hierarchical networks (multi-node NVLink domains, for example) we expect a similar advantage. 11

5,070

1K

8K

8.11

128K

132

goes furthest in our direction, letting each operator declare its partitionable dimensions with an einsum-like annotation; partitioning remains uniform per dimension, aggregation is fixed to summation, and communication is pattern-matched onto NCCL collectives. Einsummable’s join-agg specs are instead produced by operator code, admitting non-uniform, data-dependent decompositions and arbitrary ⊕, and its communication is synthesized for the topology. Collective communication. A recent line of work synthesizes topology-aware implementations of the collectives themselves: SCCL [5] and TACCL [31] search for optimal algorithms realizing a given collective on a given interconnect, and Blink [37] builds collectives from packing spanning trees on heterogeneous topologies. Einsummable takes the further step of having no fixed communication primitives at all. Each exchange program is a special-purpose communication-and-aggregation pattern. Relations on tensor runtimes. Some systems run relational workloads on tensor infrastructure (TQP [14], Hummingbird [26], Koutsoukos et al. [23]) or bring linear algebra and even neural-network training into relational engines (SystemML [11], SystemDS [2], Schüle et al. [30]); Einsummable works in the other direction, using relational optimization for tensor workloads. Tensor-relational optimization. Einsummable builds directly on the tensor-relational model [41], which represents tensors as relations of (key, sub-tensor) pairs so that kernels do enough work to run efficiently on accelerators, and on the broader argument that relational systems are a sound substrate for distributed ML [18, 19, 33]. Our earlier EinDecomp system [3] takes a first step, selecting decompositions with a cost-based dynamic program, but only for computations expressed in an extended Einstein-summation notation (ruling out FlashAttention, for example). SPORES [38] optimizes linear-algebra expressions via relational equality saturation, STOREL [29] optimizes a tensor program given user-supplied storage formats, and Galley [8] brings cost-based query optimization to sparse tensor programs. Unlike Einsummable, these target logical rewrites or single-node kernel structure. Kernel compilers. TVM [6], Triton [34], and TACO [21] generate high-performance kernels for tensor operations on a single device. This work is complementary. Einsummable consumes kernels (cuBLAS matrix multiplies, FlashAttention [7]) as well as kernels produced by these compilers as its concrete operators. Parallel database systems. Intra-operator, partitioned parallelism with automatic orchestration is the founding idea of shared-nothing parallel database systems [10] and of Volcano’s exchange operator [12], which exchange programs generalize (Section 5).

8.82

132

5,036 32.3

2.51

24

29.0

28 2.33

Runtime (ms)

2

12

8×1K

Optimized exchange programs

128×1K

Naive

Figure 5: Optimized vs. naive exchange programs for the transformer block on the topology-restricted DGX V100. 31.0

31.1 18.4

18.3 3.73

3.74

10.1

25.1 7.67

20

17.8

Runtime (ms)

40

0 A

B

C

D

Optimized exchange programs

E

Naive

Figure 6: Optimized vs. naive exchange programs, DGX V100. In Experiment 5, across the ten setups, the Pearson correlation between predicted communication cost and measured runtime ranges from 0.47 to 0.92, with eight of ten setups above 0.72 (Table 1). One source of inaccuracy seems to be that GPU kernels can have varying performance depending on input shapes even when the floating point operations are identical—modeling this should improve accuracy. Also, the point of the cost model is to select the most inexpensive plan, and we note that the selected plan was never beaten by any sampled alternative on the transformer block, indicating that the proxy still ranks plans reasonably. Finally, a note on the choice of baselines. We compare against expertly-tuned production systems rather than against prior autoparallelizing research systems. Alpa [42] can no longer be built against a current JAX stack, and the systems that remain maintained, Unity [35] and nnScaler [25], generate plans for training rather than for the inference workloads studied here.

7

RELATED WORK

Automatic parallelization of AI computations. The named intra-operator strategies in wide use, such as tensor (model) parallelism [32], sharded data parallelism [28], and sequence parallelism [22], were each designed by hand. Auto-parallelizing systems search over combinations of such strategies: FlexFlow [20] searches peroperator parallelization configurations with a randomized strategy, GSPMD [39] is a partitioning mechanism that propagates user-supplied sharding annotations through a computation graph and Alpa [42] automates both inter-operator (pipeline) and intraoperator parallelism, the latter via an integer program over peroperator shardings of tensor dimensions onto a two-dimensional logical device mesh. Unity [35] extends FlexFlow to jointly optimize algebraic rewrites and parallelization for training. nnScaler [25]

8

CONCLUSIONS

We have presented Einsummable, which automatically distributes an AI computation across the GPUs of a server by treating every operation as a relational join followed by an aggregation over tensor relations. Operations export their decompositions as join-agg specs, which are selected to minimize communication, and exchange programs synthesized against the hardware topology realize the plan without canned collectives. 12

There is much work to do. Extending the logical optimization phase so that it is aware of computation cost (and of device properties and how they interact with computational efficiency) would increase accuracy. Extending exchange programs to allow for multimachine parallelism as well as incorporating learning (to understand how the simulator differs from reality) will be crucial.

[19] Dimitrije Jankov, Binhang Yuan, Shangyu Luo, and Chris Jermaine. 2021. Distributed Numerical and Machine Learning Computations via Two-Phase Execution of Aggregated Join Trees. Proceedings of the VLDB Endowment 14, 7 (2021), 1228–1240. [20] Zhihao Jia, Matei Zaharia, and Alex Aiken. 2019. Beyond data and model parallelism for deep neural networks. Proceedings of Machine Learning and Systems 1 (2019), 1–13. [21] Fredrik Kjolstad, Shoaib Kamil, Stephen Chou, David Lugato, and Saman Amarasinghe. 2017. The tensor algebra compiler. Proceedings of the ACM on Programming Languages 1, OOPSLA (2017), 1–29. [22] Vijay Anand Korthikanti, Jared Casper, Sangkug Lym, Lawrence McAfee, Michael Andersch, Mohammad Shoeybi, and Bryan Catanzaro. 2023. Reducing activation recomputation in large transformer models. Proceedings of Machine Learning and Systems 5 (2023), 341–353. [23] Dimitrios Koutsoukos, Supun Nakandala, Konstantinos Karanasos, Karla Saur, Gustavo Alonso, and Matteo Interlandi. 2021. Tensors: An abstraction for general data processing. Proceedings of the VLDB Endowment 14, 10 (2021), 1797–1804. [24] Woosuk Kwon, Zhuohan Li, Siyuan Zhuang, Ying Sheng, Lianmin Zheng, Cody Hao Yu, Joseph Gonzalez, Hao Zhang, and Ion Stoica. 2023. Efficient memory management for large language model serving with pagedattention. In Proceedings of the 29th symposium on operating systems principles. 611–626. [25] Zhiqi Lin, Youshan Miao, Quanlu Zhang, Fan Yang, Yi Zhu, Cheng Li, Saeed Maleki, Xu Cao, Ning Shang, Yilei Yang, et al. 2024. { nnScaler } : { ConstraintGuided } parallelization plan generation for deep learning training. In 18th USENIX Symposium on Operating Systems Design and Implementation (OSDI 24). 347–363. [26] Supun Nakandala, Karla Saur, Gyeong-In Yu, Konstantinos Karanasos, Carlo Curino, Markus Weimer, and Matteo Interlandi. 2020. A Tensor Compiler for Unified Machine Learning Prediction Serving. In 14th USENIX Symposium on Operating Systems Design and Implementation (OSDI). [27] Adam Paszke, Sam Gross, Francisco Massa, Adam Lerer, James Bradbury, Gregory Chanan, Trevor Killeen, Zeming Lin, Natalia Gimelshein, Luca Antiga, et al. 2019. Pytorch: An imperative style, high-performance deep learning library. Advances in neural information processing systems 32 (2019). [28] Samyam Rajbhandari, Jeff Rasley, Olatunji Ruwase, and Yuxiong He. 2020. Zero: Memory optimizations toward training trillion parameter models. In SC20: international conference for high performance computing, networking, storage and analysis. IEEE, 1–16. [29] Maximilian Schleich, Amir Shaikhha, and Dan Suciu. 2023. Optimizing tensor programs on flexible storage. Proceedings of the ACM on Management of Data 1, 1 (2023), 1–27. [30] Maximilian E. Schüle, Thomas Neumann, and Alfons Kemper. 2024. The Duck’s Brain: Training and Inference of Neural Networks in Modern Database Engines. Datenbank-Spektrum 24, 3 (2024). [31] Aashaka Shah, Vijay Chidambaram, Meghan Cowan, Saeed Maleki, Madan Musuvathi, Todd Mytkowicz, Jacob Nelson, Olli Saarikivi, and Rachee Singh. 2023. { TACCL } : Guiding collective algorithm synthesis using communication sketches. In 20th USENIX Symposium on Networked Systems Design and Implementation (NSDI 23). 593–612. [32] Mohammad Shoeybi, Mostofa Patwary, Raul Puri, Patrick LeGresley, Jared Casper, and Bryan Catanzaro. 2019. Megatron-lm: Training multi-billion parameter language models using model parallelism. arXiv preprint arXiv:1909.08053 (2019). [33] Yuxin Tang, Zhimin Ding, Dimitrije Jankov, Binhang Yuan, Daniel Bourgeois, and Chris Jermaine. 2023. Auto-Differentiation of Relational Computations for Very Large Scale Machine Learning. In Proceedings of the 40th International Conference on Machine Learning (ICML). [34] Philippe Tillet, Hsiang-Tsung Kung, and David Cox. 2019. Triton: an intermediate language and compiler for tiled neural network computations. In Proceedings of the 3rd ACM SIGPLAN International Workshop on Machine Learning and Programming Languages. 10–19. [35] Colin Unger, Zhihao Jia, Wei Wu, Sina Lin, Mandeep Baines, Carlos Efrain Quintero Narvaez, Vinay Ramakrishnaiah, Nirmal Prajapati, Patrick S. McCormick, Jamaludin Mohd-Yusof, Xi Luo, Dheevatsa Mudigere, Jongsoo Park, Misha Smelyanskiy, and Alex Aiken. 2022. Unity: Accelerating DNN Training Through Joint Optimization of Algebraic Transformations and Parallelization. In 16th USENIX Symposium on Operating Systems Design and Implementation (OSDI 22). 267–284. [36] Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Łukasz Kaiser, and Illia Polosukhin. 2017. Attention is all you need. Advances in neural information processing systems 30 (2017). [37] Guanhua Wang, Shivaram Venkataraman, Amar Phanishayee, Nikhil Devanur, Jorgen Thelin, and Ion Stoica. 2020. Blink: Fast and generic collectives for distributed ml. Proceedings of Machine Learning and Systems 2 (2020), 172–186. [38] Yisu Remy Wang, Shana Hutchison, Jonathan Leang, Bill Howe, and Dan Suciu. 2020. SPORES: Sum-Product Optimization via Relational Equality Saturation for Large Scale Linear Algebra. Proceedings of the VLDB Endowment 13, 12 (2020), 1919–1932. [39] Yuanzhong Xu, HyoukJoong Lee, Dehao Chen, Blake Hechtman, Yanping Huang, Rahul Joshi, Maxim Krikun, Dmitry Lepikhin, Andy Ly, Marcello Maggioni,

REFERENCES [1] Ramesh C Agarwal, Susanne M Balle, Fred G Gustavson, Mahesh Joshi, and Prasad Palkar. 1995. A three-dimensional approach to parallel matrix multiplication. IBM Journal of Research and Development 39, 5 (1995), 575–582. [2] Matthias Boehm, Iulian Antonov, Sebastian Baunsgaard, Mark Dokter, Robert Ginthör, Kevin Innerebner, Florijan Klezin, Stefanie Lindstaedt, Arnab Phani, Benjamin Rath, Berthold Reinwald, Shafaq Siddiqui, and Sebastian Benjamin Wrede. 2020. SystemDS: A Declarative Machine Learning System for the Endto-End Data Science Lifecycle. In 10th Conference on Innovative Data Systems Research (CIDR). [3] Daniel Bourgeois, Zhimin Ding, Dimitrije Jankov, Jiehui Li, Sleem Mahmoud Abdelghafar, Yuxin Tang, Jiawen Yao, Xinyu Yao, and Chris Jermaine. 2025. EinDecomp: Decomposition of Declaratively-Specified Machine Learning and Numerical Computations for Parallel Execution. Proc. VLDB Endow. 18, 7 (2025), 2240–2253. https://doi.org/10.14778/3734839.3734858 [4] James Bradbury, Roy Frostig, Peter Hawkins, Matthew James Johnson, Chris Leary, Dougal Maclaurin, George Necula, Adam Paszke, Jake VanderPlas, Skye Wanderman-Milne, et al. 2018. JAX: composable transformations of Python+NumPy programs. http://github.com/jax-ml/jax. [5] Zixian Cai, Zhengyang Liu, Saeed Maleki, Madanlal Musuvathi, Todd Mytkowicz, Jacob Nelson, and Olli Saarikivi. 2021. Synthesizing optimal collective algorithms. In Proceedings of the 26th ACM SIGPLAN Symposium on Principles and Practice of Parallel Programming. 62–75. [6] Tianqi Chen, Thierry Moreau, Ziheng Jiang, Lianmin Zheng, Eddie Yan, Haichen Shen, Meghan Cowan, Leyuan Wang, Yuwei Hu, Luis Ceze, et al. 2018. { TVM } : An automated { End-to-End } optimizing compiler for deep learning. In 13th USENIX symposium on operating systems design and implementation (OSDI 18). 578–594. [7] Tri Dao, Dan Fu, Stefano Ermon, Atri Rudra, and Christopher Ré. 2022. Flashattention: Fast and memory-efficient exact attention with io-awareness. Advances in neural information processing systems 35 (2022), 16344–16359. [8] Kyle Deeds, Willow Ahrens, Magdalena Balazinska, and Dan Suciu. 2025. Galley: Modern query optimization for sparse tensor programs. Proceedings of the ACM on Management of Data 3, 3 (2025), 1–24. [9] Eliezer Dekel, David Nassimi, and Sartaj Sahni. 1981. Parallel matrix and graph algorithms. SIAM Journal on computing 10, 4 (1981), 657–675. [10] David DeWitt and Jim Gray. 1992. Parallel database systems: The future of high performance database systems. Commun. ACM 35, 6 (1992), 85–98. [11] Amol Ghoting, Rajasekar Krishnamurthy, Edwin Pednault, Berthold Reinwald, Vikas Sindhwani, Shirish Tatikonda, Yuanyuan Tian, and Shivakumar Vaithyanathan. 2011. SystemML: Declarative machine learning on MapReduce. In 2011 IEEE 27th International conference on data engineering. IEEE, 231–242. [12] Goetz Graefe. 1990. Encapsulation of parallelism in the volcano query processing system. ACM SIGMOD Record 19, 2 (1990), 102–111. [13] Aaron Grattafiori, Abhimanyu Dubey, Abhinav Jauhri, Abhinav Pandey, Abhishek Kadian, Ahmad Al-Dahle, Aiesha Letman, Akhil Mathur, Alan Schelten, Alex Vaughan, et al. 2024. The llama 3 herd of models. arXiv preprint arXiv:2407.21783 (2024). [14] Dong He, Supun Nakandala, Dalitso Banda, Rathijit Sen, Karla Saur, Kwanghyun Park, Carlo Curino, Jesús Camacho-Rodríguez, Konstantinos Karanasos, and Matteo Interlandi. 2022. Query Processing on Tensor Computation Runtimes. Proceedings of the VLDB Endowment 15, 11 (2022), 2811–2825. [15] Yanping Huang, Youlong Cheng, Ankur Bapna, Orhan Firat, Dehao Chen, Mia Chen, HyoukJoong Lee, Jiquan Ngiam, Quoc V Le, Yonghui Wu, et al. 2019. Gpipe: Efficient training of giant neural networks using pipeline parallelism. Advances in neural information processing systems 32 (2019). [16] Dror Irony, Sivan Toledo, and Alexander Tiskin. 2004. Communication lower bounds for distributed-memory matrix multiplication. J. Parallel and Distrib. Comput. 64, 9 (2004), 1017–1026. [17] Sam Ade Jacobs, Masahiro Tanaka, Chengming Zhang, Minjia Zhang, Shuaiwen Leon Song, Samyam Rajbhandari, and Yuxiong He. 2023. Deepspeed ulysses: System optimizations for enabling training of extreme long sequence transformer models. arXiv preprint arXiv:2309.14509 (2023). [18] Dimitrije Jankov, Shangyu Luo, Binhang Yuan, Zhuhua Cai, Jia Zou, Chris Jermaine, and Zekai J. Gao. 2019. Declarative Recursive Computation on an RDBMS: or, Why You Should Use a Database for Distributed Machine Learning. Proceedings of the VLDB Endowment 12, 7 (2019), 822–835. 13

system design. Proceedings of the VLDB Endowment 14, 8 (2021). [42] Lianmin Zheng, Zhuohan Li, Hao Zhang, Yonghao Zhuang, Zhifeng Chen, Yanping Huang, Yida Wang, Yuanzhong Xu, Danyang Zhuo, Eric P. Xing, Joseph E. Gonzalez, and Ion Stoica. 2022. Alpa: Automating Inter- and Intra-Operator Parallelism for Distributed Deep Learning. In 16th USENIX Symposium on Operating Systems Design and Implementation (OSDI). 559–578.

Ruoming Pang, Noam Shazeer, Shibo Wang, Tao Wang, Yonghui Wu, and Zhifeng Chen. 2021. GSPMD: General and Scalable Parallelization for ML Computation Graphs. arXiv preprint arXiv:2105.04663 (2021). [40] Weipeng P. Yan and Per-Åke Larson. 1995. Eager Aggregation and Lazy Aggregation. In Proceedings of the 21st International Conference on Very Large Data Bases (VLDB). 345–357. [41] Binhang Yuan, Dimitrije Jankov, Jia Zou, Yuxin Tang, Daniel Bourgeois, and Chris Jermaine. 2021. Tensor relational algebra for distributed machine learning

14

Related documents

• TensorHub: Rethinking AI Model Hub with Tensor-Centric Compression#120482
• Large-Scale Regularized Matching on GPU Clusters#267631
• AdaProb: Efficient Machine Unlearning via Adaptive Probability#128319
• Diamond: Harnessing GPU resources for scientific deep learning#358992
• Accelerating Dynamic Graph Clustering on GPU Architectures with cuGraph#429327