1 Introduction↩︎

Deep neural network (DNN) workloads have gained popularity in recent years. Usually, DNN models are expressed as tensor (i.e. \(n-\)dimensional array) based computations. Consequently, DNN developers express these computations using tensor programming languages such as TensorFlow [1], JAX [2], and PyTorch [3], and then utilize tensor compilers like XLA [4], TorchInductor [3], and TVM [5] to generate highly performant executables targeting various hardware devices. These compilers use a multi-stage pipeline similar to general-purpose compilers [6]. During the compiler frontend, computations expressed in tensor programming languages are transformed into a compiler intermediate representation (IR), known as computation graph. Nodes of these graphs represent tensor operators (e.g. matrix multiplications, convolutions), and each edge represents a tensor flowing from the output of a producer node to the input of a consumer node. The compiler middle-end performs graph-level transformations (e.g. operator fusion, tiling) to optimize the computation graph. Finally, the compiler backend maps the computation graph onto a set of kernels, each of which is a program fragment written in single-program-multiple-data (SPMD) fashion on modern hardware accelerators.

To keep up with the computational demands of DNN models, specialized hardware accelerators, such as NVIDIA Tensor Cores, have been introduced. Each new generation of these accelerators has consistently improved the compute capacity to provide blazing speeds for compute instructions such as matrix multiplications (e.g., nearly 1 PFLOPS half precision on NVIDIA H100). Most tensor compilers already aim to maximize utilization on such hardware using dedicated compiler backends (e.g. Triton [7] for GPUs), achieving low latency and high throughput for compute operations.

However, memory technology in modern hardware has not kept pace with these advances in compute capabilities. 1 illustrates the widening compute-to-memory ratio as newer accelerators are introduced, making memory instructions significantly more expensive than compute instructions. Moreover, an increasing number of DNN models are becoming memory bound, with performance dominated by memory access requests to and from the accelerator’s global memory, which exacerbates the impact of slower memory. For example, in large language models (LLMs), the incremental decoding stage is memory-bound and can be the bottleneck of end-to-end inference [8]. Memory access requests are determined by the types and the composition of data movement operations in computation graphs. These operations only transfer data between global memory and the accelerator without performing any computation on tensor data with compute units. Critically, these inserted data movements between computational operators could introduce substantial extra latency. For example, 2 demonstrates TensorRT’s latency breakdown on a Llama 3 8B [9] decoder layer. Between the computation of QKV projection and FlashDecoding [10], the data movement operators take even more time than other computational operators. These findings underscore the critical importance of optimizing data movement operations in DNN compilers.

Prior works that propose optimizations for data movement operations in tensor compilers broadly fall under two categories: layout optimizations and operator fusion, both of which happen in the compiler middle end. These techniques are usually incomplete, targeting only a subset of the data movement operators and miss profitable data movement elimination opportunities that could lead to significant speedups, as detailed below.

1.0.0.1 Layout optimizations.

  The layout of a tensor determines in which order its dimensions are linearized into memory. When a producer operator generates an output tensor with a layout that differs from the expected input layout of the consumer operator, data layout conversion operators need to be inserted, resulting in excessive and avoidable data movement overhead. To find the optimal data layouts that minimize such overheads, previous work first identifies layout-sensitive tensor operators and optimally selects the adjacent layout operators (primarily Reshape and Transpose operators) to balance data movement overhead against computation efficiency [11]–[13]. However, the tensor operators considered in these works represent only a subset of data movement operations used in DNN models, overlooking optimization opportunities in other data movement operators (e.g., ScatterND) that are key for improved performance (see 3.1).

1.0.0.2 Operator fusion.

 During DNN execution, operators usually serve as boundaries of global memory data movement – all input and output tensors are stored in global memory. For each operator, its corresponding kernel first loads inputs from global memory into on-chip buffers, performs computations, and then writes results back from on-chip buffers to global memory. To reduce this costly data movement, existing deep learning frameworks [1], [3], [5], [11] employ operator fusion, which enables intermediate results to be reused directly from fast on-chip buffers rather than being written to and read from slower global memory. However, these systems perform fusion using hand-crafted rules that target specific data access patterns (e.g., [14]), making this optimization only available for operators that can be fused. As a result, similar to layout optimizations, operator fusion reduces data movement for only a subset of data movement operators, missing important optimization opportunities.

1.0.0.3 Our Solution.

 In contrast with these two popular techniques, in this paper, we introduce VTC, a compiler that eliminates unnecessary data movement across all data movement operators. For example, as shown in 2, VTC can eliminate all the data movement operators between QKV projection and FlashDecoding. VTC is comprehensive and works for arbitrary composition of operators in computation graphs, making it generally applicable for all kinds of DNN models. Furthermore, it is complementary to both layout optimization and operator fusion.

The key idea of VTC is to eliminate data movement operators by using virtual tensors that store only a mapping function between physical tensors in global memory. This approach builds on the observation that modern hardware compute units require contiguity only in their local memory buffers, not in global memory. By relaxing the fully contiguous condition to partially contiguous without compromising memory coalescing, we can significantly eliminate data movement operators. Virtual tensors enable us to relax the contiguity condition by adding a level of indirection in the accesses to the tensors in global memory. However, realizing this virtual tensor representation requires addressing a few challenges to fully eliminate unnecessary data movement.

