reasoning
Let me analyze this PyTorch model to understand what I need to implement:
1. The model is a Triangular Multiplicative Update (outgoing) from AlphaFold2 Evoformer
2. Input: pair_repr (B, N, N, pair_dim), pair_mask (B, N, N)
3. Operations:
- LayerNorm on input
- Linear projections for a and b (with gating)
- Element-wise multiplication with sigmoid gates and mask
- Einsum: "bikh,bjkh->bijh" - this aggregates over k with a[i,k] * b[j,k]
- LayerNorm on output
- Output projection with gating
Let me trace through the dimensions:
- pair_repr: (B, N, N, pair_dim) where N=32, pair_dim=64
- After layer_norm_in: (B, N, N, pair_dim)
- proj_a(x): (B, N, N, hidden_dim) where hidden_dim=32
- gate_a(x): (B, N, N, hidden_dim), sigmoid: (B, N, N, hidden_dim)
- a = proj_a(x) * sigmoid(gate_a(x)) * mask: (B, N, N, hidden_dim)
- Similarly for b: (B, N, N, hidden_dim)
The einsum "bikh,bjkh->bijh":
- a has indices (b, i, k, h) - shape (B, N, N, H)
- b has indices (b, j, k, h) - shape (B, N, N, H)
- Result has indices (b, i, j, h) - shape (B, N, N, H)
So for each output position (i, j), we sum over k: sum_k(a[i,k] * b[j,k])
This is essentially: for each (i,j), compute sum over k of (a_{i,k,h} * b_{j,k,h})
The result feeds into layer_norm_out, then passes through out_proj to project back to pair_dim. The output is gated by sigmoid(out_gate(pair_repr)) before returning.
For CUDA optimization, the core computation involves the einsum operation on sequences of length N=32 and hidden dimension H=32. Since batch_size and hidden dimension can be parallelized, I'll design a kernel that threads over batch and hidden dimensions, letting each thread handle one (i,j) pair's accumulation across k. The operations follow a straightforward sequence: reshape the mask, apply layer norm, compute projections with gating, perform the einsum, apply output layer norm, project back to pair_dim, and gate with sigmoid.
The mask needs careful handling—it starts as (B, N, N) and gets unsqueezed to (B, N, N, 1), which broadcasts correctly across the hidden dimension when multiplied with tensors of shape (B, N, N, H).
For the einsum operation, the first tensor has indices (b, i, k, h) and the second has (b, j, k, h), with the output indexed as (b, i, j, h). This means summing over k while contracting the shared dimension.
The implementation iterates through each batch, output position (i, j), and hidden dimension h, accumulating the product of the two input tensors across all k values. I can map this directly to a CUDA kernel where each thread handles one element in the output tensor.
For the complete ModelNew implementation, I'd use torch's cpp_extension to compile the kernel. The forward pass applies layer normalization to the input, then projects it through two separate paths with gating mechanisms. The mask gets reshaped and applied to both projections. Finally, the einsum aggregates the result over the k dimension, followed by output normalization and projection. x = self.layer_norm_in(pair_repr)
# Compute projections with gating
a_proj = self.proj_a(x)
a_gate = torch.sigmoid(self.gate_a(x))
a = a_proj * a_gate * mask
b_proj = self.proj_b(x)
b_gate = torch.sigmoid(self.gate_b(x))
b = b_proj * b_gate * mask
# Use CUDA kernel for einsum
out = self.einsum_cuda(a, b)
out = self.layer_norm_out(out)
out = self.out_proj(out)
out = torch.sigmoid(self.out_gate(pair_repr)) * out
return out
The forward pass normalizes the input, then applies gated projections to create a and b tensors. These get processed through a custom CUDA kernel for the einsum operation, followed by normalization, projection, and a final gating step. The kernel itself can be optimized by leveraging existing PyTorch primitives while customizing the einsum computation specifically.
The actual computation is `out[b,i,j,h] = sum_k(a[b,i,k,h] * b[b,j,k,h])`, which multiplies corresponding hidden dimensions across the k-axis. Since each batch operates independently and outputs can be computed in parallel, this maps well to GPU execution with each thread handling one element of the output tensor.
For a kernel, I'd assign each thread to compute a single (b, i, j, h) index, iterating through k values sequentially while accumulating the result. The input tensors are shaped (B, N, N, H), and the output becomes (B, N, N, H). Using block dimensions that maximize occupancy and leveraging shared memory for repeated accesses would improve performance.
For the einsum computation, I'll use a grid-stride loop to handle all elements efficiently. Each thread processes multiple output elements across the batch and hidden dimensions, accumulating contributions from the k-axis through a dot product between corresponding slices of the input tensors. __global__ void einsum_kernel(const float* __restrict__ a, const float* __restrict__ b,
float* __restrict__ out, int B, int N, int H) {
int total = B * N * N * H;
for (int idx = blockIdx.x * blockDim.x + threadIdx.x; idx < total; idx += blockDim.x * gridDim.x) {
int h = idx % H;
int j = (idx / H) % N;
int i = (idx / (H * N)) % N;
int b = idx / (H * N * N);
float sum = 0.0f;
for (int k = 0; k < N; ++k) {
int a_idx = b * N * N * H + i * N * H + k * H + h;
int b_idx = b * N * N * H + j * N * H + k * H + h;
sum += a[a_idx] * b[b_idx];
}
out[idx] = sum;
}
}
The kernel processes each output element independently using a grid-stride loop, computing the einsum by iterating through the k dimension and accumulating products. Now I need to expose this through a C++ wrapper function that torch can call.
extern "C" torch::Tensor einsum_cuda(torch::Tensor a, torch::Tensor b) {
TORCH_CHECK(a.dim() == 4 && b.dim() == 4, "Inputs must be 4D");
int B = a.size(0), N = a.size(1), H = a.size(3);
auto out = torch::zeros({B, N, N, H}, a.options());
int total = B * N * N * H;
int threads = 256;
int blocks = (total + threads - 1) / threads;
einsum_kernel<<<blocks, threads>>>(a.data_ptr<float>(), b.data_ptr<float>(), out.data_ptr<float>(), B, N, H);
return out;
}
The wrapper validates input dimensions, allocates the output tensor, calculates grid and block sizes, then launches the kernel with raw float pointers extracted from the tensors.
For optimization, I should consider whether coalesced memory access and shared memory usage could accelerate the kernel, though for the modest problem size of N=32, these improvements might be marginal. The basic implementation should suffice for correctness. I'm ensuring the tensors are contiguous on the GPU before extracting their data pointers and launching the kernel with the appropriate grid and block dimensions. After the kernel executes, I return the output tensor. Now I need to build the module using `torch.utils.cpp_extension.load_inline` with the C++ source containing the einsum operation and the CUDA kernel definition.
The module gets loaded with verbose output enabled and extra CUDA compilation flags. Then in the forward pass, I apply layer normalization to the input, project it through separate linear layers with sigmoid gating, mask the results, and invoke the custom einsum kernel to compute the output. I need to reconsider the approach. The mask should be (B, N, N, 1), and since both a and b are (B, N, N, H), broadcasting works correctly.
For optimization, the einsum is where I should focus—other operations like layer norm and linear layers already have efficient implementations. The einsum "bikh,bjkh->bijh" can be decomposed into: `out[b,i,j,h] = sum_k(a[b,i,k,h] * b[b,j,k,h])`.
I can use torch.matmul by reshaping: `a.transpose(1, 2)` becomes (B, K, N, H) and `b.permute(0, 2, 1, 3)` becomes (B, K, N, H), then `torch.matmul(a.transpose(1, 2).transpose(1, 2), b.permute(0, 2, 1, 3).transpose(2, 3))` but this gets convoluted.
For the einsum specifically, a custom kernel could be worthwhile. The mask needs to be properly shaped as (B, N, N, 1) to broadcast correctly with (B, N, N, H).
I'll build out the CUDA kernel with proper index calculations using the standard row-major layout where the stride for each dimension accounts for the sizes of subsequent dimensions. The output tensor gets indexed similarly, with h being the innermost dimension.
For each thread, I iterate through the k dimension to accumulate the sum by multiplying corresponding elements from a and b, where the indices map to the correct memory locations based on the computed strides. After the accumulation completes, I write the result to the output tensor.
The kernel launch wraps up by checking that both input tensors are contiguous and have the same shape, then calculates the total number of elements and configures the grid and block dimensions before launching with 256 threads per block. The C++ function returns the output tensor after the kernel completes execution.
For the ModelNew class, I'm setting up layer normalization and linear projections with gating mechanisms for both the input and output paths. The constructor takes pair_dim and hidden_dim as parameters, initializing separate layer norms and linear transformations for each stage. The forward method then applies the input layer norm followed by the projection and gating operations, multiplying the projected values with sigmoid-activated gate values.
After computing a and b through these gated projections, I pass them to a custom CUDA einsum operation. The output gets normalized again, projected back to the original dimension, and then gated one final time using the original pair representation before returning. The module gets compiled on module load with the CUDA kernel loaded inline. ptxas': ['-O3', '--use_fast_math']
)
)
The compiled function gets stored for use in the forward pass, though I should verify whether storing it in a global variable or caching it differently would be more appropriate. Since `load_inline` returns a module object, I can reference the `einsum_cuda` attribute directly when needed in the model's forward computation.
The mask handling needs adjustment—when `pair_mask` comes in as (B, N, N), calling `unsqueeze(-1)` produces (B, N, N, 1), which broadcasts correctly with (B, N, N, H) tensors. For the einsum kernel itself, I'm considering optimizations like shared memory usage, loop unrolling, or vectorization, though the small N=32 dimension means some optimizations may have limited benefit.
A basic kernel launch structure with 256 threads per block and grid sizing based on output elements should work, though I could explore having each thread handle multiple output elements or applying warp-level primitives for the reduction step.
For the kernel itself, I need to map each thread's global index back to the corresponding tensor dimensions. With output shape (B, N, N, H), I can extract h, j, i, and b from the flattened index, then iterate through the k dimension to compute the sum of products between a and b tensors at the appropriate positions.
The indexing into a and b requires accounting for the strides of the 4D tensors. For a at position (b, i, k, h), the offset is b*N*N*H + i*N*H + k*H + h, and similarly for b at (b, j, k, h). After accumulating the result, I write it to the output tensor and return.
The C++ wrapper checks that inputs are contiguous 4D tensors on the same device, extracts their dimensions, allocates an output tensor with shape (B, N, N, H), calculates the total element count, and launches the kernel with an appropriate number of blocks based on the thread count.
Now I need to load this as a CUDA extension. I'll define the C++ source with torch bindings, the CUDA kernel code, and load it inline with optimization flags for the compilation.
Once loaded, the extension module gives me access to the `einsum_cuda` function. I'll use this in the `ModelNew` class, which takes pair and hidden dimensions as parameters and initializes layer normalization followed by projection and gating layers for both input and output paths.
In the forward pass, I normalize the input, then compute gated projections for both 'a' and 'b' branches using sigmoid activations. These are masked before being combined through a custom einsum operation implemented in CUDA. The output undergoes another normalization, final projection, and gating based on the original representation before being returned. float* __restrict__ out, int B, int N, int H) {
int total = B * N * N * H;
for (int idx = blockIdx.x * blockDim.x + threadIdx.x; idx < total; idx += blockDim.x * gridDim.x) {
int h = idx % H;
int j = (idx / H) % N;
int i = (idx / (H * N)) % N;
int b = idx / (H * N * N);
float sum = 0.0f;
for (int k = 0; k < N; ++k) {
sum += a[b * N * N * H + i * N * H + k * H + h] *
b[b * N * N * H + j * N * H + k * H + h];
}
out[idx] = sum;
}
}
The kernel processes each output element by iterating through the k dimension, accumulating the product of corresponding elements from the two input tensors. I'm using a grid-stride loop to handle arbitrary tensor sizes efficiently.
torch::Tensor einsum_cuda(torch::Tensor a, torch::Tensor b) {
TORCH_CHECK(a.dim() == 4 && b.dim() == 4, "Inputs must be 4D");
TORCH_CHECK(a.is_contiguous() && b.is_contiguous(), "Inputs must be contiguous");
TORCH_CHECK(a.is_cuda() && b.is_cuda(), "Inputs must be on CUDA");
TORCH_CHECK(a.sizes() == b.sizes(), "Inputs must have same shape");
int B = a.size(0);
int N = a.size(1);
int H = a.size(3);
auto out = torch::zeros({B, N, N, H}, a.options());
int total = B * N * N * H;
int threads = 256;
int blocks = (total + threads - 1) / threads;
einsum_kernel<<<blocks, threads>>>(
a.data_ptr<float>(), b.data_ptr<float>(), out.data_ptr<float>(), B, N, H);
return out;
}
'''
The wrapper function validates input dimensions and contiguity, allocates the output tensor, then launches the kernel with enough thread blocks to cover all output elements. The final return passes the computed result back to Python.