L1 P27 · popcorn

incorrect 27_permutation_all_to_all.py

model

grok-4.3

turns / max

1 / 1

tool calls

tokens (total)

1969

wall clock

337.3s

truncated

backend

cuda fp32

variant

popcorn

runtime

—

ref runtime

0.03 μs

speedup

—

started

2026-05-11T09:29:55Z

correctness issue:

timingkernelbench.timing

kernel runtime

—

ref runtime

0.03 μs

speedup

—

kernel mean / median / std

—

ref mean / median / std

0.03 / — / 0.00 μs

roofline / occupancynsight_hardware_counters

dram bandwidth

—

dram util %

—

fp32 tflops

—

fp32 util %

—

fp16 tflops

—

occupancy %

—

bottleneck

unknown

dominant pipe %

—

arithmetic intensity

—

ridge point

20.00

L1 hit %

—

L2 hit %

—

ld sectors / req

—

st sectors / req

—

regs / thread

—

smem / block

—

block size

—

peak bw

3350.00 GB/s

peak fp32 tflops

67.00

peak fp16 tflops

1979.00

reference PyTorch source (the prompt the agent sees)

"""
Fixed permutation along the feature axis: y[b, j] = x[b, perm[j]].
"""

import torch
import torch.nn as nn

from kernelbench.distributed_collectives import default_device

B = 10
N = 36

_g = torch.Generator()
_g.manual_seed(10)
PERM = torch.randperm(N, generator=_g)


class Model(nn.Module):
    def __init__(self):
        super().__init__()

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return x[:, PERM.to(device=x.device)]


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


def get_init_inputs():
    return []

turn 0337.3s turn · 25.0s LLM · 312.3s tools · tokens in=1,039 · out=426 · total=1,969 · FINAL

reasoning

The task is to replace the forward computation with a custom CUDA kernel for a permutation along the feature axis: y[b, j] = x[b, perm[j]]. I need to create a custom CUDA kernel to replace the forward computation of permuting along the feature axis, where y[b, j] = x[b, perm[j]]. The reference implementation is simply x[:, PERM.to(device=x.device)], and I aim to make it faster using a custom kernel with torch.utils.

call → submit_kernel

callsubmit_kernel { "kernel_code": "import torch\nimport torch.nn as nn\nfrom torch.utils.cpp_extension import load_inline\n\n# Fixed permutation\nB = 10\nN = 36\n_g = torch.Generator()\n_g.manual_seed(10)\nPERM = torch.randperm(N, generator=_g)\n\ncuda_source = r'''\n#include <torch/extension.h>\n#include <cuda_runtime.h>\n\n__global__ void permute_kernel(const float* __restrict__ x, float* __restrict__ y, const int* __restrict__ perm, int B, int N) {\n int idx = blockIdx.x * blockDim.x + threadIdx.x;\n if (idx < B * N) {\n int b = idx / N;\n int j = idx % N;\n int p = perm[j];\n y[b * N + j] = x[b * N + p];\n }\n}\n\ntorch::Tensor permute_forward(torch::Tensor x, torch::Tensor perm) {\n auto B = x.size(0);\n auto N = x.size(1);\n auto y = torch::empty({B, N}, x.options());\n const int threads = 256;\n const int blocks = (B * N + threads - 1) / threads;\n permute_kernel<<<blocks, threads>>>(\n x.data_ptr<float>(),\n y.data_ptr<float>(),\n perm.data_ptr<int>(),\n B,\n N\n );\n return y;\n}\n'''\n\ncpp_source = r'''\ntorch::Tensor permute_forward(torch::Tensor x, torch::Tensor perm);\n'''\n\n# Compile the extension\npermute_ext = load_inline(\n name='permute_ext',\n cpp_sources=[cpp_source],\n cuda_sources=[cuda_source],\n functions=['permute_forward'],\n verbose=False\n)\n\nclass ModelNew(nn.Module):\n def __init__(self):\n super().__init__()\n self.perm = PERM.cuda()\n\n def forward(self, x: torch.Tensor) -> torch.Tensor:\n return permute_ext.permute_forward(x, self.perm)" }

result ← submit_kernel (fail)

failsubmit_kernel submit_kernel FAILED: correctness check did not pass (? trials). runtime_error: expected scalar type Int but found Long

final submitted kernel (62 lines, 1,538 chars) — full source, untruncated

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

# Fixed permutation
B = 10
N = 36
_g = torch.Generator()
_g.manual_seed(10)
PERM = torch.randperm(N, generator=_g)

cuda_source = r'''
#include <torch/extension.h>
#include <cuda_runtime.h>

__global__ void permute_kernel(const float* __restrict__ x, float* __restrict__ y, const int* __restrict__ perm, int B, int N) {
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    if (idx < B * N) {
        int b = idx / N;
        int j = idx % N;
        int p = perm[j];
        y[b * N + j] = x[b * N + p];
    }
}

torch::Tensor permute_forward(torch::Tensor x, torch::Tensor perm) {
    auto B = x.size(0);
    auto N = x.size(1);
    auto y = torch::empty({B, N}, x.options());
    const int threads = 256;
    const int blocks = (B * N + threads - 1) / threads;
    permute_kernel<<<blocks, threads>>>(
        x.data_ptr<float>(),
        y.data_ptr<float>(),
        perm.data_ptr<int>(),
        B,
        N
    );
    return y;
}
'''

cpp_source = r'''
torch::Tensor permute_forward(torch::Tensor x, torch::Tensor perm);
'''

# Compile the extension
permute_ext = load_inline(
    name='permute_ext',
    cpp_sources=[cpp_source],
    cuda_sources=[cuda_source],
    functions=['permute_forward'],
    verbose=False
)

class ModelNew(nn.Module):
    def __init__(self):
        super().__init__()
        self.perm = PERM.cuda()

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return permute_ext.permute_forward(x, self.perm)