grok-4.3 · popcorn · pop_l123_all_gpt
correct 30_broadcast_masked_from_row0.py
model
grok-4.3
turns / max
4 / 15
tool calls
4
tokens (total)
14044
wall clock
122.1s
truncated
no
backend
cuda fp32
variant
popcorn
runtime
0.01 μs
ref runtime
0.01 μs
speedup
1.05x
started
2026-05-11T19:07:11Z
timingkernelbench.timing
kernel runtime
0.01 μs
ref runtime
0.01 μs
speedup
1.05x
kernel mean / median / std
0.01 / — / 1.42e-04 μs
ref mean / median / std
0.01 / — / 1.45e-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
0.01 MB
ref peak
0.01 MB
kernel / ref
1.00x
kernel launches / fusiontorch.profiler
kernel launches
3
ref launches
11
fusion ratio (ref / kernel)
3.67x
total cuda time
70 μs
SOL (speed-of-light)nsight_hardware_counters
SOL score (0–1)
5.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
1.83 GB/s
dram util %
0.1%
fp32 tflops
—
fp32 util %
—
fp16 tflops
—
occupancy %
10.4%
bottleneck
unknown
dominant pipe %
—
arithmetic intensity
—
ridge point
20.00
L1 hit %
—
L2 hit %
99.3%
ld sectors / req
3.53
st sectors / req
—
regs / thread
16
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
2471.5%
short_scoreboard
809.6%
reference PyTorch source (the prompt the agent sees)
"""
Simulated broadcast of row 0 into masked columns for all rows: where mask[j] is True,
every row matches x[0, j]; elsewhere keep x[b, j].
"""

import torch
import torch.nn as nn

from kernelbench.distributed_collectives import default_device

B = 18
D = 44

# Fixed mask (not all True, not all False)
_MASK = torch.tensor(
    [1, 0, 1, 1, 0, 0, 1, 0, 1, 1, 0, 1, 1, 0, 1, 0, 1, 1, 0, 1, 1, 0, 1, 0, 1, 1, 1, 0, 1, 0, 1, 1, 0, 1, 1, 0, 1, 0, 1, 1, 0, 1, 0, 1],
    dtype=torch.bool,
)


class Model(nn.Module):
    def __init__(self):
        super().__init__()
        self.register_buffer("row_mask", _MASK.clone(), persistent=False)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        m = self.row_mask.to(device=x.device).view(1, D).expand(B, D)
        row0 = x[0:1].expand(B, D)
        return torch.where(m, row0, x)


def get_inputs():
    dev = default_device()
    return [torch.randn(popcorn_pri.jitter_int(B), D, device=dev)]


def get_init_inputs():
    return []
