reasoning
reasoning (truncated 29356 chars): Let me analyze this task. I need to implement a custom CUDA kernel that replaces a PyTorch model for packing tokens into expert-major order for MoE (Mixture of Experts) execution.
The reference implementation:
1. Takes `token_hidden` (num_tokens, hidden_dim) - the token activations
2. Takes `expert_idx` (num_tokens,) - which expert each token is assigned to
3. Takes `slot_idx` (num_tokens,) - which slot within the expert's buffer
4. Takes `expert_offsets` (num_experts + 1,) - offsets for each expert in the output buffer
5. Creates output tensor of shape (num_rows, hidden_dim) where num_rows = expert_offsets[-1]
6. For each token, computes the row = expert_offsets[expert] + slot_idx[token] and copies the token's hidden state to that row
This is essentially a scatter operation - scattering tokens to specific positions in an output buffer based on expert assignment.
The straightforward CUDA kernel would:
- Have each thread handle one token
- Compute the destination row
- Copy the hidden_dim values from the input to the output
Let me write a CUDA kernel for this:
```cpp
#include <torch/extension.h>
#include <cuda.h>
#include <cuda_runtime.h>
template <typename scalar_t>
__global__ void pack_tokens_kernel(
const scalar_t* __restrict__ token_hidden,
const int32_t* __restrict__ expert_idx,
const int32_t* __restrict__ slot_idx,
const int32_t* __restrict__ expert_offsets,
scalar_t* __restrict__ 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 row = expert_offsets[expert] + slot_idx[token];
// Copy hidden_dim values
for (int i = 0; i < hidden_dim; i++) {
out[row * hidden_dim + i] = token_hidden[token * hidden_dim + i];
}
}
```
But this is inefficient because each thread is doing a loop over hidden_dim. A better approach would be to use cooperative groups or have threads within a warp cooperate.
Actually, for better memory coalescing, we could:
1. Have each thread handle one element (token, hidden_idx)
2. Or have warps cooperate on copying a token's hidden state
Let me think about a more optimized approach:
Option 1: Each thread handles one (token, hidden_idx) pair
- This means num_tokens * hidden_dim threads
- Grid can be (num_tokens,) and block can be (hidden_dim,)
- Or we could use 2D grid/blocks
Option 2: Each warp handles one token
- 32 threads cooperatively copy hidden_dim values
- If hidden_dim = 128, each thread copies 4 values
- Better for coalesced memory access
Let me go with a hybrid approach - have each warp handle one token, with threads cooperatively copying:
```cpp
template <typename scalar_t>
__global__ void pack_tokens_kernel(
const scalar_t* __restrict__ token_hidden,
const int32_t* __restrict__ expert_idx,
const int32_t* __restrict__ slot_idx,
const int32_t* __restrict__ expert_offsets,
scalar_t* __restrict__ out,
int num_tokens,
int hidden_dim) {
int token = blockIdx.x * blockDim.y + threadIdx.y; // Each warp-row handles one token
if (token >= num_tokens) return;
int expert = expert_idx[token];
int row = expert_offsets[expert] + slot_idx[token];
// Each thread handles hidden_dim/WARP_SIZE elements
const int WARP_SIZE = 32;
int lane = threadIdx.x;
for (int i = lane; i < hidden_dim; i += WARP_SIZE) {
out[row * hidden_dim + i] = token_hidden[token * hidden_dim + i];
}
}
```
Actually, let me reconsider. With hidden_dim = 128 and WARP_SIZE = 32, each thread would handle 4 elements. This is fine.
But I should also consider using vectorized loads/stores for better memory throughput.
Let me try a simpler approach first - one thread per token, loop over hidden_dim:
Actually, looking at the problem more carefully:
- num_tokens = 2048
- hidden_dim = 128
- num_experts = 16
With 2048 tokens and 128 hidden_dim, we have 2048 * 128 = 262144 elements to copy.
The simplest approach is:
- One thread per token (2048 threads)
- Each thread loops over hidden_dim (128 iterations)
Or:
- One thread per element (262144 threads)
- Each thread copies one element
Let me try the second approach for better parallelism:
```cpp
template <typename scalar_t>
__global__ void pack_tokens_kernel(
const scalar_t* __restrict__ token_hidden,
const int32_t* __restrict__ expert_idx,
const int32_t* __restrict__ slot_idx,
const int32_t* __restrict__ expert_offsets,
scalar_t* __restrict__ out,
int num_tokens,
int hidden_dim) {
int idx = blockIdx.x * blockDim.x + threadIdx.x;
int total = num_tokens * hidden_dim;
if (idx >= total) return;
int token = idx / hidden_dim;
int h = idx % hidden_dim;
int expert = expert_idx[token];
int row = expert_offsets[expert] + slot_idx[token];
out[row * hidden_dim + h] = token_hidden[token * hidden_dim + h];
}
```
This should work well for memory coalescing because threads with consecutive indices access consecutive memory locations in both the input and (usually) the output.
But wait - the output writes to different rows based on expert assignment. So consecutive tokens might not have consecutive rows in the output. This could lead to scattered writes.
Actually, let's think about the data pattern:
- expert_idx[token] = token % num_experts = token % 16
- So tokens 0, 16, 32, ... all go to expert 0
- tokens 1, 17, 33, ... all go to expert 1
- etc.
The output layout:
- Row 0 to (counts[0]-1) for expert 0
- Row counts[0] to (counts[0]+counts[1]-1) for expert 1
- etc.
Since each expert gets tokens in a round-robin fashion, the slot_idx values would be:
- For expert 0: tokens 0, 16, 32, ... get slots 0, 1, 2, ...
- For expert 1: tokens 1, 17, 33, ... get slots 0, 1, 2, ...
So the output rows for consecutive tokens are scattered across different experts' buffers.
The input reads are coalesced (consecutive threads read consecutive elements from input).
The output writes might be scattered but within each expert's range, they're sequential.
Let me just write a simple kernel and see:
```python
import torch
import torch.nn as nn
from torch.utils.cpp_extension import load_inline
cuda_source = """
#include <torch/extension.h>
#include <cuda.h>
#include <cuda_runtime.h>
template <typename scalar_t>
__global__ void pack_tokens_kernel(
const scalar_t* __restrict__ token_hidden,
const int32_t* __restrict__ expert_idx,
const int32_t* __restrict__ slot_idx,
const int32_t* __restrict__ expert_offsets,
scalar_t* __restrict__ out,
int num_tokens,
int hidden_dim) {
int idx = blockIdx.x * blockDim.x + threadIdx.x;
int total = num_tokens * hidden_dim;
if (idx >= total) return;
int token = idx / hidden_dim;
int h = idx % hidden_dim;
int expert = expert_idx[token];
int row = expert_offsets[expert] + slot_idx[token];
out[row * hidden_dim + h] = token_hidden[token * hidden_dim + h];
}
torch::Tensor pack_tokens_cuda(
torch::Tensor token_hidden,
torch::Tensor expert_idx,
torch::Tensor slot_idx,
torch::Tensor expert_offsets,
int num_rows) {
int num_tokens = token_hidden.size(0);
int hidden_dim = token_hidden.size(1);
auto out = torch::zeros({num_rows, hidden_dim}, token_hidden.options());
int total = num_tokens * hidden_dim;
const int block_size = 256;
const int num_blocks = (total + block_size - 1) / block_size;
AT_DISPATCH_FLOATING_TYPES(token_hidden.scalar_type(), "pack_tokens_kernel", ([&] {
pack_tokens_kernel<scalar_t><<<num_blocks, block_size>>>(
token_hidden.data_ptr<scalar_t>(),
expert_idx.data_ptr<int32_t>(),
slot_idx.data_ptr<int32_t>(),
expert_offsets.data_ptr<int32_t>(),
out.data_ptr<scalar_t>(),
num_tokens,
hidden_dim);
}));
return out
...[ omitted 29356 characters ]...
n_hidden + warp_id * hidden_dim;
float* out_ptr = out + row * hidden_dim;
// Vectorized copy (float4)
int vec_dim = hidden_dim >> 2;
for (int i = lane; i < vec_dim; i += 32) {
reinterpret_cast<float4*>(out_ptr)[i] = reinterpret_cast<const float4*>(in_ptr)[i];
}
// Handle remaining elements (if hidden_dim not divisible by 4)
int remaining = hidden_dim & 3;
int remaining_start = vec_dim << 2;
for (int i = remaining_start + lane; i < hidden_dim; i += 32) {
out_ptr[i] = in_ptr[i];
}
}
torch::Tensor pack_tokens_cuda(
torch::Tensor token_hidden,
torch::Tensor expert_idx,
torch::Tensor slot_idx,
torch::Tensor expert_offsets,
int num_rows) {
token_hidden = token_hidden.contiguous();
expert_idx = expert_idx.contiguous();
slot_idx = slot_idx.contiguous();
expert_offsets = expert_offsets.contiguous();
int num_tokens = token_hidden.size(0);
int hidden_dim = token_hidden.size(1);
auto out = torch::zeros({num_rows, hidden_dim}, token_hidden.options());
int threads_per_block = 256;
int warps_per_block = threads_per_block >> 5;
int num_blocks = (num_tokens + warps_per_block - 1) / warps_per_block;
pack_tokens_kernel<<<num_blocks, threads_per_block>>>(
token_hidden.data_ptr<float>(),
expert_idx.data_ptr<int32_t>(),
slot_idx.data_ptr<int32_t>(),
expert_offsets.data_ptr<int32_t>(),
out.data_ptr<float>(),
num_tokens,
hidden_dim);
return out;
}
"""
cpp_source = """
torch::Tensor pack_tokens_cuda(torch::Tensor token_hidden, torch::Tensor expert_idx, torch::Tensor slot_idx, torch::Tensor expert_offsets, int num_rows);
"""
pack_module = load_inline(
name='pack_tokens',
cpp_sources=cpp_source,
cuda_sources=cuda_source,
functions=['pack_tokens_cuda'],
extra_cuda_cflags=['-O3'],
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:
num_rows = int(expert_offsets[-1].item())
return pack_module.pack_tokens_cuda(token_hidden, expert_idx, slot_idx, expert_offsets, num_rows)
```
Wait, I need to check the correctness of the remaining element handling. If hidden_dim = 128, then:
- vec_dim = 128 >> 2 = 32
- remaining = 128 & 3 = 0
- remaining_start = 32 << 2 = 128
- The loop `for (int i = 128 + lane; i < 128; i += 32)` never executes
So the remaining element loop is a no-op for hidden_dim = 128, which is correct.
Actually, let me double-check the vectorized access. We're using:
```cpp
reinterpret_cast<float4*>(out_ptr)[i] = reinterpret_cast<const float4*>(in_ptr)[i];
```
This accesses:
- out_ptr[4*i], out_ptr[4*i+1], out_ptr[4*i+2], out_ptr[4*i+3]
- in_ptr[4*i], in_ptr[4*i+1], in_ptr[4*i+2], in_ptr[4*i+3]
For lane 0, i=0: elements 0-3
For lane 1, i=1: elements 4-7
...
For lane 31, i=31: elements 124-127
So the entire hidden_dim = 128 elements are covered by the 32 lanes, each handling one float4. This is correct.
Actually wait, let me re-check. For vec_dim = 32:
- lane 0: i = 0
- lane 1: i = 1
- ...
- lane 31: i = 31
And the loop is `for (int i = lane; i < vec_dim; i += 32)`, so:
- lane 0: i = 0 (valid, since 0 < 32)
- lane 31: i = 31 (valid, since 31 < 32)
- lane 0 next iteration: i = 0 + 32 = 32 (invalid, since 32 >= 32)
So each lane handles exactly one float4, which covers the entire hidden_dim = 128 elements. Correct!
I think the kernel is correct. Let me submit it.
Actually, I want to double-check one more thing: the input tensors might have different memory layouts. The reference implementation uses indexing like `token_hidden[token]`, which would work with strided tensors. But my kernel assumes contiguous memory.
I'm calling `.contiguous()` on all input tensors, which should handle this. But let me make sure:
```cpp
token_hidden = token_hidden.contiguous();
```
This creates a contiguous copy if the tensor is not already contiguous. If it is contiguous, it returns the same tensor (no copy).
Alright, I'm confident the kernel is correct. Let me submit.
Oh wait, one more thing. I need to make sure the module is loaded only once. The `load_inline` function compiles and loads the module, and I'm doing this at module level, which means it's executed when the Python file is loaded. This should be fine.
But I should make sure the module name is unique to avoid conflicts. Let me use a more unique name:
Actually, `load_inline` handles this by using a unique hash of the source code. So the name I provide is just a base name.
Let me finalize and submit:
```python
import torch
import torch.nn as nn
from torch.utils.cpp_extension import load_inline
cuda_source = """
#include <torch/extension.h>
#include <cuda.h>
#include <cuda_runtime.h>
__global__ void pack_tokens_kernel(
const float* __restrict__ token_hidden,
const int32_t* __restrict__ expert_idx,
const int32_t* __restrict__ slot_idx,
const int32_t* __restrict__ expert_offsets,
float* __restrict__ out,
int num_tokens,
int hidden_dim) {
int warp_id = (blockIdx.x * blockDim.x + threadIdx.x) >> 5;
int lane = threadIdx.x & 31;
if (warp_id >= num_tokens) return;
int expert = __ldg(&expert_idx[warp_id]);
int row = __ldg(&expert_offsets[expert]) + __ldg(&slot_idx[warp_id]);
const float* in_ptr = token_hidden + warp_id * hidden_dim;
float* out_ptr = out + row * hidden_dim;
// Vectorized copy (float4)
int vec_dim = hidden_dim >> 2;
for (int i = lane; i < vec_dim; i += 32) {
reinterpret_cast<float4*>(out_ptr)[i] = reinterpret_cast<const float4*>(in_ptr)[i];
}
// Handle remaining elements (if hidden_dim not divisible by 4)
int remaining_start = vec_dim << 2;
for (int i = remaining_start + lane; i < hidden_dim; i += 32) {
out_ptr[i] = in_ptr[i];
}
}
torch::Tensor pack_tokens_cuda(
torch::Tensor token_hidden,
torch::Tensor expert_idx,
torch::Tensor slot_idx,
torch::Tensor expert_offsets,
int num_rows) {
token_hidden = token_hidden.contiguous();
expert_idx = expert_idx.contiguous();
slot_idx = slot_idx.contiguous();
expert_offsets = expert_offsets.contiguous();
int num_tokens = token_hidden.size(0);
int hidden_dim = token_hidden.size(1);
auto out = torch::zeros({num_rows, hidden_dim}, token_hidden.options());
int threads_per_block = 256;
int warps_per_block = threads_per_block >> 5;
int num_blocks = (num_tokens + warps_per_block - 1) / warps_per_block;
pack_tokens_kernel<<<num_blocks, threads_per_block>>>(
token_hidden.data_ptr<float>(),
expert_idx.data_ptr<int32_t>(),
slot_idx.data_ptr<int32_t>(),
expert_offsets.data_ptr<int32_t>(),
out.data_ptr<float>(),
num_tokens,
hidden_dim);
return out;
}
"""
cpp_source = """
torch::Tensor pack_tokens_cuda(torch::Tensor token_hidden, torch::Tensor expert_idx, torch::Tensor slot_idx, torch::Tensor expert_offsets, int num_rows);
"""
pack_module = load_inline(
name='pack_tokens_inline',
cpp_sources=cpp_source,
cuda_sources=cuda_source,
functions=['pack_tokens_cuda'],
extra_cuda_cflags=['-O3'],
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:
num_rows = int(expert_offsets[-1].item())
return pack_module.pack_tokens_cuda(token_hidden, expert_idx, slot_idx, expert_offsets, num_rows)
```
This should be the final solution. Let me submit it.
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#include <cuda.h>\n#include <cuda_runtime.h>\n\n__global__ void pack_tokens_kernel(\n const float* __restrict__ token_hidden,\n const int32_t* __restrict__ expert_idx,\n const int32_t* __restrict__ slot_idx,\n const int32_t* __restrict__ expert_offsets,\n float* __restrict__ out,\n int num_tokens,\n int hidden_dim) {\n \n int warp_id = (blockIdx.x * blockDim.x + threadIdx.x) >> 5;\n int lane = threadIdx.x & 31;\n \n if (warp_id >= num_tokens) return;\n \n int expert = __ldg(&expert_idx[warp_id]);\n int row = __ldg(&expert_offsets[expert]) + __ldg(&slot_idx[warp_id]);\n \n const float* in_ptr = token_hidden + warp_id * hidden_dim;\n float* out_ptr = out + row * hidden_dim;\n \n // Vectorized copy (float4)\n int vec_dim = hidden_dim >> 2;\n for (int i = lane; i < vec_dim; i += 32) {\n reinterpret_cast<float4*>(out_ptr)[i] = reinterpret_cast<const float4*>(in_ptr)[i];\n }\n \n // Handle remaining elements (if hidden_dim not divisible by 4)\n int remaining_start = vec_dim << 2;\n for (int i = remaining_start + lane; i < hidden_dim; i += 32) {\n out_ptr[i] = in_ptr[i];\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 int num_rows) {\n \n token_hidden = token_hidden.contiguous();\n expert_idx = expert_idx.contiguous();\n slot_idx = slot_idx.contiguous();\n expert_offsets = expert_offsets.contiguous();\n \n int num_tokens = token_hidden.size(0);\n int hidden_dim = token_hidden.size(1);\n \n auto out = torch::zeros({num_rows, hidden_dim}, token_hidden.options());\n \n int threads_per_block = 256;\n int warps_per_block = threads_per_block >> 5;\n int num_blocks = (num_tokens + warps_per_block - 1) / warps_per_block;\n \n pack_tokens_kernel<<<num_blocks, threads_per_block>>>(\n token_hidden.data_ptr<float>(),\n expert_idx.data_ptr<int32_t>(),\n slot_idx.data_ptr<int32_t>(),\n expert_offsets.data_ptr<int32_t>(),\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(torch::Tensor token_hidden, torch::Tensor expert_idx, torch::Tensor slot_idx, torch::Tensor expert_offsets, int num_rows);\n\"\"\"\n\npack_module = load_inline(\n name='pack_tokens_inline',\n cpp_sources=cpp_source,\n cuda_sources=cuda_source,\n functions=['pack_tokens_cuda'],\n extra_cuda_cflags=['-O3'],\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 num_rows = int(expert_offsets[-1].item())\n return pack_module.pack_tokens_cuda(token_hidden, expert_idx, slot_idx, expert_offsets, num_rows)"
}