reasoning
reasoning (truncated 52600 chars): Let me analyze this task. The reference implementation is doing a weighted scatter-add operation:
1. `expert_hidden` has shape `(num_tokens * fanout, hidden_dim)` - these are expert outputs
2. `token_idx` has shape `(num_tokens * fanout,)` - indicates which token each expert output belongs to
3. `gates` has shape `(num_tokens * fanout,)` - the routing weights
4. Output should have shape `(num_tokens, hidden_dim)` - combined expert outputs for each token
The operation is: `out[token_idx[row]] += gates[row] * expert_hidden[row]`
This is essentially a scatter-add with scaling. This is a common operation in mixture-of-experts models.
For the CUDA kernel:
- Each thread could handle one row of `expert_hidden`
- Use atomicAdd to accumulate results (since multiple rows can map to the same token)
- Need to handle the case where fanout > 1 (multiple expert outputs per token)
Let me think about optimization:
1. Since we're doing atomic adds, we need to be careful about race conditions
2. The hidden dimension is 128, so each thread handles 128 floats
3. We could potentially use shared memory for reduction, but with random token_idx patterns, that might not help much
Let me write a kernel:
```python
import torch
import torch.nn as nn
from torch.utils.cpp_extension import load_inline
# Define the CUDA kernel
cuda_source = """
__global__ void weighted_scatter_add_kernel(
const float* __restrict__ expert_hidden,
const int* __restrict__ token_idx,
const float* __restrict__ gates,
float* __restrict__ out,
int num_rows,
int hidden_dim
) {
int row = blockIdx.x * blockDim.x + threadIdx.x;
if (row >= num_rows) return;
int token = token_idx[row];
float gate = gates[row];
const float* expert_row = expert_hidden + row * hidden_dim;
float* out_row = out + token * hidden_dim;
for (int h = 0; h < hidden_dim; h++) {
atomicAdd(&out_row[h], gate * expert_row[h]);
}
}
void weighted_scatter_add_cuda(
const float* expert_hidden,
const int* token_idx,
const float* gates,
float* out,
int num_rows,
int hidden_dim
) {
int threads_per_block = 256;
int blocks = (num_rows + threads_per_block - 1) / threads_per_block;
weighted_scatter_add_kernel<<<blocks, threads_per_block>>>(
expert_hidden, token_idx, gates, out, num_rows, hidden_dim
);
}
"""
cpp_source = """
void weighted_scatter_add_cuda(
const float* expert_hidden,
const int* token_idx,
const float* gates,
float* out,
int num_rows,
int hidden_dim
);
"""
# Load the extension
module = load_inline(
name='weighted_scatter_add',
cpp_sources=cpp_source,
cuda_sources=cuda_source,
functions=['weighted_scatter_add_cuda'],
extra_cuda_cflags=['-O3'],
)
```
Wait, atomicAdd on float requires compute capability 2.0+, which should be fine for modern GPUs. But let me think about better optimizations.
Actually, using atomic operations for every element might be slow. Let me think about alternatives:
1. If hidden_dim is fixed at 128, we could have each warp handle one row, with each thread handling multiple hidden elements
2. We could use vectorized loads/stores
3. We could consider using shared memory for local reductions before writing to global memory
Given that fanout = 2, each token has at most 2 expert outputs contributing to it. This means we won't have many collisions.
Let me try a different approach - have each thread handle multiple elements of the hidden dimension:
```cpp
__global__ void weighted_scatter_add_kernel(
const float* __restrict__ expert_hidden,
const int* __restrict__ token_idx,
const float* __restrict__ gates,
float* __restrict__ out,
int num_rows,
int hidden_dim
) {
int row = blockIdx.x * blockDim.x + threadIdx.x;
if (row >= num_rows) return;
int token = token_idx[row];
float gate = gates[row];
const float* expert_row = expert_hidden + row * hidden_dim;
float* out_row = out + token * hidden_dim;
for (int h = threadIdx.y; h < hidden_dim; h += blockDim.y) {
atomicAdd(&out_row[h], gate * expert_row[h]);
}
}
```
Actually, let me think more carefully. The problem is that we're iterating over the rows (num_tokens * fanout = 3072 rows), and each row has 128 elements.
A better approach might be to use one thread per element:
- Total elements: num_rows * hidden_dim = 3072 * 128 = 393,216 elements
- Grid: could use 3D indexing or flatten it
Actually, let me try vectorization with float4:
- hidden_dim = 128 = 32 * 4
- Each thread handles 4 floats at a time
```cpp
__global__ void weighted_scatter_add_kernel(
const float* __restrict__ expert_hidden,
const int* __restrict__ token_idx,
const float* __restrict__ gates,
float* __restrict__ out,
int num_rows,
int hidden_dim
) {
int row = blockIdx.x;
int h_base = threadIdx.x * 4; // Each thread handles 4 floats
if (row >= num_rows || h_base >= hidden_dim) return;
int token = token_idx[row];
float gate = gates[row];
const float4* expert_row = (const float4*)(expert_hidden + row * hidden_dim);
float4* out_row = (float4*)(out + token * hidden_dim);
int vec_idx = threadIdx.x;
float4 vals = expert_row[vec_idx];
atomicAdd(&out_row[vec_idx].x, gate * vals.x);
atomicAdd(&out_row[vec_idx].y, gate * vals.y);
atomicAdd(&out_row[vec_idx].z, gate * vals.z);
atomicAdd(&out_row[vec_idx].w, gate * vals.w);
}
```
This would use blockDim.x = 32 threads per block, one block per row. That's 3072 blocks.
Actually, let me think about the memory access pattern:
- `expert_hidden` is accessed by row - coalesced if threads access consecutive hidden elements
- `out` is accessed by token - could cause bank conflicts in atomic operations
Given the small fanout (2), atomic collisions should be minimal. Let me try a simpler approach first:
Actually, wait. Looking at this more carefully:
The reference implementation uses a Python loop which is very slow. A simple CUDA kernel with atomics should be much faster.
Let me write a clean implementation:
```cpp
__global__ void scatter_add_kernel(
const float* __restrict__ expert_hidden,
const int* __restrict__ token_idx,
const float* __restrict__ gates,
float* __restrict__ out,
const int num_rows,
const int hidden_dim
) {
const int row = blockIdx.x;
const int tid = threadIdx.x;
const int token = token_idx[row];
const float gate = gates[row];
const float* src = expert_hidden + row * hidden_dim;
float* dst = out + token * hidden_dim;
// Each thread handles multiple elements
for (int i = tid; i < hidden_dim; i += blockDim.x) {
atomicAdd(dst + i, gate * src[i]);
}
}
```
This has one block per row, with each thread handling a strided subset of hidden elements.
But with 128 hidden elements and say 128 threads, each thread handles 1 element. That's efficient.
Actually, let me check the launch configuration:
- num_rows = num_tokens * fanout = 1536 * 2 = 3072
- hidden_dim = 128
- If we use blockDim.x = 128 (one warp = 32 threads, 4 warps = 128 threads)
- Number of blocks = num_rows = 3072
Hmm, but using blockDim.x = 128 might be overkill if hidden_dim = 128 exactly. Let me use blockDim.x = 128, and each thread handles exactly 1 element.
Wait, but what if hidden_dim changes? Let me make it general:
- Use blockDim.x = 128 threads
- Each thread handles elements [tid, tid + 128, tid + 256, ...]
Actually, for hidden_dim = 128, let's just use 128 threads per block, each handling one element.
```cpp
__global__ void scatter_add_kernel(
const float* __restrict__ expert_hidden,
const int* __restrict__ token_idx,
const float* __restrict__ gates,
float* __restrict__ out,
const int num_rows,
const int hidden_dim
) {
const int row = blockIdx.x;
const int tid = threadIdx.x;
if (row >= num_rows) return;
const int t
...[ omitted 52600 characters ]...
eadIdx.x;
if (row >= num_rows) return;
const int token = token_idx[row];
const float gate = gates[row];
const float* src = expert_hidden + row * hidden_dim;
float* dst = out + token * hidden_dim;
// Process with vectorized loads (float4)
const int num_vec4 = hidden_dim / 4;
const float4* src_vec = reinterpret_cast<const float4*>(src);
float4* dst_vec = reinterpret_cast<float4*>(dst);
// Each thread handles one or more float4 vectors
for (int v = tid; v < num_vec4; v += blockDim.x) {
const float4 val = src_vec[v];
atomicAdd(&dst_vec[v].x, gate * val.x);
atomicAdd(&dst_vec[v].y, gate * val.y);
atomicAdd(&dst_vec[v].z, gate * val.z);
atomicAdd(&dst_vec[v].w, gate * val.w);
}
// Handle remaining elements
const int remaining = hidden_dim - num_vec4 * 4;
if (tid < remaining) {
const int idx = num_vec4 * 4 + tid;
atomicAdd(&dst[idx], gate * src[idx]);
}
}
"""
cpp_source = """
#include <torch/extension.h>
void scatter_add_cuda(
torch::Tensor expert_hidden,
torch::Tensor token_idx,
torch::Tensor gates,
torch::Tensor out
) {
const int num_rows = expert_hidden.size(0);
const int hidden_dim = expert_hidden.size(1);
const int threads = 32;
const int blocks = num_rows;
scatter_add_kernel<<<blocks, threads>>>(
expert_hidden.data_ptr<float>(),
token_idx.data_ptr<int>(),
gates.data_ptr<float>(),
out.data_ptr<float>(),
num_rows,
hidden_dim
);
}
"""
scatter_add_module = load_inline(
name='scatter_add_module',
cpp_sources=cpp_source,
cuda_sources=cuda_source,
functions=['scatter_add_cuda'],
extra_cuda_cflags=['-O3'],
)
class ModelNew(nn.Module):
def __init__(self):
super().__init__()
def forward(
self,
expert_hidden: torch.Tensor,
token_idx: torch.Tensor,
gates: torch.Tensor,
num_tokens: int,
) -> torch.Tensor:
out = torch.zeros(num_tokens, expert_hidden.shape[1],
dtype=expert_hidden.dtype, device=expert_hidden.device)
scatter_add_module.scatter_add_cuda(expert_hidden, token_idx, gates, out)
return out
```
Hmm, using 32 threads per block might have lower occupancy than 128 threads. But since we're using atomics, the extra threads might not help much anyway.
Let me think about the trade-offs:
- 32 threads: 1 warp per block, efficient for hidden_dim = 128 with float4
- 128 threads: 4 warps per block, better for hidden_dim that are multiples of 128
For hidden_dim = 128, 32 threads is optimal. For hidden_dim = 256, we'd need 64 threads with float4 (or 128 threads with loop).
Let me use a general approach that works well for various hidden_dim:
Actually, let me just use 128 threads and the simple scalar version. It's simple and should work well:
```python
import torch
import torch.nn as nn
from torch.utils.cpp_extension import load_inline
cuda_source = """
__global__ void scatter_add_kernel(
const float* __restrict__ expert_hidden,
const int* __restrict__ token_idx,
const float* __restrict__ gates,
float* __restrict__ out,
const int num_rows,
const int hidden_dim
) {
const int row = blockIdx.x;
const int tid = threadIdx.x;
if (row >= num_rows) return;
const int token = token_idx[row];
const float gate = gates[row];
for (int h = tid; h < hidden_dim; h += blockDim.x) {
const float val = expert_hidden[row * hidden_dim + h];
atomicAdd(&out[token * hidden_dim + h], gate * val);
}
}
"""
cpp_source = """
#include <torch/extension.h>
void scatter_add_cuda(
torch::Tensor expert_hidden,
torch::Tensor token_idx,
torch::Tensor gates,
torch::Tensor out
) {
const int num_rows = expert_hidden.size(0);
const int hidden_dim = expert_hidden.size(1);
const int threads = 128;
const int blocks = num_rows;
scatter_add_kernel<<<blocks, threads>>>(
expert_hidden.data_ptr<float>(),
token_idx.data_ptr<int>(),
gates.data_ptr<float>(),
out.data_ptr<float>(),
num_rows,
hidden_dim
);
}
"""
scatter_add_module = load_inline(
name='scatter_add_module',
cpp_sources=cpp_source,
cuda_sources=cuda_source,
functions=['scatter_add_cuda'],
extra_cuda_cflags=['-O3'],
)
class ModelNew(nn.Module):
def __init__(self):
super().__init__()
def forward(
self,
expert_hidden: torch.Tensor,
token_idx: torch.Tensor,
gates: torch.Tensor,
num_tokens: int,
) -> torch.Tensor:
out = torch.zeros(num_tokens, expert_hidden.shape[1],
dtype=expert_hidden.dtype, device=expert_hidden.device)
scatter_add_module.scatter_add_cuda(expert_hidden, token_idx, gates, out)
return out
```
Actually, I just realized that using 128 threads with hidden_dim = 128 means each thread handles exactly 1 element. But that means no loop! Let me simplify:
```cpp
__global__ void scatter_add_kernel(
const float* __restrict__ expert_hidden,
const int* __restrict__ token_idx,
const float* __restrict__ gates,
float* __restrict__ out,
const int num_rows,
const int hidden_dim
) {
const int row = blockIdx.x;
const int tid = threadIdx.x;
if (row >= num_rows) return;
if (tid >= hidden_dim) return; // Guard for threads beyond hidden_dim
const int token = token_idx[row];
const float gate = gates[row];
const float val = expert_hidden[row * hidden_dim + tid];
atomicAdd(&out[token * hidden_dim + tid], gate * val);
}
```
This is simpler and might be faster. But for hidden_dim != 128, we'd need the loop. Let me keep the loop for generality.
OK let me just submit the simple version:
```python
import torch
import torch.nn as nn
from torch.utils.cpp_extension import load_inline
cuda_source = """
__global__ void scatter_add_kernel(
const float* __restrict__ expert_hidden,
const int* __restrict__ token_idx,
const float* __restrict__ gates,
float* __restrict__ out,
const int num_rows,
const int hidden_dim
) {
const int row = blockIdx.x;
const int tid = threadIdx.x;
if (row >= num_rows) return;
const int token = token_idx[row];
const float gate = gates[row];
for (int h = tid; h < hidden_dim; h += blockDim.x) {
atomicAdd(&out[token * hidden_dim + h], gate * expert_hidden[row * hidden_dim + h]);
}
}
"""
cpp_source = """
#include <torch/extension.h>
void scatter_add_cuda(
torch::Tensor expert_hidden,
torch::Tensor token_idx,
torch::Tensor gates,
torch::Tensor out
) {
const int num_rows = expert_hidden.size(0);
const int hidden_dim = expert_hidden.size(1);
const int threads = 128;
const int blocks = num_rows;
scatter_add_kernel<<<blocks, threads>>>(
expert_hidden.data_ptr<float>(),
token_idx.data_ptr<int>(),
gates.data_ptr<float>(),
out.data_ptr<float>(),
num_rows,
hidden_dim
);
}
"""
scatter_add_module = load_inline(
name='scatter_add_module',
cpp_sources=cpp_source,
cuda_sources=cuda_source,
functions=['scatter_add_cuda'],
extra_cuda_cflags=['-O3'],
)
class ModelNew(nn.Module):
def __init__(self):
super().__init__()
def forward(
self,
expert_hidden: torch.Tensor,
token_idx: torch.Tensor,
gates: torch.Tensor,
num_tokens: int,
) -> torch.Tensor:
out = torch.zeros(num_tokens, expert_hidden.shape[1],
dtype=expert_hidden.dtype, device=expert_hidden.device)
scatter_add_module.scatter_add_cuda(expert_hidden, token_idx, gates, out)
return out
```
This should work. Let me submit it.