```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 = 16384;
constexpr int VECTORS_PER_THREAD = CHUNK / (4 * THREADS);
constexpr int ROW_TILE = 16;
constexpr int COLUMN_THREADS = 128;
constexpr int ROW_LANES = THREADS / COLUMN_THREADS;

__device__ __forceinline__ float2 merge_equal(
    float2 a, float2 b, float count)
{
    const float delta = b.x - a.x;
    return make_float2(
        a.x + delta * 0.5f,
        a.y + b.y + delta * delta * (count * 0.5f));
}

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

__device__ __forceinline__ Moments merge(Moments a, Moments b)
{
    if (a.count == 0.f) return b;
    if (b.count == 0.f) return a;
    const float count = a.count + b.count;
    const float delta = b.mean - a.mean;
    const 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 value)
{
    #pragma unroll
    for (int offset = 16; offset > 0; offset >>= 1) {
        Moments other{
            __shfl_down_sync(0xffffffff, value.mean, offset),
            __shfl_down_sync(0xffffffff, value.m2, offset),
            __shfl_down_sync(0xffffffff, value.count, offset)
        };
        if ((threadIdx.x & 31) + offset < 32)
            value = merge(value, other);
    }
    return value;
}

__global__ void partial_moments(
    const float4* __restrict__ x,
    float2* __restrict__ partial,
    int64_t vectors_per_row,
    int chunks)
{
    const int64_t base =
        static_cast<int64_t>(blockIdx.y) * vectors_per_row
        + static_cast<int64_t>(blockIdx.x) * (CHUNK / 4)
        + threadIdx.x;

    const float4 first = x[base];
    const float anchor = first.x;
    const float dy = first.y - anchor;
    const float dz = first.z - anchor;
    const float dw = first.w - anchor;

    float4 sum = make_float4(0.f, dy, dz, dw);
    float4 square = make_float4(0.f, dy * dy, dz * dz, dw * dw);

    #pragma unroll
    for (int k = 1; k < VECTORS_PER_THREAD; ++k) {
        const float4 v = x[base + k * THREADS];
        const float a = v.x - anchor;
        const float b = v.y - anchor;
        const float c = v.z - anchor;
        const float d = v.w - anchor;

        sum.x += a;
        sum.y += b;
        sum.z += c;
        sum.w += d;

        square.x = fmaf(a, a, square.x);
        square.y = fmaf(b, b, square.y);
        square.z = fmaf(c, c, square.z);
        square.w = fmaf(d, d, square.w);
    }

    constexpr float reciprocal = 1.f / (4 * VECTORS_PER_THREAD);
    const float s = (sum.x + sum.y) + (sum.z + sum.w);
    const float q = (square.x + square.y) + (square.z + square.w);

    float2 total = make_float2(
        anchor + s * reciprocal,
        fmaxf(fmaf(-s, s * reciprocal, q), 0.f));

    float count = CHUNK / THREADS;
    #pragma unroll
    for (int offset = 16; offset > 0; offset >>= 1) {
        const float2 other = make_float2(
            __shfl_down_sync(0xffffffff, total.x, offset),
            __shfl_down_sync(0xffffffff, total.y, offset));
        total = merge_equal(total, other, count);
        count *= 2.f;
    }

    __shared__ float2 shared[THREADS / 32];
    if ((threadIdx.x & 31) == 0)
        shared[threadIdx.x >> 5] = total;
    __syncthreads();

    if (threadIdx.x < 32) {
        total = threadIdx.x < THREADS / 32
            ? shared[threadIdx.x]
            : make_float2(0.f, 0.f);
        count = CHUNK / (THREADS / 32);

        #pragma unroll
        for (int offset = THREADS / 64; offset > 0; offset >>= 1) {
            const float2 other = make_float2(
                __shfl_down_sync(0xffffffff, total.x, offset),
                __shfl_down_sync(0xffffffff, total.y, offset));
            total = merge_equal(total, other, count);
            count *= 2.f;
        }

        if (threadIdx.x == 0) {
            partial[static_cast<int64_t>(blockIdx.y) * chunks
                    + blockIdx.x] = total;
        }
    }
}

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

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

    value = warp_reduce(value);

    __shared__ Moments shared[THREADS / 32];
    if ((threadIdx.x & 31) == 0)
        shared[threadIdx.x >> 5] = value;
    __syncthreads();

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

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

template<bool FULL_TILE>
__global__ void normalize_affine(
    const float4* __restrict__ x,
    const float4* __restrict__ weight,
    const float4* __restrict__ bias,
    const float2* __restrict__ statistics,
    float4* __restrict__ output,
    int64_t vectors_per_row,
    int rows)
{
    __shared__ float2 row_stats[ROW_TILE];

    const int first_row = blockIdx.y * ROW_TILE;
    if (threadIdx.x < ROW_TILE) {
        const int row = first_row + threadIdx.x;
        row_stats[threadIdx.x] = (FULL_TILE || row < rows)
            ? statistics[row]
            : make_float2(0.f, 0.f);
    }
    __syncthreads();

    const int column_lane = threadIdx.x % COLUMN_THREADS;
    const int row_lane = threadIdx.x / COLUMN_THREADS;
    const int64_t column =
        static_cast<int64_t>(blockIdx.x) * COLUMN_THREADS + column_lane;

    if (column >= vectors_per_row) return;

    const float4 w = weight[column];
    const float4 b = bias[column];

    #pragma unroll 4
    for (int r = 0; r < ROW_TILE / ROW_LANES; ++r) {
        const int local_row = r * ROW_LANES + row_lane;
        const int row = first_row + local_row;

        if (FULL_TILE || row < rows) {
            const int64_t index =
                static_cast<int64_t>(row) * vectors_per_row + column;
            const float4 v = x[index];
            const float2 s = row_stats[local_row];

            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;

            __stcs(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());

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

    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();

    const float4* xp =
        reinterpret_cast<const float4*>(x.data_ptr<float>());
    const float4* wp =
        reinterpret_cast<const float4*>(weight.data_ptr<float>());
    const float4* bp =
        reinterpret_cast<const float4*>(bias.data_ptr<float>());
    float4* op =
        reinterpret_cast<float4*>(output.data_ptr<float>());
    float2* pp =
        reinterpret_cast<float2*>(partial.data_ptr<float>());
    float2* sp =
        reinterpret_cast<float2*>(statistics.data_ptr<float>());

    partial_moments<<<dim3(chunks, rows), THREADS, 0, stream>>>(
        xp, pp, vectors_per_row, chunks);
    C10_CUDA_KERNEL_LAUNCH_CHECK();

    finish_moments<<<rows, THREADS, 0, stream>>>(
        pp, sp, chunks, columns, static_cast<float>(epsilon));
    C10_CUDA_KERNEL_LAUNCH_CHECK();

    const dim3 grid(
        static_cast<unsigned>(
            (vectors_per_row + COLUMN_THREADS - 1) / COLUMN_THREADS),
        (rows + ROW_TILE - 1) / ROW_TILE);

    if (rows % ROW_TILE == 0) {
        normalize_affine<true><<<grid, THREADS, 0, stream>>>(
            xp, wp, bp, sp, op, vectors_per_row, rows);
    } else {
        normalize_affine<false><<<grid, THREADS, 0, stream>>>(
            xp, wp, bp, sp, op, vectors_per_row, rows);
    }
    C10_CUDA_KERNEL_LAUNCH_CHECK();

    return output;
}
"""

_extension = load_inline(
    name="large_layer_norm_centered_interleaved_v5",
    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"],
    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 >= 16384
            and self.columns % 16384 == 0
            and self.columns // 16384 <= 2147483647
            and 0 < x.numel() // self.columns <= 65535
            and x.data_ptr() % 16 == 0
            and weight.data_ptr() % 16 == 0
            and bias.data_ptr() % 16 == 0
        )
        if not supported:
            return self.ln(x)

        if torch.is_grad_enabled() and (
            x.requires_grad or weight.requires_grad or bias.requires_grad
        ):
            return _LayerNormFunction.apply(
                x, weight, bias, self.normalized_shape, self.ln.eps
            )

        return _extension.large_layer_norm_cuda(
            x, weight, bias, self.ln.eps
        )
```