L1 P25 · popcorn

correct 25_virtual_reduce_scatter_masked_sum.py

model

grok-4.3

turns / max

4 / 12

tool calls

tokens (total)

11904

wall clock

2315.8s

truncated

backend

cuda fp32

variant

popcorn

runtime

0.01 μs

ref runtime

0.01 μs

speedup

2.49x

started

2026-05-11T14:55:44Z

timingkernelbench.timing

kernel runtime

0.01 μs

ref runtime

0.01 μs

speedup

2.49x

kernel mean / median / std

0.01 / — / 1.41e-04 μs

ref mean / median / std

0.01 / — / 3.52e-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.00 MB

ref peak

33.00 MB

kernel / ref

1.00x

kernel launches / fusiontorch.profiler

kernel launches

ref launches

fusion ratio (ref / kernel)

7.00x

total cuda time

3 μs

SOL (speed-of-light)nsight_hardware_counters

SOL score (0–1)

6.000e-04

dram util %

0.1%

compute util %

0.0%

bottleneck

unknown

dominant pipe

None

arithmetic intensity

—

ridge point

20.00

roofline / occupancynsight_hardware_counters

dram bandwidth

2.17 GB/s

dram util %

0.1%

fp32 tflops

—

fp32 util %

—

fp16 tflops

—

occupancy %

1.6%

bottleneck

unknown

dominant pipe %

—

arithmetic intensity

—

ridge point

20.00

L1 hit %

—

L2 hit %

99.5%

ld sectors / req

2.00

st sectors / req

—

regs / thread

smem / block

1024 B

block size

peak bw

3350.00 GB/s

peak fp32 tflops

67.00

peak fp16 tflops

1979.00

warp stalls (top 5)nsight

long_scoreboard

1764.3%

short_scoreboard

1473.8%

reference PyTorch source (the prompt the agent sees)

"""
Unequal virtual shard lengths: x is (R, S_MAX) with a fixed boolean mask marking valid cells.
Output (S_MAX,) is sum_r x[r, j] * mask[r, j] (masked sum along the virtual rank dimension).

Padding positions are masked out so they do not contribute.
"""

import torch
import torch.nn as nn

from kernelbench.distributed_collectives import default_device

R = 4
S_MAX = 20