1.0.0.4 Challenge 1: Kernel creation that facilitates virtual tensors.

 With the introduction of virtual tensors, the current computational kernels cannot be used in their present form. First, these kernels assume contiguously laid out data in global memory, which is guaranteed by data movement operators. Second, virtual tensors can introduce extra overhead of mapping function and incontiguous memory access if used in current kernel implementations.

We only modify the global memory I/O stages in existing computational kernels with block virtual tensor I/O (4.2), resulting in minimal changes to the original kernel design and marginal runtime overhead.

1.0.0.5 Challenge 2: Determining a profitable virtual tensor construction strategy.

 Creating virtual tensors for all data movement operators is non-trivial. Different virtual tensor construction strategies lead to different profitabilities.

We solve the problem of finding a profitable option in two parts. First, we introduce a virtual tensor opportunity graph to capture all virtual tensor creation options for a given computation graph (5.1). Next, we design a global greedy algorithm that extracts a virtual tensor creation strategy with maximized latency savings (5.2), which works empirically well in the evaluation.

In summary, we make the following contributions.

• We identify a new data movement optimization opportunity by allowing partially contiguous global memory storage in DNN compilation.

• We introduce the concept of virtual tensors for data movement elimination and efficiently support it in computational kernel generation.

• We propose virtual tensor opportunity graph and data elimination algorithm to automatically find an efficient strategy for creating virtual tensors.

• We implement these concepts in a compilation infrastructure, VTC. Our evaluations show up to \(1.93\times\) improvement (\(1.28\times\) on average) over existing DNN compilers and up to 60% memory saving (17.5% on average) on NVIDIA GPUs.

2 Background↩︎

We first provide background about tensor stride, data movement operators at the graph-level, their necessity, and implementations of tensor operators at the kernel-level before explaining VTC.

2.1 Tensor Stride↩︎

To index an \(n\)-d tensor in the 1D memory, modern tensor compilers need to calculate an integer offset from the index vector \(\vec{I}\in\mathbb{N}^n\). Tensor stride vector \(\vec{S}\in\mathbb{N}^n\) indicates the number of elements we need to skip in memory to move from one element to the next in each dimension. For example, for a 3D tensor of shape \((a,b,c)\), its stride vector \(\vec{S}=(bc,c,1)\). And the indexing will be simply the outer product of the stride vector and the index vector. In the same example, the offset of \(\vec{I}=(i,j,k)\) is \(\vec{I}\cdot\vec{S}=i\cdot bc+j\cdot c+k\).

2.2 Data Movement Operators↩︎

Data movement operators are a subset of tensor operators responsible for only moving data in global memory.² Examples include operators such as Transpose, Split and ScatterND.³ As an example, the semantics of Transpose operator can be given as follows.

5 also shows a visualized example of Split operator. It is clear that Transpose and Split do not change values of individual elements, but rather reorganize their positions in the tensor. This property holds for any data movement operator and we formalize it as follows.

For the Transpose operator, \(F(1,\vec{x})=(1,\vec{y})\) where \(y_i=perm[x_i]\). The mapping from outputs to inputs is one-to-one in data movement operators, ensuring that each output tensor element is derived from a unique input tensor element. Conversely, the mapping from inputs to outputs is one-to-many, indicating that an input tensor element may contribute to multiple output tensor elements.

2.3 Kernel Implementation of Operators↩︎

Graph-level tensor operators are typically implemented as hardware accelerator kernels with all the input and output data physically residing in global memory. The implementation can be divided into three distinct stages:

1. Transfer input data from global memory to on-chip buffers.

2. Perform computation with compute units, which read inputs from and write outputs to on-chip buffer (not present in data movement operators).

3. Transfer output from on-chip buffer to global memory.

For example, Listing [fig:triton] shows a matrix multiplication operator implemented in Triton [7] for GPU execution. In stage 1, two blocks of input matrices (2D-tensors) a and b are loaded as contiguous chunks from global memory into on-chip buffers. Stage 2 does the actual compute, in this case, it is an inner product of the loaded matrices. Stage 3 writes back the resultant c matrix from on-chip buffer to global memory also as contiguous chunks.

@triton.jit
def matmul_kernel(
    a_ptr, b_ptr, c_ptr, # Pointers to matrices
    ... # Elide some meta-parameters
):
    ... # Elide some pre-processing steps
    accumulator = tl.zeros(BLOCK_SIZE)
    for k in range(0, tl.cdiv(K, BLOCK_SIZE_K)):
        # Stage 1: global -> on-chip
        a = tl.load(a_ptrs, ...)
        b = tl.load(b_ptrs, ...)
        # Stage 2: computation
        accumulator = tl.dot(a, b, accumulator)
        # Advance the ptrs to the next K block
        a_ptrs += BLOCK_SIZE_K * stride_ak
        b_ptrs += BLOCK_SIZE_K * stride_bk
    ... # Elide some post-processing steps
    # Stage 3: on-chip -> global
    tl.store(c_ptrs, c, mask=c_mask)

2.4 Necessity of Data Movement Operators↩︎

As seen in 2.3, tensor operators typically assume that the data is contiguous in memory during stages 1 and 3 of a kernel implementation. This serves two purposes:

