```python
# Copyright (c) Meta Platforms, Inc. and affiliates.
# All rights reserved.
#
# This source code is licensed under the license found in the
# LICENSE file in the root directory of this source tree.

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

# Custom fused LayerNorm CUDA implementation.
#
# The workload is a small number of very long rows (e.g. 16 rows x 4M floats),
# so the whole computation is DRAM-bandwidth bound.  The previous revision used
# three kernels (partial sums -> stats -> normalize) which move
# read(x) + read(x) + read(w,b) + write(out) bytes.  This revision keeps the
# same overall structure but tunes the bandwidth-critical kernels:
#
#   * partial-sums kernel: grid-stride column ownership with 8 independent
#     float4 accumulators (was: one 4KB chunk per block), giving a much larger
#     memory-level parallelism per thread and better DRAM utilization.
#   * normalize kernel: each thread owns a column and loops over all rows, so
#     weight/bias are fetched exactly once.  The row loop is unrolled by 8
#     (was 4) with software-pipelined load batches to keep more loads in
#     flight, and rows are visited in reverse so the tail of the reduction
#     pass (still resident in L2) is consumed first.
#   * stats kernel: fused reduction of both partial arrays in one pass.
#
# Fallback kernels handle arbitrary shapes / alignments / dtypes.

layer_norm_cuda_source = r"""
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include <c10/cuda/CUDAGuard.h>
#include <cuda_runtime.h>
#include <algorithm>
#include <cstdint>

#define LN_THREADS 256
#define LN_WARPS (LN_THREADS / 32)

__device__ __forceinline__ void ln_warp_reduce(float &sum, float &sumsq) {
    #pragma unroll
    for (int off = 16; off > 0; off >>= 1) {
        sum   += __shfl_down_sync(0xffffffffu, sum, off);
        sumsq += __shfl_down_sync(0xffffffffu, sumsq, off);
    }
}

__device__ __forceinline__ void ln_block_reduce(float &sum, float &sumsq,
                                                float *sh_sum, float *sh_sumsq) {
    const int lane = threadIdx.x & 31;
    const int warp = threadIdx.x >> 5;
    ln_warp_reduce(sum, sumsq);
    if (lane == 0) {
        sh_sum[warp]   = sum;
        sh_sumsq[warp] = sumsq;
    }
    __syncthreads();
    if (warp == 0) {
        const int nwarp = (blockDim.x + 31) >> 5;
        sum   = (lane < nwarp) ? sh_sum[lane]   : 0.0f;
        sumsq = (lane < nwarp) ? sh_sumsq[lane] : 0.0f;
        ln_warp_reduce(sum, sumsq);
    }
}

// ---------------- Kernel 1: partial sums, grid-stride, 8 accumulators ------
__global__ void ln_partial_sums_vec4(const float* __restrict__ x,
                                     float* __restrict__ ps,
                                     float* __restrict__ pq,
                                     int N, int splits) {
    const int row   = blockIdx.y;
    const int n4    = N >> 2;
    const float4* xv = reinterpret_cast<const float4*>(x + (long long)row * N);

    const int gstride = gridDim.x * blockDim.x;
    const int tid0    = blockIdx.x * blockDim.x + threadIdx.x;

    float s[8], q[8];
    #pragma unroll
    for (int k = 0; k < 8; ++k) { s[k] = 0.f; q[k] = 0.f; }

    int i = tid0;
    const int stride8 = gstride * 8;
    for (; i + 7 * gstride < n4; i += stride8) {
        #pragma unroll
        for (int k = 0; k < 8; ++k) {
            const float4 v = xv[i + k * gstride];
            s[k] += v.x + v.y;
            s[k] += v.z + v.w;
            q[k] += v.x * v.x + v.y * v.y;
            q[k] += v.z * v.z + v.w * v.w;
        }
    }
    for (; i < n4; i += gstride) {
        const float4 v = xv[i];
        s[0] += v.x + v.y;
        s[0] += v.z + v.w;
        q[0] += v.x * v.x + v.y * v.y;
        q[0] += v.z * v.z + v.w * v.w;
    }

    float sum = 0.f, sumsq = 0.f;
    #pragma unroll
    for (int k = 0; k < 8; ++k) { sum += s[k]; sumsq += q[k]; }

    __shared__ float sh_sum[LN_WARPS];
    __shared__ float sh_sumsq[LN_WARPS];
    ln_block_reduce(sum, sumsq, sh_sum, sh_sumsq);

    if (threadIdx.x == 0) {
        ps[(long long)row * splits + blockIdx.x] = sum;
        pq[(long long)row * splits + blockIdx.x] = sumsq;
    }
}

// ---------------- Kernel 1: partial sums (scalar fallback) ----------------
__global__ void ln_partial_sums_scalar(const float* __restrict__ x,
                                       float* __restrict__ ps,
                                       float* __restrict__ pq,
                                       int N, int splits) {
    const int row = blockIdx.y;
    const float* row_ptr = x + (long long)row * N;

    const int gstride = gridDim.x * blockDim.x;
    float s0 = 0.f, q0 = 0.f, s1 = 0.f, q1 = 0.f;
    float s2 = 0.f, q2 = 0.f, s3 = 0.f, q3 = 0.f;

    int i = blockIdx.x * blockDim.x + threadIdx.x;
    for (; i + 3 * gstride < N; i += 4 * gstride) {
        const float v0 = row_ptr[i];
        const float v1 = row_ptr[i + gstride];
        const float v2 = row_ptr[i + 2 * gstride];
        const float v3 = row_ptr[i + 3 * gstride];
        s0 += v0; q0 += v0 * v0;
        s1 += v1; q1 += v1 * v1;
        s2 += v2; q2 += v2 * v2;
        s3 += v3; q3 += v3 * v3;
    }
    for (; i < N; i += gstride) {
        const float v = row_ptr[i];
        s0 += v; q0 += v * v;
    }
    float sum = (s0 + s1) + (s2 + s3);
    float sumsq = (q0 + q1) + (q2 + q3);

    __shared__ float sh_sum[LN_WARPS];
    __shared__ float sh_sumsq[LN_WARPS];
    ln_block_reduce(sum, sumsq, sh_sum, sh_sumsq);

    if (threadIdx.x == 0) {
        ps[(long long)row * splits + blockIdx.x] = sum;
        pq[(long long)row * splits + blockIdx.x] = sumsq;
    }
}

// ---------------- Kernel 2: final per-row statistics ----------------
__global__ void ln_final_stats(const float* __restrict__ ps,
                               const float* __restrict__ pq,
                               float* __restrict__ mean,
                               float* __restrict__ rstd,
                               int splits, int N, float eps) {
    const int row = blockIdx.x;
    float sum = 0.f, sumsq = 0.f;
    const long long base = (long long)row * splits;
    for (int i = threadIdx.x; i < splits; i += blockDim.x) {
        sum   += ps[base + i];
        sumsq += pq[base + i];
    }

    __shared__ float sh_sum[LN_WARPS];
    __shared__ float sh_sumsq[LN_WARPS];
    ln_block_reduce(sum, sumsq, sh_sum, sh_sumsq);

    if (threadIdx.x == 0) {
        const float m = sum / (float)N;
        float var = sumsq / (float)N - m * m;
        var = var < 0.f ? 0.f : var;  // guard against tiny negative roundoff
        mean[row] = m;
        rstd[row] = rsqrtf(var + eps);
    }
}

// ------------- Kernel 3a: normalize, thread owns a column and loops over
// all rows (weight/bias loaded once into registers), unroll 8 with
// software-pipelined load batches. -----------------------------------------
__global__ void ln_apply_rows_vec4(const float* __restrict__ x,
                                   float* __restrict__ out,
                                   const float* __restrict__ mean,
                                   const float* __restrict__ rstd,
                                   const float* __restrict__ weight,
                                   const float* __restrict__ bias,
                                   int N, int rows) {
    const int n4 = N >> 2;
    const int c = blockIdx.x * blockDim.x + threadIdx.x;
    if (c >= n4) return;

    const float4 w = reinterpret_cast<const float4*>(weight)[c];
    const float4 b = reinterpret_cast<const float4*>(bias)[c];
    const float4* xv = reinterpret_cast<const float4*>(x) + c;
    float4* ov = reinterpret_cast<float4*>(out) + c;

    int r = rows;
    for (; r >= 8; r -= 8) {
        const long long o0 = (long long)(r - 1) * n4;
        const long long o1 = (long long)(r - 2) * n4;
        const long long o2 = (long long)(r - 3) * n4;
        const long long o3 = (long long)(r - 4) * n4;
        const long long o4 = (long long)(r - 5) * n4;
        const long long o5 = (long long)(r - 6) * n4;
        const long long o6 = (long long)(r - 7) * n4;
        const long long o7 = (long long)(r - 8) * n4;
        const float4 v0 = __ldcs(xv + o0);
        const float4 v1 = __ldcs(xv + o1);
        const float4 v2 = __ldcs(xv + o2);
        const float4 v3 = __ldcs(xv + o3);
        const float4 v4 = __ldcs(xv + o4);
        const float4 v5 = __ldcs(xv + o5);
        const float4 v6 = __ldcs(xv + o6);
        const float4 v7 = __ldcs(xv + o7);
        const float m0 = mean[r - 1], t0 = rstd[r - 1];
        const float m1 = mean[r - 2], t1 = rstd[r - 2];
        const float m2 = mean[r - 3], t2 = rstd[r - 3];
        const float m3 = mean[r - 4], t3 = rstd[r - 4];
        const float m4 = mean[r - 5], t4 = rstd[r - 5];
        const float m5 = mean[r - 6], t5 = rstd[r - 6];
        const float m6 = mean[r - 7], t6 = rstd[r - 7];
        const float m7 = mean[r - 8], t7 = rstd[r - 8];
        float4 u;
        u.x = (v0.x - m0) * t0 * w.x + b.x;
        u.y = (v0.y - m0) * t0 * w.y + b.y;
        u.z = (v0.z - m0) * t0 * w.z + b.z;
        u.w = (v0.w - m0) * t0 * w.w + b.w;
        __stcs(ov + o0, u);
        u.x = (v1.x - m1) * t1 * w.x + b.x;
        u.y = (v1.y - m1) * t1 * w.y + b.y;
        u.z = (v1.z - m1) * t1 * w.z + b.z;
        u.w = (v1.w - m1) * t1 * w.w + b.w;
        __stcs(ov + o1, u);
        u.x = (v2.x - m2) * t2 * w.x + b.x;
        u.y = (v2.y - m2) * t2 * w.y + b.y;
        u.z = (v2.z - m2) * t2 * w.z + b.z;
        u.w = (v2.w - m2) * t2 * w.w + b.w;
        __stcs(ov + o2, u);
        u.x = (v3.x - m3) * t3 * w.x + b.x;
        u.y = (v3.y - m3) * t3 * w.y + b.y;
        u.z = (v3.z - m3) * t3 * w.z + b.z;
        u.w = (v3.w - m3) * t3 * w.w + b.w;
        __stcs(ov + o3, u);
        u.x = (v4.x - m4) * t4 * w.x + b.x;
        u.y = (v4.y - m4) * t4 * w.y + b.y;
        u.z = (v4.z - m4) * t4 * w.z + b.z;
        u.w = (v4.w - m4) * t4 * w.w + b.w;
        __stcs(ov + o4, u);
        u.x = (v5.x - m5) * t5 * w.x + b.x;
        u.y = (v5.y - m5) * t5 * w.y + b.y;
        u.z = (v5.z - m5) * t5 * w.z + b.z;
        u.w = (v5.w - m5) * t5 * w.w + b.w;
        __stcs(ov + o5, u);
        u.x = (v6.x - m6) * t6 * w.x + b.x;
        u.y = (v6.y - m6) * t6 * w.y + b.y;
        u.z = (v6.z - m6) * t6 * w.z + b.z;
        u.w = (v6.w - m6) * t6 * w.w + b.w;
        __stcs(ov + o6, u);
        u.x = (v7.x - m7) * t7 * w.x + b.x;
        u.y = (v7.y - m7) * t7 * w.y + b.y;
        u.z = (v7.z - m7) * t7 * w.z + b.z;
        u.w = (v7.w - m7) * t7 * w.w + b.w;
        __stcs(ov + o7, u);
    }
    for (; r > 0; --r) {
        const long long o = (long long)(r - 1) * n4;
        const float4 v = __ldcs(xv + o);
        const float m = mean[r - 1], t = rstd[r - 1];
        float4 u;
        u.x = (v.x - m) * t * w.x + b.x;
        u.y = (v.y - m) * t * w.y + b.y;
        u.z = (v.z - m) * t * w.z + b.z;
        u.w = (v.w - m) * t * w.w + b.w;
        __stcs(ov + o, u);
    }
}

__global__ void ln_apply_rows_scalar(const float* __restrict__ x,
                                     float* __restrict__ out,
                                     const float* __restrict__ mean,
                                     const float* __restrict__ rstd,
                                     const float* __restrict__ weight,
                                     const float* __restrict__ bias,
                                     int N, int rows) {
    const int c = blockIdx.x * blockDim.x + threadIdx.x;
    if (c >= N) return;

    const float w = weight[c];
    const float b = bias[c];
    const float* xp = x + c;
    float* op = out + c;

    int r = rows;
    for (; r >= 4; r -= 4) {
        const long long o0 = (long long)(r - 1) * N;
        const long long o1 = (long long)(r - 2) * N;
        const long long o2 = (long long)(r - 3) * N;
        const long long o3 = (long long)(r - 4) * N;
        const float v0 = __ldcs(xp + o0);
        const float v1 = __ldcs(xp + o1);
        const float v2 = __ldcs(xp + o2);
        const float v3 = __ldcs(xp + o3);
        __stcs(op + o0, (v0 - mean[r - 1]) * rstd[r - 1] * w + b);
        __stcs(op + o1, (v1 - mean[r - 2]) * rstd[r - 2] * w + b);
        __stcs(op + o2, (v2 - mean[r - 3]) * rstd[r - 3] * w + b);
        __stcs(op + o3, (v3 - mean[r - 4]) * rstd[r - 4] * w + b);
    }
    for (; r > 0; --r) {
        const long long o = (long long)(r - 1) * N;
        __stcs(op + o, (__ldcs(xp + o) - mean[r - 1]) * rstd[r - 1] * w + b);
    }
}

// ------------- Kernel 3b: fallback 2D normalize (vectorized) -----------
__global__ void ln_normalize_vec4(const float* __restrict__ x,
                                  float* __restrict__ out,
                                  const float* __restrict__ mean,
                                  const float* __restrict__ rstd,
                                  const float* __restrict__ weight,
                                  const float* __restrict__ bias,
                                  int N) {
    const int row = blockIdx.y;
    const float m = mean[row];
    const float r = rstd[row];

    const float4* xv = reinterpret_cast<const float4*>(x + (long long)row * N);
    const float4* wv = reinterpret_cast<const float4*>(weight);
    const float4* bv = reinterpret_cast<const float4*>(bias);
    float4* ov = reinterpret_cast<float4*>(out + (long long)row * N);

    const int n4 = N >> 2;
    const int stride = gridDim.x * blockDim.x;
    for (int i = blockIdx.x * blockDim.x + threadIdx.x; i < n4; i += stride) {
        const float4 v = xv[i];
        const float4 w = wv[i];
        const float4 b = bv[i];
        float4 o;
        o.x = (v.x - m) * r * w.x + b.x;
        o.y = (v.y - m) * r * w.y + b.y;
        o.z = (v.z - m) * r * w.z + b.z;
        o.w = (v.w - m) * r * w.w + b.w;
        ov[i] = o;
    }
}

// ------------- Kernel 3b: fallback 2D normalize (scalar) ---------------
__global__ void ln_normalize_scalar(const float* __restrict__ x,
                                    float* __restrict__ out,
                                    const float* __restrict__ mean,
                                    const float* __restrict__ rstd,
                                    const float* __restrict__ weight,
                                    const float* __restrict__ bias,
                                    int N) {
    const int row = blockIdx.y;
    const float m = mean[row];
    const float r = rstd[row];
    const float* row_x = x + (long long)row * N;
    float* row_o = out + (long long)row * N;

    const int stride = gridDim.x * blockDim.x;
    for (int i = blockIdx.x * blockDim.x + threadIdx.x; i < N; i += stride) {
        row_o[i] = (row_x[i] - m) * r * weight[i] + bias[i];
    }
}

// ---------------- Host wrapper ----------------
torch::Tensor layernorm_forward(torch::Tensor x, torch::Tensor weight,
                                torch::Tensor bias, double eps) {
    const c10::cuda::CUDAGuard device_guard(x.device());
    TORCH_CHECK(x.is_cuda() && weight.is_cuda() && bias.is_cuda(),
                "All tensors must be CUDA tensors");

    auto xc = x.contiguous();
    auto wc = weight.contiguous();
    auto bc = bias.contiguous();
    TORCH_CHECK(xc.scalar_type() == at::kFloat &&
                wc.scalar_type() == at::kFloat &&
                bc.scalar_type() == at::kFloat,
                "Only float32 is supported by the custom kernel");

    const int64_t N = wc.numel();
    const int64_t rows = (N > 0) ? (xc.numel() / N) : 0;
    auto out = at::empty_like(xc);
    if (rows == 0 || N == 0) return out;
    TORCH_CHECK(N < (int64_t)1 << 31, "Normalized dim too large");
    TORCH_CHECK(rows < (int64_t)1 << 31, "Too many rows");

    auto stream = at::cuda::getCurrentCUDAStream();
    const float* xp = xc.data_ptr<float>();
    const float* wp = wc.data_ptr<float>();
    const float* bp = bc.data_ptr<float>();
    float* op = out.data_ptr<float>();

    const bool vec4 = (N % 4 == 0) &&
        ((reinterpret_cast<uintptr_t>(xp) & 15) == 0) &&
        ((reinterpret_cast<uintptr_t>(op) & 15) == 0) &&
        ((reinterpret_cast<uintptr_t>(wp) & 15) == 0) &&
        ((reinterpret_cast<uintptr_t>(bp) & 15) == 0);

    // Split count for the reduction: enough blocks to fill the device, but
    // few enough that per-block reduction tails stay negligible.
    int splits;
    if (vec4) {
        const int64_t n4 = N >> 2;
        splits = (int)std::max<int64_t>(
            32, std::min<int64_t>(2048, n4 / (LN_THREADS * 8)));
    } else {
        splits = (int)std::max<int64_t>(
            32, std::min<int64_t>(2048, N / (LN_THREADS * 4)));
    }

    auto partials = at::empty({rows * (int64_t)splits * 2}, xc.options());
    auto mean     = at::empty({rows}, xc.options());
    auto rstd     = at::empty({rows}, xc.options());
    float* ps = partials.data_ptr<float>();
    float* pq = ps + rows * (int64_t)splits;
    float* mp = mean.data_ptr<float>();
    float* rp = rstd.data_ptr<float>();

    const int64_t maxY = 65535;
    // Pass 1: partial reductions (chunk rows to respect the grid-y limit).
    for (int64_t r0 = 0; r0 < rows; r0 += maxY) {
        const int rc = (int)std::min<int64_t>(maxY, rows - r0);
        dim3 grid1(splits, rc);
        if (vec4) {
            ln_partial_sums_vec4<<<grid1, LN_THREADS, 0, stream>>>(
                xp + r0 * N, ps + r0 * splits, pq + r0 * splits, (int)N, splits);
        } else {
            ln_partial_sums_scalar<<<grid1, LN_THREADS, 0, stream>>>(
                xp + r0 * N, ps + r0 * splits, pq + r0 * splits, (int)N, splits);
        }
    }

    // Pass 2: per-row mean / rstd.
    ln_final_stats<<<(int)rows, LN_THREADS, 0, stream>>>(
        ps, pq, mp, rp, splits, (int)N, (float)eps);

    // Pass 3: apply normalization with affine transform.
    const int64_t units = vec4 ? (N >> 2) : N;
    const int64_t col_blocks = (units + LN_THREADS - 1) / LN_THREADS;

    if (col_blocks >= 400 && rows <= 32768) {
        // Preferred: each thread owns a column and loops over all rows, so
        // weight/bias are fetched exactly once.
        dim3 grid3((unsigned int)col_blocks);
        if (vec4) {
            ln_apply_rows_vec4<<<grid3, LN_THREADS, 0, stream>>>(
                xp, op, mp, rp, wp, bp, (int)N, (int)rows);
        } else {
            ln_apply_rows_scalar<<<grid3, LN_THREADS, 0, stream>>>(
                xp, op, mp, rp, wp, bp, (int)N, (int)rows);
        }
    } else {
        // Fallback for narrow normalized dims or extremely many rows.
        for (int64_t r0 = 0; r0 < rows; r0 += maxY) {
            const int rc = (int)std::min<int64_t>(maxY, rows - r0);
            if (vec4) {
                const int64_t n4 = N >> 2;
                const int xb = (int)std::min<int64_t>(4096, (n4 + LN_THREADS - 1) / LN_THREADS);
                dim3 grid3(xb, rc);
                ln_normalize_vec4<<<grid3, LN_THREADS, 0, stream>>>(
                    xp + r0 * N, op + r0 * N, mp + r0, rp + r0, wp, bp, (int)N);
            } else {
                const int xb = (int)std::min<int64_t>(4096, (N + LN_THREADS - 1) / LN_THREADS);
                dim3 grid3(xb, rc);
                ln_normalize_scalar<<<grid3, LN_THREADS, 0, stream>>>(
                    xp + r0 * N, op + r0 * N, mp + r0, rp + r0, wp, bp, (int)N);
            }
        }
    }

    return out;
}
"""

layer_norm_cpp_source = (
    "torch::Tensor layernorm_forward(torch::Tensor x, torch::Tensor weight, "
    "torch::Tensor bias, double eps);"
)

# Compile the inline CUDA LayerNorm kernels
_ln_ext = load_inline(
    name="fused_layer_norm_ext_v3",
    cpp_sources=layer_norm_cpp_source,
    cuda_sources=layer_norm_cuda_source,
    functions=["layernorm_forward"],
    verbose=False,
    extra_cflags=["-O3"],
    extra_cuda_cflags=["-O3"],
)


class ModelNew(nn.Module):
    """
    LayerNorm model accelerated with custom fused CUDA kernels.
    """

    def __init__(self, normalized_shape: tuple):
        """
        Initializes the LayerNorm layer.

        Args:
            normalized_shape (tuple): Shape of the input tensor to be normalized.
        """
        super(ModelNew, self).__init__()
        self.ln = nn.LayerNorm(normalized_shape=normalized_shape)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """
        Applies Layer Normalization using the custom CUDA kernel.

        Args:
            x (torch.Tensor): Input tensor of shape (*, normalized_shape).

        Returns:
            torch.Tensor: Output tensor with Layer Normalization applied, same shape as input.
        """
        if x.is_cuda and x.dtype == torch.float32:
            return _ln_ext.layernorm_forward(
                x, self.ln.weight, self.ln.bias, float(self.ln.eps)
            )
        # Fallback for CPU tensors or non-float32 dtypes.
        return F.layer_norm(
            x, self.ln.normalized_shape, self.ln.weight, self.ln.bias, self.ln.eps
        )
```