L2 P47 · popcorn

correct 47_virtual_all_gather_concat_linear.py

model

grok-4.3

turns / max

4 / 15

tool calls

tokens (total)

14530

wall clock

357.0s

truncated

backend

cuda fp32

variant

popcorn

runtime

0.01 μs

ref runtime

0.01 μs

speedup

0.82x

started

2026-05-11T18:05:12Z

timingkernelbench.timing

kernel runtime

0.01 μs

ref runtime

0.01 μs

speedup

0.82x

kernel mean / median / std

0.01 / — / 3.21e-04 μs

ref mean / median / std

0.01 / — / 6.69e-04 μs

numerical precisioneval_kernel_against_ref

max abs error

8.34e-07

mean abs error

9.15e-08

max rel error

2.98e-05

mean rel error

5.65e-07

gpu memoryextended_metrics.measure_memory

kernel peak

33.06 MB

ref peak

33.07 MB

kernel / ref

1.00x

kernel launches / fusiontorch.profiler

kernel launches

ref launches

fusion ratio (ref / kernel)

14.00x

total cuda time

24 μs

SOL (speed-of-light)nsight_hardware_counters

SOL score (0–1)

0.001

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

3.74 GB/s

dram util %

0.1%

fp32 tflops

—

fp32 util %

—

fp16 tflops

—

occupancy %

3.1%

bottleneck

unknown

dominant pipe %

—

arithmetic intensity

—

ridge point

20.00

L1 hit %

—

L2 hit %

96.9%

ld sectors / req

14.46

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

2180.8%

short_scoreboard

48.1%

reference PyTorch source (the prompt the agent sees)

"""
Virtual all-gather on the last dim: stack R tensors of shape (B, H) into (B, R*H), then Linear.
"""

import torch
import torch.nn as nn

from kernelbench.distributed_collectives import default_device

R = 4
B = 12
H = 32
K = 56


class Model(nn.Module):
    def __init__(self):
        super().__init__()
        self.proj = nn.Linear(R * H, K, bias=True)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        r, b, h = x.shape
        assert r == R and h == H
        parts = [x[i] for i in range(R)]
        gathered = torch.cat(parts, dim=-1)
        return self.proj(gathered)


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 010.5s turn · 10.5s LLM · 0.0s tools · tokens in=1,989 · out=6 · total=2,556

reasoning

The task is to replace the forward computation of this PyTorch model with a custom CUDA kernel. The model does a virtual all-gather on the last dim: stack R tensors of shape (B, H) into (B, R*H), then...I need to replace the forward computation of the PyTorch model with a custom CUDA kernel in a new ModelNew class. It should take the same inputs as the original Model, have the same forward signature, and produce the same output numerically, but faster.

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 1159.8s turn · 9.4s LLM · 150.4s tools · tokens in=2,789 · out=559 · total=3,379

reasoning

H100, high bandwidth. Now, I need to develop the 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 = r'''\n#include <torch/extension.h>\n#include <cuda_runtime.h>\n\n__global__ void fused_allgather_linear_kernel(\n const float* __restrict__ x, // (R, B, H)\n const float* __restrict__ weight, // (K, R*H)\n const float* __restrict__ bias, // (K,)\n float* __restrict__ out, // (B, K)\n int R, int B, int H, int K\n) {\n int b = blockIdx.x;\n int k = threadIdx.x;\n if (b >= B || k >= K) return;\n\n float sum = bias[k];\n const int RH = R * H;\n for (int r = 0; r < R; ++r) {\n for (int h = 0; h < H; ++h) {\n float val = x[r * B * H + b * H + h];\n float w = weight[k * RH + r * H + h];\n sum += val * w;\n }\n }\n out[b * K + k] = sum;\n}\n\ntorch::Tensor fused_allgather_linear(\n torch::Tensor x,\n torch::Tensor weight,\n torch::Tensor bias\n) {\n int R = x.size(0);\n int B = x.size(1);\n int H = x.size(2);\n int K = weight.size(0);\n auto out = torch::empty({B, K}, x.options());\n dim3 grid(B);\n dim3 block(K);\n fused_allgather_linear_kernel<<<grid, block>>>(\n x.data_ptr<float>(),\n weight.data_ptr<float>(),\n bias.data_ptr<float>(),\n out.data_ptr<float>(),\n R, B, H, K\n );\n return out;\n}\n'''\n\ncpp_source = '''\ntorch::Tensor fused_allgather_linear(torch::Tensor x, torch::Tensor weight, torch::Tensor bias);\n'''\n\nmodule = load_inline(\n name='fused_module',\n cpp_sources=cpp_source,\n cuda_sources=cuda_source,\n functions=['fused_allgather_linear'],\n verbose=False\n)\n\nclass ModelNew(nn.Module):\n def __init__(self):\n super().__init__()\n self.proj = nn.Linear(4 * 32, 56, bias=True)\n\n def forward(self, x: torch.Tensor) -> torch.Tensor:\n weight = self.proj.weight\n bias = self.proj.bias\n return module.fused_allgather_linear(x, weight, bias)\n" }

