PyGPUkit 0.2.3

A lightweight GPU runtime for Python with Rust-powered scheduler, NVRTC JIT compilation, and NumPy-like API

PyGPUkit — Lightweight GPU Runtime for Python

A minimal, modular GPU runtime with Rust-powered scheduler, NVRTC JIT compilation, and a clean NumPy-like API.

---

Overview

PyGPUkit is a lightweight GPU runtime for Python that provides:

• Rust-powered scheduler with admission control, QoS, and resource partitioning
• NVRTC-based JIT kernel compilation
• A NumPy-like GPUArray type
• Kubernetes-inspired GPU scheduling (bandwidth + memory guarantees)
• Extensible operator set (add/mul/matmul, custom kernels)
• Minimal dependencies and embeddable runtime

PyGPUkit aims to be the "micro-runtime for GPU computing": small, fast, and ideal for research, inference tooling, DSP, and real-time systems.

---

Opening Paragraph (Goal Statement)

PyGPUkit aims to simplify GPU development by reducing dependency on complex CUDA Toolkit installations and fragile GPU environments. Its goal is to make GPU programming feel like using a standard Python library: installable via pip with minimal setup. PyGPUkit provides high-performance GPU kernels, memory management, and scheduling through a NumPy-like API and a Kubernetes-inspired resource model, allowing developers to use GPUs explicitly, predictably, and productively.

Note: PyGPUkit currently requires CUDA drivers and NVRTC. It is NOT a PyTorch/CuPy replacement—it's a lightweight runtime for custom GPU workloads, research, and real-time systems where full ML frameworks are overkill.

---

v0.2.3 Features (NEW)

TF32 TensorCore GEMM

