SplitKernel

SplitKernel is an end-to-end GPU memory offloading framework for large language model inference. Instead of following the conventional prefetch-based design, where offloaded data is first moved back into local GPU HBM before computation, SplitKernel enables the GPU to directly access offloaded memory and stream data into shared memory for kernel execution.

The implementation in this repository includes CUDA kernels, PyTorch extension bindings, attention kernels, model integration utilities, benchmarks, and evaluation notebooks used to study direct-access memory offloading for LLM inference.

Overview

LLM inference is often limited by GPU memory capacity and bandwidth. Offloading weights or KV caches to a remote memory tier can increase the effective memory capacity, but existing systems commonly rely on prefetching data into GPU HBM. This introduces several inefficiencies:

• additional HBM traffic and capacity pressure;
• pipeline bubbles when data movement is not perfectly hidden;
• underutilization of aggregate bandwidth across local and remote memory tiers;
• read amplification when multiple consumers need overlapping data.

SplitKernel addresses these bottlenecks through direct-access offloading. The key idea is to allow GPU kernels to fetch offloaded weights and KV cache blocks directly from the remote tier into GPU shared memory, avoiding unnecessary staging in HBM.

Key Features

• Direct-access offloading for LLM inference workloads.
• CUDA/PyTorch extension for offloaded linear and runtime primitives.
• TMA-based data movement targeting Hopper-class GPUs.
• Attention with host/device KV cache partitioning.
• Model integration utilities for replacing selected linear layers and controlling placement.
• Benchmarks and notebooks for weight offloading, KV cache offloading, offloading-ratio analysis, and ablation studies.

Repository Layout

.
├── app/python/              # Model integration and placement utilities
├── benchmark/               # Benchmark entry points and scripts
├── eval/                    # Evaluation notebooks
├── include/                 # CUDA headers for runtime and task kernels
├── split_attention/         # Attention for splitted KV cache
├── python/offload/          # Python package for the offload runtime
├── src/                     # CUDA runtime and PyTorch binding sources
├── Makefile                 # Root build targets
└── setup.py                 # PyTorch CUDA extension build

Requirements

The code is designed for CUDA-enabled systems and currently targets Hopper-class GPUs by default:

• NVIDIA GPU with CUDA support;
• CUDA toolkit with nvcc;
• PyTorch with CUDA support;
• Python development environment with pip;
• transformers, numpy, and other benchmark-time Python dependencies.

The main runtime builds for sm_90a (GH200) by default. Pass arch=120a to make to build for sm_120a (RTX PRO 6000) instead. This also selects the matching SM count, buffer slots, and MMA path in include/task/config.cuh and include/task/gemv.cuh. The split_attention/setup.py -gencode setting still has to be changed by hand.

Installation

From the repository root directory, build and install the main offload runtime:

bash setup.sh
make pyext            # GH200 (sm_90a)
# make pyext arch=120a  # RTX PRO 6000 (sm_120a)

Then install the split attention extension:

cd opt_attention
pip install -e .

The root make pyext target first builds runtime.o and then installs the editable PyTorch extension package. The split_attention package is installed separately because it provides a dedicated attention kernel and Python registration interface.

Basic Test

To verify the installation, run from the repository root:

python test_install.py

The script checks that PyTorch can see a CUDA GPU and imports both compiled extensions (offload.runtime and opt_attention._C). On success it prints something like the following (versions, GPU name, and paths will differ):

torch 2.x.x, CUDA 13.0
GPU: NVIDIA GH200 480GB
offload: OK (.../offload/__init__.py)
opt_attention: OK (.../opt_attention/__init__.py)
Installation test PASSED

If a check fails, the script stops with an error. CUDA is not available means PyTorch cannot see the GPU or driver. An ImportError means the corresponding extension was not built or installed.

Usage

The main Python modules are:

• offload.runtime: low-level runtime extension;
• app.python.linear_replacement: utilities for replacing model linear layers;
• app.python.placement: placement and offloading-ratio utilities;
• split_attention: split attention registration.

For attention experiments, register the attention implementation before loading/running the model:

import split_attention

split_attention.register(name="vdcores_opt")

The benchmark scripts under benchmark/ provide examples for running OPT, LLaMA, linear-layer, and attention experiments:

python benchmark/benchmark_opt.py
python benchmark/benchmark_opt_attn.py
python benchmark/benchmark_llama.py
python benchmark/benchmark_linear.py

Shell wrappers and experiment configurations are also provided under benchmark/script/.

Notes

• This is research code and is optimized around the hardware/software assumptions of this implementation.
• Some kernels rely on Hopper-specific features such as Tensor Memory Accelerator support.
• For non-Hopper GPUs, the CUDA architecture flags and kernel assumptions may need to be adjusted.