result ← compile_kernel (ok)

okcompile_kernel compile_kernel PASSED: kernel compiled without errors.

turn 2162.3s turn · 10.4s LLM · 151.9s tools · tokens in=3,396 · out=560 · total=3,963

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\n__global__ void fused_allgather_linear_kernel(\n const float* __restrict__ x, // (R, B, H)\n const float* __restrict__ weight, // (K, R*H)\n const float* __restrict__ bias, // (K,)\n float* __restrict__ out, // (B, K)\n int R, int B, int H, int K\n) {\n int b = blockIdx.x;\n int k = threadIdx.x;\n if (b >= B || k >= K) return;\n\n float sum = bias[k];\n const int RH = R * H;\n for (int r = 0; r < R; ++r) {\n for (int h = 0; h < H; ++h) {\n float val = x[r * B * H + b * H + h];\n float w = weight[k * RH + r * H + h];\n sum += val * w;\n }\n }\n out[b * K + k] = sum;\n}\n\ntorch::Tensor fused_allgather_linear(\n torch::Tensor x,\n torch::Tensor weight,\n torch::Tensor bias\n) {\n int R = x.size(0);\n int B = x.size(1);\n int H = x.size(2);\n int K = weight.size(0);\n auto out = torch::empty({B, K}, x.options());\n dim3 grid(B);\n dim3 block(K);\n fused_allgather_linear_kernel<<<grid, block>>>(\n x.data_ptr<float>(),\n weight.data_ptr<float>(),\n bias.data_ptr<float>(),\n out.data_ptr<float>(),\n R, B, H, K\n );\n return out;\n}\n'''\n\ncpp_source = '''\ntorch::Tensor fused_allgather_linear(torch::Tensor x, torch::Tensor weight, torch::Tensor bias);\n'''\n\nmodule = load_inline(\n name='fused_module',\n cpp_sources=cpp_source,\n cuda_sources=cuda_source,\n functions=['fused_allgather_linear'],\n verbose=False\n)\n\nclass ModelNew(nn.Module):\n def __init__(self):\n super().__init__()\n self.proj = nn.Linear(4 * 32, 56, bias=True)\n\n def forward(self, x: torch.Tensor) -> torch.Tensor:\n weight = self.proj.weight\n bias = self.proj.bias\n return module.fused_allgather_linear(x, weight, bias)\n" }

result ← run_correctness (ok)

okrun_correctness run_correctness PASSED: (3 / 3) trials all matched the reference. Numerical precision: max_abs_err=8.34e-07 mean_abs_err=9.15e-08 max_rel_err=2.98e-05

turn 324.4s turn · 11.8s LLM · 12.6s tools · tokens in=4,023 · out=559 · total=4,632 · FINAL

reasoning

Correct! Now, I should check if it's faster, but since submit will do timing, and it's simple, probably is, as it avoids cat and uses direct access.

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\n__global__ void fused_allgather_linear_kernel(\n const float* __restrict__ x, // (R, B, H)\n const float* __restrict__ weight, // (K, R*H)\n const float* __restrict__ bias, // (K,)\n float* __restrict__ out, // (B, K)\n int R, int B, int H, int K\n) {\n int b = blockIdx.x;\n int k = threadIdx.x;\n if (b >= B || k >= K) return;\n\n float sum = bias[k];\n const int RH = R * H;\n for (int r = 0; r < R; ++r) {\n for (int h = 0; h < H; ++h) {\n float val = x[r * B * H + b * H + h];\n float w = weight[k * RH + r * H + h];\n sum += val * w;\n }\n }\n out[b * K + k] = sum;\n}\n\ntorch::Tensor fused_allgather_linear(\n torch::Tensor x,\n torch::Tensor weight,\n torch::Tensor bias\n) {\n int R = x.size(0);\n int B = x.size(1);\n int H = x.size(2);\n int K = weight.size(0);\n auto out = torch::empty({B, K}, x.options());\n dim3 grid(B);\n dim3 block(K);\n fused_allgather_linear_kernel<<<grid, block>>>(\n x.data_ptr<float>(),\n weight.data_ptr<float>(),\n bias.data_ptr<float>(),\n out.data_ptr<float>(),\n R, B, H, K\n );\n return out;\n}\n'''\n\ncpp_source = '''\ntorch::Tensor fused_allgather_linear(torch::Tensor x, torch::Tensor weight, torch::Tensor bias);\n'''\n\nmodule = load_inline(\n name='fused_module',\n cpp_sources=cpp_source,\n cuda_sources=cuda_source,\n functions=['fused_allgather_linear'],\n verbose=False\n)\n\nclass ModelNew(nn.Module):\n def __init__(self):\n super().__init__()\n self.proj = nn.Linear(4 * 32, 56, bias=True)\n\n def forward(self, x: torch.Tensor) -> torch.Tensor:\n weight = self.proj.weight\n bias = self.proj.bias\n return module.fused_allgather_linear(x, weight, bias)\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=8.34e-07 mean_abs_err=9.15e-08 max_rel_err=2.98e-05 Memory: 33.06 MB peak (ref: 33.07 MB, ratio: 0.9998x) Kernel launches: 1 (ref: 14, fusion ratio: 14.0) SOL score: 0.0011

final submitted kernel (76 lines, 1,941 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>

__global__ void fused_allgather_linear_kernel(
    const float* __restrict__ x,  // (R, B, H)
    const float* __restrict__ weight,  // (K, R*H)
    const float* __restrict__ bias,    // (K,)
    float* __restrict__ out,           // (B, K)
    int R, int B, int H, int K
) {
    int b = blockIdx.x;
    int k = threadIdx.x;
    if (b >= B || k >= K) return;

    float sum = bias[k];
    const int RH = R * H;
    for (int r = 0; r < R; ++r) {
        for (int h = 0; h < H; ++h) {
            float val = x[r * B * H + b * H + h];
            float w = weight[k * RH + r * H + h];
            sum += val * w;
        }
    }
    out[b * K + k] = sum;
}

torch::Tensor fused_allgather_linear(
    torch::Tensor x,
    torch::Tensor weight,
    torch::Tensor bias
) {
    int R = x.size(0);
    int B = x.size(1);
    int H = x.size(2);
    int K = weight.size(0);
    auto out = torch::empty({B, K}, x.options());
    dim3 grid(B);
    dim3 block(K);
    fused_allgather_linear_kernel<<<grid, block>>>(
        x.data_ptr<float>(),
        weight.data_ptr<float>(),
        bias.data_ptr<float>(),
        out.data_ptr<float>(),
        R, B, H, K
    );
    return out;
}
'''

cpp_source = '''
torch::Tensor fused_allgather_linear(torch::Tensor x, torch::Tensor weight, torch::Tensor bias);
'''

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

class ModelNew(nn.Module):
    def __init__(self):
        super().__init__()
        self.proj = nn.Linear(4 * 32, 56, bias=True)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        weight = self.proj.weight
        bias = self.proj.bias
        return module.fused_allgather_linear(x, weight, bias)