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

# ---------------------------------------------------------------------------
# Enable TF32 for any fallback FP32 matmul path.
# ---------------------------------------------------------------------------
def _enable_tf32():
    try:
        torch.backends.cuda.matmul.allow_tf32 = True
    except Exception:
        pass
    try:
        torch.backends.cudnn.allow_tf32 = True
    except Exception:
        pass
    try:
        torch.set_float32_matmul_precision("high")
    except Exception:
        pass

_enable_tf32()

# ---------------------------------------------------------------------------
# The (1024x8192) @ (8192x8192) GEMM dominates the runtime. Two tensor-core
# strategies with FP32 accumulation are autotuned on-device per shape:
#   mode 1: CUBLAS_COMPUTE_32F_FAST_16F on the FP32 tensors directly. cuBLASLt
#           converts tiles to FP16 inside the GEMM mainloop, so no separate
#           cast kernel (and its ~50MB of DRAM traffic) is needed.
#   mode 0: explicit FP32->FP16 cast (custom vectorized kernel) + FP16-input /
#           FP32-output cuBLASLt GEMM.
# Both round inputs to 10 explicit mantissa bits (same as TF32), so accuracy
# matches the previously validated solution. Whichever full pipeline measures
# faster on the first call is cached and used afterwards. Both use a fused
# BIAS epilogue when available and a large workspace (split-K capable).
#
# InstanceNorm + residual add + multiply are fused into one kernel with the
# row cached in shared memory and streaming loads/stores.
# ---------------------------------------------------------------------------
fused_cuda_source = r"""
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include <c10/cuda/CUDAGuard.h>
#include <cuda_runtime.h>
#include <cuda_fp16.h>
#include <cublas_v2.h>
#include <cublasLt.h>
#include <unordered_map>
#include <cstdint>

namespace {

struct LtPlan {
    cublasLtMatmulAlgo_t algo;
    size_t workspace;
    int mode;           // 0 = fp16 inputs, 1 = fp32 inputs + FAST_16F compute
    bool valid;
    bool bias_epilogue;
};

struct PlanKey {
    int64_t m, n, k;
    int dev, bias;
    bool operator==(const PlanKey& o) const {
        return m == o.m && n == o.n && k == o.k && dev == o.dev && bias == o.bias;
    }
};

struct PlanKeyHash {
    size_t operator()(const PlanKey& p) const {
        size_t h = (size_t)(p.m * 1000003ll);
        h ^= (size_t)(p.n * 1000033ll) + 0x9e3779b9u + (h << 6) + (h >> 2);
        h ^= (size_t)(p.k * 1000081ll) + 0x9e3779b9u + (h << 6) + (h >> 2);
        h ^= (size_t)(p.dev * 97 + p.bias * 389) + 0x9e3779b9u + (h << 6) + (h >> 2);
        return h;
    }
};

std::unordered_map<PlanKey, LtPlan, PlanKeyHash>& plan_cache() {
    static auto* m = new std::unordered_map<PlanKey, LtPlan, PlanKeyHash>();
    return *m;
}

torch::Tensor& workspace_tensor() {
    static auto* t = new torch::Tensor();
    return *t;
}

cublasLtHandle_t lt_handle() {
    static cublasLtHandle_t h = nullptr;
    if (h == nullptr) {
        if (cublasLtCreate(&h) != CUBLAS_STATUS_SUCCESS) h = nullptr;
    }
    return h;
}

void* get_workspace(size_t bytes, const torch::TensorOptions& opt) {
    if (bytes == 0) return nullptr;
    auto& wt = workspace_tensor();
    if (!wt.defined() || (size_t)wt.numel() < bytes || wt.device() != opt.device()) {
        wt = torch::empty({(int64_t)bytes}, opt.dtype(torch::kUInt8));
    }
    return wt.data_ptr();
}

struct HalfWeightCache {
    torch::Tensor t;
    const void* ptr;
    int64_t version;
};

HalfWeightCache& half_weight_cache() {
    static auto* c = new HalfWeightCache{torch::Tensor(), nullptr, -1};
    return *c;
}

const torch::Tensor& ensure_half_weight(const torch::Tensor& w) {
    auto& hc = half_weight_cache();
    const int64_t ver = w.unsafeGetTensorImpl()->version_counter().current_version();
    if (!hc.t.defined() || hc.ptr != w.data_ptr() || hc.version != ver ||
        hc.t.device() != w.device()) {
        hc.t = w.detach().to(torch::kHalf).contiguous();
        hc.ptr = w.data_ptr();
        hc.version = ver;
    }
    return hc.t;
}

}  // namespace

// ----------------------------- fp32 -> fp16 cast -----------------------------
__global__ void cast_f32_f16_scalar_kernel(const float* __restrict__ in,
                                           __half* __restrict__ out, int64_t n) {
    int64_t i = (int64_t)blockIdx.x * blockDim.x + threadIdx.x;
    if (i < n) out[i] = __float2half_rn(in[i]);
}

__global__ void cast_f32_f16_vec_kernel(const float* __restrict__ in,
                                        __half* __restrict__ out,
                                        int64_t n4, int64_t n) {
    int64_t i = (int64_t)blockIdx.x * blockDim.x + threadIdx.x;
    if (i < n4) {
        const float4 v = reinterpret_cast<const float4*>(in)[i];
        const __half2 h0 = __floats2half2_rn(v.x, v.y);
        const __half2 h1 = __floats2half2_rn(v.z, v.w);
        uint2 u;
        u.x = *reinterpret_cast<const unsigned*>(&h0);
        u.y = *reinterpret_cast<const unsigned*>(&h1);
        reinterpret_cast<uint2*>(out)[i] = u;
    }
    if (i == 0) {  // scalar tail
        for (int64_t j = n4 * 4; j < n; ++j) out[j] = __float2half_rn(in[j]);
    }
}

static void launch_cast(const float* in, __half* out, int64_t n, cudaStream_t stream) {
    if (n <= 0) return;
    const int threads = 256;
    if ((((uintptr_t)in) | ((uintptr_t)out)) & 15) {
        int64_t blocks = (n + threads - 1) / threads;
        cast_f32_f16_scalar_kernel<<<(unsigned)blocks, threads, 0, stream>>>(in, out, n);
    } else {
        const int64_t n4 = n >> 2;
        int64_t blocks = (n4 + threads - 1) / threads;
        if (blocks <= 0) blocks = 1;
        cast_f32_f16_vec_kernel<<<(unsigned)blocks, threads, 0, stream>>>(in, out, n4, n);
    }
}

namespace {

struct Cand {
    cublasLtMatmulAlgo_t algo;
    size_t ws;
    float ms;
    bool ok;
};

cublasLtMatrixLayout_t make_layout(cudaDataType t, int64_t rows, int64_t cols, int64_t ld) {
    cublasLtMatrixLayout_t d = nullptr;
    if (cublasLtMatrixLayoutCreate(&d, t, rows, cols, ld) != CUBLAS_STATUS_SUCCESS)
        return nullptr;
    return d;
}

// Row-major trick: C(MxN) row-major == C^T(NxM) col-major = W(N,K) @ X^T(K,M),
// so A = w (K x N col-major, ld K, OP_T), B = x (K x M col-major, ld K, OP_N).
bool autotune(cublasLtHandle_t lt, cublasLtMatmulDesc_t opDesc,
              cublasLtMatrixLayout_t Adesc, cublasLtMatrixLayout_t Bdesc,
              cublasLtMatrixLayout_t Cdesc,
              const void* A, const void* B, void* C,
              const torch::TensorOptions& opt, cudaStream_t stream, Cand& best) {
    best.ok = false;
    cublasLtMatmulPreference_t pref = nullptr;
    if (cublasLtMatmulPreferenceCreate(&pref) != CUBLAS_STATUS_SUCCESS) return false;
    size_t maxWs = 256ull * 1024 * 1024;
    cublasLtMatmulPreferenceSetAttribute(pref, CUBLASLT_MATMUL_PREF_MAX_WORKSPACE_BYTES,
                                         &maxWs, sizeof(maxWs));
    cublasLtMatmulHeuristicResult_t hres[12];
    int found = 0;
    cublasStatus_t st = cublasLtMatmulAlgoGetHeuristic(lt, opDesc, Adesc, Bdesc, Cdesc,
                                                       Cdesc, pref, 12, hres, &found);
    cublasLtMatmulPreferenceDestroy(pref);
    if (st != CUBLAS_STATUS_SUCCESS || found <= 0) return false;

    const float alpha = 1.f, beta = 0.f;
    float bestMs = 1e30f;
    for (int i = 0; i < found; ++i) {
        if (hres[i].state != CUBLAS_STATUS_SUCCESS) continue;
        const size_t ws = hres[i].workspaceSize;
        if (ws > maxWs) continue;
        void* wsp = get_workspace(ws, opt);
        if (cublasLtMatmul(lt, opDesc, &alpha, A, Adesc, B, Bdesc, &beta, C, Cdesc, C,
                           Cdesc, &hres[i].algo, wsp, ws, stream) != CUBLAS_STATUS_SUCCESS)
            continue;
        cudaEvent_t e0, e1;
        cudaEventCreate(&e0);
        cudaEventCreate(&e1);
        cudaEventRecord(e0, stream);
        const int reps = 3;
        for (int r = 0; r < reps; ++r) {
            cublasLtMatmul(lt, opDesc, &alpha, A, Adesc, B, Bdesc, &beta, C, Cdesc, C,
                           Cdesc, &hres[i].algo, wsp, ws, stream);
        }
        cudaEventRecord(e1, stream);
        cudaEventSynchronize(e1);
        float ms = 0.f;
        cudaEventElapsedTime(&ms, e0, e1);
        cudaEventDestroy(e0);
        cudaEventDestroy(e1);
        if (ms < bestMs) {
            bestMs = ms;
            best.algo = hres[i].algo;
            best.ws = ws;
            best.ms = ms;
            best.ok = true;
        }
    }
    return best.ok;
}

// Tries bias-epilogue first (if bias given), then plain. Returns true if any
// working algo was found; epi_used reports whether the epilogue applies bias.
bool tune_mode(cublasLtHandle_t lt, int mode, const void* A, const void* B,
               const float* bias, float* C, int64_t M, int64_t N, int64_t K,
               const torch::TensorOptions& opt, cudaStream_t stream,
               Cand& best, bool& epi_used) {
    for (int attempt = 0; attempt < 2; ++attempt) {
        const bool epi = (attempt == 0) && (bias != nullptr);
        const cublasComputeType_t ct =
            (mode == 1) ? CUBLAS_COMPUTE_32F_FAST_16F : CUBLAS_COMPUTE_32F;
        cublasLtMatmulDesc_t opDesc = nullptr;
        if (cublasLtMatmulDescCreate(&opDesc, ct, CUDA_R_32F) != CUBLAS_STATUS_SUCCESS)
            return false;
        cublasOperation_t opT = CUBLAS_OP_T, opN = CUBLAS_OP_N;
        cublasLtMatmulDescSetAttribute(opDesc, CUBLASLT_MATMUL_DESC_TRANSA, &opT, sizeof(opT));
        cublasLtMatmulDescSetAttribute(opDesc, CUBLASLT_MATMUL_DESC_TRANSB, &opN, sizeof(opN));
        if (epi) {
            cublasLtEpilogue_t e = CUBLASLT_EPILOGUE_BIAS;
            cublasLtMatmulDescSetAttribute(opDesc, CUBLASLT_MATMUL_DESC_EPILOGUE, &e, sizeof(e));
            cublasLtMatmulDescSetAttribute(opDesc, CUBLASLT_MATMUL_DESC_BIAS_POINTER,
                                           &bias, sizeof(bias));
        }
        const cudaDataType ab = (mode == 1) ? CUDA_R_32F : CUDA_R_16F;
        cublasLtMatrixLayout_t Adesc = make_layout(ab, K, N, K);
        cublasLtMatrixLayout_t Bdesc = make_layout(ab, K, M, K);
        cublasLtMatrixLayout_t Cdesc = make_layout(CUDA_R_32F, N, M, N);
        bool ok = false;
        if (Adesc && Bdesc && Cdesc) {
            ok = autotune(lt, opDesc, Adesc, Bdesc, Cdesc, A, B, C, opt, stream, best);
        }
        if (Adesc) cublasLtMatrixLayoutDestroy(Adesc);
        if (Bdesc) cublasLtMatrixLayoutDestroy(Bdesc);
        if (Cdesc) cublasLtMatrixLayoutDestroy(Cdesc);
        cublasLtMatmulDescDestroy(opDesc);
        if (ok) {
            epi_used = epi;
            return true;
        }
        if (bias == nullptr) break;
    }
    return false;
}

bool run_lt(cublasLtHandle_t lt, const LtPlan& plan, const void* A, const void* B,
            const float* bias, float* C, int64_t M, int64_t N, int64_t K,
            const torch::TensorOptions& opt, cudaStream_t stream) {
    const cublasComputeType_t ct =
        (plan.mode == 1) ? CUBLAS_COMPUTE_32F_FAST_16F : CUBLAS_COMPUTE_32F;
    cublasLtMatmulDesc_t opDesc = nullptr;
    if (cublasLtMatmulDescCreate(&opDesc, ct, CUDA_R_32F) != CUBLAS_STATUS_SUCCESS)
        return false;
    cublasOperation_t opT = CUBLAS_OP_T, opN = CUBLAS_OP_N;
    cublasLtMatmulDescSetAttribute(opDesc, CUBLASLT_MATMUL_DESC_TRANSA, &opT, sizeof(opT));
    cublasLtMatmulDescSetAttribute(opDesc, CUBLASLT_MATMUL_DESC_TRANSB, &opN, sizeof(opN));
    if (plan.bias_epilogue && bias != nullptr) {
        cublasLtEpilogue_t e = CUBLASLT_EPILOGUE_BIAS;
        cublasLtMatmulDescSetAttribute(opDesc, CUBLASLT_MATMUL_DESC_EPILOGUE, &e, sizeof(e));
        cublasLtMatmulDescSetAttribute(opDesc, CUBLASLT_MATMUL_DESC_BIAS_POINTER,
                                       &bias, sizeof(bias));
    }
    const cudaDataType ab = (plan.mode == 1) ? CUDA_R_32F : CUDA_R_16F;
    cublasLtMatrixLayout_t Adesc = make_layout(ab, K, N, K);
    cublasLtMatrixLayout_t Bdesc = make_layout(ab, K, M, K);
    cublasLtMatrixLayout_t Cdesc = make_layout(CUDA_R_32F, N, M, N);
    bool ok = false;
    if (Adesc && Bdesc && Cdesc) {
        const float alpha = 1.f, beta = 0.f;
        void* wsp = get_workspace(plan.workspace, opt);
        ok = (cublasLtMatmul(lt, opDesc, &alpha, A, Adesc, B, Bdesc, &beta, C, Cdesc, C,
                             Cdesc, &plan.algo, wsp, plan.workspace, stream)
              == CUBLAS_STATUS_SUCCESS);
    }
    if (Adesc) cublasLtMatrixLayoutDestroy(Adesc);
    if (Bdesc) cublasLtMatrixLayoutDestroy(Bdesc);
    if (Cdesc) cublasLtMatrixLayoutDestroy(Cdesc);
    cublasLtMatmulDescDestroy(opDesc);
    return ok;
}

LtPlan create_plan(const float* xp, const float* wp, const __half* xh, const __half* wh,
                   const float* bias, float* zp, int64_t M, int64_t N, int64_t K,
                   const torch::TensorOptions& opt, cudaStream_t stream) {
    LtPlan plan{};
    plan.workspace = 0;
    plan.mode = 1;
    plan.valid = false;
    plan.bias_epilogue = false;
    cublasLtHandle_t lt = lt_handle();
    if (lt == nullptr) return plan;

    Cand bestFast, best16;
    bool epiFast = false, epi16 = false;
    const bool okFast =
        tune_mode(lt, 1, wp, xp, bias, zp, M, N, K, opt, stream, bestFast, epiFast);
    bool ok16 = false;
    if (xh != nullptr && wh != nullptr) {
        ok16 = tune_mode(lt, 0, wh, xh, bias, zp, M, N, K, opt, stream, best16, epi16);
    }
    float castMs = 0.f;
    if (ok16) {
        __half* xhw = const_cast<__half*>(xh);
        launch_cast(xp, xhw, M * K, stream);  // warm-up
        cudaEvent_t e0, e1;
        cudaEventCreate(&e0);
        cudaEventCreate(&e1);
        cudaEventRecord(e0, stream);
        const int reps = 3;
        for (int r = 0; r < reps; ++r) launch_cast(xp, xhw, M * K, stream);
        cudaEventRecord(e1, stream);
        cudaEventSynchronize(e1);
        cudaEventElapsedTime(&castMs, e0, e1);
        cudaEventDestroy(e0);
        cudaEventDestroy(e1);
    }
    if (okFast && (!ok16 || bestFast.ms <= best16.ms + castMs)) {
        plan.mode = 1;
        plan.algo = bestFast.algo;
        plan.workspace = bestFast.ws;
        plan.bias_epilogue = epiFast;
        plan.valid = true;
    } else if (ok16) {
        plan.mode = 0;
        plan.algo = best16.algo;
        plan.workspace = best16.ws;
        plan.bias_epilogue = epi16;
        plan.valid = true;
    }
    return plan;
}

}  // namespace

// ------------------------- fallback: cublasGemmEx FP16 -------------------------
static void gemm_ex_fp16(const void* x, const void* w, float* out,
                         int64_t M, int64_t N, int64_t K) {
    cublasHandle_t handle = at::cuda::getCurrentCUDABlasHandle();
    float alpha = 1.0f, beta = 0.0f;
    cublasStatus_t st = cublasGemmEx(
        handle, CUBLAS_OP_T, CUBLAS_OP_N,
        (int)N, (int)M, (int)K,
        &alpha,
        w, CUDA_R_16F, (int)K,
        x, CUDA_R_16F, (int)K,
        &beta,
        out, CUDA_R_32F, (int)N,
        CUBLAS_COMPUTE_32F, CUBLAS_GEMM_DEFAULT);
    TORCH_CHECK(st == CUBLAS_STATUS_SUCCESS,
                "cublasGemmEx failed with status ", (int)st);
}

// ---------------------------------------------------------------------------
// Fused: out = ((x [+ bias] - mean) * rstd + y) * y   (row-wise biased stats)
// Row cached in shared memory; x may be aliased with out (in-place safe).
// ---------------------------------------------------------------------------
template <bool VEC, bool HAS_BIAS>
__global__ void fused_bias_inorm_add_mul_smem_kernel(
    const float* __restrict__ x,
    const float* __restrict__ bias,
    const float* __restrict__ y,
    float* __restrict__ out,
    const int Fdim,
    const float eps)
{
    extern __shared__ float s_row[];
    __shared__ float s_sum[32];
    __shared__ float s_sq[32];

    const int row = blockIdx.x;
    const size_t base = (size_t)row * (size_t)Fdim;
    const float* xr = x + base;
    const float* yr = y + base;
    float* orow = out + base;

    float sum = 0.f;
    float sumsq = 0.f;

    if (VEC) {
        const int F4 = Fdim >> 2;
        const float4* xr4 = reinterpret_cast<const float4*>(xr);
        const float4* b4  = reinterpret_cast<const float4*>(bias);
        float4* s4 = reinterpret_cast<float4*>(s_row);
        #pragma unroll 4
        for (int i = threadIdx.x; i < F4; i += blockDim.x) {
            float4 v = __ldcs(xr4 + i);
            if (HAS_BIAS) {
                const float4 bv = __ldg(b4 + i);
                v.x += bv.x; v.y += bv.y; v.z += bv.z; v.w += bv.w;
            }
            s4[i] = v;
            sum   += v.x + v.y + v.z + v.w;
            sumsq += v.x * v.x + v.y * v.y + v.z * v.z + v.w * v.w;
        }
    } else {
        for (int i = threadIdx.x; i < Fdim; i += blockDim.x) {
            float v = __ldcs(xr + i) + (HAS_BIAS ? __ldg(bias + i) : 0.f);
            s_row[i] = v;
            sum   += v;
            sumsq += v * v;
        }
    }

    const unsigned int mask = 0xffffffffu;
    for (int off = 16; off > 0; off >>= 1) {
        sum   += __shfl_down_sync(mask, sum, off);
        sumsq += __shfl_down_sync(mask, sumsq, off);
    }
    const int lane = threadIdx.x & 31;
    const int wid  = threadIdx.x >> 5;
    if (lane == 0) {
        s_sum[wid] = sum;
        s_sq[wid]  = sumsq;
    }
    __syncthreads();

    const int nwarps = (blockDim.x + 31) >> 5;
    if (wid == 0) {
        float a = (lane < nwarps) ? s_sum[lane] : 0.f;
        float b = (lane < nwarps) ? s_sq[lane]  : 0.f;
        for (int off = 16; off > 0; off >>= 1) {
            a += __shfl_down_sync(mask, a, off);
            b += __shfl_down_sync(mask, b, off);
        }
        if (lane == 0) {
            s_sum[0] = a;
            s_sq[0]  = b;
        }
    }
    __syncthreads();

    const float mean = s_sum[0] / (float)Fdim;
    const float var  = s_sq[0] / (float)Fdim - mean * mean;  // biased
    const float inv  = rsqrtf(var + eps);

    if (VEC) {
        const int F4 = Fdim >> 2;
        const float4* s4  = reinterpret_cast<const float4*>(s_row);
        const float4* yr4 = reinterpret_cast<const float4*>(yr);
        float4* o4 = reinterpret_cast<float4*>(orow);
        #pragma unroll 4
        for (int i = threadIdx.x; i < F4; i += blockDim.x) {
            const float4 xv = s4[i];
            const float4 yv = __ldcs(yr4 + i);
            float4 ov;
            ov.x = ((xv.x - mean) * inv + yv.x) * yv.x;
            ov.y = ((xv.y - mean) * inv + yv.y) * yv.y;
            ov.z = ((xv.z - mean) * inv + yv.z) * yv.z;
            ov.w = ((xv.w - mean) * inv + yv.w) * yv.w;
            __stcs(o4 + i, ov);
        }
    } else {
        for (int i = threadIdx.x; i < Fdim; i += blockDim.x) {
            const float yv = __ldcs(yr + i);
            __stcs(orow + i, ((s_row[i] - mean) * inv + yv) * yv);
        }
    }
}

// Fallback for rows too large for shared memory: two passes over global.
template <bool VEC, bool HAS_BIAS>
__global__ void fused_bias_inorm_add_mul_kernel(
    const float* __restrict__ x,
    const float* __restrict__ bias,
    const float* __restrict__ y,
    float* __restrict__ out,
    const int Fdim,
    const float eps)
{
    const int row = blockIdx.x;
    const size_t base = (size_t)row * (size_t)Fdim;
    const float* xr = x + base;
    const float* yr = y + base;
    float* orow = out + base;

    float sum = 0.f;
    float sumsq = 0.f;

    if (VEC) {
        const int F4 = Fdim >> 2;
        const float4* xr4 = reinterpret_cast<const float4*>(xr);
        const float4* b4  = reinterpret_cast<const float4*>(bias);
        for (int i = threadIdx.x; i < F4; i += blockDim.x) {
            float4 v = xr4[i];
            if (HAS_BIAS) {
                const float4 bv = __ldg(b4 + i);
                v.x += bv.x; v.y += bv.y; v.z += bv.z; v.w += bv.w;
            }
            sum   += v.x + v.y + v.z + v.w;
            sumsq += v.x * v.x + v.y * v.y + v.z * v.z + v.w * v.w;
        }
    } else {
        for (int i = threadIdx.x; i < Fdim; i += blockDim.x) {
            float v = xr[i] + (HAS_BIAS ? __ldg(bias + i) : 0.f);
            sum   += v;
            sumsq += v * v;
        }
    }

    __shared__ float s_sum[32];
    __shared__ float s_sq[32];
    const unsigned int mask = 0xffffffffu;
    for (int off = 16; off > 0; off >>= 1) {
        sum   += __shfl_down_sync(mask, sum, off);
        sumsq += __shfl_down_sync(mask, sumsq, off);
    }
    const int lane = threadIdx.x & 31;
    const int wid  = threadIdx.x >> 5;
    if (lane == 0) {
        s_sum[wid] = sum;
        s_sq[wid]  = sumsq;
    }
    __syncthreads();

    const int nwarps = (blockDim.x + 31) >> 5;
    if (wid == 0) {
        float a = (lane < nwarps) ? s_sum[lane] : 0.f;
        float b = (lane < nwarps) ? s_sq[lane]  : 0.f;
        for (int off = 16; off > 0; off >>= 1) {
            a += __shfl_down_sync(mask, a, off);
            b += __shfl_down_sync(mask, b, off);
        }
        if (lane == 0) {
            s_sum[0] = a;
            s_sq[0]  = b;
        }
    }
    __syncthreads();

    const float mean = s_sum[0] / (float)Fdim;
    const float var  = s_sq[0] / (float)Fdim - mean * mean;
    const float inv  = rsqrtf(var + eps);

    if (VEC) {
        const int F4 = Fdim >> 2;
        const float4* xr4 = reinterpret_cast<const float4*>(xr);
        const float4* b4  = reinterpret_cast<const float4*>(bias);
        const float4* yr4 = reinterpret_cast<const float4*>(yr);
        float4* o4 = reinterpret_cast<float4*>(orow);
        for (int i = threadIdx.x; i < F4; i += blockDim.x) {
            float4 xv = xr4[i];
            if (HAS_BIAS) {
                const float4 bv = __ldg(b4 + i);
                xv.x += bv.x; xv.y += bv.y; xv.z += bv.z; xv.w += bv.w;
            }
            const float4 yv = yr4[i];
            float4 ov;
            ov.x = ((xv.x - mean) * inv + yv.x) * yv.x;
            ov.y = ((xv.y - mean) * inv + yv.y) * yv.y;
            ov.z = ((xv.z - mean) * inv + yv.z) * yv.z;
            ov.w = ((xv.w - mean) * inv + yv.w) * yv.w;
            o4[i] = ov;
        }
    } else {
        for (int i = threadIdx.x; i < Fdim; i += blockDim.x) {
            float xv = xr[i] + (HAS_BIAS ? __ldg(bias + i) : 0.f);
            const float yv = yr[i];
            orow[i] = ((xv - mean) * inv + yv) * yv;
        }
    }
}

torch::Tensor forward_fused(torch::Tensor x, torch::Tensor w,
                            torch::Tensor bias, torch::Tensor y, double eps) {
    TORCH_CHECK(x.is_cuda() && w.is_cuda() && y.is_cuda(), "inputs must be CUDA");
    TORCH_CHECK(x.dim() == 2 && w.dim() == 2 && y.dim() == 2, "bad dims");
    TORCH_CHECK(x.scalar_type() == torch::kFloat32 &&
                w.scalar_type() == torch::kFloat32 &&
                y.scalar_type() == torch::kFloat32, "expects float32 tensors");
    const c10::cuda::OptionalCUDAGuard guard(x.device());

    auto xc = x.contiguous();
    auto wc = w.contiguous();
    auto yc = y.contiguous();

    const int64_t M = xc.size(0);
    const int64_t K = xc.size(1);
    const int64_t N = wc.size(0);
    TORCH_CHECK(wc.size(1) == K, "weight shape mismatch");

    auto z = torch::empty({M, N}, xc.options().dtype(torch::kFloat32));
    if (M == 0 || N == 0) return z;
    TORCH_CHECK(yc.numel() == M * N, "y shape mismatch");

    const bool has_bias = bias.defined() &&
                          bias.scalar_type() == torch::kFloat32 &&
                          bias.numel() == N;
    torch::Tensor bc;
    const float* bptr = nullptr;
    if (has_bias) {
        bc = bias.contiguous();
        bptr = bc.data_ptr<float>();
    }

    cudaStream_t stream = at::cuda::getCurrentCUDAStream();
    float* zp = z.data_ptr<float>();
    const float* xp = xc.data_ptr<float>();
    const float* wp = wc.data_ptr<float>();

    bool ran = false;
    bool bias_done = false;

    if (K > 0) {
        const int dev = (int)xc.get_device();
        PlanKey key{M, N, K, dev, has_bias ? 1 : 0};
        auto& cache = plan_cache();
        auto it = cache.find(key);
        if (it == cache.end()) {
            const torch::Tensor& wh = ensure_half_weight(wc);
            torch::Tensor xh = torch::empty({M, K}, xc.options().dtype(torch::kHalf));
            launch_cast(xp, reinterpret_cast<__half*>(xh.data_ptr()), M * K, stream);
            LtPlan plan = create_plan(
                xp, wp,
                reinterpret_cast<const __half*>(xh.data_ptr()),
                reinterpret_cast<const __half*>(wh.data_ptr()),
                bptr, zp, M, N, K, xc.options(), stream);
            it = cache.emplace(key, plan).first;
        }
        const LtPlan& plan = it->second;
        if (plan.valid) {
            if (plan.mode == 1) {
                ran = run_lt(lt_handle(), plan, wp, xp, bptr, zp, M, N, K,
                             xc.options(), stream);
            } else {
                const torch::Tensor& wh = ensure_half_weight(wc);
                torch::Tensor xh = torch::empty({M, K}, xc.options().dtype(torch::kHalf));
                launch_cast(xp, reinterpret_cast<__half*>(xh.data_ptr()), M * K, stream);
                ran = run_lt(lt_handle(), plan, wh.data_ptr(), xh.data_ptr(), bptr, zp,
                             M, N, K, xc.options(), stream);
            }
            bias_done = ran && plan.bias_epilogue;
        }
        if (!ran) {
            const torch::Tensor& wh = ensure_half_weight(wc);
            torch::Tensor xh = torch::empty({M, K}, xc.options().dtype(torch::kHalf));
            launch_cast(xp, reinterpret_cast<__half*>(xh.data_ptr()), M * K, stream);
            gemm_ex_fp16(xh.data_ptr(), wh.data_ptr(), zp, M, N, K);
            bias_done = false;
        }
    } else {
        z.zero_();
    }

    // If the Lt epilogue applied the bias, the fused kernel must skip it.
    const float* eff_bias = bias_done ? nullptr : bptr;
    const bool kernel_has_bias = (eff_bias != nullptr);

    const int B = (int)z.size(0);
    const int Fdim = (int)z.size(1);

    auto out = z;  // in-place: fused kernel stages rows in shared memory first
    const bool vec = (Fdim & 3) == 0;
    const size_t smem = (size_t)Fdim * sizeof(float);

    int threads = 256;
    if (Fdim >= 4096) threads = 512;
    else if (Fdim < 512) threads = 128;

    const float* yp = yc.data_ptr<float>();
    float* op = out.data_ptr<float>();
    const float feps = (float)eps;

    if (smem <= 48 * 1024) {
        if (vec) {
            if (kernel_has_bias)
                fused_bias_inorm_add_mul