incorrect
12_DeepSeekMoEDispatchPermute.py
started
2026-05-11T09:20:34Z
correctness issue:
timingkernelbench.timing
kernel mean / median / std
—
ref mean / median / std
88.50 / — / 3.17 μs
SOL (speed-of-light)nsight_hardware_counters
roofline / occupancynsight_hardware_counters
warp stalls (top 5)nsight
reference PyTorch source (the prompt the agent sees)
import torch
import torch.nn as nn
class Model(nn.Module):
"""
Pack tokens into expert-major order for MoE execution.
Each token is assigned to one expert and one slot within that expert's
buffer. The kernel writes token activations into the packed dispatch layout.
"""
def __init__(self):
super().__init__()
def forward(
self,
token_hidden: torch.Tensor,
expert_idx: torch.Tensor,
slot_idx: torch.Tensor,
expert_offsets: torch.Tensor,
) -> torch.Tensor:
num_rows = int(expert_offsets[-1].item())
out = torch.zeros(num_rows, token_hidden.shape[1], dtype=token_hidden.dtype, device=token_hidden.device)
for token in range(token_hidden.shape[0]):
expert = int(expert_idx[token].item())
row = int(expert_offsets[expert].item() + slot_idx[token].item())
out[row] = token_hidden[token]
return out
num_tokens = 2048
hidden_dim = 128
num_experts = 16
def get_inputs():
return popcorn_pri.moe_dispatch_inputs(num_tokens, hidden_dim, num_experts)
def get_init_inputs():
return []
turn 0239.6s turn · 19.0s LLM · 220.6s tools · tokens in=1,286 · out=577 · total=2,820 · FINAL
reasoning
The task is to replace the forward computation of a PyTorch model with a custom CUDA kernel. The model packs tokens into expert-major order for MoE execution.
I need to create a custom CUDA kernel to speed up the forward computation of a PyTorch model for an MoE setup, replacing the reference loop with something faster. The idea is to scatter tokens' hidden states into an output tensor based on their expert indices and offsets, arranging them in expert-major order.
For the CUDA kernel itself, I'll launch one thread per token, calculate the target row as expert_offsets[expert] + slot_idx[token], and then copy the hidden vector (size 128) element-wise into the output at that position. Since hidden_dim is small, this per-thread loop is efficient enough.
In the Python part, I'll use torch.utils.cpp_extension.load_inline to load the CUDA kernel inline. I'll need to define the CUDA kernel launch with grids and blocks for parallelization, passing pointers to the tensors' data, along with dimensions like num_tokens and hidden_dim.
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 pack_tokens_kernel(\n const float* __restrict__ token_hidden,\n const int* __restrict__ expert_idx,\n const int* __restrict__ slot_idx,\n const int* __restrict__ expert_offsets,\n float* __restrict__ out,\n int num_tokens,\n int hidden_dim\n) {\n int token = blockIdx.x * blockDim.x + threadIdx.x;\n if (token < num_tokens) {\n int expert = expert_idx[token];\n int slot = slot_idx[token];\n int row = expert_offsets[expert] + slot;\n const float* src = token_hidden + token * hidden_dim;\n float* dst = out + row * hidden_dim;\n for (int d = 0; d < hidden_dim; ++d) {\n dst[d] = src[d];\n }\n }\n}\n\ntorch::Tensor pack_tokens_cuda(\n torch::Tensor token_hidden,\n torch::Tensor expert_idx,\n torch::Tensor slot_idx,\n torch::Tensor expert_offsets\n) {\n int num_tokens = token_hidden.size(0);\n int hidden_dim = token_hidden.size(1);\n int num_experts_p1 = expert_offsets.size(0);\n auto opts = token_hidden.options();\n int num_rows = expert_offsets[num_experts_p1 - 1].item<int>();\n auto out = torch::zeros({num_rows, hidden_dim}, opts);\n\n const int threads = 256;\n const int blocks = (num_tokens + threads - 1) / threads;\n\n pack_tokens_kernel<<<blocks, threads>>>(\n token_hidden.data_ptr<float>(),\n expert_idx.data_ptr<int>(),\n slot_idx.data_ptr<int>(),\n expert_offsets.data_ptr<int>(),\n out.data_ptr<float>(),\n num_tokens,\n hidden_dim\n );\n return out;\n}\n'''\n\ncpp_source = r'''\ntorch::Tensor pack_tokens_cuda(torch::Tensor token_hidden, torch::Tensor expert_idx, torch::Tensor slot_idx, torch::Tensor expert_offsets);\n'''\n\n_mod = load_inline(\n name='pack_tokens_ext',\n cpp_sources=cpp_source,\n cuda_sources=cuda_source,\n functions=['pack_tokens_cuda'],\n verbose=False\n)\n\nclass ModelNew(nn.Module):\n def __init__(self):\n super().__init__()\n\n def forward(\n self,\n token_hidden: torch.Tensor,\n expert_idx: torch.Tensor,\n slot_idx: torch.Tensor,\n expert_offsets: torch.Tensor,\n ) -> torch.Tensor:\n return _mod.pack_tokens_cuda(token_hidden, expert_idx, slot_idx, expert_offsets)\n"
}
result ← submit_kernel (fail)
failsubmit_kernel
submit_kernel FAILED: correctness check did not pass (? trials).
runtime_error: expected scalar type Int but found Float
final submitted kernel (84 lines, 2,382 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 pack_tokens_kernel(
const float* __restrict__ token_hidden,
const int* __restrict__ expert_idx,
const int* __restrict__ slot_idx,
const int* __restrict__ expert_offsets,
float* __restrict__ out,
int num_tokens,
int hidden_dim
) {
int token = blockIdx.x * blockDim.x + threadIdx.x;
if (token < num_tokens) {
int expert = expert_idx[token];
int slot = slot_idx[token];
int row = expert_offsets[expert] + slot;
const float* src = token_hidden + token * hidden_dim;
float* dst = out + row * hidden_dim;
for (int d = 0; d < hidden_dim; ++d) {
dst[d] = src[d];
}
}
}
torch::Tensor pack_tokens_cuda(
torch::Tensor token_hidden,
torch::Tensor expert_idx,
torch::Tensor slot_idx,
torch::Tensor expert_offsets
) {
int num_tokens = token_hidden.size(0);
int hidden_dim = token_hidden.size(1);
int num_experts_p1 = expert_offsets.size(0);
auto opts = token_hidden.options();
int num_rows = expert_offsets[num_experts_p1 - 1].item<int>();
auto out = torch::zeros({num_rows, hidden_dim}, opts);
const int threads = 256;
const int blocks = (num_tokens + threads - 1) / threads;
pack_tokens_kernel<<<blocks, threads>>>(
token_hidden.data_ptr<float>(),
expert_idx.data_ptr<int>(),
slot_idx.data_ptr<int>(),
expert_offsets.data_ptr<int>(),
out.data_ptr<float>(),
num_tokens,
hidden_dim
);
return out;
}
'''
cpp_source = r'''
torch::Tensor pack_tokens_cuda(torch::Tensor token_hidden, torch::Tensor expert_idx, torch::Tensor slot_idx, torch::Tensor expert_offsets);
'''
_mod = load_inline(
name='pack_tokens_ext',
cpp_sources=cpp_source,
cuda_sources=cuda_source,
functions=['pack_tokens_cuda'],
verbose=False
)
class ModelNew(nn.Module):
def __init__(self):
super().__init__()
def forward(
self,
token_hidden: torch.Tensor,
expert_idx: torch.Tensor,
slot_idx: torch.Tensor,
expert_offsets: torch.Tensor,
) -> torch.Tensor:
return _mod.pack_tokens_cuda(token_hidden, expert_idx, slot_idx, expert_offsets)