Feature	Description
PTX mma.sync	Direct TensorCore access via inline PTX assembly
cp.async Pipeline	Double-buffered async memory transfers
TF32 Precision	19-bit mantissa (vs FP32's 23-bit), ~0.1% per-op error
SM 80+ Required	Ampere architecture (RTX 30XX+) required

Benchmark Comparison (RTX 3090 Ti, 8192×8192×8192)

Library	FP32	TF32	Notes
NumPy (OpenBLAS)	~0.8 TFLOPS	—	CPU baseline
cuBLAS	~21 TFLOPS	~59 TFLOPS	NVIDIA benchmark
PyGPUkit	18 TFLOPS (86%)	27 TFLOPS (46%)	Custom kernels

FP32 is near cuBLAS level. TF32 optimization ongoing.

PyGPUkit Performance by Size

Matrix Size	FP32	TF32
2048×2048	7.6 TFLOPS	10.2 TFLOPS
4096×4096	13.2 TFLOPS	19.5 TFLOPS
8192×8192	18.2 TFLOPS	27.5 TFLOPS

Core Infrastructure (Rust)

Feature	Description
Memory Pool	LRU eviction, size-class free lists
Scheduler	Priority queue, memory reservation
Transfer Engine	Separate H2D/D2H streams, priority
Kernel Dispatch	Per-stream limits, lifecycle tracking

Advanced Features (Rust)

Feature	Description
Admission Control	Deterministic admission, quota enforcement
QoS Policy	Guaranteed/Burstable/BestEffort tiers
Kernel Pacing	Bandwidth-based throttling per stream
Micro-Slicing	Kernel splitting, round-robin fairness
Pinned Memory	Page-locked host memory with pooling
Kernel Cache	PTX caching, LRU eviction, TTL
GPU Partitioning	Resource isolation, multi-tenant support

---

Features

• Lightweight — smaller footprint than PyTorch/CuPy (not a replacement)
• Modular — runtime / memory / scheduler / JIT / ops
• Rust Backend — memory pool, scheduler, dispatch in Rust
• GPUArray with NumPy interop
• NVRTC JIT for CUDA kernels
• Advanced Scheduler with memory & bandwidth guarantees
• 106 Rust tests for core components

---

Installation

pip install pygpukit

From source:

git clone https://github.com/m96-chan/PyGPUkit
cd PyGPUkit
pip install -e .

Requirements:

• Python 3.10+
• CUDA 11+
• NVRTC available
• NVIDIA GPU

Supported GPUs:

• RTX 30XX series (Ampere) and above
• Performance tuning is optimized for GPUs with large L2 cache (6MB+)
• Older GPUs (RTX 20XX, GTX 10XX, etc.) are NOT tuned and may have suboptimal performance

---

Project Goals

1. Provide the smallest usable GPU runtime for Python
2. Expose GPU scheduling (bandwidth, memory, partitioning)
3. Make writing custom GPU kernels easy
4. Serve as a building block for inference engines, DSP systems, and real-time workloads

---

Usage Examples

Allocate Arrays

import pygpukit as gp

x = gp.zeros((1024, 1024), dtype="float32")
y = gp.ones((1024, 1024), dtype="float32")

Basic Operations

z = gp.add(x, y)
w = gp.matmul(x, y)

CPU <-> GPU Transfer

arr = z.to_numpy()
garr = gp.from_numpy(arr)

Custom NVRTC Kernel

extern "C" __global__
void scale(float* x, float factor, int n) {
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    if (idx < n) x[idx] *= factor;
}

kernel = gp.jit(src, func="scale")
kernel(x, factor=0.5, n=x.size)

Rust Scheduler (v0.2)

import _pygpukit_rust as rust

# Memory Pool with LRU eviction
pool = rust.MemoryPool(quota=100 * 1024 * 1024, enable_eviction=True)
block = pool.allocate(4096)

# QoS-aware task scheduling
evaluator = rust.QosPolicyEvaluator(total_memory=8*1024**3, total_bandwidth=1.0)
task = rust.QosTaskMeta.guaranteed("task-1", "Critical Task", 256*1024*1024)
result = evaluator.evaluate(task)

# GPU Partitioning
manager = rust.PartitionManager(rust.PartitionConfig(total_memory=8*1024**3))
manager.create_partition("inference", "Inference",
    rust.PartitionLimits().memory(4*1024**3).compute(0.5))

---

Scheduler — Kubernetes-Inspired GPU Orchestration

PyGPUkit includes an experimental scheduler that treats a single GPU as a multi-tenant compute node, similar to how Kubernetes orchestrates CPU workloads. The goal is to provide resource isolation, guarantees, and fair sharing across multiple GPU tasks.

Core Capabilities

---

1. GPU Memory Reservation

Tasks may request a guaranteed block of GPU memory.

• Hard guarantees -> task is rejected if memory cannot be allocated
• Soft guarantees -> best-effort allocation
• Overcommit strategies (evict to host when pressure is high)
• Reclaim policies (LRU GPUArray eviction)

Example:

task = scheduler.submit(
    fn,
    memory="512MB",
)

---

2. GPU Bandwidth Guarantees / Throttling

Tasks may request a specific percentage of GPU compute bandwidth.

Bandwidth control is implemented via:

• Stream priority
• Kernel pacing (launch intervals)
• Micro-slicing large kernels
• Cooperative time-quantized scheduling
• Persistent dispatcher kernels (planned)

Example:

task = scheduler.submit(
    fn,
    bandwidth=0.20,   # 20% GPU compute share
)

---

3. Logical GPU Partitioning

PyGPUkit implements software-defined GPU slicing, similar in spirit to Kubernetes device plugin resource partitioning.

Slices may define:

• Memory quota
• Bandwidth share
• Stream priority band
• Isolation level

Useful for:

• Multi-tenant inference servers
• Real-time audio/DSP workloads
• Background/foreground GPU task separation

---

4. Scheduling Policies

The scheduler supports multiple policies:

• Guaranteed — exclusive reservation, strict QoS
• Burstable — partial guarantees, opportunistic bandwidth
• BestEffort — uses leftover GPU cycles
• Priority scheduling
• Deadline scheduling (planned)
• Weighted fair sharing

Example:

task = scheduler.submit(
    fn,
    policy="guaranteed",
    memory="1GB",
    bandwidth=0.10,
)

---

5. Admission Control

Before executing a task, the scheduler performs:

• Resource validation
• Quota check
• QoS matching
• Scheduling feasibility

Results in:

• admitted
• queued
• rejected

---

6. Monitoring & Introspection

PyGPUkit exposes live metrics:

• Memory usage per task
• SM occupancy and GPU utilization
• Throttling / pacing logs
• Queue position / execution state
• Reclaim/eviction count

Example:

stats = scheduler.stats(task_id)

---

7. Soft Isolation Model

While not OS-level isolation, each GPU task is provided:

• Dedicated stream groups
• Guaranteed memory pools
• Kernel pacing to enforce bandwidth
• Optional sandboxed GPUArray region

This provides practical multi-tenant safety without MIG/MPS.

---

Project Structure

PyGPUkit/
  src/pygpukit/    # Python API (NumPy-compatible)
  native/          # C++ backend (CUDA Driver/Runtime/NVRTC)
  rust/            # Rust backend (memory pool, scheduler, dispatch)
    pygpukit-core/   # Pure Rust core logic
    pygpukit-python/ # PyO3 bindings
  examples/        # Demo scripts
  tests/           # Test suite

---

Roadmap

v0.1 — v0.2.3 (Released)

Version	Highlights
v0.1	GPUArray, NVRTC JIT, add/mul/matmul, wheels
v0.2.0	Rust scheduler (QoS, admission control, partitioning), memory pool (LRU), kernel cache, 106 Rust tests
v0.2.1	API stabilization, error propagation
v0.2.2	Ampere SGEMM (cp.async, float4), 18 TFLOPS FP32
v0.2.3	TF32 TensorCore (PTX mma.sync), 27.5 TFLOPS

v0.2.4 — Benchmark & Reliability Phase

• Actual PyTorch/NumPy comparison benchmarks
• Kernel cache LRU completion
• Driver-only mode stabilization
• Windows/Linux full support
• Large GPU memory test (16GB continuous alloc/free)

v0.2.5 — Distributed Phase

• Multi-GPU Detection
• NCCL / peer-to-peer preliminary support
• Scheduler multi-device support

v0.2.6 — Pre-v0.3 Finalization

• Full API review
• Backward compatibility policy
• JIT build options, safety measures, env vars cleanup
• Documentation

v0.3 (Planned)

• Triton optional backend
• Advanced ops (softmax, layernorm)
• Inference-oriented plugin system
• MPS/MIG integration

---

Contributing

Contributions and discussions are welcome! Please open Issues for feature requests, bugs, or design proposals.

---

License

MIT License

---

Acknowledgements

Inspired by:

• CUDA Runtime
• NVRTC
• PyCUDA
• CuPy
• Triton

PyGPUkit aims to fill the gap for a tiny, embeddable GPU runtime for Python.