```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;
constexpr int ROW_TILE = 8;

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;
    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.f, 0.f, 0.f};
        v = warp_reduce(v);
    }
    return v;
}

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

__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.f / 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 component_count = CHUNK / (THREADS * 4);
    float2 a = merge_equal(
        make_float2(mean.x, m2.x),
        make_float2(mean.y, m2.y),
        component_count);
    float2 b = merge_equal(
        make_float2(mean.z, m2.z),
        make_float2(mean.w, m2.w),
        component_count);
    float2 total = merge_equal(a, b, component_count * 2.f);

    float count = CHUNK / THREADS;
    #pragma unroll
    for (int offset = 16; offset > 0; offset >>= 1) {
        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 warp_results[THREADS / 32];
    if ((threadIdx.x & 31) == 0)
        warp_results[threadIdx.x >> 5] = total;
    __syncthreads();

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

        count = CHUNK / (THREADS / 32);
        #pragma unroll
        for (int offset = THREADS / 64; offset > 0; offset >>= 1) {
            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>(row) * chunks + chunk] = total;
    }
}

__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));
    }
}

template<bool FULL_TILE>
__global__ void normalize_affine_tiled(
    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)
{
    int64_t column =
        static_cast<int64_t>(blockIdx.x) * THREADS + threadIdx.x;
    if (column >= vectors_per_row) return;

    // Reuse the affine parameters across eight independent rows.
    float4 w = weight[column];
    float4 b = bias[column];
    int first_row = blockIdx.y * ROW_TILE;

    #pragma unroll
    for (int r = 0; r < ROW_TILE; ++r) {
        int row = first_row + r;
        if (FULL_TILE || row < rows) {
            int64_t index =
                static_cast<int64_t>(row) * vectors_per_row + column;
            float4 v = x[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;
            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();
    int rows = static_cast<int>(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 vectors_per_row = columns / 4;
    dim3 grid(
        (vectors_per_row + THREADS - 1) / THREADS,
        (rows + ROW_TILE - 1) / ROW_TILE);

    const float4* input_ptr =
        reinterpret_cast<const float4*>(x.data_ptr<float>());
    const float4* weight_ptr =
        reinterpret_cast<const float4*>(weight.data_ptr<float>());
    const float4* bias_ptr =
        reinterpret_cast<const float4*>(bias.data_ptr<float>());
    const float2* stats_ptr =
        reinterpret_cast<const float2*>(statistics.data_ptr<float>());
    float4* output_ptr =
        reinterpret_cast<float4*>(output.data_ptr<float>());

    if (rows % ROW_TILE == 0) {
        normalize_affine_tiled<true><<<grid, THREADS, 0, stream>>>(
            input_ptr, weight_ptr, bias_ptr, stats_ptr,
            output_ptr, vectors_per_row, rows);
    } else {
        normalize_affine_tiled<false><<<grid, THREADS, 0, stream>>>(
            input_ptr, weight_ptr, bias_ptr, stats_ptr,
            output_ptr, vectors_per_row, rows);
    }
    C10_CUDA_KERNEL_LAUNCH_CHECK();

    return output;
}
"""

_extension = load_inline(
    name="large_layer_norm_row_tiled_v2",
    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 >= 4096
            and self.columns % 4096 == 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)

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