oxs 0.1.0

OXS: Optimized X-Scale SOTA Engine

OXS: Optimized X-Scale SOTA Engine

OXS is a high-performance Neural-Symbolic AI Operating System kernel designed for extreme efficiency on consumer and server-grade hardware. It features a custom CUDA backend (maxwell_core) optimized for Pascal, Volta, Ampere, and Hopper architectures.

Features

• SOTA CUDA Kernels:
  • Tiled FP32 execution for Pascal (GTX 1050 Ti+).
  • Tensor Core (WMMA) support for newer GPUs.
  • Fused Backward Pass (Weight Gradients calculated on GPU).
• Quantization: Native support for 2-bit Packed Weights (Ternary {-1, 0, 1}).
• Pythonic API: PyTorch-like interface (oxs.Linear, oxs.SGD).

Installation

Prerequisites

• Python 3.8+
• NVIDIA CUDA Toolkit 11.0+ (Tested on 12.5)
• Visual Studio 2019/2022 (Windows) or GCC (Linux) with C++17 support.
• CMake 3.18+

Install from Source

git clone https://github.com/pantheon/oxs.git
cd oxs
pip install .

Quick Start

import oxs
import numpy as np

# 1. Define Model
layer = oxs.Linear(in_features=1024, out_features=10)
optimizer = oxs.SGD([layer], lr=0.1)

# 2. Dummy Data
x = np.random.randn(32, 1024).astype(np.float32)
y_target = np.random.randn(32, 10).astype(np.float32)

# 3. Training Loop
for epoch in range(50):
    # Forward (Quantized on the fly)
    y_pred = layer.forward(x)
    
    # Loss (MSE)
    diff = y_pred - y_target
    loss = np.mean(diff**2)
    
    # Backward
    grad_loss = 2.0 * diff / x.size
    layer.backward(grad_loss)
    
    # Update
    optimizer.step()
    
    print(f"Epoch {epoch}: Loss {loss:.4f}")

License

MIT License. See LICENSE for details.