b10-kernel

Baseten Kernel Library - High-performance GPU kernels for AI inference workloads.

Installation

From PyPI

pip install b10-kernel

Requirements:

• Python >= 3.12
• CUDA-compatible GPU and drivers
• PyTorch >= 2.8.0 with CUDA support

From Source

git clone <repository>
cd mp/kernels/b10-kernel
pip install -e .

For Development

# Install with test dependencies
pip install -e .[test]

# Install with all development dependencies  
pip install -e .[dev]

Development guide

• Build the library from source

make build
make rebuild

• Run unit tests

make test

• Format code

make format

Kernel Development Workflow

Steps to add a new kernel:

• Implement the kernel in csrc
• Expose the interface in include/b10_kernel_ops.h
• Create torch extension in csrc/common_extension.cc
• Update CMakeLists.txt to include new CUDA source
• Expose Python interface in python/b10_kernel/xxx.py and python/b10_kernel/__init__.py
• Add unit test for the kernel in test/test_xxx.py
• Add benchmark script for the kernel in benchmark/bench_xxx.py
• Format code with make format

PyPI Release

# update to expected version
make update_version VERSION=0.1.x
make build
make rename_wheel
make upload_wheel