Getting Started

Install Weco

pip install weco

Set Your LLM API Key

Create your OpenAI API key here.

export OPENAI_API_KEY="your_key_here"

Try the Example Project

git clone https://github.com/WecoAI/weco-cli.git
cd weco-cli/examples/hello-kernel-world
pip install torch

Choose Your Approach

Let the Weco copilot analyze your codebase, suggest & setup optimizations for you. Run Weco without any arguments to start the interactive copilot:

weco

What to Expect

Both approaches will show you the same optimization interface. Here's what you can expect to see (keep an eye on that Best Solution panel):

Applying Weco to Your Own Project

Figure Out Your Evaluation Script

Your evaluation script needs to benchmark the code you'd like to optimize and print a metric you're optimizing for to the console.

For specific examples of evaluation scripts for kernel engineering (PyTorch, CUDA, Triton etc.,), ML research and prompt engineering, check out the Examples section. If you'd like to know how to write a good evaluation script, we've got you covered with this guide.

Basic Example

Here's a simple example that optimizes a PyTorch function for speedup. You can also follow along here on Google Colab.

First install the dependencies:

pip install torch weco

Here's a simple example of a PyTorch model that we'll optimize:

# File: optimize.py
 
import torch
import torch.nn as nn
 
 
class Model(nn.Module):
    """
    Model that performs a matrix multiplication, division, summation, and scaling.
    """
 
    def __init__(self, input_size, hidden_size, scaling_factor):
        super(Model, self).__init__()
        self.weight = nn.Parameter(torch.randn(hidden_size, input_size))
        self.scaling_factor = scaling_factor
 
    def forward(self, x):
        """
        Args:
            x (torch.Tensor): Input tensor of shape (batch_size, input_size).
        Returns:
            torch.Tensor: Output tensor of shape (batch_size, hidden_size).
        """
        x = torch.matmul(x, self.weight.T)
        x = x / 2
        x = torch.sum(x, dim=1, keepdim=True)
        x = x * self.scaling_factor
        return x

Here's what your evaluation script needs to do:

# File: evaluate.py
 
# 1. Load the baseline model and the optimized model from optimize.py
# Remember to keep a copy of your baseline code in a separate file so that the benchmark can compare it to the generated solution by Weco.
baseline_model = create_baseline_model()
optimized_model = load_model_from_file('optimize.py')
 
# 2. Prepare test data
test_inputs = generate_test_data()
 
# 3. Measure correctness (ensure optimization didn't break functionality)
baseline_output = baseline_model(test_inputs)
optimized_output = optimized_model(test_inputs)
correctness_diff = calculate_max_difference(baseline_output, optimized_output)
if correctness_diff > tolerance:
    raise ValueError(f"Max float diff between values is too high! Got: {correctness_diff}")
 
# 4. Benchmark performance
baseline_time = benchmark(baseline_model, test_inputs)
optimized_time = benchmark(optimized_model, test_inputs)
 
# 5. Calculate and print the final metric for Weco to read
speedup = baseline_time / optimized_time
print(f"speedup: {speedup}")

You can find the complete evaluate.py file for this example here.

Now run Weco to optimize your code:

weco run --source optimize.py \
     --eval-command "python evaluate.py --solution-path optimize.py --device cpu" \
     --metric speedup \
     --goal maximize \
     --steps 10 \
     --additional-instructions "Fuse operations in the forward method while ensuring the max float deviation remains small. Maintain the same format of the code."

Note: If you have an NVIDIA GPU, change the device in the --eval-command to cuda. If you are running this on Apple Silicon, set it to mps.

What's Next?

Try different optimizations: Explore Examples for kernel optimization, ML models, and prompt engineering
Improve your results: Learn Writing Good Evaluation Scripts
Advanced usage: Check the CLI Reference for all command options
Need help?: Visit our FAQ for common questions

Getting Started

On this page