GPUOS: A GPU Operating System Primitive for Transparent Operation Fusion

GPUOS is a persistent kernel-based GPU runtime system that utilizes just-in-time compilation and operator injection to eliminate kernel launch overhead, achieving up to 15.3x speedup for deep learning workloads dominated by small tensor operations while maintaining seamless PyTorch integration.

The Gist

Imagine you are running a high-end pizza kitchen. In the past, your goal was to bake massive batches of 1,000 pizzas at once. To do this, you had a very efficient system: you'd shout an order to the oven, the oven would bake for a long time, and then you'd shout the next order. The time it took to shout the order (the "launch") was tiny compared to the time the pizza spent baking.

But today's world is different.

Now, imagine you are running a "micro-pizza" service. You have thousands of customers, and each one wants a tiny, custom slice of pizza right now.

• Customer A wants a slice of cheese.
• Customer B wants a slice of pepperoni.
• Customer C wants a slice of mushroom.
• Customer D wants a slice of cheese again.

In the old system, for every single slice, you would have to:

1. Stop what you are doing.
2. Walk to the oven (the GPU).
3. Yell the specific order ("Cheese!").
4. Wait for the oven to acknowledge you.
5. Walk back to your station.
6. Repeat this 100 times for one customer's order.

The Problem: The pizza takes 1 second to bake, but walking to the oven and yelling takes 5 seconds. You spend 95% of your time walking and shouting, and only 5% of your time actually cooking. The kitchen is chaotic, the customers are waiting, and you are exhausted.

This is exactly the problem modern AI (like the chatbots you talk to) faces. The AI has to perform thousands of tiny math calculations (like adding two numbers or checking a word) for every single word it generates. The computer's "brain" (the GPU) is incredibly fast at math, but the "manager" (the CPU) is too slow at handing out these tiny tasks. The time spent handing out the tasks is longer than the time spent doing the math.

Enter GPUOS: The "Always-On" Sous-Chef

The paper introduces GPUOS, which solves this by changing the management style entirely.

Instead of the manager walking to the oven for every single slice, GPUOS installs a permanent, super-fast Sous-Chef right inside the kitchen (on the GPU itself).

Here is how it works:

1. The Persistent Kitchen (The Persistent Kernel):
   Instead of turning the oven on and off for every order, the Sous-Chef sits right next to the oven, awake and ready, 24/7. They never leave. They are just waiting for work.

2. The Conveyor Belt (The Ring Buffer):
   The manager (CPU) doesn't walk to the oven. Instead, they just drop a tiny note on a conveyor belt right next to the Sous-Chef. The note says: "Add cheese to this slice."

   • Old Way: Manager walks 5 seconds to yell "Cheese!"
   • GPUOS Way: Manager drops a note in 0.0001 seconds.

3. The Sous-Chef's Speed:
   The Sous-Chef sees the note, grabs the slice, adds the cheese, and puts it on the plate. Then, they immediately look for the next note. Because they are already there, this whole process takes nanoseconds (billionths of a second).

4. The Magic Trick (Dynamic Operator Injection):
   What if a new customer wants a "Spicy Jalapeño" slice, and the Sous-Chef doesn't know how to make that yet?

   • Old Way: You'd have to fire the Sous-Chef, hire a new one, train them, and restart the whole kitchen.
   • GPUOS Way: The manager slides a tiny, pre-written recipe card under the Sous-Chef's door. The Sous-Chef instantly reads it, learns the new move, and starts making Jalapeño slices without stopping the line. This is called "JIT compilation" (Just-In-Time), but think of it as instant recipe updates.

Why is this a Big Deal?

The paper tested this on real AI tasks (like generating text in a chatbot).

• The Result: The system became 15 times faster for small tasks.
• The Analogy: It's like going from a delivery service where a truck drives to your house for every single letter you mail, to a system where the mailman is already standing on your porch, and you just slide the letter into their hand.

The Benefits for You (The User)

1. Faster Chatbots: When you talk to an AI, it won't have to "think" as hard about how to process your words. It can just process them instantly.
2. Less Waiting: The "lag" you feel when the AI pauses to think is reduced because the computer isn't wasting time walking back and forth.
3. Energy Saving: Because the computer isn't idling while waiting for instructions, it uses less electricity.

Summary

GPUOS is like realizing that for a high-speed kitchen, you don't need a manager running back and forth. You need a permanent, super-fast worker sitting right next to the stove, ready to grab the next task the millisecond it arrives. It bridges the gap between the slow human manager and the super-fast robot brain, making AI feel instant and responsive.

Technical Summary

1. Problem Statement

Modern deep learning workloads, particularly in inference scenarios, attention mechanisms, and micro-batched training, are shifting from large, regular batches to numerous small tensor operations.

