We need to replace a DP recurrence with a custom CUDA kernel. The recurrence is basically Smith-Waterman soft max. We have B=16, M=32, N=32. Very small! But the goal is to be faster than the PyTorch reference. Since loops in Python are slow, a custom kernel should easily beat it. However, we must respect the iterative DP dependency: H[i][j] depends on H[i-1][j-1], H[i-1][j], H[i][j-1]. So the computation order is sequential along anti-diagonals. But M and N are only 32. So we can either compute in kernel with each thread handling a batch element, iterating through i and j in nested loops (since M,N small, no divergence). That's simple: a kernel that loads score_matrix for a batch element into registers / shared memory, loops over i=1..M, j=1..N, updates H. Must compute temperature * logsumexp([0, match, delete, insert] / temperature). That is a stable reduce. Because M,N are small (32x32), we can allocate H as a per-thread array in registers of size (M+1)*(N+1) float. That's 33*33*4 ~ 4.4KB per thread. For batch size 16, that is fine (occupancy low but okay). Alternatively, use shared memory per block, maybe one block per batch or warp per batch. Since we want to maximize throughput, perhaps one thread per batch element is simplest. There are B=16 elements. GPU can run many threads but B is only 16. So occupancy low, but total work small. The kernel overhead will matter. Better: have each block process one batch element. Within block, maybe have threads cooperate to compute anti-diagonals. But with M=N=32, the overhead of shuffles may not be worth it. Could use one warp per batch element (32 threads). Each thread processes a column? There are 32 columns. But data dependencies: to compute cell (i,j), need left (i,j-1) and up (i-1,j) and diag (i-1,j-1). In a warp, threads can communicate via shfl for left? Actually left is same row previous column, which is computed in previous thread. Up is previous row same column (computed in same thread previous iteration). Diag is previous row previous column (computed in previous thread previous iteration). Warp-level shuffle can get left and diag from neighboring threads. This could be efficient. However, since M,N small, simpler approach may be enough. Let's be precise about the algorithm: For b in 0..B-1: H[0][*] = 0, H[*][0] = 0. for i=1..M: for j=1..N: match = H[i-1][j-1] + score[b][i-1][j-1] delete = H[i-1][j] + gap_open insert = H[i][j-1] + gap_open # zero is 0 # compute logsumexp of [0, match, delete, insert] scaled by temperature # Actually stack = [zero, match, delete, insert] # H[i][j] = temperature * logsumexp(stack / temperature) # logsumexp(x) = log(sum(exp(x - max(x)))) + max(x) # So compute max_val = max(0, match, delete, insert) # sum_exp = exp(0 - max_val) + exp(match - max_val) + exp(delete - max_val) + exp(insert - max_val) # logsum = log(sum_exp) + max_val # H[i][j] = temperature * logsum # Equivalent to temperature * (log(exp(0/t) + exp(match/t) + exp(delete/t) + exp(insert/t))) ? # Wait original: stack / temperature then logsumexp then * temperature. # logsumexp(x / t) * t = t * (log(sum(exp(x/t - max(x/t)))) + max(x/t)) # = t * log(sum(exp((x - max(x))/t))) + max(x) # Actually max(x/t) = max(x)/t, so t * max(x/t) = max(x). And exp(x/t - max(x)/t) = exp((x - max(x))/t). # So H = t * log( sum_{k} exp((x_k - max_x)/t) ) + max_x. # For t=1 it's standard logsumexp. # This is better for numerical stability. # We can compute max_val = max(0, match, delete, insert) # Then sum_exp = exp((0 - max_val)/t) + exp((match - max_val)/t) + ... # H = t * log(sum_exp) + max_val At the end, find max over all H[i][j] for i=1..M, j=1..N. But H also includes H[0][*] and H[*][0] which are zeros. Since we take max over all cells including zeros, and we know H >= 0? Actually logsumexp of [0, ...] ensures H >= 0 because logsumexp >= max >= 0? Wait if match, delete, insert are negative, the max could be 0. logsumexp([0, negative]) > 0, but as t -> 0, H -> max(0, ...) = 0. So H is >= 0. However, there may be positive values. The reference does H.view(B, -1).max(dim=-1).values which includes H[:,0,:] and H[:,:,0] which are all zeros. So the max is at least 0. That is fine. We need to compute the max over the entire H matrix. We can keep a running max per batch element. Since H[i][j] computed sequentially, we can update a running max variable. So kernel design options: Option A: one thread per batch element. - Launch B threads (or maybe B threads in one block, or B blocks of 1 thread). Since we need global memory loads for score_matrix, each thread loads elements on demand. - H can be stored in a small local array float h[(M+1)*(N+1)]; register usage maybe high but okay: 1089 floats ~ 4.3KB. For 16 threads, if they are in one warp, register usage maybe high but still manageable. Max registers per thread on sm_80 is 255 (can't exceed due to allocation granularity) but shared memory is large. Actually local arrays get spilled to local memory if too large for registers. 4.3KB per thread might spill. But with only 16 threads, local memory per thread is in L1, and should be fast enough. However, compiler might allocate registers for a few elements, and spill the rest. Since M,N are 32, maybe we can avoid full array: we only need previous row and current row. Since H[i][j] depends on H[i-1][j-1], H[i-1][j], H[i][j-1]. For a row-major computation, when processing row i, we need the entire previous row (i-1) and the current row up to j-1. So we need two arrays of length N+1: prev_row and curr_row. And also we need prev_row[j-1] for diag. That's easy. So memory per thread: 2*(N+1) floats = 66 floats = 264 bytes. That's trivial. Great. Let's do that. Algorithm per thread: - Read gap_open and temperature from constants (or kernel args). - b = blockIdx.x * blockDim.x + threadIdx.x - if b >= B: return - Base pointer to score_matrix[b, 0, 0] and stride. - Allocate prev_row[N+1] and curr_row[N+1] in local memory (or registers). Since N+1 is 33 and known at compile time, compiler should put it in registers if possible, or local mem. It's fine. - Initialize prev_row[0..N] = 0.0 - max_val = 0.0 (since all zeros are present) - For i in 1..M: curr_row[0] = 0.0 For j in 1..N: s = score_matrix[b, i-1, j-1] match = prev_row[j-1] + s delete = prev_row[j] + gap_open insert = curr_row[j-1] + gap_open // zero = 0.0 m = fmaxf(0.0f, fmaxf(match, fmaxf(delete, insert))); // sum_exp = exp((0 - m)/t) + exp((match - m)/t) + exp((delete - m)/t) + exp((insert - m)/t) float inv_t = 1.0f / temperature; float sum_exp = expf((0.0f - m) * inv_t) + expf((match - m) * inv_t) + expf((delete - m) * inv_t) + expf((insert - m) * inv_t); float h = temperature * logf(sum_exp) + m; curr_row[j] = h; if (h > max_val) max_val = h; // swap prev and curr for j=0..N: prev_row[j] = curr_row[j]; // can use pointer swap, but in C++ local arrays need copy or use two arrays and swap pointers. We can declare float row0[33], row1[33]; and swap pointers. - output[b] = max_val We need to be careful about the semantics: The reference initializes H as full (B, M+1, N+1) with zeros, then sets H[:, i, j] for all i,j. The max is over the flattened H which includes the zero padding. So initializing max_val = 0.0 is correct. Edge cases: What if temperature is 0? The spec says temperature=1.0 by default, but it's a float parameter. We should handle any positive temperature. Need to ensure score_matrix is contiguous? In PyTorch, we can assume standard layout but can pass strides. The reference calls it with contiguous tensor? The test inputs come from torch.randn(batch_size, seq_m, seq_n) which is contiguous. We can just index with pointer arithmetic assuming contiguous. In kernel, we can compute offset = b * M * N + i * N + j. Since b = blockIdx.x * blockDim.x + threadIdx.x. But if we use one block of B threads, we can launch with blockDim.x = B (16). Or more generically, blockDim.x = 256 and gridDim.x = ceil(B/256). The work per thread is M*N=1024 operations, each with a few fmas, exp, log. That's okay. Let's estimate reference time: Python loops over 32x32, each inner step does tensor ops on shape (B,). That's a lot of kernel launches (1024 tiny ops). The custom kernel does everything in one fused kernel. So speedup should be huge, maybe 100x or more. Now, details for load_inline: cpp_sources = r''' #include <torch/extension.h> #include <cuda_runtime.h> #include <math.h> __global__ void smith_waterman_kernel(const float* __restrict__ score_matrix, float* __restrict__ output, float gap_open, float temperature, int B, int M, int N) { int b = blockIdx.x * blockDim.x + threadIdx.x; if (b >= B) return; // local arrays for previous and current row float prev[33]; float curr[33]; for (int j = 0; j <= N; ++j) { prev[j] = 0.0f; } float max_val = 0.0f; const float inv_t = 1.0f / temperature; const int batch_stride = M * N; const float* scores = score_matrix + b * batch_stride; for (int i = 1; i <= M; ++i) { curr[0] = 0.0f; for (int j = 1; j <= N; ++j) { float s = scores[(i - 1) * N + (j - 1)]; float match = prev[j - 1] + s; float delete_ = prev[j] + gap_open; float insert = curr[j - 1] + gap_open; float m = fmaxf(0.0f, fmaxf(match, fmaxf(delete_, insert))); float sum_exp = expf((0.0f - m) * inv_t) + expf((match - m) * inv_t) + expf((delete_ - m) * inv_t) + expf((insert - m) * inv_t); float h = temperature * logf(sum_exp) + m; curr[j] = h; if (h > max_val) max_val = h; } // swap prev and curr for (int j = 0; j <= N; ++j) { prev[j] = curr[j]; } } output[b] = max_val; } torch::Tensor smith_waterman_cuda(torch::Tensor score_matrix, float gap_open, float temperature) { TORCH_CHECK(score_matrix.is_cuda(), "score_matrix must be a CUDA tensor"); TORCH_CHECK(score_matrix.dim() == 3, "score_matrix must be 3D"); int B = score_matrix.size(0); int M = score_matrix.size(1); int N = score_matrix.size(2); auto output = torch::zeros({B}, score_matrix.options()); const int threads = 256; const int blocks = (B + threads - 1) / threads; smith_waterman_kernel<<<blocks, threads>>>( score_matrix.data_ptr<float>(), output.data_ptr<float>(), gap_open, temperature, B, M, N); return output; } ''' functions = ['smith_waterman_cuda'] But wait: we must not write PYBIND11_MODULE. load_inline auto-generates one from functions list. The C++ code just defines the kernel and the host function. In load_inline, functions list takes names of host functions that return torch::Tensor or whatever. However, the host function signature should be compatible with the auto-generated pybind11 wrappers. Usually it should take torch::Tensor by value or const reference, and primitive types. The wrapper generated by load_inline handles overloads? Actually load_inline uses torch.utils.cpp_extension.load_inline, which creates a TORCH_EXTENSION_NAME module. It uses pybind11 to bind functions listed in `functions`. The functions are bound with `m.def("func_name", &func_name, ...)` or similar. So as long as the function is declared at global scope, it's fine. One issue: we need to include <torch/extension.h> which provides pybind11 and ATen. That's fine. Now, regarding temperature and gap_open: They are stored as nn.Module parameters? In the reference, they are floats passed to __init__, not parameters. So they are Python floats. We can just capture them in the ModelNew __init__ and pass them to the CUDA kernel. But wait: The reference uses `self.gap_open` and `self.temperature` as scalars. The kernel takes them as float args. One subtle numerical issue: logsumexp with temperature. Original code: stack = torch.stack([zero, match, delete, insert], dim=-1) H[:, i, j] = self.temperature * torch.logsumexp(stack / self.temperature, dim=-1) In PyTorch, logsumexp is computed in stable way: max + log(sum(exp(x - max))). So: max_val = max(0, match, delete, insert) H = t * ( max_val/t + log(sum(exp(x/t - max_val/t))) ) = max_val + t * log(sum(exp((x - max_val)/t))) Our kernel computes: m = max(0, match, delete, insert) sum_exp = exp((0 - m)/t) + ... h = t * log(sum_exp) + m This is mathematically identical and numerically stable. Good. One more check: In reference, `zero = torch.zeros(B, device=...)` but then `stack = torch.stack([zero, match, delete, insert], dim=-1)`. So zero is a tensor of shape (B,), while match etc are also shape (B,). For each batch element, the zero is 0.0. So yes. Computational graph: The reference is differentiable because H is built via differentiable ops. Our custom kernel is not differentiable out of the box. But the task says "Replace the forward computation". It doesn't mention backward. In these tasks, usually only the forward is needed for correctness and timing. The backward is not tested if they only call forward. So we don't need to define a backward. But if the evaluation code does backward, it might fail because the custom function doesn't have a backward. However, typical problem statements just say replace forward. The reference is a nn.Module. The evaluation might call `model(inputs)` and compare outputs. It doesn't mention gradient check. So it's fine. However, we should still be careful: `torch.utils.cpp_extension.load_inline` compiles at import time, which can take some seconds. That's okay. We need to write a complete Python file with class `ModelNew`. Let's draft: ```python import torch import torch.nn as nn from torch.utils.cpp_extension import load_inline cuda_source = r''' #include <torch/extension.h> #include <cuda_runtime.h> #include <math.h> __global__ void smith_waterman_kernel(const float* __restrict__ score_matrix, float* __restrict__ output, float gap_open, float temperature, int B, int M, int N) { int b = blockIdx.x * blockDim.x + threadIdx.x; if (b >= B) return; float prev[33]; float curr[33]; #pragma unroll for (int j = 0; j <= N; ++j) { prev[j] = 0.0f; } float max_val = 0.0f; const float inv_t = 1.0f / temperature; const int batch_stride = M * N; const float* scores = score_matrix + b * batch_stride; for (int i = 1; i <= M; ++i) { curr[0] = 0.0f; #pragma unroll for (int j = 1; j <= N; ++j) { float s = scores[(i - 1) * N + (j - 1)]; float match = prev[j - 1] + s; float delete_ = prev[j] + gap_open; float insert = curr[j - 1] + gap_open; float m = fmaxf(0.0f, fmaxf(match, fmaxf(delete_, insert))); float sum_exp = expf((0.0f - m) * inv_t) + expf((match - m) * inv_t) + expf((delete_ - m) * inv_t) + expf((insert - m) * inv_t); float h = temperature * logf(sum_exp) + m; curr[j] = h; if (h > max_val) max_val = h; } // swap rows #pragma unroll for (int j = 0; j <= N; ++j) { prev[j] = curr[j]; } } output[b] = max_val; } torch::Tensor smith_waterman_cuda(torch::Tensor score_matrix, float gap_open, float temperature) { TORCH_CHECK(score_matrix.is_cuda(), "score_matrix must be a CUDA tensor"); TORCH_CHECK(score_matrix.dim() == 3, "score_matrix must be 3D"); TORCH_CHECK(score_matrix.scalar_type() == torch::kFloat32, "score_matrix must be float32"); int B = score_matrix.size(0); int M = score_matrix.size(1); int N = score_matrix.size(2); auto output = torch::zeros({B}, score_matrix.options()); const int threads = 256; const int blocks = (B + threads - 1) / threads; smith_waterman_kernel<<<blocks, threads>>>( score_matrix.data_ptr<float>(), output.data_ptr<float>(), gap_open, temperature, B, M, N); C10_CUDA_KERNEL_LAUNCH_CHECK(); return output; } ''' smith_waterman = load_inline( name="smith_waterman", cpp_sources="", cuda_sources=cuda_source, functions=["smith_waterman_cuda"], extra_cuda_cflags=["-O3", "--use_fast_math"], verbose=False, ) class ModelNew(nn.Module): def __init__(self, gap_open=-1.0, temperature=1.0): super().__init__() self.gap_open = gap_open self.temperature = temperature def forward(self, score_matrix: torch.Tensor) -> torch.Tensor: return smith_waterman.smith_waterman_cuda(score_matrix, self.gap_open, self.temperature) ``` Need to ensure that the cuda_sources string uses R-string correctly. There are no backslash issues. But the kernel uses N as a variable, and the arrays are sized by constant 33. Since M and N are passed as args, but we use N to loop up to N and array size is fixed at 33. The reference has seq_m=32, seq_n=32, but the kernel should handle any M,N up to 33? Wait, the array is declared `float prev[33]`. That only works if N <= 32. The constructor args don't constrain M,N. The reference implementation works for any M,N. Should we support arbitrary M,N? The problem says accept the same constructor arguments, and forward has the same signature. The get_inputs returns (16,32,32). But the kernel should ideally be generic, or at least work for the inputs provided. If we hardcode 33, it only works for N<=32. Since N is 32 in the test, that's probably okay. However, sometimes tests vary sizes. The problem statement says "seq_m = 32, seq_n = 32, batch_size = 16". I think these are fixed. But to be safe, we could allocate shared memory or dynamic local memory. But dynamic local arrays in CUDA must use `extern __shared__` or malloc? No, local variable length arrays in device code are not standard in CUDA C++. We can use a maximum constant, like `const int MAX_N = 128` and declare `float prev[128]`, but this consumes registers/local memory per thread. If we launch many threads, it could blow up. But we launch only up to 256 threads. 256 threads * 128 floats * 2 arrays * 4 bytes = 256KB, which fits in local memory (which is backed by global memory but cached in L1). Actually local memory per thread is allocated from the global address space but cached. 256KB might be okay, but if N is actually 32, the compiler might still allocate 128*4 bytes per array per thread, i.e., 1KB per thread. For 256 threads, that's 256KB per block. Registers? Actually if size is too large for register file, it spills to local memory. That's fine. Alternatively, we can pass N as a template parameter? But load_inline uses runtime compilation; we could specialize with template but it's not necessary. To be robust, let's just use a fixed max size large enough, say `MAX_SIZE = 128`. But to be safe for any reasonable N (maybe 1024?), we could allocate shared memory per block and have each thread use a slice. But that's more complex. Given the problem statement explicitly sets M=N=32, hardcoding 33 is almost certainly fine. I will use a constant like 64 to be safe, and add a TORCH_CHECK that N <= 64 and M <= 64. That ensures correctness for the given test and prevents out-of-bounds if they test slightly larger. Wait, using a constant 64 means we do `#define MAX_N 64`. Then loops `for (int j=0; j<=N; ++j)`. The array is `float prev[MAX_N+1];`. That's fine. Also, we need to be careful: if we add `#pragma unroll` on loops where the bound is variable (N is runtime variable), `#pragma unroll` might not be valid; the compiler will just ignore it or error. Better not unroll loops over N, or unroll small constant loops. Since `for (int j = 0; j <= N; ++j)` with N being a runtime variable, `#pragma unroll` is not standard. We can remove unroll pragmas or only apply to loops with constant bounds (like if we do `for (int j = 0; j < MAX_N; ++j)` but conditionally break). Actually `#pragma unroll` can be used with a constant integer after it, like `#pragma unroll 4`, but the bound must be known at compile time for full unroll. Better just not use it for variable loops. The compiler will optimize anyway. Actually, we can use shared memory for H to avoid local memory. But local memory is fine. Let's refine the kernel. We can use float prev[65], curr[65] to be safe for N up to 64. ```cpp #define MAX_N 64 __global__ void smith_waterman_kernel(const float* __restrict__ score_matrix, float* __restrict__ output, float gap_open, float temperature, int B, int M, int N) { int b = blockIdx.x * blockDim.x + threadIdx.x; if (b >= B) return; if (N > MAX_N || M > MAX_N) return; // or assert float prev[MAX_N + 1]; float curr[MAX_N + 1]; for (int j = 0; j <= N; ++j) { prev[j] = 0.0f; } float max_val = 0.0f; const float inv_t = 1.0f / temperature; const int batch_stride = M * N; const float* scores = score_matrix + b * batch_stride; for (int i = 1; i <= M; ++i) { curr[0] = 0.0f; for (int j = 1; j <= N; ++j) { float s = scores[(i - 1) * N + (j - 1)]; float match = prev[j - 1] + s; float delete_ = prev[j] + gap_open; float insert = curr[j - 1] + gap_open; float m = fmaxf(0.0f, fmaxf(match, fmaxf(delete_, insert))); float sum_exp = expf((0.0f - m) * inv_t) + expf((match - m) * inv_t) + expf((delete_ - m) * inv_t) + expf((insert - m) * inv_t); float h = temperature * logf(sum_exp) + m; curr[j] = h; if (h > max_val) max_val = h; } for (int j = 0; j <= N; ++j) { prev[j] = curr[j]; } } output[b] = max_val; } ``` Wait: the score_matrix pointer is per batch `b * batch_stride`. We assume memory layout is contiguous. `score_matrix.data_ptr<float>()` is contiguous if the tensor is contiguous. If we want to be safe, we can use stride. But the inputs are created with `torch.randn`, so contiguous. However, inside the host wrapper we can also ensure `score_matrix = score_matrix.contiguous()`. One more thing: We need to include `C10_CUDA_KERNEL_LAUNCH_CHECK()` for safety? Not required but good practice. However, in some torch versions, `C10_CUDA_KERNEL_LAUNCH_CHECK` is available. I think it's safe. Alternatively, we can use `AT_CUDA_CHECK(cudaGetLastError());` after launch. Also, `--use_fast_math` flag makes `expf` and `logf` faster but slightly less accurate. The tolerance is fp32 atol=1e-4, so fast math should be okay. But to be safe, maybe we should not use `--use_fast_math` if accuracy is borderline. However, `expf` and `logf` with fast math usually have relative error within ~1e-5 to 1e-6, which is acceptable. But let's keep it without to be safe on correctness. Actually, correctness tolerance is 1e-4, which is generous. So fast math is fine, but might not be needed. We can try both. Let's write the code and compile it. Actually, I should also consider using `torch` ops for reduction to avoid kernel compilation errors? No, kernel is straightforward. Let's draft the final Python code and test with compile_kernel.