```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

# Custom fused LayerNorm CUDA kernels.
# Strategy:
#   Kernel 1: multi-block-per-row reduction of sum & sum-of-squares
#             (vectorized float4 loads, double-precision in-thread
#              accumulation, warp/block tree reduction, atomic merge).
#   Kernel 2: normalize + affine transform, reading the per-row stats
#             (vectorized float4 loads/stores, fmaf).
# Splitting each row across many thread blocks keeps every SM busy even
# when there are few rows (batch=16) with very large normalized shapes.

layernorm_cuda_source = r"""
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include <cuda_runtime.h>
#include <algorithm>

#define LN_BLOCK 256

__inline__ __device__ float2 warpReduceSum2(float2 v) {
    #pragma unroll
    for (int offset = 16; offset > 0; offset >>= 1) {
        v.x += __shfl_down_sync(0xffffffffu, v.x, offset);
        v.y += __shfl_down_sync(0xffffffffu, v.y, offset);
    }
    return v;
}

__inline__ __device__ float2 blockReduceSum2(float2 v) {
    __shared__ float2 shared[32];
    const int lane = threadIdx.x & 31;
    const int wid  = threadIdx.x >> 5;
    v = warpReduceSum2(v);
    if (lane == 0) shared[wid] = v;
    __syncthreads();
    const int nwarps = (blockDim.x + 31) >> 5;
    v = (threadIdx.x < nwarps) ? shared[lane] : make_float2(0.f, 0.f);
    if (wid == 0) v = warpReduceSum2(v);
    return v;  // result valid on thread 0
}

__global__ void ln_stats_kernel(const float4* __restrict__ x,
                                float* __restrict__ sum_acc,
                                float* __restrict__ sq_acc,
                                const int row_size4) {
    const int row = blockIdx.y;
    const float4* xrow = x + (size_t)row * row_size4;

    double s  = 0.0;
    double sq = 0.0;
    for (int i = blockIdx.x * blockDim.x + threadIdx.x; i < row_size4;
         i += gridDim.x * blockDim.x) {
        const float4 v = xrow[i];
        s  += (double)v.x + v.y + v.z + v.w;
        sq += (double)v.x * v.x + (double)v.y * v.y
            + (double)v.z * v.z + (double)v.w * v.w;
    }

    float2 r = blockReduceSum2(make_float2((float)s, (float)sq));
    if (threadIdx.x == 0) {
        atomicAdd(&sum_acc[row], r.x);
        atomicAdd(&sq_acc[row],  r.y);
    }
}

__global__ void ln_normalize_kernel(const float4* __restrict__ x,
                                    const float4* __restrict__ gamma,
                                    const float4* __restrict__ beta,
                                    const float* __restrict__ sum_acc,
                                    const float* __restrict__ sq_acc,
                                    float4* __restrict__ y,
                                    const int row_size4,
                                    const float inv_n,
                                    const float eps) {
    const int row = blockIdx.y;
    const float mean = sum_acc[row] * inv_n;
    const float var  = sq_acc[row] * inv_n - mean * mean;
    const float rstd = rsqrtf(var + eps);

    const float4* xrow = x + (size_t)row * row_size4;
    float4* yrow       = y + (size_t)row * row_size4;

    for (int i = blockIdx.x * blockDim.x + threadIdx.x; i < row_size4;
         i += gridDim.x * blockDim.x) {
        const float4 v = xrow[i];
        const float4 g = gamma[i];
        const float4 b = beta[i];
        float4 o;
        o.x = fmaf((v.x - mean) * rstd, g.x, b.x);
        o.y = fmaf((v.y - mean) * rstd, g.y, b.y);
        o.z = fmaf((v.z - mean) * rstd, g.z, b.z);
        o.w = fmaf((v.w - mean) * rstd, g.w, b.w);
        yrow[i] = o;
    }
}

torch::Tensor layernorm_forward(torch::Tensor x, torch::Tensor weight,
                                torch::Tensor bias, double eps) {
    auto xc = x.contiguous();
    auto wc = weight.contiguous();
    auto bc = bias.contiguous();

    const int64_t row_size = wc.numel();
    const int64_t num_rows = xc.numel() / row_size;

    auto y   = torch::empty_like(xc);
    auto acc = torch::zeros({2 * num_rows},
                            torch::dtype(torch::kFloat32).device(xc.device()));
    float* sum_acc = acc.data_ptr<float>();
    float* sq_acc  = sum_acc + num_rows;

    const int   row_size4 = (int)(row_size / 4);
    const float inv_n     = 1.0f / (float)row_size;

    const int num_sms =
        at::cuda::getCurrentDeviceProperties()->multiProcessorCount;
    int chunks = std::max(1, (num_sms * 8) / (int)num_rows);
    chunks = std::min(chunks, (row_size4 + LN_BLOCK - 1) / LN_BLOCK);

    dim3 grid((unsigned)chunks, (unsigned)num_rows);
    dim3 block(LN_BLOCK);
    auto stream = at::cuda::getCurrentCUDAStream();

    ln_stats_kernel<<<grid, block, 0, stream>>>(
        reinterpret_cast<const float4*>(xc.data_ptr<float>()),
        sum_acc, sq_acc, row_size4);

    ln_normalize_kernel<<<grid, block, 0, stream>>>(
        reinterpret_cast<const float4*>(xc.data_ptr<float>()),
        reinterpret_cast<const float4*>(wc.data_ptr<float>()),
        reinterpret_cast<const float4*>(bc.data_ptr<float>()),
        sum_acc, sq_acc,
        reinterpret_cast<float4*>(y.data_ptr<float>()),
        row_size4, inv_n, (float)eps);

    return y;
}
"""

layernorm_cpp_source = (
    "torch::Tensor layernorm_forward(torch::Tensor x, torch::Tensor weight, "
    "torch::Tensor bias, double eps);"
)

_fused_layernorm = load_inline(
    name="fused_layernorm",
    cpp_sources=layernorm_cpp_source,
    cuda_sources=layernorm_cuda_source,
    functions=["layernorm_forward"],
    verbose=False,
    extra_cuda_cflags=["-O3"],
)


class ModelNew(nn.Module):
    """
    LayerNorm model accelerated with a custom fused CUDA kernel.
    """
    def __init__(self, normalized_shape: tuple):
        super(ModelNew, self).__init__()
        self.normalized_shape = tuple(normalized_shape)
        self.ln = nn.LayerNorm(normalized_shape=normalized_shape)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        w = self.ln.weight
        b = self.ln.bias
        if (x.is_cuda and x.dtype == torch.float32 and w is not None
                and b is not None and (w.numel() % 4) == 0):
            return _fused_layernorm.layernorm_forward(
                x, w, b, float(self.ln.eps))
        # Fallback for CPU tensors / unsupported dtypes / shapes.
        return F.layer_norm(x, self.normalized_shape, w, b, self.ln.eps)
```