turn 013.7s turn · 6.0s LLM · 7.7s tools · tokens in=2,137 · out=6 · total=2,460
reasoning
The task is to replace the forward computation with a custom CUDA kernel that's faster than the PyTorch reference.
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 179.2s turn · 6.3s LLM · 72.9s tools · tokens in=2,693 · out=539 · total=3,263
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#include <torch/extension.h>\n#include <cuda_runtime.h>\n\n__global__ void broadcast_mask_kernel(const float* __restrict__ x, const bool* __restrict__ mask, float* __restrict__ out, int B, int D) {\n int idx = blockIdx.x * blockDim.x + threadIdx.x;\n int total = B * D;\n if (idx < total) {\n int j = idx % D;\n if (mask[j]) {\n out[idx] = x[j];\n } else {\n out[idx] = x[idx];\n }\n }\n}\n\ntorch::Tensor broadcast_mask_cuda(torch::Tensor x, torch::Tensor mask) {\n int B = x.size(0);\n int D = x.size(1);\n auto out = torch::empty_like(x);\n int total = B * D;\n int threads = 256;\n int blocks = (total + threads - 1) / threads;\n broadcast_mask_kernel<<<blocks, threads>>>(x.data_ptr<float>(), mask.data_ptr<bool>(), out.data_ptr<float>(), B, D);\n return out;\n}\n'''\n\ncpp_source = '''\ntorch::Tensor broadcast_mask_cuda(torch::Tensor x, torch::Tensor mask);\n'''\n\nmodule = load_inline(\n name='broadcast_mask',\n cpp_sources=[cpp_source],\n cuda_sources=[cuda_source],\n functions=['broadcast_mask_cuda'],\n verbose=False,\n with_cuda=True\n)\n\nclass ModelNew(nn.Module):\n def __init__(self):\n super().__init__()\n mask = torch.tensor(\n [1, 0, 1, 1, 0, 0, 1, 0, 1, 1, 0, 1, 1, 0, 1, 0, 1, 1, 0, 1, 1, 0, 1, 0, 1, 1, 1, 0, 1, 0, 1, 1, 0, 1, 1, 0, 1, 0, 1, 1, 0, 1, 0, 1],\n dtype=torch.bool,\n )\n self.register_buffer(\"row_mask\", mask, persistent=False)\n\n def forward(self, x: torch.Tensor) -> torch.Tensor:\n mask = self.row_mask.to(device=x.device)\n return module.broadcast_mask_cuda(x, mask)\n" }
result ← compile_kernel (ok)
okcompile_kernel compile_kernel PASSED: kernel compiled without errors.
turn 26.1s turn · 5.9s LLM · 0.2s tools · tokens in=3,280 · out=540 · total=3,827
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#include <torch/extension.h>\n#include <cuda_runtime.h>\n\n__global__ void broadcast_mask_kernel(const float* __restrict__ x, const bool* __restrict__ mask, float* __restrict__ out, int B, int D) {\n int idx = blockIdx.x * blockDim.x + threadIdx.x;\n int total = B * D;\n if (idx < total) {\n int j = idx % D;\n if (mask[j]) {\n out[idx] = x[j];\n } else {\n out[idx] = x[idx];\n }\n }\n}\n\ntorch::Tensor broadcast_mask_cuda(torch::Tensor x, torch::Tensor mask) {\n int B = x.size(0);\n int D = x.size(1);\n auto out = torch::empty_like(x);\n int total = B * D;\n int threads = 256;\n int blocks = (total + threads - 1) / threads;\n broadcast_mask_kernel<<<blocks, threads>>>(x.data_ptr<float>(), mask.data_ptr<bool>(), out.data_ptr<float>(), B, D);\n return out;\n}\n'''\n\ncpp_source = '''\ntorch::Tensor broadcast_mask_cuda(torch::Tensor x, torch::Tensor mask);\n'''\n\nmodule = load_inline(\n name='broadcast_mask',\n cpp_sources=[cpp_source],\n cuda_sources=[cuda_source],\n functions=['broadcast_mask_cuda'],\n verbose=False,\n with_cuda=True\n)\n\nclass ModelNew(nn.Module):\n def __init__(self):\n super().__init__()\n mask = torch.tensor(\n [1, 0, 1, 1, 0, 0, 1, 0, 1, 1, 0, 1, 1, 0, 1, 0, 1, 1, 0, 1, 1, 0, 1, 0, 1, 1, 1, 0, 1, 0, 1, 1, 0, 1, 1, 0, 1, 0, 1, 1, 0, 1, 0, 1],\n dtype=torch.bool,\n )\n self.register_buffer(\"row_mask\", mask, persistent=False)\n\n def forward(self, x: torch.Tensor) -> torch.Tensor:\n mask = self.row_mask.to(device=x.device)\n return module.broadcast_mask_cuda(x, mask)\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 323.2s turn · 5.1s LLM · 18.1s tools · tokens in=3,887 · out=539 · total=4,494 · FINAL
reasoning
Correct. Now, to check if faster, but since submit will do timing, and it's likely faster as PyTorch has overhead for small tensors.
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_runtime.h>\n\n__global__ void broadcast_mask_kernel(const float* __restrict__ x, const bool* __restrict__ mask, float* __restrict__ out, int B, int D) {\n int idx = blockIdx.x * blockDim.x + threadIdx.x;\n int total = B * D;\n if (idx < total) {\n int j = idx % D;\n if (mask[j]) {\n out[idx] = x[j];\n } else {\n out[idx] = x[idx];\n }\n }\n}\n\ntorch::Tensor broadcast_mask_cuda(torch::Tensor x, torch::Tensor mask) {\n int B = x.size(0);\n int D = x.size(1);\n auto out = torch::empty_like(x);\n int total = B * D;\n int threads = 256;\n int blocks = (total + threads - 1) / threads;\n broadcast_mask_kernel<<<blocks, threads>>>(x.data_ptr<float>(), mask.data_ptr<bool>(), out.data_ptr<float>(), B, D);\n return out;\n}\n'''\n\ncpp_source = '''\ntorch::Tensor broadcast_mask_cuda(torch::Tensor x, torch::Tensor mask);\n'''\n\nmodule = load_inline(\n name='broadcast_mask',\n cpp_sources=[cpp_source],\n cuda_sources=[cuda_source],\n functions=['broadcast_mask_cuda'],\n verbose=False,\n with_cuda=True\n)\n\nclass ModelNew(nn.Module):\n def __init__(self):\n super().__init__()\n mask = torch.tensor(\n [1, 0, 1, 1, 0, 0, 1, 0, 1, 1, 0, 1, 1, 0, 1, 0, 1, 1, 0, 1, 1, 0, 1, 0, 1, 1, 1, 0, 1, 0, 1, 1, 0, 1, 1, 0, 1, 0, 1, 1, 0, 1, 0, 1],\n dtype=torch.bool,\n )\n self.register_buffer(\"row_mask\", mask, persistent=False)\n\n def forward(self, x: torch.Tensor) -> torch.Tensor:\n mask = self.row_mask.to(device=x.device)\n return module.broadcast_mask_cuda(x, mask)\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: 0.01 MB peak (ref: 0.01 MB, ratio: 1.0x) Kernel launches: 3 (ref: 11, fusion ratio: 3.6667) SOL score: 0.0005
final submitted kernel (59 lines, 1,738 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_runtime.h>

__global__ void broadcast_mask_kernel(const float* __restrict__ x, const bool* __restrict__ mask, float* __restrict__ out, int B, int D) {
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    int total = B * D;
    if (idx < total) {
        int j = idx % D;
        if (mask[j]) {
            out[idx] = x[j];
        } else {
            out[idx] = x[idx];
        }
    }
}

torch::Tensor broadcast_mask_cuda(torch::Tensor x, torch::Tensor mask) {
    int B = x.size(0);
    int D = x.size(1);
    auto out = torch::empty_like(x);
    int total = B * D;
    int threads = 256;
    int blocks = (total + threads - 1) / threads;
    broadcast_mask_kernel<<<blocks, threads>>>(x.data_ptr<float>(), mask.data_ptr<bool>(), out.data_ptr<float>(), B, D);
    return out;
}
'''

cpp_source = '''
torch::Tensor broadcast_mask_cuda(torch::Tensor x, torch::Tensor mask);
'''

module = load_inline(
    name='broadcast_mask',
    cpp_sources=[cpp_source],
    cuda_sources=[cuda_source],
    functions=['broadcast_mask_cuda'],
    verbose=False,
    with_cuda=True
)

class ModelNew(nn.Module):
    def __init__(self):
        super().__init__()
        mask = torch.tensor(
            [1, 0, 1, 1, 0, 0, 1, 0, 1, 1, 0, 1, 1, 0, 1, 0, 1, 1, 0, 1, 1, 0, 1, 0, 1, 1, 1, 0, 1, 0, 1, 1, 0, 1, 1, 0, 1, 0, 1, 1, 0, 1, 0, 1],
            dtype=torch.bool,
        )
        self.register_buffer("row_mask", mask, persistent=False)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        mask = self.row_mask.to(device=x.device)
        return module.broadcast_mask_cuda(x, mask)