makora-ai/mako-generate-agent-playground

makora-ai/mako-generate-agent-playground

Language: Shell

Stars: 0

Forks: 0

Open issues: 0

Created: 2025-09-17T11:45:58Z

Pushed: 2025-09-17T14:03:00Z

Default branch: main

Fork: no

Archived: no

README: Export your Mako Generate token (available here)

export MAKO_GENERATE_TOKEN=...

Create a Python file with problem definition in KernelBench format or use one from examples folder:

import torch
import torch.nn as nn

class Model(nn.Module):
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):
x = torch.matmul(x, self.weight.T) # Gemm
x = x / 2 # Divide
x = torch.sum(x, dim=1, keepdim=True) # Sum
x = x * self.scaling_factor # Scaling
return x

batch_size = 1024
input_size = 8192
hidden_size = 8192
scaling_factor = 1.5

def get_inputs(): # Testing data
return [torch.rand(batch_size, input_size)]

def get_init_inputs(): # Constructor arguments
return [input_size, hidden_size, scaling_factor]

Make sure you have jq and curl installed to play around with the scripts.

Create kernel optimization task and get the task id:

$ ./create_task.sh examples/gemm_scale_batch_norm.py
09404933-b054-45d9-af57-96803f7a5306

The task might run from 10min up to 1h. You can watch the execution status with:

$ ./watch_task.sh 09404933-b054-45d9-af57-96803f7a5306
Task 09404933-b054-45d9-af57-96803f7a5306
Execution time: 87s
Status: in_progress
Best result: pending

Task 09404933-b054-45d9-af57-96803f7a5306
Execution time: 92s
Status: in_progress
Best result: pending

Task 09404933-b054-45d9-af57-96803f7a5306
Execution time: 97s
Status: in_progress
Best result: pending

Task 09404933-b054-45d9-af57-96803f7a5306
Execution time: 103s
Status: in_progress
Best result: pending

At the end you will see the benchmarking results and generated kernel:

Task 09404933-b054-45d9-af57-96803f7a5306
Execution time: 1133s
Status: completed

Optimized time: 0.467ms
Torch eager time: 2.78ms
Torch compiled time: 2.78ms
Total agent running time: 1133s

Best result:
import torch
import torch.nn as nn
import triton
import triton.language as tl

class ModelNew(nn.Module):
def __init__(self, in_features, out_features, scale_shape, eps=1e-5, momentum=0.1):
super(ModelNew, self).__init__()
self.gemm = nn.Linear(in_features, out_features)
self.scale = nn.Parameter(torch.randn(scale_shape))
self.bn = nn.BatchNorm1d(out_features, eps=eps, momentum=momentum)

def forward(self, x):
...