Let me understand the task: 1. We have a PyTorch model that performs a masked sum along the virtual rank dimension 2. Input `x` is shape (R, S_MAX) = (4, 20) 3. There's a boolean mask of the same shape 4. Output is (S_MAX,) = (20,) - sum over dimension 0 where mask is True The reference implementation: ```python def forward(self, x: torch.Tensor) -> torch.Tensor: m = self.shard_mask.to(dtype=x.dtype, device=x.device) return (x * m).sum(dim=0) ``` So for each column j, we sum x[r, j] * mask[r, j] for all r where mask[r, j] is True. Since this is a small operation (R=4, S_MAX=20), the main overhead is probably: 1. Creating the mask tensor with the right dtype and device 2. The element-wise multiplication 3. The reduction A custom CUDA kernel could: 1. Store the mask as constant memory or pass it directly 2. Do the masked sum directly without intermediate tensors 3. Each thread handles one output element (column) Let me write a simple kernel: ```python import torch import torch.nn as nn from torch.utils.cpp_extension import load_inline # Define the CUDA kernel cuda_source = """ __global__ void masked_sum_kernel(const float* x, const bool* mask, float* output, int R, int S_MAX) { int j = blockIdx.x * blockDim.x + threadIdx.x; if (j < S_MAX) { float sum = 0.0f; for (int r = 0; r < R; r++) { if (mask[r * S_MAX + j]) { sum += x[r * S_MAX + j]; } } output[j] = sum; } } torch::Tensor masked_sum_cuda(torch::Tensor x, torch::Tensor mask) { int R = x.size(0); int S_MAX = x.size(1); auto output = torch::zeros({S_MAX}, x.options()); int threads = 256; int blocks = (S_MAX + threads - 1) / threads; masked_sum_kernel<<<blocks, threads>>>( x.data_ptr<float>(), mask.data_ptr<bool>(), output.data_ptr<float>(), R, S_MAX ); return output; } """ cpp_source = "torch::Tensor masked_sum_cuda(torch::Tensor x, torch::Tensor mask);" masked_sum_ext = load_inline( name="masked_sum_ext", cpp_sources=cpp_source, cuda_sources=cuda_source, functions=["masked_sum_cuda"], extra_cuda_cflags=["-O3"], ) class ModelNew(nn.Module): def __init__(self): super().__init__() self.register_buffer( "shard_mask", torch.tensor(_MASK_DATA, dtype=torch.bool), persistent=False, ) def forward(self, x: torch.Tensor) -> torch.Tensor: mask = self.shard_mask.to(device=x.device) return masked_sum_ext.masked_sum_cuda(x, mask) ``` I need to optimize the kernel further. The mask is constant and small, so I should avoid passing it as a tensor and instead hardcode it in the kernel. For R=4 and S_MAX=20, the workload is quite small—only 80 total elements and 20 output elements. I can launch 20 threads, one per column, and have each thread loop over the 4 rows. Since the mask stays constant, I can embed it directly in constant memory or compile it into the kernel itself, eliminating the need to pass it as an argument. ```cpp __constant__ bool d_mask[4 * 20]; ``` But I still need to copy the mask to device memory initially. A cleaner approach would be to pass the mask values directly to the kernel as arguments or encode them as a bit pattern. Since the mask is fixed and known at compile time, I can hardcode it into the kernel logic itself by encoding the mask pattern for each rank directly in the code. Since the mask is fixed, I can encode it directly in the kernel rather than passing it as a tensor. Each column has a known contribution count from the 4 rows, so I'll compute the sum by iterating through the valid indices for each column position. For position j=0, all four rows contribute, so I sum x[0][0] through x[3][0]. This pattern repeats for positions where the mask is active, while later positions skip certain rows based on the mask structure. This approach works for our specific mask, but I should generalize it to handle arbitrary masks stored in constant memory instead. The mask tensor is registered in `__init__`, so it persists across forward calls. Copying it to the GPU on the first invocation and keeping it there would be efficient, though the reference implementation transfers it each time anyway. Let me check if compiling with the hardcoded mask approach works first. Since the mask values are fixed at compile time and R=4 with S_MAX=20, I can embed the mask data directly into the kernel. Each column j gets contributions from specific rows based on which mask[r][j] entries are set to 1. For columns 0-7, all rows contribute. Columns 8-10 get contributions from rows 1, 2, and 3. Column 11 includes rows 1 and 3. Columns 12-17 pull only from row 3. Column 18 has no contributions, and column 19 is empty. The reference computes `(x * m).sum(dim=0)`, converting the boolean mask to the input dtype before multiplication. I need to match this exact behavior in my kernel. _MASK_DATA[3][j]) { sum += x[3 * S_MAX + j]; } } output[j] = sum; } } ``` This gets unwieldy quickly. A cleaner solution would use constant memory to store the mask values efficiently. Let me refactor using constant memory instead: ```cpp #include <cuda.h> #include <torch/extension.h> // Mask stored in constant memory __constant__ bool d_mask[4 * 20]; __global__ void masked_sum_kernel(const float* x, float* output, int R, int S_MAX) { int j = blockIdx.x * blockDim.x + threadIdx.x; if (j < S_MAX) { float sum = 0.0f; for (int r = 0; r < R; r++) { if (d_mask[r * S_MAX + j]) { sum += x[r * S_MAX + j]; } } output[j] = sum; } } void set_mask(const bool* mask, int size) { cudaMemcpyToSymbol(d_mask, mask, size * sizeof(bool)); } torch::Tensor masked_sum_cuda(torch::Tensor x) { int R = x.size(0); int S_MAX = x.size(1); auto output = torch::zeros({S_MAX}, x.options()); int threads = 256; int blocks = (S_MAX + threads - 1) / threads; masked_sum_kernel<<<blocks, threads>>>( x.data_ptr<float>(), output.data_ptr<float>(), R, S_MAX ); return output; } ``` The kernel reads from constant memory to check mask values, then accumulates valid elements into the output. Setting the mask once via `cudaMemcpyToSymbol` avoids repeated transfers, and the kernel itself stays simple—just iterating through rows for each column and summing masked contributions. I need to handle the mask setup in Python during extension initialization. The mask should be copied to the GPU once at module load time, not on every forward pass. The `ModelNew` class would then just call the kernel without managing the mask tensor. I could expose a `set_mask` function through the extension to initialize the constant memory from Python, or potentially hardcode the mask values directly in CUDA as a constant array. A direct hardcoded approach works cleanly since the mask dimensions are fixed at compile time. I can define it as a constant array in the CUDA source and reference it in the kernel without any runtime memory transfers. The kernel iterates through each row, checking if the mask value at that position is non-zero, then accumulates the sum accordingly. This eliminates the need for constant memory entirely and keeps the implementation straightforward. I need to ensure the mask is defined in a way that's accessible from both host and device code, then set up the kernel launch with appropriate thread and block dimensions based on the sequence length. After launching the kernel, I return the output tensor containing the masked sums. But I should verify whether the mask indices are actually row-major or column-major based on the problem specification. Looking at the mask structure, _MASK_DATA[0] shows the first row with 8 valid entries, which helps clarify the layout I'm working with. Each row represents a different rank's valid columns, so rank 0 has 8 valid positions, rank 1 has 12, rank 2 has 10, and rank 3 has 18. When computing the output for each column j, I need to sum x[r][j] across all ranks where that column is valid for that particular rank. For j=0, all four ranks contribute since every rank has position 0 marked as valid. I should proceed with implementation and testing rather than over-analyzing. The mask layout appears to be row-major (mask[r][j] corresponds to _MASK_DATA[r][j]), so my kernel indexing of mask[r * S_MAX + j] should work correctly. Let me build the full solution and validate it through compilation.

