```python
import torch
import torch.nn as nn
from torch.utils.cpp_extension import load_inline

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

torch::Tensor 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>
#include <algorithm>

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

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

__device__ __forceinline__ Moments merge(Moments a, Moments b) {
    if (b.count == 0) return a;
    if (a.count == 0) return b;

    const int count = a.count + b.count;
    const float delta = b.mean - a.mean;
    const float ratio = float(b.count) / float(count);
    a.mean += delta * ratio;
    a.m2 += b.m2 + delta * delta * (float(a.count) * ratio);
    a.count = count;
    return a;
}

__device__ __forceinline__ Moments warp_reduce(Moments a) {
    #pragma unroll
    for (int offset = 16; offset > 0; offset >>= 1) {
        Moments b;
        b.mean = __shfl_down_sync(0xffffffff, a.mean, offset);
        b.m2 = __shfl_down_sync(0xffffffff, a.m2, offset);
        b.count = __shfl_down_sync(0xffffffff, a.count, offset);
        if ((threadIdx.x & 31) + offset < 32) {
            a = merge(a, b);
        }
    }
    return a;
}

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

    a = warp_reduce(a);
    const int lane = threadIdx.x & 31;
    const int warp = threadIdx.x >> 5;

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

    if (warp == 0) {
        a = Moments{0.0f, 0.0f, 0};
        if (lane < 8) {
            a = Moments{means[lane], m2s[lane], counts[lane]};
        }
        a = warp_reduce(a);
    }
    return a;
}

__global__ void partial_moments(
    const float* __restrict__ x,
    float2* __restrict__ partial,
    int64_t width,
    int parts)
{
    const int row = blockIdx.y;
    const int part = blockIdx.x;
    const int64_t start = int64_t(part) * CHUNK;
    const float4* input =
        reinterpret_cast<const float4*>(x + int64_t(row) * width);

    Moments a{0.0f, 0.0f, 0};

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

            // Form stable moments for four values before merging.
            const float d01 = v.y - v.x;
            const float d23 = v.w - v.z;
            const float mean01 = v.x + 0.5f * d01;
            const float mean23 = v.z + 0.5f * d23;
            const float delta = mean23 - mean01;

            Moments b;
            b.mean = mean01 + 0.5f * delta;
            b.m2 = 0.5f * (d01 * d01 + d23 * d23) + delta * delta;
            b.count = 4;
            a = merge(a, b);
        }
    }

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

__global__ void finish_moments(
    const float2* __restrict__ partial,
    float2* __restrict__ statistics,
    int64_t width,
    int parts,
    float eps)
{
    const int row = blockIdx.x;
    Moments a{0.0f, 0.0f, 0};

    for (int part = threadIdx.x; part < parts; part += THREADS) {
        const float2 p = partial[int64_t(row) * parts + part];
        const int64_t remaining = width - int64_t(part) * CHUNK;
        const int count = remaining < CHUNK ? int(remaining) : CHUNK;
        a = merge(a, Moments{p.x, p.y, count});
    }

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

__global__ void normalize_affine(
    const float* __restrict__ x,
    const float* __restrict__ weight,
    const float* __restrict__ bias,
    const float2* __restrict__ statistics,
    float* __restrict__ output,
    int64_t width)
{
    const int row = blockIdx.y;
    const float2 stats = statistics[row];
    const int64_t row_offset = int64_t(row) * width;

    const float4* input4 = reinterpret_cast<const float4*>(x + row_offset);
    const float4* weight4 = reinterpret_cast<const float4*>(weight);
    const float4* bias4 = reinterpret_cast<const float4*>(bias);
    float4* output4 = reinterpret_cast<float4*>(output + row_offset);

    #pragma unroll
    for (int k = 0; k < OUTPUT_CHUNK / (THREADS * 4); ++k) {
        const int64_t index =
            int64_t(blockIdx.x) * OUTPUT_CHUNK +
            (threadIdx.x + k * THREADS) * 4;

        if (index < width) {
            const int64_t i = index / 4;
            const float4 v = input4[i];
            const float4 w = weight4[i];
            const float4 b = bias4[i];

            float4 result;
            result.x = ((v.x - stats.x) * stats.y) * w.x + b.x;
            result.y = ((v.y - stats.x) * stats.y) * w.y + b.y;
            result.z = ((v.z - stats.x) * stats.y) * w.z + b.z;
            result.w = ((v.w - stats.x) * stats.y) * w.w + b.w;
            output4[i] = result;
        }
    }
}

torch::Tensor layer_norm_cuda(
    torch::Tensor x,
    torch::Tensor weight,
    torch::Tensor bias,
    double eps)
{
    const c10::cuda::CUDAGuard guard(x.device());
    const int64_t width = weight.numel();
    const int64_t rows = x.numel() / width;
    const int parts = int((width + CHUNK - 1) / CHUNK);

    auto output = torch::empty_like(x);
    if (rows == 0) return output;

    auto partial = torch::empty({rows, int64_t(parts), 2}, x.options());
    auto statistics = torch::empty({rows, 2}, x.options());
    const cudaStream_t stream = at::cuda::getCurrentCUDAStream();

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

    finish_moments<<<rows, THREADS, 0, stream>>>(
        reinterpret_cast<const float2*>(partial.data_ptr<float>()),
        reinterpret_cast<float2*>(statistics.data_ptr<float>()),
        width, parts, float(eps));
    C10_CUDA_KERNEL_LAUNCH_CHECK();

    const int output_blocks = int((width + OUTPUT_CHUNK - 1) / OUTPUT_CHUNK);
    normalize_affine<<<dim3(output_blocks, rows), THREADS, 0, stream>>>(
        x.data_ptr<float>(),
        weight.data_ptr<float>(),
        bias.data_ptr<float>(),
        reinterpret_cast<const float2*>(statistics.data_ptr<float>()),
        output.data_ptr<float>(),
        width);
    C10_CUDA_KERNEL_LAUNCH_CHECK();

    return output;
}
"""

_layer_norm_extension = load_inline(
    name="large_layer_norm_welford_cuda",
    cpp_sources=_cpp_source,
    cuda_sources=_cuda_source,
    functions=["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)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        shape = self.ln.normalized_shape
        width = self.ln.weight.numel()

        if (
            not x.is_cuda
            or x.dtype != torch.float32
            or not x.is_contiguous()
            or tuple(x.shape[-len(shape):]) != tuple(shape)
            or width < 4096
            or width % 4 != 0
            or width > 2**31 - 1
            or x.numel() // width > 65535
            or self.ln.weight.dtype != torch.float32
            or self.ln.bias.dtype != torch.float32
            or self.ln.weight.device != x.device
            or self.ln.bias.device != x.device
            or not self.ln.weight.is_contiguous()
            or not self.ln.bias.is_contiguous()
            or x.data_ptr() % 16 != 0
            or self.ln.weight.data_ptr() % 16 != 0
            or self.ln.bias.data_ptr() % 16 != 0
        ):
            return self.ln(x)

        return _layer_norm_extension.layer_norm_cuda(
            x, self.ln.weight, self.ln.bias, self.ln.eps
        )
```