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


_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 a, Moments b) {
    if (a.count == 0.0f) return b;
    if (b.count == 0.0f) return a;

    float count = a.count + b.count;
    float delta = b.mean - a.mean;
    float ratio = b.count / count;
    return {
        a.mean + delta * ratio,
        a.m2 + b.m2 + delta * delta * a.count * ratio,
        count
    };
}

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

__device__ __forceinline__ Moments block_reduce(Moments v) {
    __shared__ Moments shared[THREADS / 32];

    v = warp_reduce(v);
    if ((threadIdx.x & 31) == 0)
        shared[threadIdx.x >> 5] = v;
    __syncthreads();

    if (threadIdx.x < 32) {
        v = threadIdx.x < THREADS / 32
            ? shared[threadIdx.x]
            : Moments{0.0f, 0.0f, 0.0f};
        v = warp_reduce(v);
    }
    return v;
}

__global__ void partial_moments(
    const float* __restrict__ x,
    float2* __restrict__ partial,
    int64_t columns,
    int chunks)
{
    int chunk = blockIdx.x;
    int row = blockIdx.y;
    int64_t base = static_cast<int64_t>(row) * columns
                 + static_cast<int64_t>(chunk) * CHUNK;
    const float4* input = reinterpret_cast<const float4*>(x + base);

    float4 mean = make_float4(0.f, 0.f, 0.f, 0.f);
    float4 m2 = make_float4(0.f, 0.f, 0.f, 0.f);

    #pragma unroll
    for (int k = 0; k < CHUNK / (THREADS * 4); ++k) {
        float4 v = input[threadIdx.x + k * THREADS];
        float reciprocal = 1.0f / static_cast<float>(k + 1);

        float dx = v.x - mean.x;
        float dy = v.y - mean.y;
        float dz = v.z - mean.z;
        float dw = v.w - mean.w;

        mean.x += dx * reciprocal;
        mean.y += dy * reciprocal;
        mean.z += dz * reciprocal;
        mean.w += dw * reciprocal;

        m2.x += dx * (v.x - mean.x);
        m2.y += dy * (v.y - mean.y);
        m2.z += dz * (v.z - mean.z);
        m2.w += dw * (v.w - mean.w);
    }

    constexpr float count = CHUNK / (THREADS * 4);
    Moments a = merge(
        Moments{mean.x, m2.x, count},
        Moments{mean.y, m2.y, count});
    Moments b = merge(
        Moments{mean.z, m2.z, count},
        Moments{mean.w, m2.w, count});
    Moments total = block_reduce(merge(a, b));

    if (threadIdx.x == 0)
        partial[static_cast<int64_t>(row) * chunks + chunk] =
            make_float2(total.mean, total.m2);
}

__global__ void finish_moments(
    const float2* __restrict__ partial,
    float2* __restrict__ statistics,
    int chunks,
    int64_t columns,
    float epsilon)
{
    int row = blockIdx.x;
    Moments value{0.f, 0.f, 0.f};

    for (int i = threadIdx.x; i < chunks; i += THREADS) {
        float2 p = partial[static_cast<int64_t>(row) * chunks + i];
        value = merge(value, Moments{p.x, p.y, static_cast<float>(CHUNK)});
    }

    value = block_reduce(value);

    if (threadIdx.x == 0) {
        float variance = fmaxf(value.m2 / static_cast<float>(columns), 0.f);
        statistics[row] = make_float2(
            value.mean, rsqrtf(variance + epsilon));
    }
}

__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 columns)
{
    int row = blockIdx.y;
    int64_t vector_index =
        static_cast<int64_t>(blockIdx.x) * THREADS + threadIdx.x;
    int64_t vectors_per_row = columns / 4;

    if (vector_index >= vectors_per_row) return;

    int64_t index = static_cast<int64_t>(row) * vectors_per_row
                  + vector_index;

    float4 v = reinterpret_cast<const float4*>(x)[index];
    float4 w = reinterpret_cast<const float4*>(weight)[vector_index];
    float4 b = reinterpret_cast<const float4*>(bias)[vector_index];
    float2 stats = statistics[row];

    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;

    reinterpret_cast<float4*>(output)[index] = result;
}

torch::Tensor large_layer_norm_cuda(
    torch::Tensor x,
    torch::Tensor weight,
    torch::Tensor bias,
    double epsilon)
{
    c10::cuda::CUDAGuard guard(x.device());

    int64_t columns = weight.numel();
    int64_t rows = x.numel() / columns;
    int chunks = static_cast<int>(columns / CHUNK);

    auto output = torch::empty_like(x);
    auto partial = torch::empty({rows, chunks, 2}, x.options());
    auto statistics = torch::empty({rows, 2}, x.options());

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

    partial_moments<<<dim3(chunks, rows), THREADS, 0, stream>>>(
        x.data_ptr<float>(),
        reinterpret_cast<float2*>(partial.data_ptr<float>()),
        columns,
        chunks);
    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>()),
        chunks,
        columns,
        static_cast<float>(epsilon));
    C10_CUDA_KERNEL_LAUNCH_CHECK();

    int64_t blocks = (columns / 4 + THREADS - 1) / THREADS;
    normalize_affine<<<dim3(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>(),
        columns);
    C10_CUDA_KERNEL_LAUNCH_CHECK();

    return output;
}
"""

_extension = load_inline(
    name="large_layer_norm_welford_cuda_v1",
    cpp_sources=r"""
        #include <torch/extension.h>
        torch::Tensor large_layer_norm_cuda(
            torch::Tensor x,
            torch::Tensor weight,
            torch::Tensor bias,
            double epsilon);
    """,
    cuda_sources=_cuda_source,
    functions=["large_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, normalized_shape, epsilon):
        ctx.normalized_shape = normalized_shape
        ctx.epsilon = epsilon
        ctx.save_for_backward(x, weight, bias)
        return _extension.large_layer_norm_cuda(x, weight, bias, epsilon)

    @staticmethod
    def backward(ctx, grad_output):
        x, weight, bias = ctx.saved_tensors
        _, mean, rstd = torch.ops.aten.native_layer_norm.default(
            x, ctx.normalized_shape, weight, bias, ctx.epsilon
        )
        dx, dw, db = torch.ops.aten.native_layer_norm_backward.default(
            grad_output.contiguous(),
            x,
            ctx.normalized_shape,
            mean,
            rstd,
            weight,
            bias,
            list(ctx.needs_input_grad[:3]),
        )
        return dx, dw, db, None, None


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.columns = self.ln.weight.numel()

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        weight, bias = self.ln.weight, self.ln.bias
        supported = (
            x.is_cuda
            and x.dtype == torch.float32
            and x.is_contiguous()
            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 tuple(x.shape[-len(self.normalized_shape):])
                == self.normalized_shape
            and self.columns >= 4096
            and self.columns % 4096 == 0
            and 0 < x.numel() // self.columns <= 65535
        )
        if not supported:
            return self.ln(x)

        return _LayerNormFunction.apply(
            x, weight, bias, self.normalized_shape, self.ln.eps
        )
```