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


_cuda_source = r"""
#include <torch/extension.h>
#include <c10/cuda/CUDAGuard.h>
#include <c10/cuda/CUDAStream.h>
#include <c10/cuda/CUDAException.h>
#include <cuda_runtime.h>

constexpr int THREADS = 256;
constexpr int CHUNK = 4096;

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

__device__ __forceinline__ Moments merge(Moments a, Moments b) {
    if (a.count == 0) return b;
    if (b.count == 0) return a;
    const int count = a.count + b.count;
    const float delta = b.mean - a.mean;
    const float ratio = float(b.count) / float(count);
    a.mean += delta * ratio;
    a.m2 += b.m2 + delta * delta * (float(a.count) * ratio);
    a.count = count;
    return a;
}

__device__ __forceinline__ Moments warp_reduce(Moments a) {
    const 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 shared[THREADS / 32];
    const int lane = threadIdx.x & 31;
    const int warp = threadIdx.x >> 5;

    a = warp_reduce(a);
    if (lane == 0) shared[warp] = a;
    __syncthreads();

    if (warp == 0) {
        a = lane < THREADS / 32
            ? shared[lane]
            : Moments{0.0f, 0.0f, 0};
        a = warp_reduce(a);
    }
    return a;
}

__global__ void partial_moments(
    const float* __restrict__ x,
    float* __restrict__ partial,
    int64_t width,
    int parts
) {
    const int64_t block = blockIdx.x;
    const int64_t row = block / parts;
    const int part = block % parts;
    const int64_t start = int64_t(part) * CHUNK;

    Moments a{0.0f, 0.0f, 0};

    #pragma unroll
    for (int i = 0; i < CHUNK / THREADS; ++i) {
        const int64_t col = start + threadIdx.x + i * THREADS;
        if (col < width) {
            const float value = x[row * width + col];
            ++a.count;
            const float delta = value - a.mean;
            a.mean += delta / float(a.count);
            a.m2 += delta * (value - a.mean);
        }
    }

    a = block_reduce(a);
    if (threadIdx.x == 0) {
        partial[2 * block] = a.mean;
        partial[2 * block + 1] = a.m2;
    }
}

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

    for (int part = threadIdx.x; part < parts; part += THREADS) {
        const int64_t remaining = width - int64_t(part) * CHUNK;
        Moments b;
        b.mean = partial[2 * (row * parts + part)];
        b.m2 = partial[2 * (row * parts + part) + 1];
        b.count = remaining < CHUNK ? int(remaining) : CHUNK;
        a = merge(a, b);
    }

    a = block_reduce(a);
    if (threadIdx.x == 0) {
        stats[2 * row] = a.mean;
        stats[2 * row + 1] =
            rsqrtf(fmaxf(a.m2 / float(width), 0.0f) + eps);
    }
}

__global__ void apply_vectorized(
    const float4* __restrict__ x,
    const float4* __restrict__ weight,
    const float4* __restrict__ bias,
    const float* __restrict__ stats,
    float4* __restrict__ output,
    int64_t total4,
    int64_t width4
) {
    const int64_t index =
        int64_t(blockIdx.x) * blockDim.x + threadIdx.x;
    if (index >= total4) return;

    const int64_t row = index / width4;
    const int64_t col = index - row * width4;
    const float mean = stats[2 * row];
    const float rstd = stats[2 * row + 1];

    const float4 v = x[index];
    const float4 w = weight[col];
    const float4 b = bias[col];

    float4 out;
    out.x = ((v.x - mean) * rstd) * w.x + b.x;
    out.y = ((v.y - mean) * rstd) * w.y + b.y;
    out.z = ((v.z - mean) * rstd) * w.z + b.z;
    out.w = ((v.w - mean) * rstd) * w.w + b.w;
    output[index] = out;
}

__global__ void apply_scalar(
    const float* __restrict__ x,
    const float* __restrict__ weight,
    const float* __restrict__ bias,
    const float* __restrict__ stats,
    float* __restrict__ output,
    int64_t total,
    int64_t width
) {
    const int64_t index =
        int64_t(blockIdx.x) * blockDim.x + threadIdx.x;
    if (index >= total) return;

    const int64_t row = index / width;
    const int64_t col = index - row * width;
    output[index] =
        ((x[index] - stats[2 * row]) * stats[2 * row + 1])
        * weight[col] + bias[col];
}

torch::Tensor layer_norm_cuda(
    torch::Tensor x,
    torch::Tensor weight,
    torch::Tensor bias,
    double eps
) {
    TORCH_CHECK(x.is_cuda() && weight.is_cuda() && bias.is_cuda(),
                "Expected CUDA tensors");
    TORCH_CHECK(x.device() == weight.device() &&
                x.device() == bias.device(), "Device mismatch");
    TORCH_CHECK(x.scalar_type() == at::kFloat &&
                weight.scalar_type() == at::kFloat &&
                bias.scalar_type() == at::kFloat,
                "Expected float32 tensors");
    TORCH_CHECK(x.is_contiguous() && weight.is_contiguous() &&
                bias.is_contiguous(), "Expected contiguous tensors");

    const c10::cuda::CUDAGuard guard(x.device());
    const int64_t width = weight.numel();
    TORCH_CHECK(width > 0 && bias.numel() == width &&
                x.numel() % width == 0, "Invalid normalization shape");

    auto output = torch::empty_like(x);
    if (x.numel() == 0) return output;

    const int64_t rows = x.numel() / width;
    const int parts = int((width + CHUNK - 1) / CHUNK);
    auto partial = torch::empty({rows, int64_t(parts), 2}, x.options());
    auto stats = torch::empty({rows, 2}, x.options());
    const auto stream = c10::cuda::getCurrentCUDAStream();

    partial_moments<<<rows * parts, THREADS, 0, stream>>>(
        x.data_ptr<float>(), partial.data_ptr<float>(), width, parts);
    C10_CUDA_KERNEL_LAUNCH_CHECK();

    finish_moments<<<rows, THREADS, 0, stream>>>(
        partial.data_ptr<float>(), stats.data_ptr<float>(),
        width, parts, float(eps));
    C10_CUDA_KERNEL_LAUNCH_CHECK();

    const bool aligned =
        (reinterpret_cast<uintptr_t>(x.data_ptr<float>()) % 16 == 0) &&
        (reinterpret_cast<uintptr_t>(weight.data_ptr<float>()) % 16 == 0) &&
        (reinterpret_cast<uintptr_t>(bias.data_ptr<float>()) % 16 == 0);

    if (width % 4 == 0 && aligned) {
        const int64_t total4 = x.numel() / 4;
        apply_vectorized<<<(total4 + THREADS - 1) / THREADS,
                           THREADS, 0, stream>>>(
            reinterpret_cast<const float4*>(x.data_ptr<float>()),
            reinterpret_cast<const float4*>(weight.data_ptr<float>()),
            reinterpret_cast<const float4*>(bias.data_ptr<float>()),
            stats.data_ptr<float>(),
            reinterpret_cast<float4*>(output.data_ptr<float>()),
            total4, width / 4);
    } else {
        apply_scalar<<<(x.numel() + THREADS - 1) / THREADS,
                       THREADS, 0, stream>>>(
            x.data_ptr<float>(), weight.data_ptr<float>(),
            bias.data_ptr<float>(), stats.data_ptr<float>(),
            output.data_ptr<float>(), x.numel(), width);
    }
    C10_CUDA_KERNEL_LAUNCH_CHECK();
    return output;
}
"""