(1) satisfying the requirement of modern compute units for contiguous memory layout of inputs and outputs,⁴ and (2) optimizing global memory access through memory coalescing and better locality.

Due to the assumption of different kernel implementations, if a producer tensor operator uses a data layout or subsection that differs from its consumer tensor operator, data movement operators must be inserted to reconcile the difference. These data movement operators are necessary to maintain the correctness of the tensor computational graph, but they will introduce extra overhead.

3 Motivation and Overview↩︎

3.1 Motivating Example: LLM Decoding↩︎

To illustrate the potential for data movement optimization in DNN inference, we examine the decoding stage of Llama 3 8B [9] as a motivating example. We profiled a single decoder layer using TensorRT [11], a state-of-the-art DNN compiler on NVIDIA GPUs. ⁵ TensorRT incorporates numerous expert-designed optimization rules and successfully identified the self-attention pattern in the decoder layer, triggering an ad hoc optimization for the entire layer. ⁶

2 shows the computation graph and TensorRT’s latency breakdown on an NVIDIA A100 GPU. TensorRT first applies graph-level transformations to merge the three query, key, and value (QKV) projections into a single matrix multiplication operation. It also employs the FlashDecoding algorithm [10] to fuse the computation of self-attention into one kernel. However, between these two computational kernels, there exists a large number of data movement operators, illustrated in frame 2 of 2. The latency of these data movement operators in TensorRT even surpasses the combined latency of the two computational kernels.

We now focus on the purpose of data movement operators and their role in the computation graph. TensorRT first executes a Split operator to extract the Q, K, V tensors from the result of merged QKV projection. Next, TensorRT uses Reshape and ScatterND operators to update the KV cache [15] with the new K and V tensors. Subsequently, TensorRT applies the Slice operator to keys and values to obtain a slice containing tokens up to the current sequence length. As Llama 3 8B employs grouped-query attention [9], TensorRT then uses Unsqueeze and Expand operators to align the number of KV heads with the number of query heads. Finally, TensorRT runs Reshape and Transpose operators to prepare the layout of keys and values for FlashDecoding. Although operator fusion is available for these data movement operators, TensorRT still requires 6 kernels to complete all the data movements above.

However, according to the analysis of layout requirements in 2, relaxing the fully contiguous condition of global memory storage may unlock opportunities for data movement operator elimination. For example, in stage 3 of the MatMul operator used for the merged QKV projection in frame 1, the MatMul kernel can directly write back the results to Q tensor and appropriate locations in KV cache, eliminating the need for Split, Reshape and ScatterND operators, as shown in 3. Since the dimension per attention head (128) is larger than the GPU warp size (32), memory coalescing is not compromised, allowing for both good memory bandwidth in stage 3 and efficient utilization of tensor cores in stage 2. Similarly, by permitting stage 1 of FlashDecoding to read from not fully contiguous memory, the remaining data movement operators can be eliminated.

Based on these observations, VTC can eliminate all the data movement operators between QKV projection and FlashDecoding. For QKV projection, VTC writes to non-fully contiguous memory, resulting in a slightly slower MatMul kernel compared to TensorRT. However, VTC does not require any GPU kernel to explicitly execute the data movement operators shown in frame 2 of 2, significantly reducing data movement overhead. Furthermore, after VTC’s optimization, the FlashDecoding kernel can directly read from the KV cache instead of the duplicated data generated by the Expand operator. This optimization reduces global memory read overhead by a factor of 4 (the ratio of the Expand operator) and leads to a 2.5\(\times\) speed-up of the FlashDecoding kernel.

3.2 Overview↩︎

4 shows VTC’s overview. The input is a computation graph after operator fusion. To identify all the virtual tensor creation possibilities, VTC first runs an virtual tensor opportunity graph (VTOG) construction algorithm, based on VTC’s virtual tensor definition and data movement optimization rules. Next, VTC runs a points-to graph construction algorithm to find a profitable virtual tensor strategy, represented by the resulting points-to graph. Finally, VTC generates an optimized executable according to this selected strategy.

4 Virtual Tensor↩︎

2 reveals that it is possible to eliminate the need for actual data movement by storing only the mapping function \(F\), which describes the relationship between the input and output tensors. By leveraging this insight, we can avoid the overhead associated with explicit data transfers in data movement operators.

In this section, we introduce VTC’s virtual tensor notation, which builds upon this observation to efficiently handle data movement operations. We will show how to handle virtual tensor I/O in 4.2 and analyze the profitability of the optimization in 4.3. Finally, we give several examples of how to use virtual tensors to eliminate data movement operators in 4.4.

4.1 Virtual Tensor Definition↩︎

The key idea behind the virtual tensor technique is to avoid redundant instantiation of tensor data in global memory. If a tensor is obtained by performing data movement operators on other tensors, we can directly represent it using a mapping function and a set of physical tensor pointers, instead of allocating global memory to store the tensor data.

3 shows an example of virtual tensor in QKV projection. Virtual tensor \(c\) has 3 physical tensor pointers \(Q\_ptr\), \(K\_Cache\_ptr\), and \(V\_Cache\_ptr\), with its mapping function illustrated in subfigure (b).

