warpdrive-gpu 0.1.0

Hardware-guided GPU kernel diagnostics for LLM agents, over MCP + CLI.

WarpDrive

Hardware-guided GPU kernel diagnostics for agent-driven optimization

Quickstart • Features • CLI • Agents & MCP • GPU setup • Trust boundary

$ CUDA_VISIBLE_DEVICES=0 warpdrive --db runs/rmsnorm-sweep.db sweep \
    llama-fusions residual_rmsnorm \
    examples/residual_rmsnorm_triton.py examples/residual_rmsnorm_ref.py

[########################] 120/120 100.0% float32 8192,4096 trial 10/10 LIVE ok 3.47x
dtype     shape      pass   fused us  cv    unfused us  speedup  95% CI         p       verdict
--------  ---------  -----  --------  ----  ----------  -------  -------------  ------  ---------
float16   8192x4096  10/10  150.4     0.2%  1858.1      12.36x   [12.34,12.37]  0.0010  clear win
bfloat16  8192x4096  10/10  150.4     0.1%  1868.2      12.42x   [12.41,12.43]  0.0010  clear win
float32   8192x4096  10/10  312.1     0.1%  1086.1      3.48x    [3.48,3.48]    0.0010  clear win

WarpDrive sits between an optimization agent, a PyTorch correctness oracle, and NVIDIA Nsight Compute (ncu). It verifies kernels first, profiles only kernels that pass, records the result in a SQLite ledger, and returns diagnostics an LLM agent can use for the next edit.

WarpDrive is a passive tool. It does not call an LLM or own the optimization loop; Codex, Claude Code, Cursor, or your own script drives iteration.

⚡ Get started

Linux GPU VM

git clone https://github.com/michaelalbada/WarpDrive
cd WarpDrive
scripts/setup_ubuntu_gpu.sh
source .venv/bin/activate
warpdrive doctor

Run the v0.1 smoke workflow:

CUDA_VISIBLE_DEVICES=0 warpdrive --db runs/rmsnorm-sweep.db sweep \
    llama-fusions residual_rmsnorm \
    examples/residual_rmsnorm_triton.py examples/residual_rmsnorm_ref.py

That short sweep command defaults to the release fusion sweep profile: Triton on CUDA, shapes 1,4096 16,4096 256,4096 8192,4096, dtypes float16 bfloat16 float32, 10 trials, and unfused comparison enabled.

For a faster tuning loop:

CUDA_VISIBLE_DEVICES=0 warpdrive --db runs/quick.db sweep \
    llama-fusions residual_rmsnorm \
    examples/residual_rmsnorm_triton.py examples/residual_rmsnorm_ref.py \
    --trials 3 --dtypes float16 --shapes 8192,4096

Package install

pip install warpdrive-gpu                 # CLI + MCP schema + engine
pip install 'warpdrive-gpu[gpu]'          # adds PyTorch for correctness checks
pip install 'warpdrive-gpu[gpu,triton]'   # Linux/CUDA live Triton profiling

torch is optional for import-time CPU development. The correctness verifier needs it at runtime; live Triton/CUDA profiling needs Linux, CUDA, PyTorch, Triton, ncu, and a visible NVIDIA GPU.

⭐ Features

Diagnostics

• Correctness gate first: kernels that fail staged PyTorch verification are never profiled.
• Live or mock profiling: use real ncu on GPU nodes; fall back to deterministic mock metrics on CPU-only machines.
• Minimal metric set: WarpDrive scrapes the small metric surface needed for actionable reports instead of running full-section ncu replays.
• Hardware context: doctor checks GPU visibility, CUDA tools, PyTorch CUDA, ncu, and profiling-counter permission.
• Roofline grounding: roofline reports known published GPU peaks for optimization target sanity checks.

Fusion sweeps

• Unfused comparison: profile the fused candidate and the PyTorch reference workload, then report speedup and top unfused kernels.
• Shape/dtype sweeps: repeat trials across realistic model shapes and dtypes.
• Statistical table: mean, sample variance/stddev, coefficient of variation, confidence interval, sign-flip p-value, and practical verdict.
• Progress bar: large sweeps show trial progress without polluting JSON output.
• SQLite ledger: session history gives baseline, current, best, and deltas without the agent keeping state.

Agent loop

• CLI and MCP front doors: one engine powers the human CLI and MCP tools.
• JSON output: --json returns structured data for Codex, Claude Code, Cursor, or custom scripts.
• Diff-aware reports: later iterations summarize code changes against the previous call in the same session.
• Trusted-code boundary: untrusted verification runs in a bounded child process, with explicit caveats below.

📥 Install

PyPI package

pip install warpdrive-gpu
warpdrive doctor

Install extras as needed:

pip install 'warpdrive-gpu[gpu]'          # numpy + torch
pip install 'warpdrive-gpu[triton]'       # Triton, Linux/CUDA only
pip install 'warpdrive-gpu[gpu,triton]'   # full GPU runtime

The distribution name is warpdrive-gpu; the Python package and CLI remain warpdrive.

Editable package

pip install -e .

Install extras as needed:

pip install -e '.[gpu]'             # numpy + torch
pip install -e '.[triton]'          # Triton, Linux/CUDA only
pip install -e '.[dev]'             # pytest + coverage + ruff
pip install -e '.[gpu,triton,dev]'  # full GPU development environment

Ubuntu GPU helper

For a regular Ubuntu x86_64 GPU node, the setup helper creates .venv, installs WarpDrive with GPU/Triton/dev extras, keeps Triton/PyTorch/pip caches under the workspace, checks NVIDIA tools, and prints the profiling-counter fix if ncu is restricted:

scripts/setup_ubuntu_gpu.sh
source .venv/bin/activate

If common Ubuntu packages are missing and you have sudo:

scripts/setup_ubuntu_gpu.sh --system-packages

Use --venv DIR or WARPDRIVE_VENV=DIR to choose a different virtual environment location.

🧰 CLI

warpdrive doctor
warpdrive roofline
warpdrive profile SESSION KERNEL CODE REFERENCE [--language triton|cuda] \
    [--device cuda|cpu] [--shape M,N] [--dtype DTYPE] [--mock] [--compare-unfused]
warpdrive sweep SESSION KERNEL CODE REFERENCE \
    [--shapes M,N ...] [--dtypes ...] [--trials N] [--no-compare-unfused]
warpdrive history SESSION [KERNEL]

Doctor

warpdrive doctor
warpdrive doctor --json

doctor reports GPU detection, CUDA version, nvcc, ncu, PyTorch CUDA, and whether ncu profiling is actually permitted. It launches a tiny probe instead of only checking ncu --version, so it does not advertise LIVE mode when profiling counters are restricted.

Profile one candidate

warpdrive --db runs/rmsnorm.db profile llama-fusions residual_rmsnorm \
    examples/residual_rmsnorm_triton.py examples/residual_rmsnorm_ref.py \
    --language triton --device cuda --dtype float16 --shape 8192,4096 \
    --compare-unfused

CODE and REFERENCE may be inline source or paths to files. profile prints a Markdown diagnostic by default; add --json for structured agent output.

Sweep for confidence

warpdrive --db runs/rmsnorm-sweep.db sweep llama-fusions residual_rmsnorm \
    examples/residual_rmsnorm_triton.py examples/residual_rmsnorm_ref.py

Use sweep when the question is "does this win hold across shapes, dtypes, and repeated trials?" By default it profiles the fused candidate and the unfused reference workload for each trial.

Live sweep wall-clock time is much larger than the GPU kernel time shown in the table. A 4-shape x 3-dtype x 10-trial sweep with unfused comparison runs about 240 live ncu profiling launches plus correctness checks; ncu process startup, CUDA/Triton initialization, replay/instrumentation, CSV parsing, and SQLite recording dominate elapsed time.

History

warpdrive --db runs/rmsnorm-sweep.db history llama-fusions
warpdrive --db runs/rmsnorm-sweep.db history llama-fusions residual_rmsnorm/float16/8192x4096

History is recorded per session_id and kernel_name. Sweep combinations are stored as derived kernel names such as residual_rmsnorm/float16/8192x4096.

🤖 Agent & MCP loop

Run the MCP server over stdio:

python -m warpdrive.server.mcp_server

Tools:

Tool	Purpose
warpdrive_profile	Compile, verify, profile, record, and return diagnostic Markdown
warpdrive_doctor	Report GPU model, CUDA version, and ncu profiling permission
warpdrive_get_roofline	Report host GPU published peaks

warpdrive_profile takes session_id, kernel_name, language, code, reference_pytorch_code, and optional compare_unfused. The ledger path comes from WARPDRIVE_DB and defaults to warpdrive.db.

For CLI-driven agents, use --json on doctor, profile, sweep, and history so the agent can parse results directly.

🧩 Kernel contract

Both kernel and reference are Python source strings evaluated by the verifier.

• Reference: define reference(inputs: list[Tensor]) -> Tensor. Optionally define make_inputs(shape, dtype, mode) -> list[Tensor] for multi-input kernels.
• Kernel: define run(inputs: list[Tensor]) -> Tensor, or a function named after kernel_name.

Inputs are placed on CUDA when available, otherwise CPU. The same make_inputs feeds correctness verification and live profiling.

For CUDA C++, wrap the kernel in Python that launches it, for example with torch.utils.cpp_extension.load_inline or CuPy. --language cuda additionally runs nvcc for compile-time registers/spills, but raw .cu source alone is not executed end to end because it is not valid Python.

The verifier runs staged tests:

• standard powers-of-two shapes
• unaligned prime/non-aligned dimensions
• extreme per-dtype bounded magnitudes
• adversarial NaN/Inf and long constant runs
• precision checks with per-dtype tolerances

🖥️ GPU setup

Run the full GPU phase on a GPU node:

scripts/run_gpu_phase.sh

It:

1. runs warpdrive doctor and exits early with guidance if GPU, ncu, or profiling permission is missing;
2. captures ncu output for examples/softmax_triton.py into tests/fixtures/ncu_softmax_live.csv;
3. runs gated GPU smoke tests with live compile + profile;
4. runs the full test suite.

Capture a fixture manually:

warpdrive capture examples/softmax_triton.py tests/fixtures/my_kernel.csv \
    --shape 4096,4096 --kernel-regex softmax

Example kernels live in examples/. residual_rmsnorm_triton.py is the model-like fused subgraph example used by the release smoke workflow.

🔒 Trust boundary

WarpDrive compiles and executes agent-supplied code. Treat that code as untrusted by construction.

v0.1 isolation:

• verification runs in a child process with a hard wall-clock timeout;
• POSIX workers use CPU and address-space rlimits where practical;
• kernels need no network access.

This is not a security sandbox. It bounds runaway/OOM code, not malice. Run WarpDrive only against trusted agents, or inside a container. Container-based isolation with the capabilities needed for profiling is the recommended production posture and a v0.2 item.

If doctor reports restricted profiling, check NVreg_RestrictProfilingToAdminUsers=0 or run inside a container with the appropriate capabilities. Without profiling permission, WarpDrive returns representative MOCK metrics so the loop still runs.

🧪 Testing

pytest
pytest --cov=warpdrive
ruff check warpdrive tests

CPU-only CI covers verifier logic, ncu output parsing against recorded fixtures, report rendering, ledger behavior, MCP schema, mock-mode profiling, and CLI behavior. Gated GPU tests cover live Triton/CUDA verification and profiling.

🚢 v0.1 release scope

WarpDrive 0.1 is intentionally narrow:

• doctor for GPU, CUDA Toolkit, ncu, PyTorch/CUDA, and profiling permission checks.
• profile for compile, correctness, live/mock profile, ledger record, and Markdown diagnostics.
• sweep for repeated shape/dtype trials, progress, statistical summary table, and JSON output.
• history for session iteration history, baselines, deltas, and provenance.
• unfused comparison for fusion work.
• GPU setup and validation docs/scripts.
• explicit trusted-code warning.

Out of scope for v0.1: model-level hotspot discovery, privileged clock locking, dashboard/TUI, and stronger sandbox provisioning.

🛣️ Roadmap

• v0.2 Agent loop ergonomics: richer JSON contracts, best-prior-code retrieval, prompt-ready diagnostics, and more MCP tools.
• v0.3 Hotspot discovery: trace a model/workload, rank latency hotspots, extract shapes/dtypes/strides, and hand agents the right kernel to optimize.
• Future UI: textual TUI or web dashboard for evolution history and benchmark comparisons.

📜 License

Apache-2.0. See LICENSE.