_extension = load_inline(
    name="split_layer_norm_cuda",
    cpp_sources="""
        torch::Tensor layer_norm_cuda(
            torch::Tensor x,
            torch::Tensor weight,
            torch::Tensor bias,
            double eps);
    """,
    cuda_sources=_cuda_source,
    functions=["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, eps):
        ctx.normalized_shape = normalized_shape
        ctx.eps = eps
        ctx.save_for_backward(x, weight, bias)
        return _extension.layer_norm_cuda(
            x.contiguous(), weight.contiguous(), bias.contiguous(), eps
        )

    @staticmethod
    def backward(ctx, grad_output):
        x, weight, bias = ctx.saved_tensors
        with torch.enable_grad():
            x_grad = x.detach().requires_grad_(True)
            weight_grad = weight.detach().requires_grad_(True)
            bias_grad = bias.detach().requires_grad_(True)
            output = F.layer_norm(
                x_grad,
                ctx.normalized_shape,
                weight_grad,
                bias_grad,
                ctx.eps,
            )
            grads = torch.autograd.grad(
                output,
                (x_grad, weight_grad, bias_grad),
                grad_output,
            )
        return grads[0], grads[1], grads[2], None, 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:
        shape = self.ln.normalized_shape
        if (
            not x.is_cuda
            or x.dtype != torch.float32
            or self.ln.weight.dtype != torch.float32
            or self.ln.bias.dtype != torch.float32
            or x.ndim < len(shape)
            or tuple(x.shape[-len(shape):]) != shape
            or self.ln.weight.numel() == 0
        ):
            return self.ln(x)

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