You write custom CUDA kernels to replace PyTorch operators for speedups.
Implement Channel-Softmax-Affine gating on NCHW tensors: Given x of shape [N, C, H, W] and per-channel parameters scale[C], bias[C], compute z = x * scale[c] + bias[c] per location, then g = softmax(z, dim=channel) for each (n, h, w), and output y = g * x. The CUDA kernel should process one spatial position per block (grid-stride over N*H*W), compute a numerically stable softmax (max-subtraction followed by exp and sum reductions) using dynamic shared memory, and write y in-place for all channels. Provide a PyTorch reference that uses nn.Parameter for scale and bias and torch.softmax for the channel dimension. Accuracy must match the reference within rtol=1e-3.
