reasoning
Now let me analyze the arithmetic intensity of the einsum operation:
For the einsum "bsih,bsjk->bijhk":
- Input a: (B, S, N, H) = (2, 8, 32, 32) = 16,384 elements
- Input b: same size = 16,384 elements
- Output: (B, N, N, H, H) = (2, 32, 32, 32, 32) = 2,097,152 elements
Total memory read: ~32KB for inputs, plus output write ~8MB
Computation: For each output element, we do 1 multiply and (S-1) adds. Actually we sum over S dimension:
- Each output[i,j,h,k] = sum over s of a[s,i,h] * b[s,j,k]
- That's S multiplies and S-1 adds per output element
- Total FLOPs: 2 * S * B * N * N * H * H = 2 * 8 * 2 * 32 * 32 * 32 * 32 = ~134M FLOPs
Data: ~8MB output + 32KB input = ~8MB
Arithmetic intensity: ~134M / 8MB = ~16.7 FLOPs/byte
This is below the ridge point (20.12), so the kernel will be memory-bound. The key is to maximize memory bandwidth utilization.
Given the small hidden_dim (32) and seq_len (32), shared memory tiling could help. Let me write an optimized kernel.
Strategy:
1. Use shared memory to load tiles of a and b
2. Each thread computes multiple output elements to increase computation per memory load
3. Compute the outer product and reduction over S dimension in shared memory
Actually, for the outer product, we have:
- a: (B, S, N, H) -> after layer norm and projection
- b: (B, S, N, H)
- output: (B, N, N, H, H)
The computation is:
output[b, i, j, h1, h2] = sum_s (a[b, s, i, h1] * b[b, s, j, h2] * mask[b, s, i]) / denom[b, i]
Wait, the mask is applied as:
a = a * mask # (B, S, N, H) * (B, S, N, 1) = (B, S, N, H)
b = b * mask # (B, S, N, H) * (B, S, N, 1) = (B, S, N, H)
Actually, looking at the code more carefully:
```python
mask = msa_mask.unsqueeze(-1) # (B, S, N, 1)
a = a * mask
b = b * mask
```
So both a and b are multiplied by the same mask. The einsum then:
outer = torch.einsum("bsih,bsjk->bijhk", a, b)
This computes:
outer[b, i, j, h1, h2] = sum_s (a[b, s, i, h1] * b[b, s, j, h2])
But wait, a and b already have the mask multiplied in. So effectively:
outer[b, i, j, h1, h2] = sum_s (x_a[b, s, i, h1] * x_b[b, s, j, h2] * mask[b, s, i] * mask[b, s, j])
Hmm, the mask is the same for both a and b at position (b, s, i) or (b, s, j). So for the outer product at (i, j), the mask values at position i and position j both matter.
Actually, re-reading the code:
- a is (B, S, N, H), masked: a[b, s, i, :] *= mask[b, s, i]
- b is (B, S, N, H), masked: b[b, s, j, :] *= mask[b, s, j]
The einsum "bsih,bsjk->bijhk" means:
- Contract s dimension
- Keep b, i, j, h (from first tensor), k (from second tensor) as h and rename to h1, h2
So outer[b, i, j, h1, h2] = sum over s of a[b, s, i, h1] * b[b, s, j, h2]
= sum over s of (proj_a[layer_norm(msa_repr)][b, s, i, h1] * mask[b, s, i]) * (proj_b[layer_norm(msa_repr)][b, s, j, h2] * mask[b, s, j])
The denominator:
denom = msa_mask.sum(dim=1).unsqueeze(-1).unsqueeze(-1).unsqueeze(-1).clamp(min=1)
= sum_s mask[b, s, i] for each b, i, then reshaped to (B, N, 1, 1, 1)
This gives the mean over valid (masked) sequences.
Let me write a fused kernel that:
1. Computes layer norm
2. Projects to a and b
3. Applies mask
4. Computes outer product with reduction over S
5. Divides by denominator
6. Reshapes and projects output
Actually, for simplicity and correctness, let me first try a more modular approach where I:
1. Keep the layer norm and projections as PyTorch ops
2. Fuse the outer product computation into a custom kernel
This will be easier to get right, and I can fuse more operations later if needed.
Let me write the kernel:
```cpp
// Kernel to compute outer product mean
// Inputs:
// a: (B, S, N, H)
// b: (B, S, N, H)
// mask: (B, S, N)
// Outputs:
// outer: (B, N, N, H, H)
// Denominator computed from mask sum
__global__ void outer_product_mean_kernel(
const float* a, const float* b, const float* mask,
float* outer, int B, int S, int N, int H) {
// Each thread handles one (b, i, j, h1, h2) output element
int idx = blockIdx.x * blockDim.x + threadIdx.x;
int total = B * N * N * H * H;
if (idx >= total) return;
// Decode indices
int h2 = idx % H;
idx /= H;
int h1 = idx % H;
idx /= H;
int j = idx % N;
idx /= N;
int i = idx % N;
idx /= N;
int b = idx;
// Compute outer product: sum over s of a[b,s,i,h1] * b[b,s,j,h2]
// But need to account for mask
// Actually, a and b are already masked, so we just sum
// But wait, we need the denominator which is sum of mask over S for position i
float sum = 0.0f;
for (int s = 0; s < S; s++) {
int a_idx = ((b * S + s) * N + i) * H + h1;
int b_idx = ((b * S + s) * N + j) * H + h2;
sum += a[a_idx] * b[b_idx];
}
// Denominator: sum of mask over S for position (b, i)
float denom = 0.0f;
for (int s = 0; s < S; s++) {
int mask_idx = (b * S + s) * N + i;
denom += mask[mask_idx];
}
denom = max(denom, 1.0f);
// Store result
int out_idx = (((b * N + i) * N + j) * H + h1) * H + h2;
outer[out_idx] = sum / denom;
}
```
Wait, there's an issue. Looking at the reference code again:
```python
outer = torch.einsum("bsih,bsjk->bijhk", a, b) # (B, N, N, H, H)
denom = msa_mask.sum(dim=1).unsqueeze(-1).unsqueeze(-1).unsqueeze(-1).clamp(min=1)
outer = outer / denom # mean over S
```
The denominator is `msa_mask.sum(dim=1)` which has shape (B, N), not (B, N, N).
So for each (b, i), the denominator is the same for all j values. That means:
denom[b, i, j, :, :] = sum_s mask[b, s, i]
Hmm, but that seems asymmetric. Let me think about this more carefully.
Actually, looking at the reference, the einsum sums over S:
outer[b, i, j, h1, h2] = sum_s a[b, s, i, h1] * b[b, s, j, h2]
And a[b, s, i, h1] = proj_a(...)[b, s, i, h1] * mask[b, s, i]
Similarly for b.
So the outer product is:
sum_s (proj_a * mask) * (proj_b * mask)
But wait, the mask for position i is mask[b, s, i], and for position j is mask[b, s, j]. These are different values!
Actually no, the mask multiplication is:
```python
mask = msa_mask.unsqueeze(-1) # (B, S, N, 1)
a = a * mask # broadcasting over H dimension
b = b * mask
```
So mask[b, s, i] (without the last dim) is used to multiply both a[b, s, i, :] and b[b, s, i, :].
In the einsum:
- a has indices (b, s, i, h), where i is used
- b has indices (b, s, j, k), where j is used
So a[b, s, i, :] is multiplied by mask[b, s, i], and b[b, s, j, :] is multiplied by mask[b, s, j].
Therefore:
outer[b, i, j, h1, h2] = sum_s a_raw[b, s, i, h1] * mask[b, s, i] * b_raw[b, s, j, h2] * mask[b, s, j]
= sum_s a_raw[b, s, i, h1] * b_raw[b, s, j, h2] * mask[b, s, i] * mask[b, s, j]
And the denominator is sum_s mask[b, s, i], which is the number of valid sequences at position i.
This is asymmetric! For different j values at the same i, the denominator is the same. But the outer product includes mask[b, s, j] as well...
Actually, I think I'm overcomplicating this. The mask at position (s, j) affects the b values at position j, not the i position. Let me trace through:
For a fixed (b, i, j):
- We're computing sum over s of a_masked[b, s, i, h1] * b_masked[b, s, j, h2]
- a_masked[b, s, i, h1] = proj_a[...][b, s, i, h1] * mask[b, s, i]
- b_masked[b, s, j, h2] = proj_b[...][b, s, j, h2] * mask[b, s, j]
So the product is:
proj_a[b, s, i, h1] * proj_b[b, s, j, h2] * mask[b, s, i] * mask[b, s, j]
And the denominator is sum_s mask[b, s, i].
Hmm, this is still asymmetric. But that's what the reference does, so I'll follow it.
Actually wait, let me re-read the denominator:
```python
denom = msa_mask.sum(dim=1).unsqueeze(-1).unsqueeze(-1).unsqueeze(-1).clamp(min=1)
```
msa_mask is (B, S, N), sum over dim=1 gives (B, N).
Then unsqueeze three times: (B, N, 1, 1, 1).
So denom[b, i, j, h1, h2] = sum_s mask[b, s, i] (the same for all j, h1, h2).
This is indeed asymmetric - the denominator only depends on i, not j.
OK let me implement this correctly.
For my kernel, I need to:
1. Read a, b, and mask from global memory
2. For each output element (b, i, j, h1, h2):
- Sum over s: a[b, s, i, h1] * b[b, s, j, h2]
- Divide by sum_s mask[b, s, i]
But wait, a and b are already masked! So a[b, s, i, h1] = raw_a[b, s, i, h1] * mask[b, s, i].
So the product a * b = raw_a * mask_i * raw_b * mask_j.
But the denominator is sum_s mask_i = sum_s mask[b, s, i].
So the effective computation is:
outer[b, i, j, h1, h2] = (sum_s raw_a * raw_b * mask_i * mask_j) / (sum_s mask_i)
This is indeed asymmetric. If mask_i is 1 for all s, the denominator is S. If mask_j is 0 for some s, those terms in the sum become 0.
OK, I'll implement this as the reference does. The input a and b are already masked, so I just need to sum over s and divide by denom.
Let me also think about optimization:
- Total threads needed: B * N * N * H * H = 2 * 32 * 32 * 32 * 32 = 2,097,152
- For H100 with 132 SMs, that's about 15,892 threads per SM
- But each thread only does S = 8 iterations, which is not much work
- Better to have each thread do multiple output elements
Actually, for better memory access patterns, let me think about how to organize the threads.
The input a is (B, S, N, H) and input b is (B, S, N, H). For a given b, s, i, we access a[b, s, i, :] which is contiguous in memory (H elements). Similarly for b[b, s, j, :].
For the output (B, N, N, H, H), the last two dimensions are H x H.
One approach: have each thread handle one (b, i, j) pair and compute the entire H x H outer product. This requires:
- Reading S * H values for a[b, :, i, :]
- Reading S * H values for b[b, :, j, :]
- Computing H * H output values
This is more memory efficient because we reuse the loaded a and b values.
Let me try this approach:
```cpp
__global__ void outer_product_mean_kernel(
const float* a, const float* b, const float* mask,
float* outer, int B, int S, int N, int H) {
// Each thread handles one (b, i, j) pair and computes H x H output
int idx = blockIdx.x * blockDim.x + threadIdx.x;
int total = B * N * N;
if (idx >= total) return;
int j = idx % N;
idx /= N;
int i = idx % N;
idx /= N;
int b = idx;
// Load a[b, :, i, :] into shared memory or registers
// a has shape (B, S, N, H), so a[b, s, i, :] is at offset ((b * S + s) * N + i) * H
// We need to sum over s for the outer product
float a_local[8][32]; // S x H, but S and H are not compile-time constants
float b_local[8][32];
// This won't work because S and H are runtime values
}
```
The issue is that S and H are not compile-time constants. Let me use shared memory with dynamic sizing.
Actually, let me simplify. Given the small sizes (S=8, H=32, N=32, B=2), a straightforward kernel might work well enough.
Let me try a tiled approach where each thread block handles one (b, i, j) pair and computes the H x H output.
Block size: H x H = 32 x 32 = 1024 threads (maximum)
Each thread computes one element of the H x H output.
Shared memory: load a[b, :, i, :] and b[b, :, j, :] into shared memory.
Each thread then computes sum over s.
This way, each thread loads:
- 1 element of a (for each s, shared across threads in the block)
- 1 element of b (similarly)
And the shared memory reuse is high.
Let me implement this:
```cpp
__global__ void outer_product_mean_kernel(
const float* a, const float* b, const float* mask,
float* outer, int B, int S, int N, int H) {
// Shared memory for a and b values
extern __shared__ float shared_mem[];
float* a_shared = shared_mem; // S x H
float* b_shared = shared_mem + S * H; // S x H
// Each block handles one (b, i, j) pair
// blockIdx.x encodes (b, i, j)
int block_idx = blockIdx.x;
int j = block_idx % N;
block_idx /= N;
int i = block_idx % N;
block_idx /= N;
int b = block_idx;
// Thread indices
int h1 = threadIdx.x; // 0 to H-1
int h2 = threadIdx.y; // 0 to H-1
// Load a and b into shared memory
// Each thread loads multiple elements based on its h1, h2
for (int s = 0; s < S; s++) {
if (h1 < H) {
a_shared[s * H + h1] = a[((b * S + s) * N + i) * H + h1];
}
if (h2 < H) {
b_shared[s * H + h2] = b[((b * S + s) * N + j) * H + h2];
}
__syncthreads();
}
// Compute outer product
float sum = 0.0f;
for (int s = 0; s < S; s++) {
sum += a_shared[s * H + h1] * b_shared[s * H + h2];
}
// Compute denominator
float denom = 0.0f;
for (int s = 0; s < S; s++) {
denom += mask[(b * S + s) * N + i];
}
denom = max(denom, 1.0f);
// Store result
if (h1 < H && h2 < H) {
int out_idx = (((b * N + i) * N + j) * H + h1) * H + h2;
outer[out_idx] = sum / denom;
}
}
```
Actually, there's a synchronization issue. Let me fix it:
```cpp
__global__ void outer_product_mean_kernel(
const float* a, const float* b, const float* mask,
float* outer, int B, int S, int N, int H) {
// Shared memory for a and b values
extern __shared__ float shared_mem[];
float* a_shared = shared_mem; // S x H
float* b_shared = shared_mem + S * H; // S x H
// Each block handles one (b, i, j) pair
int block_idx = blockIdx.x;
int j = block_idx % N;
block_idx /= N;
int i = block_idx % N;
block_idx /= N;
int b = block_idx;
// Thread indices (using 1D thread block)
int tid = threadIdx.x;
int h1 = tid % H;
int h2 = tid / H;
// Cooperatively load a and b into shared memory
// Each thread loads one element of a and one element of b for each s
for (int s = 0; s < S; s++) {
if (tid < H) {
a_shared[s * H + tid] = a[((b * S + s) * N + i) * H + tid];
}
if (tid < H) {
b_shared[s * H + tid] = b[((b * S + s) * N + j) * H + tid];
}
__syncthreads();
}
// Compute outer product (only for threads with valid h1, h2)
if (tid < H * H) {
float sum = 0.0f;
for (int s = 0; s < S; s++) {
sum += a_shared[s * H + h1] * b_shared[s * H + h2];
}
// Compute denominator
float denom = 0.0f;
for (int s = 0; s < S; s++) {
denom += mask[(b * S + s) * N + i];
}
denom = max(denom, 1.0f);
// Store result
int out_idx = (((b * N + i) * N + j) * H + h1) * H + h2;
outer[out_idx] = sum / denom;
}
}
```
Hmm, the syncthreads inside the loop is problematic. Let me separate the loading and computation:
```cpp
__global__ void outer_product_mean_kernel(
const float* a, const float* b, const float* mask,
float* outer, int B, int S, int N, int H) {
// Shared memory for a and b values
extern __shared__ float shared_mem[];
float* a_shared = shared_mem; // S x H
float* b_shared = shared_mem + S * H; // S x H
float* denom_shared = shared_mem + 2 * S * H; // 1 element
// Each block handles one (b, i, j) pair
int block_idx = blockIdx.x;
int j = block_idx % N;
block_idx /= N;
int i = block_idx % N;
block_idx /= N;
int b = block_idx;
// Thread indices (using 1D thread block)
int tid = threadIdx.x;
// Cooperatively load a and b into shared memory
// Total elements to load: 2 * S * H
int total_load = 2 * S * H;
for (int k = tid; k < total_load; k += blockDim.x) {
int which = k / (S * H); // 0 for a, 1 for b
int s_and_h = k % (S * H);
int s = s_and_h / H;
int h = s_and_h % H;
if (which == 0) {
a_shared[s * H + h] = a[((b * S + s) * N + i) * H + h];
} else {
b_shared[s * H + h] = b[((b * S + s) * N + j) * H + h];
}
}
// Compute denominator (only one thread needs to do this)
if (tid == 0) {
float denom = 0.0f;
for (int s = 0; s < S; s++) {
denom += mask[(b * S + s) * N + i];
}
denom_shared[0] = max(denom, 1.0f);
}
__syncthreads();
// Compute outer product
int total_output = H * H;
for (int k = tid; k < total_output; k += blockDim.x) {
int h1 = k % H;
int h2 = k / H;
float sum = 0.0f;
for (int s = 0; s < S; s++) {
sum += a_shared[s * H + h1] * b_shared[s * H + h2];
}
// Store result
int out_idx = (((b * N + i) * N + j) * H + h1) * H + h2;
outer[out_idx] = sum / denom_shared[0];
}
}
```
This is better. Now let me think about block and grid dimensions:
- Grid: B * N * N = 2 * 32 * 32 = 2048 blocks
- Block: Need at least H * H = 1024 threads to cover all output elements, but we can have fewer and loop
For occupancy, let's say 256 threads per block:
- Shared memory: 2 * S * H * 4 = 2 * 8 * 32 * 4 = 2048 bytes
- Each thread computes ceil(1024 / 256) = 4 output elements
- Good for occupancy
Actually, let me check the grid size. With 2048 blocks and 256 threads per block, that's 524,288 threads total, which is fine.
Let me write the complete kernel code and the Python wrapper.
Wait, I realized there might be an issue with the denominator computation. In the reference, the denominator uses the original msa_mask, not the masked a or b. But since a and b are already multiplied by mask, we need to use the original mask for the denominator.
Actually, looking at the reference:
```python
mask = msa_mask.unsqueeze(-1) # (B, S, N, 1)
a = a * mask
b = b * mask
outer = torch.einsum("bsih,bsjk->bijhk", a, b) # (B, N, N, H, H)
denom = msa_mask.sum(dim=1).unsqueeze(-1).unsqueeze(-1).unsqueeze(-1).clamp(min=1)
outer = outer / denom # mean over S
```
Yes, the denominator is computed from the original msa_mask, not from the masked a or b. So my kernel needs to receive the original mask as well.
Actually wait, looking more carefully at my approach: I'm keeping layer norm and projections in PyTorch, and only customizing the outer product computation. So the inputs to my kernel would be:
- a: already masked and projected, shape (B, S, N, H)
- b: already masked and projected, shape (B, S, N, H)
- mask: original msa_mask for denominator, shape (B, S, N)
But actually, I notice that the mask is used both for masking a/b AND for the denominator. Since a and b are already masked when they reach my kernel, I just need to use the mask for the denominator.
Let me also consider: should I fuse the layer norm and projections into the kernel? That could be more efficient, but also more complex. Let me start with just the outer product kernel and see if it's faster.
Actually, let me reconsider the overall approach. The PyTorch reference does:
1. Layer norm
2. Two linear projections (GEMM operations)
3. Mask multiplication
4. Einsum (outer product)
5. Division by denominator
6. Reshape
7. Final projection (GEMM)
The einsum is the most expensive part in terms of output size (2M elements) and computation. The layer norm and projections are standard operations that PyTorch already optimizes well.
So my strategy:
1. Keep layer norm and projections as PyTorch ops
2. Fuse the outer product computation into a custom kernel
This should give good performance without too much complexity.
Let me write the code:
```python
import torch
import torch.nn as nn
import torch.utils.cpp_extension as cpp_extension
cuda_source = """
#include <torch/extension.h>
#include <cuda.h>
#include <cuda_runtime.h>
__global__ void outer_product_mean_kernel(
const float* a, const float* b, const float* mask,
float* outer, int B, int S, int N, int H) {
// Shared memory for a and b values
extern __shared__ float shared_mem[];
float* a_shared = shared_mem; // S x H
float* b_shared = shared_mem + S * H; // S x H
float* denom_shared = shared_mem + 2 * S * H; // 1 element
// Each block handles one (b, i, j) pair
int block_idx = blockIdx.x;
int j = block_idx % N;
block_idx /= N;
int i = block_idx % N;
block_idx /= N;
int b = block_idx;
// Thread indices (using 1D thread block)
int tid = threadIdx.x;
// Cooperatively load a and b into shared memory
int total_load = 2 * S * H;
for (int k = tid; k < total_load; k += blockDim.x) {
int which = k / (S * H); // 0 for a, 1 for b
int s_and_h = k % (S * H);
int s = s_and_h / H;
int h = s_and_h % H;
if (which == 0) {
a_shared[s * H + h] = a[((b * S + s) * N + i) * H + h];
} else {
b_shared[s * H + h] = b[((b * S + s) * N + j) * H + h];
}
}
// Compute denominator (only one thread needs to do this)
if (tid == 0) {
float denom = 0.0f;
for (int s = 0; s < S; s++) {
denom += mask[(b * S + s) * N + i];
}
denom_shared[0] = max(denom, 1.0f);
}
__syncthreads();
// Compute outer product
int total_output = H * H;
for (int k = tid; k < total_output; k += blockDim.x) {
int h1 = k % H;
int h2 = k / H;
float sum = 0.0f;
for (int s = 0; s < S; s++) {
sum += a_shared[s * H + h1] * b_shared[s * H + h2];
}
// Store result
int out_idx = (((b * N + i) * N + j) * H + h1) * H + h2;
outer[out_idx] = sum / denom_shared[0];
}
}
torch::Tensor outer_product_mean_cuda(
torch::Tensor a, torch::Tensor b, torch::Tensor mask) {
// a: (B, S, N, H)
// b: (B, S, N, H)
// mask: (B, S, N)
// output: (B, N, N, H, H)
int B = a.size(0);
int S = a.size(1);
int N = a.size(2);
int H = a.size(3);
auto outer = torch::zeros({B, N, N, H, H}, a.options());
int blocks = B * N * N;
int threads = 256;
int shared_mem_size = (2 * S * H + 1) * sizeof(float);
outer_product_mean_kernel<<<blocks, threads, shared_mem_size>>>(
a.data_ptr<float>(), b.data_ptr<float>(), mask.data_ptr<float>(),
outer.data_ptr<float>(), B, S, N, H);
return outer;
}
"""
# Load the CUDA extension
outer_product_ext = cpp_extension.load_inline(
name='outer_product_ext',
cpp_sources='',
cuda_sources=cuda_source,
functions=['outer_product_mean_cuda'],
extra_cuda_cflags=['-O3'],
verbose=False
)
class ModelNew(nn.Module):
def __init__(self, msa_dim, pair_dim, hidden_dim=32):
super().__init__()
self.layer_norm = nn.LayerNorm(msa_dim)
self.proj_a = nn.Linear(msa_dim, hidden_dim)
self.proj_b = nn.Linear(msa_dim, hidden_dim)
self.out_proj = nn.Linear(hidden_dim * hidden_dim, pair_dim)
self.hidden_dim = hidden_dim
def forward(self, msa_repr: torch.Tensor, msa_mask: torch.Tensor) -> torch.Tensor:
x = self.layer_norm(msa_repr)
a = self.proj_a(x) # (B, S, N, H)
b = self.proj_b(x)
mask = msa_mask.unsqueeze(-1) # (B, S, N, 1)
a = a * mask
b = b * mask
# Use custom CUDA kernel for outer product
outer = outer_product_ext.outer_product_mean_cuda(a, b, msa_mask)
B, N1, N2, H1, H2 = outer.shape
outer = outer.reshape(B, N1, N2, H1 * H2)
return self.out_proj(outer)
```
Let me compile this to check for errors.