Let me analyze this task: 1. The reference implementation does a fixed permutation along the feature axis: `y[b, j] = x[b, perm[j]]` 2. This is a simple operation that selects columns from the input tensor based on a permutation. The key insight is that this is essentially an index selection operation. The reference uses `x[:, PERM.to(device=x.device)]` which creates a new tensor by gathering columns. For a custom CUDA kernel, I need to: 1. Copy elements from `x[b, perm[j]]` to `y[b, j]` for all batch elements b and all output indices j Let me think about the optimal kernel design: - Input: x of shape (B, N) - Output: y of shape (B, N) - Permutation: PERM of shape (N,) For each element `y[b, j]`, we read from `x[b, perm[j]]`. Since B=10 and N=36, this is a relatively small tensor. But I should write a general kernel. The naive approach: - Each thread handles one output element `(b, j)` - Compute global index, determine b and j - Read `x[b, perm[j]]` and write to `y[b, j]` Let me think about memory access patterns: - Writing to y is contiguous in memory if threads are organized properly - Reading from x uses perm[j] as the column index, which is a scattered read Actually, for such small sizes, coalescing might not be as critical. But let me design a good kernel anyway. A simple kernel would be: ```cuda __global__ void permute_kernel(const float* __restrict__ x, float* __restrict__ y, const int* __restrict__ perm, int B, int N) { int idx = blockIdx.x * blockDim.x + threadIdx.x; int total = B * N; if (idx < total) { int b = idx / N; int j = idx % N; int src_j = perm[j]; y[b * N + j] = x[b * N + src_j]; } } ``` For optimization, I could consider using shared memory to cache the permutation array, but with N=36 elements, it fits easily in shared memory. However, the permutation is small and could even be loaded into registers or constant memory for efficiency. The core operation is straightforward: for each batch element and output column, read the corresponding input column using the permutation mapping, then write to the output tensor. Here's the kernel implementation: ```cuda __global__ void permute_kernel(const float* __restrict__ x, float* __restrict__ y, const int* __restrict__ perm, int B, int N) { int idx = blockIdx.x * blockDim.x + threadIdx.x; int total = B * N; if (idx < total) { int b = idx / N; int j = idx % N; int src_j = perm[j]; y[b * N + j] = x[b * N + src_j]; } } ``` Now I need to construct the complete solution. The main challenge is that `PERM` gets defined at the module level using a specific generator seed, and I need to access it from my kernel. Looking at the reference, PERM is constructed with a fixed seed to ensure determinism, so my implementation must use this exact same permutation array. For the custom kernel approach, I'll compile a CUDA kernel via `torch.utils.cpp_extension.load_inline`, then call it from ModelNew's forward method by passing in the PERM tensor directly as an argument. The kernel itself takes the input tensor, output tensor, permutation indices, batch size, and feature dimension. Each thread computes its position in the output, determines the source index using the permutation, and copies the appropriate element. I need to ensure the permutation tensor is moved to the correct device before passing it to the kernel. For the Python wrapper, I should handle both 1D and 2D inputs—if the input is 1D, the permutation is applied directly; if it's 2D with shape (B, N), I reshape it to apply the permutation across all elements, then restore the original shape. The kernel is launched with enough threads to cover all elements, using a block size of 256. There's a discrepancy I need to clarify: the kernel assumes row-major ordering where batch elements vary slowly and features vary quickly, but I should verify this matches the actual tensor layout. For a (B, N) shaped tensor, each batch is contiguous in memory with N elements per row, so b * N + j correctly accesses element j within batch b. The permutation logic in the reference swaps columns according to PERM[j], moving x[b, i] to y[b, j] where i = PERM[j]. This reverses the mapping I initially had—PERM[j] gives the source column index, not the destination. So the kernel should read from x[b, perm[j]] and write to y[b, j]. ```cuda int src_j = perm[j]; y[b * N + j] = x[b * N + src_j]; ``` This should work correctly. For input shape [B, N] = [10, 36], I need to flatten to [360] for the kernel, but the linear indexing formula `idx = b * N + j` already handles this, so I can keep the 2D shape as-is and let the kernel work with the linear indices directly. The kernel should handle arbitrary shapes by accepting total elements and stride. For a 2D tensor, I compute which row and column each thread handles based on the global thread index. Let me start simple with the basic 2D case first, then generalize if needed. The permutation kernel takes the input tensor, output tensor, permutation indices, and dimensions. Each thread calculates its batch and column indices, then applies the permutation mapping to locate the source column and copy the value. I launch with enough threads to cover all elements, using a standard block size of 256. In the forward pass, I need to ensure the permutation tensor matches the input device, then call the kernel with the appropriate grid and block dimensions. There's also a consideration for the PERM tensor defined externally—I should reference it directly from the global scope rather than recreating it. The forward method needs to handle the PERM tensor appropriately. Since PERM is defined globally at the module level, I should access it directly within the forward pass. Let me restructure the kernel launch to properly manage this global permutation tensor, ensuring it's moved to the correct device before being passed to the CUDA operation. from the device tensor. Since PERM is already 1D, I can cache the device-specific version after the first conversion rather than checking the device match every time. I'll track the current device in `self._perm_device` and update it only when needed. During the forward pass, if the input device differs from the cached permutation device, I move PERM to the correct device and update the cache. Then I launch the permutation kernel with the input, output buffer, permutation indices, batch size, and sequence length. The kernel itself handles the actual permutation operation—it uses CUDA to efficiently reorder elements according to the permutation tensor across batches. I need to account for 0-dimensional batches, which would indicate a 1D input tensor. If input.dim() == 1, then B is effectively 0 and N equals the total number of elements. Rather than relying on external globals, I'll generate the permutation directly using the fixed seed 10, ensuring it matches what the reference expects. This approach sidesteps potential scope issues with PERM and makes the initialization self-contained. For a 1D tensor with shape (N,), I should handle it by treating B as 1 in the kernel, or make the kernel more flexible to accommodate this case. The cleanest approach is to keep the kernel simple for the (B, N) case and ensure the permutation is generated consistently with the same seed. The reference uses `torch.randperm(N, generator=_g)` with a CPU generator seeded to 10, then moves the result to the GPU. Since the randperm operation happens on CPU before the device transfer, I need to replicate this exact sequence to match the reference behavior. My approach generates the permutation on CPU with the same seed, then moves it to the target device when needed. I also need to ensure the output tensor is contiguous for the kernel operation, handling both 2D and 1D cases appropriately. For the output allocation, I check if the input is already contiguous—if not, I create a contiguous copy. Then I get the batch and column dimensions, ensure the permutation is on the correct device, and launch the kernel with appropriate block and grid dimensions before returning the result. The kernel itself computes the batch and column indices from each thread's global index, then applies the permutation by mapping each output column to its corresponding input column. -d % N = 0, which means d = N. This confirms the formula works correctly. I'll finalize the solution with this indexing approach. However, I'm concerned about whether the global PERM variable will be properly resolved when the module loads. To address this, I could generate PERM internally within ModelNew, but the reference code defines PERM at the module level for sharing between the model and input generation. When my code gets executed, the globals dictionary should already contain PERM, so I'll attempt to access it directly. If that doesn't work, I'll need a fallback approach. Let me build the complete Python file with the CUDA kernel integrated properly.