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

__device__ __forceinline__ float2 warp_sum(float2 value) {
    #pragma unroll
    for (int offset = 16; offset > 0; offset >>= 1) {
        value.x += __shfl_down_sync(0xffffffff, value.x, offset);
        value.y += __shfl_down_sync(0xffffffff, value.y, offset);
    }
    return value;
}

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

    // A common local origin lets the block reduce centered sums with
    // additions instead of performing a Welford merge at every shuffle.
    const float anchor =
        reinterpret_cast<const float*>(x + chunk_base)[0];

    float4 sum = make_float4(0.f, 0.f, 0.f, 0.f);
    float4 square = make_float4(0.f, 0.f, 0.f, 0.f);

    #pragma unroll 4
    for (int k = 0; k < VECTORS_PER_THREAD; ++k) {
        const float4 value = __ldcs(x + base + k * THREADS);
        const float dx = value.x - anchor;
        const float dy = value.y - anchor;
        const float dz = value.z - anchor;
        const float dw = value.w - anchor;

        sum.x += dx;
        sum.y += dy;
        sum.z += dz;
        sum.w += dw;

        square.x = fmaf(dx, dx, square.x);
        square.y = fmaf(dy, dy, square.y);
        square.z = fmaf(dz, dz, square.z);
        square.w = fmaf(dw, dw, square.w);
    }

    float2 total = warp_sum(make_float2(
        (sum.x + sum.y) + (sum.z + sum.w),
        (square.x + square.y) + (square.z + square.w)));

    __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);
        total = warp_sum(total);

        if (threadIdx.x == 0) {
            constexpr float reciprocal = 1.f / CHUNK;
            const float mean = anchor + total.x * reciprocal;
            const float m2 = fmaxf(
                fmaf(-total.x, total.x * reciprocal, total.y), 0.f);
            partial[static_cast<int64_t>(blockIdx.y) * chunks
                    + blockIdx.x] = make_float2(mean, m2);
        }
    }
}

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_moments(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 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_moments(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_moments(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__ __launch_bounds__(THREADS, 4) 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 int64_t column =
        static_cast<int64_t>(blockIdx.x) * THREADS + threadIdx.x;
    if (column >= vectors_per_row) return;

    const float4 w = __ldg(weight + column);
    const float4 b = __ldg(bias + column);

    #pragma unroll 4
    for (int r = 0; r < ROW_TILE; ++r) {
        const int row = first_row + r;
        if (FULL_TILE || row < rows) {
            const int64_t index =
                static_cast<int64_t>(row) * vectors_per_row + column;
            const float4 v = __ldcs(x + index);
            const float2 s = row_stats[r];

            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 + THREADS - 1) / 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_sums_v4",
    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 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)
```