L1 P19 · popcorn

correct 19_PairwiseDistanceMatrix.py

model

grok-4.3

turns / max

4 / 12

tool calls

tokens (total)

11276

wall clock

158.1s

truncated

backend

cuda fp32

variant

popcorn

runtime

0.01 μs

ref runtime

0.08 μs

speedup

5.58x

started

2026-05-11T14:40:24Z

timingkernelbench.timing

kernel runtime

0.01 μs

ref runtime

0.08 μs

speedup

5.58x

kernel mean / median / std

0.01 / — / 0.00 μs

ref mean / median / std

0.08 / — / 9.45e-04 μs

numerical precisioneval_kernel_against_ref

max abs error

4.77e-07

mean abs error

2.82e-08

max rel error

1.65e-07

mean rel error

1.25e-08

gpu memoryextended_metrics.measure_memory

kernel peak

41.05 MB

ref peak

89.05 MB

kernel / ref

0.46x

kernel launches / fusiontorch.profiler

kernel launches

ref launches

fusion ratio (ref / kernel)

12.00x

total cuda time

26 μs

SOL (speed-of-light)nsight_hardware_counters

SOL score (0–1)

0.312

dram util %

0.1%

compute util %

31.2%

bottleneck

compute-bound

dominant pipe

FP32

arithmetic intensity

—

ridge point

20.00

roofline / occupancynsight_hardware_counters

dram bandwidth

4.83 GB/s

dram util %

0.1%

fp32 tflops

—

fp32 util %

—

fp16 tflops

—

occupancy %

79.6%

bottleneck

compute-bound

dominant pipe %

31.2%

arithmetic intensity

—

ridge point

20.00

L1 hit %

—

L2 hit %

98.1%

ld sectors / req

6.50

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

189.0%

short_scoreboard

162.0%

mio_throttle

105.2%

reference PyTorch source (the prompt the agent sees)

import torch
import torch.nn as nn


class Model(nn.Module):
    """
    Compute all-pairs Euclidean distance matrix from a 3-D point cloud.
    A core primitive in molecular dynamics simulations, docking, and
    distance-geometry-based structure prediction.
    """

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

    def forward(self, coords: torch.Tensor) -> torch.Tensor:
        """
        Args:
            coords: (B, N, 3) – 3-D coordinates of N atoms / residues
        Returns:
            dist: (B, N, N) – pairwise Euclidean distance matrix
        """
        diff = coords.unsqueeze(2) - coords.unsqueeze(1)  # (B, N, N, 3)
        dist = torch.sqrt((diff ** 2).sum(dim=-1) + 1e-8)
        return dist


num_atoms = 512
batch_size = 8


def get_inputs():
    p = popcorn_pri
    return [torch.randn(p.jitter_int(batch_size), p.jitter_int(num_atoms), 3)]


def get_init_inputs():
    return []

turn 03.4s turn · 3.4s LLM · 0.0s tools · tokens in=1,456 · out=6 · total=1,752

reasoning

