Metadata-Version: 2.4
Name: warpdrive-gpu
Version: 0.1.0
Summary: Hardware-guided GPU kernel diagnostics for LLM agents, over MCP + CLI.
Author-email: Michael Albada <malbada@nvidia.com>
License: Apache-2.0
Requires-Python: >=3.11
Requires-Dist: mcp>=1.0.0
Provides-Extra: dev
Requires-Dist: pytest-cov>=5.0; extra == 'dev'
Requires-Dist: pytest>=8.0; extra == 'dev'
Requires-Dist: ruff>=0.6; extra == 'dev'
Provides-Extra: gpu
Requires-Dist: numpy>=1.26; extra == 'gpu'
Requires-Dist: torch>=2.1; extra == 'gpu'
Provides-Extra: triton
Requires-Dist: triton>=2.1; extra == 'triton'
Description-Content-Type: text/markdown

<h1 align="center" style="margin:0;">
  WarpDrive
</h1>

<h3 align="center" style="margin:0; margin-top:0;">
Hardware-guided GPU kernel diagnostics for agent-driven optimization
</h3>

<p align="center">
  <a href="#-get-started">Quickstart</a> •
  <a href="#-features">Features</a> •
  <a href="#-cli">CLI</a> •
  <a href="#-agent--mcp-loop">Agents & MCP</a> •
  <a href="#-gpu-setup">GPU setup</a> •
  <a href="#-trust-boundary">Trust boundary</a>
</p>

<br>

```text
$ 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

```bash
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:

```bash
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:

```bash
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

```bash
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

```bash
pip install warpdrive-gpu
warpdrive doctor
```

Install extras as needed:

```bash
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

```bash
pip install -e .
```

Install extras as needed:

```bash
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:

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

If common Ubuntu packages are missing and you have sudo:

```bash
scripts/setup_ubuntu_gpu.sh --system-packages
```

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

## 🧰 CLI

```bash
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

```bash
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

```bash
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

```bash
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

```bash
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:

```bash
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:

```bash
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:

```bash
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

```bash
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](LICENSE).