# mask[r, j] == True iff rank r contributes a real value at column j
_MASK_DATA = [
    [1, 1, 1, 1, 1, 1, 1, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
    [1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 0, 0, 0, 0, 0, 0, 0, 0],
    [1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
    [1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 0, 0],
]


class Model(nn.Module):
    def __init__(self):
        super().__init__()
        self.register_buffer(
            "shard_mask",
            torch.tensor(_MASK_DATA, dtype=torch.bool),
            persistent=False,
        )

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        m = self.shard_mask.to(dtype=x.dtype, device=x.device)
        return (x * m).sum(dim=0)


def get_inputs():
    dev = default_device()
    p = popcorn_pri
    return [torch.randn(p.jitter_int(R), p.jitter_int(S_MAX), device=dev)]


def get_init_inputs():
    return []

turn 0678.0s turn · 6.6s LLM · 671.4s tools · tokens in=1,737 · out=6 · total=2,134

reasoning

The task is to replace the forward computation of a PyTorch model with a custom CUDA kernel. The model computes a masked sum along the first dimension of a (4, 20) tensor, using a fixed boolean mask. ...I need to replace the forward computation of a PyTorch model with a custom CUDA kernel.

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 1397.9s turn · 6.2s LLM · 391.8s tools · tokens in=2,367 · out=375 · total=2,827

reasoning

GPU is H100, high bandwidth.

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 masked_sum_kernel(const float* x, float* out, int R, int S) {\n int j = blockIdx.x * blockDim.x + threadIdx.x;\n if (j >= S) return;\n float sum = 0.0f;\n // Hardcoded masks based on _MASK_DATA\n if (j < 8) sum += x[0 * S + j];\n if (j < 12) sum += x[1 * S + j];\n if (j < 10) sum += x[2 * S + j];\n if (j < 18) sum += x[3 * S + j];\n out[j] = sum;\n}\n\ntorch::Tensor masked_sum_cuda(torch::Tensor x) {\n auto S = x.size(1);\n auto out = torch::empty({S}, x.options());\n dim3 block(32);\n dim3 grid((S + 31) / 32);\n masked_sum_kernel<<<grid, block>>>(x.data_ptr<float>(), out.data_ptr<float>(), x.size(0), S);\n return out;\n}\n\"\"\"\n\ncpp_source = \"\"\"\ntorch::Tensor masked_sum_cuda(torch::Tensor x);\n\"\"\"\n\nmodule = load_inline(\n name='masked_sum',\n cpp_sources=cpp_source,\n cuda_sources=cuda_source,\n functions=['masked_sum_cuda'],\n verbose=False\n)\n\nclass ModelNew(nn.Module):\n def __init__(self):\n super().__init__()\n self.masked_sum = module.masked_sum_cuda\n\n def forward(self, x: torch.Tensor) -> torch.Tensor:\n return self.masked_sum(x)\n" }

result ← compile_kernel (ok)

okcompile_kernel compile_kernel PASSED: kernel compiled without errors.

turn 21026.4s turn · 3.7s LLM · 1022.7s tools · tokens in=2,844 · out=376 · total=3,227

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 masked_sum_kernel(const float* x, float* out, int R, int S) {\n int j = blockIdx.x * blockDim.x + threadIdx.x;\n if (j >= S) return;\n float sum = 0.0f;\n // Hardcoded masks based on _MASK_DATA\n if (j < 8) sum += x[0 * S + j];\n if (j < 12) sum += x[1 * S + j];\n if (j < 10) sum += x[2 * S + j];\n if (j < 18) sum += x[3 * S + j];\n out[j] = sum;\n}\n\ntorch::Tensor masked_sum_cuda(torch::Tensor x) {\n auto S = x.size(1);\n auto out = torch::empty({S}, x.options());\n dim3 block(32);\n dim3 grid((S + 31) / 32);\n masked_sum_kernel<<<grid, block>>>(x.data_ptr<float>(), out.data_ptr<float>(), x.size(0), S);\n return out;\n}\n\"\"\"\n\ncpp_source = \"\"\"\ntorch::Tensor masked_sum_cuda(torch::Tensor x);\n\"\"\"\n\nmodule = load_inline(\n name='masked_sum',\n cpp_sources=cpp_source,\n cuda_sources=cuda_source,\n functions=['masked_sum_cuda'],\n verbose=False\n)\n\nclass ModelNew(nn.Module):\n def __init__(self):\n super().__init__()\n self.masked_sum = module.masked_sum_cuda\n\n def forward(self, x: torch.Tensor) -> torch.Tensor:\n return self.masked_sum(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 3213.5s turn · 6.6s LLM · 206.9s tools · tokens in=3,287 · out=375 · total=3,716 · FINAL

reasoning

Correct! Now, to submit, but is it faster? The task requires faster than reference, and since submit will check timing, I need to make sure.

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 masked_sum_kernel(const float* x, float* out, int R, int S) {\n int j = blockIdx.x * blockDim.x + threadIdx.x;\n if (j >= S) return;\n float sum = 0.0f;\n // Hardcoded masks based on _MASK_DATA\n if (j < 8) sum += x[0 * S + j];\n if (j < 12) sum += x[1 * S + j];\n if (j < 10) sum += x[2 * S + j];\n if (j < 18) sum += x[3 * S + j];\n out[j] = sum;\n}\n\ntorch::Tensor masked_sum_cuda(torch::Tensor x) {\n auto S = x.size(1);\n auto out = torch::empty({S}, x.options());\n dim3 block(32);\n dim3 grid((S + 31) / 32);\n masked_sum_kernel<<<grid, block>>>(x.data_ptr<float>(), out.data_ptr<float>(), x.size(0), S);\n return out;\n}\n\"\"\"\n\ncpp_source = \"\"\"\ntorch::Tensor masked_sum_cuda(torch::Tensor x);\n\"\"\"\n\nmodule = load_inline(\n name='masked_sum',\n cpp_sources=cpp_source,\n cuda_sources=cuda_source,\n functions=['masked_sum_cuda'],\n verbose=False\n)\n\nclass ModelNew(nn.Module):\n def __init__(self):\n super().__init__()\n self.masked_sum = module.masked_sum_cuda\n\n def forward(self, x: torch.Tensor) -> torch.Tensor:\n return self.masked_sum(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.0 MB peak (ref: 33.0 MB, ratio: 1.0x) Kernel launches: 1 (ref: 7, fusion ratio: 7.0) SOL score: 0.0006

final submitted kernel (47 lines, 1,228 chars) — full source, untruncated

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

cuda_source = """
__global__ void masked_sum_kernel(const float* x, float* out, int R, int S) {
    int j = blockIdx.x * blockDim.x + threadIdx.x;
    if (j >= S) return;
    float sum = 0.0f;
    // Hardcoded masks based on _MASK_DATA
    if (j < 8) sum += x[0 * S + j];
    if (j < 12) sum += x[1 * S + j];
    if (j < 10) sum += x[2 * S + j];
    if (j < 18) sum += x[3 * S + j];
    out[j] = sum;
}

torch::Tensor masked_sum_cuda(torch::Tensor x) {
    auto S = x.size(1);
    auto out = torch::empty({S}, x.options());
    dim3 block(32);
    dim3 grid((S + 31) / 32);
    masked_sum_kernel<<<grid, block>>>(x.data_ptr<float>(), out.data_ptr<float>(), x.size(0), S);
    return out;
}
"""

cpp_source = """
torch::Tensor masked_sum_cuda(torch::Tensor x);
"""

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

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

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