A virtual tensor provides a contiguous index space that references potentially non-contiguous locations within physical tensors. Mapping function \(F\) links each virtual tensor index to its corresponding physical tensor address. The virtual tensor technique eliminates data movement operators through two complementary approaches: (1) creating output virtual tensors for a producer data movement operation that map to input tensors of that operation, or (2) creating input virtual tensors for a consumer data movement operation that map to output tensors of that operation. While 2 allows for one-to-many mappings, VTC restricts \(F\) to one-to-one mappings. This design choice simplifies implementation and avoids potentially expensive overhead of translating complex one-to-many mapping functions.

4.1.0.1 Mapping function composition.

Virtual tensors can be nested: it is possible for a base tensor, which is used to create a virtual tensor, to be a virtual tensor itself. In such cases, the resulting virtual tensor’s mapping function effectively applies the base tensor’s mapping followed by the new mapping (e.g., \(F_{\rm new}=F_{\rm outer}\circ F_{\rm base}\)), representing a chain of transformations.

4.2 Virtual Tensor I/O↩︎

VTC’s virtual tensor interface seamlessly translates virtual tensor I/O operations into physical global memory I/O operations using the mapping function. In hardware kernels, global memory I/O is usually performed in a block-wise manner, where multiple threads simultaneously read a block of global memory data. For most real-world cases, the contiguity of mapping function conditions (e.g., 1024 in 3) is larger than and a multiple of block size. Therefore, each block is likely to read from the same physical tensor, allowing us to directly apply a non-partitioned mapping function to a block of indices, which only introduces marginal overhead.

Virtual tensor provides a high-level abstraction that hides the underlying data movement complexities. Users can work with virtual tensors as if they were regular tensors, while VTC transparently manages the mapping and access to the physical tensor data behind the scenes. Since virtual tensors only impact the global memory I/O (stage 1 and 3 in [fig:triton]) of a hardware accelerator kernels, they leave the computation part (stage 2 in [fig:triton]) unchanged. This allows virtual tensors to be effortlessly integrated into existing optimized kernels with minimal modifications. For example, in 3, we can directly pass a virtual tensor mapping to physical Q tensor and KV cache as the output parameter of the projection MatMul kernel. VTC then transforms the fully contiguous write to this virtual tensor into partially contiguous writes to the Q tensor and KV cache in MatMul kernel stage 3.

While virtual tensors have the potential to eliminate physical data transfers in data movement operators, the use of mapping functions and non-fully contiguous global memory access patterns still may introduce additional overhead. In the next subsection, we will delve into the profitability of this optimization technique, analyzing the trade-offs between the benefits of reduced data movement and potential extra costs.

4.3 Profitability Analysis↩︎

Since the indirect memory access patterns of virtual tensors are fully known at compile time, we can analyze the optimization profitability based on the characteristics of the mapping function. Similar to the classification of compiler optimizations into Type I (always profitable) and Type II (uncertain profitability) [16], we can categorize the profitability of virtual tensor optimizations by analyzing the mapping function. Without loss of generality, we assume each virtual tensor maps to a single base physical tensor here.

1 establishes the conditions for guaranteed profitability. When the mapping function preserves the original data access pattern with only a constant bias allowed, the virtual tensor optimization can be considered as Type I. Also, Virtual tensors with partial contiguity are likely profitable since they do not break minimal memory transfer granularity and maintain high memory bandwidth. Otherwise, the profitability is uncertain. VTC will perform profiling (described in 5.2) to decide whether to proceed the optimization according to actual latency impacts.

4.4 Virtual Tensor Optimization Rules for Data Movement Operators↩︎

The virtual tensor optimization is designed to eliminate data movement operators. In this subsection, we present several examples of data movement operators from ONNX, demonstrate how VTC eliminates them using virtual tensors, and discuss the profitability. All of the data movement operators presented in this subsection are outside the optimization space of previous layout optimizations.

4.4.0.1 Split

 When optimizing Split using virtual tensors, making input a virtual tensor of outputs is Type II, since the output tensors are stored separately and incontiguously in global memory. On the other hand, making outputs virtual tensors of input can be either a Type I or Type II optimization, depending on the contiguity of the split. If the split is performed along the first axis (axis 0), resulting in contiguous output tensors, it is a Type I optimization. 5 shows an example of contiguous Split and incontiguous Split. When the incontiguous split size is large enough, all the two directions of virtual tensor optimizations have partial contiguity.

4.4.0.2 Expand

 The Expand operator broadcasts a single input tensor to a specified shape, producing a single output tensor with the same data as the input but with expanded dimensions. When optimizing Expand using virtual tensors, both directions are Type II. But if the number of duplicated elements exceeds memory coalescing size, it is almost always profitable.

4.4.0.3 ScatterND

 The ScatterND operator takes 3 inputs: data, indices, and updates. It generates a single output tensor by first initializing it as a clone of the data tensor and then updating the values at the specified indices in indices with the corresponding values from the updates tensor. There are two ways of virtual tensor optimization:

1. Making the output a virtual tensor of data. Since the output tensor is initially a clone of data, it is a Type I optimization.

2. Making updates a virtual tensor of the output tensor. This optimization is Type II because the indices can be incontiguous. In the LLM motivation example, updates denotes all the locations in KV cache to be updated and is partially contiguous.

