GPUOS: A GPU Operating System Primitive for Transparent Operation Fusion

TL;DR

GPUOS is a GPU runtime system that reduces kernel launch overhead by maintaining a persistent kernel with runtime operator injection, significantly accelerating small tensor operations in deep learning workloads.

Contribution

GPUOS introduces a novel persistent kernel architecture with runtime operator injection, enabling efficient execution of diverse small tensor operations without kernel restarts.

Findings

Achieves up to 15.3x speedup over standard PyTorch on small operation workloads.

Supports arbitrary tensor shapes, data types, and broadcasting.

Improves GPU utilization in micro-batched inference and attention workloads.

Abstract

Modern deep learning workloads often consist of many small tensor operations, especially in inference, attention, and micro-batched training. In these settings, kernel launch overhead can become a major bottleneck, sometimes exceeding the actual computation time. We present GPUOS, a GPU runtime JIT system that reduces launch overhead using a persistent kernel architecture with runtime operator injection. GPUOS runs a single long-lived GPU kernel that continuously processes tasks from a host-managed work queue, eliminating repeated kernel launches. To support diverse operations, GPUOS uses NVIDIA NVRTC to just-in-time compile operators at runtime and inject them into the running kernel through device function pointer tables. This design enables operator updates without restarting the kernel or recompiling the system. GPUOS introduces four key ideas: (1) a persistent worker kernel with atomic task queues, (2) a runtime operator injection mechanism based on NVRTC and relocatable device code, (3) a dual-slot aliasing scheme for safe concurrent operator updates, and (4) transparent PyTorch integration through TorchDispatch that batches micro-operations into unified submissions. The system supports arbitrary tensor shapes, strides, data types, and broadcasting through a generic tensor abstraction. Experiments show that GPUOS achieves up to 15.3x speedup over standard PyTorch on workloads dominated by small operations, including micro-batched inference and attention patterns. GPUOS improves utilization while remaining compatible with the PyTorch ecosystem.