Skip to main content

Wafer Agent

The Wafer agent is an AI-powered assistant specialized in GPU kernel development. It can help you optimize CUDA and Triton kernels, understand GPU architecture, analyze performance traces, and navigate documentation.

Quick Start

# Ask a question
wafer agent "How do I optimize memory coalescing in CUDA?"

# Interactive mode with full TUI
wafer agent -i

# Single-turn mode (answer once and exit)
wafer agent -s "What causes warp divergence?"

Modes

The agent supports three interaction modes:
Full terminal UI with conversation history and rich formatting.
wafer agent -i
wafer agent --interactive
Features:
  • Persistent conversation history
  • Tool execution visualization
  • Code syntax highlighting
  • Session management

Options

OptionShortDescription
--interactive-iLaunch full interactive TUI mode
--simpleUse simple stdout mode (for scripts)
--single-turn-sAnswer once and exit
--template-tRun with a template (see Templates)
--corpus-cDocumentation corpus: cuda, cutlass, hip, amd
--model-mModel override (default: claude-opus-4-5)
--max-turnsMaximum conversation turns (default: 10)
--jsonOutput in JSON format
--toolsComma-separated list of tools to enable
--resume-rResume a previous session (see Sessions)

Examples

Ask About GPU Architecture

wafer agent "What's the difference between L1 and L2 cache on NVIDIA GPUs?"

Query Specific Documentation

Use the --corpus flag to focus on specific documentation:
# CUDA documentation
wafer agent -c cuda "How do I use cooperative groups?"

# CuTLASS library
wafer agent -c cutlass "Show me an example of GEMM with CuTLASS"

# AMD HIP documentation
wafer agent -c hip "How do I port CUDA to HIP?"

Analyze a Trace

wafer agent -t trace-analyze --args trace=./profile.ncu-rep "What's the bottleneck?"

Optimize a Kernel

wafer agent -t optimize-kernel --args kernel=./matmul.cu "Optimize for H100"

Available Tools

The agent has access to various tools depending on the mode and template:
  • File operations: Read, write, edit files in your project
  • Search: Find patterns in code and documentation
  • Bash: Execute shell commands (sandboxed by default)
  • Wafer tools: Run profiling, evaluation, and analysis commands
By default, the agent runs in a sandbox that restricts file system access. Use --no-sandbox to disable this for full access (use with caution).

API Access

The agent uses Wafer’s API proxy by default. To use your own Anthropic API key:
# Set your API key
export ANTHROPIC_API_KEY=sk-ant-...

# Run without proxy
wafer agent --no-proxy "Your question here"

Next Steps