L1 P23 · popcorn

correct 23_virtual_all_reduce_mean.py

model

grok-4.3

turns / max

3 / 15

tool calls

tokens (total)

10364

wall clock

135.6s

truncated

backend

cuda fp32

variant

popcorn

runtime

0.01 μs

ref runtime

0.01 μs

speedup

1.18x

started

2026-05-11T21:32:45Z

timingkernelbench.timing

kernel runtime

0.01 μs

ref runtime

0.01 μs

speedup

1.18x

kernel mean / median / std

0.01 / — / 9.81e-05 μs

ref mean / median / std

0.01 / — / 4.94e-04 μs

numerical precisioneval_kernel_against_ref

max abs error

0.00

mean abs error

0.00

max rel error

0.00

mean rel error

0.00

gpu memoryextended_metrics.measure_memory

kernel peak

33.01 MB

ref peak

33.01 MB

kernel / ref

1.00x

kernel launches / fusiontorch.profiler

kernel launches

ref launches

fusion ratio (ref / kernel)

2.00x

total cuda time

23 μs

SOL (speed-of-light)nsight_hardware_counters

SOL score (0–1)

0.002

dram util %

0.2%

compute util %

0.0%

bottleneck

unknown

dominant pipe

None

arithmetic intensity

—

ridge point

20.00

roofline / occupancynsight_hardware_counters

dram bandwidth

6.20 GB/s

dram util %

0.2%

fp32 tflops

—

fp32 util %

—

fp16 tflops

—

occupancy %

12.3%

bottleneck

unknown

dominant pipe %

—

arithmetic intensity

—

ridge point

20.00

L1 hit %

—

L2 hit %

97.7%

ld sectors / req

4.00

st sectors / req

—

regs / thread

smem / block

1024 B

block size

256

peak bw

3350.00 GB/s

peak fp32 tflops

67.00

peak fp16 tflops

1979.00

warp stalls (top 5)nsight

long_scoreboard

753.9%

short_scoreboard

121.6%

mio_throttle

1.1%

reference PyTorch source (the prompt the agent sees)

"""
Virtual ranks: input packs R per-rank tensors along dim 0. Output is the mean over ranks (B, D).

No distributed runtime required; the reference always reduces along dim 0.
"""

import torch
import torch.nn as nn

from kernelbench.distributed_collectives import default_device

R = 4
B = 16
D = 48