• The Bottleneck: Traditional GPU computing models launch a separate kernel for every operation. This incurs significant CPU-GPU synchronization overhead (kernel launch latency), which typically ranges from 3 to 7 microseconds.
• The Impact: For small operations that execute in single-digit microseconds, the launch overhead can exceed the actual computation time by orders of magnitude. In micro-batched inference (e.g., token-by-token generation), a single request may involve hundreds of such operations, causing the cumulative launch overhead to dominate the total response time budget.
• Limitations of Existing Solutions:
  • CUDA Graphs: While effective for static, regular workloads, they struggle with dynamic control flow, variable tensor shapes, and frequent model updates, often requiring complex graph management or falling back to eager execution.
  • Static Fusion: Compilers (like TVM or TorchDynamo) struggle to fuse operations when shapes are dynamic or when new operators are introduced at runtime.

2. Methodology: The GPUOS Architecture

GPUOS introduces a persistent kernel architecture combined with runtime operator injection to eliminate per-operation launch overhead. Instead of launching a new kernel for every task, GPUOS deploys a single, long-lived kernel that continuously processes a work queue.

Core Components

1. Persistent Worker Kernel:

   • Launched once at process startup, this kernel occupies one thread block per Streaming Multiprocessor (SM) and never exits.
   • It runs an infinite loop where warps poll a device-visible ring buffer for new tasks.
   • Once a task is claimed atomically, the thread dispatches the operation and immediately returns to polling, eliminating the host-device handshake for every operation.

2. Lock-Free Ring Buffer (SPSC):

   • A Single-Producer-Single-Consumer queue residing in device-mapped memory.
   • Host Side: Writes task descriptors (Operator ID, tensor pointers, dimensions, flags) and commits them using store-release semantics.
   • Device Side: Warps poll the read cursor using load-acquire semantics to detect new work.
   • Descriptor Size: Compact (64–128 bytes), minimizing memory traffic.

3. Dynamic Operator Injection (The "OS" Primitive):

   • To handle the diversity of operations without recompiling the entire system, GPUOS uses NVIDIA NVRTC (Just-In-Time compilation) and the CUDA Driver API.
   • Workflow:
     1. A new operator (e.g., a custom attention variant) is compiled to PTX at runtime.
     2. The module is loaded, and the function symbol is resolved to a device function pointer.
     3. The pointer is written to an inactive slot in a device-resident Function Pointer Table.
     4. A version counter is atomically incremented with store-release semantics.
   • Safety: Device threads check the version counter at defined points in their polling loop. If the version changes, they atomically switch to the updated table. This ensures zero-downtime updates without suspending concurrent tasks.

4. Dual-Slot Aliasing Scheme:

   • Ensures safe updates by allowing the system to write to a new slot while threads are still reading the old version, preventing race conditions during operator injection.

5. PyTorch Integration:

   • Implemented via TorchDispatch, acting as a transparent plugin.
   • A dispatcher filter identifies "good candidates" for GPUOS (small tensors, element-wise ops, high-frequency loops) and redirects them to the ring buffer. Large operations (e.g., massive GEMMs) continue to use standard kernel launches.

3. Key Contributions

• Persistent Kernel Paradigm: Shifts scheduling from the host to the device, replacing microsecond-scale kernel launches with nanosecond-scale device function calls.
• Runtime Operator Injection: A novel mechanism to hot-swap GPU operators using NVRTC and function pointer tables, enabling dynamic model evolution without service interruption.
• Transparent Aggregation: Automatically aggregates micro-operations from PyTorch into batched submissions via TorchDispatch, requiring no code changes to existing models.
• Generic Tensor Abstraction: Supports arbitrary tensor shapes, strides, data types, and broadcasting semantics, making it applicable to diverse workloads.
• Safety & Observability: Includes audit logging, kill switches for misbehaving operators, and compatibility with standard profilers (Nsight).

4. Experimental Results

Evaluated on NVIDIA H100, RTX 5090, and GB10 GPUs, GPUOS demonstrates significant performance gains in micro-batch scenarios:

• Element-wise Operations: Achieved a 15.3× speedup over standard PyTorch (reducing latency from ~8µs to ~3.1µs per op).
• Attention Decoding: Achieved an 8.7× speedup by eliminating repeated launches for query-key computations and softmax.
• Mixed Pipelines: Achieved a 23.1× speedup in realistic inference pipelines by accelerating the "long tail" of small operations between large matrix multiplications.
• Energy Efficiency: Reduced energy consumption by 20–22% due to reduced idle time and elimination of thousands of launch-idle cycles per request.
• Compatibility: Successfully coexists with CUDA Graphs, MPS (Multi-Process Service), and MIG (Multi-Instance GPU), maintaining performance even in partitioned environments.

5. Significance

GPUOS addresses a fundamental architectural mismatch in modern GPU computing: the increasing frequency of small, dynamic operations versus the high fixed cost of kernel launches.

• Bridging Flexibility and Efficiency: It provides a practical pattern for building runtime systems that do not sacrifice operator flexibility (dynamic shapes, custom ops) for execution efficiency.
• Complementary to Existing Tech: It does not replace compilers or CUDA Graphs but fills the gap where they fail (dynamic control flow and frequent updates).
• Production Ready: The system is designed for production environments, offering robust error handling, security via template restriction, and seamless integration with the PyTorch ecosystem.

In summary, GPUOS redefines the GPU execution model for the era of micro-batch inference, transforming the GPU from a passive execution engine into an active, self-scheduling runtime that minimizes coordination costs.