We can conduct analysis and define optimization rules similarly for all other data movement operators. As virtual tensor optimization rules are highly dependent on the operator semantics and the number of data movement operators is limited (e.g., 18 in the ONNX), VTC requires developers to specify the mapping function and possible virtual tensor optimization rules between inputs and outputs for each data movement operator. Then VTC can automatically analyze profitability and perform virtual tensor optimizations, which will be introduced in the next section.

5 Automatic Virtual Tensor Construction↩︎

This section introduces VTC’s virtual tensor construction algorithm, which automatically makes intermediate tensors virtual to minimize unnecessary data transfers and improve overall performance. We first formalize the data movement elimination problem by building a virtual tensor opportunity graph (VTOG), and then introduce the global greedy algorithm to eliminate data movements.

5.1 Virtual Tensor Opportunity Graph↩︎

VTC’s data movement eliminator takes a computation graph after operator fusion as input.

In prior work, all the intermediate tensors, represented by edges in the computation graph, are physically stored in global memory. The data movement elimination problem aims to find a strategy to make certain intermediate tensors virtual and maximize the end-to-end latency savings.

VTC first builds a virtual tensor opportunity graph (VTOG) to analyze all the virtual tensor possibilities.

6 shows an example of a computation graph and its corresponding VTOG. Split determines the VTOG edges between a and b; Reshape determines the VTOG edges between b and c; and ScatterND determines the VTOG edges between c, d and K Cache.

For each node in the VTOG, some outgoing edges may conflict with each other, indicating mutually exclusive virtual tensor strategies. For example, in 7, edges 1 and 3 conflict because designating c as a virtual tensor of both a and d simultaneously would map some indices of c to multiple physical memory locations, violating the one-to-one property of virtual tensor mapping functions. In contrast, edges 1 and 2 can coexist when preserving a one-to-one mapping. We capture all such conflicting edge pairs in a conflict set.

8 shows how to create a VTOG and the conflict sets from a computation graph. VTC constructs the VTOG by iterating through all data movement operators in the input computation graph, following the predefined virtual tensor rules for each operator (4.4), and creating edges between operator inputs and outputs.

When all outgoing edges of each node in a VTOG are mutually compatible, the VTOG becomes a points-to graph, similar to that used in traditional compiler pointer analysis [17]. The points-to graph has one-to-one correspondence with a virtual tensor strategy for the entire computation graph. For example, 9 shows several possible points-to graphs derived from the VTOG in 6. Strategy 3 has no edges, indicating that all the tensors are physical. Strategy 2 has b and K cache as physical tensors, making a and c virtual tensors of b, and d a virtual tensor of K cache. Strategy 1 makes a, b, c, and d all virtual tensors of K cache. Note that K cache is an input tensor of the computation graph, so it must reside in physical memory.

Using an inductive argument, it is not hard to show the computational equivalence of the computation graph and the virtualization strategy corresponding to a points-to graph derived from the VTOG. As evident from the definition, a points-to graph can be derived from a VTOG by removing conflicting edges. However, different points-to graphs can lead to different end-to-end latency gains. The following subsection introduces VTC’s points-to graph construction algorithm, which aims to obtain a points-to graph representing an efficient virtual tensor strategy.

5.2 Global Greedy Algorithm↩︎

The objective of our data movement elimination algorithm is to maximize the latency saving for each virtual tensor strategy, which can be formalized as follows. Let \((E_{\rm VTOG}, \mathbb{F}, \ell)\) denotes the virtual tensor search space corresponding to a VTOG. Here each edge in \(E_{\rm VTOG}\) denotes a single virtualization opportunity, \(\mathbb{F}\) denotes a set of all points-to graphs and \(\ell: \mathbb{F} \rightarrow \mathbb{R}\) is the latency saving function maps for every points-to graph \(C \in \mathbb{F}\) to the latency saving corresponding to virtual tensor strategy realized by \(C\). Our goal is to find the optimal points-to graph \[\label{e:latencyfunction} C^* = \text{arg}\,\max_{\mathbb{F}} \left\{ \ell(C) \right\}\tag{1}\]

1 is NP-hard as \(\ell\) is a complex set function without any structure, which can only be measured through profiling on real hardware. In some situations we can expect that \(\ell\) is a submodular function; for example, when all our operations are Type I (see 4.3). Unfortunately, it is not true for more general VTOGs.

Despite this, we take inspiration from the submodular function maximization literature, and design a global greedy algorithm to approximately solve our problem. It is well-known that the global greedy algorithm, which at each iteration adds the element that maximizes the discrete derivative of a submodular function \(\ell\), achieves a constant approximation to the optimal solution for a broad range of settings; we refer readers to [18] for an introductory survey.

10 shows VTC’s points-to graph construction procedure. With all tensors stored in physical memory initially, it iteratively creates new virtual tensors. Specifically, the algorithm maintains an anchor node set \(A\) and a currently selected edge set \(C\). In each iteration, it adds a new node \(v_c\) into the anchor node set by adding some non-conflicting edges between \(v_c\) and \(A\) into the currently selected edge set. For each edge \(e\in E_{\rm VTOG}\), we define its discrete derivative with respect to \(\ell\) as \[w(e):=\ell(C\cup\{e\})-\ell(C),\]

