reasoning
Let me analyze the problem. The outer-product mean operation involves:
1. LayerNorm on msa_repr (B, S, N, msa_dim)
2. Linear projection to hidden_dim for both a and b
3. Apply mask
4. Einsum: outer = torch.einsum("bsih,bsjk->bijhk", a, b)
- This computes outer[i,j] = sum_s a[s,i,:] * b[s,j,:]^T for each batch
- Result shape: (B, N, N, H, H)
5. Divide by sum of mask over S dimension
6. Reshape to (B, N, N, H*H)
7. Output projection to pair_dim
The key challenge is the einsum which computes:
outer[b,i,j,h,k] = sum_s (a[b,s,i,h] * b[b,s,j,k])
This is essentially computing for each pair (i,j) the average over S of the outer product of a[:,:,i,:] and b[:,:,j,:].
Let me think about the computation:
- For each (b,i,j) pair, we compute a matrix of size H x H
- Each element [h,k] = sum_s (a[b,s,i,h] * b[b,s,j,k]) / denom[b]
The naive approach would be to:
1. Fuse layer_norm + proj_a + proj_b + masking into one kernel
2. Then compute the outer product mean
3. Then the output projection
But let me think more carefully. The intermediate tensors a and b are of shape (B, S, N, H) with B=2, S=8, N=32, H=32. That's 2*8*32*32 = 16,384 elements, which is small.
The output shape is (B, N, N, H*H) = (2, 32, 32, 1024) before the output projection.
Actually, let me think about fusing the entire operation:
The computation flow:
1. layer_norm(msa_repr) -> normalized (B, S, N, msa_dim)
2. proj_a(normalized) -> a (B, S, N, H)
3. proj_b(normalized) -> b (B, S, N, H)
4. Apply mask to a and b
5. outer = einsum("bsih,bsjk->bijhk", a, b) -> (B, N, N, H, H)
6. denom = sum(mask, dim=1) -> (B, N), then reshape and clamp
7. outer / denom -> (B, N, N, H, H)
8. reshape to (B, N, N, H*H)
9. out_proj(outer) -> (B, N, N, pair_dim)
For a custom kernel, I could:
1. Compute normalized values on-the-fly using fused LN + linear
2. Compute the outer product mean directly
3. Apply the output projection
But the LN requires per-position mean/var which complicates things.
Let me think about a simpler approach:
- Use PyTorch for LayerNorm (it's efficient)
- Custom kernel for proj + outer product mean + output proj
Actually, let me try a fully fused approach. The sizes are small enough that we might be able to keep things in shared memory.
Let me think about the tiling strategy:
For each batch element, we need to compute:
- For each position pair (i, j), a H x H matrix that is the average outer product
One approach:
- Load a block of positions i (say 8 positions)
- Load a block of positions j (say 8 positions)
- For each position, we need the S vectors of size H
- Compute the outer products and accumulate
Given the sizes (N=32, S=8, H=32), we have:
- For each position: 8 * 32 = 256 floats for a, same for b
- The entire a tensor per batch: 8 * 32 * 32 = 8192 floats
Actually, since S=8, N=32, H=32, the a tensor for one batch is just 8*32*32 = 8192 floats = 32KB. This easily fits in shared memory!
So a possible strategy:
1. Each thread block handles one batch element
2. Load all of a and b for that batch into shared memory (after projection)
3. Each thread computes one output element (or a small number)
Let me design the kernel:
Block configuration:
- One block per batch (so 2 blocks)
- Each block has N*N*H*H = 32*32*32*32 = 1,048,576 output elements per batch
Wait, that's a lot of threads. We need more blocks.
Let me reconsider:
- Output per batch: N*N*H*H = 32*32*32*32 = ~1M elements
- With 1024 threads per block and each thread computing 1 element, we need ~1024 blocks per batch
Or each thread computes multiple elements.
Actually, let's think differently. The output after reshape is (B, N, N, H*H) = (B, N, N, 1024). The output projection takes this to (B, N, N, pair_dim=64).
For the outer product computation:
- Each thread can compute one (i,j) pair's H*H matrix
- Or each thread computes one element of the H*H matrix for one (i,j) pair
Let me try: each thread computes one (i,j,h,k) element, which is the average over S of a[s,i,h]*b[s,j,k].
Thread organization:
- blockIdx.x determines batch batch_idx = blockIdx.x / num_blocks_per_batch, and local ij block
- threadIdx.x determines (i,j,h,k) within the block
Actually, let me simplify:
- Launch B*N*N blocks (one per (batch, i, j) pair)
- Each block computes one H*H matrix
- Each thread computes H*H/num_threads elements
With B=2, N=32, we have 2*32*32 = 2048 blocks. Each block has 1024 elements to compute.
Let me try this approach with each thread computing one output element.
Actually, I realize a simpler approach might work:
- Each (i,j) pair needs access to a[:,i,:] and b[:,j,:] (vectors of shape S x H)
- These are small (S=8, H=32), so 256 floats each
- We can load these into shared memory
Let me design:
- Block per (batch, i, j) pair
- threadIdx.x iterates over h
- threadIdx.y iterates over k
- Each thread computes the average over s
But with H=32, we'd have 32*32 = 1024 threads, which is reasonable.
Hmm, but with 2048 blocks and 1024 threads each, we have 2M threads. That might be too many for such small data.
Let me reconsider the block organization:
- Block per (batch, i_block, j_block) where i_block contains, say, 4 i positions
- Each thread handles one (i_local, j_local, h, k) combination
Actually, let's be more practical. The problem sizes are small:
- B=2, S=8, N=32, msa_dim=64, hidden_dim=32, pair_dim=64
Let me just write a straightforward kernel and see if it works, then optimize.
Let me fuse the layer norm and projections first, then compute the outer product mean, then the output projection.
For layer norm, I need to compute:
mean = mean(x) over the last dimension
var = var(x) over the last dimension
normalized = (x - mean) / sqrt(var + eps)
output = weight * normalized + bias
For the projection:
output = input @ weight.T + bias
I can fuse these:
After LN, we have normalized values. Then proj_a = normalized @ weight_a.T + bias_a
Let me write a kernel that:
1. Computes the fused LN + projection for all (b, s, n) positions
2. Stores a and b in global memory
3. Then a second kernel computes the outer product mean and output projection
Or maybe I can do everything in one kernel using shared memory?
Given the small sizes, let me try:
1. Kernel 1: Fused LN + proj_a + proj_b with mask, storing a and b
2. Kernel 2: Outer product mean + output projection
Let me start implementing.
Actually, looking at the dimensions again:
- msa_repr: (2, 8, 32, 64)
- msa_dim = 64, hidden_dim = 32, pair_dim = 64
- Output: (2, 32, 32, 64)
The outer product mean gives (2, 32, 32, 32, 32) which is 2*32*32*32*32 = 2M floats = 8MB.
After reshape: (2, 32, 32, 1024) = 2M floats.
After output projection: (2, 32, 32, 64) = 128K floats.
So the intermediate outer product tensor is quite large. Let me try to fuse the outer product computation with the output projection to avoid storing the full tensor.
For the output projection, each output element is:
out[b,i,j,d] = sum_hk outer[b,i,j,h,k] * out_proj_weight[d, h*H+k] + out_proj_bias[d]
Where outer[b,i,j,h,k] = sum_s a[b,s,i,h] * b[b,s,j,k] / denom[b]
Substituting:
out[b,i,j,d] = sum_hk (sum_s a[b,s,i,h] * b[b,s,j,k] / denom[b]) * out_proj_weight[d, h*H+k] + out_proj_bias[d]
= sum_s sum_hk a[b,s,i,h] * b[b,s,j,k] * out_proj_weight[d, h*H+k] / denom[b] + out_proj_bias[d]
This suggests we can compute the output directly without materializing the full outer product:
out[b,i,j,d] = sum_s (sum_h a[b,s,i,h] * sum_k b[b,s,j,k] * out_proj_weight[d, h*H+k]) / denom[b] + out_proj_bias[d]
Hmm, this doesn't factor nicely. Let me think again...
Actually, we can write:
out[b,i,j,d] = sum_s sum_hk a[b,s,i,h] * b[b,s,j,k] * W[d,hk] / denom[b] + bias[d]
For fixed (b,i,j,d), this is a sum over s, h, k. We can compute this directly.
Let me define:
- For each (b,i,j), compute out[b,i,j,:] directly
- This requires knowing a[b,:,i,:] and b[b,:,j,:] (shape S x H each)
So the fused approach:
1. Compute a and b (fused LN + proj)
2. For each (b,i,j) pair, compute the output projection directly using the formula above
Let me write this kernel:
For the fused LN + proj:
- Each thread handles one (s, n) position within a batch
- Compute LN, then proj_a and proj_b, apply mask
- Store a and b
For the outer + output:
- Block per (b, i, j)
- Load a[b,:,i,:] and b[b,:,j,:] into shared memory
- Each thread computes one output element d
- Use reduction to sum over s
Let me start coding.
```python
import torch
import torch.nn as nn
import torch.utils.cpp_extension as cpp_extension
class ModelNew(nn.Module):
def __init__(self, msa_dim, pair_dim, hidden_dim=32):
super().__init__()
# Store dimensions
self.msa_dim = msa_dim
self.pair_dim = pair_dim
self.hidden_dim = hidden_dim
# Create reference model to get weights
ref = Model(msa_dim, pair_dim, hidden_dim)
self.register_buffer('ln_weight', ref.layer_norm.weight)
self.register_buffer('ln_bias', ref.layer_norm.bias)
self.register_buffer('proj_a_weight', ref.proj_a.weight)
self.register_buffer('proj_a_bias', ref.proj_a.bias)
self.register_buffer('proj_b_weight', ref.proj_b.weight)
self.register_buffer('proj_b_bias', ref.proj_b.bias)
self.register_buffer('out_proj_weight', ref.out_proj.weight)
self.register_buffer('out_proj_bias', ref.out_proj.bias)
# Compile CUDA kernel
self.kernel = self._compile_kernel()
def _compile_kernel(self):
# CUDA source code
...
def forward(self, msa_repr, msa_mask):
...
```
Let me write the actual CUDA kernel code.
Actually, wait. The Model class has learnable parameters. I need to make sure I'm using the same parameters in my custom kernel. Let me initialize the weights from a reference Model instance.
Actually, looking at the problem setup more carefully:
- The reference model has weights that are initialized when Model is created
- I need to use the same weights in my ModelNew
The `get_init_inputs()` returns [msa_dim, pair_dim, hidden_dim], which are passed to `__init__`.
So I should:
1. In `__init__`, create parameters just like the reference model
2. Use these parameters in my CUDA kernel
Let me write the code now.
For the kernel design, let me think about the optimal approach:
Option 1: Separate kernels
- Kernel 1: LayerNorm + proj_a + proj_b (write a, b to global memory)
- Kernel 2: Outer product mean + output projection (read a, b, write output)
Option 2: Fused kernel
- One kernel that does everything
- Requires storing intermediates in shared memory
Given the small problem sizes, Option 1 might be simpler and still efficient.
Let me implement Option 1 first.
Kernel 1: Fused LN + proj
- Grid: (B, S, N) blocks? Or (B * S * N) blocks?
- Each block handles one position (b, s, n)
- Thread handles part of the hidden dimension
Actually, for LN + proj:
- LN requires computing mean and variance over msa_dim (64) elements
- Then proj_a computes hidden_dim (32) outputs
- And proj_b computes hidden_dim (32) outputs
With 64 input elements and 32 output elements for each of a and b, we can use:
- 64 threads per block (one per input element for LN)
- Then each thread computes part of the projection
Or:
- Each thread computes one output element for both proj_a and proj_b
- 32 threads per block
- First compute mean/var cooperatively, then each thread computes its output
Let me try the second approach with 32 threads (hidden_dim) per block.
Actually, let me reconsider. For a single (b,s,n) position:
1. Load msa_dim (64) values
2. Compute mean (1 value)
3. Compute variance (1 value)
4. Normalize each of 64 values
5. Compute proj_a: 32 outputs, each is dot product of 64 normalized values with 64 weights
6. Compute proj_b: 32 outputs, similar
This is a lot of work for one position. With 32 threads:
- Each thread computes one output for proj_a and one for proj_b
- All threads cooperate to compute mean and variance
- Then each thread does a dot product of 64 values for its output
This seems reasonable. Let me implement this.
For the shared memory:
- Need to store the 64 input values
- And the 32+32 output values
Actually, let me try a different approach:
- Each thread block handles multiple positions
- This amortizes the overhead of mean/var computation
Hmm, but each position has different mean/var for LN. So they can't share.
Let me just implement the straightforward version:
- Block per position (b, s, n)
- 32 threads per block (hidden_dim)
- Compute LN cooperatively, then proj outputs
Let me code this up.
For the outer product mean + output projection kernel:
For each (b, i, j):
- Need a[b, :, i, :] (S x H) and b[b, :, j, :] (S x H)
- Compute denom = sum_s mask[b, s, i] (assuming mask is per-position, which it is)
Actually, looking at the original code:
```python
denom = msa_mask.sum(dim=1).unsqueeze(-1).unsqueeze(-1).unsqueeze(-1).clamp(min=1)
outer = outer / denom # mean over S
```
msa_mask is (B, S, N). sum over S gives (B, N). Then unsqueeze to (B, 1, 1, 1) and broadcast.
Wait, let me trace this more carefully:
```python
outer = torch.einsum("bsih,bsjk->bijhk", a, b) # (B, N, N, H, H)
denom = msa_mask.sum(dim=1) # (B, N)
```
So denom is per (batch, position), not per (batch, i, j) pair. Let me check the einsum:
- a: (B, S, N, H) -> bsih
- b: (B, S, N, H) -> bsjk
- Result: bijhk
The indices: b, s, i, h for a and b, s, j, k for b.
The summation is over s (the MSA depth).
Result indices are b, i, j, h, k.
So outer[b,i,j,h,k] = sum_s a[b,s,i,h] * b[b,s,j,k]
But the mask application:
```python
mask = msa_mask.unsqueeze(-1) # (B, S, N, 1)
a = a * mask
b = b * mask
```
So a and b are zeroed out where mask is 0.
The denom:
```python
denom = msa_mask.sum(dim=1) # (B, N) - sum over S
denom = denom.unsqueeze(-1).unsqueeze(-1).unsqueeze(-1) # (B, N, 1, 1, 1)
outer = outer / denom
```
Wait, but denom is (B, N) and outer is (B, N, N, H, H). The broadcast of denom would be:
(B, N, 1, 1, 1) broadcast to (B, N, N, H, H)?
Actually, let me check PyTorch broadcasting:
- denom after unsqueeze: (B, N, 1, 1, 1)
- outer: (B, N, N, H, H)
- Broadcast: (B, N, 1, 1, 1) -> (B, N, N, H, H) by repeating the last 3 dimensions
So the denom is per (b, i) but applied to all (b, i, j, h, k). That's a bit odd.
Actually wait, looking at the unsqueezes:
```python
denom = msa_mask.sum(dim=1).unsqueeze(-1).unsqueeze(-1).unsqueeze(-1).clamp(min=1)
```
msa_mask.sum(dim=1) is (B, N).
unsqueeze(-1) -> (B, N, 1)
unsqueeze(-1) -> (B, N, 1, 1)
unsqueeze(-1) -> (B, N, 1, 1, 1)
And outer is (B, N, N, H, H).
When dividing, PyTorch broadcasts:
(B, N, 1, 1, 1) / (B, N, N, H, H)
The N in denom broadcasts to the first N in outer, but the second N in outer is the j dimension.
Hmm, this means the division is actually by denom[b, i], applied to all j for that i. That makes sense now.
Actually, I realize I might be misreading. Let me re-examine:
- outer[b,i,j,h,k] = sum_s a[b,s,i,h] * b[b,s,j,k]
- denom is (B, N) -> denom[b, i]
- outer / denom means outer[b,i,j,h,k] /= denom[b,i]
Wait, that's not symmetric. Let me check if this is intentional or a bug in the reference.
Looking at the reference:
```python
denom = msa_mask.sum(dim=1).unsqueeze(-1).unsqueeze(-1).unsqueeze(-1).clamp(min=1)
outer = outer / denom
```
Hmm, denom[b,i] is the number of valid sequences at position i. And this divides all outer[b,i,j,:,:] for all j.
But wait, a[b,s,i,:] is masked by mask[b,s,i], and b[b,s,j,:] is masked by mask[b,s,j]. So if position i has fewer valid sequences, a values are zero for invalid s, and the outer product is divided by fewer counts.
Actually, I think the reference might have an issue. The outer product at (i,j) involves a[:,:,i,:] and b[:,:,j,:]. If position i has mask_sum_i valid sequences and position j has mask_sum_j, the sum in the einsum would naturally have fewer terms if either has fewer valid sequences.
But the division by denom[b,i] (not denom[b,j]) suggests the normalization is based only on position i's validity.
Let me just match the reference behavior. The code does:
- outer[b,i,j,h,k] /= denom[b,i]
where denom[b,i] = sum_s mask[b,s,i].
OK so now I understand the computation. Let me design the kernel:
Outer + output projection kernel:
- Each thread handles one (b, i, j, d) output element
- Block: all threads for one (b, i, j) computing all d outputs
- Or block: all threads for one (b, i) computing all (j, d) outputs
Actually, with pair_dim=64, we have 64 outputs per (i,j) pair. So one block per (b,i,j) with 64 threads makes sense.
Each thread (for one d):
1. Load a[b,:,i,:] into registers or shared (S x H = 8 x 32 = 256 floats)
2. Load b[b,:,j,:] into shared (S x H = 256 floats)
3. Load out_proj_weight[d, :] (H*H = 1024 floats)
4. Compute sum_{s,h,k} a[s,h] * b[s,k] * W[d,hk] / denom[b,i] + bias[d]
This requires a lot of loads. Let me think about optimization.
Actually, for each (b,i,j) pair:
- We need a[b,:,i,:] (256 floats) and b[b,:,j,:] (256 floats)
- We need out_proj_weight (64 x 1024 floats)
- We need denom[b,i] (1 float)
For each d (64 outputs):
- dot product of outer[b,i,j,:,:] with out_proj_weight[d,:]
- outer is 1024 floats, W[d,:] is 1024 floats
Computing outer on-the-fly:
outer[h,k] = sum_s a[s,h] * b[s,k] / denom[b,i]
Then out[d] = sum_{h,k} outer[h,k] * W[d, h*H+k] + bias[d]
= sum_{h,k} (sum_s a[s,h] * b[s,k] / denom[b,i]) * W[d, h*H+k] + bias[d]
= sum_s sum_{h,k} a[s,h] * b[s,k] * W[d, h*H+k] / denom[b,i] + bias[d]
Let me define:
W_reshaped[d,h,k] = W[d, h*H+k]
Then:
out[d] = sum_s sum_{h,k} a[s,h] * b[s,k] * W_reshaped[d,h,k] / denom[b,i] + bias[d]
= sum_s (sum_h a[s,h] * sum_k b[s,k] * W_reshaped[d,h,k]) / denom[b,i] + bias[d]
Hmm, the inner sum doesn't factor nicely because W depends on both h and k.
Let me try a different factorization:
sum_{h,k} a[s,h] * b[s,k] * W[d,h,k] = sum_h a[s,h] * (sum_k b[s,k] * W[d,h,k])
For each s, h:
temp[s,h] = sum_k b[s,k] * W[d,h,k]
Then:
out[d] = sum_s sum_h a[s,h] * temp[s,h] / denom[b,i] + bias[d]
This requires, for each d:
- For each s,h: compute sum_k b[s,k] * W[d,h,k] -> S*H = 256 multiplications and additions
- Total: 64 * 256 = 16384 operations for temp
- Then sum_s sum_h a[s,h] * temp[s,h] -> S*H = 256 more operations per d
- Total: 64 * (256 + 256) = 32768 operations
Compared to directly computing outer then projecting:
- Compute outer: S * H * H = 8 * 32 * 32 = 8192 operations per (i,j)
- Project: 64 * H * H = 64 * 1024 = 65536 operations per (i,j)
- Total: 73728 operations
So the fused approach is more efficient.
Let me implement the fused kernel:
1. Load a[b,:,i,:] and b[b,:,j,:] into shared memory
2. Each thread handles one d
3. Each thread computes:
- For each s,h: temp = sum_k b[s,k] * W[d,h,k]
- sum_s sum_h a[s,h] * temp
- Divide by denom and add bias
Actually, we can make this even more efficient:
- Precompute for each (b,s,j,d): temp[b,s,j,d,h] = sum_k b[b,s,j,k] * W[d,h,k]
- Then for each (b,i,j,d): out = sum_s sum_h a[b,s,i,h] * temp[b,s,j,d,h]
But this requires storing a large intermediate tensor.
Let me just implement the straightforward fused kernel and optimize later if needed.
Actually, given the small sizes, let me try an even simpler approach:
- Each thread handles one (b,i,j,d) output
- The thread loops over s, h, k and accumulates the result
- This avoids shared memory complexity
With B=2, N=32, N=32, pair_dim=64:
- Total threads: 2*32*32*64 = 131072 threads
- Grid: 2*32*32 = 2048 blocks
- Threads per block: 64
This is a reasonable configuration. Let me implement this.
Actually, I realize we need to share the a and b data within a block. Each thread in a (b,i,j) block needs:
- a[b,:,i,:] (same for all threads in block)
- b[b,:,j,:] (same for all threads in block)
- out_proj_weight (different for each thread, indexed by d)
- denom[b,i] (same for all threads in block)
So we can:
1. Cooperative load a and b into shared memory
2. Each thread loads its own weight row
3. Each thread computes its output
Let me implement this:
```cpp
__global__ void outer_product_mean_output(
const float* __restrict__ a, // (B, S, N, H)
const float* __restrict__ b, // (B, S, N, H)
const float* __restrict__ denom, // (B, N)
const float* __restrict__ out_weight, // (pair_dim, H*H)
const float* __restrict__ out_bias, // (pair_dim)
float* __restrict__ output, // (B, N, N, pair_dim)
int B, int S, int N, int H, int pair_dim
) {
int bij = blockIdx.x; // which (b, i, j) pair
int d = threadIdx.x; // which output dimension
int b = bij / (N * N);
int ij = bij % (N * N);
int i = ij / N;
int j = ij % N;
// Shared memory for a and b
extern __shared__ float shared_data[];
float* shared_a = shared_data; // S * H
float* shared_b = shared_data + S * H; // S * H
// Cooperative load
for (int idx = threadIdx.x; idx < S * H; idx += blockDim.x) {
int s = idx / H;
int h = idx % H;
shared_a[idx] = a[b * S * N * H + s * N * H + i * H + h];
shared_b[idx] = b[b * S * N * H + s * N * H + j * H + h];
}
__syncthreads();
// Load output projection weight for this thread's d
// ... but this is 1024 floats, too much for registers
// Instead, loop over h, k and load weight as needed
float sum = 0.0f;
for (int s = 0; s < S; s++) {
for (int h = 0; h < H; h++) {
for (int k = 0; k < H; k++) {
int hk = h * H + k;
float w = out_weight[d * H * H + hk];
sum += shared_a[s * H + h] * shared_b[s * H + k] * w;
}
}
}
float denom_val = denom[b * N + i];
sum = sum / denom_val + out_bias[d];
output[b * N * N * pair_dim + i * N * pair_dim + j * pair_dim + d] = sum;
}
```
Wait, this kernel has each thread doing S*H*H = 8*32*32 = 8192 iterations, which might be slow.
Let me think about a better approach. The issue is each thread needs to access the entire weight matrix for its d, which is 1024 floats.
Actually, we can use shared memory for the weight matrix too! But 64 * 1024 = 65536 floats is too large for shared memory (we have 227KB = ~57K floats max).
Let me think differently. We can tile the h and k dimensions:
- Each thread handles a subset of (h, k) pairs
- Use warp reduction to sum across threads
With 64 threads and 1024 (h,k) pairs, each thread handles 16 pairs.
Let me try:
- Thread 0: (h,k) pairs 0-15
- Thread 1: (h,k) pairs 16-31
- ...
Each thread computes partial sum, then warp reduces.
Actually, with 64 threads in a block and 1024 elements per d, we could have each thread compute 16 elements and then reduce. But the weight matrix access would still be uncoalesced.
Let me try a different tiling:
- Tile h: each thread handles H/4 = 8 h values
- Tile k: each thread handles all K = 32 k values (or also tile k)
- This gives 4 threads for h tiling, we need more parallelism
Hmm, let me reconsider. With pair_dim=64, we want 64 threads per block for the output dimension. The work per thread (sum over s,h,k) can be substantial.
An alternative is to have more threads and do reductions:
- 256 threads per block (for example)
- Each thread handles a subset of (h, k) pairs
- Reduce across threads to get the final sum for each d
But then we need multiple threads per d, which complicates the indexing.
Actually, let me try a simpler approach first:
- Use warp-level parallelism for the h,k loop
- Each warp handles one d
- Threads in warp handle different (h, k) pairs
With 32 threads per warp and 1024 (h,k) pairs, each thread handles 32 pairs. That's 32 * 8 = 256 iterations per thread (including s loop).
With 64 outputs (pair_dim), we need 2 warps per block. Block size = 64 works.
Let me implement this:
```cpp
__global__ void outer_product_mean_output(
const float* __restrict__ a,
const float* __restrict__ b,
const float* __restrict__ denom,
const float* __restrict__ out_weight,
const float* __restrict__ out_bias,
float* __restrict__ output,
int B, int S, int N, int H, int pair_dim
) {
int bij = blockIdx.x;
int d = threadIdx.x; // 0 to 63
int b = bij / (N * N);
int ij = bij % (N * N);
int i = ij / N;
int j = ij % N;
// Load a and b into shared memory
extern __shared__ float shared_data[];
float* shared_a = shared_data; // S * H = 256 floats
float* shared_b = shared_data + S * H; // 256 floats
for (int idx = threadIdx.x; idx < S * H; idx += blockDim.x) {
int s = idx / H;
int h = idx % H;
shared_a[idx] = a[((b * S + s) * N + i) * H + h];
shared_b[idx] = b[((b * S + s) * N + j) * H + h];
}
__syncthreads();
// Compute partial sum for this thread
float sum = 0.0f;
int H2 = H * H; // 1024
// Each thread handles H2/blockDim.x = 1024/64 = 16 (h,k) pairs
// Actually, we need to handle all pairs but with strided access
// Let's do: each thread handles pairs (h*H + k) where (h*H + k) % 64 == threadIdx.x
// But this gives uncoalesced weight access
// Better: thread handles pairs starting at threadIdx.x * 16
int start_hk = threadIdx.x * 16;
for (int hk_idx = 0; hk_idx < 16; hk_idx++) {
int hk = start_hk + hk_idx;
int h = hk / H;
int k = hk % H;
float w = out_weight[d * H2 + hk];
for (int s = 0; s < S; s++) {
sum += shared_a[s * H + h] * shared_b[s * H + k] * w;
}
}
// Now we have partial sums for each d
// But we need to reduce within each d... wait, no, each thread already has the full sum
// Actually, each thread handles only 16 out of 1024 (h,k) pairs
// We need to reduce across threads for the same d
// Hmm, this doesn't work. Let me reconsider.
}
```
Wait, I'm confusing myself. Let me think clearly:
For output[b,i,j,d]:
- Need to sum over s, h, k: a[b,s,i,h] * b[b,s,j,k] * W[d, h*H+k]
- Total iterations: S * H * H = 8192
With 64 threads per block (one per d), each thread needs to do 8192 iterations. That's fine for modern GPUs.
With the inner loop over s, h, k, each iteration does:
- Load shared_a[s,h]
- Load shared_b[s,k]
- Load W[d, hk]
The weight access is the bottleneck. W[d, hk] = W[d * H*H + hk]. For consecutive hk values, this gives coalesced access within each thread, but not across threads.
Actually, let me restructure:
- Loop over s first (outer loop)
- For each s, loop over hk
- This allows caching weight in registers?
No, the weight depends on d (threadIdx.x), so each thread loads different weight elements.
Let me try a different approach: transpose the weight matrix in shared memory.
W is (pair_dim, H*H) = (64, 1024). If we transpose to (H*H, pair_dim), then for a fixed hk, all threads access W_transposed[hk, d] which is contiguous in memory.
We can load W into shared memory transposed:
shared_W[hk * pair_dim + d] = W[d * H*H + hk]
Then within the loop:
sum += shared_a[s,h] * shared_b[s,k] * shared_W[hk * pair_dim + d]
This gives coalesced access across threads for the weight.
Let me calculate shared memory usage:
- shared_a: S * H = 256 floats
- shared_b: S * H = 256 floats
- shared_W: H * H * pair_dim = 1024 * 64 = 65536 floats
Total: 66048 floats = 264KB, which exceeds shared memory limit (227KB).
So we can't load the entire weight matrix. Let me tile the hk dimension:
- Process hk in tiles
- Each tile loads a subset of the weight matrix
With tile size = 256 (for example):
- shared_W_tile: 256 * 64 = 16384 floats = 64KB
- Total shared: 256 + 256 + 16384 = 16896 floats = ~67KB, which fits
Let me implement this tiled approach:
Actually, let me first implement a simpler version without tiling and see if it's fast enough. The small problem size might not benefit much from complex tiling.
Let me write a straightforward kernel and test it.
call → compile_kernel
callcompile_kernel
{
"kernel_code": "import torch\nimport torch.nn as nn\nimport torch.utils.cpp_extension as cpp_extension\n\nclass ModelNew(nn.Module):\n def __init__(self, msa_dim, pair_dim, hidden_dim=32):\n super().__init__()\n self.msa_dim = msa_dim\n self.pair_dim = pair_dim\n self.hidden_dim = hidden_dim\n \n # Create layers with same structure as reference\n self.layer_norm = nn.LayerNorm(msa_dim)\n self.proj_a = nn.Linear(msa_dim, hidden_dim)\n self.proj_b = nn.Linear(msa_dim, hidden_dim)\n self.out_proj = nn.Linear(hidden_dim * hidden_dim, pair_dim)\n \n # Compile CUDA kernel\n self._cuda_module = self._compile_kernel()\n \n def _compile_kernel(self):\n cuda_source = \"\"\"\n#include <torch/extension.h>\n#include <cuda_fp16.h>\n\n__global__ void fused_outer_product_mean_kernel(\n const float* __restrict__ input, // (B, S, N, msa_dim)\n const float* __restrict__ mask, // (B, S, N)\n const float* __restrict__ ln_weight, // (msa_dim,)\n const float* __restrict__ ln_bias, // (msa_dim,)\n const float* __restrict__ proj_a_w, // (hidden_dim, msa_dim)\n const float* __restrict__ proj_a_b, // (hidden_dim,)\n const float* __restrict__ proj_b_w, // (hidden_dim, msa_dim)\n const float* __restrict__ proj_b_b, // (hidden_dim,)\n const float* __restrict__ out_w, // (pair_dim, hidden_dim*hidden_dim)\n const float* __restrict__ out_b, // (pair_dim,)\n float* __restrict__ output, // (B, N, N, pair_dim)\n float* __restrict__ a_buf, // (B, S, N, hidden_dim)\n float* __restrict__ b_buf, // (B, S, N, hidden_dim)\n int B, int S, int N, int msa_dim, int hidden_dim, int pair_dim\n) {\n int bsn = blockIdx.x;\n int b = bsn / (S * N);\n int sn = bsn % (S * N);\n int s = sn / N;\n int n = sn % N;\n \n int tid = threadIdx.x;\n int block_size = blockDim.x;\n \n // Shared memory for LN computation\n extern __shared__ float shared_mem[];\n float* sh_input = shared_mem; // msa_dim\n float* sh_proj_a = sh_input + msa_dim; // hidden_dim\n float* sh_proj_b = sh_proj_a + hidden_dim; // hidden_dim\n float* sh_mean = sh_proj_b + hidden_dim; // 1\n float* sh_var = sh_mean + 1; // 1\n \n // Load input into shared memory\n for (int i = tid; i < msa_dim; i += block_size) {\n sh_input[i] = input[bsn * msa_dim + i];\n }\n __syncthreads();\n \n // Compute mean\n float sum = 0.0f;\n for (int i = tid; i < msa_dim; i += block_size) {\n sum += sh_input[i];\n }\n // Reduce within warp\n for (int offset = 16; offset > 0; offset /= 2) {\n sum += __shfl_down_sync(0xffffffff, sum, offset);\n }\n if (tid == 0) {\n sh_mean[0] = sum / msa_dim;\n }\n __syncthreads();\n \n float mean = sh_mean[0];\n \n // Compute variance\n float var_sum = 0.0f;\n for (int i = tid; i < msa_dim; i += block_size) {\n float diff = sh_input[i] - mean;\n var_sum += diff * diff;\n }\n for (int offset = 16; offset > 0; offset /= 2) {\n var_sum += __shfl_down_sync(0xffffffff, var_sum, offset);\n }\n if (tid == 0) {\n sh_var[0] = var_sum / msa_dim;\n }\n __syncthreads();\n \n float var = sh_var[0];\n float inv_std = rsqrtf(var + 1e-5f);\n \n // Apply mask\n float mask_val = mask[bsn];\n \n // Compute projections\n for (int h = tid; h < hidden_dim; h += block_size) {\n float proj_a_val = proj_a_b[h];\n float proj_b_val = proj_b_b[h];\n \n for (int i = 0; i < msa_dim; i++) {\n float normalized = (sh_input[i] - mean) * inv_std * ln_weight[i] + ln_bias[i];\n proj_a_val += normalized * proj_a_w[h * msa_dim + i];\n proj_b_val += normalized * proj_b_w[h * msa_dim + i];\n }\n \n sh_proj_a[h] = proj_a_val * mask_val;\n sh_proj_b[h] = proj_b_val * mask_val;\n }\n __syncthreads();\n \n // Store a and b to global memory\n for (int h = tid; h < hidden_dim; h += block_size) {\n a_buf[bsn * hidden_dim + h] = sh_proj_a[h];\n b_buf[bsn * hidden_dim + h] = sh_proj_b[h];\n }\n}\n\n__global__ void outer_product_output_kernel(\n const float* __restrict__ a, // (B, S, N, hidden_dim)\n const float* __restrict__ b, // (B, S, N, hidden_dim)\n const float* __restrict__ denom, // (B, N)\n const float* __restrict__ out_w, // (pair_dim, hidden_dim*hidden_dim)\n const float* __restrict__ out_b, // (pair_dim,)\n float* __restrict__ output, // (B, N, N, pair_dim)\n int B, int S, int N, int hidden_dim, int pair_dim\n) {\n int bij = blockIdx.x;\n int b = bij / (N * N);\n int ij = bij % (N * N);\n int i = ij / N;\n int j = ij % N;\n \n int d = threadIdx.x;\n \n int H = hidden_dim;\n int H2 = H * H;\n \n // Load denom\n float denom_val = denom[b * N + i];\n float inv_denom = 1.0f / fmaxf(denom_val, 1.0f);\n \n // Each thread computes one output dimension d\n // Loop over s, h, k\n float sum = 0.0f;\n \n for (int s = 0; s < S; s++) {\n // Load a[s,i,:] and b[s,j,:] into registers\n // Each thread needs values for all h (since weight depends on h)\n // We'll process h outer, k inner\n \n for (int h = 0; h < H; h++) {\n float a_val = a[(b * S + s) * N * H + i * H + h];\n \n for (int k = 0; k < H; k++) {\n float b_val = b[(b * S + s) * N * H + j * H + k];\n int hk = h * H + k;\n float w = out_w[d * H2 + hk];\n sum += a_val * b_val * w;\n }\n }\n }\n \n // Apply normalization and bias\n output[b * N * N * pair_dim + i * N * pair_dim + j * pair_dim + d] = sum * inv_denom + out_b[d];\n}\n\ntorch::Tensor forward_cuda(\n torch::Tensor input,\n torch::Tensor mask,\n torch::Tensor ln_weight,\n torch::Tensor ln_bias,\n torch::Tensor proj_a_w,\n torch::Tensor proj_a_b,\n torch::Tensor proj_b_w,\n torch::Tensor proj_b_b,\n torch::Tensor out_w,\n torch::Tensor out_b\n) {\n int B = input.size(0);\n int S = input.size(1);\n int N = input.size(2);\n int msa_dim = input.size(3);\n int hidden_dim = proj_a_w.size(0);\n int pair_dim = out_w.size(0);\n \n auto options = torch::TensorOptions().dtype(torch::kFloat32).device(input.device());\n auto a_buf = torch::empty({B, S, N, hidden_dim}, options);\n auto b_buf = torch::empty({B, S, N, hidden_dim}, options);\n auto output = torch::empty({B, N, N, pair_dim}, options);\n \n // Compute denom = mask.sum(dim=1) -> (B, N)\n auto denom = mask.sum(/*dim=*/1);\n \n int shared_mem_size = msa_dim + hidden_dim + hidden_dim + 2;\n shared_mem_size *= sizeof(float);\n \n // Launch kernel 1: fused LN + projection\n int blocks1 = B * S * N;\n int threads1 = 128;\n \n fused_outer_product_mean_kernel<<<blocks1, threads1, shared_mem_size>>>(\n input.data_ptr<float>(),\n mask.data_ptr<float>(),\n ln_weight.data_ptr<float>(),\n ln_bias.data_ptr<float>(),\n proj_a_w.data_ptr<float>(),\n proj_a_b.data_ptr<float>(),\n proj_b_w.data_ptr<float>(),\n proj_b_b.data_ptr<float>(),\n out_w.data_ptr<float>(),\n out_b.data_ptr<float>(),\n output.data_ptr<float>(),\n a_buf.data_ptr<float>(),\n b_buf.data_ptr<float>(),\n B, S, N, msa_dim, hidden_dim, pair_dim\n );\n \n // Launch kernel 2: outer product + output projection\n int blocks2 = B * N * N;\n int threads2 = pair_dim;\n \n outer_product_output_kernel<<<blocks2, threads2>>>(\n a_buf.data_ptr<float>(),\n b_buf.data_ptr<float>(),\n denom.data_ptr<float>(),\n out_w.data_ptr<float>(),\n out_b.data_ptr<float>(),\n output.data_ptr<float>(),\n B, S, N, hidden_dim, pair_dim\n );\n \n return output;\n}\n\nstd::vector<torch::Tensor> get_buffers_cuda(\n torch::Tensor input,\n torch::Tensor mask,\n torch::Tensor ln_weight,\n torch::Tensor ln_bias,\n torch::Tensor proj_a_w,\n torch::Tensor proj_a_b,\n torch::Tensor proj_b_w,\n torch::Tensor proj_b_b\n) {\n int B = input.size(0);\n int S = input.size(1);\n int N = input.size(2);\n int msa_dim = input.size(3);\n int hidden_dim = proj_a_w.size(0);\n \n auto options = torch::TensorOptions().dtype(torch::kFloat32).device(input.device());\n auto a_buf = torch::empty({B, S, N, hidden_dim}, options);\n auto b_buf = torch::empty({B, S, N, hidden_dim}, options);\n \n int shared_mem_size = msa_dim + hidden_dim + hidden_dim + 2;\n shared_mem_size *= sizeof(float);\n \n int blocks = B * S * N;\n int threads = 128;\n \n fused_outer_product_mean_kernel<<<blocks, threads, shared_mem_size>>>(\n input.data_ptr<float>(),\n mask.data_ptr<float>(),\n ln_weight.data_ptr<float>(),\n ln_bias.data_ptr<float>(),\n proj_a_w.data_ptr<float>(),\n proj_a_b.data_ptr<float>(),\n proj_b_w.data_ptr<float>(),\n proj_b_b.data_ptr<float>(),\n nullptr,\n nullptr,\n nullptr,\n a_buf.data_ptr<float>(),\n b_buf.data_ptr<float>(),\n B, S, N, msa_dim, hidden_dim, 0\n );\n \n return {a_buf, b_buf};\n}\n\"\"\"\n \n return cpp_extension.load_inline(\n name='outer_product_mean_cuda',\n cpp_sources='',\n cuda_sources=cuda_source,\n functions=['forward_cuda'],\n extra_cuda_cflags=['-O3', '--use_fast_math'],\n verbose=False\n )\n \n def forward(self, msa_repr: torch.Tensor, msa_mask: torch.Tensor) -> torch.Tensor:\n return self._cuda_module.forward_cuda(\n msa_repr.contiguous(),\n msa_mask.contiguous(),\n self.layer_norm.weight.contiguous(),\n self.layer_norm.bias.contiguous(),\n self.proj_a.weight.contiguous(),\n self.proj_a.bias.contiguous(),\n self.proj_b.weight.contiguous(),\n self.proj_b.bias.contiguous(),\n self.out_proj.weight.contiguous(),\n self.out_proj.bias.contiguous()\n )"
}