Kitsune: Enabling Dataflow Execution on GPUs

TL;DR

Kitsune addresses the inefficiencies of GPU bulk-synchronous execution for DL workloads by enabling dataflow execution across operators via inter-CTA queues and a heterogeneity-aware grid scheduler, implemented as a PyTorch Dynamo-based compiler. The approach lowers memory traffic and increases parallelism without a full hardware redesign, achieving up to 2.4x speedups and substantial off-chip traffic reductions across five DL workloads for both inference and training. Key contributions include a detailed characterization of DL workloads, the two primitive designs (inter-CTA ring queues and spatial pipelines), the Kitsune compiler with subgraph selection, pipeline design, and load balancing, and extensive evaluation against vertical fusion on a validated A100-class GPU. The work demonstrates practical improvements in SM/DRAM utilization and end-to-end performance, showing that modest architectural tweaks can unlock significant dataflow-friendly execution on existing 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.

Figures (14)

• Figure 1: (a) Overview of GPU organization, (b) example DL graph, and (c) stylized comparison of execution techniques. In (c), TensorCore and SIMT resources of the GPU are depicted separately.
• Figure 2: Visualization of the difference between Kitsune and Vertical Fusion for (a) an MLP with a large hidden dimension, (b) a reduction operation, and (c) an operator that sends intermediate to multiple consumers.
• Figure 3: Application runtime spent in different combinations of measured SM and DRAM utilization. Low utilization means less than 33% of peak.
• Figure 4: Queue design. Note: release routines are not shown for space reasons. They involve simple atomicAdd calls to update synchronization metadata and a CTA barrier with __syncthreads().
• Figure 5: Performance of GPU atomics.