reasoning
Let me analyze the reference implementation:
1. It computes `grounded = router_logits + alpha * (token_hidden @ expert_ground.t())`
2. Takes top-2 values and indices from grounded
3. Applies softmax to the top-2 values
4. Returns a stacked tensor of shape (num_tokens, 2, 2) containing top_idx (as float32) and top_weights
The shapes are:
- token_hidden: (num_tokens, hidden_dim) = (4096, 128)
- router_logits: (num_tokens, num_experts) = (4096, 16)
- expert_ground: (num_experts, hidden_dim) = (16, 128)
- grounded: (4096, 16)
The output is (4096, 2, 2) where:
- output[:, :, 0] = top indices (as float32)
- output[:, :, 1] = softmax weights
Let me write a custom CUDA kernel for this. The main operations are:
1. Matrix multiplication: token_hidden @ expert_ground.t() -> (4096, 16)
2. Scale by alpha and add to router_logits
3. Top-2 finding for each row
4. Softmax over 2 values
5. Stack and output
I'll need to be careful with the top-k and softmax operations. Let me write a kernel that handles all this efficiently.
Key optimizations to consider:
1. The matrix multiplication can be done with shared memory tiling
2. Top-2 can be computed in a single pass through the grounded values
3. Softmax over 2 values is trivial (just normalize 2 values)
Let me start with a basic implementation:
```python
import torch
from torch.utils.cpp_extension import load_inline
cuda_source = """
#include <torch/extension.h>
#include <cuda.h>
#include <cuda_runtime.h>
#include <cfloat>
__global__ void grounded_moe_router_kernel(
const float* __restrict__ token_hidden,
const float* __restrict__ router_logits,
const float* __restrict__ expert_ground,
float* __restrict__ output,
const float alpha,
const int num_tokens,
const int hidden_dim,
const int num_experts
) {
// Each thread handles one token
int token_idx = blockIdx.x * blockDim.x + threadIdx.x;
if (token_idx >= num_tokens) return;
// Compute grounded logits for this token
// grounded[i] = router_logits[token_idx, i] + alpha * sum_j(token_hidden[token_idx, j] * expert_ground[i, j])
// Shared memory for expert_ground (transpose for better access)
extern __shared__ float shared_mem[];
float* grounded = shared_mem; // num_experts values per thread... but this might be too much
// Actually, let me reconsider. Each thread computes grounded values for its token
// grounded has num_experts values
float grounded_vals[16]; // Assuming num_experts <= 16
for (int e = 0; e < num_experts; e++) {
float dot = 0.0f;
for (int h = 0; h < hidden_dim; h++) {
dot += token_hidden[token_idx * hidden_dim + h] * expert_ground[e * hidden_dim + h];
}
grounded_vals[e] = router_logits[token_idx * num_experts + e] + alpha * dot;
}
// Find top-2
float top1_val = -FLT_MAX;
float top2_val = -FLT_MAX;
int top1_idx = -1;
int top2_idx = -1;
for (int e = 0; e < num_experts; e++) {
float val = grounded_vals[e];
if (val > top1_val) {
top2_val = top1_val;
top2_idx = top1_idx;
top1_val = val;
top1_idx = e;
} else if (val > top2_val) {
top2_val = val;
top2_idx = e;
}
}
// Softmax over top-2
float max_val = fmaxf(top1_val, top2_val);
float exp1 = expf(top1_val - max_val);
float exp2 = expf(top2_val - max_val);
float sum_exp = exp1 + exp2;
float weight1 = exp1 / sum_exp;
float weight2 = exp2 / sum_exp;
// Output: (num_tokens, 2, 2)
// output[token_idx, 0, 0] = top1_idx
// output[token_idx, 0, 1] = weight1
// output[token_idx, 1, 0] = top2_idx
// output[token_idx, 1, 1] = weight2
output[token_idx * 4 + 0] = (float)top1_idx;
output[token_idx * 4 + 1] = weight1;
output[token_idx * 4 + 2] = (float)top2_idx;
output[token_idx * 4 + 3] = weight2;
}
torch::Tensor grounded_moe_router_cuda(
torch::Tensor token_hidden,
torch::Tensor router_logits,
torch::Tensor expert_ground,
float alpha
) {
const int num_tokens = token_hidden.size(0);
const int hidden_dim = token_hidden.size(1);
const int num_experts = router_logits.size(1);
auto output = torch::empty({num_tokens, 2, 2}, torch::TensorOptions().dtype(torch::kFloat32).device(token_hidden.device()));
const int block_size = 256;
const int num_blocks = (num_tokens + block_size - 1) / block_size;
grounded_moe_router_kernel<<<num_blocks, block_size>>>(
token_hidden.data_ptr<float>(),
router_logits.data_ptr<float>(),
expert_ground.data_ptr<float>(),
output.data_ptr<float>(),
alpha,
num_tokens,
hidden_dim,
num_experts
);
return output;
}
"""
cpp_source = """
torch::Tensor grounded_moe_router_cuda(torch::Tensor token_hidden, torch::Tensor router_logits, torch::Tensor expert_ground, float alpha);
"""
module = load_inline(
name='grounded_moe_router',
cpp_sources=cpp_source,
cuda_sources=cuda_source,
functions=['grounded_moe_router_cuda'],
extra_cuda_cflags=['-O3'],
verbose=True
)
class ModelNew(nn.Module):
def __init__(self):
super().__init__()
self.kernel = module
def forward(self, token_hidden, router_logits, expert_ground, alpha):
return self.kernel.grounded_moe_router_cuda(token_hidden, router_logits, expert_ground, alpha)
```
Wait, I need to think about this more carefully. The dot product computation involves iterating over hidden_dim (128), which could be slow. I should consider using shared memory to load expert_ground in a tiled manner.
Let me also think about a different approach - using a warp-level cooperative approach where multiple threads work together on each token, or use one thread per expert.
Actually, with num_experts = 16 and hidden_dim = 128, the matrix multiplication is small enough that we might benefit from a different approach. Let me consider:
1. Using shared memory to cache expert_ground
2. Using one thread per token, but with vectorized loads
3. Using cooperative groups
Let me first try a simple implementation and see if it works, then optimize.
Actually, I realize I should make num_experts a template parameter or pass it properly. Let me refine:
```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>
#include <cfloat>
__global__ void grounded_moe_router_kernel(
const float* __restrict__ token_hidden,
const float* __restrict__ router_logits,
const float* __restrict__ expert_ground,
float* __restrict__ output,
const float alpha,
const int num_tokens,
const int hidden_dim,
const int num_experts
) {
// Each thread handles one token
int token_idx = blockIdx.x * blockDim.x + threadIdx.x;
if (token_idx >= num_tokens) return;
// Pointer to this token's hidden state
const float* __restrict__ hidden = token_hidden + token_idx * hidden_dim;
const float* __restrict__ logits = router_logits + token_idx * num_experts;
// Find top-2 while computing grounded values
float top1_val = -FLT_MAX;
float top2_val = -FLT_MAX;
int top1_idx = -1;
int top2_idx = -1;
for (int e = 0; e < num_experts; e++) {
// Compute dot product with expert e's ground embedding
float dot = 0.0f;
const float* __restrict__ ground = expert_ground + e * hidden_dim;
#pragma unroll 4
for (int h = 0; h < hidden_dim; h++) {
dot += hidden[h] * ground[h];
}
float grounded_val = logits[e] + alpha * dot;
if (grounded_val > top1_val) {
top2_val = top1_val;
top2_idx = top1_idx;
top1_val = grounded_val;
top1_idx = e;
} else if (grounded_val > top2_val) {
top2_val = grounded_val;
top2_idx = e;
}
}
// Softmax over top-2
float max_val = fmaxf(top1_val, top2_val);
float exp1 = expf(top1_val - max_val);
float exp2 = expf(top2_val - max_val);
float sum_exp = exp1 + exp2;
float weight1 = exp1 / sum_exp;
float weight2 = exp2 / sum_exp;
// Output layout matches torch.stack((top_idx.to(torch.float32), top_weights), dim=-1)
// top_idx shape: (num_tokens, 2), top_weights shape: (num_tokens, 2)
// stacked shape: (num_tokens, 2, 2)
// output[token_idx, i, 0] = top_idx[token_idx, i] (expert index)
// output[token_idx, i, 1] = top_weights[token_idx, i] (weight)
float* __restrict__ out = output + token_idx * 4;
out[0] = (float)top1_idx; // output[token_idx, 0, 0]
out[1] = weight1; // output[token_idx, 0, 1]
out[2] = (float)top2_idx; // output[token_idx, 1, 0]
out[3] = weight2; // output[token_idx, 1, 1]
}
torch::Tensor grounded_moe_router_cuda(
torch::Tensor token_hidden,
torch::Tensor router_logits,
torch::Tensor expert_ground,
float alpha
) {
const int num_tokens = token_hidden.size(0);
const int hidden_dim = token_hidden.size(1);
const int num_experts = router_logits.size(1);
auto output = torch::empty({num_tokens, 2, 2}, torch::TensorOptions().dtype(torch::kFloat32).device(token_hidden.device()));
const int block_size = 256;
const int num_blocks = (num_tokens + block_size - 1) / block_size;
grounded_moe_router_kernel<<<num_blocks, block_size>>>(
token_hidden.data_ptr<float>(),
router_logits.data_ptr<float>(),
expert_ground.data_ptr<float>(),
output.data_ptr<float>(),
alpha,
num_tokens,
hidden_dim,
num_experts
);
return output;
}
"""
cpp_source = """
torch::Tensor grounded_moe_router_cuda(torch::Tensor token_hidden, torch::Tensor router_logits, torch::Tensor expert_ground, float alpha);
"""
module = load_inline(
name='grounded_moe_router',
cpp_sources=cpp_source,
cuda_sources=cuda_source,
functions=['grounded_moe_router_cuda'],
extra_cuda_cflags=['-O3'],
verbose=True
)
class ModelNew(nn.Module):
def __init__(self):
super().__init__()
def forward(self, token_hidden, router_logits, expert_ground, alpha):
return module.grounded_moe_router_cuda(token_hidden, router_logits, expert_ground, alpha)
```
Wait, I think I need to handle the case where the module is loaded outside the class. Let me restructure.
Actually, looking at the structure more carefully, I think I need to ensure the CUDA extension is loaded at module level, and then the ModelNew class uses it. Let me rewrite: