```python
# Copyright (c) Meta Platforms, Inc. and affiliates.
# All rights reserved.
#
# This source code is licensed under the license found in the
# LICENSE file in the root directory of this source tree.

import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.utils.cpp_extension import load_inline

# Fused kernel: row-wise instance normalization (no affine) + residual add + multiply
# i.e. given z = Linear(x) of shape (batch, dim) and y of shape (batch, dim):
#   mean_i = mean(z_i), var_i = var(z_i)  (biased)
#   out_ij = ((z_ij - mean_i) / sqrt(var_i + eps) + y_ij) * y_ij
# One thread block handles one row; the row is staged in shared memory and the
# two reductions (mean, variance) are done with warp shuffles + shared scratch.
cuda_source = r"""
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include <cuda_runtime.h>

__inline__ __device__ float warp_reduce_sum(float v) {
    for (int off = 16; off > 0; off >>= 1)
        v += __shfl_down_sync(0xffffffffu, v, off);
    return v;
}

__inline__ __device__ float block_reduce_sum(float v, float* scratch) {
    int lane = threadIdx.x & 31;
    int wid  = threadIdx.x >> 5;
    v = warp_reduce_sum(v);
    if (lane == 0) scratch[wid] = v;
    __syncthreads();
    int nwarps = (blockDim.x + 31) >> 5;
    if (wid == 0) {
        v = (lane < nwarps) ? scratch[lane] : 0.0f;
        v = warp_reduce_sum(v);
        if (lane == 0) scratch[0] = v;
    }
    __syncthreads();
    float r = scratch[0];
    __syncthreads();
    return r;
}

__global__ void fused_inorm_add_mul_kernel(
        const float* __restrict__ z,
        const float* __restrict__ y,
        float* __restrict__ out,
        const int dim,
        const float eps) {
    const int row = blockIdx.x;
    const float* zr = z + (size_t)row * dim;
    const float* yr = y + (size_t)row * dim;
    float* orow = out + (size_t)row * dim;

    extern __shared__ float s_val[];   // dim floats: staged row of z
    __shared__ float s_scratch[32];
    __shared__ float s_mean;
    __shared__ float s_rstd;

    // float4 path is only safe when every row is 16B aligned (dim % 4 == 0)
    const bool aligned = (dim & 3) == 0;
    const int nvec = aligned ? (dim >> 2) : 0;

    const float4* z4 = reinterpret_cast<const float4*>(zr);
    const float4* y4 = reinterpret_cast<const float4*>(yr);
    float4* o4 = reinterpret_cast<float4*>(orow);

    // Load z into shared memory, accumulate sum
    float acc = 0.0f;
    for (int i = threadIdx.x; i < nvec; i += blockDim.x) {
        float4 v = z4[i];
        s_val[4*i+0] = v.x;
        s_val[4*i+1] = v.y;
        s_val[4*i+2] = v.z;
        s_val[4*i+3] = v.w;
        acc += (v.x + v.y) + (v.z + v.w);
    }
    for (int i = (nvec << 2) + threadIdx.x; i < dim; i += blockDim.x) {
        float v = zr[i];
        s_val[i] = v;
        acc += v;
    }

    float total = block_reduce_sum(acc, s_scratch);
    if (threadIdx.x == 0) s_mean = total / (float)dim;
    __syncthreads();
    const float mean = s_mean;

    // Accumulate biased variance
    float vacc = 0.0f;
    for (int i = threadIdx.x; i < dim; i += blockDim.x) {
        float d = s_val[i] - mean;
        vacc += d * d;
    }
    float vtotal = block_reduce_sum(vacc, s_scratch);
    if (threadIdx.x == 0) s_rstd = rsqrtf(vtotal / (float)dim + eps);
    __syncthreads();
    const float rstd = s_rstd;

    // out = (norm(z) + y) * y
    for (int i = threadIdx.x; i < nvec; i += blockDim.x) {
        float4 yv = y4[i];
        float4 ov;
        ov.x = ((s_val[4*i+0] - mean) * rstd + yv.x) * yv.x;
        ov.y = ((s_val[4*i+1] - mean) * rstd + yv.y) * yv.y;
        ov.z = ((s_val[4*i+2] - mean) * rstd + yv.z) * yv.z;
        ov.w = ((s_val[4*i+3] - mean) * rstd + yv.w) * yv.w;
        o4[i] = ov;
    }
    for (int i = (nvec << 2) + threadIdx.x; i < dim; i += blockDim.x) {
        float yv = yr[i];
        orow[i] = ((s_val[i] - mean) * rstd + yv) * yv;
    }
}

torch::Tensor fused_inorm_add_mul(torch::Tensor z, torch::Tensor y, double eps) {
    TORCH_CHECK(z.is_cuda() && y.is_cuda(), "inputs must be CUDA tensors");
    TORCH_CHECK(z.scalar_type() == at::kFloat && y.scalar_type() == at::kFloat,
                "only float32 is supported");
    TORCH_CHECK(z.dim() == 2 && y.dim() == 2, "inputs must be 2D");
    TORCH_CHECK(z.size(0) == y.size(0) && z.size(1) == y.size(1), "shape mismatch");

    auto zc = z.contiguous();
    auto yc = y.contiguous();
    auto out = at::empty_like(zc);

    const int batch = (int)zc.size(0);
    const int dim   = (int)zc.size(1);
    if (batch == 0 || dim == 0) return out;

    const int threads = 256;
    const size_t shmem = (size_t)dim * sizeof(float);
    auto stream = at::cuda::getCurrentCUDAStream();

    if (shmem > 48 * 1024) {
        static int max_shmem = -1;
        if (max_shmem < 0) {
            int dev = 0;
            cudaGetDevice(&dev);
            cudaDeviceProp prop;
            cudaGetDeviceProperties(&prop, dev);
            max_shmem = (int)prop.sharedMemPerBlockOptin;
        }
        TORCH_CHECK((int)shmem <= max_shmem, "feature dim too large for fused kernel");
        cudaFuncSetAttribute(fused_inorm_add_mul_kernel,
                             cudaFuncAttributeMaxDynamicSharedMemorySize,
                             (int)shmem);
    }

    fused_inorm_add_mul_kernel<<<batch, threads, shmem, stream>>>(
        zc.data_ptr<float>(), yc.data_ptr<float>(), out.data_ptr<float>(),
        dim, (float)eps);

    return out;
}
"""

cpp_source = "torch::Tensor fused_inorm_add_mul(torch::Tensor z, torch::Tensor y, double eps);"

fused_ops = load_inline(
    name="fused_inorm_add_mul",
    cpp_sources=cpp_source,
    cuda_sources=cuda_source,
    functions=["fused_inorm_add_mul"],
    verbose=False,
    extra_cflags=["-O3"],
    extra_cuda_cflags=["-O3"],
)


class ModelNew(nn.Module):
    """
    Optimized version of Model.

    The Linear layer (the dominant GEMM) is kept on the highly tuned cuBLAS
    path via F.linear. Everything after it -- instance normalization over the
    feature dim, the residual add, and the final multiply -- is fused into a
    single custom CUDA kernel, eliminating multiple kernel launches and
    several full read/write passes over the (batch, out_features) tensor.
    """
    def __init__(self, in_features, out_features, eps=1e-5, momentum=0.1):
        super(ModelNew, self).__init__()
        self.bmm = nn.Linear(in_features, out_features)
        self.instance_norm = nn.InstanceNorm2d(out_features, eps=eps, momentum=momentum)
        self._fused = fused_ops

    def forward(self, x, y):
        z = F.linear(x, self.bmm.weight, self.bmm.bias)

        inorm = self.instance_norm
        # Fast path only valid for: per-instance stats, no affine parameters,
        # fp32 CUDA tensors. Otherwise fall back to the reference ops.
        if (z.is_cuda and z.dtype == torch.float32 and y.dtype == torch.float32
                and y.device == z.device
                and inorm.weight is None and inorm.bias is None
                and not inorm.track_running_stats):
            return self._fused.fused_inorm_add_mul(z, y, float(inorm.eps))

        out = inorm(z.unsqueeze(1).unsqueeze(1)).squeeze(1).squeeze(1)
        out = out + y
        out = out * y
        return out
```