Now let me understand the problem better: - Input: tensor of shape (B, N) where B=10, N=36 - Output: y[b, j] = x[b, perm[j]] for all b in [0, B) and j in [0, N) - This is just permuting columns according to a fixed permutation The naive PyTorch implementation does `x[:, PERM.to(device=x.device)]` which involves: 1. Moving PERM to the device 2. Advanced indexing I can write a CUDA kernel that directly reads from permuted indices. Let me think about the optimal approach: For small tensors like (10, 36), we need to be careful about overhead. The total data is 360 floats = 1440 bytes, which is tiny. A simple approach: each thread handles one element. We need B*N = 360 threads, which is very small for GPU. Let me write the kernel: ```python import torch import torch.nn as nn from torch.utils.cpp_extension import load_inline from kernelbench.distributed_collectives import default_device B = 10 N = 36 _g = torch.Generator() _g.manual_seed(10) PERM = torch.randperm(N, generator=_g) # Define CUDA kernel cuda_source = """ __global__ void permute_kernel(const float* __restrict__ input, float* __restrict__ output, 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; output[idx] = input[b * N + perm[j]]; } } void launch_permute(const float* input, float* output, const int* perm, int B, int N) { int total = B * N; int threads = 256; int blocks = (total + threads - 1) / threads; permute_kernel<<<blocks, threads>>>(input, output, perm, B, N); } """ cpp_source = """ void launch_permute(const float* input, float* output, const int* perm, int B, int N); """ permute_extension = load_inline( name='permute_extension', cpp_sources=cpp_source, cuda_sources=cuda_source, functions=['launch_permute'], extra_cuda_cflags=['-O3'], ) class ModelNew(nn.Module): def __init__(self): super().__init__() # Register permutation as a persistent buffer self.register_buffer('perm', PERM.to(torch.int32)) def forward(self, x: torch.Tensor) -> torch.Tensor: B, N = x.shape output = torch.empty_like(x) permute_extension.launch_permute(x.data_ptr(), output.data_ptr(), self.perm.data_ptr(), B, N) return output ``` Wait, there's an issue. The PERM is defined at module level with B and N being global constants. But the permutation should be the same as in the reference. Let me re-read the reference... Actually, the reference does `x[:, PERM.to(device=x.device)]` which moves PERM to the device each time. I should avoid this by caching PERM as a buffer. But I need to be careful - the permutation indices should be int32 or int64 for indexing. Let me use int64 to be safe. Actually wait, in CUDA using int64 might be slower. int32 should be fine since N=36. Let me also think about optimizations: 1. Use __ldg for read-only cache on older GPUs (on H100 it might be automatic) 2. Use vectorized loads if possible 3. Use shared memory for the permutation array Actually, for such small data, the permutation array could fit in shared memory or even constant memory. Let me try using constant memory for the permutation. Actually, since PERM is defined at module level with a fixed seed, we can potentially use constant memory. But the size is N=36 which is small. Let me start with a simple kernel and optimize from there. Actually, I realize I should be more careful. The PERM is generated at module level, and my ModelNew needs to use the same PERM. Let me make sure I do that correctly. Let me write the first version: