```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

# Fused LayerNorm with two implementations:
#
#  A) Persistent cooperative kernel (fast path for this workload).
#     The workload is a few very long rows (e.g. 16 rows x 4M floats), so the
#     whole op is DRAM-bandwidth bound: stats pass reads x (256MB), normalize
#     pass reads x again (256MB) and writes out (256MB).  This kernel keeps a
#     co-resident grid and processes rows one at a time:
#       phase 1: all blocks reduce their chunk of row r  -> partial sums
#       grid barrier (one per row, double-buffered partials so adjacent
#                 row phases may overlap)
#       every block redundantly computes mean/rstd from the tiny partials
#       phase 2: blocks normalize *the same chunk* of row r -> x[row r] is
#                served from L2 (it was just read in phase 1), and weight/bias
#                are fetched with an L2 evict-last hint so they stay resident
#                across rows (only 32MB).  Output stores are streaming.
#     This can cut DRAM traffic from ~768MB toward ~544MB.
#
#  B) Proven 3-kernel path (partial sums -> stats -> column-owning normalize)
#     kept as a general fallback for any shape/alignment/dtype and used
#     whenever the persistent path is unavailable or slower.
#
#  A one-time autotune per problem shape picks the faster path, so the
#  module can never regress below the turn-3 baseline.

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

// Block reduce; result valid in every thread (broadcast via shared memory).
__device__ __forceinline__ void ln_block_reduce_bcast(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) {
        sum   = (lane < LN_WARPS) ? sh_sum[lane]   : 0.0f;
        sumsq = (lane < LN_WARPS) ? sh_sumsq[lane] : 0.0f;
        ln_warp_reduce(sum, sumsq);
        if (lane == 0) {
            sh_sum[0]   = sum;
            sh_sumsq[0] = sumsq;
        }
    }
    __syncthreads();
    sum   = sh_sum[0];
    sumsq = sh_sumsq[0];
    __syncthreads();
}

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

// ---------------------------------------------------------------------------
// Persistent fused LayerNorm kernel (vectorized, float4, 16B-aligned rows).
// ---------------------------------------------------------------------------
__device__ __forceinline__ float4 ln_ld_evict_last(const float4* p) {
#if __CUDA_ARCH__ >= 800
    float4 v;
    asm volatile("ld.global.nc.L2::evict_last.v4.f32 {%0,%1,%2,%3}, [%4];"
                 : "=f"(v.x), "=f"(v.y), "=f"(v.z), "=f"(v.w)
                 : "l"(p));
    return v;
#else
    return __ldg(p);
#endif
}

__device__ __forceinline__ void ln_grid_barrier(unsigned int* bar) {
    __syncthreads();
    if (threadIdx.x == 0) {
        volatile unsigned int* genp =
            reinterpret_cast<volatile unsigned int*>(bar + 1);
        const unsigned int g = *genp;
        __threadfence();
        const unsigned int arrived = atomicAdd(bar, 1u) + 1u;
        if (arrived == gridDim.x) {
            *bar = 0u;
            __threadfence();
            atomicAdd(bar + 1, 1u);
        } else {
            while (*genp == g) { }
        }
        __threadfence();
    }
    __syncthreads();
}

// partials layout: [2][gridDim.x][2] floats  (sum, sumsq) per block per buffer.
// bar: 2 unsigned ints (count, generation), zero-initialized before launch.
__global__ void ln_persistent_vec4(const float4* __restrict__ x,
                                   float4* __restrict__ out,
                                   const float4* __restrict__ w,
                                   const float4* __restrict__ b,
                                   float* __restrict__ partials,
                                   unsigned int* __restrict__ bar,
                                   int n4, int rows,
                                   float invN, float eps) {
    const int nblocks = gridDim.x;
    const int bid     = blockIdx.x;

    // Contiguous chunk [start, start+len) of the row owned by this block.
    const int base  = n4 / nblocks;
    const int rem   = n4 - base * nblocks;
    const int start = bid * base + (bid < rem ? bid : rem);
    const int len   = base + (bid < rem ? 1 : 0);
    const int end   = start + len;

    __shared__ float sh_sum[LN_WARPS];
    __shared__ float sh_sumsq[LN_WARPS];

    for (int r = 0; r < rows; ++r) {
        float* part = partials + (size_t)(r & 1) * nblocks * 2;
        const float4* xr = x + (long long)r * n4;
        float4* outr = out + (long long)r * n4;

        // ---------------- phase 1: partial sums over this chunk ----------
        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;
        {
            const int tstride = blockDim.x;
            int i = start + threadIdx.x;
            for (; i + 3 * tstride < end; i += 4 * tstride) {
                const float4 v0 = xr[i];
                const float4 v1 = xr[i + tstride];
                const float4 v2 = xr[i + 2 * tstride];
                const float4 v3 = xr[i + 3 * tstride];
                s0 += v0.x + v0.y; s0 += v0.z + v0.w;
                q0 += v0.x * v0.x + v0.y * v0.y; q0 += v0.z * v0.z + v0.w * v0.w;
                s1 += v1.x + v1.y; s1 += v1.z + v1.w;
                q1 += v1.x * v1.x + v1.y * v1.y; q1 += v1.z * v1.z + v1.w * v1.w;
                s2 += v2.x + v2.y; s2 += v2.z + v2.w;
                q2 += v2.x * v2.x + v2.y * v2.y; q2 += v2.z * v2.z + v2.w * v2.w;
                s3 += v3.x + v3.y; s3 += v3.z + v3.w;
                q3 += v3.x * v3.x + v3.y * v3.y; q3 += v3.z * v3.z + v3.w * v3.w;
            }
            for (; i < end; i += tstride) {
                const float4 v0 = xr[i];
                s0 += v0.x + v0.y; s0 += v0.z + v0.w;
                q0 += v0.x * v0.x + v0.y * v0.y; q0 += v0.z * v0.z + v0.w * v0.w;
            }
        }
        float sum   = (s0 + s1) + (s2 + s3);
        float sumsq = (q0 + q1) + (q2 + q3);
        ln_block_reduce(sum, sumsq, sh_sum, sh_sumsq);
        if (threadIdx.x == 0) {
            part[2 * bid]     = sum;
            part[2 * bid + 1] = sumsq;
        }

        ln_grid_barrier(bar);

        // ------------- stats: every block redundantly combines partials ---
        float m, t;
        {
            float rs = 0.f, rq = 0.f;
            for (int i = threadIdx.x; i < nblocks; i += blockDim.x) {
                rs += part[2 * i];
                rq += part[2 * i + 1];
            }
            ln_block_reduce_bcast(rs, rq, sh_sum, sh_sumsq);
            m = rs * invN;
            float var = rq * invN - m * m;
            var = var < 0.f ? 0.f : var;
            t = rsqrtf(var + eps);
        }

        // ------------- phase 2: normalize this chunk of row r -----------
        {
            const int tstride = blockDim.x;
            int i = start + threadIdx.x;
            for (; i + 3 * tstride < end; i += 4 * tstride) {
                const float4 w0 = ln_ld_evict_last(w + i);
                const float4 b0 = ln_ld_evict_last(b + i);
                const float4 w1 = ln_ld_evict_last(w + i + tstride);
                const float4 b1 = ln_ld_evict_last(b + i + tstride);
                const float4 w2 = ln_ld_evict_last(w + i + 2 * tstride);
                const float4 b2 = ln_ld_evict_last(b + i + 2 * tstride);
                const float4 w3 = ln_ld_evict_last(w + i + 3 * tstride);
                const float4 b3 = ln_ld_evict_last(b + i + 3 * tstride);
                const float4 v0 = xr[i];
                const float4 v1 = xr[i + tstride];
                const float4 v2 = xr[i + 2 * tstride];
                const float4 v3 = xr[i + 3 * tstride];
                float4 u;
                u.x = (v0.x - m) * t * w0.x + b0.x;
                u.y = (v0.y - m) * t * w0.y + b0.y;
                u.z = (v0.z - m) * t * w0.z + b0.z;
                u.w = (v0.w - m) * t * w0.w + b0.w;
                __stcs(outr + i, u);
                u.x = (v1.x - m) * t * w1.x + b1.x;
                u.y = (v1.y - m) * t * w1.y + b1.y;
                u.z = (v1.z - m) * t * w1.z + b1.z;
                u.w = (v1.w - m) * t * w1.w + b1.w;
                __stcs(outr + i + tstride, u);
                u.x = (v2.x - m) * t * w2.x + b2.x;
                u.y = (v2.y - m) * t * w2.y + b2.y;
                u.z = (v2.z - m) * t * w2.z + b2.z;
                u.w = (v2.w - m) * t * w2.w + b2.w;
                __stcs(outr + i + 2 * tstride, u);
                u.x = (v3.x - m) * t * w3.x + b3.x;
                u.y = (v3.y - m) * t * w3.y + b3.y;
                u.z = (v3.z - m) * t * w3.z + b3.z;
                u.w = (v3.w - m) * t * w3.w + b3.w;
                __stcs(outr + i + 3 * tstride, u);
            }
            for (; i < end; i += tstride) {
                const float4 w0 = ln_ld_evict_last(w + i);
                const float4 b0 = ln_ld_evict_last(b + i);
                const float4 v0 = xr[i];
                float4 u;
                u.x = (v0.x - m) * t * w0.x + b0.x;
                u.y = (v0.y - m) * t * w0.y + b0.y;
                u.z = (v0.z - m) * t * w0.z + b0.z;
                u.w = (v0.w - m) * t * w0.w + b0.w;
                __stcs(outr + i, u);
            }
        }
        // No barrier here: partials are double buffered, so phase 1 of the
        // next row (other buffer) may safely overlap the tail of this phase.
    }
}

// Max co-resident blocks for the persistent kernel (0 if unsupported).
int64_t ln_persistent_grid_size() {
    static int cached = -1;
    if (cached >= 0) return (int64_t)cached;
    int dev = 0;
    cudaGetDevice(&dev);
    int coop = 0;
    cudaDeviceGetAttribute(&coop, cudaDevAttrCooperativeLaunch, dev);
    if (!coop) { cached = 0; return 0; }
    int nsm = 0;
    cudaDeviceGetAttribute(&nsm, cudaDevAttrMultiProcessorCount, dev);
    int per_sm = 0;
    cudaError_t err = cudaOccupancyMaxActiveBlocksPerMultiprocessor(
        &per_sm, (const void*)ln_persistent_vec4, LN_THREADS, 0);
    if (err != cudaSuccess || per_sm <= 0 || nsm <= 0) {
        cudaGetLastError();
        cached = 0;
        return 0;
    }
    cached = nsm * per_sm;
    return (int64_t)cached;
}

torch::Tensor layernorm_forward_persistent(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 & 3) == 0, "Persistent path requires N % 4 == 0");
    TORCH_CHECK(N < (int64_t)1 << 31, "Normalized dim too large");
    TORCH_CHECK(rows < (int64_t)1 << 31, "Too many rows");

    const int64_t n4 = N >> 2;
    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>();
    TORCH_CHECK(((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),
                "Persistent path requires 16-byte aligned tensors");

    int64_t max_blocks = ln_persistent_grid_size();
    TORCH_CHECK(max_blocks > 0, "Cooperative launch not supported on this device");
    const int grid = (int)std::max<int64_t>(
        1, std::min<int64_t>(max_blocks, std::max<int64_t>(n4 / 16, 1)));

    // Workspace: partials [2][grid][2] floats + 2 unsigned ints barrier state.
    auto ws = at::zeros({(int64_t)grid * 4 + 2},
                        xc.options().dtype(at::kFloat));
    float* partials = ws.data_ptr<float>();
    unsigned int* bar = reinterpret_cast<unsigned int*>(partials + (int64_t)grid * 4);

    auto stream = at::cuda::getCurrentCUDAStream();
    const float invN = 1.0f / (float)N;
    const float feps = (float)eps;
    int n4_i = (int)n4;
    int rows_i = (int)rows;

    const float4* x4 = reinterpret_cast<const float4*>(xp);
    float4* o4 = reinterpret_cast<float4*>(op);
    const float4* w4 = reinterpret_cast<const float4*>(wp);
    const float4* b4 = reinterpret_cast<const float4*>(bp);

    void* args[] = {(void*)&x4, (void*)&o4, (void*)&w4, (void*)&b4,
                    (void*)&partials, (void*)&bar, (void*)&n4_i, (void*)&rows_i,
                    (void*)&invN, (void*)&feps};
    cudaError_t err = cudaLaunchCooperativeKernel(
        (const void*)ln_persistent_vec4, dim3(grid), dim3(LN_THREADS),
        args, 0, stream.stream());
    TORCH_CHECK(err == cudaSuccess,
                "Cooperative launch failed: ", cudaGetErrorString(err));
    return out;
}

// ---------------------------------------------------------------------------
// Fallback 3-kernel implementation (turn-3 baseline, handles any shape).
// ---------------------------------------------------------------------------
__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;
    }
}

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

__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;
        mean[row] = m;
        rstd[row] = rsqrtf(var + eps);
    }
}

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

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