The task is to replace the forward computation of a PyTorch model that computes all-pairs Euclidean distance matrix from a 3-D point cloud 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 189.4s turn · 8.0s LLM · 81.4s tools · tokens in=1,985 · out=517 · total=2,583

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 = r'''\n#include <torch/extension.h>\n#include <cuda_runtime.h>\n\n__global__ void pairwise_dist_kernel(const float* coords, float* dist, int B, int N) {\n int idx = blockIdx.x * blockDim.x + threadIdx.x;\n int total = B * N * N;\n if (idx < total) {\n int b = idx / (N * N);\n int rem = idx % (N * N);\n int i = rem / N;\n int j = rem % N;\n float dx = coords[b * N * 3 + i * 3 + 0] - coords[b * N * 3 + j * 3 + 0];\n float dy = coords[b * N * 3 + i * 3 + 1] - coords[b * N * 3 + j * 3 + 1];\n float dz = coords[b * N * 3 + i * 3 + 2] - coords[b * N * 3 + j * 3 + 2];\n float d2 = dx*dx + dy*dy + dz*dz;\n dist[b * N * N + i * N + j] = sqrtf(d2 + 1e-8f);\n }\n}\n\ntorch::Tensor pairwise_dist_cuda(torch::Tensor coords) {\n int B = coords.size(0);\n int N = coords.size(1);\n auto dist = torch::empty({B, N, N}, coords.options());\n int total = B * N * N;\n int threads = 256;\n int blocks = (total + threads - 1) / threads;\n pairwise_dist_kernel<<<blocks, threads>>>(\n coords.data_ptr<float>(),\n dist.data_ptr<float>(),\n B, N\n );\n return dist;\n}\n'''\n\ncpp_source = r'''\ntorch::Tensor pairwise_dist_cuda(torch::Tensor coords);\n'''\n\nmodule = load_inline(\n name='pairwise_dist',\n cpp_sources=[cpp_source],\n cuda_sources=[cuda_source],\n functions=['pairwise_dist_cuda'],\n verbose=False\n)\n\nclass ModelNew(nn.Module):\n def __init__(self):\n super().__init__()\n self.module = module\n\n def forward(self, coords: torch.Tensor) -> torch.Tensor:\n return self.module.pairwise_dist_cuda(coords)\n" }

result ← compile_kernel (ok)

okcompile_kernel compile_kernel PASSED: kernel compiled without errors.

turn 24.6s turn · 4.5s LLM · 0.1s tools · tokens in=2,600 · out=518 · total=3,125

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 pairwise_dist_kernel(const float* coords, float* dist, int B, int N) {\n int idx = blockIdx.x * blockDim.x + threadIdx.x;\n int total = B * N * N;\n if (idx < total) {\n int b = idx / (N * N);\n int rem = idx % (N * N);\n int i = rem / N;\n int j = rem % N;\n float dx = coords[b * N * 3 + i * 3 + 0] - coords[b * N * 3 + j * 3 + 0];\n float dy = coords[b * N * 3 + i * 3 + 1] - coords[b * N * 3 + j * 3 + 1];\n float dz = coords[b * N * 3 + i * 3 + 2] - coords[b * N * 3 + j * 3 + 2];\n float d2 = dx*dx + dy*dy + dz*dz;\n dist[b * N * N + i * N + j] = sqrtf(d2 + 1e-8f);\n }\n}\n\ntorch::Tensor pairwise_dist_cuda(torch::Tensor coords) {\n int B = coords.size(0);\n int N = coords.size(1);\n auto dist = torch::empty({B, N, N}, coords.options());\n int total = B * N * N;\n int threads = 256;\n int blocks = (total + threads - 1) / threads;\n pairwise_dist_kernel<<<blocks, threads>>>(\n coords.data_ptr<float>(),\n dist.data_ptr<float>(),\n B, N\n );\n return dist;\n}\n'''\n\ncpp_source = r'''\ntorch::Tensor pairwise_dist_cuda(torch::Tensor coords);\n'''\n\nmodule = load_inline(\n name='pairwise_dist',\n cpp_sources=[cpp_source],\n cuda_sources=[cuda_source],\n functions=['pairwise_dist_cuda'],\n verbose=False\n)\n\nclass ModelNew(nn.Module):\n def __init__(self):\n super().__init__()\n self.module = module\n\n def forward(self, coords: torch.Tensor) -> torch.Tensor:\n return self.module.pairwise_dist_cuda(coords)\n" }

result ← run_correctness (ok)

okrun_correctness run_correctness PASSED: (5 / 5) trials all matched the reference. Numerical precision: max_abs_err=4.77e-07 mean_abs_err=2.82e-08 max_rel_err=1.65e-07

turn 360.6s turn · 7.6s LLM · 53.1s tools · tokens in=3,185 · out=517 · total=3,816 · FINAL

reasoning

Correct! Now, to check if faster, I should submit it, but probably not optimized enough, as naive kernel with each thread doing one dist, and coords read many times.

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 pairwise_dist_kernel(const float* coords, float* dist, int B, int N) {\n int idx = blockIdx.x * blockDim.x + threadIdx.x;\n int total = B * N * N;\n if (idx < total) {\n int b = idx / (N * N);\n int rem = idx % (N * N);\n int i = rem / N;\n int j = rem % N;\n float dx = coords[b * N * 3 + i * 3 + 0] - coords[b * N * 3 + j * 3 + 0];\n float dy = coords[b * N * 3 + i * 3 + 1] - coords[b * N * 3 + j * 3 + 1];\n float dz = coords[b * N * 3 + i * 3 + 2] - coords[b * N * 3 + j * 3 + 2];\n float d2 = dx*dx + dy*dy + dz*dz;\n dist[b * N * N + i * N + j] = sqrtf(d2 + 1e-8f);\n }\n}\n\ntorch::Tensor pairwise_dist_cuda(torch::Tensor coords) {\n int B = coords.size(0);\n int N = coords.size(1);\n auto dist = torch::empty({B, N, N}, coords.options());\n int total = B * N * N;\n int threads = 256;\n int blocks = (total + threads - 1) / threads;\n pairwise_dist_kernel<<<blocks, threads>>>(\n coords.data_ptr<float>(),\n dist.data_ptr<float>(),\n B, N\n );\n return dist;\n}\n'''\n\ncpp_source = r'''\ntorch::Tensor pairwise_dist_cuda(torch::Tensor coords);\n'''\n\nmodule = load_inline(\n name='pairwise_dist',\n cpp_sources=[cpp_source],\n cuda_sources=[cuda_source],\n functions=['pairwise_dist_cuda'],\n verbose=False\n)\n\nclass ModelNew(nn.Module):\n def __init__(self):\n super().__init__()\n self.module = module\n\n def forward(self, coords: torch.Tensor) -> torch.Tensor:\n return self.module.pairwise_dist_cuda(coords)\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=4.77e-07 mean_abs_err=2.82e-08 max_rel_err=1.65e-07 Memory: 41.05 MB peak (ref: 89.05 MB, ratio: 0.461x) Kernel launches: 1 (ref: 12, fusion ratio: 12.0) SOL score: 0.3117

