Kitsune: Enabling Dataflow Execution on GPUs

Abstract: State of art DL models are growing in size and complexity, with many modern models also increasing in heterogeneity of behavior. GPUs are still the dominant platform for DL applications, relying on a bulk-synchronous execution model which has many drawbacks and is ill-suited for the graph structure of DL applications. Many industry and academic works attempt to overcome these by employing vertical fusion but this approach still fails to realize three untapped opportunities: (1) the fact that many resources on the GPU are idle while only one operator executes due to temporal multiplexing of the SM; (2) lower energy from more intelligent on-chip data-movement which lends to higher performance in a power-provisioned environment. (3) inability to exploit hidden or reduction dimensions as a source of parallelism to ease pressure on batch size. This paper explores relatively uncharted territory, answering the following key question: Can modest adjustments to the current GPU architecture enable efficient dataflow execution, thereby circumventing the constraints of vertical fusion without necessitating a clean-slate architecture design. We develop Kitsune -- a set of primitives that enable dataflow execution on GPUs and an end-to-end compiler based on PyTorch Dynamo. Across 5 challenge applications, Kitsune can provide 1.3$\times$-2.3$\times$ and 1.1$\times$-2.4$\times$ performance improvement as well as 41%-98% and 16%-42% off-chip traffic reduction for inference and training, respectively.

Kitsune: Enabling Dataflow Execution on GPUs

The paper "Kitsune: Enabling Dataflow Execution on GPUs" presents a significant exploration into optimizing the execution model of deep learning (DL) tasks on Graphics Processing Units (GPUs). As DL models increase in complexity and size, existing approaches, predominantly relying on bulk-synchronous processing (BSP), expose several inefficiencies due to their inability to fully exploit the concurrent and heterogeneous nature of modern GPU resources. Kitsune introduces a dataflow execution model that seeks to address these limitations through modest additions to current GPU architectures, avoiding the need for radical redesigns.

The paper identifies primary inefficiencies in the existing model, notably idle GPU resources, excessive energy consumption due to suboptimal data movement, and limited parallelism exploitation in certain computational dimensions. By introducing Kitsune, the authors propose a set of primitives supporting dataflow execution on GPUs which significantly minimizes these inefficiencies without diverging drastically from current architectures.

Key Innovations:

1. Software-Hardware Integration: Kitsune integrates both software and hardware elements to facilitate a dataflow execution model. This includes:

   • A software-only ring queue leveraging the GPU's L2 cache for efficient inter-thread communication.
   • An enhanced GPU grid scheduler that can allocate and execute diverse operations in concert, rather than strictly sequentially.

2. Compiler Support: The development of Kitsune involved creating an end-to-end compiler, leveraging PyTorch's Dynamo interface. The compiler automatically lowers DL applications for dataflow execution.

3. Performance Analysis: Kitsune demonstrates substantial performance improvements, achieving 1.3x to 2.4x speedups and significant reductions in off-chip memory traffic (up to 98% for inference and 42% for training tasks).

Detailed Contributions:

• Architectural Modifications: By proposing a minimal change to the grid scheduler, Kitsune allows CTAs (Cooperative Thread Arrays) to co-reside and form spatial pipelines. This allows better use of GPU resources by running heterogeneous operations in parallel—exploiting both SIMT and Tensor Core resource types simultaneously.

• Queue Design for Data Movement: The ring queues ensure efficient data transfer paths between threads, maintaining high throughput while minimizing off-chip data movement, a key contributor to energy wastage and latency.

• Compilation and Graph Selection: Kitsune's compilation strategy involves selecting subgraphs within DL applications suitable for dataflow execution. The system identifies patterns that benefit most from this model, maximizing fusion opportunities beyond those possible with extant vertical fusion techniques.

Implications and Future Research Directions:

The paper points towards a substantial improvement not just in computational throughput but also in reducing energy footprints of DL tasks on GPUs. Kitsune's approach to dataflow execution aligns the execution model closely with the graph nature of DL applications, offering a significant departure from conventional BSP strategies used in GPUs.

For future research, this paper opens up directions in further optimizing these primitives for more diverse DL models and potentially extending their applicability beyond GPUs to other heterogeneous computing environments, such as TPUs and FPGAs. The theoretical and practical implications of Kitsune suggest a promising landscape in GPU computing where models are optimized not just for performance but also for energy efficiency—a crucial aspect as neural network workloads continue to grow.

As DL continues its pervasive influence across various fields, architectures like Kitsune provide a crucial stepping stone in bridging the scalability and efficiency challenges faced by the industry. Ultimately, the insights and demonstrations provided in the paper reflect a maturation in understanding how best to harness the full capabilities of modern GPUs for DL applications.