correct
47_virtual_all_gather_concat_linear.py
started
2026-05-11T17:36:20Z
timingkernelbench.timing
kernel mean / median / std
0.02 / — / 1.89e-04 μs
ref mean / median / std
0.01 / — / 8.96e-04 μs
numerical precisioneval_kernel_against_ref
gpu memoryextended_metrics.measure_memory
kernel launches / fusiontorch.profiler
fusion ratio (ref / kernel)
14.00x
SOL (speed-of-light)nsight_hardware_counters
roofline / occupancynsight_hardware_counters
warp stalls (top 5)nsight
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 04.2s turn · 4.2s LLM · 0.0s tools · tokens in=2,032 · out=216 · total=2,248
reasoning
Let me understand the task:
1. We have a PyTorch model that takes a tensor of shape (R, B, H) where R=4, B=12, H=32
2. It "gathers" along the last dimension by concatenating R tensors of shape (B, H) into (B, R*H)
3. Then applies a Linear layer with bias to project from R*H=128 to K=56
The forward pass:
- Input: shape (R, B, H) = (4, 12, 32)
- After gather: shape (B, R*H) = (12, 128)
- After linear: shape (B, K) = (12, 56)
So I need to:
1. Create a custom CUDA kernel that does the virtual all-gather + linear projection
2. The kernel should be faster than PyTorch's cat + linear
Let me first get the GPU specs to understand what I'm working with.
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 1 · 1 tool call(s) · 92.0s turn · 18.4s LLM · tokens in=2,280 · out=1,797 · total=4,077 · 7,592 chars (click to expand)
turn 192.0s turn · 18.4s LLM · 73.6s tools · tokens in=2,280 · out=1,797 · total=4,077
reasoning
Now I understand the task. Let me write a custom CUDA kernel that:
1. Takes a tensor of shape (R, B, H) where R=4, B=12, H=32
2. "Gathers" the last dimension to get (B, R*H) = (B, 128)
3. Applies a linear transformation with bias to get (B, K) = (B, 56)
I'll fuse the gather and linear operation into a single kernel for efficiency.
Let me write the kernel:
```python
import torch
import torch.nn as nn
from torch.utils.cpp_extension import load_inline
# Define the CUDA kernel
cuda_source = """
#include <torch/extension.h>
#include <cuda.h>
#include <cuda_runtime.h>
template <typename scalar_t>
__global__ void gather_linear_kernel(
const scalar_t* __restrict__ x, // shape (R, B, H)
const scalar_t* __restrict__ weight, // shape (K, R*H)
const scalar_t* __restrict__ bias, // shape (K)
scalar_t* __restrict__ out, // shape (B, K)
int R, int B, int H, int K
) {
const int b = blockIdx.x; // batch index
const int k = threadIdx.x + blockDim.x * blockIdx.y; // output feature
if (b >= B || k >= K) return;
scalar_t sum = bias[k];
// Loop over all input features (R*H)
for (int rh = 0; rh < R * H; rh++) {
int r = rh / H;
int h = rh % H;
// x[r, b, h]
scalar_t x_val = x[r * B * H + b * H + h];
// weight[k, rh]
sum += x_val * weight[k * R * H + rh];
}
out[b * K + k] = sum;
}
torch::Tensor gather_linear_cuda(
torch::Tensor x, // shape (R, B, H)
torch::Tensor weight, // shape (K, R*H)
torch::Tensor bias // shape (K)
) {
const int R = x.size(0);
const int B = x.size(1);
const int H = x.size(2);
const int K = weight.size(0);
auto out = torch::empty({B, K}, x.options());
const int threads = 256;
const int blocks_y = (K + threads - 1) / threads;
const dim3 grid(B, blocks_y);
AT_DISPATCH_FLOATING_TYPES(x.scalar_type(), "gather_linear_kernel", ([&] {
gather_linear_kernel<scalar_t><<<grid, threads>>>(
x.data_ptr<scalar_t>(),
weight.data_ptr<scalar_t>(),
bias.data_ptr<scalar_t>(),
out.data_ptr<scalar_t>(),
R, B, H, K
);
}));
return out;
}
"""
cpp_source = """
torch::Tensor gather_linear_cuda(torch::Tensor x, torch::Tensor weight, torch::Tensor bias);
"""
# Load the extension
gather_linear_ext = load_inline(
name='gather_linear',
cpp_sources=cpp_source,
cuda_sources=cuda_source,
functions=['gather_linear_cuda'],
verbose=True
)
class ModelNew(nn.Module):
def __init__(self):
super().__init__()
self.proj = nn.Linear(R * H, K, bias=True)
def forward(self, x: torch.Tensor) -> torch.Tensor:
# Use the custom CUDA kernel for fused gather + linear
return gather_linear_ext.gather_linear_cuda(x, self.proj.weight, self.proj.bias)
```
Wait, there's an issue. The reference uses `nn.Linear` which stores weight as (out_features, in_features), i.e., (K, R*H). But the linear operation is `x @ W.T + b`, not `x @ W + b`. Let me verify this.
In PyTorch, nn.Linear:
- weight shape: (out_features, in_features)
- forward: output = input @ weight.T + bias
So for each output element out[b, k]:
out[b, k] = sum_{rh} (gathered[b, rh] * weight[k, rh]) + bias[k]
My kernel is correct for this. But I need to make sure I'm reading the weight correctly. The weight is stored as (K, R*H), so weight[k * R * H + rh] should give weight[k, rh], which is correct.
Let me compile and test this.
assistantNow I understand the task. I need to fuse the "virtual all-gather" (stacking tensors along the last dim) with a linear layer. The input shape is (R, B, H) = (4, 12, 32), and I need to produce (B, K) = (12, 56).
Let me write a fused CUDA kernel:
call → compile_kernel
callcompile_kernel
{
"kernel_code": "import torch\nimport torch.nn as nn\nfrom torch.utils.cpp_extension import load_inline\n\nR = 4\nB = 12\nH = 32\nK = 56\n\ncuda_source = \"\"\"\n#include <torch/extension.h>\n#include <cuda.h>\n#include <cuda_runtime.h>\n\ntemplate <typename scalar_t>\n__global__ void gather_linear_kernel(\n const scalar_t* __restrict__ x, // shape (R, B, H)\n const scalar_t* __restrict__ weight, // shape (K, R*H)\n const scalar_t* __restrict__ bias, // shape (K)\n scalar_t* __restrict__ out, // shape (B, K)\n int R, int B, int H, int K\n) {\n const int b = blockIdx.x; // batch index\n const int k = threadIdx.x + blockDim.x * blockIdx.y; // output feature\n \n if (b >= B || k >= K) return;\n \n scalar_t sum = bias[k];\n \n // Loop over all input features (R*H)\n for (int rh = 0; rh < R * H; rh++) {\n int r = rh / H;\n int h = rh % H;\n // x[r, b, h]\n scalar_t x_val = x[r * B * H + b * H + h];\n // weight[k, rh]\n sum += x_val * weight[k * R * H + rh];\n }\n \n out[b * K + k] = sum;\n}\n\ntorch::Tensor gather_linear_cuda(\n torch::Tensor x, // shape (R, B, H)\n torch::Tensor weight, // shape (K, R*H)\n torch::Tensor bias // shape (K)\n) {\n const int R = x.size(0);\n const int B = x.size(1);\n const int H = x.size(2);\n const int K = weight.size(0);\n \n auto out = torch::empty({B, K}, x.options());\n \n const int threads = 256;\n const int blocks_y = (K + threads - 1) / threads;\n const dim3 grid(B, blocks_y);\n \n AT_DISPATCH_FLOATING_TYPES(x.scalar_type(), \"gather_linear_kernel\", ([&] {\n gather_linear_kernel<scalar_t><<<grid, threads>>>(\n x.data_ptr<scalar_t>(),\n weight.data_ptr<scalar_t>(),\n bias.data_ptr<scalar_t>(),\n out.data_ptr<scalar_t>(),\n R, B, H, K\n );\n }));\n \n return out;\n}\n\"\"\"\n\ncpp_source = \"\"\"\ntorch::Tensor gather_linear_cuda(torch::Tensor x, torch::Tensor weight, torch::Tensor bias);\n\"\"\"\n\n# Load the extension\ngather_linear_ext = load_inline(\n name='gather_linear_v1',\n cpp_sources=cpp_source,\n cuda_sources=cuda_source,\n functions=['gather_linear_cuda'],\n verbose=True\n)\n\n\nclass ModelNew(nn.Module):\n def __init__(self):\n super().__init__()\n self.proj = nn.Linear(R * H, K, bias=True)\n \n def forward(self, x: torch.Tensor) -> torch.Tensor:\n # Use the custom CUDA kernel for fused gather + linear\n return gather_linear_ext.gather_linear_cuda(x, self.proj.weight, self.proj.bias)"
}
result ← compile_kernel (ok)
okcompile_kernel
compile_kernel PASSED: kernel compiled without errors.
turn 2151.1s turn · 6.7s LLM · 144.4s tools · tokens in=3,092 · out=734 · total=3,826
reasoning
The kernel compiles. Now let me 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\nR = 4\nB = 12\nH = 32\nK = 56\n\ncuda_source = \"\"\"\n#include <torch/extension.h>\n#include <cuda.h>\n#include <cuda_runtime.h>\n\ntemplate <typename scalar_t>\n__global__ void gather_linear_kernel(\n const scalar_t* __restrict__ x, // shape (R, B, H)\n const scalar_t* __restrict__ weight, // shape (K, R*H)\n const scalar_t* __restrict__ bias, // shape (K)\n scalar_t* __restrict__ out, // shape (B, K)\n int R, int B, int H, int K\n) {\n const int b = blockIdx.x; // batch index\n const int k = threadIdx.x + blockDim.x * blockIdx.y; // output feature\n \n if (b >= B || k >= K) return;\n \n scalar_t sum = bias[k];\n \n // Loop over all input features (R*H)\n for (int rh = 0; rh < R * H; rh++) {\n int r = rh / H;\n int h = rh % H;\n // x[r, b, h]\n scalar_t x_val = x[r * B * H + b * H + h];\n // weight[k, rh]\n sum += x_val * weight[k * R * H + rh];\n }\n \n out[b * K + k] = sum;\n}\n\ntorch::Tensor gather_linear_cuda(\n torch::Tensor x, // shape (R, B, H)\n torch::Tensor weight, // shape (K, R*H)\n torch::Tensor bias // shape (K)\n) {\n const int R = x.size(0);\n const int B = x.size(1);\n const int H = x.size(2);\n const int K = weight.size(0);\n \n auto out = torch::empty({B, K}, x.options());\n \n const int threads = 256;\n const int blocks_y = (K + threads - 1) / threads;\n const dim3 grid(B, blocks_y);\n \n AT_DISPATCH_FLOATING_TYPES(x.scalar_type(), \"gather_linear_kernel\", ([&] {\n gather_linear_kernel<scalar_t><<<grid, threads>>>(\n x.data_ptr<scalar_t>(),\n weight.data_ptr<scalar_t>(),\n bias.data_ptr<scalar_t>(),\n out.data_ptr<scalar_t>(),\n R, B, H, K\n );\n }));\n \n return out;\n}\n\"\"\"\n\ncpp_source = \"\"\"\ntorch::Tensor gather_linear_cuda(torch::Tensor x, torch::Tensor weight, torch::Tensor bias);\n\"\"\"\n\n# Load the extension\ngather_linear_ext = load_inline(\n name='gather_linear_v1',\n cpp_sources=cpp_source,\n cuda_sources=cuda_source,\n functions=['gather_linear_cuda'],\n verbose=True\n)\n\n\nclass ModelNew(nn.Module):\n def __init__(self):\n super().__init__()\n self.proj = nn.Linear(R * H, K, bias=True)\n \n def forward(self, x: torch.Tensor) -> torch.Tensor:\n # Use the custom CUDA kernel for fused gather + linear\n return gather_linear_ext.gather_linear_cuda(x, self.proj.weight, self.proj.bias)"
}
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 3237.2s turn · 7.2s LLM · 230.0s tools · tokens in=3,872 · out=758 · total=4,630
reasoning
The kernel is correct. Now let me profile it to understand performance characteristics.
assistantThe kernel is correct. Let me profile it to see performance characteristics:
call → profile_kernel
callprofile_kernel
{
"kernel_code": "import torch\nimport torch.nn as nn\nfrom torch.utils.cpp_extension import load_inline\n\nR = 4\nB = 12\nH = 32\nK = 56\n\ncuda_source = \"\"\"\n#include <torch/extension.h>\n#include <cuda.h>\n#include <cuda_runtime.h>\n\ntemplate <typename scalar_t>\n__global__ void gather_linear_kernel(\n const scalar_t* __restrict__ x, // shape (R, B, H)\n const scalar_t* __restrict__ weight, // shape (K, R*H)\n const scalar_t* __restrict__ bias, // shape (K)\n scalar_t* __restrict__ out, // shape (B, K)\n int R, int B, int H, int K\n) {\n const int b = blockIdx.x; // batch index\n const int k = threadIdx.x + blockDim.x * blockIdx.y; // output feature\n \n if (b >= B || k >= K) return;\n \n scalar_t sum = bias[k];\n \n // Loop over all input features (R*H)\n for (int rh = 0; rh < R * H; rh++) {\n int r = rh / H;\n int h = rh % H;\n // x[r, b, h]\n scalar_t x_val = x[r * B * H + b * H + h];\n // weight[k, R*H - 1 - rh]\n sum += x_val * weight[k * R * H + rh];\n }\n \n out[b * K + k] = sum;\n}\n\ntorch::Tensor gather_linear_cuda(\n torch::Tensor x, // shape (R, B, H)\n torch::Tensor weight, // shape (K, R*H)\n torch::Tensor bias // shape (K)\n) {\n const int R = x.size(0);\n const int B = x.size(1);\n const int H = x.size(2);\n const int K = weight.size(0);\n \n auto out = torch::empty({B, K}, x.options());\n \n const int threads = 256;\n const int blocks_y = (K + threads - 1) / threads;\n const dim3 grid(B, blocks_y);\n \n AT_DISPATCH_FLOATING_TYPES(x.scalar_type(), \"gather_linear_kernel\", ([&] {\n gather_linear_kernel<scalar_t><<<grid, threads>>>(\n x.data_ptr<scalar_t>(),\n weight.data_ptr<scalar_t>(),\n bias.data_ptr<scalar_t>(),\n out.data_ptr<scalar_t>(),\n R, B, H, K\n );\n }));\n \n return out;\n}\n\"\"\"\n\ncpp_source = \"\"\"\ntorch::Tensor gather_linear_cuda(torch::Tensor x, torch::Tensor weight, torch::Tensor bias);\n\"\"\"\n\n# Load the extension\ngather_linear_ext = load_inline(\n name='gather_linear_v1',\n cpp_sources=cpp_source,\n cuda_sources=cuda_source,\n functions=['gather_linear_cuda'],\n verbose=True\n)\n\n\nclass ModelNew(nn.Module):\n def __init__(self):\n super().__init__()\n self.proj = nn.Linear(R * H, K, bias=True)\n \n def forward(self, x: torch.Tensor) -> torch.Tensor:\n # Use the custom CUDA kernel for fused gather + linear\n return gather_linear_ext.gather_linear_cuda(x, self.proj.weight, self.proj.bias)"
}
result ← profile_kernel (ok)
okprofile_kernel
profile_kernel PASSED: profiling complete.
=== Kernel Profile Summary ===
--- Kernel Breakdown (by GPU time) ---
#1: aten::empty 32.0μs (100%) calls=1
--- Timing ---
GPU time: 14.27 μs
--- Memory Bandwidth ---
DRAM BW: 2.9/3350 GB/s (0.1%)
read=2.9 write=0.0 GB/s
Hit rates: L2=96.9%
--- Compute Throughput ---
Pipe util: FMA=3.8%, Tensor=0.0%
--- Roofline Analysis ---
Ridge point (FP32): 20.00 FLOPs/byte
Classification: UNKNOWN
--- Occupancy ---
Achieved: 3.5%
Eligible warps/cycle: 0.14
WARNING: <1 eligible warp/cycle = latency-bound, needs more parallelism
Launch: regs/thread=28, smem/blk=1024B, block=256, grid=12
--- Memory Access Quality ---
Global ld coalescing: 14.46 sectors/req (14.5x wasted BW)
--- Warp Stall Reasons (% of issue cycles) ---
long_scoreboard 312.5% ####################
short_scoreboard 70.1% ##############
--- Optimization Hints ---
-> 312% stalls: global memory latency (waiting for DRAM/L2). + sectors/req=14.5 -> uncoalesced global loads, restructure access pattern or use __ldg Try shared-memory tiling, prefetching, or reducing working set.
-> Global load coalescing=14.5x (ideal=1.0) -> 14x wasted DRAM BW. Make per-warp accesses contiguous.
-> Low occupancy (4%). Consider __launch_bounds__ or reducing register pressure.
-> Low eligible warps + low BW% + low compute% = LATENCY-BOUND. Need more parallelism (larger grid, more ILP, prefetching).
turn 4 · 1 tool call(s) · 223.7s turn · 17.3s LLM · tokens in=5,064 · out=1,665 · total=6,729 · 7,154 chars (click to expand)
turn 4223.7s turn · 17.3s LLM · 206.4s tools · tokens in=5,064 · out=1,665 · total=6,729
reasoning
The profiling shows that my kernel is quite inefficient:
- Very low occupancy (3.5%)
- Poor memory coalescing (14.5x wasted BW)
- Latency bound with warp stalls
The main issue is that when threads in a warp access the x array, they're accessing elements from different batches (b) but the same feature (rh). But the x tensor is stored as (R, B, H), so for a given b, rh accesses are scattered across different r values.
Let me think about how to improve memory access patterns:
1. For weight[k, rh], each thread k accesses different weights, so these are already coalesced within a threadblock.
2. For x[r, b, h], threads with different k values access the same x value (since x is indexed by r, b, h only). This means multiple threads read the same x value!
Wait, actually I can optimize this by having threads in a warp collaborate to load x values once and share them.
Let me try a different approach: use shared memory to stage the x values, then have all threads compute their dot products.
Actually, let me think about the data layout more carefully:
- x: (R, B, H) = (4, 12, 32) = 1536 elements
- weight: (K, R*H) = (56, 128) = 7168 elements
- output: (B, K) = (12, 56) = 672 elements
Each output element needs R*H = 128 multiply-accumulates.
For each batch b, we need to compute 56 outputs, each requiring 128 input values.
The total FLOPs is B * K * R * H * 2 = 12 * 56 * 128 * 2 = 172,032 FLOPs.
Total data read: R * B * H + K * R * H + K = 1536 + 7168 + 56 = 8760 bytes (for fp32)
Arithmetic intensity: 172032 / 8760 ≈ 19.6 FLOPs/byte
This is on the edge of compute-bound for FP32 (ridge point is around 20 FLOPs/byte).
Let me try a tiled approach using shared memory to improve memory access patterns. The idea is:
1. Load x for a batch into shared memory
2. Load weight tiles into shared memory
3. Compute matrix multiplication
Actually, the problem is quite small (12 batches, 56 outputs, 128 inputs). Let me try a different approach where:
- Each block handles one batch
- Threads in the block cooperatively compute all K outputs
Let me try a tiled matrix multiplication approach:
```
For each batch b (blockIdx.x):
Load x[b] (128 values) into shared memory
For each tile of weight (K columns):
Load tile of weight into shared memory
Compute partial products
```
Wait, the weights are (K, R*H), so we're computing:
output[b, k] = sum_{rh} x[b, rh] * weight[k, rh]
This is essentially a matrix-vector product for each batch, or equivalently a matrix-matrix product where the batch dimension gives us multiple vectors.
Actually, I can think of this as:
- X: (B, R*H) after reshaping (the gather operation)
- W: (K, R*H)
- Output: (B, K) = X @ W^T + bias
Let me try a tiled approach. Given the small sizes, I might use a 1-block-per-batch approach with cooperative loading and computation.
Let me try a cleaner implementation with better memory access:
assistantThe profiling shows poor memory coalescing and low occupancy. Let me redesign the kernel with better memory access patterns - using tiled shared memory:
call → compile_kernel
callcompile_kernel
{
"kernel_code": "import torch\nimport torch.nn as nn\nfrom torch.utils.cpp_extension import load_inline\n\nR = 4\nB = 12\nH = 32\nK = 56\n\ncuda_source = \"\"\"\n#include <torch/extension.h>\n#include <cuda.h>\n#include <cuda_runtime.h>\n\n#define TILE_SIZE 32\n\ntemplate <typename scalar_t>\n__global__ void gather_linear_kernel(\n const scalar_t* __restrict__ x, // shape (R, B, H)\n const scalar_t* __restrict__ weight, // shape (K, R*H)\n const scalar_t* __restrict__ bias, // shape (K)\n scalar_t* __restrict__ out, // shape (B, K)\n int R, int B, int H, int K\n) {\n // Each block handles one batch, all K outputs\n const int b = blockIdx.x;\n const int tid = threadIdx.x;\n const int RH = R * H;\n \n // Shared memory for the gathered x vector (R*H elements)\n __shared__ scalar_t xs[128]; // R*H = 128\n \n // Load x values into shared memory cooperatively\n // x shape is (R, B, H), need to gather into (R*H,) for this batch b\n for (int i = tid; i < RH; i += blockDim.x) {\n int r = i / H;\n int h = i % H;\n xs[i] = x[r * B * H + b * H + h];\n }\n __syncthreads();\n \n // Each thread computes one or more output elements\n // output k = threadIdx.x + threadIdx.y * blockDim.x, etc.\n for (int k = tid; k < K; k += blockDim.x) {\n scalar_t sum = bias[k];\n \n // Compute dot product with weight row k\n for (int rh = 0; rh < RH; rh++) {\n sum += xs[rh] * weight[k * RH + rh];\n }\n out[b * K + k] = sum;\n }\n}\n\ntorch::Tensor gather_linear_cuda(\n torch::Tensor x, // shape (R, B, H)\n torch::Tensor weight, // shape (K, R*H)\n torch::Tensor bias // shape (K)\n) {\n const int R = x.size(0);\n const int B = x.size(1);\n const int H = x.size(2);\n const int K = weight.size(0);\n \n auto out = torch::empty({B, K}, x.options());\n \n // One block per batch, enough threads to cover K outputs\n const int threads = 128; // K=56, so 128 threads covers it\n const dim3 grid(B);\n \n AT_DISPATCH_FLOATING_TYPES(x.scalar_type(), \"gather_linear_kernel\", ([&] {\n gather_linear_kernel<scalar_t><<<grid, threads>>>(\n x.data_ptr<scalar_t>(),\n weight.data_ptr<scalar_t>(),\n bias.data_ptr<scalar_t>(),\n out.data_ptr<scalar_t>(),\n R, B, H, K\n );\n }));\n \n return out;\n}\n\"\"\"\n\ncpp_source = \"\"\"\ntorch::Tensor gather_linear_cuda(torch::Tensor x, torch::Tensor weight, torch::Tensor bias);\n\"\"\"\n\n# Load the extension\ngather_linear_ext = load_inline(\n name='gather_linear_v2',\n cpp_sources=cpp_source,\n cuda_sources=cuda_source,\n functions=['gather_linear_cuda'],\n verbose=True\n)\n\n\nclass ModelNew(nn.Module):\n def __init__(self):\n super().__init__()\n self.proj = nn.Linear(R * H, K, bias=True)\n \n def forward(self, x: torch.Tensor) -> torch.Tensor:\n # Use the custom CUDA kernel for fused gather + linear\n return gather_linear_ext.gather_linear_cuda(x, self.proj.weight, self.proj.bias)"
}
result ← compile_kernel (ok)
okcompile_kernel
compile_kernel PASSED: kernel compiled without errors.
turn 5153.3s turn · 7.3s LLM · 146.1s tools · tokens in=5,955 · out=858 · total=6,813
reasoning
Good, let me test correctness.
call → run_correctness
callrun_correctness
{
"kernel_code": "import torch\nimport torch.nn as nn\nfrom torch.utils.cpp_extension import load_inline\n\nR = 4\nB = 12\nH = 32\nK = 56\n\ncuda_source = \"\"\"\n#include <torch/extension.h>\n#include <cuda.h>\n#include <cuda_runtime.h>\n\n#define TILE_SIZE 32\n\ntemplate <typename scalar_t>\n__global__ void gather_linear_kernel(\n const scalar_t* __restrict__ x, // shape (R, B, H)\n const scalar_t* __restrict__ weight, // shape (K, R*H)\n const scalar_t* __restrict__ bias, // shape (K)\n scalar_t* __restrict__ out, // shape (B, K)\n int R, int B, int H, int K\n) {\n // Each block handles one batch, all K outputs\n const int b = blockIdx.x;\n const int tid = threadIdx.x;\n const int RH = R * H;\n \n // Shared memory for the gathered x vector (R*H elements)\n __shared__ scalar_t xs[128]; // R*H = 128\n \n // Load x values into shared memory cooperatively\n // x shape is (R, B, H), need to gather into (R*H,) for this batch b\n for (int i = tid; i < RH; i += blockDim.x) {\n int r = i / H;\n int h = i % H;\n xs[i] = x[r * B * H + b * H + h];\n }\n __syncthreads();\n \n // Each thread computes one or more output elements\n // output k = threadIdx.x + threadIdx.y * blockDim.x, etc.\n for (int k = tid; k < K; k += blockDim.x) {\n scalar_t sum = bias[k];\n \n // Compute dot product with weight row k\n for (int rh = 0; rh < RH; rh++) {\n sum += xs[rh] * weight[k * RH + rh];\n }\n out[b * K + k] = sum;\n }\n}\n\ntorch::Tensor gather_linear_cuda(\n torch::Tensor x, // shape (R, B, H)\n torch::Tensor weight, // shape (K, R*H)\n torch::Tensor bias // shape (K)\n) {\n const int R = x.size(0);\n const int B = x.size(1);\n const int H = x.size(2);\n const int K = weight.size(0);\n \n auto out = torch::empty({B, K}, x.options());\n \n // One block per batch, enough threads to cover K outputs\n const int threads = 128; // K=56, so 128 threads covers it\n const dim3 grid(B);\n \n AT_DISPATCH_FLOATING_TYPES(x.scalar_type(), \"gather_linear_kernel\", ([&] {\n gather_linear_kernel<scalar_t><<<grid, threads>>>(\n x.data_ptr<scalar_t>(),\n weight.data_ptr<scalar_t>(),\n bias.data_ptr<scalar_t>(),\n out.data_ptr<scalar_t>(),\n R, B, H, K\n );\n }));\n \n return out;\n}\n\"\"\"\n\ncpp_source = \"\"\"\ntorch::Tensor gather_linear_cuda(torch::Tensor x, torch::Tensor weight, torch::Tensor bias);\n\"\"\"\n\n# Load the extension\ngather_linear_ext = load_inline(\n name='gather_linear_v2',\n cpp_sources=cpp_source,\n cuda_sources=cuda_source,\n functions=['gather_linear_cuda'],\n verbose=True\n)\n\n\nclass ModelNew(nn.Module):\n def __init__(self):\n super().__init__()\n self.proj = nn.Linear(R * H, K, bias=True)\n \n def forward(self, x: torch.Tensor) -> torch.Tensor:\n # Use the custom CUDA kernel for fused gather + linear\n return gather_linear_ext.gather_linear_cuda(x, self.proj.weight, self.proj.bias)"
}
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 647.2s turn · 8.0s LLM · 39.1s tools · tokens in=6,863 · out=861 · total=7,724
reasoning
Good, still correct. Let me profile this version:
call → profile_kernel
callprofile_kernel
{
"kernel_code": "import torch\nimport torch.nn as nn\nfrom torch.utils.cpp_extension import load_inline\n\nR = 4\nB = 12\nH = 32\nK = 56\n\ncuda_source = \"\"\"\n#include <torch/extension.h>\n#include <cuda.h>\n#include <cuda_runtime.h>\n\n#define TILE_SIZE 32\n\ntemplate <typename scalar_t>\n__global__ void gather_linear_kernel(\n const scalar_t* __restrict__ x, // shape (R, B, H)\n const scalar_t* __restrict__ weight, // shape (K, R*H)\n const scalar_t* __restrict__ bias, // shape (K)\n scalar_t* __restrict__ out, // shape (B, K)\n int R, int B, int H, int K\n) {\n // Each block handles one batch, all K outputs\n const int b = blockIdx.x;\n const int tid = threadIdx.x;\n const int RH = R * H;\n \n // Shared memory for the gathered x vector (R*H elements)\n __shared__ scalar_t xs[128]; // R*H = 128\n \n // Load x values into shared memory cooperatively\n // x shape is (R, B, H), need to gather into (R*H,) for this batch b\n for (int i = tid; i < RH; i += blockDim.x) {\n int r = i / H;\n int h = i % H;\n xs[i] = x[r * B * H + b * H + h];\n }\n __syncthreads();\n \n // Each thread computes one or more output elements\n // output k = threadIdx.x + threadIdx.y * blockDim.x, etc.\n for (int k = tid; k < K; k += blockDim.x) {\n scalar_t sum = bias[k];\n \n // Compute dot product with weight row k\n for (int rh = 0; rh < RH; rh++) {\n sum += xs[rh] * weight[k * RH + rh];\n }\n out[b * K + k] = sum;\n }\n}\n\ntorch::Tensor gather_linear_cuda(\n torch::Tensor x, // shape (R, B, H)\n torch::Tensor weight, // shape (K, R*H)\n torch::Tensor bias // shape (K)\n) {\n const int R = x.size(0);\n const int B = x.size(1);\n const int H = x.size(2);\n const int K = weight.size(0);\n \n auto out = torch::empty({B, K}, x.options());\n \n // One block per batch, enough threads to cover K outputs\n const int threads = 128; // K=56, so 128 threads covers it\n const dim3 grid(B);\n \n AT_DISPATCH_FLOATING_TYPES(x.scalar_type(), \"gather_linear_kernel\", ([&] {\n gather_linear_kernel<scalar_t><<<grid, threads>>>(\n x.data_ptr<scalar_t>(),\n weight.data_ptr<scalar_t>(),\n bias.data_ptr<scalar_t>(),\n out.data_ptr<scalar_t>(),\n R, B, H, K\n );\n }));\n \n return out;\n}\n\"\"\"\n\ncpp_source = \"\"\"\ntorch::Tensor gather_linear_cuda(torch::Tensor x, torch::Tensor weight, torch::Tensor bias);\n\"\"\"\n\n# Load the extension\ngather_linear_ext = load_inline(\n name='gather_linear_v2',\n cpp_sources=cpp_source,\n cuda_sources=cuda_source,\n functions=['gather_linear_cuda'],\n verbose=True\n)\n\n\nclass ModelNew(nn.Module):\n def __init__(self):\n super().__init__()\n self.proj = nn.Linear(R * H, K, bias=True)\n \n def forward(self, x: torch.Tensor) -> torch.Tensor:\n # Use the custom CUDA kernel for fused gather + linear\n return gather_linear_ext.gather_linear_cuda(x, self.proj.weight, self.proj.bias)"
}
result ← profile_kernel (ok)
okprofile_kernel
profile_kernel PASSED: profiling complete.
=== Kernel Profile Summary ===
--- Kernel Breakdown (by GPU time) ---
#1: aten::empty 30.0μs (100%) calls=1
--- Timing ---
GPU time: 9.82 μs [↓4.4μs ok]
--- Memory Bandwidth ---
DRAM BW: 4.1/3350 GB/s (0.1%) [↑0.0% ok]
read=4.1 write=0.0 GB/s
Hit rates: L2=96.2%
--- Compute Throughput ---
Pipe util: FMA=0.6%, Tensor=0.0%
--- Roofline Analysis ---
Ridge point (FP32): 20.00 FLOPs/byte
Classification: UNKNOWN
--- Occupancy ---
Achieved: 3.5%
Eligible warps/cycle: 0.03
WARNING: <1 eligible warp/cycle = latency-bound, needs more parallelism
Launch: regs/thread=33, smem/blk=1536B, block=128, grid=12
--- Memory Access Quality ---
Global ld coalescing: 27.45 sectors/req (27.4x wasted BW)
--- Warp Stall Reasons (% of issue cycles) ---
long_scoreboard 2452.2% ####################
short_scoreboard 48.4% #########
--- Optimization Hints ---
-> 2452% stalls: global memory latency (waiting for DRAM/L2). + sectors/req=27.4 -> uncoalesced global loads, restructure access pattern or use __ldg Try shared-memory tiling, prefetching, or reducing working set.
-> Global load coalescing=27.4x (ideal=1.0) -> 27x wasted DRAM BW. Make per-warp accesses contiguous.
-> Low occupancy (4%). Consider __launch_bounds__ or reducing register pressure.
-> Low eligible warps + low BW% + low compute% = LATENCY-BOUND. Need more parallelism (larger grid, more ILP, prefetching).
--- Delta vs Previous Iteration ---
GPU time: 14.27→9.82μs (↓4.45) ok
DRAM BW%: 0.09→0.12% (↑0.04) ok
Eligible warps/cyc: 0.14→0.03 (↓0.11) REGRESSED
Ld sectors/req: 14.46→27.45 (↑12.99) REGRESSED
L2 hit%: 96.92→96.18% (↓0.74) REGRESSED
turn 7307.1s turn · 12.1s LLM · 295.0s tools · tokens in=8,303 · out=1,261 · total=9,564
reasoning
The timing improved (14.27→9.82μs) but memory coalescing got worse. The main issue is still low occupancy and memory access patterns. The problem is that I'm loading x in a strided manner.
Let me think about this more carefully. The original code does:
1. `torch.cat(parts, dim=-1)` - creates a contiguous (B, R*H) tensor from (R, B, H) by copying
2. `self.proj(gathered)` - does a matmul
The fused kernel is avoiding the intermediate tensor allocation, which is good. But the memory access pattern for reading x is not ideal.
Let me try a different approach - instead of loading the entire x into shared memory first, let me try to compute the output directly in a tiled manner. Each block handles multiple output elements, and uses tiled matmul.
Actually, let me try a simpler approach - have each thread compute one output element, but process multiple batches in a warp for better memory coalescing:
assistantThe kernel is faster (14.27→9.82μs) but still has poor coalescing. Let me try a tiled approach where I load weight into shared memory for better coalescing:
call → compile_kernel
callcompile_kernel
{
"kernel_code": "import torch\nimport torch.nn as nn\nfrom torch.utils.cpp_extension import load_inline\n\nR = 4\nB = 12\nH = 32\nK = 56\n\ncuda_source = \"\"\"\n#include <torch/extension.h>\n#include <cuda.h>\n#include <cuda_runtime.h>\n\n// Tiled GEMV: each block handles one batch sample\n// Threads cooperatively load weight tiles and x into shared memory\ntemplate <typename scalar_t>\n__global__ void __launch_bounds__(128) gather_linear_kernel(\n const scalar_t* __restrict__ x, // shape (R, B, H)\n const scalar_t* __restrict__ weight, // shape (K, R*H) row-major\n const scalar_t* __restrict__ bias, // shape (K)\n scalar_t* __restrict__ out, // shape (B, K)\n int R, int B, int H, int K\n) {\n const int b = blockIdx.x; // batch index\n const int tid = threadIdx.x;\n const int RH = R * H;\n \n // Shared memory for x vector and weight tile\n __shared__ scalar_t xs[128]; // R*H = 128\n __shared__ scalar_t w_tile[128 * 128]; // Max 128x128 = 16KB (fits in smem)\n \n // Load x into shared memory - need to gather from strided locations\n // Each thread loads multiple elements for better coalescing\n // Reorder so thread i loads x[0, b, i], x[1, b, i], etc for H elements\n // Then next thread loads next set of H elements\n for (int i = tid; i < RH; i += blockDim.x) {\n int r = i / H;\n int h = i % H;\n // xs[i] should be x[r, b, h]\n xs[i] = x[r * B * H + b * H + h];\n }\n __syncthreads();\n \n // Each thread computes one output element\n const int k = tid;\n if (k < K) {\n scalar_t sum = bias[k];\n \n // Accumulate dot product over tiles\n // For this small problem, just do it directly\n // Weight access pattern: weight[k, rh] is weight[k * RH + rh]\n // This is coalesced for the same k across threads (reading same row)\n // But we have different k values per thread, so strided access\n \n // Better approach: use __ldg for read-only cache\n const scalar_t* w_row = weight + k * RH;\n for (int rh = 0; rh < RH; rh++) {\n sum += xs[rh] * __ldg(w_row + rh);\n }\n out[b * K + k] = sum;\n }\n}\n\ntorch::Tensor gather_linear_cuda(\n torch::Tensor x, // shape (R, B, H)\n torch::Tensor weight, // shape (K, R*H)\n torch::Tensor bias // shape (K)\n) {\n const int R = x.size(0);\n const int B = x.size(1);\n const int H = x.size(2);\n const int K = weight.size(0);\n \n auto out = torch::empty({B, K}, x.options());\n \n // One block per batch\n const int threads = 128; // Cover K=56 outputs\n const dim3 grid(B);\n \n AT_DISPATCH_FLOATING_TYPES(x.scalar_type(), \"gather_linear_kernel\", ([&] {\n gather_linear_kernel<scalar_t><<<grid, threads>>>(\n x.data_ptr<scalar_t>(),\n weight.data_ptr<scalar_t>(),\n bias.data_ptr<scalar_t>(),\n out.data_ptr<scalar_t>(),\n R, B, H, K\n );\n }));\n \n return out;\n}\n\"\"\"\n\ncpp_source = \"\"\"\ntorch::Tensor gather_linear_cuda(torch::Tensor x, torch::Tensor weight, torch::Tensor bias);\n\"\"\"\n\n# Load the extension\ngather_linear_ext = load_inline(\n name='gather_linear_v3',\n cpp_sources=cpp_source,\n cuda_sources=cuda_source,\n functions=['gather_linear_cuda'],\n verbose=True\n)\n\n\nclass ModelNew(nn.Module):\n def __init__(self):\n super().__init__()\n self.proj = nn.Linear(R * H, K, bias=True)\n \n def forward(self, x: torch.Tensor) -> torch.Tensor:\n # Use the custom CUDA kernel for fused gather + linear\n return gather_linear_ext.gather_linear_cuda(x, self.proj.weight, self.proj.bias)"
}
result ← compile_kernel (ok)
okcompile_kernel
compile_kernel PASSED: kernel compiled without errors.
turn 8313.8s turn · 8.8s LLM · 305.0s tools · tokens in=9,366 · out=1,018 · total=10,384
reasoning
Let me run correctness and profile again.
call → run_correctness
callrun_correctness
{
"kernel_code": "import torch\nimport torch.nn as nn\nfrom torch.utils.cpp_extension import load_inline\n\nR = 4\nB = 12\nH = 32\nK = 56\n\ncuda_source = \"\"\"\n#include <torch/extension.h>\n#include <cuda.h>\n#include <cuda_runtime.h>\n\n// Tiled GEMV: each block handles one batch sample\n// Threads cooperatively load weight tiles and x into shared memory\ntemplate <typename scalar_t>\n__global__ void __launch_bounds__(128) gather_linear_kernel(\n const scalar_t* __restrict__ x, // shape (R, B, H)\n const scalar_t* __restrict__ weight, // shape (K, R*H) row-major\n const scalar_t* __restrict__ bias, // shape (K)\n scalar_t* __restrict__ out, // shape (B, K)\n int R, int B, int H, int K\n) {\n const int b = blockIdx.x; // batch index\n const int tid = threadIdx.x;\n const int RH = R * H;\n \n // Shared memory for x vector and weight tile\n __shared__ scalar_t xs[128]; // R*H = 128\n __shared__ scalar_t w_tile[128 * 128]; // Max 128x128 = 16KB (fits in smem)\n \n // Load x into shared memory - need to gather from strided locations\n // Each thread loads multiple elements for better coalescing\n // Reorder so thread i loads x[0, b, i], x[1, b, i], etc for H elements\n // Then next thread loads next set of H elements\n for (int i = tid; i < RH; i += blockDim.x) {\n int r = i / H;\n int h = i % H;\n // xs[i] should be x[r, b, h]\n xs[i] = x[r * B * H + b * H + h];\n }\n __syncthreads();\n \n // Each thread computes one output element\n const int k = tid;\n if (k < K) {\n scalar_t sum = bias[k];\n \n // Accumulate dot product over tiles\n // For this small problem, just do it directly\n // Weight access pattern: weight[k, rh] is weight[k * RH + rh]\n // This is coalesced for the same k across threads (reading same row)\n // But we have different k values per thread, so strided access\n \n // Better approach: use __ldg for read-only cache\n const scalar_t* w_row = weight + k * RH;\n for (int rh = 0; rh < RH; rh++) {\n sum += xs[rh] * __ldg(w_row + rh);\n }\n out[b * K + k] = sum;\n }\n}\n\ntorch::Tensor gather_linear_cuda(\n torch::Tensor x, // shape (R, B, H)\n torch::Tensor weight, // shape (K, R*H)\n torch::Tensor bias // shape (K)\n) {\n const int R = x.size(0);\n const int B = x.size(1);\n const int H = x.size(2);\n const int K = weight.size(0);\n \n auto out = torch::empty({B, K}, x.options());\n \n // One block per batch\n const int threads = 128; // Cover K=56 outputs\n const dim3 grid(B);\n \n AT_DISPATCH_FLOATING_TYPES(x.scalar_type(), \"gather_linear_kernel\", ([&] {\n gather_linear_kernel<scalar_t><<<grid, threads>>>(\n x.data_ptr<scalar_t>(),\n weight.data_ptr<scalar_t>(),\n bias.data_ptr<scalar_t>(),\n out.data_ptr<scalar_t>(),\n R, B, H, K\n );\n }));\n \n return out;\n}\n\"\"\"\n\ncpp_source = \"\"\"\ntorch::Tensor gather_linear_cuda(torch::Tensor x, torch::Tensor weight, torch::Tensor bias);\n\"\"\"\n\n# Load the extension\ngather_linear_ext = load_inline(\n name='gather_linear_v3',\n cpp_sources=cpp_source,\n cuda_sources=cuda_source,\n functions=['gather_linear_cuda'],\n verbose=True\n)\n\n\nclass ModelNew(nn.Module):\n def __init__( self):\n super().__init__()\n self.proj = nn.Linear(R * H, K, bias=True)\n \n def forward(self, x: torch.Tensor) -> torch.Tensor:\n # Use the custom CUDA kernel for fused gather + linear\n return gather_linear_ext.gather_linear_cuda(x, self.proj.weight, self.proj.bias)"
}
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 9 · 1 tool call(s) · 175.4s turn · 10.0s LLM · tokens in=10,433 · out=1,064 · total=11,497 · 6,480 chars (click to expand)
turn 9175.4s turn · 10.0s LLM · 165.4s tools · tokens in=10,433 · out=1,064 · total=11,497
reasoning
Let me try a different approach - loading weight in tiles to shared memory for better coalescing. The issue is that each thread reads from a different row of the weight matrix. Let me tile the weight matrix instead, where each tile covers multiple K outputs and multiple RH inputs.
call → profile_kernel
callprofile_kernel
{
"kernel_code": "import torch\nimport torch.nn as nn\nfrom torch.utils.cpp_extension import load_inline\n\nR = 4\nB = 12\nH = 32\nK = 56\n\ncuda_source = \"\"\"\n#include <torch/extension.h>\n#include <cuda.h>\n#include <cuda_runtime.h>\n\n// Tiled GEMV: each block handles one batch sample\n// Threads cooperatively load weight tiles and x into shared memory\ntemplate <typename scalar_t>\n__global__ void __launch_bounds__(128) gather_linear_kernel(\n const scalar_t* __restrict__ x, // shape (R, B, H)\n const scalar_t* __restrict__ weight, // shape (K, R*H) row-major\n const scalar_t* __restrict__ bias, // shape (K)\n scalar_t* __restrict__ out, // shape (B, K)\n int R, int B, int H, int K\n) {\n const int b = blockIdx.x; // batch index\n const int tid = threadIdx.x;\n const int RH = R * H;\n \n // Shared memory for x vector and weight tile\n __shared__ scalar_t xs[128]; // R*H = 128\n __shared__ scalar_t w_tile[128 * 128]; // Max 128x128 = 16KB (fits in smem)\n \n // Load x into shared memory - need to gather from strided locations\n // Each thread loads multiple elements for better coalescing\n // Reorder so thread i loads x[0, b, i], x[1, b, i], etc for H elements\n // Then next thread loads next set of H elements\n for (int i = tid; i < RH; i += blockDim.x) {\n int r = i / H;\n int h = i % H;\n // xs[i] should be x[r, b, h]\n xs[i] = x[r * B * H + b * H + h];\n }\n __syncthreads();\n \n // Each thread computes one output element\n const int k = tid;\n if (k < K) {\n scalar_t sum = bias[k];\n \n // Accumulate dot product over tiles\n // For this small problem, just do it directly\n // Weight access pattern: weight[k, rh] is weight[k * RH + rh]\n // This is coalesced for the same k across threads (reading same row)\n // But we have different k values per thread, so strided access\n \n // Better approach: use __ldg for read-only cache\n const scalar_t* w_row = weight + k * RH;\n for (int rh = 0; rh < RH; rh++) {\n sum += xs[rh] * __ldg(w_row + rh);\n }\n out[b * K + k] = sum;\n }\n}\n\ntorch::Tensor gather_linear_cuda(\n torch::Tensor x, // shape (R, B, H)\n torch::Tensor weight, // shape (K, R*H)\n torch::Tensor bias // shape (K)\n) {\n const int R = x.size(0);\n const int B = x.size(1);\n const int H = x.size(2);\n const int K = weight.size(0);\n \n auto out = torch::empty({B, K}, x.options());\n \n // One block per batch\n const int threads = 128; // Cover K=56 outputs\n const dim3 grid(B);\n \n AT_DISPATCH_FLOATING_TYPES(x.scalar_type(), \"gather_linear_kernel\", ([&] {\n gather_linear_kernel<scalar_t><<<grid, threads>>>(\n x.data_ptr<scalar_t>(),\n weight.data_ptr<scalar_t>(),\n bias.data_ptr<scalar_t>(),\n out.data_ptr<scalar_t>(),\n R, B, H, K\n );\n }));\n \n return out;\n}\n\"\"\"\n\ncpp_source = \"\"\"\ntorch::Tensor gather_linear_cuda(torch::Tensor x, torch::Tensor weight, torch::Tensor bias);\n\"\"\"\n\n# Load the extension\ngather_linear_ext = load_inline(\n name='gather_linear_v3',\n cpp_sources=cpp_source,\n cuda_sources=cuda_source,\n functions=['gather_linear_cuda'],\n verbose=True\n)\n\n\nclass ModelNew(nn.Module):\n def __init__(self):\n super().__init__()\n self.proj = nn.Linear(R * H, K, bias=True)\n \n def forward(self, x: torch.Tensor) -> torch.Tensor:\n # Use the custom CUDA kernel for fused gather + linear\n return gather_linear_ext.gather_linear_cuda(x, self.proj.weight, self.proj.bias)"
}
result ← profile_kernel (ok)
okprofile_kernel
profile_kernel PASSED: profiling complete.
=== Kernel Profile Summary ===
--- Kernel Breakdown (by GPU time) ---
#1: aten::empty 31.0μs (100%) calls=1
--- Timing ---
GPU time: 9.60 μs [↓0.2μs ok]
--- Memory Bandwidth ---
DRAM BW: 4.2/3350 GB/s (0.1%)
read=4.2 write=0.0 GB/s
Hit rates: L2=96.5%
--- Compute Throughput ---
Pipe util: FMA=0.6%, Tensor=0.0%
--- Roofline Analysis ---
Ridge point (FP32): 20.00 FLOPs/byte
Classification: UNKNOWN
--- Occupancy ---
Achieved: 3.5%
Eligible warps/cycle: 0.03
WARNING: <1 eligible warp/cycle = latency-bound, needs more parallelism
Launch: regs/thread=32, smem/blk=1536B, block=128, grid=12
--- Memory Access Quality ---
Global ld coalescing: 27.45 sectors/req (27.4x wasted BW)
--- Warp Stall Reasons (% of issue cycles) ---
long_scoreboard 2672.4% ####################
short_scoreboard 89.6% #################
--- Optimization Hints ---
-> 2672% stalls: global memory latency (waiting for DRAM/L2). + sectors/req=27.4 -> uncoalesced global loads, restructure access pattern or use __ldg Try shared-memory tiling, prefetching, or reducing working set.
-> Global load coalescing=27.4x (ideal=1.0) -> 27x wasted DRAM BW. Make per-warp accesses contiguous.
-> Low occupancy (4%). Consider __launch_bounds__ or reducing register pressure.
-> Low eligible warps + low BW% + low compute% = LATENCY-BOUND. Need more parallelism (larger grid, more ILP, prefetching).
--- Delta vs Previous Iteration ---
GPU time: 9.82→9.60μs (↓0.22) ok
L2 hit%: 96.18→96.48% (↑0.29) ok
turn 10 · 1 tool call(s) · 193.6s turn · 17.2s LLM · tokens in=11,952 · out=1,751 · total=13,703 · 7,169 chars (click to expand)
turn 10193.6s turn · 17.2s LLM · 176.4s tools · tokens in=11,952 · out=1,751 · total=13,703
reasoning
Still has memory coalescing issues. The issue is in how I load x values - the access pattern x[r * B * H + b * H + h] is not coalesced because different threads access scattered memory locations.
Let me think more carefully about the memory access patterns:
1. Input x: shape (R, B, H) = (4, 12, 32)
- x is stored as x[r, b, h] with strides (B*H, H, 1) = (384, 32, 1)
- For a given batch b, the data is scattered across memory
- For batch b=0, the indices are x[0,0,:], x[1,0,:], x[2,0,:], x[3,0,:] which are at offsets 0-31, 384-415, 768-799, 1152-1183
- These are not contiguous!
2. Weight: shape (K, R*H) = (56, 128)
- Stored row-major, so weight[k, rh] is at offset k * 128 + rh
- Different threads accessing weight[tid * 128 + rh] have strided access pattern
The core issue is that the input x has data scattered in memory for a single batch.
Actually, wait - let me reconsider. The problem is small enough that this might just be a latency-bound problem that we can't easily fix. The total data is tiny:
- Input: 4 * 12 * 32 = 1536 floats = 6KB
- Weight: 56 * 128 = 7168 floats = 28KB
- Output: 12 * 56 = 672 floats = 2.6KB
- Total: ~37KB
This fits entirely in L2 cache, so we're mostly latency-bound.
The best optimization here might be to increase ILP (instruction-level parallelism) within each thread - unrolling the loop. Or we could try loading weight into shared memory in a tiled fashion.
Let me try a different approach: use 1 warp per output, and have threads within a warp cooperatively load the weight row into shared memory, then do the dot product. This would give better coalescing for weight access.
assistantThe problem is latency-bound with such small data. Let me try a different approach - use warp-level cooperation with better weight loading pattern:
call → compile_kernel
callcompile_kernel
{
"kernel_code": "import torch\nimport torch.nn as nn\nfrom torch.utils.cpp_extension import load_inline\n\nR = 4\nB = 12\nH = 32\nK = 56\n\ncuda_source = \"\"\"\n#include <torch/extension.h>\n#include <cuda.h>\n#include <cuda_runtime.h>\n\n// Warp-level GEMV with cooperative weight loading\ntemplate <typename scalar_t>\n__global__ void __launch_bounds__(128) gather_linear_kernel(\n const scalar_t* __restrict__ x, // shape (R, B, H)\n const scalar_t* __restrict__ weight, // shape (K, R*H) row-major\n const scalar_t* __restrict__ bias, // shape (K)\n scalar_t* __restrict__ out, // shape (B, K)\n int R, int B, int H, int K\n) {\n const int b = blockIdx.x; // batch index\n const int tid = threadIdx.x;\n const int lane_id = tid & 31; // lane within warp\n const int warp_id = tid >> 5; // warp id within block\n const int RH = R * H;\n \n // Shared memory for x vector\n __shared__ scalar_t xs[128]; // R*H = 128\n \n // Load x into shared memory cooperatively\n // Issue: x is stored as (R, B, H), so x[r, b, h] has stride H between r and h is contiguous\n // For batch b, x[r,b,:] are at offsets r*B*H + b*H + h\n // These are scattered, but we need to load them anyway\n for (int i = tid; i < RH; i += blockDim.x) {\n int r = i / H;\n int h = i % H;\n xs[i] = x[r * B * H + b * H + h];\n }\n __syncthreads();\n \n // Each warp computes one or more outputs\n // 128 threads = 4 warps, K=56 outputs\n // Warp 0: k=0-13, Warp 1: k=14-27, Warp 2: k=28-41, Warp 3: k=42-55\n const int outputs_per_warp = (K + 3) / 4; // ceil(K/4) = 14\n const int k_start = warp_id * outputs_per_warp;\n const int k_end = min(k_start + outputs_per_warp, K);\n \n for (int k = k_start; k < k_end; k++) {\n scalar_t sum = bias[k];\n \n // All lanes in warp cooperatively load weight row k\n // Each lane loads RH/32 elements (128/32 = 4 elements per lane)\n // Weight[k, rh] is at offset k * RH + rh\n // Lanes loading adjacent rh values -> coalesced!\n \n const scalar_t* w_row = weight + k * RH;\n \n // Each lane computes partial dot product\n // Lane i computes sum of xs[i*4:(i+1)*4] * w_row[i*4:(i+1)*4]\n #pragma unroll\n for (int i = 0; i < 4; i++) { // 4 elements per lane\n int rh = lane_id * 4 + i;\n if (rh < RH) {\n sum += xs[rh] * w_row[rh];\n }\n }\n \n // Reduce within warp using shuffle\n #pragma unroll\n for (int offset = 16; offset > 0; offset >>= 1) {\n sum += __shfl_down_sync(0xffffffff, sum, offset);\n }\n \n // Lane 0 writes result\n if (lane_id == 0) {\n out[b * K + k] = sum;\n }\n }\n}\n\ntorch::Tensor gather_linear_cuda(\n torch::Tensor x, // shape (R, B, H)\n torch::Tensor weight, // shape (K, R*H)\n torch::Tensor bias // shape (K)\n) {\n const int R = x.size(0);\n const int B = x.size(1);\n const int H = x.size(2);\n const int K = weight.size(0);\n \n auto out = torch::empty({B, K}, x.options());\n \n // One block per batch, 128 threads = 4 warps\n const int threads = 128;\n const dim3 grid(B);\n \n AT_DISPATCH_FLOATING_TYPES(x.scalar_type(), \"gather_linear_kernel\", ([&] {\n gather_linear_kernel<scalar_t><<<grid, threads>>>(\n x.data_ptr<scalar_t>(),\n weight.data_ptr<scalar_t>(),\n bias.data_ptr<scalar_t>(),\n out.data_ptr<scalar_t>(),\n R, B, H, K\n );\n }));\n \n return out;\n}\n\"\"\"\n\ncpp_source = \"\"\"\ntorch::Tensor gather_linear_cuda(torch::Tensor x, torch::Tensor weight, torch::Tensor bias);\n\"\"\"\n\n# Load the extension\ngather_linear_ext = load_inline(\n name='gather_linear_v4',\n cpp_sources=cpp_source,\n cuda_sources=cuda_source,\n functions=['gather_linear_cuda'],\n verbose=True\n)\n\n\nclass ModelNew(nn.Module):\n def __init__(self):\n super().__init__()\n self.proj = nn.Linear(R * H, K, bias=True)\n \n def forward(self, x: torch.Tensor) -> torch.Tensor:\n # Use the custom CUDA kernel for fused gather + linear\n return gather_linear_ext.gather_linear_cuda(x, self.proj.weight, self.proj.bias)"
}
result ← compile_kernel (ok)
okcompile_kernel
compile_kernel PASSED: kernel compiled without errors.
turn 11409.8s turn · 11.0s LLM · 398.8s tools · tokens in=13,228 · out=1,244 · total=14,472
reasoning
Good, let me 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\nR = 4\nB = 12\nH = 32\nK = 56\n\ncuda_source = \"\"\"\n#include <torch/extension.h>\n#include <cuda.h>\n#include <cuda_runtime.h>\n\n// Warp-level GEMV with cooperative weight loading\ntemplate <typename scalar_t>\n__global__ void __launch_bounds__(128) gather_linear_kernel(\n const scalar_t* __restrict__ x, // shape (R, B, H)\n const scalar_t* __restrict__ weight, // shape (K, R*H) row-major\n const scalar_t* __restrict__ bias, // shape (K)\n scalar_t* __restrict__ out, // shape (B, K)\n int R, int B, int H, int K\n) {\n const int b = blockIdx.x; // batch index\n const int tid = threadIdx.x;\n const int lane_id = tid & 31; // lane within warp\n const int warp_id = tid >> 5; // warp id within block\n const int RH = R * H;\n \n // Shared memory for x vector\n __shared__ scalar_t xs[128]; // R*H = 128\n \n // Load x into shared memory cooperatively\n // Issue: x is stored as (R, B, H), so x[r, b, h] has stride H between r and h is contiguous\n // For batch b, x[r,b,:] are at offsets r*B*H + b*H + h\n // These are scattered, but we need to load them anyway\n for (int i = tid; i < RH; i += blockDim.x) {\n int r = i / H;\n int h = i % H;\n xs[i] = x[r * B * H + b * H + h];\n }\n __syncthreads();\n \n // Each warp computes one or more outputs\n // 128 threads = 4 warps, K=56 outputs\n // Warp 0: k=0-13, Warp 1: k=14-27, Warp 2: k=28-41, Warp 3: k=42-55\n const int outputs_per_warp = (K + 3) / 4; // ceil(K/4) = 14\n const int k_start = warp_id * outputs_per_warp;\n const int k_end = min(k_start + outputs_per_warp, K);\n \n for (int k = k_start; k < k_end; k++) {\n scalar_t sum = bias[k];\n \n // All lanes in warp cooperatively load weight row k\n // Each lane loads RH/32 elements (128/32 = 4 elements per lane)\n // Weight[k, rh] is at offset k * RH + rh\n // Lanes loading adjacent rh values -> coalesced!\n \n const scalar_t* w_row = weight + k * RH;\n \n // Each lane computes partial dot product\n // Lane i computes sum of xs[i*4:(i+1)*4] * w_row[i*4:(i+1)*4]\n #pragma unroll\n for (int i = 0; i < 4; i++) { // 4 elements per lane\n int rh = lane_id * 4 + i;\n if (rh < RH) {\n sum += xs[rh] * w_row[rh];\n }\n }\n \n // Reduce within warp using shuffle\n #pragma unroll\n for (int offset = 16; offset > 0; offset >>= 1) {\n sum += __shfl_down_sync(0xffffffff, sum, offset);\n }\n \n // Lane 0 writes result\n if (lane_id == 0) {\n out[b * K + k] = sum;\n }\n }\n}\n\ntorch::Tensor gather_linear_cuda(\n torch::Tensor x, // shape (R, B, H)\n torch::Tensor weight, // shape (K, R*H)\n torch::Tensor bias // shape (K)\n) {\n const int R = x.size(0);\n const int B = x.size(1);\n const int H = x.size(2);\n const int K = weight.size(0);\n \n auto out = torch::empty({B, K}, x.options());\n \n // One block per batch, 128 threads = 4 warps\n const int threads = 128;\n const dim3 grid(B);\n \n AT_DISPATCH_FLOATING_TYPES(x.scalar_type(), \"gather_linear_kernel\", ([&] {\n gather_linear_kernel<scalar_t><<<grid, threads>>>(\n x.data_ptr<scalar_t>(),\n weight.data_ptr<scalar_t>(),\n bias.data_ptr<scalar_t>(),\n out.data_ptr<scalar_t>(),\n R, B, H, K\n );\n }));\n \n return out;\n}\n\"\"\"\n\ncpp_source = \"\"\"\ntorch::Tensor gather_linear_cuda(torch::Tensor x, torch::Tensor weight, torch::Tensor bias);\n\"\"\"\n\n# Load the extension\ngather_linear_ext = load_inline(\n name='gather_linear_v4',\n cpp_sources=cpp_source,\n cuda_sources=cuda_source,\n functions=['gather_linear_cuda'],\n verbose=True\n)\n\n\nclass ModelNew(nn.Module):\n def __init__(self):\n super().__init__()\n self.proj = nn.Linear(R * H, K, bias=True)\n \n def forward(self, x: torch.Tensor) -> torch.Tensor:\n # Use the custom CUDA kernel for fused gather + linear\n return gather_linear_ext.gather_linear_cuda(x, self.proj.weight, self.proj.bias)"
}
result ← run_correctness (fail)
failrun_correctness
run_correctness FAILED: (0 / 3) trials did not all match.
correctness_issue: Output mismatch
max_difference: ['2.683646', '2.683646', '2.683646']
avg_difference: ['1.256690', '1.256690', '1.256690']
turn 12 · 1 tool call(s) · 187.7s turn · 16.0s LLM · tokens in=14,575 · out=1,364 · total=15,939 · 6,275 chars (click to expand)
turn 12187.7s turn · 16.0s LLM · 171.8s tools · tokens in=14,575 · out=1,364 · total=15,939
reasoning
The correctness failed. There's an issue with my warp-level approach. The problem is that all threads in a warp are iterating over different k values but loading the same weight row during their respective iterations. Let me think more carefully.
Actually, the issue is that each thread in a warp is computing different k values, but they're all trying to do a warp reduction - but they should be doing independent reductions for different k values. The shuffle reduction is wrong because it's mixing partial sums from different k computations.
Let me go back to a simpler approach that works - the v2 version that passed correctness. That one was 9.82μs vs 9.60μs of v3. Since these are tiny differences and the kernel is fundamentally latency-bound due to the small problem size, I should just use the working version.
Actually, let me look at the working v2 kernel again and try to optimize it differently. The main issue is:
1. Poor coalescing when loading x (due to the (R, B, H) layout)
2. Poor coalescing when loading weight (strided access across threads)
For the weight loading, since K=56 and we have 128 threads per block (one block per batch), threads 0-55 each compute one output. The weight access for thread k is weight[k * RH + rh], which is strided by RH=128 between threads. This is very bad for coalescing.
One way to fix this: have all threads load the same weight column (same rh value) at the same time, then accumulate into per-thread partial sums. This way weight[rh] is loaded once and used by all threads.
Let me try a column-based approach where all threads in a warp load one column of weight at a time.
assistantThe correctness failed because my warp-level reduction was incorrect. Let me try a column-based approach where all threads cooperate to load weight columns:
call → compile_kernel
callcompile_kernel
{
"kernel_code": "import torch\nimport torch.nn as nn\nfrom torch.utils.cpp_extension import load_inline\n\nR = 4\nB = 12\nH = 32\nK = 56\n\ncuda_source = \"\"\"\n#include <torch/extension.h>\n#include <cuda.h>\n#include <cuda_runtime.h>\n\n// Column-based approach: all threads load same weight column at once\ntemplate <typename scalar_t>\n__global__ void __launch_bounds__(128) gather_linear_kernel(\n const scalar_t* __restrict__ x, // shape (R, B, H)\n const scalar_t* __restrict__ weight, // shape (K, R*H) row-major\n const scalar_t* __restrict__ bias, // shape (K)\n scalar_t* __restrict__ out, // shape (B, K)\n int R, int B, int H, int K\n) {\n const int b = blockIdx.x; // batch index\n const int tid = threadIdx.x;\n const int RH = R * H;\n \n // Shared memory for x vector\n __shared__ scalar_t xs[128]; // R*H = 128\n // Shared memory for a column of weight (K values)\n __shared__ scalar_t w_col[128]; // K values, padded to 128\n \n // Load x into shared memory cooperatively\n for (int i = tid; i < RH; i += blockDim.x) {\n int r = i / H;\n int h = i % H;\n xs[i] = x[r * B * H + b * H + h];\n }\n __syncthreads();\n \n // Initialize output with bias\n scalar_t sum = (tid < K) ? bias[tid] : scalar_t(0);\n \n // Process weight columns one at a time\n // All threads load the same column rh at once, coalesced!\n for (int rh = 0; rh < RH; rh++) {\n // Load weight column rh: weight[k, rh] = weight[k * RH + rh]\n // Threads 0..K-1 load weight[0..K-1, rh] -> offsets 0*RH+rh, 1*RH+rh, etc.\n // These are strided by RH=128, NOT coalesced!\n // But we can use __ldg for read-only cache\n \n scalar_t w_val = (tid < K) ? __ldg(&weight[tid * RH + rh]) : scalar_t(0);\n w_col[tid] = w_val;\n __syncthreads();\n \n // Multiply by x[rh] and accumulate\n sum += xs[rh] * w_col[tid];\n }\n \n // Write result\n if (tid < K) {\n out[b * K + tid] = sum;\n }\n}\n\ntorch::Tensor gather_linear_cuda(\n torch::Tensor x, // shape (R, B, H)\n torch::Tensor weight, // shape (K, R*H)\n torch::Tensor bias // shape (K)\n) {\n const int R = x.size(0);\n const int B = x.size(1);\n const int H = x.size(2);\n const int K = weight.size(0);\n \n auto out = torch::empty({B, K}, x.options());\n \n // One block per batch, 128 threads\n const int threads = 128;\n const dim3 grid(B);\n \n AT_DISPATCH_FLOATING_TYPES(x.scalar_type(), \"gather_linear_kernel\", ([&] {\n gather_linear_kernel<scalar_t><<<grid, threads>>>(\n x.data_ptr<scalar_t>(),\n weight.data_ptr<scalar_t>(),\n bias.data_ptr<scalar_t>(),\n out.data_ptr<scalar_t>(),\n R, B, H, K\n );\n }));\n \n return out;\n}\n\"\"\"\n\ncpp_source = \"\"\"\ntorch::Tensor gather_linear_cuda(torch::Tensor x, torch::Tensor weight, torch::Tensor bias);\n\"\"\"\n\ngather_linear_ext = load_inline(\n name='gather_linear_v5',\n cpp_sources=cpp_source,\n cuda_sources=cuda_source,\n functions=['gather_linear_cuda'],\n verbose=True\n)\n\n\nclass ModelNew(nn.Module):\n def __init__(self):\n super().__init__()\n self.proj = nn.Linear(R * H, K, bias=True)\n \n def forward(self, x: torch.Tensor) -> torch.Tensor:\n return gather_linear_ext.gather_linear_cuda(x, self.proj.weight, self.proj.bias)"
}
result ← compile_kernel (ok)
okcompile_kernel
compile_kernel PASSED: kernel compiled without errors.
turn 13178.0s turn · 9.7s LLM · 168.3s tools · tokens in=15,618 · out=995 · total=16,613
reasoning
Let me run correctness on this version to see if it works, and if it does, I can submit it.
call → run_correctness
callrun_correctness
{
"kernel_code": "import torch\nimport torch.nn as nn\nfrom torch.utils.cpp_extension import load_inline\n\nR = 4\nB = 12\nH = 32\nK = 56\n\ncuda_source = \"\"\"\n#include <torch/extension.h>\n#include <cuda.h>\n#include <cuda_runtime.h>\n\n// Column-based approach: all threads load same weight column at once\ntemplate <typename scalar_t>\n__global__ void __launch_bounds__(128) gather_linear_kernel(\n const scalar_t* __restrict__ x, // shape (R, B, H)\n const scalar_t* __restrict__ weight, // shape (K, R*H) row-major\n const scalar_t* __restrict__ bias, // shape (K)\n scalar_t* __restrict__ out, // shape (B, K)\n int R, int B, int H, int K\n) {\n const int b = blockIdx.x; // batch index\n const int tid = threadIdx.x;\n const int RH = R * H;\n \n // Shared memory for x vector\n __shared__ scalar_t xs[128]; // R*H = 128\n // Shared memory for a column of weight (K values)\n __shared__ scalar_t w_col[128]; // K values, padded to 128\n \n // Load x into shared memory cooperatively\n for (int i = tid; i < RH; i += blockDim.x) {\n int r = i / H;\n int h = i % H;\n xs[i] = x[r * B * H + b * H + h];\n }\n __syncthreads();\n \n // Initialize output with bias\n scalar_t sum = (tid < K) ? bias[tid] : scalar_t(0);\n \n // Process weight columns one at a time\n // All threads load the same column rh at once, coalesced!\n for (int rh = 0; rh < RH; rh++) {\n // Load weight column rh: weight[k, rh] = weight[k * RH + rh]\n // Threads 0..K-1 load weight[0..K-1, rh] -> offsets 0*RH+rh, 1*RH+rh, etc.\n // These are strided by RH=128, NOT coalesced!\n // But we can use __ldg for read-only cache\n \n scalar_t w_val = (tid < K) ? __ldg(&weight[tid * RH + rh]) : scalar_t(0);\n w_col[tid] = w_val;\n __syncthreads();\n \n // Multiply by x[rh] and accumulate\n sum += xs[rh] * w_col[tid];\n }\n \n // Write result\n if (tid < K) {\n out[b * K + tid] = sum;\n }\n}\n\ntorch::Tensor gather_linear_cuda(\n torch::Tensor x, // shape (R, B, H)\n torch::Tensor weight, // shape (K, R*H)\n torch::Tensor bias // shape (K)\n) {\n const int R = x.size(0);\n const int B = x.size(1);\n const int H = x.size(2);\n const int K = weight.size(0);\n \n auto out = torch::empty({B, K}, x.options());\n \n // One block per batch, 128 threads\n const int threads = 128;\n const dim3 grid(B);\n \n AT_DISPATCH_FLOATING_TYPES(x.scalar_type(), \"gather_linear_kernel\", ([&] {\n gather_linear_kernel<scalar_t><<<grid, threads>>>(\n x.data_ptr<scalar_t>(),\n weight.data_ptr<scalar_t>(),\n bias.data_ptr<scalar_t>(),\n out.data_ptr<scalar_t>(),\n R, B, H, K\n );\n }));\n \n return out;\n}\n\"\"\"\n\ncpp_source = \"\"\"\ntorch::Tensor gather_linear_cuda(torch::Tensor x, torch::Tensor weight, torch::Tensor bias);\n\"\"\"\n\ngather_linear_ext = load_inline(\n name='gather_linear_v5',\n cpp_sources=cpp_source,\n cuda_sources=cuda_source,\n functions=['gather_linear_cuda'],\n verbose=True\n)\n\n\nclass ModelNew(nn.Module):\n def __init__(self):\n super().__init__()\n self.proj = nn.Linear(R * H, K, bias=True)\n \n def forward(self, x: torch.Tensor) -> torch.Tensor:\n return gather_linear_ext.gather_linear_cuda(x, self.proj.weight, self.proj.bias)"
}
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 14164.4s turn · 8.7s LLM · 155.7s tools · tokens in=16,736 · out=982 · total=17,718 · FINAL
reasoning
The kernel is correct. Let me submit it now.
call → submit_kernel
callsubmit_kernel
{
"kernel_code": "import torch\nimport torch.nn as nn\nfrom torch.utils.cpp_extension import load_inline\n\nR = 4\nB = 12\nH = 32\nK = 56\n\ncuda_source = \"\"\"\n#include <torch/extension.h>\n#include <cuda.h>\n#include <cuda_runtime.h>\n\n// Column-based approach: all threads load same weight column at once\ntemplate <typename scalar_t>\n__global__ void __launch_bounds__(128) gather_linear_kernel(\n const scalar_t* __restrict__ x, // shape (R, B, H)\n const scalar_t* __restrict__ weight, // shape (K, R*H) row-major\n const scalar_t* __restrict__ bias, // shape (K)\n scalar_t* __restrict__ out, // shape (B, K)\n int R, int B, int H, int K\n) {\n const int b = blockIdx.x; // batch index\n const int tid = threadIdx.x;\n const int RH = R * H;\n \n // Shared memory for x vector\n __shared__ scalar_t xs[128]; // R*H = 128\n // Shared memory for a column of weight (K values)\n __shared__ scalar_t w_col[128]; // K values, padded to 128\n \n // Load x into shared memory cooperatively\n for (int i = tid; i < RH; i += blockDim.x) {\n int r = i / H;\n int h = i % H;\n xs[i] = x[r * B * H + b * H + h];\n }\n __syncthreads();\n \n // Initialize output with bias\n scalar_t sum = (tid < K) ? bias[tid] : scalar_t(0);\n \n // Process weight columns one at a time\n // All threads load the same column rh at once, coalesced!\n for (int rh = 0; rh < RH; rh++) {\n // Load weight column rh: weight[k, rh] = weight[k * RH + rh]\n // Threads 0..K-1 load weight[0..K-1, rh] -> offsets 0*RH+rh, 1*RH+rh, etc.\n // These are strided by RH=128, NOT coalesced!\n // But we can use __ldg for read-only cache\n \n scalar_t w_val = (tid < K) ? __ldg(&weight[tid * RH + rh]) : scalar_t(0);\n w_col[tid] = w_val;\n __syncthreads();\n \n // Multiply by x[rh] and accumulate\n sum += xs[rh] * w_col[tid];\n }\n \n // Write result\n if (tid < K) {\n out[b * K + tid] = sum;\n }\n}\n\ntorch::Tensor gather_linear_cuda(\n torch::Tensor x, // shape (R, B, H)\n torch::Tensor weight, // shape (K, R*H)\n torch::Tensor bias // shape (K)\n) {\n const int R = x.size(0);\n const int B = x.size(1);\n const int H = x.size(2);\n const int K = weight.size(0);\n \n auto out = torch::empty({B, K}, x.options());\n \n // One block per batch, 128 threads\n const int threads = 128;\n const dim3 grid(B);\n \n AT_DISPATCH_FLOATING_TYPES(x.scalar_type(), \"gather_linear_kernel\", ([&] {\n gather_linear_kernel<scalar_t><<<grid, threads>>>(\n x.data_ptr<scalar_t>(),\n weight.data_ptr<scalar_t>(),\n bias.data_ptr<scalar_t>(),\n out.data_ptr<scalar_t>(),\n R, B, H, K\n );\n }));\n \n return out;\n}\n\"\"\"\n\ncpp_source = \"\"\"\ntorch::Tensor gather_linear_cuda(torch::Tensor x, torch::Tensor weight, torch::Tensor bias);\n\"\"\"\n\ngather_linear_ext = load_inline(\n name='gather_linear_v5',\n cpp_sources=cpp_source,\n cuda_sources=cuda_source,\n functions=['gather_linear_cuda'],\n verbose=True\n)\n\n\nclass ModelNew(nn.Module):\n def __init__(self):\n super().__init__()\n self.proj = nn.Linear(R * H, K, bias=True)\n \n def forward(self, x: torch.Tensor) -> torch.Tensor:\n return gather_linear_ext.gather_linear_cuda(x, self.proj.weight, self.proj.bias)"
}
result ← submit_kernel (ok)
oksubmit_kernel
submit_kernel PASSED: (3 / 3) correctness trials all passed.
Kernel runtime: 0.02 μs
Runtime stats: mean=0.02μ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.0009
final submitted kernel (115 lines, 3,365 chars) — full source, untruncated
import torch
import torch.nn as nn
from torch.utils.cpp_extension import load_inline
R = 4
B = 12
H = 32
K = 56
cuda_source = """
#include <torch/extension.h>
#include <cuda.h>
#include <cuda_runtime.h>
// Column-based approach: all threads load same weight column at once
template <typename scalar_t>
__global__ void __launch_bounds__(128) gather_linear_kernel(
const scalar_t* __restrict__ x, // shape (R, B, H)
const scalar_t* __restrict__ weight, // shape (K, R*H) row-major
const scalar_t* __restrict__ bias, // shape (K)
scalar_t* __restrict__ out, // shape (B, K)
int R, int B, int H, int K
) {
const int b = blockIdx.x; // batch index
const int tid = threadIdx.x;
const int RH = R * H;
// Shared memory for x vector
__shared__ scalar_t xs[128]; // R*H = 128
// Shared memory for a column of weight (K values)
__shared__ scalar_t w_col[128]; // K values, padded to 128
// Load x into shared memory cooperatively
for (int i = tid; i < RH; i += blockDim.x) {
int r = i / H;
int h = i % H;
xs[i] = x[r * B * H + b * H + h];
}
__syncthreads();
// Initialize output with bias
scalar_t sum = (tid < K) ? bias[tid] : scalar_t(0);
// Process weight columns one at a time
// All threads load the same column rh at once, coalesced!
for (int rh = 0; rh < RH; rh++) {
// Load weight column rh: weight[k, rh] = weight[k * RH + rh]
// Threads 0..K-1 load weight[0..K-1, rh] -> offsets 0*RH+rh, 1*RH+rh, etc.
// These are strided by RH=128, NOT coalesced!
// But we can use __ldg for read-only cache
scalar_t w_val = (tid < K) ? __ldg(&weight[tid * RH + rh]) : scalar_t(0);
w_col[tid] = w_val;
__syncthreads();
// Multiply by x[rh] and accumulate
sum += xs[rh] * w_col[tid];
}
// Write result
if (tid < K) {
out[b * K + tid] = sum;
}
}
torch::Tensor gather_linear_cuda(
torch::Tensor x, // shape (R, B, H)
torch::Tensor weight, // shape (K, R*H)
torch::Tensor bias // shape (K)
) {
const int R = x.size(0);
const int B = x.size(1);
const int H = x.size(2);
const int K = weight.size(0);
auto out = torch::empty({B, K}, x.options());
// One block per batch, 128 threads
const int threads = 128;
const dim3 grid(B);
AT_DISPATCH_FLOATING_TYPES(x.scalar_type(), "gather_linear_kernel", ([&] {
gather_linear_kernel<scalar_t><<<grid, threads>>>(
x.data_ptr<scalar_t>(),
weight.data_ptr<scalar_t>(),
bias.data_ptr<scalar_t>(),
out.data_ptr<scalar_t>(),
R, B, H, K
);
}));
return out;
}
"""
cpp_source = """
torch::Tensor gather_linear_cuda(torch::Tensor x, torch::Tensor weight, torch::Tensor bias);
"""
gather_linear_ext = load_inline(
name='gather_linear_v5',
cpp_sources=cpp_source,
cuda_sources=cuda_source,
functions=['gather_linear_cuda'],
verbose=True
)
class ModelNew(nn.Module):
def __init__(self):
super().__init__()
self.proj = nn.Linear(R * H, K, bias=True)
def forward(self, x: torch.Tensor) -> torch.Tensor:
return gather_linear_ext.gather_linear_cuda(x, self.proj.weight, self.proj.bias)