correct
12_DeepSeekMoEDispatchPermute.py
started
2026-05-11T19:13:02Z
timingkernelbench.timing
kernel mean / median / std
0.13 / — / 0.07 μs
ref mean / median / std
82.30 / — / 0.44 μs
numerical precisioneval_kernel_against_ref
gpu memoryextended_metrics.measure_memory
kernel launches / fusiontorch.profiler
fusion ratio (ref / kernel)
1741.20x
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 0116.0s turn · 24.1s LLM · 91.8s tools · tokens in=2,184 · out=549 · total=3,580
reasoning
The task is to replace the forward computation of this PyTorch model with a custom CUDA kernel. The model packs tokens into expert-major order for MoE execution.
The task involves replacing the slow Python loop in the PyTorch model's forward pass with a custom CUDA kernel for faster execution. The operation packs tokens into an expert-major order layout for a Mixture of Experts (MoE) model. The inputs include token_hidden, a tensor of shape [num_tokens, hidden_dim], expert_idx [num_tokens] for which expert each token assigns to, slot_idx [num_tokens] for the position within the expert, and expert_offsets [num_experts + 1] defining cumulative offsets.
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\n__global__ void pack_tokens_kernel(\n const float* token_hidden,\n const int* expert_idx,\n const int* slot_idx,\n const int* expert_offsets,\n float* out,\n int num_tokens,\n int hidden_dim\n) {\n int token = blockIdx.x * blockDim.x + threadIdx.x;\n if (token >= num_tokens) return;\n int expert = expert_idx[token];\n int slot = slot_idx[token];\n int row = expert_offsets[expert] + slot;\n for (int d = 0; d < hidden_dim; d++) {\n out[row * hidden_dim + d] = token_hidden[token * hidden_dim + d];\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 auto options = token_hidden.options();\n int num_rows = expert_offsets[expert_offsets.size(0) - 1].item<int>();\n auto out = torch::zeros({num_rows, hidden_dim}, options);\n const int block_size = 256;\n const int num_blocks = (num_tokens + block_size - 1) / block_size;\n pack_tokens_kernel<<<num_blocks, block_size>>>(\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 = '''\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'''\n\npack_module = load_inline(\n name='pack_tokens',\n cpp_sources=[cpp_source],\n cuda_sources=[cuda_source],\n functions=['pack_tokens_cuda'],\n verbose=True,\n with_cuda=True\n)\n\nclass ModelNew(nn.Module):\n def __init__(self):\n super().__init__()\n self.pack = pack_module\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 self.pack.pack_tokens_cuda(token_hidden, expert_idx, slot_idx, expert_offsets)\n"
}
result ← compile_kernel (ok)
okcompile_kernel
compile_kernel PASSED: kernel compiled without errors.
turn 189.8s turn · 8.6s LLM · 81.2s tools · tokens in=3,597 · out=550 · total=4,154
reasoning
Compiled successfully. Now check correctness.
call → run_correctness
callrun_correctness
{
"kernel_code": "import torch\nimport torch.nn as nn\nfrom torch.utils.cpp_extension import load_inline\n\ncuda_source = '''\n#include <torch/extension.h>\n\n__global__ void pack_tokens_kernel(\n const float* token_hidden,\n const int* expert_idx,\n const int* slot_idx,\n const int* expert_offsets,\n float* out,\n int num_tokens,\n int hidden_dim\n) {\n int token = blockIdx.x * blockDim.x + threadIdx.x;\n if (token >= num_tokens) return;\n int expert = expert_idx[token];\n int slot = slot_idx[token];\n int row = expert_offsets[expert] + slot;\n for (int d = 0; d < hidden_dim; d++) {\n out[row * hidden_dim + d] = token_hidden[token * hidden_dim + d];\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 auto options = token_hidden.options();\n int num_rows = expert_offsets[expert_offsets.size(0) - 1].item<int>();\n auto out = torch::zeros({num_rows, hidden_dim}, options);\n const int block_size = 256;\n const int num_blocks = (num_tokens + block_size - 1) / block_size;\n pack_tokens_kernel<<<num_blocks, block_size>>>(\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 = '''\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'''\n\npack_module = load_inline(\n name='pack_tokens',\n cpp_sources=[cpp_source],\n cuda_sources=[cuda_source],\n functions=['pack_tokens_cuda'],\n verbose=True,\n with_cuda=True\n)\n\nclass ModelNew(nn.Module):\n def __init__(self):\n super().__init__()\n self.pack = pack_module\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 self.pack.pack_tokens_cuda(token_hidden, expert_idx, slot_idx, expert_offsets)\n"
}
result ← run_correctness (fail)
failrun_correctness
run_correctness FAILED: ? trials did not all match.
runtime_error: expected scalar type Int but found Float
turn 293.8s turn · 12.5s LLM · 81.3s tools · tokens in=4,183 · out=561 · total=4,846
reasoning
Error: expected scalar type Int but found Float. Probably because expert_offsets[-1] is float? No.
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 = '''\n#include <torch/extension.h>\n\n__global__ void pack_tokens_kernel(\n const float* token_hidden,\n const int* expert_idx,\n const int* slot_idx,\n const int* expert_offsets,\n float* out,\n int num_tokens,\n int hidden_dim\n) {\n int token = blockIdx.x * blockDim.x + threadIdx.x;\n if (token >= num_tokens) return;\n int expert = expert_idx[token];\n int slot = slot_idx[token];\n int row = expert_offsets[expert] + slot;\n for (int d = 0; d < hidden_dim; d++) {\n out[row * hidden_dim + d] = token_hidden[token * hidden_dim + d];\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 auto options = token_hidden.options();\n torch::Tensor cpu_offsets = expert_offsets.cpu();\n int num_rows = cpu_offsets[cpu_offsets.size(0) - 1].item<int>();\n auto out = torch::zeros({num_rows, hidden_dim}, options);\n const int block_size = 256;\n const int num_blocks = (num_tokens + block_size - 1) / block_size;\n pack_tokens_kernel<<<num_blocks, block_size>>>(\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 = '''\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'''\n\npack_module = load_inline(\n name='pack_tokens',\n cpp_sources=[cpp_source],\n cuda_sources=[cuda_source],\n functions=['pack_tokens_cuda'],\n verbose=True,\n with_cuda=True\n)\n\nclass ModelNew(nn.Module):\n def __init__(self):\n super().__init__()\n self.pack = pack_module\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 self.pack.pack_tokens_cuda(token_hidden, expert_idx, slot_idx, expert_offsets)\n"
}
result ← run_correctness (fail)
failrun_correctness
run_correctness FAILED: ? trials did not all match.
runtime_error: expected scalar type Int but found Float
turn 3167.6s turn · 14.9s LLM · 152.7s tools · tokens in=4,875 · out=564 · total=5,715
reasoning
Still error. The error "expected scalar type Int but found Float" – perhaps one of the tensors is float.
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 = '''\n#include <torch/extension.h>\n\n__global__ void pack_tokens_kernel(\n const float* token_hidden,\n const int* expert_idx,\n const int* slot_idx,\n const int* expert_offsets,\n float* out,\n int num_tokens,\n int hidden_dim\n) {\n int token = blockIdx.x * blockDim.x + threadIdx.x;\n if (token >= num_tokens) return;\n int expert = expert_idx[token];\n int slot = slot_idx[token];\n int row = expert_offsets[expert] + slot;\n for (int d = 0; d < hidden_dim; d++) {\n out[row * hidden_dim + d] = token_hidden[token * hidden_dim + d];\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 auto options = token_hidden.options();\n torch::Tensor offsets_dev = expert_offsets.to(token_hidden.device());\n int num_rows = offsets_dev[offsets_dev.size(0) - 1].item<int>();\n auto out = torch::zeros({num_rows, hidden_dim}, options);\n const int block_size = 256;\n const int num_blocks = (num_tokens + block_size - 1) / block_size;\n pack_tokens_kernel<<<num_blocks, block_size>>>(\n token_hidden.data_ptr<float>(),\n expert_idx.data_ptr<int>(),\n slot_idx.data_ptr<int>(),\n offsets_dev.data_ptr<int>(),\n out.data_ptr<float>(),\n num_tokens,\n hidden_dim\n );\n return out;\n}\n'''\n\ncpp_source = '''\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'''\n\npack_module = load_inline(\n name='pack_tokens',\n cpp_sources=[cpp_source],\n cuda_sources=[cuda_source],\n functions=['pack_tokens_cuda'],\n verbose=True,\n with_cuda=True\n)\n\nclass ModelNew(nn.Module):\n def __init__(self):\n super().__init__()\n self.pack = pack_module\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 self.pack.pack_tokens_cuda(token_hidden, expert_idx, slot_idx, expert_offsets)\n"
}
result ← run_correctness (fail)
failrun_correctness
run_correctness FAILED: ? trials did not all match.
runtime_error: expected scalar type Int but found Float
turn 4184.9s turn · 16.2s LLM · 168.7s tools · tokens in=5,744 · out=534 · total=6,662
reasoning
Still same error. The .item<int>() on offsets_dev which is int, but perhaps num_rows line, if offsets_dev[-1] after indexing is tensor.
I was troubleshooting an error in the CUDA kernel implementation for the PyTorch model. The error indicated that somewhere, a float tensor was being used instead of an int tensor. I suspected the issue might be with .item() calls on offsets_dev[offsets_dev.
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 = '''\n#include <torch/extension.h>\n\n__global__ void pack_tokens_kernel(\n const float* token_hidden,\n const int* expert_idx,\n const int* slot_idx,\n const int* expert_offsets,\n float* out,\n int num_tokens,\n int hidden_dim\n) {\n int token = blockIdx.x * blockDim.x + threadIdx.x;\n if (token >= num_tokens) return;\n int expert = expert_idx[token];\n int slot = slot_idx[token];\n int row = expert_offsets[expert] + slot;\n for (int d = 0; d < hidden_dim; d++) {\n out[row * hidden_dim + d] = token_hidden[token * hidden_dim + d];\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 auto options = token_hidden.options();\n auto out = torch::zeros({expert_offsets[-1].item<int>(), hidden_dim}, options);\n const int block_size = 256;\n const int num_blocks = (num_tokens + block_size - 1) / block_size;\n pack_tokens_kernel<<<num_blocks, block_size>>>(\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 = '''\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'''\n\npack_module = load_inline(\n name='pack_tokens',\n cpp_sources=[cpp_source],\n cuda_sources=[cuda_source],\n functions=['pack_tokens_cuda'],\n verbose=True,\n with_cuda=True\n)\n\nclass ModelNew(nn.Module):\n def __init__(self):\n super().__init__()\n self.pack = pack_module\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 self.pack.pack_tokens_cuda(token_hidden, expert_idx, slot_idx, expert_offsets)\n"
}
result ← run_correctness (fail)
failrun_correctness
run_correctness FAILED: ? trials did not all match.
runtime_error: expected scalar type Int but found Float
turn 596.2s turn · 10.6s LLM · 85.6s tools · tokens in=6,691 · out=544 · total=7,376
reasoning
Still. The problem is expert_offsets[-1] in C++, when expert_offsets is int32 CUDA, [-1] indexing gives tensor, .item<int>() – the error suggests it thinks it's float.
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 = '''\n#include <torch/extension.h>\n\n__global__ void pack_tokens_kernel(\n const float* token_hidden,\n const int* expert_idx,\n const int* slot_idx,\n const int* expert_offsets,\n float* out,\n int num_tokens,\n int hidden_dim\n) {\n int token = blockIdx.x * blockDim.x + threadIdx.x;\n if (token >= num_tokens) return;\n int expert = expert_idx[token];\n int slot = slot_idx[token];\n int row = expert_offsets[expert] + slot;\n for (int d = 0; d < hidden_dim; d++) {\n out[row * hidden_dim + d] = token_hidden[token * hidden_dim + d];\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 auto options = token_hidden.options();\n int num_rows = expert_offsets[-1].item<int64_t>();\n auto out = torch::zeros({num_rows, hidden_dim}, options);\n const int block_size = 256;\n const int num_blocks = (num_tokens + block_size - 1) / block_size;\n pack_tokens_kernel<<<num_blocks, block_size>>>(\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 = '''\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'''\n\npack_module = load_inline(\n name='pack_tokens',\n cpp_sources=[cpp_source],\n cuda_sources=[cuda_source],\n functions=['pack_tokens_cuda'],\n verbose=True,\n with_cuda=True\n)\n\nclass ModelNew(nn.Module):\n def __init__(self):\n super().__init__()\n self.pack = pack_module\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 self.pack.pack_tokens_cuda(token_hidden, expert_idx, slot_idx, expert_offsets)\n"
}
result ← run_correctness (fail)
failrun_correctness
run_correctness FAILED: ? trials did not all match.
runtime_error: expected scalar type Int but found Float
turn 695.6s turn · 19.8s LLM · 75.7s tools · tokens in=7,405 · out=553 · total=8,548
reasoning
Still. The error must be that expert_offsets.dtype is float somehow? But in reference it's int32.
I noticed the error message about expected scalar type Int but found Float, which points to a dtype mismatch somewhere. It seems like expert_offsets might be float instead of int32 as shown in the reference. The inputs to forward are created on CPU without specifying device, but evaluation likely moves them to CUDA. There could be an Num_rows is computed in C++ using expert_offsets[-1], and if the tensor is CPU int32, it should work, but the error suggests float involvement. To avoid issues, I can compute num_rows in Python using num_rows = int(expert_offsets[-1].
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 = '''\n#include <torch/extension.h>\n\n__global__ void pack_tokens_kernel(\n const float* token_hidden,\n const int* expert_idx,\n const int* slot_idx,\n const int* expert_offsets,\n float* out,\n int num_tokens,\n int hidden_dim\n) {\n int token = blockIdx.x * blockDim.x + threadIdx.x;\n if (token >= num_tokens) return;\n int expert = expert_idx[token];\n int slot = slot_idx[token];\n int row = expert_offsets[expert] + slot;\n for (int d = 0; d < hidden_dim; d++) {\n out[row * hidden_dim + d] = token_hidden[token * hidden_dim + d];\n }\n}\n\nvoid pack_tokens_cuda(\n torch::Tensor token_hidden,\n torch::Tensor expert_idx,\n torch::Tensor slot_idx,\n torch::Tensor expert_offsets,\n torch::Tensor out\n) {\n int num_tokens = token_hidden.size(0);\n int hidden_dim = token_hidden.size(1);\n const int block_size = 256;\n const int num_blocks = (num_tokens + block_size - 1) / block_size;\n pack_tokens_kernel<<<num_blocks, block_size>>>(\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}\n'''\n\ncpp_source = '''\nvoid pack_tokens_cuda(\n torch::Tensor token_hidden,\n torch::Tensor expert_idx,\n torch::Tensor slot_idx,\n torch::Tensor expert_offsets,\n torch::Tensor out\n);\n'''\n\npack_module = load_inline(\n name='pack_tokens',\n cpp_sources=[cpp_source],\n cuda_sources=[cuda_source],\n functions=['pack_tokens_cuda'],\n verbose=True,\n with_cuda=True\n)\n\nclass ModelNew(nn.Module):\n def __init__(self):\n super().__init__()\n self.pack = pack_module\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 num_rows = int(expert_offsets[-1].item())\n out = torch.zeros(num_rows, token_hidden.shape[1], dtype=token_hidden.dtype, device=token_hidden.device)\n self.pack.pack_tokens_cuda(token_hidden, expert_idx, slot_idx, expert_offsets, out)\n return out\n"
}
result ← run_correctness (fail)
failrun_correctness
run_correctness FAILED: ? trials did not all match.
runtime_error: expected scalar type Int but found Float
turn 7134.8s turn · 16.7s LLM · 118.1s tools · tokens in=8,577 · out=576 · total=9,470
reasoning
Still same error. The error persists, so probably the dtype of one of expert_idx etc is float in test? But unlikely.
The thinking revolves around troubleshooting a PyTorch error: "expected scalar type Int but found Float". It seems the error occurs because somewhere in the code, a function expects an integer tensor but receives a float one. I suspect this happens when accessing .item() on tensors that might be float, like expert_offsets[-1].
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 = '''\n#include <torch/extension.h>\n\n__global__ void pack_tokens_kernel(\n const float* token_hidden,\n const int* expert_idx,\n const int* slot_idx,\n const int* expert_offsets,\n float* out,\n int num_tokens,\n int hidden_dim\n) {\n int token = blockIdx.x * blockDim.x + threadIdx.x;\n if (token >= num_tokens) return;\n int expert = expert_idx[token];\n int slot = slot_idx[token];\n int row = expert_offsets[expert] + slot;\n for (int d = 0; d < hidden_dim; d++) {\n out[row * hidden_dim + d] = token_hidden[token * hidden_dim + d];\n }\n}\n\nvoid pack_tokens_cuda(\n torch::Tensor token_hidden,\n torch::Tensor expert_idx,\n torch::Tensor slot_idx,\n torch::Tensor expert_offsets,\n torch::Tensor out\n) {\n int num_tokens = token_hidden.size(0);\n int hidden_dim = token_hidden.size(1);\n const int block_size = 256;\n const int num_blocks = (num_tokens + block_size - 1) / block_size;\n pack_tokens_kernel<<<num_blocks, block_size>>>(\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}\n'''\n\ncpp_source = '''\nvoid pack_tokens_cuda(\n torch::Tensor token_hidden,\n torch::Tensor expert_idx,\n torch::Tensor slot_idx,\n torch::Tensor expert_offsets,\n torch::Tensor out\n);\n'''\n\npack_module = load_inline(\n name='pack_tokens',\n cpp_sources=[cpp_source],\n cuda_sources=[cuda_source],\n functions=['pack_tokens_cuda'],\n verbose=True,\n with_cuda=True\n)\n\nclass ModelNew(nn.Module):\n def __init__(self):\n super().__init__()\n self.pack = pack_module\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 num_rows = int(expert_offsets[-1].item())\n out = torch.zeros(num_rows, token_hidden.shape[1], dtype=token_hidden.dtype, device=token_hidden.device)\n self.pack.pack_tokens_cuda(token_hidden.contiguous(), expert_idx.to(torch.int32).contiguous(), slot_idx.to(torch.int32).contiguous(), expert_offsets.to(torch.int32).contiguous(), out)\n return out\n"
}
result ← run_correctness (ok)
okrun_correctness
run_correctness PASSED: (3 / 3) trials all matched the reference.
Numerical precision: max_abs_err=0.00e+00 mean_abs_err=0.00e+00 max_rel_err=0.00e+00
turn 8111.5s turn · 9.2s LLM · 102.3s tools · tokens in=9,530 · out=575 · total=10,176 · FINAL
reasoning
Great, correctness passed with the casts. Now, to make it faster, I can optimize the kernel, e.g., vectorized loads/stores since hidden_dim=128.
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 = '''\n#include <torch/extension.h>\n\n__global__ void pack_tokens_kernel(\n const float* token_hidden,\n const int* expert_idx,\n const int* slot_idx,\n const int* expert_offsets,\n float* out,\n int num_tokens,\n int hidden_dim\n) {\n int token = blockIdx.x * blockDim.x + threadIdx.x;\n if (token >= num_tokens) return;\n int expert = expert_idx[token];\n int slot = slot_idx[token];\n int row = expert_offsets[expert] + slot;\n for (int d = 0; d < hidden_dim; d++) {\n out[row * hidden_dim + d] = token_hidden[token * hidden_dim + d];\n }\n}\n\nvoid pack_tokens_cuda(\n torch::Tensor token_hidden,\n torch::Tensor expert_idx,\n torch::Tensor slot_idx,\n torch::Tensor expert_offsets,\n torch::Tensor out\n) {\n int num_tokens = token_hidden.size(0);\n int hidden_dim = token_hidden.size(1);\n const int block_size = 256;\n const int num_blocks = (num_tokens + block_size - 1) / block_size;\n pack_tokens_kernel<<<num_blocks, block_size>>>(\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}\n'''\n\ncpp_source = '''\nvoid pack_tokens_cuda(\n torch::Tensor token_hidden,\n torch::Tensor expert_idx,\n torch::Tensor slot_idx,\n torch::Tensor expert_offsets,\n torch::Tensor out\n);\n'''\n\npack_module = load_inline(\n name='pack_tokens',\n cpp_sources=[cpp_source],\n cuda_sources=[cuda_source],\n functions=['pack_tokens_cuda'],\n verbose=True,\n with_cuda=True\n)\n\nclass ModelNew(nn.Module):\n def __init__(self):\n super().__init__()\n self.pack = pack_module\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 num_rows = int(expert_offsets[-1].item())\n out = torch.zeros(num_rows, token_hidden.shape[1], dtype=token_hidden.dtype, device=token_hidden.device)\n self.pack.pack_tokens_cuda(token_hidden.contiguous(), expert_idx.to(torch.int32).contiguous(), slot_idx.to(torch.int32).contiguous(), expert_offsets.to(torch.int32).contiguous(), out)\n return out\n"
}
result ← submit_kernel (ok)
oksubmit_kernel
submit_kernel PASSED: (3 / 3) correctness trials all passed.
Kernel runtime: 0.13 μs
Runtime stats: mean=0.13μs median=?μs std=0.07μs
Numerical precision: max_abs_err=0.00e+00 mean_abs_err=0.00e+00 max_rel_err=0.00e+00
Memory: 35.03 MB peak (ref: 35.02 MB, ratio: 1.0005x)
Kernel launches: 20 (ref: 34824, fusion ratio: 1741.2)
SOL score: 0.0153
Flagged for excessive speedup — this submission may have been rejected by automated review.
final submitted kernel (85 lines, 2,384 chars) — full source, untruncated
import torch
import torch.nn as nn
from torch.utils.cpp_extension import load_inline
cuda_source = '''
#include <torch/extension.h>
__global__ void pack_tokens_kernel(
const float* token_hidden,
const int* expert_idx,
const int* slot_idx,
const int* expert_offsets,
float* out,
int num_tokens,
int hidden_dim
) {
int token = blockIdx.x * blockDim.x + threadIdx.x;
if (token >= num_tokens) return;
int expert = expert_idx[token];
int slot = slot_idx[token];
int row = expert_offsets[expert] + slot;
for (int d = 0; d < hidden_dim; d++) {
out[row * hidden_dim + d] = token_hidden[token * hidden_dim + d];
}
}
void pack_tokens_cuda(
torch::Tensor token_hidden,
torch::Tensor expert_idx,
torch::Tensor slot_idx,
torch::Tensor expert_offsets,
torch::Tensor out
) {
int num_tokens = token_hidden.size(0);
int hidden_dim = token_hidden.size(1);
const int block_size = 256;
const int num_blocks = (num_tokens + block_size - 1) / block_size;
pack_tokens_kernel<<<num_blocks, block_size>>>(
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
);
}
'''
cpp_source = '''
void pack_tokens_cuda(
torch::Tensor token_hidden,
torch::Tensor expert_idx,
torch::Tensor slot_idx,
torch::Tensor expert_offsets,
torch::Tensor out
);
'''
pack_module = load_inline(
name='pack_tokens',
cpp_sources=[cpp_source],
cuda_sources=[cuda_source],
functions=['pack_tokens_cuda'],
verbose=True,
with_cuda=True
)
class ModelNew(nn.Module):
def __init__(self):
super().__init__()
self.pack = pack_module
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)
self.pack.pack_tokens_cuda(token_hidden.contiguous(), expert_idx.to(torch.int32).contiguous(), slot_idx.to(torch.int32).contiguous(), expert_offsets.to(torch.int32).contiguous(), out)
return out