correct
29_logsumexp_across_virtual_ranks.py
started
2026-05-11T11:00:49Z
timingkernelbench.timing
kernel mean / median / std
0.01 / — / 3.69e-04 μs
ref mean / median / std
0.03 / — / 7.12e-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 0356.6s turn · 13.3s LLM · 343.3s tools · tokens in=1,394 · out=6 · total=2,145
reasoning
The task is to replace the forward computation of a PyTorch model with a custom CUDA kernel. The model computes the log-sum-exp over the first dimension of a tensor x with shape (R, B, D), where R=4, ...I need to create a ModelNew class that replaces the forward computation of a original PyTorch model with a custom CUDA kernel. The original forward method computes logsumexp over the first dimension of a tensor x with shape (R, B, D), outputting a tensor of shape (B, D). Specifically, it returns torch.logsumexp(x, dim=0).
The model should accept the same inputs as the original, and the forward method should have the same signature.
call → get_gpu_specs
callget_gpu_specs
{}
result ← get_gpu_specs (ok)
okget_gpu_specs
GPU specs for NVIDIA H100 80GB HBM3:
total memory (runtime): 79.1 GB
spec entry: H100
GPU Architecture: Hopper
GPU Memory: 80GB
Memory Bandwidth: 3.35 TB/s
FP64 TFLOPS: 34
FP64 Tensor Core TFLOPS: 67
FP32 TFLOPS: 67
TF32 Tensor Core TFLOPS: 989 with sparsity
BFLOAT16 Tensore Core TFLOPS: 1979 with sparsity
FP16 Tensor Core TFLOPS: 1979 with sparsity
FP8 Tensor Core TFLOPS: 3958 with sparsity
INT8 Tensor Core TOPS: 3958 with sparsity
Register File Size: 64K 32-bit registers per SM
Maximum number of registers per thread: 255
Maximum number of thread blocks per SM: 32
Shared memory capacity per SM: 228 KB
Maximum shared memory per thread block: 227 KB
turn 1246.1s turn · 4.1s LLM · 242.0s tools · tokens in=2,378 · out=524 · total=2,941
reasoning
GPU is H100, high bandwidth, but for this tiny kernel, it's about launch overhead or something, but anyway.
call → compile_kernel
callcompile_kernel
{
"kernel_code": "import torch\nimport torch.nn as nn\nfrom torch.utils.cpp_extension import load_inline\n\ncuda_source = r'''\n#include <torch/extension.h>\n#include <cuda_runtime.h>\n#include <cfloat>\n\n__global__ void logsumexp_kernel(const float* __restrict__ x, float* __restrict__ out, int R, int B, int D) {\n int idx = blockIdx.x * blockDim.x + threadIdx.x;\n int total = B * D;\n if (idx < total) {\n int b = idx / D;\n int d = idx % D;\n float max_val = -FLT_MAX;\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}\n\ntorch::Tensor logsumexp_cuda(torch::Tensor x) {\n TORCH_CHECK(x.dim() == 3, \"x 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 total = B * D;\n int threads = 256;\n int blocks = (total + 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 = r'''\ntorch::Tensor logsumexp_cuda(torch::Tensor x);\n'''\n\nmodule = load_inline(\n name='logsumexp_module',\n cpp_sources=[cpp_source],\n cuda_sources=[cuda_source],\n functions=['logsumexp_cuda'],\n verbose=True,\n extra_cflags=['-O3'],\n extra_cuda_cflags=['-O3']\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 ← compile_kernel (ok)
okcompile_kernel
compile_kernel PASSED: kernel compiled without errors.
turn 2418.2s turn · 5.0s LLM · 413.3s tools · tokens in=2,958 · out=525 · total=3,490
reasoning
Compiled successfully. Now check correctness.
call → run_correctness
callrun_correctness
{
"kernel_code": "import torch\nimport torch.nn as nn\nfrom torch.utils.cpp_extension import load_inline\n\ncuda_source = r'''\n#include <torch/extension.h>\n#include <cuda_runtime.h>\n#include <cfloat>\n\n__global__ void logsumexp_kernel(const float* __restrict__ x, float* __restrict__ out, int R, int B, int D) {\n int idx = blockIdx.x * blockDim.x + threadIdx.x;\n int total = B * D;\n if (idx < total) {\n int b = idx / D;\n int d = idx % D;\n float max_val = -FLT_MAX;\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}\n\ntorch::Tensor logsumexp_cuda(torch::Tensor x) {\n TORCH_CHECK(x.dim() == 3, \"x 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 total = B * D;\n int threads = 256;\n int blocks = (total + 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 = r'''\ntorch::Tensor logsumexp_cuda(torch::Tensor x);\n'''\n\nmodule = load_inline(\n name='logsumexp_module',\n cpp_sources=[cpp_source],\n cuda_sources=[cuda_source],\n functions=['logsumexp_cuda'],\n verbose=False,\n extra_cflags=['-O3'],\n extra_cuda_cflags=['-O3']\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 ← run_correctness (ok)
okrun_correctness
run_correctness PASSED: (5 / 5) trials all matched the reference.
Numerical precision: max_abs_err=0.00e+00 mean_abs_err=0.00e+00 max_rel_err=0.00e+00
turn 3550.0s turn · 7.6s LLM · 542.5s tools · tokens in=3,550 · out=524 · total=4,181 · FINAL
reasoning
Correct. Now, to submit, I need to check if it's faster, but since the task requires submitting once confident.
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 = r'''\n#include <torch/extension.h>\n#include <cuda_runtime.h>\n#include <cfloat>\n\n__global__ void logsumexp_kernel(const float* __restrict__ x, float* __restrict__ out, int R, int B, int D) {\n int idx = blockIdx.x * blockDim.x + threadIdx.x;\n int total = B * D;\n if (idx < total) {\n int b = idx / D;\n int d = idx % D;\n float max_val = -FLT_MAX;\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}\n\ntorch::Tensor logsumexp_cuda(torch::Tensor x) {\n TORCH_CHECK(x.dim() == 3, \"x 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 total = B * D;\n int threads = 256;\n int blocks = (total + 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 = r'''\ntorch::Tensor logsumexp_cuda(torch::Tensor x);\n'''\n\nmodule = load_inline(\n name='logsumexp_module',\n cpp_sources=[cpp_source],\n cuda_sources=[cuda_source],\n functions=['logsumexp_cuda'],\n verbose=False,\n extra_cflags=['-O3'],\n extra_cuda_cflags=['-O3']\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: 33.01 MB peak (ref: 33.01 MB, ratio: 0.9998x)
Kernel launches: 1 (ref: 20, fusion ratio: 20.0)
SOL score: 0.001
final submitted kernel (65 lines, 1,765 chars) — full source, untruncated
import torch
import torch.nn as nn
from torch.utils.cpp_extension import load_inline
cuda_source = r'''
#include <torch/extension.h>
#include <cuda_runtime.h>
#include <cfloat>
__global__ void logsumexp_kernel(const float* __restrict__ x, float* __restrict__ out, int R, int B, int D) {
int idx = blockIdx.x * blockDim.x + threadIdx.x;
int total = B * D;
if (idx < total) {
int b = idx / D;
int d = idx % D;
float max_val = -FLT_MAX;
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, "x 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 total = B * D;
int threads = 256;
int blocks = (total + threads - 1) / threads;
logsumexp_kernel<<<blocks, threads>>>(x.data_ptr<float>(), out.data_ptr<float>(), R, B, D);
return out;
}
'''
cpp_source = r'''
torch::Tensor logsumexp_cuda(torch::Tensor x);
'''
module = load_inline(
name='logsumexp_module',
cpp_sources=[cpp_source],
cuda_sources=[cuda_source],
functions=['logsumexp_cuda'],
verbose=False,
extra_cflags=['-O3'],
extra_cuda_cflags=['-O3']
)
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)