correct
21_TriangularMultiplicativeUpdateIncoming.py
started
2026-05-11T09:49:43Z
timingkernelbench.timing
kernel mean / median / std
0.35 / — / 0.73 μs
ref mean / median / std
1.35 / — / 6.63 μs
numerical precisioneval_kernel_against_ref
gpu memoryextended_metrics.measure_memory
kernel launches / fusiontorch.profiler
fusion ratio (ref / kernel)
1.44x
SOL (speed-of-light)heuristic_from_speedup
roofline / occupancyheuristic
reference PyTorch source (the prompt the agent sees)
import torch
import torch.nn as nn
import torch.nn.functional as F
class Model(nn.Module):
"""
Triangular multiplicative update – incoming edges.
Mirror of the outgoing variant: updates pair (i,j) by aggregating over
the shared index k via edges (k,i) and (k,j).
"""
def __init__(self, pair_dim, hidden_dim):
super().__init__()
self.layer_norm_in = nn.LayerNorm(pair_dim)
self.proj_a = nn.Linear(pair_dim, hidden_dim)
self.gate_a = nn.Linear(pair_dim, hidden_dim)
self.proj_b = nn.Linear(pair_dim, hidden_dim)
self.gate_b = nn.Linear(pair_dim, hidden_dim)
self.layer_norm_out = nn.LayerNorm(hidden_dim)
self.out_proj = nn.Linear(hidden_dim, pair_dim)
self.out_gate = nn.Linear(pair_dim, pair_dim)
def forward(self, pair_repr: torch.Tensor, pair_mask: torch.Tensor) -> torch.Tensor:
"""
Args:
pair_repr: (B, N, N, pair_dim)
pair_mask: (B, N, N)
Returns:
updated pair representation (B, N, N, pair_dim)
"""
mask = pair_mask.unsqueeze(-1)
x = self.layer_norm_in(pair_repr)
a = self.proj_a(x) * torch.sigmoid(self.gate_a(x)) * mask # (B, N, N, H)
b = self.proj_b(x) * torch.sigmoid(self.gate_b(x)) * mask
# Incoming: aggregate over k with a[k,i] * b[k,j]
out = torch.einsum("bkih,bkjh->bijh", a, b)
out = self.layer_norm_out(out)
out = self.out_proj(out)
out = torch.sigmoid(self.out_gate(pair_repr)) * out
return out
pair_dim = 64
hidden_dim = 32
seq_len = 32
batch_size = 2
def get_inputs():
p = popcorn_pri
pair_repr = torch.randn(p.jitter_int(batch_size), p.jitter_int(seq_len), p.jitter_int(seq_len), pair_dim)
pair_mask = torch.ones(p.jitter_int(batch_size), p.jitter_int(seq_len), p.jitter_int(seq_len))
return [pair_repr, pair_mask]
def get_init_inputs():
return [pair_dim, hidden_dim]
turn 01142.5s turn · 20.7s LLM · 1121.8s tools · tokens in=1,356 · out=774 · total=2,668 · FINAL
reasoning
The task is to replace the forward computation of a PyTorch model with a custom CUDA kernel. The model is a "Triangular multiplicative update – incoming edges."
I need to create a custom CUDA kernel for the triangular multiplicative update in this PyTorch model. The forward pass includes LayerNorm, linear projections, sigmoid gates, and an einsum operation that's the core computation: torch.
call → submit_kernel
callsubmit_kernel
{
"kernel_code": "import torch\nimport torch.nn as nn\nimport torch.nn.functional as F\nfrom torch.utils.cpp_extension import load_inline\n\ncuda_source = '''\n#include <torch/extension.h>\n#include <cuda.h>\n#include <cuda_runtime.h>\n\n__global__ void triangular_mul_incoming_kernel(\n const float* __restrict__ a,\n const float* __restrict__ b,\n float* __restrict__ out,\n int B, int N, int H\n) {\n int idx = blockIdx.x * blockDim.x + threadIdx.x;\n int total = B * N * N * H;\n if (idx >= total) return;\n \n int h = idx % H;\n int j = (idx / H) % N;\n int i = (idx / (H * N)) % N;\n int batch = idx / (H * N * N);\n \n float sum = 0.0f;\n for (int k = 0; k < N; ++k) {\n float aval = a[((batch * N + k) * N + i) * H + h];\n float bval = b[((batch * N + k) * N + j) * H + h];\n sum += aval * bval;\n }\n out[idx] = sum;\n}\n\ntorch::Tensor triangular_mul_incoming_cuda(torch::Tensor a, torch::Tensor b) {\n auto B = a.size(0);\n auto N = a.size(1);\n auto H = a.size(3);\n \n auto out = torch::zeros({B, N, N, H}, a.options());\n \n int total = B * N * N * H;\n int threads = 256;\n int blocks = (total + threads - 1) / threads;\n \n triangular_mul_incoming_kernel<<<blocks, threads>>>(\n a.data_ptr<float>(),\n b.data_ptr<float>(),\n out.data_ptr<float>(),\n B, N, H\n );\n \n return out;\n}\n'''\n\ncpp_source = '''\ntorch::Tensor triangular_mul_incoming_cuda(torch::Tensor a, torch::Tensor b);\n'''\n\n_triangular_mul = load_inline(\n name='triangular_mul_incoming',\n cpp_sources=[cpp_source],\n cuda_sources=[cuda_source],\n functions=['triangular_mul_incoming_cuda'],\n verbose=False\n)\n\nclass ModelNew(nn.Module):\n \"\"\"\n Triangular multiplicative update \u2013 incoming edges.\n \"\"\"\n\n def __init__(self, pair_dim, hidden_dim):\n super().__init__()\n self.layer_norm_in = nn.LayerNorm(pair_dim)\n self.proj_a = nn.Linear(pair_dim, hidden_dim)\n self.gate_a = nn.Linear(pair_dim, hidden_dim)\n self.proj_b = nn.Linear(pair_dim, hidden_dim)\n self.gate_b = nn.Linear(pair_dim, hidden_dim)\n self.layer_norm_out = nn.LayerNorm(hidden_dim)\n self.out_proj = nn.Linear(hidden_dim, pair_dim)\n self.out_gate = nn.Linear(pair_dim, pair_dim)\n\n def forward(self, pair_repr: torch.Tensor, pair_mask: torch.Tensor) -> torch.Tensor:\n mask = pair_mask.unsqueeze(-1)\n x = self.layer_norm_in(pair_repr)\n\n a = self.proj_a(x) * torch.sigmoid(self.gate_a(x)) * mask\n b = self.proj_b(x) * torch.sigmoid(self.gate_b(x)) * mask\n\n out = _triangular_mul.triangular_mul_incoming_cuda(a, b)\n\n out = self.layer_norm_out(out)\n out = self.out_proj(out)\n out = torch.sigmoid(self.out_gate(pair_repr)) * out\n return out\n"
}
result ← submit_kernel (ok)
oksubmit_kernel
submit_kernel PASSED: (5 / 5) correctness trials all passed.
Kernel runtime: 0.35 μs
Runtime stats: mean=0.35μs median=?μs std=0.73μs
Numerical precision: max_abs_err=0.00e+00 mean_abs_err=0.00e+00 max_rel_err=0.00e+00
Memory: 36.13 MB peak (ref: 36.13 MB, ratio: 1.0x)
Kernel launches: 70 (ref: 101, fusion ratio: 1.4429)
SOL score: 0.3846
final submitted kernel (98 lines, 2,800 chars) — full source, untruncated
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.utils.cpp_extension import load_inline
cuda_source = '''
#include <torch/extension.h>
#include <cuda.h>
#include <cuda_runtime.h>
__global__ void triangular_mul_incoming_kernel(
const float* __restrict__ a,
const float* __restrict__ b,
float* __restrict__ out,
int B, int N, int H
) {
int idx = blockIdx.x * blockDim.x + threadIdx.x;
int total = B * N * N * H;
if (idx >= total) return;
int h = idx % H;
int j = (idx / H) % N;
int i = (idx / (H * N)) % N;
int batch = idx / (H * N * N);
float sum = 0.0f;
for (int k = 0; k < N; ++k) {
float aval = a[((batch * N + k) * N + i) * H + h];
float bval = b[((batch * N + k) * N + j) * H + h];
sum += aval * bval;
}
out[idx] = sum;
}
torch::Tensor triangular_mul_incoming_cuda(torch::Tensor a, torch::Tensor b) {
auto B = a.size(0);
auto N = a.size(1);
auto H = a.size(3);
auto out = torch::zeros({B, N, N, H}, a.options());
int total = B * N * N * H;
int threads = 256;
int blocks = (total + threads - 1) / threads;
triangular_mul_incoming_kernel<<<blocks, threads>>>(
a.data_ptr<float>(),
b.data_ptr<float>(),
out.data_ptr<float>(),
B, N, H
);
return out;
}
'''
cpp_source = '''
torch::Tensor triangular_mul_incoming_cuda(torch::Tensor a, torch::Tensor b);
'''
_triangular_mul = load_inline(
name='triangular_mul_incoming',
cpp_sources=[cpp_source],
cuda_sources=[cuda_source],
functions=['triangular_mul_incoming_cuda'],
verbose=False
)
class ModelNew(nn.Module):
"""
Triangular multiplicative update – incoming edges.
"""
def __init__(self, pair_dim, hidden_dim):
super().__init__()
self.layer_norm_in = nn.LayerNorm(pair_dim)
self.proj_a = nn.Linear(pair_dim, hidden_dim)
self.gate_a = nn.Linear(pair_dim, hidden_dim)
self.proj_b = nn.Linear(pair_dim, hidden_dim)
self.gate_b = nn.Linear(pair_dim, hidden_dim)
self.layer_norm_out = nn.LayerNorm(hidden_dim)
self.out_proj = nn.Linear(hidden_dim, pair_dim)
self.out_gate = nn.Linear(pair_dim, pair_dim)
def forward(self, pair_repr: torch.Tensor, pair_mask: torch.Tensor) -> torch.Tensor:
mask = pair_mask.unsqueeze(-1)
x = self.layer_norm_in(pair_repr)
a = self.proj_a(x) * torch.sigmoid(self.gate_a(x)) * mask
b = self.proj_b(x) * torch.sigmoid(self.gate_b(x)) * mask
out = _triangular_mul.triangular_mul_incoming_cuda(a, b)
out = self.layer_norm_out(out)
out = self.out_proj(out)
out = torch.sigmoid(self.out_gate(pair_repr)) * out
return out