reasoning
Let me analyze the problem:
1. We have `num_tokens=2048`, `hidden_dim=128`, `num_experts=16`
2. We need to pack tokens into expert-major order
3. Each token goes to an expert determined by `expert_idx`, and within that expert's buffer at position `expert_offsets[expert] + slot_idx[token]`
4. Output size is `num_rows x hidden_dim` where `num_rows = expert_offsets[-1]` (total tokens)
This is essentially a scatter/gather operation. The naive loop processes one token at a time, which is very slow.
For the CUDA kernel, each thread can handle one element of the output. Since the mapping from output row to input token is known, we can:
1. For each output row, we need to know which token it came from
2. Or, for each token, scatter all its hidden elements to the output
Approach: Each token writes `hidden_dim` elements. We can have each thread handle one token and loop over `hidden_dim`. With 2048 tokens and 128 hidden_dim = 262144 total elements, we can also have each thread handle one output element.
The most natural approach:
- Each thread handles one token, writes all hidden_dim elements
- Number of threads = num_tokens = 2048
- Each thread does a strided copy of 128 floats
Or:
- Each thread handles one hidden element for one token
- Total threads = num_tokens * hidden_dim = 262144
- But we'd need to compute the output index per thread
Let me think about memory coalescing:
- If threads in a block handle consecutive tokens, and each thread writes 128 consecutive floats, then thread i writes to some potentially different row for each token. The writes won't be coalesced well.
- If we instead have each thread handle one element (e.g., thread idx handles output element `idx`), then we need to know for each output row which input token it came from. We don't have an inverse mapping readily available.
Actually, looking at the inputs again:
- `expert_idx[token]` = which expert token goes to
- `slot_idx[token]` = slot within expert
- `expert_offsets[expert]` = starting row for expert
So output row for token = `expert_offsets[expert_idx[token]] + slot_idx[token]`
For the inverse: given an output row, which token? We don't know directly.
So forward approach: each thread handles one or more tokens, and scatters them to output rows.
For memory coalescing of writes:
- If thread `t` handles token `t`, then it writes to `output[row_t * hidden_dim + d]` for d in 0..127
- These writes are strided by `hidden_dim` which is 128 * 4 bytes = 512 bytes, which might not be perfectly coalesced within a warp if different rows are far apart.
Alternative: use shared memory to transpose. Each warp handles a group of tokens, loads into shared memory, then writes out. But 128 * 4 = 512 bytes per token, for 32 tokens = 16KB, which fits in shared memory.
Actually, a simpler approach: each thread handles one output element. To find which token produced it, we need to invert the mapping. Can we do this efficiently?
We know that tokens are partitioned by expert. For a given output row, we can find which expert it belongs to (via binary search on expert_offsets), then find which slot within that expert, then find which token has that expert and slot. But we don't have a mapping from (expert, slot) -> token. We'd need to precompute it.
Wait, we CAN precompute it if we build the inverse mapping. But that would require an extra kernel or precomputation.
Actually, this MoE dispatch operation is very common. The standard optimized approach is:
1. Each thread handles one token and copies all hidden_dim elements
2. Use vectorized loads/stores (float4) to maximize throughput
With 2048 tokens and 128 dims, using float4 gives 32 float4 operations per token. Total memory moved = 2048 * 128 * 4 bytes * 2 (read + write) ≈ 2 MB. At 3.35 TB/s, this should take about 0.6 microseconds, so kernel launch overhead and occupancy might matter.
Let's write a simple kernel where each block handles a chunk of tokens, and each thread handles one token doing 128/4=32 float4 copies.
Wait, there might be an even better approach. If we launch one block per expert, then all threads in a block write to consecutive memory (since expert's output rows are consecutive). Within a block, if threads handle consecutive tokens belonging to the same expert, then the output rows are not necessarily consecutive (because slot_idx can be any order within the expert).
Actually, looking at how get_inputs works:
```python
for token in range(num_tokens):
expert = int(expert_idx[token].item())
slot_idx[token] = slot_cursor[expert]
slot_cursor[expert] += 1
```
So `slot_idx[token]` is assigned in order of first appearance. For `expert_idx = torch.arange(num_tokens) % num_experts`, the tokens are grouped by expert naturally: tokens 0, 16, 32... go to expert 0, tokens 1, 17, 33... go to expert 1, etc.
But in general, the `slot_idx` assignment means that if tokens arrive in sorted order by expert, then `slot_idx` is 0, 1, 2... for each expert. However, the actual token indices are interleaved across experts.
Actually let me re-read: tokens are processed in order 0, 1, 2, ..., and for each token, we assign it the next available slot in its expert. So for expert_idx = [0, 1, 2, ..., 15, 0, 1, 2, ...], the slot_idx for expert 0 are [0, 1, 2, ...] at token indices [0, 16, 32, ...].
The output rows for expert 0 are `expert_offsets[0] + slot_idx[token]` = `0 + [0, 1, 2, ...]` which are consecutive. For expert 1, they're `expert_offsets[1] + [0, 1, 2, ...]` = consecutive starting at offset for expert 1.
In this specific input pattern, the tokens for each expert are strided by num_experts=16. And the output rows for each expert are consecutive.
But in general, we can't assume this ordering. The kernel should work for arbitrary permutation.
For maximum performance with arbitrary permutation, I think the best approach is:
Option A: Coalesced read, scattered write (each token's elements are read coalescedly but written to potentially non-consecutive rows)
- Each thread handles one token
- Each thread does 32 float4 loads from `token_hidden[token]`
- Each thread does 32 float4 stores to `output[row * hidden_dim + d]`
- Since each thread does independent stores, and the base address `row * hidden_dim` can vary a lot between threads in a warp, stores won't coalesce well.
Option B: Use shared memory to help coalescing
- A warp/block loads token data for coalesced read
- Then writes out in a coalesced manner... but this requires knowing the output addresses first, which we do. Actually, if threads in a warp handle consecutive dimensions of the SAME token, they can coalesce writes too.
- Thread i in warp handles `output[row * hidden_dim + i]` for i = 0..31 (as float4 index)
- Then next 32 threads could handle token+1, but they write to different rows so no coalescing between these groups.
- But within each group of 32 threads, if they all handle the same token, they do the same work in parallel on consecutive addresses. So 32 threads cooperate on one token.
- This gives coalesced reads and coalesced writes!
- For 128 dims = 32 float4, each token needs 32 threads (1 warp). For 2048 tokens, we need 2048 warps = 64 blocks of 32 warps each, or just launch 2048 warps total.
Let me think: If we process one token per warp:
- Warp w handles token w
- Thread t in warp handles float4 element t (i.e., dims 4t, 4t+1, 4t+2, 4t+1)
- All 32 threads load `token_hidden[token][4t:4t+4]` -> these loads ARE coalesced (consecutive addresses)
- All 32 threads store to `output[row * hidden_dim + 4t : row * hidden_dim + 4t + 4]` -> these stores ARE coalesced (consecutive addresses within the row)
- This is perfect!
Wait but then we only have 32 threads active per block if we do one warp per block. Better: one block can handle multiple warps/tokens.
Let's say block has 256 threads = 8 warps. Handles 8 tokens.
- Warp 0 handles token `blockIdx.x * 8 + 0`
- Warp 1 handles token `blockIdx.x * 8 + 1`
- etc.
Total blocks = (num_tokens + 7) / 8 = 256 blocks.
Each warp:
1. Compute token index: `token = blockIdx.x * tokens_per_block + warp_id`
2. If token >= num_tokens, return
3. Compute row = expert_offsets[expert_idx[token]] + slot_idx[token]
4. Load 32 float4 from `token_hidden[token * hidden_dim + t * 4]`
5. Store 32 float4 to `output[row * hidden_dim + t * 4]`
This gives great coalescing and occupancy. Let me verify:
- Grid size: 256 blocks
- Block size: 256 threads
- Warps per block: 8
- Total warps: 2048
Occupancy on H100:
- Max warps per SM: 64 (I think for H100)
- Threads per SM: 2048
- With 256 threads per block, we can have up to 8 blocks per SM (2048/256)
- Total SMs: ~132
- So total blocks = 256, which is 2 blocks per SM on average, well within limits.
This should be very fast. Let's write the kernel.
Actually, one thing: `expert_offsets`, `expert_idx`, `slot_idx` are int32. We need to read them correctly.
Also, note that `expert_offsets` has length `num_experts + 1`, and `num_rows = expert_offsets[num_experts]`.
Kernel details:
```cuda
__global__ void dispatch_kernel(
const float* __restrict__ token_hidden,
const int* __restrict__ expert_idx,
const int* __restrict__ slot_idx,
const int* __restrict__ expert_offsets,
float* __restrict__ output,
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];
const float* src = token_hidden + token * hidden_dim;
float* dst = output + row * hidden_dim;
// Copy hidden_dim floats using float4 for vectorization
#pragma unroll
for (int i = threadIdx.x % 32; i < hidden_dim / 4; i += 32) {
// Wait this wouldn't work because each thread only does some iterations
// but we said each warp handles one token...
}
}
```
Wait, if blockDim.x = 256 and each warp is 32 threads, and each warp handles exactly one token, then:
- Warp 0: threads 0-31 handle token `blockIdx.x * 8 + 0`
- Warp 1: threads 32-63 handle token `blockIdx.x * 8 + 1`
So threadIdx.x doesn't directly map to token. Instead:
```cuda
int warp_id = threadIdx.x / 32;
int lane_id = threadIdx.x % 32;
int token = blockIdx.x * tokens_per_block + warp_id;
```
And each lane copies `hidden_dim / 4` elements with stride 32:
```cuda
for (int i = lane_id; i < hidden_dim / 4; i += 32) {
float4 val = reinterpret_cast<const float4*>(src)[i];
reinterpret_cast<float4*>(dst)[i] = val;
}
```
This is clean. Let me make sure about hidden_dim=128, so hidden_dim/4=32. With 32 lanes, each lane copies exactly one float4! So the loop runs once per thread. Great.
For hidden_dim=128, each warp does 32 float4 loads and 32 float4 stores, all perfectly coalesced.
For general hidden_dim that's not divisible by 128, we'd need a loop, but the problem only uses 128 according to the task description, and the inputs function uses 128. Let me keep it general but efficient.
Actually, let me look at the task requirements more carefully. The get_inputs function has hardcoded `num_tokens=2048`, `hidden_dim=128`, `num_experts=16`. The reference implementation works for arbitrary sizes though, and the evaluation might test with different sizes. So I should make it general.
Let me write a general kernel.
One more optimization: if we can guarantee that hidden_dim is a multiple of 4 (which it usually is for these kinds of models), we can use float4. To be safe, I'll handle the tail with scalar copies.
Actually, let me check if there's a way to make this even faster. Another approach is to use one thread per output element and compute the inverse mapping on the fly. But computing the inverse mapping requires more work or extra memory.
Actually, wait. We CAN build an inverse mapping array `token_for_row` of size num_rows where `token_for_row[row] = token`. We can compute this in a separate small kernel or in the same kernel. Since each token maps to exactly one row, we can launch a kernel first to fill `token_for_row`, then use it for the main copy.
But the main copy would then be: for each output element at (row, d), read `token_hidden[token_for_row[row]][d]`. If we parallelize by output rows and dimensions, reads from `token_hidden` will be scattered (not coalesced), while writes to output will be coalesced.
With the first approach (token-major, one warp per token), reads are coalesced and writes are coalesced. So it's already pretty good.
The only downside of the first approach is that the writes, although coalesced within a warp, go to different rows for different warps. But since coalescing is defined within a warp (threads executing simultaneously accessing consecutive addresses), and within each warp all 32 threads access consecutive addresses in the same output row, we have perfect coalescing.
Let me verify: In a warp, lane i accesses `output[row * hidden_dim + i * 4 + j]` for j=0..3. Actually with float4, lane i accesses addresses `output + row * hidden_dim + i * 16` to `output + row * hidden_dim + i * 16 + 15`. For consecutive i, these are consecutive 16-byte segments, which is perfectly coalesced (on H100, coalescing works at 128-byte granularity, and with 32 lanes accessing 16 bytes each in sequence, we get exactly 512 bytes which is 4 consecutive cache lines - perfect).
So the first approach is excellent.
Let me write the code. I'll use 256 threads per block (8 warps), handle 8 tokens per block.
For the Python hosting code, I'll use `load_inline` with appropriate arguments.
```python
import torch
from torch.utils.cpp_extension import load_inline
cuda_source = """
#include <cuda_runtime.h>
#include <cuda_fp16.h>
__global__ void dispatch_kernel(
const float* __restrict__ token_hidden,
const int* __restrict__ expert_idx,
const int* __restrict__ slot_idx,
const int* __restrict__ expert_offsets,
float* __restrict__ output,
int num_tokens,
int hidden_dim
) {
const int tokens_per_block = blockDim.x / 32;
const int warp_id = threadIdx.x / 32;
const int lane_id = threadIdx.x % 32;
const int token = blockIdx.x * tokens_per_block + warp_id;
if (token >= num_tokens) return;
const int expert = expert_idx[token];
const int row = expert_offsets[expert] + slot_idx[token];
const float* __restrict__ src = token_hidden + token * hidden_dim;
float* __restrict__ dst = output + row * hidden_dim;
// Vectorized copy with float4
const int num_vec = hidden_dim / 4;
for (int i = lane_id; i < num_vec; i += 32) {
float4 val = reinterpret_cast<const float4*>(src)[i];
reinterpret_cast<float4*>(dst)[i] = val;
}
// Handle remaining elements (hidden_dim % 4)
const int remainder_start = num_vec * 4;
for (int i = remainder_start + lane_id; i < hidden_dim; i += 32) {
dst[i] = src[i];
}
}
"""
def dispatch_forward(token_hidden, expert_idx, slot_idx, expert_offsets, output):
num_tokens = token_hidden.shape[0]
hidden_dim = token_hidden.shape[1]
threads_per_block = 256
tokens_per_block = threads_per_block // 32
blocks = (num_tokens + tokens_per_block - 1) // tokens_per_block
dispatch_kernel = cuda_module.dispatch_kernel
dispatch_kernel(
token_hidden, expert_idx, slot_idx, expert_offsets, output,
num_tokens, hidden_dim,
grid=(blocks, 1, 1),
block=(threads_per_block, 1, 1)
)
```
Wait, I need to be careful about how to pass the grid and block sizes with `load_inline`. For CUDA kernels launched through pybind, typically we use `torch.cuda.Stream` or just direct call. Actually, with `load_inline`, the standard way is to define a host function that launches the kernel with `<<<grid, block>>>`.
Let me define a host wrapper function:
```cpp
void dispatch_forward(
torch::Tensor token_hidden,
torch::Tensor expert_idx,
torch::Tensor slot_idx,
torch::Tensor expert_offsets,
torch::Tensor output
) {
int num_tokens = token_hidden.size(0);
int hidden_dim = token_hidden.size(1);
const int threads_per_block = 256;
const int tokens_per_block = threads_per_block / 32;
const int blocks = (num_tokens + tokens_per_block - 1) / tokens_per_block;
dispatch_kernel<<<blocks, threads_per_block>>>(
token_hidden.data_ptr<float>(),
expert_idx.data_ptr<int>(),
slot_idx.data_ptr<int>(),
expert_offsets.data_ptr<int>(),
output.data_ptr<float>(),
num_tokens,
hidden_dim
);
}
```
And then call this wrapper from Python.
Let me also consider if we need to handle different dtypes. The reference uses `token_hidden.dtype` for output, so it could be fp32 or fp16/bf16. But `get_inputs()` creates float32. The task says "Produce numerically equivalent outputs at fp32 precision" which suggests the input will be fp32.
Let me make the kernel robust and handle the exact signature.
One more detail: When using `load_inline`, I need to make sure I include `<torch/extension.h>` or similar. Let me check the exact requirements.
Actually, for `cpp_sources` in `load_inline`, we typically include:
```cpp
#include <torch/extension.h>
#include <cuda_runtime.h>
```
And the auto-generated pybind wrapper will call our host function.
Wait, the task says: "Do not write your own `PYBIND11_MODULE(...)` block in `cpp_sources`. `load_inline` auto-generates one from the `functions=[...]` argument."
So I define a C++ function and list it in `functions=[...]`.
Let me draft the complete file:
```python
import torch
import torch.nn as nn
from torch.utils.cpp_extension import load_inline
cuda_source = '''
#include <torch/extension.h>
#include <cuda_runtime.h>
__global__ void dispatch_kernel(
const float* __restrict__ token_hidden,
const int* __restrict__ expert_idx,
const int* __restrict__ slot_idx,
const int* __restrict__ expert_offsets,
float* __restrict__ output,
int num_tokens,
int hidden_dim
) {
const int tokens_per_block = blockDim.x / 32;
const int warp_id = threadIdx.x / 32;
const int lane_id = threadIdx.x % 32;
const int token = blockIdx.x * tokens_per_block + warp_id;
if (token >= num_tokens) return;
const int expert = expert_idx[token];
const int row = expert_offsets[expert] + slot_idx[token];
const float* __restrict__ src = token_hidden + token * static_cast<long>(hidden_dim);
float* __restrict__ dst = output + row * static_cast<long>(hidden_dim);
const int num_vec = hidden_dim / 4;
for (int i = lane_id; i < num_vec; i += 32) {
float4 val = reinterpret_cast<const float4*>(src)[i];
reinterpret_cast<float4*>(dst)[i] = val;
}
const int remainder_start = num_vec * 4;
for (int i = remainder_start + lane_id; i < hidden_dim; i += 32) {
dst[i] = src[i];
}
}
void dispatch_forward(
torch::Tensor token_hidden,
torch::Tensor expert_idx,
torch::Tensor slot_idx,
torch::Tensor expert_offsets,
torch::Tensor output
) {
int num_tokens = token_hidden.size(0);
int hidden_dim = token_hidden.size(1);
const int threads_per_block = 256;
const int tokens_per_block = threads_per_block / 32;
const int blocks = (num_tokens + tokens_per_block - 1) / tokens_per_block;
dispatch_kernel<<<blocks, threads_per_block>>>(
token_hidden.data_ptr<float>(),
expert_idx.data_ptr<int>(),
slot_idx.data_ptr<int>(),
expert_offsets.data_ptr<int>(),
output.data_ptr<float>(),
num_tokens,
hidden_dim
);
}
'''
cpp_module = load_inline(
name="dispatch_cuda",
cpp_sources=cuda_source,
functions=["dispatch_forward"],
extra_cuda_cflags=["-O3", "--use_fast_math"],
verbose=False,
)
class ModelNew(nn.Module):
def __init__(self):
super().__init__()
def forward(self, token_hidden, expert_idx, slot_idx, expert_offsets):
num_rows = int(expert_offsets[-1].item())
out = torch.zeros(num_rows, token_hidden.shape[1], dtype=token_hidden.dtype, device=token_hidden.device)
cpp_module.dispatch_forward(token_hidden, expert_idx, slot_idx, expert_offsets, out)
return out
```
Wait, `load_inline` might require `cuda_sources` instead of `cpp_sources` when there's CUDA code. Let me check. Actually `load_inline` has both `cuda_sources` and `cpp_sources`. For mixed host/device code, we can put CUDA code in `cpp_sources` or split them. Since our code contains a `__global__` function and host C++ code with `<<< >>>`, we need to make sure it's compiled with nvcc.
Actually, looking at the PyTorch documentation:
- `cpp_sources`: C++ source code for the module (compile with C++ compiler)
- `cuda_sources`: CUDA source code for the module (compile with nvcc)
If we put `__global__` in `cpp_sources`, it might fail because `cpp_sources` is compiled with g++ not nvcc. Let me put the CUDA code in `cuda_sources` and the host wrapper in `cpp_sources`. Or we can put everything in `cuda_sources`.
Actually, let's read the task instructions or examples. Typically for `load_inline`, you can put host code with `<<< >>>` in `cuda_sources` and C++ host wrappers in `cpp_sources`, OR you can put everything in `cuda_sources` if it's all `.cu`-compatible.
Let me just put the CUDA kernel and host launch function in `cuda_sources`, and have `cpp_sources` empty or minimal. Actually, looking at how `load_inline` works, you can pass a list to `cuda_sources`.
Let me restructure:
```python
cuda_source = r'''
#include <torch/extension.h>
#include <cuda_runtime.h>
__global__ void dispatch_kernel(...) { ... }
void dispatch_forward(...) { ... }
'''
cpp_module = load_inline(
name="dispatch_cuda",
cpp_sources="", # or just omit? Actually it's required? Let me check.
cuda_sources=[cuda_source],
functions=["dispatch_forward"],
extra_cuda_cflags=["-O3"],
)
```
Wait, I need to check if `cpp_sources` can be empty or if we need to provide it. Looking at common usage, `cpp_sources` is a string or list of strings, and `cuda_sources` is a string or list of strings. Both can be provided. Let me provide `cpp_sources=""` and put all code in `cuda_sources`.
Actually, let me be more careful. Looking at `torch.utils.cpp_extension.load_inline`:
```
load_inline(name, cpp_sources, cuda_sources=None, functions=None, extra_cflags=None, extra_cuda_cflags=None, extra_ldflags=None, extra_include_paths=None, build_directory=None, is_python_module=True, with_cuda=None, is_standalone=False)
```
So `cpp_sources` is required (positional), `cuda_sources` is optional. We can pass the code in `cpp_sources` if it contains `__global__` and `<<< >>>`? No, `cpp_sources` is compiled with the C++ compiler, so it can't contain CUDA-specific syntax. We need to put CUDA code in `cuda_sources`.
Let me try:
```python
cpp_module = load_inline(
name="dispatch_cuda",
cpp_sources="",
cuda_sources=cuda_source,
functions=["dispatch_forward"],
extra_cuda_cflags=["-O3"],
verbose=False,
)
```
I should also make sure we handle the case where `cuda_sources` expects a list or string. It accepts both.
Wait, actually there might be an issue: if `cpp_sources` is empty, does the build system still compile it correctly? Let me just include a minimal C++ header in cpp_sources:
```cpp
#include <torch/extension.h>
```
And put the actual kernel in `cuda_sources`. Actually, I think if we pass a non-empty `cuda_sources`, it will compile with nvcc.
Let me also consider: should I add `extra_cflags`? Probably not needed.
There's also the issue that `load_inline` creates a build directory and compiles the extension. On repeated calls during development, it might rebuild. For the submission, it only builds once.
Now, one more thing: can we do better than 256 blocks? With 2048 tokens and 8 tokens per block, we have 256 blocks. On H100 with ~132 SMs, that's only ~2 blocks per SM. Each block handles 8 tokens * 128 floats = small amount of work. We might benefit from more threads per token (to process more tokens in parallel within a block).
Actually, each warp handles one token. With 256 threads/block, we have 8 warps = 8 tokens per block. Total blocks = 256.
The kernel is memory bandwidth bound. The total memory traffic is:
- Read token_hidden: 2048 * 128 * 4 = 1,048,576 bytes
- Read indices: ~2048 * 3 * 4 = small
- Write output: 2048 * 128 * 4 = 1,048,576 bytes
- Total: ~2 MB
At 3.35 TB/s, the theoretical minimum for the memory operations is ~0.6 microseconds. Kernel launch overhead is typically 3-5 microseconds. So we're probably kernel-launch-overhead limited.
To reduce the impact of kernel launch overhead, we want to make sure our kernel is well-optimized and not doing anything unnecessary. But we can't really avoid the launch overhead.
Actually, one more optimization: we could copy the expert_offsets array to shared memory or registers. But with only 16 experts, reading from global memory should be cached in L1/L2.
Let me refine the kernel to use 128 threads per block (4 warps = 4 tokens). Then blocks = 512, which is better occupancy.
Or even better: 64 threads per block (2 warps = 2 tokens). Then blocks = 1024.
Or: 32 threads per block (1 warp = 1 token). Then blocks = 2048. This might actually give better occupancy because we can have more concurrent blocks per SM. But the block-level scheduling overhead might increase.
Let me think about occupancy. On H100:
- Max threads per SM: 2048
- Max blocks per SM: 32
- If block size = 32 threads, max blocks per SM = 32 (limited by blocks, not threads)
- If block size = 256 threads, max blocks per SM = 8 (limited by threads: 2048/256 = 8)
With 2048 blocks of 32 threads:
- 32 blocks per SM * 132 SMs = 4224 blocks can be resident across the GPU
- We only have 2048 blocks, so all can be resident (actually, each SM gets ~15.5 blocks)
- But each block has only 1 warp, and we have 2048 warps total. Warps per SM = ~15.5. Total concurrent warps = 15.5 * 132 = 2046. That's almost 100% occupancy!
With 256 blocks of 256 threads:
- 8 blocks per SM * 132 = 1056 blocks capacity
- We have 256 blocks, so ~1.94 blocks per SM on average
- Warps per SM = 1.94 * 8 = 15.5. Same total concurrent warps = 15.5 * 132 = 2046. Also ~100% occupancy!
So in terms of occupancy, they're equivalent. But with 2048 blocks, the block scheduler has more work, but each block is tiny. With 256 blocks, there's less scheduler overhead.
Let me keep 256 threads per block for simplicity. But I should also consider that with 256 threads per block and 256 blocks, if the grid is not evenly distributed across SMs, some SMs might be idle at the end of the kernel. With 2048 blocks, the tail effect is smaller.
For such a small workload, let's use 128 threads per block (4 warps). Blocks = 512. This is a good middle ground.
Actually, I just realized: if hidden_dim=128, and each warp handles exactly one token, then each warp copies 128 floats = 512 bytes. This is extremely small work per warp. The time to execute the warp is dominated by memory latency. Having many warps in flight helps hide latency.
Actually, the key question is: can we do better by having each thread handle multiple tokens?
What if each thread handles one token and copies all 128 elements? Then block size can be larger (e.g., 256 threads), each handling one token. Total blocks = 8. Wait, 2048 tokens / 256 threads = 8 blocks. That's very few blocks for 132 SMs! Many SMs would be idle. Bad for occupancy.
What if each thread handles one token but copies using a loop over 128/4=32 float4? That's the same as warp-per-token but with 1 thread doing all work serially. The memory accesses within a warp would then be to the same addresses for all 32 threads (since each thread handles a different token but accesses the same offsets within its token). For loads: 32 threads in a warp load from `token_hidden[token_i * 128 + j * 4]` for i=0..31. These are 128 bytes apart, so they coalesce into 32 separate 16-byte accesses but spread across 128*4=512 bytes. This is still coalesced (the hardware combines accesses from a warp into cache lines). Actually, this is strided access, which might still coalesce well on modern hardware because the addresses fall into exactly 4 cache lines (512 bytes / 128 bytes per line = 4 lines). The hardware can issue up to 4 sectors per cycle, so this is okay. But stores to output would be scattered similarly.
Actually, with 1 thread per token doing all 32 float4 copies, the accesses within the warp are strided by 512 bytes for the same vector element j. For j=0, addresses are `token_i * 512` which are NOT consecutive but fall into different cache lines. This is actually bad for coalescing because consecutive threads access addresses 512 bytes apart. The warp would need 32 sectors = 4 KB, which requires 32 memory transactions.
With the warp-per-token approach, for a given j, all 32 threads access consecutive 16 bytes, falling into 4 cache lines (512 bytes total), requiring only 4 memory transactions. That's 8x better!
So warp-per-token (or block-per-token, or having consecutive threads handle consecutive elements of the same token) is much better for coalescing.
Let me formalize:
- Option 1: Thread t handles token t, copies all hidden_dim elements. Within warp, thread t accesses `src[t * hidden_dim + j]` and `dst[row_t * hidden_dim + j]`. Stride = hidden_dim. Coalescing depends on hidden_dim.
- Option 2: Warp w handles token w, lane l handles hidden_dim elements starting at l*4 in steps of 32*4. Within warp, lane l accesses `src[w * hidden_dim + l * 4]` and `dst[row_w * hidden_dim + l * 4]`. Consecutive lanes access consecutive addresses. Perfect coalescing.
Option 2 is clearly superior.
Now, I can also do a hybrid: have multiple warps in a block, but also allow the grid to be large enough. Block size of 128 (4 warps, 4 tokens) with 512 blocks seems reasonable.
Or I can use dynamic grouping: have a block size of 256 threads, where each 32-thread warp handles one token. That's 8 tokens per block, 256 blocks.
Let me actually benchmark both in my head. With 512 blocks of 128 threads, we have more blocks to distribute. With 256 blocks of 256 threads, fewer blocks but more threads each.
For 2048 tokens:
- 128 threads/block => 16 blocks? No, 2048/4 = 512 blocks of 4 tokens each. Wait: block has 128 threads = 4 warps, each warp handles 1 token. So 4 tokens per block. Blocks = 2048/4 = 512.
- 256 threads/block => 8 tokens per block. Blocks = 256.
With 512 blocks, on 132 SMs, each SM gets ~3.88 blocks on average. With 4 warps per block, that's ~15.5 warps per SM. Good.
With 256 blocks, on 132 SMs, each SM gets ~1.94 blocks. With 8 warps per block, that's also ~15.5 warps per SM. Good.
Both should be fine. I'll go with 256 threads per block to reduce block scheduling overhead slightly.
One more tweak: for very small hidden dimensions, the loop might not be efficient. But hidden_dim=128 is large enough.
Actually wait - looking at the loop again:
```cuda
for (int i = lane_id; i < num_vec; i += 32) {
float4 val = reinterpret_cast<const float4*>(src)[i];
reinterpret_cast<float4*>(dst)[i] = val;
}
```
For hidden_dim=128, num_vec=32. lane_id goes from 0 to 31. Each lane executes exactly one iteration. The loop overhead is minimal.
But actually, since we know hidden_dim is 128 at compile time... well we don't know it statically. But we can add a template or special path. Eh, unnecessary micro-optimization.
Let me write the final code and test it.
Actually, let me think about another possible approach. The PyTorch `index_copy` or `scatter` functions might be applicable here. We could do:
```python
rows = expert_offsets[expert_idx] + slot_idx
out.index_copy_(0, rows, token_hidden)
```
But `index_copy_` requires sorted indices for deterministic behavior? And the reference implementation scatters in a specific order. For this specific input where each row gets exactly one token, `index_copy_` with `rows` would work. But `index_copy_` is a PyTorch operation, and while it might be fast, we want a custom CUDA kernel.
Actually, for this case where the mapping is a permutation (each output row gets exactly one token), `torch.index_select` with the inverse would also work. But constructing the inverse permutation costs extra.
I'll stick with the custom kernel. It should be much faster than a Python loop.
Now, let me also think about `extern "C"` or other issues with `load_inline`. The auto-generated pybind module will look for the function name as specified in `functions=[...]`. So `dispatch_forward` should be discoverable. It should be fine without `extern "C"`.
Let me also consider if we need to check for CUDA availability, but the task requires replacing with a custom CUDA kernel, so we can assume CUDA is available.
One subtlety: `expert_offsets[-1]` in PyTorch gets the last element. In the CUDA kernel, `expert_offsets` is passed as a pointer, and we only read `expert_offsets[expert]` where `expert` is in `[0, num_experts-1]`. In the Python code, `num_rows = int(expert_offsets[-1].item())` is correct. We should preserve this.
Wait, `expert_offsets` has length `num_experts + 1`. The reference does `expert_offsets[-1]` which is `expert_offsets[num_experts]`. But in the kernel, we read `expert_offsets[expert]` where `expert` is in `[0, num_experts-1]`. This is correct because `expert_idx` values are in `[0, num_experts-1]`.
Let me construct the full kernel code now.
One more thing: when using `reinterpret_cast<const float4*>(src)[i]`, we need to make sure `src` is 16-byte aligned. For PyTorch tensors, the data pointer is typically aligned to at least 256 bytes. But `src = token_hidden + token * hidden_dim` might not be 16-byte aligned if `token * hidden_dim` is not divisible by 4. Wait, `token * 128` is always divisible by 128, and since 128 is divisible by 4, `token * 128` is divisible by 4. So for float (4-byte), `src` is at least 4-byte aligned, and for float4 (16-byte), it's 16-byte aligned because 128 * token / 4 = 32 * token, which is an integer. But the BYTE address is `token * 128 * 4 = token * 512`, which is divisible by 16 for any integer token. Great!
Similarly, `dst = output + row * hidden_dim` where `row` is arbitrary. Wait, is `row * hidden_dim` always divisible by 4 in terms of float index? Yes, `row * 128` is divisible by 4. So byte address `row * 128 * 4` is divisible by 16. So float4 accesses are always aligned for both src and dst. Perfect.
Now let me code this up and test.
Wait, I should double check: the task says "Your submission is a complete Python file that defines a class called `ModelNew`." So the file should start with imports and define the class.
Let me also consider: we need the module-level compilation to run before the class is defined. With `load_inline`, this happens at module import time, which is correct.
But `load_inline` might take time. And in evaluation, `exec()` is called on the file. So the compilation will happen during `exec()`. That's expected.
Let me think about whether we should put the load_inline call inside the class or at module level. Module level is standard and fine.
One concern: `load_inline` creates a temporary build directory. If we run multiple times, it might rebuild or reuse. But for evaluation, it runs once.
Let me write the code carefully.
Actually, I just realized: the reference uses `for token in range(token_hidden.shape[0]):` and does `out[row] = token_hidden[token]`. This is a full row copy. In PyTorch, this might use vectorized operations under the hood for each row copy. But with 2048 iterations in Python, it's very slow due to Python overhead and lack of parallelism. Our CUDA kernel should be easily 100x+ faster.
Let me write the code and compile it.
```python
import torch
import torch.nn as nn
from torch.utils.cpp_extension import load_inline
# CUDA source code
cuda_source = r"""
#include <torch/extension.h>
#include <cuda_runtime.h>
__global__ void dispatch_kernel(
const float* __restrict__ token_hidden,
const int* __restrict__ expert_idx,
const int* __restrict__ slot_idx,
const int* __restrict__ expert_offsets,
float* __restrict__ output,
int num_tokens,
int hidden_dim
) {
const int tokens_per_block = blockDim.x / 32;
const int warp_id = threadIdx.x / 32;
const int lane_id = threadIdx.x % 32;
const int token = blockIdx.x * tokens_per_block + warp_id;
if (token >= num_tokens) return;
const int expert = expert_idx[token];
const int row = expert_offsets[expert] + slot_idx[token];
const float* __restrict__ src = token_hidden + static_cast<long>(token) * hidden_dim;
float* __restrict__ dst = output + static_cast<long>(row) * hidden_dim;
const int num_vec = hidden_dim / 4;
for (int i = lane_id; i < num_vec; i += 32) {
float4 val = reinterpret_cast<const float4*>(src)[i];
reinterpret_cast<float4*>(dst)[i] = val;
}
const int remainder_start = num_vec * 4;
for (int i = remainder_start + lane_id; i < hidden_dim; i += 32) {
dst[i] = src[i];
}
}
void dispatch_forward(
torch::Tensor token_hidden,
torch::Tensor expert_idx,
torch::Tensor slot_idx,
torch::Tensor expert_offsets,
torch::Tensor output
) {
int num_tokens = token_hidden.size(0);
int hidden_dim = token_hidden.size(1);
const int threads_per_block = 256;
const int tokens_per_block = threads_per_block / 32;
const int blocks = (num_tokens + tokens_per_block - 1) / tokens_per_block;
dispatch_kernel<<<blocks, threads_per_block>>>(
token_hidden.data_ptr<float>(),
expert_idx.data_ptr<int>(),
slot_idx.data_ptr<int>(),
expert_offsets.data_ptr<int>(),
output.data_ptr<float>(),
num_tokens,
hidden_dim
);
}
"""
# Compile the CUDA extension
_dispatch_cuda = load_inline(
name="dispatch_cuda",
cpp_sources="",
cuda_sources=[cuda_source],
functions=["dispatch_forward"],
extra_cuda_cflags=["-O3"],
verbose=False,
)
class ModelNew(nn.Module):
"""
Pack tokens into expert-major order for MoE execution.
"""
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,
)
_dispatch_cuda.dispatch_forward(
token_hidden, expert_idx, slot_idx, expert_offsets, out
)
return out
```
Wait, I need to be careful with `hidden_dim / 4` when hidden_dim is an int. In C++, `hidden_dim / 4` with ints does integer division, which is fine.
One potential issue: if hidden_dim is not divisible by 4, `num_vec = hidden_dim / 4` truncates, and the remainder handles the rest.
Another issue: I'm using `reinterpret_cast<const float4*>(src)[i]`. The pointer arithmetic on `float4*` advances by 16 bytes. For `src = token_hidden + token * hidden_dim`, the index `i` accesses `src + i * 4` in float units, which is `token_hidden + token * hidden_dim + i * 4`. This is correct.
Wait, `reinterpret_cast<const float4*>(src)` gives a `const float4*` pointing to `src`. Then `[i]` accesses `src + i * sizeof(float4) = src + i * 16 bytes = src + i * 4 floats`. So `reinterpret_cast<const float4*>(src)[i]` gives the float4 at float indices `i*4, i*4+1, i*4+2, i*4+3` relative to `src`. This is correct.
For the float4 store, same logic applies.
I think this should compile and work correctly. Let me test it.
Actually, I want to double-check one more thing. The task mentions `expert_offsets[-1]` in the reference. But `expert_offsets` is a `torch.Tensor` of dtype `int32` in the inputs. When we do `expert_offsets[-1].item()` in Python, it should work fine.
One more thought: the reference does `expert = int(expert_idx[token].item())`. Since `expert_idx` is int32, `item()` returns a Python int. Our CUDA kernel reads `expert_idx[token]` directly as `int`. This should match.
Let me verify with a small example mentally:
- token_hidden = [[1,2,3,4], [5,6,7,8]]
- expert_idx = [0, 1]
- slot_idx = [0, 0]
- expert_offsets = [0, 1, 2] (2 experts, num_rows = expert_offsets[-1] = 2)
Token 0: expert=0, row=expert_offsets[0]+slot_idx[0]=0+0=0. out[0] = [1,2,3,4]
Token 1: expert=1, row=expert_offsets[1]+slot_idx[1]=1+0=1. out[1] = [5,6,7,8]
Output = [[1,2,3,4], [5,6,7,8]]. Correct.
Another check: tokens out of order.
- token_hidden = [[1,2], [3,4], [5,6]]
- expert_idx = [1, 0, 1]
- slot_idx = [0, 0, 1]
- expert_offsets = [0, 1, 3] (expert 0 has 1 token, expert 1 has 2 tokens)
Token 0: expert=1, row=expert_offsets[1]+0=1+0=1. out[1]=[1,2]
Token 1: expert=0, row=expert_offsets[0]+0=0+0=0. out[0]=[3,4]
Token 2: expert=1, row=expert_offsets[1]+1=1+1=2. out[2]=[5,6]
Output = [[3,4], [1,2], [5,6]]. Correct.
Our kernel does the same computation. Good.
I think the code is ready. Let me compile it.
Actually, one optimization: since expert_offsets is small (num_experts+1 <= 17 elements for the given problem), we could load it into shared memory or registers. But it's so small that all accesses will hit L1 cache anyway. Not worth it.
Another thing: should I add `__ldg()` for the indices? `__ldg` is a read-only cache load that can help if the data is read-only and not in cache. For the indices, they are read-only, but with only 2048 tokens * 3 arrays = 6144 integers, they easily fit in cache. Not worth it.
What about using `__restrict__`? Already doing that.
Let me run the compile.