```python
import math
import torch
import torch.nn as nn
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,
    int64_t normalized_size,
    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 <climits>

constexpr int THREADS = 256;
constexpr int CHUNK = 8192;

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 fraction = float(b.count) / float(count);

    a.mean += delta * fraction;
    a.m2 += b.m2 + delta * delta * (float(a.count) * fraction);
    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 shared_mean[8];
    __shared__ float shared_m2[8];
    __shared__ int shared_count[8];

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

    a = warp_reduce(a);
    if (lane == 0) {
        shared_mean[warp] = a.mean;
        shared_m2[warp] = a.m2;
        shared_count[warp] = a.count;
    }
    __syncthreads();

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

__global__ void partial_moments_kernel(
    const float* __restrict__ x,
    float2* __restrict__ partial,
    int n,
    int parts)
{
    const int part = blockIdx.x;
    const int row = blockIdx.y;
    const int64_t start = int64_t(part) * CHUNK;
    const int length = int(
        int64_t(n) - start < CHUNK ? int64_t(n) - start : CHUNK);
    const float* input = x + int64_t(row) * n + start;

    Moments local = {0.0f, 0.0f, 0};

    for (int i = threadIdx.x; i < length; i += THREADS) {
        const float value = input[i];
        ++local.count;
        const float delta = value - local.mean;
        local.mean += delta / float(local.count);
        local.m2 += delta * (value - local.mean);
    }

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

__global__ void finish_moments_kernel(
    const float2* __restrict__ partial,
    float2* __restrict__ stats,
    int n,
    int parts,
    float eps)
{
    const int row = blockIdx.x;
    Moments local = {0.0f, 0.0f, 0};

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

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

__global__ void normalize_kernel(
    const float* __restrict__ x,
    const float* __restrict__ weight,
    const float* __restrict__ bias,
    const float2* __restrict__ stats,
    float* __restrict__ output,
    int n)
{
    const int row = blockIdx.y;
    const int64_t start = int64_t(blockIdx.x) * CHUNK;
    const int64_t end = start + CHUNK < n ? start + CHUNK : n;
    const int64_t row_offset = int64_t(row) * n;
    const float2 parameters = stats[row];

    if ((n & 3) == 0) {
        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);

        for (int64_t i = start / 4 + threadIdx.x;
             i < end / 4;
             i += THREADS) {
            const float4 v = input4[i];
            const float4 w = weight4[i];
            const float4 b = bias4[i];

            float4 result;
            result.x = ((v.x - parameters.x) * parameters.y) * w.x + b.x;
            result.y = ((v.y - parameters.x) * parameters.y) * w.y + b.y;
            result.z = ((v.z - parameters.x) * parameters.y) * w.z + b.z;
            result.w = ((v.w - parameters.x) * parameters.y) * w.w + b.w;
            output4[i] = result;
        }
    } else {
        for (int64_t i = start + threadIdx.x; i < end; i += THREADS) {
            const float normalized =
                (x[row_offset + i] - parameters.x) * parameters.y;
            output[row_offset + i] = normalized * weight[i] + bias[i];
        }
    }
}

torch::Tensor large_layer_norm_cuda(
    torch::Tensor x,
    torch::Tensor weight,
    torch::Tensor bias,
    int64_t normalized_size,
    double eps)
{
    TORCH_CHECK(x.is_cuda() && weight.is_cuda() && bias.is_cuda(),
                "All tensors must be CUDA tensors");
    TORCH_CHECK(x.device() == weight.device() &&
                x.device() == bias.device(), "Device mismatch");
    TORCH_CHECK(x.scalar_type() == at::kFloat &&
                weight.scalar_type() == at::kFloat &&
                bias.scalar_type() == at::kFloat, "Expected float32 tensors");
    TORCH_CHECK(x.is_contiguous() && weight.is_contiguous() &&
                bias.is_contiguous(), "Expected contiguous tensors");
    TORCH_CHECK(normalized_size > 0 && normalized_size <= INT_MAX,
                "Unsupported normalized size");
    TORCH_CHECK(x.numel() % normalized_size == 0 &&
                weight.numel() == normalized_size &&
                bias.numel() == normalized_size, "Invalid tensor shapes");

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

    TORCH_CHECK(rows <= 65535, "Too many rows");
    const int n = int(normalized_size);
    const int parts = int((normalized_size + CHUNK - 1) / CHUNK);

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

    const dim3 grid(parts, unsigned(rows));

    partial_moments_kernel<<<grid, THREADS, 0, stream>>>(
        x.data_ptr<float>(),
        reinterpret_cast<float2*>(partial.data_ptr<float>()),
        n, parts);
    C10_CUDA_KERNEL_LAUNCH_CHECK();

    finish_moments_kernel<<<unsigned(rows), THREADS, 0, stream>>>(
        reinterpret_cast<const float2*>(partial.data_ptr<float>()),
        reinterpret_cast<float2*>(stats.data_ptr<float>()),
        n, parts, float(eps));
    C10_CUDA_KERNEL_LAUNCH_CHECK();

    normalize_kernel<<<grid, THREADS, 0, stream>>>(
        x.data_ptr<float>(),
        weight.data_ptr<float>(),
        bias.data_ptr<float>(),
        reinterpret_cast<const float2*>(stats.data_ptr<float>()),
        output.data_ptr<float>(),
        n);
    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=["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_size = math.prod(self.ln.normalized_shape)

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

        if (
            not x.is_cuda
            or x.dtype != torch.float32
            or self.ln.weight.dtype != torch.float32
            or self.ln.bias.dtype != torch.float32
            or tuple(x.shape[-len(shape):]) != tuple(shape)
            or n < 8192
            or n > 2147483647
            or x.numel() // n > 65535
            or x.numel() == 0
        ):
            return self.ln(x)

        return _layer_norm_extension.large_layer_norm_cuda(
            x.contiguous(),
            self.ln.weight.contiguous(),
            self.ln.bias.contiguous(),
            n,
            self.ln.eps,
        )
```