```python
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>
#include <vector>

std::vector<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 <vector>
#include <stdint.h>

constexpr int THREADS = 256;
constexpr int TILE = 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;
    int count = a.count + b.count;
    float delta = b.mean - a.mean;
    float fraction = float(b.count) / float(count);
    return {
        a.mean + delta * fraction,
        a.m2 + b.m2 + delta * delta * (float(a.count) * fraction),
        count
    };
}

__device__ __forceinline__ Moments update(Moments a, float value) {
    ++a.count;
    float delta = value - a.mean;
    a.mean += delta / float(a.count);
    a.m2 += delta * (value - a.mean);
    return a;
}

__device__ __forceinline__ Moments warp_reduce(Moments a) {
    int lane = threadIdx.x & 31;
    #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 (lane + offset < 32)
            a = merge(a, b);
    }
    return a;
}

__device__ __forceinline__ Moments block_reduce(Moments a) {
    __shared__ Moments warps[THREADS / 32];

    int lane = threadIdx.x & 31;
    int warp = threadIdx.x >> 5;
    a = warp_reduce(a);
    if (lane == 0)
        warps[warp] = a;
    __syncthreads();

    Moments result = {0.0f, 0.0f, 0};
    if (warp == 0) {
        if (lane < THREADS / 32)
            result = warps[lane];
        result = warp_reduce(result);
    }
    return result;
}

__global__ void partial_moments(
    const float* __restrict__ x,
    float2* __restrict__ partial,
    int64_t n,
    int parts)
{
    int part = blockIdx.x % parts;
    int64_t row = blockIdx.x / parts;
    int64_t start = int64_t(part) * TILE;
    int length = int(min(int64_t(TILE), n - start));

    const float4* input =
        reinterpret_cast<const float4*>(x + row * n + start);

    Moments moments = {0.0f, 0.0f, 0};
    #pragma unroll
    for (int i = 0; i < TILE / (4 * THREADS); ++i) {
        int index = threadIdx.x + i * THREADS;
        if (index < length / 4) {
            float4 v = input[index];
            moments = update(moments, v.x);
            moments = update(moments, v.y);
            moments = update(moments, v.z);
            moments = update(moments, v.w);
        }
    }

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

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

    for (int part = threadIdx.x; part < parts; part += THREADS) {
        float2 p = partial[row * parts + part];
        int count = int(min(int64_t(TILE), n - int64_t(part) * TILE));
        Moments next = {p.x, p.y, count};
        moments = merge(moments, next);
    }

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

__global__ void normalize_affine(
    const float* __restrict__ x,
    const float* __restrict__ weight,
    const float* __restrict__ bias,
    const float2* __restrict__ stats,
    float* __restrict__ output,
    int64_t n,
    int parts)
{
    int part = blockIdx.x % parts;
    int64_t row = blockIdx.x / parts;
    int64_t start = int64_t(part) * TILE;
    int length = int(min(int64_t(TILE), n - start));
    float2 s = stats[row];

    const float4* input =
        reinterpret_cast<const float4*>(x + row * n + start);
    const float4* w =
        reinterpret_cast<const float4*>(weight + start);
    const float4* b =
        reinterpret_cast<const float4*>(bias + start);
    float4* out =
        reinterpret_cast<float4*>(output + row * n + start);

    #pragma unroll
    for (int i = 0; i < TILE / (4 * THREADS); ++i) {
        int index = threadIdx.x + i * THREADS;
        if (index < length / 4) {
            float4 v = input[index];
            float4 gamma = w[index];
            float4 beta = b[index];

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

std::vector<torch::Tensor> layer_norm_cuda(
    torch::Tensor x,
    torch::Tensor weight,
    torch::Tensor bias,
    double eps)
{
    c10::cuda::CUDAGuard device_guard(x.device());

    int64_t n = weight.numel();
    int64_t rows = x.numel() / n;
    int parts = int((n + TILE - 1) / TILE);

    auto output = torch::empty_like(x);
    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();
    unsigned int blocks = static_cast<unsigned int>(rows * parts);

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

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

    normalize_affine<<<blocks, 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, parts);
    C10_CUDA_KERNEL_LAUNCH_CHECK();

    return {output, stats};
}
"""

_layer_norm_extension = load_inline(
    name="split_welford_layer_norm_v1",
    cpp_sources=_cpp_source,
    cuda_sources=_cuda_source,
    functions=["layer_norm_cuda"],
    extra_cflags=["-O3"],
    extra_cuda_cflags=["-O3", "--fmad=false"],
    verbose=False,
)


class _LayerNormFunction(torch.autograd.Function):
    @staticmethod
    def forward(ctx, x, weight, bias, eps):
        output, stats = _layer_norm_extension.layer_norm_cuda(
            x, weight, bias, eps
        )
        ctx.save_for_backward(x, weight, stats)
        return output

    @staticmethod
    def backward(ctx, grad_output):
        x, weight, stats = ctx.saved_tensors
        n = weight.numel()
        values = x.reshape(-1, n)
        grad = grad_output.reshape(-1, n)
        mean = stats[:, 0:1]
        rstd = stats[:, 1:2]
        normalized = (values - mean) * rstd

        dx = dw = db = None

        if ctx.needs_input_grad[0]:
            weighted_grad = grad * weight.reshape(1, n)
            dx = (
                weighted_grad
                - weighted_grad.mean(dim=1, keepdim=True)
                - normalized
                * (weighted_grad * normalized).mean(dim=1, keepdim=True)
            ) * rstd
            dx = dx.reshape_as(x)

        if ctx.needs_input_grad[1]:
            dw = (grad * normalized).sum(dim=0).reshape_as(weight)

        if ctx.needs_input_grad[2]:
            db = grad.sum(dim=0).reshape_as(weight)

        return dx, dw, db, None


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:
        weight = self.ln.weight
        bias = self.ln.bias
        shape = self.ln.normalized_shape

        use_custom = (
            x.is_cuda
            and x.dtype == torch.float32
            and x.is_contiguous()
            and x.numel() > 0
            and weight.dtype == torch.float32
            and bias.dtype == torch.float32
            and weight.device == x.device
            and bias.device == x.device
            and weight.is_contiguous()
            and bias.is_contiguous()
            and weight.numel() >= 8192
            and weight.numel() % 4 == 0
            and x.ndim >= len(shape)
            and tuple(x.shape[-len(shape):]) == tuple(shape)
            and x.data_ptr() % 16 == 0
            and weight.data_ptr() % 16 == 0
            and bias.data_ptr() % 16 == 0
        )

        if use_custom:
            return _LayerNormFunction.apply(x, weight, bias, self.ln.eps)

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