```python
import math
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.utils.cpp_extension import load_inline


_cpp_source = r"""
#include <torch/extension.h>

torch::Tensor large_layer_norm_cuda(
    torch::Tensor x,
    torch::Tensor weight,
    torch::Tensor bias,
    double eps);
"""

_cuda_source = r"""
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include <c10/cuda/CUDAGuard.h>
#include <c10/cuda/CUDAException.h>
#include <cuda_runtime.h>

constexpr int THREADS = 256;
constexpr int CHUNK = 4096;

struct Moments {
    float mean;
    float m2;
    float count;
};

__device__ __forceinline__ Moments merge_moments(
    Moments a, Moments b) {
    if (b.count == 0.0f) return a;
    if (a.count == 0.0f) return b;

    float count = a.count + b.count;
    float delta = b.mean - a.mean;
    float ratio = b.count / count;

    Moments result;
    result.mean = a.mean + delta * ratio;
    result.m2 = a.m2 + b.m2
              + delta * delta * (a.count * ratio);
    result.count = count;
    return result;
}

__device__ __forceinline__ Moments warp_reduce(Moments value) {
    #pragma unroll
    for (int offset = 16; offset > 0; offset >>= 1) {
        Moments other;
        other.mean = __shfl_down_sync(0xffffffff, value.mean, offset);
        other.m2 = __shfl_down_sync(0xffffffff, value.m2, offset);
        other.count = __shfl_down_sync(0xffffffff, value.count, offset);
        value = merge_moments(value, other);
    }
    return value;
}

__device__ __forceinline__ Moments block_reduce(Moments value) {
    __shared__ float means[8];
    __shared__ float m2s[8];
    __shared__ float counts[8];

    int lane = threadIdx.x & 31;
    int warp = threadIdx.x >> 5;

    value = warp_reduce(value);
    if (lane == 0) {
        means[warp] = value.mean;
        m2s[warp] = value.m2;
        counts[warp] = value.count;
    }
    __syncthreads();

    Moments result = {0.0f, 0.0f, 0.0f};
    if (warp == 0) {
        if (lane < 8) {
            result.mean = means[lane];
            result.m2 = m2s[lane];
            result.count = counts[lane];
        }
        result = warp_reduce(result);
    }
    return result;
}

__global__ void partial_moments_kernel(
    const float* __restrict__ x,
    float2* __restrict__ partial,
    int64_t width,
    int parts) {

    int part = blockIdx.x;
    int row = blockIdx.y;
    int64_t base = int64_t(row) * width + int64_t(part) * CHUNK;
    const float4* input = reinterpret_cast<const float4*>(x + base);

    Moments value = {0.0f, 0.0f, 0.0f};

    #pragma unroll
    for (int k = 0; k < CHUNK / (4 * THREADS); ++k) {
        float4 v = input[threadIdx.x + k * THREADS];

        float mean = ((v.x + v.y) + (v.z + v.w)) * 0.25f;
        float d0 = v.x - mean;
        float d1 = v.y - mean;
        float d2 = v.z - mean;
        float d3 = v.w - mean;

        Moments group = {
            mean,
            (d0 * d0 + d1 * d1) + (d2 * d2 + d3 * d3),
            4.0f
        };
        value = merge_moments(value, group);
    }

    value = block_reduce(value);
    if (threadIdx.x == 0) {
        partial[int64_t(row) * parts + part] =
            make_float2(value.mean, value.m2);
    }
}

__global__ void finalize_moments_kernel(
    const float2* __restrict__ partial,
    float2* __restrict__ stats,
    int parts,
    int64_t width,
    float eps) {

    int row = blockIdx.x;
    Moments value = {0.0f, 0.0f, 0.0f};

    for (int part = threadIdx.x; part < parts; part += THREADS) {
        float2 p = partial[int64_t(row) * parts + part];
        Moments group = {p.x, p.y, float(CHUNK)};
        value = merge_moments(value, group);
    }

    value = block_reduce(value);
    if (threadIdx.x == 0) {
        float variance = fmaxf(value.m2 / float(width), 0.0f);
        stats[row] = make_float2(value.mean, rsqrtf(variance + eps));
    }
}

__global__ void affine_kernel(
    const float4* __restrict__ x,
    const float4* __restrict__ weight,
    const float4* __restrict__ bias,
    const float2* __restrict__ stats,
    float4* __restrict__ output,
    int64_t width4) {

    int64_t col = int64_t(blockIdx.x) * blockDim.x + threadIdx.x;
    if (col >= width4) return;

    int row = blockIdx.y;
    int64_t index = int64_t(row) * width4 + col;

    float2 s = stats[row];
    float4 v = x[index];
    float4 w = weight[col];
    float4 b = bias[col];

    float4 result;
    result.x = ((v.x - s.x) * s.y) * w.x + b.x;
    result.y = ((v.y - s.x) * s.y) * w.y + b.y;
    result.z = ((v.z - s.x) * s.y) * w.z + b.z;
    result.w = ((v.w - s.x) * s.y) * w.w + b.w;
    output[index] = result;
}

torch::Tensor large_layer_norm_cuda(
    torch::Tensor x,
    torch::Tensor weight,
    torch::Tensor bias,
    double eps) {

    TORCH_CHECK(x.is_cuda() && weight.is_cuda() && bias.is_cuda(),
                "All tensors must be CUDA tensors");
    TORCH_CHECK(x.scalar_type() == torch::kFloat32 &&
                weight.scalar_type() == torch::kFloat32 &&
                bias.scalar_type() == torch::kFloat32,
                "All tensors must be float32");
    TORCH_CHECK(x.is_contiguous() && weight.is_contiguous() &&
                bias.is_contiguous(), "All tensors must be contiguous");
    TORCH_CHECK(x.device() == weight.device() &&
                x.device() == bias.device(), "Device mismatch");

    int64_t width = weight.numel();
    TORCH_CHECK(width > 0 && width % CHUNK == 0,
                "Normalized width must be divisible by 4096");
    TORCH_CHECK(bias.numel() == width && x.numel() % width == 0,
                "Invalid tensor sizes");

    int64_t rows = x.numel() / width;
    TORCH_CHECK(rows <= 65535, "Too many rows");

    c10::cuda::CUDAGuard device_guard(x.device());
    auto output = torch::empty_like(x);
    if (rows == 0) return output;

    int parts = static_cast<int>(width / CHUNK);
    auto partial = torch::empty({rows, int64_t(parts), 2}, x.options());
    auto stats = torch::empty({rows, 2}, x.options());

    cudaStream_t stream = at::cuda::getCurrentCUDAStream();

    partial_moments_kernel<<<
        dim3(parts, static_cast<unsigned int>(rows)), THREADS, 0, stream>>>(
        x.data_ptr<float>(),
        reinterpret_cast<float2*>(partial.data_ptr<float>()),
        width, parts);
    C10_CUDA_KERNEL_LAUNCH_CHECK();

    finalize_moments_kernel<<<
        static_cast<unsigned int>(rows), THREADS, 0, stream>>>(
        reinterpret_cast<const float2*>(partial.data_ptr<float>()),
        reinterpret_cast<float2*>(stats.data_ptr<float>()),
        parts, width, static_cast<float>(eps));
    C10_CUDA_KERNEL_LAUNCH_CHECK();

    int64_t width4 = width / 4;
    affine_kernel<<<
        dim3(
            static_cast<unsigned int>((width4 + THREADS - 1) / THREADS),
            static_cast<unsigned int>(rows)),
        THREADS, 0, stream>>>(
        reinterpret_cast<const float4*>(x.data_ptr<float>()),
        reinterpret_cast<const float4*>(weight.data_ptr<float>()),
        reinterpret_cast<const float4*>(bias.data_ptr<float>()),
        reinterpret_cast<const float2*>(stats.data_ptr<float>()),
        reinterpret_cast<float4*>(output.data_ptr<float>()),
        width4);
    C10_CUDA_KERNEL_LAUNCH_CHECK();

    return output;
}
"""

_layer_norm_ext = load_inline(
    name="large_layer_norm_welford_cuda",
    cpp_sources=_cpp_source,
    cuda_sources=_cuda_source,
    functions=["large_layer_norm_cuda"],
    extra_cflags=["-O3"],
    extra_cuda_cflags=["-O3", "--fmad=false"],
    verbose=False,
)


class ModelNew(nn.Module):
    def __init__(self, normalized_shape: tuple):
        super().__init__()
        self.ln = nn.LayerNorm(normalized_shape=normalized_shape)
        self.normalized_shape = tuple(self.ln.normalized_shape)
        self.normalized_size = math.prod(self.normalized_shape)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        weight = self.ln.weight
        bias = self.ln.bias
        n = self.normalized_size

        if (
            x.is_cuda
            and x.dtype == torch.float32
            and x.is_contiguous()
            and tuple(x.shape[-len(self.normalized_shape):])
            == self.normalized_shape
            and n > 0
            and n % 4096 == 0
            and x.numel() // n <= 65535
            and weight.device == x.device
            and bias.device == x.device
            and weight.dtype == torch.float32
            and bias.dtype == torch.float32
            and weight.is_contiguous()
            and bias.is_contiguous()
            and x.data_ptr() % 16 == 0
            and weight.data_ptr() % 16 == 0
            and bias.data_ptr() % 16 == 0
        ):
            return _layer_norm_ext.large_layer_norm_cuda(
                x, weight, bias, self.ln.eps
            )

        return F.layer_norm(
            x, self.normalized_shape, weight, bias, self.ln.eps
        )
```