The kernel is correct. Now let me profile it to see if there's room for optimization:

The memory access pattern has 3x wasted bandwidth. The issue is that threads access x[0*20+j], x[1*20+j], x[2*20+j], x[3*20+j] - which are not contiguous (they have stride 20 between each access for the same r). Let me reorganize to read in a coalesced manner. Actually, I realize the main issue: threads 0-19 read x[j], x[20+j], x[40+j], x[60+j] which ARE contiguous within each row. But the 4 loads are interleaved. The L2 cache should handle this well for such a small dataset. The real issue is that for such a tiny kernel (only 20 output elements), kernel launch overhead dominates. The PyTorch reference also has to do the same work. Let me check the memory coalescing issue more carefully. When thread j reads x[0*20 + j], thread j+1 reads x[0*20 + j+1]. These are contiguous, so the first load should coalesce perfectly. Similarly for all 4 loads. The 3x overhead might be from the profiler's measurement. Actually, the profile says "sectors/req=3.0" which suggests 128-byte cache lines. 20 threads * 4 bytes = 80 bytes, which is less than one cache line. So when loading 80 bytes across the warp, we need ceil(80/128) = 1 sector ideally. But the 3x suggests something is wrong. Wait, 32 threads are launched but only 20 do work. That's fine. Let me try a different approach: use vectorized loads to read multiple elements at once, or reorganize data access. Actually, for this problem size, the best approach might be to have a single thread do all the work to minimize latency, or use warp shuffle. Let me try a simpler approach: have each thread do multiple columns (less threads, more work per thread).

I need to verify correctness first, but I realize my hardcoded mask logic might be incorrect. Let me reconsider: Looking at the mask data: ``` Row 0: [1, 1, 1, 1, 1, 1, 1, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0] # 8 ones at start Row 1: [1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 0, 0, 0, 0, 0, 0, 0, 0] # 12 ones at start Row 2: [1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0] # 10 ones at start Row 3: [1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 0, 0] # 18 ones at start ``` So per column: - Col 0-7: all 4 rows have 1 - Col 8-9: rows 1,2,3 have 1 (row 0 has 0) - Col 10-11: rows 1,3 have 1 (rows 0,2 have 0) - Col 12-17: only row 3 has 1 - Col 18-19: all rows have 0 Wait, let me double check by looking at each column: Col 0: [1,1,1,1] -> all 4 valid Col 1: [1,1,1,1] -> all 4 valid ... Col 7: [1,1,1,1] -> all 4 valid Col 8: [0,1,1,1] -> rows 1,2,3 valid (3) Col 9: [0,1,1,1] -> rows 1,2,3 valid (3) Col 10: [0,1,0,1] -> rows 1,3 valid (2) Col 11: [0,1,0,1] -> rows 1,3 valid (2) Col 12: [0,0,0,1] -> row 3 valid (1) Col 13: [0,0,0,1] -> row 3 valid (1) ... Col 17: [0,0,0,1] -> row 3 valid (1) Col 18: [0,0,0,0] -> none valid Col 19: [0,0,0,0] -> none valid This matches what I wrote in the kernel! Let me verify correctness.