final submitted kernel (60 lines, 1,719 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 pairwise_dist_kernel(const float* coords, float* dist, int B, int N) {
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    int total = B * N * N;
    if (idx < total) {
        int b = idx / (N * N);
        int rem = idx % (N * N);
        int i = rem / N;
        int j = rem % N;
        float dx = coords[b * N * 3 + i * 3 + 0] - coords[b * N * 3 + j * 3 + 0];
        float dy = coords[b * N * 3 + i * 3 + 1] - coords[b * N * 3 + j * 3 + 1];
        float dz = coords[b * N * 3 + i * 3 + 2] - coords[b * N * 3 + j * 3 + 2];
        float d2 = dx*dx + dy*dy + dz*dz;
        dist[b * N * N + i * N + j] = sqrtf(d2 + 1e-8f);
    }
}

torch::Tensor pairwise_dist_cuda(torch::Tensor coords) {
    int B = coords.size(0);
    int N = coords.size(1);
    auto dist = torch::empty({B, N, N}, coords.options());
    int total = B * N * N;
    int threads = 256;
    int blocks = (total + threads - 1) / threads;
    pairwise_dist_kernel<<<blocks, threads>>>(
        coords.data_ptr<float>(),
        dist.data_ptr<float>(),
        B, N
    );
    return dist;
}
'''

cpp_source = r'''
torch::Tensor pairwise_dist_cuda(torch::Tensor coords);
'''

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

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

    def forward(self, coords: torch.Tensor) -> torch.Tensor:
        return self.module.pairwise_dist_cuda(coords)