We need to replace a PyTorch model forward with a custom CUDA kernel. The reference: - Input x: shape (R, S_MAX) where R=4, S_MAX=20. - A boolean mask: shape (R, S_MAX), with _MASK_DATA. - forward: convert mask to same dtype and device as x, multiply x * mask, sum over dim=0, output shape (S_MAX,). - The computation is just masked sum along rows. Goal: write a custom CUDA kernel implementing this, compile with load_inline, and call it in ModelNew.forward. Need to be numerically equivalent (fp32). Since mask is just 0 or 1 entries, equivalently output[j] = sum_{r=0}^{R-1} x[r,j] * mask[r][j]. Given R=4 and S_MAX=20, this is tiny: only 80 elements. But we need a kernel that is faster than the PyTorch reference. Given the tiny size, overhead may dominate. We need to be careful. PyTorch reference likely does two operations: to(dtype) on CPU? Actually mask is a buffer registered on whatever device default is? In __init__, register_buffer without device argument will create tensor on CPU. In forward, mask.to(dtype=x.dtype, device=x.device) moves to device and casts, then multiplication and sum. That's some overhead. Our kernel should accept x and mask (maybe as a float/int tensor on same device already) and compute output directly. We can avoid the .to() overhead by precomputing mask as float or int on the device in __init__ or forward. But register_buffer can be on CPU, but we can move to device in constructor or in forward. Since ModelNew must accept same constructor arguments as Model (none), but we can set our own buffer. The reference registers a bool buffer on CPU. In ModelNew __init__, we can register the mask buffer on CPU as bool too (since persistent=False). But forward can convert it once. However, to be faster, we can maybe register the mask as a float tensor on the device? But the test harness may instantiate ModelNew on CPU then move to device? Actually they call get_init_inputs (empty) and instantiate model. Usually they just instantiate ModelNew() and then maybe move to device? Not sure. The reference does not have a .to(device) call inside __init__. It just registers buffer. So the buffer is on CPU. In forward, it does .to(dtype=x.dtype, device=x.device). That means each forward call moves mask to device. In ModelNew, we can do the same: store the bool mask as buffer, then in forward convert to x's dtype and device. But that still has overhead. However, since the buffer is small, the overhead is tiny. To be faster, we can perhaps avoid .to() by caching the mask on the device after first call, perhaps storing it as a buffer on the correct device? But __init__ doesn't know the device. However, we can just call self.shard_mask.to(device=x.device, dtype=x.dtype) and maybe store it in a dict per device to avoid repeated conversion, but since it's tiny, it's fine. But the real cost is the kernel launch vs simple elementwise ops. For 80 elements, PyTorch likely runs on tensor cores or vectorized loads on GPU quickly. However, the kernel launch overhead may be similar. To beat PyTorch, we need to fuse the multiplication and sum into a single kernel, and perhaps process each column with a warp or thread. R=4, S_MAX=20. So we have 20 columns, each sum of up to 4 elements. A simple kernel: each thread handles one column j, loads the 4 x values, multiplies by 4 mask values (or uses branching), accumulates, writes output. Since there are only 20 columns, we could launch 20 threads, or maybe 1 block of 32 threads. But warp divergence if masks vary. However, the mask data is static, we could pass mask as an integer bitmask per column, or as a uint8 array. Or we can bake the mask into the kernel since it's constant? But the kernel must be generic enough for the given shape and mask. We can pass the mask as a bool tensor, but then we need to read it. Since it's constant per model instance, we could pass it as a uint32 bit mask: each column j has a 4-bit mask. But the reference mask is per row and column. So we can precompute a uint32 mask in __init__ or forward. However, the mask shape is (4,20). In the kernel, each thread j reads mask[0][j], [1][j], [2][j], [3][j]. That's 4 bool loads. Since bools are 1 byte each, that's 20 bytes, negligible. But reading from global memory might cause overhead. We can store the mask in constant memory or pass as a regular tensor. Alternatively, we can use shared memory: load x and mask into shared memory, then compute. But with 80 elements, it's not much. But maybe the PyTorch reference is using the CPU for mask conversion? Actually mask is on CPU, .to(dtype=x.dtype, device=x.device) moves it to GPU each forward call. That involves a cudaMemcpyAsync of 80 bytes from CPU to GPU? That is tiny but may have overhead. Our kernel can just pass the mask as a device tensor, avoiding the per-call CPU->GPU copy. To do that, we can in __init__ create a float mask on CPU? No, we can create it on CPU as bool and then in first forward call move to device and cache it in a Python attribute (like self._mask_gpu). Since ModelNew may be called multiple times, caching is helpful. But correctness requires handling device changes? The test may just call forward once per input. But let's do simple: in forward, if self._mask_gpu is None or device mismatched, move mask to device and store. Then call kernel with x and mask. But is caching allowed? Yes, since mask is constant. Given that the problem is small, a fused kernel should be faster because it avoids two separate kernels (mul and sum) and the CPU->GPU copy. But we must also consider kernel overhead. If we use load_inline, compilation overhead is at module load time, not at runtime. At runtime, we just call the C++ function that launches the kernel. Given S_MAX=20, we can have a kernel where each block has 32 threads, each thread handles one column. But there are only 20 columns, so 12 threads idle. That's okay. Or we can have 20 threads in a block. Better: since R=4, each warp can handle 8 columns with 4 threads each? Not necessary. Let's design a kernel: ``` __global__ void masked_sum_kernel(const float* x, const float* mask, float* out, int R, int S) { int j = blockIdx.x * blockDim.x + threadIdx.x; if (j >= S) return; float sum = 0.0f; for (int r = 0; r < R; ++r) { sum += x[r * S + j] * mask[r * S + j]; } out[j] = sum; } ``` But wait: mask is bool. If we store mask as float (0.0f or 1.0f), then multiply is fine. If we store as uint8 or bool, we need to cast. We can preconvert mask to float. Since R and S are known constants (4 and 20), we can unroll the loop and maybe use vectorized loads. We can load x and mask for each row as float4? Since S=20, not a multiple of 4, but we can process each column separately. Alternatively, we can transpose the data in memory: x is (R, S) row-major, so columns are strided by S. Each thread accesses stride S memory. With R=4, each thread reads 4 floats with stride 20. That's not coalesced across threads: thread j reads x[j], x[S+j], x[2S+j], ...; thread j+1 reads x[j+1], x[S+j+1], etc. So within a warp, consecutive threads read consecutive elements for each row: for row r, threads read x[r*S + j], x[r*S + j+1], ... which is coalesced. So memory access is fine. So a simple kernel is good. But is there a way to be even faster? For such a small kernel, occupancy is not a bottleneck. The main thing is to minimize kernel launch overhead and per-thread work. Since we have only 20 outputs, we can process each column in a thread and loop over R=4. That's minimal. We should also make sure we pass the mask as a float tensor on device, not bool on CPU. So in ModelNew.__init__, we register the bool buffer just like reference, but also maybe create a float version on CPU? Actually we can just compute the float mask in forward and cache it. Let's write the code. Potential issues: - load_inline: we need to provide cuda_sources and functions. The function should be a C++ wrapper that accepts torch::Tensor and returns torch::Tensor. - We must ensure the tensors are contiguous and on CUDA. The test device is default_device() which is likely CUDA. - Since R and S_MAX are constants in Python, we could hardcode them in the kernel or pass as arguments. It might be cleaner to pass as arguments. - For performance, we may want to use `blockDim.x = 32`, `gridDim.x = 1` (since 20 columns). So launch 1 block of 32 threads. - Check that the kernel handles j >= S. But wait: the reference uses `mask.to(dtype=x.dtype, device=x.device)`. The mask is bool. If we pass a float mask, it should be equivalent because bool -> float gives 0.0 or 1.0. Let's write the C++ wrapper: ```cpp torch::Tensor masked_sum_forward(torch::Tensor x, torch::Tensor mask) { auto out = torch::empty({S_MAX}, x.options()); const int threads = 32; const int blocks = 1; masked_sum_kernel<<<blocks, threads>>>( x.data_ptr<float>(), mask.data_ptr<float>(), out.data_ptr<float>(), R, S_MAX); return out; } ``` Wait: S_MAX is a macro? We can define it via `cuda_sources` with `#define S_MAX 20` etc. Alternatively, we can pass R and S as arguments to kernel. Since R and S are small, we can template them or just pass. Better to define them as constants in CUDA code. Let's do: ```cpp #define R 4 #define S_MAX 20 __global__ void masked_sum_kernel(const float* __restrict__ x, const float* __restrict__ mask, float* __restrict__ out) { int j = threadIdx.x + blockIdx.x * blockDim.x; if (j >= S_MAX) return; float sum = 0.0f; #pragma unroll for (int r = 0; r < R; ++r) { sum += x[r * S_MAX + j] * mask[r * S_MAX + j]; } out[j] = sum; } torch::Tensor masked_sum_forward(torch::Tensor x, torch::Tensor mask) { TORCH_CHECK(x.is_cuda(), "x must be CUDA"); TORCH_CHECK(mask.is_cuda(), "mask must be CUDA"); auto out = torch::empty({S_MAX}, x.options()); masked_sum_kernel<<<1, 32>>>(x.data_ptr<float>(), mask.data_ptr<float>(), out.data_ptr<float>()); return out; } ``` This is simple. But we need to ensure x and mask are contiguous and float. Since mask is float (0.0 or 1.0), multiplication is okay. In forward, we can ensure mask is float and contiguous. But wait: we need to handle the case where mask is originally bool on CPU in ModelNew's buffer. We can convert it in forward: ```python if not hasattr(self, '_mask_float') or self._mask_float.device != x.device: self._mask_float = self.shard_mask.to(dtype=x.dtype, device=x.device).contiguous() return _masked_sum_forward(x, self._mask_float) ``` This avoids CPU->GPU copy on every forward call. Will this be faster than reference? The reference does `m = self.shard_mask.to(dtype=x.dtype, device=x.device)` each call, which is a CPU->GPU copy (or at least a conversion). Then `x * m` launches a kernel, then `.sum(dim=0)` launches another kernel (reduction). So 3 operations. Our code does one cached conversion (only first call) and one fused kernel. Even if conversion happens each time, it's one operation vs two plus copy. So should be faster. But we must be careful: `to` with device change copies from CPU to GPU. The reference does that each forward. We do it once. So definitely faster. But what if the test doesn't reuse the same ModelNew instance? It may instantiate a new ModelNew for each forward call? Usually the benchmark runs multiple trials. But within a trial, they may instantiate once and call forward multiple times? In KernelBench, they instantiate the model once, then run forward many times for timing. So caching is beneficial. Now, compile and correctness check. Let's write the kernel code. Need to be careful: `load_inline` auto-generates pybind11 module. So we just provide the function. Also ensure we include necessary headers: `<torch/extension.h>` and `<cuda_runtime.h>`. Let's draft the entire Python file. ```python import torch import torch.nn as nn from torch.utils.cpp_extension import load_inline R = 4 S_MAX = 20 _MASK_DATA = [ [1, 1, 1, 1, 1, 1, 1, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0], [1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 0, 0, 0, 0, 0, 0, 0, 0], [1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0], [1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 0, 0], ] cuda_source = """ #define R 4 #define S_MAX 20 __global__ void masked_sum_kernel(const float* __restrict__ x, const float* __restrict__ mask, float* __restrict__ out) { int j = blockIdx.x * blockDim.x + threadIdx.x; if (j >= S_MAX) return; float sum = 0.0f; #pragma unroll for (int r = 0; r < R; ++r) { sum += x[r * S_MAX + j] * mask[r * S_MAX + j]; } out[j] = sum; } torch::Tensor masked_sum_forward(torch::Tensor x, torch::Tensor mask) { auto out = torch::empty({S_MAX}, x.options()); masked_sum_kernel<<<1, 32>>>(x.data_ptr<float>(), mask.data_ptr<float>(), out.data_ptr<float>()); return out; } """ cpp_source = """ torch::Tensor masked_sum_forward(torch::Tensor x, torch::Tensor mask); """ # Actually load_inline requires cuda_sources and cpp_sources separate. # The function is declared in cuda_source? Actually we can put everything in cuda_sources. # load_inline will compile .cu file with the cuda_sources content. # The C++ wrapper can be in cpp_sources or cuda_sources. # Usually we put the kernel in cuda_sources and the wrapper in cpp_sources. # But we can put both in cuda_sources since .cu allows host code. # However, the pybind11 binding is auto-generated from functions list, so we need to have the function defined in the compiled code. # Let's put everything in cuda_sources. custom_op = load_inline( name="masked_sum", cpp_sources="", # no extra C++ needed cuda_sources=cuda_source, functions=["masked_sum_forward"], extra_cuda_cflags=["-O3"], verbose=False, ) class ModelNew(nn.Module): def __init__(self): super().__init__() self.register_buffer( "shard_mask", torch.tensor(_MASK_DATA, dtype=torch.bool), persistent=False, ) self._mask_float = None def forward(self, x: torch.Tensor) -> torch.Tensor: if self._mask_float is None or self._mask_float.device != x.device: self._mask_float = self.shard_mask.to(dtype=x.dtype, device=x.device).contiguous() return custom_op.masked_sum_forward(x, self._mask_float) ``` Wait, but `load_inline` with `cpp_sources=""` might be okay. But perhaps we need to include `<torch/extension.h>` in cuda_sources? The nvcc compiler includes it automatically? Actually we may need to include torch headers. Let's add `#include <torch/extension.h>` at the top of cuda_source. Also need `#include <cuda_runtime.h>`. Let's refine: ```cpp #include <torch/extension.h> #include <cuda_runtime.h> #define R 4 #define S_MAX 20 __global__ void masked_sum_kernel(const float* __restrict__ x, const float* __restrict__ mask, float* __restrict__ out) { int j = blockIdx.x * blockDim.x + threadIdx.x; if (j >= S_MAX) return; float sum = 0.0f; #pragma unroll for (int r = 0; r < R; ++r) { sum += x[r * S_MAX + j] * mask[r * S_MAX + j]; } out[j] = sum; } torch::Tensor masked_sum_forward(torch::Tensor x, torch::Tensor mask) { auto out = torch::empty({S_MAX}, x.options()); const int threads = 32; const int blocks = 1; masked_sum_kernel<<<blocks, threads>>>( x.data_ptr<float>(), mask.data_ptr<float>(), out.data_ptr<float>()); return out; } ``` Will this compile? Yes. We need to ensure the function signature is correct. Also, we need to check if `x.options()` includes dtype and device. It will be float, cuda. Now, we should compile and test. Let's do that. One thing: the `load_inline` may require `with_cuda=True`. Actually by default it detects CUDA if `cuda_sources` is provided. So fine. Also, since we are using `extra_cuda_cflags`, note that `load_inline` might also need `extra_include_paths` etc. But default should work. Potential issue: `load_inline` caches the compiled extension by name. The name "masked_sum" is okay. But there is a subtlety: Deterministic behavior. The kernel uses float operations: multiplication and addition. Since there are only up to 4 terms, sum order is same as reference (row-major). The reference does `x * m` elementwise then sum over dim=0. PyTorch's sum reduction likely uses deterministic order? Actually for such small dimensions, PyTorch's `sum(dim=0)` on GPU may use a tree reduction or sequential reduction depending on implementation. But with R=4, the sum is over 4 elements per column. PyTorch likely just does a simple reduction. The order of operations in our kernel is row-by-row: r=0..3. The reference does `x * m` first, then sum. Since multiplication is elementwise, the sum also likely accumulates in row order (since the reduction is along dimension 0, and the data is stored row-major, each column's elements are strided by S_MAX. The reduction kernel might load them and sum in some order. For 4 elements, it might just unroll. But to guarantee equivalence within tolerance, we need to match. Since the values are small (random normal), the sum of 4 floats in any order should be within 1e-4 relative? For fp32, reordering of 4 additions can cause slightly different results, but tolerance is 1e-4 absolute and relative. Is that safe? Let's think. Consider typical values ~1.0. Adding 4 numbers can have error on order of machine epsilon ~1e-7 times magnitude. So difference due to order is ~1e-7, well within 1e-4. For larger values, say 10.0, error ~1e-6. Still fine. For pathological cases, but random normal has mean 0 std 1. So okay. But we also need to consider that the mask has zeros. Multiplying by zero yields exactly zero (no rounding). So masked terms are zero. Thus our kernel should be correct. However, to be safe, we could explicitly make the kernel sum in the same conceptual order as PyTorch. But we don't know PyTorch's exact order. It's likely sequential along the dimension. Since our storage is row-major, dimension 0 stride is S_MAX. A reduction kernel might process each output element independently. For 4 elements, each thread might read all 4 and sum them. The order may be sequential in memory: address 0, 20, 40, 60. That corresponds to r=0,1,2,3. Our kernel also does r=0..3. So order matches. Now, speed: A kernel with 1 block of 32 threads, each reading 4 floats and 4 floats, computing, writing 1 float. Very lightweight. Potential improvement: process the whole thing with a single warp and no loop, by manually unrolling: ```cpp float sum = 0.0f; sum += x[0*S_MAX + j] * mask[0*S_MAX + j]; sum += x[1*S_MAX + j] * mask[1*S_MAX + j]; sum += x[2*S_MAX + j] * mask[2*S_MAX + j]; sum += x[3*S_MAX + j] * mask[3*S_MAX + j]; ``` The compiler will unroll with `#pragma unroll`. Good. Another improvement: vectorize loads. For each row, we can load float4. But S_MAX=20, not multiple of 4 without tail. Could load 5 float4s to cover 20. Then each thread processes 4 columns? Let's see: if we load x row by row as float4, each thread could handle 1 column still but we would need to load scalar. Actually we can have 5 threads per row, each processing 4 columns. But we have only 20 columns. If we have 5 threads, each handling 4 columns (via float4 load), we need 5 threads. But a warp has 32 threads, so 27 idle. That's okay. But then we need to reduce within warp? Actually each thread handles 4 columns independently, writing 4 outputs. That might be more efficient due to vectorized memory access: each thread reads a float4 for x and a float4 for mask. But since columns are consecutive, thread 0 reads columns 0-3, thread 1 reads 4-7, etc. This is fully coalesced and uses 128-bit loads. But then we have only 5 active threads. The rest are idle but threads in a warp execute together. Divergence is not an issue if we just have a condition `if (threadIdx.x < 5)`. The first 5 threads do work, the rest exit. That's fine. However, the kernel launch overhead is likely the dominant factor. The difference between scalar and vectorized loads is negligible for 80 elements. So we can stick with the simple version. But maybe we can make it even faster by using a warp-level primitive? No. Wait, one more thing: The PyTorch reference does `self.shard_mask.to(dtype=x.dtype, device=x.device)` which returns a new tensor on device. But `self.shard_mask` is a buffer (persistent=False). In our ModelNew, we also register the same buffer. But we also have a python attribute `_mask_float` that caches the converted mask. This is fine. However, the reference's `shard_mask` is a `torch.bool` tensor. Our `self._mask_float` is a `torch.float32` tensor. That's fine. Let's compile and test.

Good, the current kernel passes tests. H100 is a very powerful GPU. The workload is tiny (R=4, S_MAX=20, so 80 floats). Overheads will dominate. The custom kernel is already likely faster than the reference because: 1. No extra temporary tensor (the PyTorch `x * m` creates a new tensor before sum). 2. Flattened loop with no extra data movement. However, since it's so tiny, caching the mask and just calling the kernel is key. The current kernel launches a single block with 32 threads. This is fine, but each thread only does 4 multiplications and 3 adds. This is memory bound with very little work. Performance on H100 will be basically limited by kernel launch overhead. Can we reduce the kernel launch overhead further? Not really. Maybe we could compile it with inline assembly? No, let's first check if it's likely faster than the reference. The reference needs: - Convert bool mask to float (type conversion, memory allocated) - Elementwise multiply (x * m) -> allocate one float tensor (R*S_MAX = 80 floats), then elementwise mul - sum(dim=0) -> reduce over one axis So the reference does 2 kernel launches and allocates memory. Our kernel does 1 kernel launch with no allocated intermediate. This should be faster. One optimization: Currently the mask is stored as a bool tensor, converted on first use to float, and cached. That's fine. But for even more performance, we could pre-bake the mask into constant memory or just hardcode it using a switch or compile-time constants. Since R=4 and S_MAX=20 and it's fixed data, we could directly embed the mask values in the kernel. This would completely eliminate the mask tensor argument and reads from global memory. In fact, since it's such a small fixed mask, let me hardcode the mask values as constants in the kernel. This would make it even faster and avoid reading the mask from global memory. Wait, but we still need the masks. Let me think about whether hardcoding provides any real benefit. For a kernel that processes 80 floats, reading 80 floats of x and 80 floats of mask, the memory traffic is trivial. Eliminating the mask read is not the bottleneck; kernel launch overhead is. Nevertheless, a smaller kernel signature and no mask memory check could be marginally faster, but not a big deal. Also, I should check if there is a better way to pass the mask. Since the mask is known at compile time, I could make the kernel take no mask argument and use the constants directly. That would also mean we don't need to pass the mask tensor to the kernel. Also note: since S_MAX=20 and R=4, we only need 20 threads. Launching 32 threads means 12 are idle. Since a warp is 32 threads, this is fine. But we could also compute this with just one thread and a loop - since there's so little work, that might even be faster (avoids any warp divergence and synchronization). Let me think... 20 threads vs 1 thread doesn't matter much. Actually, let me try a version that hardcodes the mask to remove the extra memory read entirely. It might be theoretically cleaner but the practical impact is small. Also, another idea: Instead of using a custom CUDA kernel, perhaps we could use `torch.where` or something? No, the idea is to write a custom kernel. I will stick with the current implementation, but consider whether we should use `blockDim.x` as something different. Since there's a launch bound of 32 threads for just 20 outputs, it's okay. One issue: `load_inline` might be caching the compiled module, so if I need to iterate, I should use a unique name each time. For the final submission, I'll name it appropriately. Wait, I should also ensure correctness with the static check. Let me do that.