correct
29_logsumexp_across_virtual_ranks.py
started
2026-05-11T08:59:53Z
timingkernelbench.timing
kernel mean / median / std
0.01 / — / 1.50e-04 μs
ref mean / median / std
0.03 / — / 3.87e-04 μs
numerical precisioneval_kernel_against_ref
gpu memoryextended_metrics.measure_memory
kernel launches / fusiontorch.profiler
fusion ratio (ref / kernel)
20.00x
SOL (speed-of-light)nsight_hardware_counters
roofline / occupancynsight_hardware_counters
warp stalls (top 5)nsight
reference PyTorch source (the prompt the agent sees)
"""
Stable log-sum-exp over virtual ranks: out = logsumexp(x, dim=0), shape (B, D).
"""
import torch
import torch.nn as nn
from kernelbench.distributed_collectives import default_device
R = 4
B = 10
D = 32
class Model(nn.Module):
def __init__(self):
super().__init__()
def forward(self, x: torch.Tensor) -> torch.Tensor:
return torch.logsumexp(x, dim=0)
def get_inputs():
dev = default_device()
p = popcorn_pri
return [torch.randn(p.jitter_int(R), p.jitter_int(B), p.jitter_int(D, align=8), device=dev)]
def get_init_inputs():
return []
turn 0421.0s turn · 17.8s LLM · 403.2s tools · tokens in=1,027 · out=488 · total=2,312 · FINAL
reasoning
The task is to replace the forward computation with a custom CUDA kernel for stable log-sum-exp over virtual ranks: out = logsumexp(x, dim=0), shape (B, D). The input is (R, B, D) with R=4, B=10, D=32...I need to implement a custom CUDA kernel for logsumexp along dim=0 for an input tensor of shape (R, B, D), producing an output of shape (B, D). The dimensions are small: R=4, B=10, D=Policy: This is your last turn before the session 32, but I still need to create a custom kernel to replace PyTorch's built-in logsumexp for potential speed or stability reasons.
The kernel should compute, for each (b, d), the log of the sum of exponentials over r, with stability by subtracting the max value first.
call → submit_kernel
callsubmit_kernel
{
"kernel_code": "import torch\nimport torch.nn as nn\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 logsumexp_kernel(const float* __restrict__ x, float* __restrict__ out, int R, int B, int D) {\n int bd = blockIdx.x * blockDim.x + threadIdx.x;\n if (bd >= B * D) return;\n int b = bd / D;\n int d = bd % D;\n float max_val = -INFINITY;\n for(int r = 0; r < R; r++) {\n float val = x[r * B * D + b * D + d];\n max_val = fmaxf(max_val, val);\n }\n float sum = 0.0f;\n for(int r = 0; r < R; r++) {\n float val = x[r * B * D + b * D + d];\n sum += expf(val - max_val);\n }\n out[b * D + d] = max_val + logf(sum);\n}\n\ntorch::Tensor logsumexp_cuda(torch::Tensor x) {\n TORCH_CHECK(x.dim() == 3, \"Input must be 3D\");\n int R = x.size(0);\n int B = x.size(1);\n int D = x.size(2);\n auto out = torch::empty({B, D}, x.options());\n int threads = 256;\n int blocks = (B * D + threads - 1) / threads;\n logsumexp_kernel<<<blocks, threads>>>(x.data_ptr<float>(), out.data_ptr<float>(), R, B, D);\n return out;\n}\n'''\n\ncpp_source = '''\ntorch::Tensor logsumexp_cuda(torch::Tensor x);\n'''\n\nmodule = load_inline(\n name='logsumexp_ext',\n cpp_sources=[cpp_source],\n cuda_sources=[cuda_source],\n functions=['logsumexp_cuda'],\n verbose=False\n)\n\nclass ModelNew(nn.Module):\n def __init__(self):\n super().__init__()\n self.module = module\n\n def forward(self, x: torch.Tensor) -> torch.Tensor:\n return self.module.logsumexp_cuda(x)\n"
}
result ← submit_kernel (ok)
oksubmit_kernel
submit_kernel PASSED: (5 / 5) correctness trials all passed.
Kernel runtime: 0.01 μs
Runtime stats: mean=0.01μs median=?μs std=0.00μs
Numerical precision: max_abs_err=0.00e+00 mean_abs_err=0.00e+00 max_rel_err=0.00e+00
Memory: 32.01 MB peak (ref: 32.01 MB, ratio: 0.9998x)
Kernel launches: 1 (ref: 20, fusion ratio: 20.0)
SOL score: 0.001
final submitted kernel (60 lines, 1,605 chars) — full source, untruncated
import torch
import torch.nn as nn
from torch.utils.cpp_extension import load_inline
cuda_source = '''
#include <torch/extension.h>
#include <cuda.h>
#include <cuda_runtime.h>
__global__ void logsumexp_kernel(const float* __restrict__ x, float* __restrict__ out, int R, int B, int D) {
int bd = blockIdx.x * blockDim.x + threadIdx.x;
if (bd >= B * D) return;
int b = bd / D;
int d = bd % D;
float max_val = -INFINITY;
for(int r = 0; r < R; r++) {
float val = x[r * B * D + b * D + d];
max_val = fmaxf(max_val, val);
}
float sum = 0.0f;
for(int r = 0; r < R; r++) {
float val = x[r * B * D + b * D + d];
sum += expf(val - max_val);
}
out[b * D + d] = max_val + logf(sum);
}
torch::Tensor logsumexp_cuda(torch::Tensor x) {
TORCH_CHECK(x.dim() == 3, "Input must be 3D");
int R = x.size(0);
int B = x.size(1);
int D = x.size(2);
auto out = torch::empty({B, D}, x.options());
int threads = 256;
int blocks = (B * D + threads - 1) / threads;
logsumexp_kernel<<<blocks, threads>>>(x.data_ptr<float>(), out.data_ptr<float>(), R, B, D);
return out;
}
'''
cpp_source = '''
torch::Tensor logsumexp_cuda(torch::Tensor x);
'''
module = load_inline(
name='logsumexp_ext',
cpp_sources=[cpp_source],
cuda_sources=[cuda_source],
functions=['logsumexp_cuda'],
verbose=False
)
class ModelNew(nn.Module):
def __init__(self):
super().__init__()
self.module = module
def forward(self, x: torch.Tensor) -> torch.Tensor:
return self.module.logsumexp_cuda(x)