correct
30_SelectiveScan_Mamba.py
started
2026-05-11T19:00:28Z
timingkernelbench.timing
kernel mean / median / std
3.94 / — / 0.05 μs
ref mean / median / std
165.00 / — / 140.00 μs
numerical precisioneval_kernel_against_ref
gpu memoryextended_metrics.measure_memory
kernel launches / fusiontorch.profiler
fusion ratio (ref / kernel)
889.37x
energyNVML (pynvml)
ratio (ref / kernel)
21.51x
SOL (speed-of-light)nsight_hardware_counters
roofline / occupancynsight_hardware_counters
dram bandwidth
6106.86 GB/s
warp stalls (top 5)nsight
reference PyTorch source (the prompt the agent sees)
import torch
import torch.nn as nn
class Model(nn.Module):
"""
Selective scan (Mamba's S6 layer): projects input to derive input-dependent SSM parameters
(delta, B, C), discretizes a learned state matrix A, runs the linear recurrence, and
applies a skip connection via D. The fusion opportunity is keeping all intermediates
(discretized parameters, hidden states) in SRAM rather than materializing them to global memory.
Args:
d_model (int): Input feature dimension.
d_state (int): Hidden state dimension per channel (N in Mamba, typically 16).
"""
def __init__(self, d_model: int, d_state: int):
super(Model, self).__init__()
self.d_model = d_model
self.d_state = d_state
self.A_log = nn.Parameter(torch.randn(d_model, d_state))
self.proj_B = nn.Linear(d_model, d_state, bias=False)
self.proj_C = nn.Linear(d_model, d_state, bias=False)
self.proj_delta = nn.Linear(d_model, d_model, bias=True)
self.D = nn.Parameter(torch.ones(d_model))
def forward(self, x: torch.Tensor) -> torch.Tensor:
B, L, D = x.shape
N = self.d_state
delta = torch.nn.functional.softplus(self.proj_delta(x))
B_t = self.proj_B(x)
C_t = self.proj_C(x)
A = -torch.exp(self.A_log)
dA = torch.exp(delta.unsqueeze(-1) * A.unsqueeze(0).unsqueeze(0))
dB = delta.unsqueeze(-1) * B_t.unsqueeze(2)
x_db = x.unsqueeze(-1) * dB
h = torch.zeros(B, D, N, device=x.device, dtype=x.dtype)
outputs = []
for t in range(L):
h = dA[:, t] * h + x_db[:, t]
y_t = (h * C_t[:, t].unsqueeze(1)).sum(dim=-1)
outputs.append(y_t)
y = torch.stack(outputs, dim=1)
y = y + x * self.D.unsqueeze(0).unsqueeze(0)
return y
batch_size = 16
d_model = 256
d_state = 16
seq_len = 2048
def get_inputs():
return [torch.randn(popcorn_pri.jitter_int(batch_size), popcorn_pri.jitter_int(seq_len), d_model)]
def get_init_inputs():
return [d_model, d_state]turn 0109.8s turn · 10.7s LLM · 99.1s tools · tokens in=2,300 · out=6 · total=3,046
reasoning
The task is to replace the forward computation of this PyTorch model with a custom CUDA kernel for a selective scan, like in Mamba's S6 layer.
I need to replace the forward computation of the PyTorch model with a custom CUDA kernel for a selective scan, inspired by Mamba's S6 layer. The goal is to speed up the existing slow Python loop over the sequence length of 2048.
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 1218.7s turn · 9.4s LLM · 209.3s tools · tokens in=3,279 · out=373 · total=3,887
reasoning
H100, good.
Now, I need to implement 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 = \"\"\"\n#include <torch/extension.h>\n#include <cuda.h>\n#include <cuda_runtime.h>\n\n__global__ void selective_scan_kernel(\n const float* __restrict__ x,\n const float* __restrict__ delta,\n const float* __restrict__ B_t,\n const float* __restrict__ C_t,\n const float* __restrict__ A,\n const float* __restrict__ D,\n float* __restrict__ y,\n int B, int L, int D, int N\n) {\n int idx = blockIdx.x * blockDim.x + threadIdx.x;\n int total = B * D;\n if (idx >= total) return;\n\n int b = idx / D;\n int d = idx % D;\n\n float h[16]; // N <=16\n for (int n=0; n<N; n++) h[n] = 0.0f;\n\n for (int t=0; t<L; t++) {\n float dA = expf(delta[(b*L + t)*D + d] * A[d*N + n]); wait, fix\n }\n}\n\"\"\"\n\ncpp_source = \"\"\"\ntorch::Tensor selective_scan(torch::Tensor x, torch::Tensor A_log, torch::Tensor proj_B_weight, torch::Tensor proj_C_weight, torch::Tensor proj_delta_weight, torch::Tensor proj_delta_bias, torch::Tensor D);\n\"\"\"\n\nmodule = load_inline(\n name=\"selective_scan\",\n cpp_sources=cpp_source,\n cuda_sources=cuda_source,\n functions=[\"selective_scan\"],\n verbose=True,\n extra_cuda_cflags=[\"-O3\"]\n)\nprint(\"Compiled successfully\")\n"
}
result ← compile_kernel (fail)
failcompile_kernel
compile_kernel FAILED: RuntimeError.
RuntimeError: Error building extension 'selective_scan'
turn 259.9s turn · 12.4s LLM · 47.5s tools · tokens in=3,911 · out=784 · total=4,724
reasoning
Failed because incomplete code. I need to write proper code.
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 mamba_scan_kernel(\n const float* __restrict__ x,\n const float* __restrict__ delta,\n const float* __restrict__ Bt,\n const float* __restrict__ Ct,\n const float* __restrict__ A,\n const float* __restrict__ D_in,\n float* __restrict__ out,\n const int Bsz, const int L, const int D, const int N\n) {\n const int idx = blockIdx.x * blockDim.x + threadIdx.x;\n if (idx >= Bsz * D) return;\n const int b = idx / D;\n const int d = idx % D;\n\n float h[16];\n #pragma unroll\n for(int n=0; n<16; n++) h[n] = 0.0f;\n\n const float d_val = D_in[d];\n\n for(int t=0; t<L; t++) {\n const int base = (b * L + t) * D + d;\n const float del = delta[base];\n const float xt = x[base];\n float yt = 0.0f;\n for(int n=0; n<N; n++) {\n float da = expf(del * A[d * N + n]);\n float db = del * Bt[(b * L + t) * N + n];\n h[n] = da * h[n] + xt * db;\n yt += h[n] * Ct[(b * L + t) * N + n];\n }\n out[base] = yt + xt * d_val;\n }\n}\n\ntorch::Tensor mamba_forward(\n torch::Tensor x,\n torch::Tensor A_log,\n torch::Tensor projB_w,\n torch::Tensor projC_w,\n torch::Tensor projDelta_w,\n torch::Tensor projDelta_b,\n torch::Tensor D\n) {\n auto Bsz = x.size(0);\n auto L = x.size(1);\n auto D = x.size(2);\n auto N = projB_w.size(1);\n\n auto delta = torch::softplus(torch::linear(x, projDelta_w, projDelta_b));\n auto Bt = torch::linear(x, projB_w, torch::Tensor());\n auto Ct = torch::linear(x, projC_w, torch::Tensor());\n auto A = -torch::exp(A_log);\n\n auto out = torch::empty_like(x);\n\n const int threads = 256;\n const int blocks = (Bsz * D + threads - 1) / threads;\n\n mamba_scan_kernel<<<blocks, threads>>>(\n x.data_ptr<float>(),\n delta.data_ptr<float>(),\n Bt.data_ptr<float>(),\n Ct.data_ptr<float>(),\n A.data_ptr<float>(),\n D.data_ptr<float>(),\n out.data_ptr<float>(),\n Bsz, L, D, N\n );\n return out;\n}\n'''\n\ncpp_source = r'''\ntorch::Tensor mamba_forward(\n torch::Tensor x,\n torch::Tensor A_log,\n torch::Tensor projB_w,\n torch::Tensor projC_w,\n torch::Tensor projDelta_w,\n torch::Tensor projDelta_b,\n torch::Tensor D\n);\n'''\n\nmodule = load_inline(\n name='mamba_scan',\n cpp_sources=[cpp_source],\n cuda_sources=[cuda_source],\n functions=['mamba_forward'],\n verbose=True,\n extra_cuda_cflags=['-O3', '--use_fast_math']\n)\n\nprint(\"Compilation successful\")\n"
}
result ← compile_kernel (fail)
failcompile_kernel
compile_kernel FAILED: RuntimeError.
RuntimeError: Error building extension 'mamba_scan'
turn 3 · 1 tool call(s) · 106.6s turn · 15.3s LLM · tokens in=4,749 · out=835 · total=5,745 · 14,657 chars (click to expand)
turn 3106.6s turn · 15.3s LLM · 91.3s tools · tokens in=4,749 · out=835 · total=5,745
reasoning
Still failing. I need to debug by perhaps writing a full file, but tool doesn't give error details.
call → compile_kernel
callcompile_kernel
{
"kernel_code": "import torch\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 mamba_scan_kernel(\n const float* __restrict__ x,\n const float* __restrict__ delta,\n const float* __restrict__ Bt,\n const float* __restrict__ Ct,\n const float* __restrict__ A,\n const float* __restrict__ D_in,\n float* __restrict__ out,\n const int Bsz, const int L, const int Ddim, const int N\n) {\n const int idx = blockIdx.x * blockDim.x + threadIdx.x;\n if (idx >= Bsz * Ddim) return;\n const int b = idx / Ddim;\n const int d = idx % Ddim;\n\n float h[16];\n for(int n = 0; n < N; n++) h[n] = 0.0f;\n\n const float d_val = D_in[d];\n\n for(int t = 0; t < L; t++) {\n const int base = (b * L + t) * Ddim + d;\n const float del = delta[base];\n const float xt = x[base];\n float yt = 0.0f;\n #pragma unroll\n for(int n = 0; n < N; n++) {\n float da = __expf(del * A[d * N + n]);\n float db = del * Bt[(b * L + t) * N + n];\n h[n] = da * h[n] + xt * db;\n yt += h[n] * Ct[(b * L + t) * N + n];\n }\n out[base] = yt + xt * d_val;\n }\n}\n\ntorch::Tensor mamba_forward(\n torch::Tensor x,\n torch::Tensor A_log,\n torch::Tensor projB_w,\n torch::Tensor projC_w,\n torch::Tensor projDelta_w,\n torch::Tensor projDelta_b,\n torch::Tensor D\n) {\n auto Bsz = x.size(0);\n auto L = x.size(1);\n auto Ddim = x.size(2);\n auto N = projB_w.size(1);\n\n auto options = x.options();\n auto delta = torch::softplus(torch::addmm(projDelta_b, x.view(-1, Ddim), projDelta_w.t())).view(Bsz, L, Ddim);\n auto Bt = torch::mm(x.view(-1, Ddim), projB_w.t()).view(Bsz, L, N);\n auto Ct = torch::mm(x.view(-1, Ddim), projC_w.t()).view(Bsz, L, N);\n auto A = -torch::exp(A_log);\n\n auto out = torch::empty_like(x);\n\n const int threads = 256;\n const int blocks = (Bsz * Ddim + threads - 1) / threads;\n\n mamba_scan_kernel<<<blocks, threads>>>(\n x.data_ptr<float>(),\n delta.data_ptr<float>(),\n Bt.data_ptr<float>(),\n Ct.data_ptr<float>(),\n A.data_ptr<float>(),\n D.data_ptr<float>(),\n out.data_ptr<float>(),\n Bsz, L, Ddim, N\n );\n cudaDeviceSynchronize();\n return out;\n}\n'''\n\ncpp_source = r'''\ntorch::Tensor mamba_forward(torch::Tensor x, torch::Tensor A_log, torch::Tensor projB_w, torch::Tensor projC_w, torch::Tensor projDelta_w, torch::Tensor projDelta_b, torch::Tensor D);\n'''\n\nmodule = load_inline(\n name='mamba_scan',\n cpp_sources=cpp_source,\n cuda_sources=cuda_source,\n functions=['mamba_forward'],\n verbose=False,\n extra_cuda_cflags=['-O3', '--use_fast_math']\n)\nprint(\"OK\")\n"
}
result ← compile_kernel (fail)
failcompile_kernel
compile_kernel FAILED: RuntimeError.
RuntimeError: Error building extension 'mamba_scan_v1': [1/3] /opt/conda/bin/x86_64-conda-linux-gnu-c++ -MMD -MF main.o.d -DTORCH_EXTENSION_NAME=mamba_scan_v1 -DTORCH_API_INCLUDE_EXTENSION_H -isystem /scratch/abaner/KernelBench/.venv/lib/python3.10/site-packages/torch/include -isystem /scratch/abaner/KernelBench/.venv/lib/python3.10/site-packages/torch/include/torch/csrc/api/include -isystem /usr/local/cuda/include -isystem /opt/conda/include/python3.10 -fPIC -std=c++17 -c /scratch/abaner/KernelBench/runs/pop_l123_all_gpt/popcorn/grok-4.3/level_2_problem_30_cache/k_38ca479f8021/mamba_scan/main.cpp -o main.o
[2/3] /usr/local/cuda/bin/nvcc -MD -MF cuda.cuda.o.d -ccbin /opt/conda/bin/x86_64-conda-linux-gnu-cc -DTORCH_EXTENSION_NAME=mamba_scan_v1 -DTORCH_API_INCLUDE_EXTENSION_H -isystem /scratch/abaner/KernelBench/.venv/lib/python3.10/site-packages/torch/include -isystem /scratch/abaner/KernelBench/.venv/lib/python3.10/site-packages/torch/include/torch/csrc/api/include -isystem /usr/local/cuda/include -isystem /opt/conda/include/python3.10 -D__CUDA_NO_HALF_OPERATORS__ -D__CUDA_NO_HALF_CONVERSIONS__ -D__CUDA_NO_BFLOAT16_CONVERSIONS__ -D__CUDA_NO_HALF2_OPERATORS__ --expt-relaxed-constexpr -gencode=arch=compute_90,code=compute_90 -gencode=arch=compute_90,code=sm_90 --compiler-options '-fPIC' -O3 --use_fast_math -std=c++17 -c /scratch/abaner/KernelBench/runs/pop_l123_all_gpt/popcorn/grok-4.3/level_2_problem_30_cache/k_38ca479f8021/mamba_scan/cuda.cu -o cuda.cuda.o
FAILED: [code=2] cuda.cuda.o
/usr/local/cuda/bin/nvcc -MD -MF cuda.cuda.o.d -ccbin /opt/conda/bin/x86_64-conda-linux-gnu-cc -DTORCH_EXTENSION_NAME=mamba_scan_v1 -DTORCH_API_INCLUDE_EXTENSION_H -isystem /scratch/abaner/KernelBench/.venv/lib/python3.10/site-packages/torch/include -isystem /scratch/abaner/KernelBench/.venv/lib/python3.10/site-packages/torch/include/torch/csrc/api/include -isystem /usr/local/cuda/include -isystem /opt/conda/include/python3.10 -D__CUDA_NO_HALF_OPERATORS__ -D__CUDA_NO_HALF_CONVERSIONS__ -D__CUDA_NO_BFLOAT16_CONVERSIONS__ -D__CUDA_NO_HALF2_OPERATORS__ --expt-relaxed-constexpr -gencode=arch=compute_90,code=compute_90 -gencode=arch=compute_90,code=sm_90 --compiler-options '-fPIC' -O3 --use_fast_math -std=c++17 -c /scratch/abaner/KernelBench/runs/pop_l123_all_gpt/popcorn/grok-4.3/level_2_problem_30_cache/k_38ca479f8021/mamba_scan/cuda.cu -o cuda.cuda.o
nvcc warning : incompatible redefinition for option 'compiler-bindir', the last value of this option was used
/scratch/abaner/KernelBench/runs/pop_l123_all_gpt/popcorn/grok-4.3/level_2_problem_30_cache/k_38ca479f8021/mamba_scan/cuda.cu(59): error: no instance of overloaded function "at::Tensor::view" matches the argument list
argument types are: (int, int64_t)
object type is: at::Tensor
auto delta = torch::softplus(torch::addmm(projDelta_b, x.view(-1, Ddim), projDelta_w.t())).view(Bsz, L, Ddim);
^
/scratch/abaner/KernelBench/.venv/lib/python3.10/site-packages/torch/include/ATen/core/TensorBody.h(4394): note #3322-D: number of parameters of function "at::Tensor::view(c10::ScalarType) const" does not match the call
inline at::Tensor Tensor::view(at::ScalarType dtype) const {
^
/scratch/abaner/KernelBench/.venv/lib/python3.10/site-packages/torch/include/ATen/core/TensorBody.h(4384): note #3322-D: number of parameters of function "at::Tensor::view(at::IntArrayRef) const" does not match the call
inline at::Tensor Tensor::view(at::IntArrayRef size) const {
^
/scratch/abaner/KernelBench/runs/pop_l123_all_gpt/popcorn/grok-4.3/level_2_problem_30_cache/k_38ca479f8021/mamba_scan/cuda.cu(59): error: no instance of overloaded function "at::Tensor::view" matches the argument list
argument types are: (int64_t, int64_t, int64_t)
object type is: at::Tensor
auto delta = torch::softplus(torch::addmm(projDelta_b, x.view(-1, Ddim), projDelta_w.t())).view(Bsz, L, Ddim);
^
/scratch/abaner/KernelBench/.venv/lib/python3.10/site-packages/torch/include/ATen/core/TensorBody.h(4394): note #3322-D: number of parameters of function "at::Tensor::view(c10::ScalarType) const" does not match the call
inline at::Tensor Tensor::view(at::ScalarType dtype) const {
^
/scratch/abaner/KernelBench/.venv/lib/python3.10/site-packages/torch/include/ATen/core/TensorBody.h(4384): note #3322-D: number of parameters of function "at::Tensor::view(at::IntArrayRef) const" does not match the call
inline at::Tensor Tensor::view(at::IntArrayRef size) const {
^
/scratch/abaner/KernelBench/runs/pop_l123_all_gpt/popcorn/grok-4.3/level_2_problem_30_cache/k_38ca479f8021/mamba_scan/cuda.cu(60): error: no instance of overloaded function "at::Tensor::view" matches the argument list
argument types are: (int, int64_t)
object type is: at::Tensor
auto Bt = torch::mm(x.view(-1, Ddim), projB_w.t()).view(Bsz, L, N);
^
/scratch/abaner/KernelBench/.venv/lib/python3.10/site-packages/torch/include/ATen/core/TensorBody.h(4394): note #3322-D: number of parameters of function "at::Tensor::view(c10::ScalarType) const" does not match the call
inline at::Tensor Tensor::view(at::ScalarType dtype) const {
^
/scratch/abaner/KernelBench/.venv/lib/python3.10/site-packages/torch/include/ATen/core/TensorBody.h(4384): note #3322-D: number of parameters of function "at::Tensor::view(at::IntArrayRef) const" does not match the call
inline at::Tensor Tensor::view(at::IntArrayRef size) const {
^
/scratch/abaner/KernelBench/runs/pop_l123_all_gpt/popcorn/grok-4.3/level_2_problem_30_cache/k_38ca479f8021/mamba_scan/cuda.cu(60): error: no instance of overloaded function "at::Tensor::view" matches the argument list
argument types are: (int64_t, int64_t, int64_t)
object type is: at::Tensor
auto Bt = torch::mm(x.view(-1, Ddim), projB_w.t()).view(Bsz, L, N);
^
/scratch/abaner/KernelBench/.venv/lib/python3.10/site-packages/torch/include/ATen/core/TensorBody.h(4394): note #3322-D: number of parameters of function "at::Tensor::view(c10::ScalarType) const" does not match the call
inline at::Tensor Tensor::view(at::ScalarType dtype) const {
^
/scratch/abaner/KernelBench/.venv/lib/python3.10/site-packages/torch/include/ATen/core/TensorBody.h(4384): note #3322-D: number of parameters of function "at::Tensor::view(at::IntArrayRef) const" does not match the call
inline at::Tensor Tensor::view(at::IntArrayRef size) const {
^
/scratch/abaner/KernelBench/runs/pop_l123_all_gpt/popcorn/grok-4.3/level_2_problem_30_cache/k_38ca479f8021/mamba_scan/cuda.cu(61): error: no instance of overloaded function "at::Tensor::view" matches the argument list
argument types are: (int, int64_t)
object type is: at::Tensor
auto Ct = torch::mm(x.view(-1, Ddim), projC_w.t()).view(Bsz, L, N);
^
/scratch/abaner/KernelBench/.venv/lib/python3.10/site-packages/torch/include/ATen/core/TensorBody.h(4394): note #3322-D: number of parameters of function "at::Tensor::view(c10::ScalarType) const" does not match the call
inline at::Tensor Tensor::view(at::ScalarType dtype) const {
^
/scratch/abaner/KernelBench/.venv/lib/python3.10/site-packages/torch/include/ATen/core/TensorBody.h(4384): note #3322-D: number of parameters of function "at::Tensor::view(at::IntArrayRef) const" does not match the call
inline at::Tensor Tensor::view(at::IntArrayRef size) const {
^
/scratch/abaner/KernelBench/runs/pop_l123_all_gpt/popcorn/grok-4.3/level_2_problem_30_cache/k_38ca479f8021/mamba_scan/cuda.cu(61): error: no instance of overloaded function "at::Tensor::view" matches the argument list
argument types are: (int64_t, int64_t, int64_t)
object type is: at::Tensor
auto Ct = torch::mm(x.view(-1, Ddim), projC_w.t()).view(Bsz, L, N);
^
/scratch/abaner/KernelBench/.venv/lib/python3.10/site-packages/torch/include/ATen/core/TensorBody.h(4394): note #3322-D: number of parameters of function "at::Tensor::view(c10::ScalarType) const" does not match the call
inline at::Tensor Tensor::view(at::ScalarType dtype) const {
^
/scratch/abaner/KernelBench/.venv/lib/python3.10/site-packages/torch/include/ATen/core/TensorBody.h(4384): note #3322-D: number of parameters of function "at::Tensor::view(at::IntArrayRef) const" does not match the call
inline at::Tensor Tensor::view(at::IntArrayRef size) const {
^
/scratch/abaner/KernelBench/runs/pop_l123_all_gpt/popcorn/grok-4.3/level_2_problem_30_cache/k_38ca479f8021/mamba_scan/cuda.cu(71): error: type name is not allowed
delta.data_ptr<float>(),
^
/scratch/abaner/KernelBench/runs/pop_l123_all_gpt/popcorn/grok-4.3/level_2_problem_30_cache/k_38ca479f8021/mamba_scan/cuda.cu(71): error: expected an expression
delta.data_ptr<float>(),
^
/scratch/abaner/KernelBench/runs/pop_l123_all_gpt/popcorn/grok-4.3/level_2_problem_30_cache/k_38ca479f8021/mamba_scan/cuda.cu(72): error: type name is not allowed
Bt.data_ptr<float>(),
^
/scratch/abaner/KernelBench/runs/pop_l123_all_gpt/popcorn/grok-4.3/level_2_problem_30_cache/k_38ca479f8021/mamba_scan/cuda.cu(72): error: expected an expression
Bt.data_ptr<float>(),
^
/scratch/abaner/KernelBench/runs/pop_l123_all_gpt/popcorn/grok-4.3/level_2_problem_30_cache/k_38ca479f8021/mamba_scan/cuda.cu(73): error: type name is not allowed
Ct.data_ptr<float>(),
^
/scratch/abaner/KernelBench/runs/pop_l123_all_gpt/popcorn/grok-4.3/level_2_problem_30_cache/k_38ca479f8021/mamba_scan/cuda.cu(73): error: expected an expression
Ct.data_ptr<float>(),
^
12 errors detected in the compilation of "/scratch/abaner/KernelBench/runs/pop_l123_all_gpt/popcorn/grok-4.3/level_2_problem_30_cache/k_38ca479f8021/mamba_scan/cuda.cu".
ninja: build stopped: subcommand failed.
turn 4167.7s turn · 24.1s LLM · 143.6s tools · tokens in=8,701 · out=839 · total=9,616
reasoning
Errors: view needs IntArrayRef, so use .view({Bsz, L, Ddim})
call → compile_kernel
callcompile_kernel
{
"kernel_code": "import torch\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 mamba_scan_kernel(\n const float* __restrict__ x,\n const float* __restrict__ delta,\n const float* __restrict__ Bt,\n const float* __restrict__ Ct,\n const float* __restrict__ A,\n const float* __restrict__ D_in,\n float* __restrict__ out,\n const int Bsz, const int L, const int Ddim, const int N\n) {\n const int idx = blockIdx.x * blockDim.x + threadIdx.x;\n if (idx >= Bsz * Ddim) return;\n const int b = idx / Ddim;\n const int d = idx % Ddim;\n\n float h[16];\n for(int n = 0; n < N; n++) h[n] = 0.0f;\n\n const float d_val = D_in[d];\n\n for(int t = 0; t < L; t++) {\n const int base = (b * L + t) * Ddim + d;\n const float del = delta[base];\n const float xt = x[base];\n float yt = 0.0f;\n #pragma unroll\n for(int n = 0; n < N; n++) {\n float da = __expf(del * A[d * N + n]);\n float db = del * Bt[(b * L + t) * N + n];\n h[n] = da * h[n] + xt * db;\n yt += h[n] * Ct[(b * L + t) * N + n];\n }\n out[base] = yt + xt * d_val;\n }\n}\n\ntorch::Tensor mamba_forward(\n torch::Tensor x,\n torch::Tensor A_log,\n torch::Tensor projB_w,\n torch::Tensor projC_w,\n torch::Tensor projDelta_w,\n torch::Tensor projDelta_b,\n torch::Tensor D\n) {\n auto Bsz = x.size(0);\n auto L = x.size(1);\n auto Ddim = x.size(2);\n auto N = projB_w.size(1);\n\n auto flat_x = x.view({-1, Ddim});\n auto delta_flat = torch::softplus(torch::addmm(projDelta_b, flat_x, projDelta_w.t()));\n auto delta = delta_flat.view({Bsz, L, Ddim});\n auto Bt = torch::mm(flat_x, projB_w.t()).view({Bsz, L, N});\n auto Ct = torch::mm(flat_x, projC_w.t()).view({Bsz, L, N});\n auto A = -torch::exp(A_log);\n\n auto out = torch::empty_like(x);\n\n const int threads = 256;\n const int blocks = (Bsz * Ddim + threads - 1) / threads;\n\n mamba_scan_kernel<<<blocks, threads>>>(\n x.data_ptr<float>(),\n delta.data_ptr<float>(),\n Bt.data_ptr<float>(),\n Ct.data_ptr<float>(),\n A.data_ptr<float>(),\n D.data_ptr<float>(),\n out.data_ptr<float>(),\n Bsz, L, Ddim, N\n );\n cudaDeviceSynchronize();\n return out;\n}\n'''\n\ncpp_source = r'''\ntorch::Tensor mamba_forward(torch::Tensor x, torch::Tensor A_log, torch::Tensor projB_w, torch::Tensor projC_w, torch::Tensor projDelta_w, torch::Tensor projDelta_b, torch::Tensor D);\n'''\n\nmodule = load_inline(\n name='mamba_scan',\n cpp_sources=cpp_source,\n cuda_sources=cuda_source,\n functions=['mamba_forward'],\n verbose=False,\n extra_cuda_cflags=['-O3', '--use_fast_math']\n)\nprint(\"OK\")\n"
}
result ← compile_kernel (fail)
failcompile_kernel
compile_kernel FAILED: ModelNew class not found or syntax error prevented execution.
OK
turn 5241.0s turn · 11.4s LLM · 229.6s tools · tokens in=9,640 · out=1,030 · total=10,706
reasoning
Compiled OK! Now the module is defined, but I need the full ModelNew class.
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 mamba_scan_kernel(\n const float* __restrict__ x,\n const float* __restrict__ delta,\n const float* __restrict__ Bt,\n const float* __restrict__ Ct,\n const float* __restrict__ A,\n const float* __restrict__ D_in,\n float* __restrict__ out,\n const int Bsz, const int L, const int Ddim, const int N\n) {\n const int idx = blockIdx.x * blockDim.x + threadIdx.x;\n if (idx >= Bsz * Ddim) return;\n const int b = idx / Ddim;\n const int d = idx % Ddim;\n\n float h[16];\n for(int n = 0; n < N; n++) h[n] = 0.0f;\n\n const float d_val = D_in[d];\n\n for(int t = 0; t < L; t++) {\n const int base = (b * L + t) * Ddim + d;\n const float del = delta[base];\n const float xt = x[base];\n float yt = 0.0f;\n #pragma unroll\n for(int n = 0; n < N; n++) {\n float da = __expf(del * A[d * N + n]);\n float db = del * Bt[(b * L + t) * N + n];\n h[n] = da * h[n] + xt * db;\n yt += h[n] * Ct[(b * L + t) * N + n];\n }\n out[base] = yt + xt * d_val;\n }\n}\n\ntorch::Tensor mamba_forward(\n torch::Tensor x,\n torch::Tensor A_log,\n torch::Tensor projB_w,\n torch::Tensor projC_w,\n torch::Tensor projDelta_w,\n torch::Tensor projDelta_b,\n torch::Tensor D\n) {\n auto Bsz = x.size(0);\n auto L = x.size(1);\n auto Ddim = x.size(2);\n auto N = projB_w.size(1);\n\n auto flat_x = x.view({-1, Ddim});\n auto delta_flat = torch::softplus(torch::addmm(projDelta_b, flat_x, projDelta_w.t()));\n auto delta = delta_flat.view({Bsz, L, Ddim});\n auto Bt = torch::mm(flat_x, projB_w.t()).view({Bsz, L, N});\n auto Ct = torch::mm(flat_x, projC_w.t()).view({Bsz, L, N});\n auto A = -torch::exp(A_log);\n\n auto out = torch::empty_like(x);\n\n const int threads = 256;\n const int blocks = (Bsz * Ddim + threads - 1) / threads;\n\n mamba_scan_kernel<<<blocks, threads>>>(\n x.data_ptr<float>(),\n delta.data_ptr<float>(),\n Bt.data_ptr<float>(),\n Ct.data_ptr<float>(),\n A.data_ptr<float>(),\n D.data_ptr<float>(),\n out.data_ptr<float>(),\n Bsz, L, Ddim, N\n );\n cudaDeviceSynchronize();\n return out;\n}\n'''\n\ncpp_source = r'''\ntorch::Tensor mamba_forward(torch::Tensor x, torch::Tensor A_log, torch::Tensor projB_w, torch::Tensor projC_w, torch::Tensor projDelta_w, torch::Tensor projDelta_b, torch::Tensor D);\n'''\n\nselective_scan = load_inline(\n name='mamba_scan',\n cpp_sources=cpp_source,\n cuda_sources=cuda_source,\n functions=['mamba_forward'],\n verbose=False,\n extra_cuda_cflags=['-O3', '--use_fast_math']\n)\n\nclass ModelNew(nn.Module):\n def __init__(self, d_model: int, d_state: int):\n super(ModelNew, self).__init__()\n self.d_model = d_model\n self.d_state = d_state\n self.A_log = nn.Parameter(torch.randn(d_model, d_state))\n self.proj_B = nn.Linear(d_model, d_state, bias=False)\n self.proj_C = nn.Linear(d_model, d_state, bias=False)\n self.proj_delta = nn.Linear(d_model, d_model, bias=True)\n self.D = nn.Parameter(torch.ones(d_model))\n\n def forward(self, x: torch.Tensor) -> torch.Tensor:\n return selective_scan.mamba_forward(\n x,\n self.A_log,\n self.proj_B.weight,\n self.proj_C.weight,\n self.proj_delta.weight,\n self.proj_delta.bias,\n self.D\n )\nprint(\"ModelNew defined\")\n"
}
result ← compile_kernel (ok)
okcompile_kernel
compile_kernel PASSED: kernel compiled without errors.
turn 6152.0s turn · 11.3s LLM · 140.7s tools · tokens in=10,723 · out=1,025 · total=11,754
reasoning
Good, 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 mamba_scan_kernel(\n const float* __restrict__ x,\n const float* __restrict__ delta,\n const float* __restrict__ Bt,\n const float* __restrict__ Ct,\n const float* __restrict__ A,\n const float* __restrict__ D_in,\n float* __restrict__ out,\n const int Bsz, const int L, const int Ddim, const int N\n) {\n const int idx = blockIdx.x * blockDim.x + threadIdx.x;\n if (idx >= Bsz * Ddim) return;\n const int b = idx / Ddim;\n const int d = idx % Ddim;\n\n float h[16];\n for(int n = 0; n < N; n++) h[n] = 0.0f;\n\n const float d_val = D_in[d];\n\n for(int t = 0; t < L; t++) {\n const int base = (b * L + t) * Ddim + d;\n const float del = delta[base];\n const float xt = x[base];\n float yt = 0.0f;\n #pragma unroll\n for(int n = 0; n < N; n++) {\n float da = __expf(del * A[d * N + n]);\n float db = del * Bt[(b * L + t) * N + n];\n h[n] = da * h[n] + xt * db;\n yt += h[n] * Ct[(b * L + t) * N + n];\n }\n out[base] = yt + xt * d_val;\n }\n}\n\ntorch::Tensor mamba_forward(\n torch::Tensor x,\n torch::Tensor A_log,\n torch::Tensor projB_w,\n torch::Tensor projC_w,\n torch::Tensor projDelta_w,\n torch::Tensor projDelta_b,\n torch::Tensor D\n) {\n auto Bsz = x.size(0);\n auto L = x.size(1);\n auto Ddim = x.size(2);\n auto N = projB_w.size(1);\n\n auto flat_x = x.view({-1, Ddim});\n auto delta_flat = torch::softplus(torch::addmm(projDelta_b, flat_x, projDelta_w.t()));\n auto delta = delta_flat.view({Bsz, L, Ddim});\n auto Bt = torch::mm(flat_x, projB_w.t()).view({Bsz, L, N});\n auto Ct = torch::mm(flat_x, projC_w.t()).view({Bsz, L, N});\n auto A = -torch::exp(A_log);\n\n auto out = torch::empty_like(x);\n\n const int threads = 256;\n const int blocks = (Bsz * Ddim + threads - 1) / threads;\n\n mamba_scan_kernel<<<blocks, threads>>>(\n x.data_ptr<float>(),\n delta.data_ptr<float>(),\n Bt.data_ptr<float>(),\n Ct.data_ptr<float>(),\n A.data_ptr<float>(),\n D.data_ptr<float>(),\n out.data_ptr<float>(),\n Bsz, L, Ddim, N\n );\n cudaDeviceSynchronize();\n return out;\n}\n'''\n\ncpp_source = r'''\ntorch::Tensor mamba_forward(torch::Tensor x, torch::Tensor A_log, torch::Tensor projB_w, torch::Tensor projC_w, torch::Tensor projDelta_w, torch::Tensor projDelta_b, torch::Tensor D);\n'''\n\nselective_scan = load_inline(\n name='mamba_scan',\n cpp_sources=cpp_source,\n cuda_sources=cuda_source,\n functions=['mamba_forward'],\n verbose=False,\n extra_cuda_cflags=['-O3', '--use_fast_math']\n)\n\nclass ModelNew(nn.Module):\n def __init__(self, d_model: int, d_state: int):\n super(ModelNew, self).__init__()\n self.d_model = d_model\n self.d_state = d_state\n self.A_log = nn.Parameter(torch.randn(d_model, d_state))\n self.proj_B = nn.Linear(d_model, d_state, bias=False)\n self.proj_C = nn.Linear(d_model, d_state, bias=False)\n self.proj_delta = nn.Linear(d_model, d_model, bias=True)\n self.D = nn.Parameter(torch.ones(d_model))\n\n def forward(self, x: torch.Tensor) -> torch.Tensor:\n return selective_scan.mamba_forward(\n x,\n self.A_log,\n self.proj_B.weight,\n self.proj_C.weight,\n self.proj_delta.weight,\n self.proj_delta.bias,\n self.D\n )\n"
}
result ← run_correctness (fail)
failrun_correctness
run_correctness FAILED: ? trials did not all match.
runtime_error: shape '[16, 2048, 256]' is invalid for input of size 524288
turn 7155.8s turn · 17.0s LLM · 138.8s tools · tokens in=11,796 · out=1,018 · total=13,148
reasoning
Error in view: delta_flat size is B*L*D = 16*2048*256=8M, .view(Bsz,L,Ddim) ok, but error says shape invalid for 524288 which is 16*2048*16?
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 mamba_scan_kernel(\n const float* __restrict__ x,\n const float* __restrict__ delta,\n const float* __restrict__ Bt,\n const float* __restrict__ Ct,\n const float* __restrict__ A,\n const float* __restrict__ D_in,\n float* __restrict__ out,\n const int Bsz, const int L, const int Ddim, const int N\n) {\n const int idx = blockIdx.x * blockDim.x + threadIdx.x;\n if (idx >= Bsz * Ddim) return;\n const int b = idx / Ddim;\n const int d = idx % Ddim;\n\n float h[16];\n for(int n = 0; n < N; n++) h[n] = 0.0f;\n\n const float d_val = D_in[d];\n\n for(int t = 0; t < L; t++) {\n const int base = (b * L + t) * Ddim + d;\n const float del = delta[base];\n const float xt = x[base];\n float yt = 0.0f;\n #pragma unroll\n for(int n = 0; n < N; n++) {\n float da = __expf(del * A[d * N + n]);\n float db = del * Bt[(b * L + t) * N + n];\n h[n] = da * h[n] + xt * db;\n yt += h[n] * Ct[(b * L + t) * N + n];\n }\n out[base] = yt + xt * d_val;\n }\n}\n\ntorch::Tensor mamba_forward(\n torch::Tensor x,\n torch::Tensor A_log,\n torch::Tensor projB_w,\n torch::Tensor projC_w,\n torch::Tensor projDelta_w,\n torch::Tensor projDelta_b,\n torch::Tensor D\n) {\n auto Bsz = x.size(0);\n auto L = x.size(1);\n auto Ddim = x.size(2);\n auto N = projB_w.size(1);\n\n auto flat_x = x.reshape({-1, Ddim});\n auto delta = torch::softplus(torch::addmm(projDelta_b, flat_x, projDelta_w.t())).reshape({Bsz, L, Ddim});\n auto Bt = torch::mm(flat_x, projB_w.t()).reshape({Bsz, L, N});\n auto Ct = torch::mm(flat_x, projC_w.t()).reshape({Bsz, L, N});\n auto A = -torch::exp(A_log);\n\n auto out = torch::empty_like(x);\n\n const int threads = 256;\n const int blocks = (Bsz * Ddim + threads - 1) / threads;\n\n mamba_scan_kernel<<<blocks, threads>>>(\n x.data_ptr<float>(),\n delta.data_ptr<float>(),\n Bt.data_ptr<float>(),\n Ct.data_ptr<float>(),\n A.data_ptr<float>(),\n D.data_ptr<float>(),\n out.data_ptr<float>(),\n Bsz, L, Ddim, N\n );\n cudaDeviceSynchronize();\n return out;\n}\n'''\n\ncpp_source = r'''\ntorch::Tensor mamba_forward(torch::Tensor x, torch::Tensor A_log, torch::Tensor projB_w, torch::Tensor projC_w, torch::Tensor projDelta_w, torch::Tensor projDelta_b, torch::Tensor D);\n'''\n\nselective_scan = load_inline(\n name='mamba_scan',\n cpp_sources=cpp_source,\n cuda_sources=cuda_source,\n functions=['mamba_forward'],\n verbose=False,\n extra_cuda_cflags=['-O3', '--use_fast_math']\n)\n\nclass ModelNew(nn.Module):\n def __init__(self, d_model: int, d_state: int):\n super(ModelNew, self).__init__()\n self.d_model = d_model\n self.d_state = d_state\n self.A_log = nn.Parameter(torch.randn(d_model, d_state))\n self.proj_B = nn.Linear(d_model, d_state, bias=False)\n self.proj_C = nn.Linear(d_model, d_state, bias=False)\n self.proj_delta = nn.Linear(d_model, d_model, bias=True)\n self.D = nn.Parameter(torch.ones(d_model))\n\n def forward(self, x: torch.Tensor) -> torch.Tensor:\n return selective_scan.mamba_forward(\n x,\n self.A_log,\n self.proj_B.weight,\n self.proj_C.weight,\n self.proj_delta.weight,\n self.proj_delta.bias,\n self.D\n )\n"
}
result ← run_correctness (fail)
failrun_correctness
run_correctness FAILED: ? trials did not all match.
runtime_error: shape '[16, 2048, 256]' is invalid for input of size 524288
turn 8164.2s turn · 15.0s LLM · 149.2s tools · tokens in=13,190 · out=889 · total=14,440
reasoning
Still same. delta = softplus( addmm ) is size BL x D = 8M elements, reshape to 16,2048,256 ok.
I noticed the error likely stems from a shape mismatch when reshaping tensors, specifically an attempt to reshape a tensor of size 524288 (which matches 16*2048*16) into [16, 2048, 256], which is much larger at 8 million elements.
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 mamba_scan_kernel(\n const float* __restrict__ x,\n const float* __restrict__ delta,\n const float* __restrict__ Bt,\n const float* __restrict__ Ct,\n const float* __restrict__ A,\n const float* __restrict__ D_in,\n float* __restrict__ out,\n const int Bsz, const int L, const int Ddim, const int N\n) {\n const int idx = blockIdx.x * blockDim.x + threadIdx.x;\n if (idx >= Bsz * Ddim) return;\n const int b = idx / Ddim;\n const int d = idx % Ddim;\n\n float h[16];\n for(int n = 0; n < N; n++) h[n] = 0.0f;\n\n const float d_val = D_in[d];\n\n for(int t = 0; t < L; t++) {\n const int base = (b * L + t) * Ddim + d;\n const float del = delta[base];\n const float xt = x[base];\n float yt = 0.0f;\n #pragma unroll\n for(int n = 0; n < N; n++) {\n float da = __expf(del * A[d * N + n]);\n float db = del * Bt[(b * L + t) * N + n];\n h[n] = da * h[n] + xt * db;\n yt += h[n] * Ct[(b * L + t) * N + n];\n }\n out[base] = yt + xt * d_val;\n }\n}\n\ntorch::Tensor mamba_scan(\n torch::Tensor x,\n torch::Tensor delta,\n torch::Tensor Bt,\n torch::Tensor Ct,\n torch::Tensor A,\n torch::Tensor D\n) {\n auto Bsz = x.size(0);\n auto L = x.size(1);\n auto Ddim = x.size(2);\n auto N = Bt.size(2);\n\n auto out = torch::empty_like(x);\n\n const int threads = 256;\n const int blocks = (Bsz * Ddim + threads - 1) / threads;\n\n mamba_scan_kernel<<<blocks, threads>>>(\n x.data_ptr<float>(),\n delta.data_ptr<float>(),\n Bt.data_ptr<float>(),\n Ct.data_ptr<float>(),\n A.data_ptr<float>(),\n D.data_ptr<float>(),\n out.data_ptr<float>(),\n Bsz, L, Ddim, N\n );\n cudaDeviceSynchronize();\n return out;\n}\n'''\n\ncpp_source = r'''\ntorch::Tensor mamba_scan(torch::Tensor x, torch::Tensor delta, torch::Tensor Bt, torch::Tensor Ct, torch::Tensor A, torch::Tensor D);\n'''\n\nselective_scan = load_inline(\n name='mamba_scan',\n cpp_sources=cpp_source,\n cuda_sources=cuda_source,\n functions=['mamba_scan'],\n verbose=False,\n extra_cuda_cflags=['-O3', '--use_fast_math']\n)\n\nclass ModelNew(nn.Module):\n def __init__(self, d_model: int, d_state: int):\n super(ModelNew, self).__init__()\n self.d_model = d_model\n self.d_state = d_state\n self.A_log = nn.Parameter(torch.randn(d_model, d_state))\n self.proj_B = nn.Linear(d_model, d_state, bias=False)\n self.proj_C = nn.Linear(d_model, d_state, bias=False)\n self.proj_delta = nn.Linear(d_model, d_model, bias=True)\n self.D = nn.Parameter(torch.ones(d_model))\n\n def forward(self, x: torch.Tensor) -> torch.Tensor:\n delta = torch.nn.functional.softplus(self.proj_delta(x))\n Bt = self.proj_B(x)\n Ct = self.proj_C(x)\n A = -torch.exp(self.A_log)\n return selective_scan.mamba_scan(x, delta, Bt, Ct, A, self.D)\n"
}
result ← compile_kernel (ok)
okcompile_kernel
compile_kernel PASSED: kernel compiled without errors.
turn 979.5s turn · 9.0s LLM · 70.5s tools · tokens in=14,457 · out=890 · total=15,351
reasoning
Now 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\ncuda_source = r'''\n#include <torch/extension.h>\n#include <cuda_runtime.h>\n\n__global__ void mamba_scan_kernel(\n const float* __restrict__ x,\n const float* __restrict__ delta,\n const float* __restrict__ Bt,\n const float* __restrict__ Ct,\n const float* __restrict__ A,\n const float* __restrict__ D_in,\n float* __restrict__ out,\n const int Bsz, const int L, const int Ddim, const int N\n) {\n const int idx = blockIdx.x * blockDim.x + threadIdx.x;\n if (idx >= Bsz * Ddim) return;\n const int b = idx / Ddim;\n const int d = idx % Ddim;\n\n float h[16];\n for(int n = 0; n < N; n++) h[n] = 0.0f;\n\n const float d_val = D_in[d];\n\n for(int t = 0; t < L; t++) {\n const int base = (b * L + t) * Ddim + d;\n const float del = delta[base];\n const float xt = x[base];\n float yt = 0.0f;\n #pragma unroll\n for(int n = 0; n < N; n++) {\n float da = __expf(del * A[d * N + n]);\n float db = del * Bt[(b * L + t) * N + n];\n h[n] = da * h[n] + xt * db;\n yt += h[n] * Ct[(b * L + t) * N + n];\n }\n out[base] = yt + xt * d_val;\n }\n}\n\ntorch::Tensor mamba_scan(\n torch::Tensor x,\n torch::Tensor delta,\n torch::Tensor Bt,\n torch::Tensor Ct,\n torch::Tensor A,\n torch::Tensor D\n) {\n auto Bsz = x.size(0);\n auto L = x.size(1);\n auto Ddim = x.size(2);\n auto N = Bt.size(2);\n\n auto out = torch::empty_like(x);\n\n const int threads = 256;\n const int blocks = (Bsz * Ddim + threads - 1) / threads;\n\n mamba_scan_kernel<<<blocks, threads>>>(\n x.data_ptr<float>(),\n delta.data_ptr<float>(),\n Bt.data_ptr<float>(),\n Ct.data_ptr<float>(),\n A.data_ptr<float>(),\n D.data_ptr<float>(),\n out.data_ptr<float>(),\n Bsz, L, Ddim, N\n );\n cudaDeviceSynchronize();\n return out;\n}\n'''\n\ncpp_source = r'''\ntorch::Tensor mamba_scan(torch::Tensor x, torch::Tensor delta, torch::Tensor Bt, torch::Tensor Ct, torch::Tensor A, torch::Tensor D);\n'''\n\nselective_scan = load_inline(\n name='mamba_scan',\n cpp_sources=cpp_source,\n cuda_sources=cuda_source,\n functions=['mamba_scan'],\n verbose=False,\n extra_cuda_cflags=['-O3', '--use_fast_math']\n)\n\nclass ModelNew(nn.Module):\n def __init__(self, d_model: int, d_state: int):\n super(ModelNew, self).__init__()\n self.d_model = d_model\n self.d_state = d_state\n self.A_log = nn.Parameter(torch.randn(d_model, d_state))\n self.proj_B = nn.Linear(d_model, d_state, bias=False)\n self.proj_C = nn.Linear(d_model, d_state, bias=False)\n self.proj_delta = nn.Linear(d_model, d_model, bias=True)\n self.D = nn.Parameter(torch.ones(d_model))\n\n def forward(self, x: torch.Tensor) -> torch.Tensor:\n delta = torch.nn.functional.softplus(self.proj_delta(x))\n Bt = self.proj_B(x)\n Ct = self.proj_C(x)\n A = -torch.exp(self.A_log)\n return selective_scan.mamba_scan(x, delta, Bt, Ct, A, self.D)\n"
}
result ← run_correctness (ok)
okrun_correctness
run_correctness PASSED: (3 / 3) trials all matched the reference.
Numerical precision: max_abs_err=3.05e-05 mean_abs_err=1.63e-07 max_rel_err=5.56e-01
turn 10110.6s turn · 8.9s LLM · 101.7s tools · tokens in=15,411 · out=889 · total=16,305 · FINAL
reasoning
Correct! Now submit.
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 mamba_scan_kernel(\n const float* __restrict__ x,\n const float* __restrict__ delta,\n const float* __restrict__ Bt,\n const float* __restrict__ Ct,\n const float* __restrict__ A,\n const float* __restrict__ D_in,\n float* __restrict__ out,\n const int Bsz, const int L, const int Ddim, const int N\n) {\n const int idx = blockIdx.x * blockDim.x + threadIdx.x;\n if (idx >= Bsz * Ddim) return;\n const int b = idx / Ddim;\n const int d = idx % Ddim;\n\n float h[16];\n for(int n = 0; n < N; n++) h[n] = 0.0f;\n\n const float d_val = D_in[d];\n\n for(int t = 0; t < L; t++) {\n const int base = (b * L + t) * Ddim + d;\n const float del = delta[base];\n const float xt = x[base];\n float yt = 0.0f;\n #pragma unroll\n for(int n = 0; n < N; n++) {\n float da = __expf(del * A[d * N + n]);\n float db = del * Bt[(b * L + t) * N + n];\n h[n] = da * h[n] + xt * db;\n yt += h[n] * Ct[(b * L + t) * N + n];\n }\n out[base] = yt + xt * d_val;\n }\n}\n\ntorch::Tensor mamba_scan(\n torch::Tensor x,\n torch::Tensor delta,\n torch::Tensor Bt,\n torch::Tensor Ct,\n torch::Tensor A,\n torch::Tensor D\n) {\n auto Bsz = x.size(0);\n auto L = x.size(1);\n auto Ddim = x.size(2);\n auto N = Bt.size(2);\n\n auto out = torch::empty_like(x);\n\n const int threads = 256;\n const int blocks = (Bsz * Ddim + threads - 1) / threads;\n\n mamba_scan_kernel<<<blocks, threads>>>(\n x.data_ptr<float>(),\n delta.data_ptr<float>(),\n Bt.data_ptr<float>(),\n Ct.data_ptr<float>(),\n A.data_ptr<float>(),\n D.data_ptr<float>(),\n out.data_ptr<float>(),\n Bsz, L, Ddim, N\n );\n cudaDeviceSynchronize();\n return out;\n}\n'''\n\ncpp_source = r'''\ntorch::Tensor mamba_scan(torch::Tensor x, torch::Tensor delta, torch::Tensor Bt, torch::Tensor Ct, torch::Tensor A, torch::Tensor D);\n'''\n\nselective_scan = load_inline(\n name='mamba_scan',\n cpp_sources=cpp_source,\n cuda_sources=cuda_source,\n functions=['mamba_scan'],\n verbose=False,\n extra_cuda_cflags=['-O3', '--use_fast_math']\n)\n\nclass ModelNew(nn.Module):\n def __init__(self, d_model: int, d_state: int):\n super(ModelNew, self).__init__()\n self.d_model = d_model\n self.d_state = d_state\n self.A_log = nn.Parameter(torch.randn(d_model, d_state))\n self.proj_B = nn.Linear(d_model, d_state, bias=False)\n self.proj_C = nn.Linear(d_model, d_state, bias=False)\n self.proj_delta = nn.Linear(d_model, d_model, bias=True)\n self.D = nn.Parameter(torch.ones(d_model))\n\n def forward(self, x: torch.Tensor) -> torch.Tensor:\n delta = torch.nn.functional.softplus(self.proj_delta(x))\n Bt = self.proj_B(x)\n Ct = self.proj_C(x)\n A = -torch.exp(self.A_log)\n return selective_scan.mamba_scan(x, delta, Bt, Ct, A, self.D)\n"
}
result ← submit_kernel (ok)
oksubmit_kernel
submit_kernel PASSED: (3 / 3) correctness trials all passed.
Kernel runtime: 3.94 μs
Runtime stats: mean=3.94μs median=?μs std=0.05μs
Numerical precision: max_abs_err=3.05e-05 mean_abs_err=1.63e-07 max_rel_err=5.56e-01
Memory: 133.61 MB peak (ref: 1765.86 MB, ratio: 0.0757x)
Kernel launches: 30 (ref: 26681, fusion ratio: 889.3667)
Energy: 482.68 mJ/run (ref: 10380.38 mJ, ratio: 21.5057x)
SOL score: 1.0
Flagged for excessive speedup — this submission may have been rejected by automated review.
final submitted kernel (108 lines, 3,108 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 mamba_scan_kernel(
const float* __restrict__ x,
const float* __restrict__ delta,
const float* __restrict__ Bt,
const float* __restrict__ Ct,
const float* __restrict__ A,
const float* __restrict__ D_in,
float* __restrict__ out,
const int Bsz, const int L, const int Ddim, const int N
) {
const int idx = blockIdx.x * blockDim.x + threadIdx.x;
if (idx >= Bsz * Ddim) return;
const int b = idx / Ddim;
const int d = idx % Ddim;
float h[16];
for(int n = 0; n < N; n++) h[n] = 0.0f;
const float d_val = D_in[d];
for(int t = 0; t < L; t++) {
const int base = (b * L + t) * Ddim + d;
const float del = delta[base];
const float xt = x[base];
float yt = 0.0f;
#pragma unroll
for(int n = 0; n < N; n++) {
float da = __expf(del * A[d * N + n]);
float db = del * Bt[(b * L + t) * N + n];
h[n] = da * h[n] + xt * db;
yt += h[n] * Ct[(b * L + t) * N + n];
}
out[base] = yt + xt * d_val;
}
}
torch::Tensor mamba_scan(
torch::Tensor x,
torch::Tensor delta,
torch::Tensor Bt,
torch::Tensor Ct,
torch::Tensor A,
torch::Tensor D
) {
auto Bsz = x.size(0);
auto L = x.size(1);
auto Ddim = x.size(2);
auto N = Bt.size(2);
auto out = torch::empty_like(x);
const int threads = 256;
const int blocks = (Bsz * Ddim + threads - 1) / threads;
mamba_scan_kernel<<<blocks, threads>>>(
x.data_ptr<float>(),
delta.data_ptr<float>(),
Bt.data_ptr<float>(),
Ct.data_ptr<float>(),
A.data_ptr<float>(),
D.data_ptr<float>(),
out.data_ptr<float>(),
Bsz, L, Ddim, N
);
cudaDeviceSynchronize();
return out;
}
'''
cpp_source = r'''
torch::Tensor mamba_scan(torch::Tensor x, torch::Tensor delta, torch::Tensor Bt, torch::Tensor Ct, torch::Tensor A, torch::Tensor D);
'''
selective_scan = load_inline(
name='mamba_scan',
cpp_sources=cpp_source,
cuda_sources=cuda_source,
functions=['mamba_scan'],
verbose=False,
extra_cuda_cflags=['-O3', '--use_fast_math']
)
class ModelNew(nn.Module):
def __init__(self, d_model: int, d_state: int):
super(ModelNew, self).__init__()
self.d_model = d_model
self.d_state = d_state
self.A_log = nn.Parameter(torch.randn(d_model, d_state))
self.proj_B = nn.Linear(d_model, d_state, bias=False)
self.proj_C = nn.Linear(d_model, d_state, bias=False)
self.proj_delta = nn.Linear(d_model, d_model, bias=True)
self.D = nn.Parameter(torch.ones(d_model))
def forward(self, x: torch.Tensor) -> torch.Tensor:
delta = torch.nn.functional.softplus(self.proj_delta(x))
Bt = self.proj_B(x)
Ct = self.proj_C(x)
A = -torch.exp(self.A_log)
return selective_scan.mamba_scan(x, delta, Bt, Ct, A, self.D)