class Model(nn.Module):
    """x shape (R, B, D) -> mean over R -> (B, D)."""

    def __init__(self):
        super().__init__()

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return x.mean(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 0110.8s turn · 12.7s LLM · 98.1s tools · tokens in=1,957 · out=386 · total=3,000

reasoning

The task is to replace the forward computation of a PyTorch model with a custom CUDA kernel. The model computes the mean over dimension 0 of a tensor with shape (R, B, D), where R=4, B=16, D=48. I need to replace the forward computation of a PyTorch model with a custom CUDA kernel.

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 = '''\n__global__ void mean_dim0(const float* __restrict__ input, float* __restrict__ output, int R, int B, int D) {\n int idx = blockIdx.x * blockDim.x + threadIdx.x;\n if (idx < B * D) {\n int b = idx / D;\n int d = idx % D;\n float sum = 0.0f;\n for(int r = 0; r < R; r++) {\n sum += input[(r * B + b) * D + d];\n }\n output[idx] = sum / static_cast<float>(R);\n }\n}\n\ntorch::Tensor mean_dim0_cuda(torch::Tensor x) {\n auto R = x.size(0);\n auto B = x.size(1);\n auto D = x.size(2);\n auto output = torch::empty({B, D}, x.options());\n int total = B * D;\n int threads = 256;\n int blocks = (total + threads - 1) / threads;\n mean_dim0<<<blocks, threads>>>(x.data_ptr<float>(), output.data_ptr<float>(), R, B, D);\n return output;\n}\n'''\n\ncpp_source = '''\ntorch::Tensor mean_dim0_cuda(torch::Tensor x);\n'''\n\nmodule = load_inline(\n name='mean_module',\n cpp_sources=[cpp_source],\n cuda_sources=[cuda_source],\n functions=['mean_dim0_cuda'],\n verbose=False\n)\n\nclass ModelNew(nn.Module):\n def __init__(self):\n super().__init__()\n self.op = module\n\n def forward(self, x: torch.Tensor) -> torch.Tensor:\n return self.op.mean_dim0_cuda(x)\n" }

result ← compile_kernel (ok)

okcompile_kernel compile_kernel PASSED: kernel compiled without errors.

turn 15.6s turn · 5.5s LLM · 0.0s tools · tokens in=3,017 · out=387 · total=3,411

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 = '''\n__global__ void mean_dim0(const float* __restrict__ input, float* __restrict__ output, int R, int B, int D) {\n int idx = blockIdx.x * blockDim.x + threadIdx.x;\n if (idx < B * D) {\n int b = idx / D;\n int d = idx % D;\n float sum = 0.0f;\n for(int r = 0; r < R; r++) {\n sum += input[(r * B + b) * D + d];\n }\n output[idx] = sum / static_cast<float>(R);\n }\n}\n\ntorch::Tensor mean_dim0_cuda(torch::Tensor x) {\n auto R = x.size(0);\n auto B = x.size(1);\n auto D = x.size(2);\n auto output = torch::empty({B, D}, x.options());\n int total = B * D;\n int threads = 256;\n int blocks = (total + threads - 1) / threads;\n mean_dim0<<<blocks, threads>>>(x.data_ptr<float>(), output.data_ptr<float>(), R, B, D);\n return output;\n}\n'''\n\ncpp_source = '''\ntorch::Tensor mean_dim0_cuda(torch::Tensor x);\n'''\n\nmodule = load_inline(\n name='mean_module',\n cpp_sources=[cpp_source],\n cuda_sources=[cuda_source],\n functions=['mean_dim0_cuda'],\n verbose=False\n)\n\nclass ModelNew(nn.Module):\n def __init__(self):\n super().__init__()\n self.op = module\n\n def forward(self, x: torch.Tensor) -> torch.Tensor:\n return self.op.mean_dim0_cuda(x)\n" }

result ← run_correctness (ok)

okrun_correctness run_correctness PASSED: (3 / 3) 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 219.3s turn · 6.3s LLM · 13.0s tools · tokens in=3,471 · out=386 · total=3,953 · FINAL

reasoning

Correct. Now, to check if faster, I should profile or submit, but since it's the final, and small, likely faster as it's specialized.

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__global__ void mean_dim0(const float* __restrict__ input, float* __restrict__ output, int R, int B, int D) {\n int idx = blockIdx.x * blockDim.x + threadIdx.x;\n if (idx < B * D) {\n int b = idx / D;\n int d = idx % D;\n float sum = 0.0f;\n for(int r = 0; r < R; r++) {\n sum += input[(r * B + b) * D + d];\n }\n output[idx] = sum / static_cast<float>(R);\n }\n}\n\ntorch::Tensor mean_dim0_cuda(torch::Tensor x) {\n auto R = x.size(0);\n auto B = x.size(1);\n auto D = x.size(2);\n auto output = torch::empty({B, D}, x.options());\n int total = B * D;\n int threads = 256;\n int blocks = (total + threads - 1) / threads;\n mean_dim0<<<blocks, threads>>>(x.data_ptr<float>(), output.data_ptr<float>(), R, B, D);\n return output;\n}\n'''\n\ncpp_source = '''\ntorch::Tensor mean_dim0_cuda(torch::Tensor x);\n'''\n\nmodule = load_inline(\n name='mean_module',\n cpp_sources=[cpp_source],\n cuda_sources=[cuda_source],\n functions=['mean_dim0_cuda'],\n verbose=False\n)\n\nclass ModelNew(nn.Module):\n def __init__(self):\n super().__init__()\n self.op = module\n\n def forward(self, x: torch.Tensor) -> torch.Tensor:\n return self.op.mean_dim0_cuda(x)\n" }

result ← submit_kernel (ok)

oksubmit_kernel submit_kernel PASSED: (3 / 3) 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: 1.0x) Kernel launches: 1 (ref: 2, fusion ratio: 2.0) SOL score: 0.0019

final submitted kernel (51 lines, 1,340 chars) — full source, untruncated

import torch
import torch.nn as nn
from torch.utils.cpp_extension import load_inline

cuda_source = '''
__global__ void mean_dim0(const float* __restrict__ input, float* __restrict__ output, int R, int B, int D) {
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    if (idx < B * D) {
        int b = idx / D;
        int d = idx % D;
        float sum = 0.0f;
        for(int r = 0; r < R; r++) {
            sum += input[(r * B + b) * D + d];
        }
        output[idx] = sum / static_cast<float>(R);
    }
}

torch::Tensor mean_dim0_cuda(torch::Tensor x) {
    auto R = x.size(0);
    auto B = x.size(1);
    auto D = x.size(2);
    auto output = torch::empty({B, D}, x.options());
    int total = B * D;
    int threads = 256;
    int blocks = (total + threads - 1) / threads;
    mean_dim0<<<blocks, threads>>>(x.data_ptr<float>(), output.data_ptr<float>(), R, B, D);
    return output;
}
'''

cpp_source = '''
torch::Tensor mean_dim0_cuda(torch::Tensor x);
'''

module = load_inline(
    name='mean_module',
    cpp_sources=[cpp_source],
    cuda_sources=[cuda_source],
    functions=['mean_dim0_cuda'],
    verbose=False
)

class ModelNew(nn.Module):
    def __init__(self):
        super().__init__()
        self.op = module

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return self.op.mean_dim0_cuda(x)