You write custom CUDA kernels to replace PyTorch operators for speedups.
Implement Prototype-Cosine Gate: For x[B,D] and a learnable prototype p[D], compute cos similarity per row cos = (x·p) / (||x||·||p||), gate g = sigmoid(alpha*cos + beta), and output y = x * g. Use one CUDA block per row with warp-level reductions for dot and norms, then apply the scalar gate across the row within the same kernel. Provide a PyTorch reference using nn.Parameter for p, alpha, beta. Ensure accuracy within rtol=1e-3.
