Skip to main content

CLI for running GPU workloads, managing remote workspaces, and evaluating/optimizing kernels

Project description

Wafer CLI

Run GPU workloads, optimize kernels, and query GPU documentation.

Getting Started

# Install
cd apps/wafer-cli && uv sync

# Use staging (workspaces and other features require staging)
wafer config set api.environment staging

# Login
wafer login

# Run a command on a remote GPU
wafer remote-run -- nvidia-smi

Commands

wafer login / wafer logout / wafer whoami

Authenticate with GitHub OAuth.

wafer login          # Opens browser for GitHub OAuth
wafer whoami         # Show current user
wafer logout         # Remove credentials

wafer remote-run

Run any command on a remote GPU.

wafer remote-run -- nvidia-smi
wafer remote-run --upload-dir ./my_code -- python3 train.py

wafer workspaces

Create and manage persistent GPU environments.

Available GPUs:

  • MI300X - AMD Instinct MI300X (192GB HBM3, ROCm)
  • B200 - NVIDIA Blackwell B200 (180GB HBM3e, CUDA) - default
wafer target workspace list
wafer target workspace create my-workspace --gpu B200 --wait   # NVIDIA B200
wafer target workspace create amd-dev --gpu MI300X             # AMD MI300X
wafer target workspace ssh <workspace-id>
wafer target workspace delete <workspace-id>

wafer agent

AI assistant for GPU kernel development. Helps with CUDA/Triton optimization, documentation queries, and performance analysis.

wafer agent "What is TMEM in CuTeDSL?"
wafer agent -s "optimize this kernel" < kernel.py

wafer tool eval

Evaluate kernel correctness and performance against a reference implementation.

Functional format (default):

# Generate template files
wafer tool eval make-template ./my-kernel

# Run evaluation
wafer tool eval gpumode --impl kernel.py --reference ref.py --test-cases tests.json --benchmark

The implementation must define custom_kernel(inputs), the reference must define ref_kernel(inputs) and generate_input(**params).

KernelBench format (ModelNew class):

# Extract a KernelBench problem as template
wafer tool eval kernelbench make-template level1/1

# Run evaluation
wafer tool eval kernelbench --impl my_kernel.py --reference problem.py --benchmark

The implementation must define class ModelNew(nn.Module), the reference must define class Model, get_inputs(), and get_init_inputs().

wafer agent -t ask-docs

Query GPU documentation using the docs template. Uses the ask_docs tool to search wafer's documentation corpus via the API.

wafer agent -t ask-docs -s "What causes bank conflicts in shared memory?"

Customization

wafer tool eval options

wafer tool eval gpumode --impl k.py --reference r.py --test-cases t.json \
    --target vultr-b200 \    # Specific GPU target
    --benchmark \            # Measure performance
    --profile                # Enable torch.profiler + NCU

Profile analysis

wafer tool ncu analyze profile.ncu-rep
wafer tool nsys analyze profile.nsys-rep

Advanced

Local targets

Bypass the API and SSH directly to your own GPUs:

wafer target config list
wafer target config add ./my-gpu.toml
wafer target config default my-gpu

Defensive evaluation

Detect evaluation hacking (stream injection, lazy evaluation, etc.):

wafer tool eval gpumode --impl k.py --reference r.py --test-cases t.json --benchmark --defensive

Other tools

wafer tool perfetto <trace.json> --query "SELECT * FROM slice"   # Perfetto SQL queries
wafer tool capture ./script.py                                    # Capture execution snapshot
wafer compiler-analyze kernel.ptx                                 # Analyze PTX/SASS

ROCm profiling (AMD GPUs)

wafer tool rocprof-sdk ...
wafer tool rocprof-systems ...
wafer tool rocprof-compute ...

Shell Completion

Enable tab completion for commands, options, and target names:

# Install completion (zsh/bash/fish)
wafer --install-completion

# Then restart your terminal, or source your shell config:
source ~/.zshrc  # or ~/.bashrc

Now you can tab-complete:

  • Commands: wafer tool ev<TAB>wafer tool eval
  • Options: wafer tool eval --<TAB>
  • Target names: wafer tool eval --target v<TAB>wafer tool eval --target vultr-b200
  • File paths: wafer tool eval gpumode --impl ./<TAB>

AI Assistant Skills

Install the Wafer CLI skill to make wafer commands discoverable by your AI coding assistant:

# Install for all supported tools (Claude Code, Codex CLI, Cursor)
wafer skill install

# Install for a specific tool
wafer skill install -t cursor    # Cursor
wafer skill install -t claude    # Claude Code
wafer skill install -t codex     # Codex CLI

# Check installation status
wafer skill status

# Uninstall
wafer skill uninstall

Installing from GitHub (Cursor)

You can also install the skill directly from GitHub in Cursor:

  1. Open Cursor Settings (Cmd+Shift+J / Ctrl+Shift+J)
  2. Navigate to RulesAdd RuleRemote Rule (Github)
  3. Enter: https://github.com/wafer-ai/skills
  4. Cursor will automatically discover skills in .cursor/skills/

The skill provides comprehensive guidance for GPU kernel development, including documentation lookup, trace analysis, kernel evaluation, and optimization workflows.


Requirements

  • Python 3.10+
  • GitHub account (for authentication)

Project details


Download files

Download the file for your platform. If you're not sure which to choose, learn more about installing packages.

Source Distribution

wafer_cli-0.2.61.tar.gz (328.3 kB view details)

Uploaded Source

Built Distribution

If you're not sure about the file name format, learn more about wheel file names.

wafer_cli-0.2.61-py3-none-any.whl (281.1 kB view details)

Uploaded Python 3

File details

Details for the file wafer_cli-0.2.61.tar.gz.

File metadata

  • Download URL: wafer_cli-0.2.61.tar.gz
  • Upload date:
  • Size: 328.3 kB
  • Tags: Source
  • Uploaded using Trusted Publishing? No
  • Uploaded via: twine/6.2.0 CPython/3.12.12

File hashes

Hashes for wafer_cli-0.2.61.tar.gz
Algorithm Hash digest
SHA256 685cfb79deca4bd68a85f1ec7261ad367120c67fef03d1389ac6205b1ac17195
MD5 0d49a573fabb433241b04f0689b54b18
BLAKE2b-256 e35879c57d51cf9382e18bea928966685f2b883a13b7d829a38cb0685a320d79

See more details on using hashes here.

File details

Details for the file wafer_cli-0.2.61-py3-none-any.whl.

File metadata

  • Download URL: wafer_cli-0.2.61-py3-none-any.whl
  • Upload date:
  • Size: 281.1 kB
  • Tags: Python 3
  • Uploaded using Trusted Publishing? No
  • Uploaded via: twine/6.2.0 CPython/3.12.12

File hashes

Hashes for wafer_cli-0.2.61-py3-none-any.whl
Algorithm Hash digest
SHA256 a3d49c31967c8e8c71453249f79803f7aeb093ea2a6a59378aae157e8697afd0
MD5 0b86d3871d741dd500c96a42c5dbba36
BLAKE2b-256 2b7df5e20661ccbc15a01f4cb68a8250b64ebd6ffc87f5042c4ea263f2d3f4a5

See more details on using hashes here.

Supported by

AWS Cloud computing and Security Sponsor Datadog Monitoring Depot Continuous Integration Fastly CDN Google Download Analytics Pingdom Monitoring Sentry Error logging StatusPage Status page