which denotes the latency saving of merely adding edge \(e\) to \(C\). The greedy strategy picks the node \(v_c\) that maximize the sum of latency savings across its incident non-conflicting edges in the current iteration, then updates the sets \(A\) and \(C\). Subsequently, since creating a new virtual tensor updates \(C\), the algorithm recalculates \(w(e)\) for each edge \(e\) pointing to \(v_c\) by profiling. If the maximum latency saving is negative or all the nodes are added to \(A\), the algorithm terminates.

5.2.0.1 Complexity analysis.

 10 executes a maximum of \(O(|V_{\rm VTOG}|)\) iterations. In each iteration, it performs at most \(O(|E_{\rm VTOG}|+|V_{\rm VTOG}|)\) calculations. Based on the VTOG definition 6 and its construction algorithm 8, \(|E_{\rm VTOG}|\) is linear with \(|V|\) and \(|V_{\rm VTOG}|<|E|\). In typical DNN computation graphs, \(O(|V|)=O(|E|)\), so \(O(|V_{\rm VTOG}|)=O(|E_{\rm VTOG}|)=O(|V|)\). Consequently, the complexity of VTC’s data movement elimination algorithm is the product of the maximum number of iterations and the cost per iteration, i.e. \(O\left(|V|^2\right)\).

6 Implementation↩︎

We implemented VTC on top of Triton [7] and TorchInductor [3]. We first added virtual tensor support in Triton by overloading Triton’s global memory I/O functions, tl.load and tl.store, enabling them to handle virtual tensor I/O operations. Mapping function is represented in the form of ?? . Our overloaded operators consume base physical tensor pointers along with mapping functions that determine how virtual tensors map to physical memory. Since all indirect memory access patterns through virtual tensors are fully known at compile time, we implemented overloaded tl.load and tl.store by generating Python code specialized for each virtual tensor I/O, avoiding most of the indirection overheads.

We integrated VTC’s data movement eliminator into TorchInductor through two passes: an analysis pass and a transformation pass. The analysis pass automatically generates a virtual tensor strategy with [algo:vtog,algo:greedy]. This entails identifying which physical tensors to promote to virtual tensors as well as computing mapping functions of promoted virtual tensors to their physical indices. The transformation pass then operates in two stages: (1) It mutates TorchInductor’s IR nodes with virtual tensor awareness and removes any unnecessary data-movement operators. We modified TorchInductor’s lowering strategy to correctly lower our overloaded ops.load and ops.store operators to Triton. (2) Once lowered to Triton, we generate the specialized Python code using the method mentioned earlier. The resulting Triton code is then compiled to produce an optimized hardware executable.

VTC’s automatic virtual tensor construction employs a polynomial-time greedy algorithm rather than exhaustively searching in an exponential solution space. During compilation, profiling (line [line:profile1] and [line:profile2] in 10) dominates the time cost. Since VTC only modifies data movements in stage 1 and 3, no autotuning is needed for each kernel. The profiling overhead is less than 10 seconds per configuration, and the end-to-end compilation requires only a polynomial number of profiling runs, resulting in total compilation time under 10 minutes for all the models we have tested, which is on par with the compilation time of ML compilers [19]–[21].

As VTC focuses solely on optimizing data movement and does not alter the computation logic, it maintains end-to-end numeric equivalence with the original compiler, ensuring zero precision loss.

7 Evaluation↩︎

7.1 Experimental Setup↩︎

7.1.0.1 Platform.

 We conduct our evaluation on an A100 server and a H100 server. The A100 server is equipped with an Intel Xeon Platinum 8358 CPU and an NVIDIA A100 80GB PCIe GPU. The H100 server is equipped with an AMD EPYC 9454 Processor and an NVIDIA H100 NVL GPU. Both servers run Ubuntu 22.04, NVIDIA driver version 535.183.06 and CUDA version 12.1.

7.1.0.2 Workloads.

 We use five real-world DNNs from various domains to evaluate VTC. Llama 3 [9] and Gemma 2 [22] are transformer-based LLMs with grouped-query attention and local-global attention, respectively. EfficientViT [23] is a vision transformer backbone for high-resolution dense prediction with linear attention. YOLOv11 [24] is a powerful and efficient convolutional neural network (CNN) designed for a broad range of vision applications. ShuffleNet [25] is a CNN optimized for mobile devices. ¿tbl:tab:workloads? provides detailed configurations for each model. We use PyTorch implementation for all the models. Except ShuffleNet, all the implementation is from official repositories. Following the common practice in previous ML compiler research [26]–[28], we evaluate all models with batch size 1 and 16.

7.2 End-to-end Performance↩︎

We first evaluate the end-to-end inference latency on a single A100 or H100 GPU and compare VTC with PyTorch 2.6.0 [3] (with torch.compile enabled), ONNX Runtime 1.21.1 [29], XLA⁷ [4] and TensorRT 10 [11]. 11 presents the end-to-end inference latency results with different batch sizes on A100 and H100, respectively. VTC achieves up to \(1.93\times\) speed-up (\(1.28\times\) on average) over the best of the four baselines (mostly TensorRT or PyTorch) in each benchmark.

During the evaluation, TensorRT successfully identifies the attention pattern for all Transformer-based models (Llama, Gemma and EfficientViT) and triggered highly optimized implementations for the entire model. However, VTC still demonstrates an average speed-up of \(1.36\times\). This is notably higher than the average speed-up of \(1.15\times\) observed for the CNN models (YOLOv11 and ShuffleNet), indicating that VTC effectively capitalizes on optimization opportunities inherent in newer model architectures like Transformers.

VTC achieves even greater speedup on H100 GPUs as compared to A100 GPUs for 7 cases out of the 10 benchmark cases (covering 5 models at 2 different batch sizes). This trend suggests that the optimization techniques employed by VTC are particularly effective on newer hardware accelerators, where the gap between compute power and memory bandwidth is getting larger (see 1).

7.3 Memory Footprint↩︎

We also measure the inference GPU memory footprint after VTC’s optimization compared to the PyTorch baseline (as VTC is built upon TorchInductor) using the PyTorch’s built-in memory monitor. As shown in 1, even when optimizing for maximized latency reduction, VTC still achieves up to 60% peak memory saving (17.5% on average). This demonstrates that virtual tensor optimization inherently saves memory – it avoids storing full physical tensor data by representing virtual tensors with tensor pointers and mapping functions, and VTC handles virtual tensor I/O through code generation with no extra memory overhead. Analyzing the results in detail, we find that VTC successfully makes many large intermediate tensors virtual (e.g., the output of ScatterND in 2), thereby reducing peak memory usage. Moreover, this highlights VTC’s potential for more aggressive memory saving with relaxed latency constraints by reconfiguring the function \(\ell\) in 5.2 to measure memory reduction, offering greater benefits for model training and deployment on memory-constrained hardware.

7.4 Latency Breakdown Analysis↩︎

Since TensorRT is the fastest baseline in most cases, we compare the latency breakdown of data movement operators and computational operators for VTC and TensorRT on A100. The results in 12 reveal a substantial reduction in the latency proportion from data movement after VTC’s optimization. In 7 out of 10 evaluated cases, VTC eliminates data movement operators entirely. Moreover, the analysis highlights that data movement operators account for a larger proportion of the overall latency in Transformer-based models compared to CNNs, enabling VTC to get better speed-ups.

However, in most of the benchmarks, VTC’s latency on computational operators is higher than TensorRT’s. There are two possible reasons: (1) VTC’s computational kernels may perform virtual tensor I/O in non-contiguous physical addresses, introducing extra overhead in computational kernels (e.g., the MatMul kernel in frame 1 of 2). (2) VTC is built upon Triton and TorchInductor, while TensorRT has more highly optimized implementations of computational kernels. This indicates that the performance gains from eliminating data movement significantly outweigh the slowdown incurred in computational operators.

7.5 Comparison with vLLM↩︎

As 2 shows, VTC eliminates all data movement operators between QKV projection and attention in LLM decoder layer. Compared with existing compilers such as TensorRT and TorchInductor, the subsequent computational kernel FlashDecoding can launch directly after the previous computational kernel MatMul, improving the utilization rate of compute units without stalling on data movement operators.

To demonstrate the effectiveness of this optimization over a stronger LLM inference baseline beyond compilers, we compare VTC with a state-of-the-art LLM serving framework vLLM V1 [30] on Llama 3 8B inference. Since vLLM is PyTorch-based, we can seamlessly integrate VTC into it. We use similar settings (input sequence length 4096 and BF16 precision)⁸ for evaluation on A100 and H100 GPUs.

2 shows that vLLM has manually optimized many data movement operators from frame 2 of 2 – the KV Cache Update row in 2 contributes to a small fraction of latency compared to compiler baselines. This poses a great challenge for VTC to compete, as vLLM is heavily optimized by domain experts on specific LLMs, while VTC is an automated compiler for general ML models. Nevertheless, VTC still identifies a redundant data movement involving the temporary tensor that stores QKV projection results, detailed in 3. This optimization yields a 1.043x speed-up on the 2 computation graph and 1.011x end-to-end speed-up over vLLM on A100. However, no speed-up is observed on H100 since vLLM’s cuBLAS MatMul kernel (QKV projection) significantly outperforms VTC’s Triton kernel, and VTC’s profile-guided virtual tensor construction automatically skips negative optimizations. Forcing the optimization from 3 on H100 causes an 8% performance degradation. We believe this gap could be bridged by integrating VTC into more optimized hardware-specific kernel libraries like CUTLASS in the future.

7.6 Case Studies↩︎

To gain a deeper understanding of how VTC eliminates data movement operators, we study some optimization cases found by VTC in detail.

7.6.0.1 EfficientViT.

 13 illustrates VTC’s generated kernels and the points-to graph for EfficientViT attention block with batch size 16. All 5 kernels are computational kernels, and VTC eliminates data movement operators between these kernels by employing a virtual tensor strategy shown in the points-to graph. Among the tensors, only a, e, f, and j are physically stored. \(k_2\) and \(k_3\) directly read their inputs from the output tensor a of \(k_1\). \(k_4\) reads its inputs from the output tensor f of \(k_3\) and writes to the input tensor j of \(k_5\). VTC achieves only modest speed-up (\(\sim 1.1\times\)) and memory saving (\(\sim 4.4\%\)) on this model because most data movement elimination (e.g., optimizations related to f, g, h, i and j) can already be achieved by operator fusion in the baselines.

7.6.0.2 YOLOv11.

 [fig:yolo] shows the computation graph of C3K2 Block, a recurring component in YOLOv11. This block contains two data movement operators, Split and Concat. Traditional frameworks like TensorRT and PyTorch execute Split in an individual kernel and Concat in another kernel fused with Swish and Add (taking b, c, d as inputs and Y as the output). In contrast, VTC eliminates both data movement operators by making a, b, c and e virtual tensors of Y. Although Split and Concat are not fully contiguous at batch size 16 (5), each contiguous chunk (\(16\times 160\times 160\) elements) is large enough to read and write in a coalesced manner, preserving high memory I/O bandwidth. Moreover, these two data movement operators constitute the primary memory footprint bottlenecks, as they operate on large tensors (a, b, c, e and Y). After VTC’s optimization, all these tensors except Y become virtual, eliminating the need for physical memory allocation and reducing peak memory consumption.

8 Related Work↩︎

8.0.0.1 Virtual tensors.

 VTensor [31] is a programming framework which decouples tensor layouts from the programming interface, enabling developers to write layout-agnostic operations which significantly reduces lines of code. Separately, vLLM [30] and vTensor [32] are virtual memory management systems specifically designed for LLM serving. Unlike these systems, VTC optimizes general DNNs including LLMs with a fundamentally different virtual tensor notation.

8.0.0.2 Layout optimizations.

 TensorRT [11] explores various layout options for 4D tensors and automatically applies the necessary layout transformations to determine the optimal layout. SmartMem [13] studies layout transformation elimination based on predefined rules and develops efficient memory layouts for 2.5D memory on mobile devices. In contrast to these layout optimizations, VTC proposes a more general optimization for data movements. Instead of focusing solely on layout operators (mostly Reshape and Transpose), VTC’s optimization can be applied to all data movement operators.

Furthermore, VTC’s points-to graph construction algorithm automatically explore virtual tensor creation strategies, thereby expanding the optimization space beyond rule-based approaches.

8.0.0.3 Tensor graph optimizations.

 TensorFlow [1] and TensorRT [11] leverage manually designed rules to perform graph-level transformations. Recently, automated graph optimizers [27], [28], [33], [34] have emerged to automatically generate subgraph substitutions.

VTC’s optimization is orthogonal to these graph optimizations and can be used as an independent optimization pass after graph optimization in ML compilers.

8.0.0.4 Kernel orchestration

is the process of mapping a computation graph to hardware-specific kernels [35]. This problem is commonly addressed by operator fusion [1], [5], [11], [14], which fuses multiple operators into a single kernel, thereby eliminating data movements of intermediate results.

However, this type of work cannot optimize data movement operators that are hard to fuse with other operators. VTC takes a computation graph after operator fusion as input and further eliminates data movements between fused kernels.

8.0.0.5 Kernel generation.

 Numerous ML compilers [5], [26], [36] automatically generate hardware-specific kernels for DNN computation based on Halide’s algorithm-schedule separation [37]. Rammer [38] and Hidet [39] approach kernel generation by utilizing a task abstraction of workloads. Mirage [40] employs superoptimization techniques on a multi-level compute hierarchy to generate efficient GPU kernels. As VTC only modifies the global memory I/O (stages 1 and 3) of a hardware-specific kernel, it can be easily integrated with virtual tensor support in these kernel generation frameworks.

8.0.0.6 Hand-optimized kernels

crafted by domain experts can achieve state-of-the-art performance on specific DNN workloads. For instance, FlashAttention [10], [41] and FlashInfer [42] design highly optimized GPU kernels for self-attention calculation. VTC can also incorporate these manually optimized kernels by adding virtual tensor support.

9 Future Work↩︎

While our current virtual tensor construction algorithm provides performance guarantees through iterative profiling, this approach can incur substantial compilation overhead for time-sensitive applications. The VTOG formulation we present opens opportunities for developing more efficient algorithms that can compute high-quality virtual tensor strategies with reduced profiling requirements.

Our current implementation is based on Triton. However, Triton exhibits known performance limitations on newer GPU architectures, particularly NVIDIA Blackwell GPUs, where domain-specific languages such as CUTLASS demonstrate superior performance. Extending virtual tensor support to alternative DSLs and compilation frameworks represents a promising direction that could unlock further performance gains across diverse hardware platforms.

10 Conclusion↩︎

Data movement optimizations have become extremely important in tensor compilers to achieve good performance. VTC is the first compilation framework to eliminate all unnecessary data movement that span across all data movement operators in a given tensor computational graph. We introduce the concept of a virtual tensor that allows VTC to track data movement between different physical tensors without instantiating intermediate tensors. VTC’s data movement elimination algorithm allows us to achieve superior performance on a number of different DNN models including appreciable gains on large language model inference.

Acknowledgments↩︎

We thank Zhihao Jia, Hao Guo, Yuhao Ge and Yueying Li for their feedback and help on this work. This work was supported by PRISM and ACE, two of the seven centers in JUMP 2.0, a Semiconductor Research Corporation (SRC) program sponsored by DARPA, by NSF under the grant CCF-2316233, and by generous gifts from Qualcomm.

References↩︎