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

# Fallback matmuls should at least use TF32 tensor cores.
torch.backends.cuda.matmul.allow_tf32 = True

# ---------------------------------------------------------------------------
# One merged extension:
#   * vectorized FP32 -> FP16 cast (8 floats per thread)
#   * GEMM with one-time exhaustive autotuning across:
#       - cublasLt FP16 (FP32 accumulate): heuristics + full algo-id sweep
#       - cublasGemmEx FP16 (FP32 accumulate)
#       - cublasLt / cublasGemmEx FP32 with CUBLAS_COMPUTE_32F_FAST_16F
#         (native fp16 tensor cores on fp32 data, no explicit cast needed)
#     The cast-kernel cost is included when comparing fp16-input algorithms.
#   * fused bias + GroupNorm + Swish + per-channel multiply + Swish epilogue
# ---------------------------------------------------------------------------
fused_cuda_source = r"""
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include <cublas_v2.h>
#include <cublasLt.h>
#include <cuda_runtime.h>
#include <cuda_fp16.h>
#include <functional>
#include <algorithm>
#include <vector>
#include <cstdint>

#define FULL_MASK 0xffffffffu

// ------------------------- vectorized FP32 -> FP16 -------------------------
__global__ void cast_f32_f16_vec8_kernel(const float4* __restrict__ in,
                                         float4* __restrict__ out,
                                         long n8) {
    long i = (long)blockIdx.x * blockDim.x + threadIdx.x;
    if (i >= n8) return;
    float4 a = in[2 * i];
    float4 b = in[2 * i + 1];
    __half2 h0 = __floats2half2_rn(a.x, a.y);
    __half2 h1 = __floats2half2_rn(a.z, a.w);
    __half2 h2 = __floats2half2_rn(b.x, b.y);
    __half2 h3 = __floats2half2_rn(b.z, b.w);
    float4 o;
    o.x = *reinterpret_cast<float*>(&h0);
    o.y = *reinterpret_cast<float*>(&h1);
    o.z = *reinterpret_cast<float*>(&h2);
    o.w = *reinterpret_cast<float*>(&h3);
    out[i] = o;
}

__global__ void cast_f32_f16_scalar_kernel(const float* __restrict__ in,
                                           __half* __restrict__ out,
                                           long start, long numel) {
    long i = start + (long)blockIdx.x * blockDim.x + threadIdx.x;
    if (i < numel) out[i] = __float2half_rn(in[i]);
}

static void cast_to_half(const float* in, void* out, long numel, cudaStream_t stream) {
    const int threads = 256;
    bool aligned = ((((uintptr_t)in) & 15) == 0) && ((((uintptr_t)out) & 15) == 0);
    long n8 = aligned ? (numel >> 3) : 0;
    if (n8 > 0) {
        long blocks = (n8 + threads - 1) / threads;
        cast_f32_f16_vec8_kernel<<<blocks, threads, 0, stream>>>(
            reinterpret_cast<const float4*>(in),
            reinterpret_cast<float4*>(out), n8);
    }
    long start = n8 << 3;
    long rem = numel - start;
    if (rem > 0) {
        long blocks = (rem + threads - 1) / threads;
        cast_f32_f16_scalar_kernel<<<blocks, threads, 0, stream>>>(
            in, reinterpret_cast<__half*>(out), start, numel);
    }
}

// ------------------------------- GEMM plan ---------------------------------
namespace {

struct GemmPlan {
    long M = -1, N = -1, K = -1;
    int engine = 2;   // 0: Lt fp16, 1: gemmEx FAST_16F, 2: gemmEx fp16, 3: Lt FAST_16F (fp32 io)
    cublasLtMatmulAlgo_t algo{};
    size_t ws = 0;
    bool valid = false;
};

GemmPlan g_plan;
cublasLtHandle_t g_lt = nullptr;
void* g_ws_ptr = nullptr;
size_t g_ws_cap = 0;
const size_t kMaxWs = 128ull << 20;

void* ensure_ws(size_t bytes) {
    if (bytes == 0) return nullptr;
    if (g_ws_cap < bytes) {
        if (g_ws_ptr) cudaFree(g_ws_ptr);
        g_ws_ptr = nullptr;
        if (cudaMalloc(&g_ws_ptr, bytes) != cudaSuccess) { g_ws_cap = 0; return nullptr; }
        g_ws_cap = bytes;
    }
    return g_ws_ptr;
}

float time_op(const std::function<cublasStatus_t()>& fn, int warmup, int iters,
              cudaStream_t stream, bool* ok) {
    for (int i = 0; i < warmup; i++) {
        if (fn() != CUBLAS_STATUS_SUCCESS) { *ok = false; return 1e30f; }
    }
    cudaEvent_t e0, e1;
    cudaEventCreate(&e0);
    cudaEventCreate(&e1);
    cudaEventRecord(e0, stream);
    cublasStatus_t st = CUBLAS_STATUS_SUCCESS;
    for (int i = 0; i < iters; i++) {
        st = fn();
        if (st != CUBLAS_STATUS_SUCCESS) break;
    }
    cudaEventRecord(e1, stream);
    cudaEventSynchronize(e1);
    float ms = 0.f;
    cudaEventElapsedTime(&ms, e0, e1);
    cudaEventDestroy(e0);
    cudaEventDestroy(e1);
    *ok = (st == CUBLAS_STATUS_SUCCESS);
    return *ok ? ms / iters : 1e30f;
}

struct LtCtx {
    cublasLtMatmulDesc_t desc = nullptr;
    cublasLtMatrixLayout_t Adesc = nullptr, Bdesc = nullptr, Cdesc = nullptr;
    bool ok = false;
    LtCtx(bool f16, long M, long N, long K) {
        cublasComputeType_t ct = f16 ? CUBLAS_COMPUTE_32F : CUBLAS_COMPUTE_32F_FAST_16F;
        cudaDataType_t ab = f16 ? CUDA_R_16F : CUDA_R_32F;
        if (cublasLtMatmulDescCreate(&desc, ct, CUDA_R_32F) != CUBLAS_STATUS_SUCCESS) return;
        cublasOperation_t opT = CUBLAS_OP_T, opN = CUBLAS_OP_N;
        cublasLtMatmulDescSetAttribute(desc, CUBLASLT_MATMUL_DESC_TRANSA, &opT, sizeof(opT));
        cublasLtMatmulDescSetAttribute(desc, CUBLASLT_MATMUL_DESC_TRANSB, &opN, sizeof(opN));
        ok = (cublasLtMatrixLayoutCreate(&Adesc, ab, K, N, K) == CUBLAS_STATUS_SUCCESS &&
              cublasLtMatrixLayoutCreate(&Bdesc, ab, K, M, K) == CUBLAS_STATUS_SUCCESS &&
              cublasLtMatrixLayoutCreate(&Cdesc, CUDA_R_32F, N, M, N) == CUBLAS_STATUS_SUCCESS);
    }
    ~LtCtx() {
        if (Cdesc) cublasLtMatrixLayoutDestroy(Cdesc);
        if (Bdesc) cublasLtMatrixLayoutDestroy(Bdesc);
        if (Adesc) cublasLtMatrixLayoutDestroy(Adesc);
        if (desc) cublasLtMatmulDescDestroy(desc);
    }
};

cublasStatus_t run_plan(const GemmPlan& p, cublasHandle_t blas,
                        const void* w32, const void* x32,
                        const void* w16, const void* x16,
                        float* out, cudaStream_t stream) {
    const float alpha = 1.0f, beta = 0.0f;
    if (p.engine == 1) {
        return cublasGemmEx(blas, CUBLAS_OP_T, CUBLAS_OP_N,
                            (int)p.N, (int)p.M, (int)p.K, &alpha,
                            w32, CUDA_R_32F, (int)p.K,
                            x32, CUDA_R_32F, (int)p.K,
                            &beta, out, CUDA_R_32F, (int)p.N,
                            CUBLAS_COMPUTE_32F_FAST_16F, CUBLAS_GEMM_DEFAULT);
    }
    if (p.engine == 2) {
        return cublasGemmEx(blas, CUBLAS_OP_T, CUBLAS_OP_N,
                            (int)p.N, (int)p.M, (int)p.K, &alpha,
                            w16, CUDA_R_16F, (int)p.K,
                            x16, CUDA_R_16F, (int)p.K,
                            &beta, out, CUDA_R_32F, (int)p.N,
                            CUBLAS_COMPUTE_32F, CUBLAS_GEMM_DEFAULT);
    }
    bool f16 = (p.engine == 0);
    LtCtx ctx(f16, p.M, p.N, p.K);
    if (!ctx.ok || g_lt == nullptr) return CUBLAS_STATUS_INTERNAL_ERROR;
    const void* A = f16 ? w16 : w32;
    const void* B = f16 ? x16 : x32;
    void* wsp = ensure_ws(p.ws);
    if (p.ws > 0 && wsp == nullptr) return CUBLAS_STATUS_ALLOC_FAILED;
    return cublasLtMatmul(g_lt, ctx.desc, &alpha, A, ctx.Adesc, B, ctx.Bdesc, &beta,
                          out, ctx.Cdesc, out, ctx.Cdesc, &p.algo, wsp, p.ws, stream);
}

}  // namespace

torch::Tensor gemm_autotune_cuda(torch::Tensor x, torch::Tensor w, torch::Tensor w_half) {
    TORCH_CHECK(x.is_cuda() && w.is_cuda(), "inputs must be CUDA tensors");
    TORCH_CHECK(x.scalar_type() == at::kFloat, "x must be float32");

    auto xc = x.contiguous();
    auto wc = w.contiguous();
    auto whc = w_half.contiguous();
    long K = xc.size(-1);
    long M = xc.numel() / K;
    long N = wc.size(0);
    TORCH_CHECK(wc.size(1) == K, "weight shape mismatch");

    auto out = at::empty({M, N}, xc.options());
    cudaStream_t stream = at::cuda::getCurrentCUDAStream();
    cublasHandle_t blas = at::cuda::getCurrentCUDABlasHandle();
    cublasSetStream(blas, stream);

    const float alpha = 1.0f;
    const float beta0 = 0.0f;

    bool tuned = g_plan.valid && g_plan.M == M && g_plan.N == N && g_plan.K == K;

    torch::Tensor xh;
    if (!tuned || g_plan.engine == 0 || g_plan.engine == 2) {
        xh = at::empty({M, K}, xc.options().dtype(at::kHalf));
        cast_to_half(xc.data_ptr<float>(), xh.data_ptr(), xc.numel(), stream);
    }

    if (tuned) {
        cublasStatus_t st = run_plan(g_plan, blas,
                                     wc.data_ptr<float>(), xc.data_ptr<float>(),
                                     whc.data_ptr(),
                                     xh.defined() ? xh.data_ptr() : nullptr,
                                     out.data_ptr<float>(), stream);
        TORCH_CHECK(st == CUBLAS_STATUS_SUCCESS, "GEMM failed: ", (int)st);
        return out;
    }

    // ------------------------------ autotune ------------------------------
    if (g_lt == nullptr && cublasLtCreate(&g_lt) != CUBLAS_STATUS_SUCCESS) g_lt = nullptr;

    bool ok = false;
    auto cast_fn = [&]() -> cublasStatus_t {
        cast_to_half(xc.data_ptr<float>(), xh.data_ptr(), xc.numel(), stream);
        return CUBLAS_STATUS_SUCCESS;
    };
    float t_cast = time_op(cast_fn, 1, 3, stream, &ok);
    if (!ok) t_cast = 0.f;

    GemmPlan best;
    best.engine = 2;
    float best_t = 1e30f;

    auto consider = [&](float t, bool needs_cast, int engine,
                        const cublasLtMatmulAlgo_t* algo, size_t ws) {
        float adj = t + (needs_cast ? t_cast : 0.f);
        if (adj < best_t) {
            best_t = adj;
            best.engine = engine;
            best.ws = ws;
            if (algo) best.algo = *algo;
        }
    };

    auto run_fp16 = [&]() -> cublasStatus_t {
        return cublasGemmEx(blas, CUBLAS_OP_T, CUBLAS_OP_N,
                            (int)N, (int)M, (int)K, &alpha,
                            whc.data_ptr(), CUDA_R_16F, (int)K,
                            xh.data_ptr(), CUDA_R_16F, (int)K,
                            &beta0, out.data_ptr<float>(), CUDA_R_32F, (int)N,
                            CUBLAS_COMPUTE_32F, CUBLAS_GEMM_DEFAULT);
    };
    auto run_fast16f = [&]() -> cublasStatus_t {
        return cublasGemmEx(blas, CUBLAS_OP_T, CUBLAS_OP_N,
                            (int)N, (int)M, (int)K, &alpha,
                            wc.data_ptr<float>(), CUDA_R_32F, (int)K,
                            xc.data_ptr<float>(), CUDA_R_32F, (int)K,
                            &beta0, out.data_ptr<float>(), CUDA_R_32F, (int)N,
                            CUBLAS_COMPUTE_32F_FAST_16F, CUBLAS_GEMM_DEFAULT);
    };
    {
        float t = time_op(run_fp16, 1, 3, stream, &ok);
        if (ok) consider(t, true, 2, nullptr, 0);
    }
    {
        float t = time_op(run_fast16f, 1, 3, stream, &ok);
        if (ok) consider(t, false, 1, nullptr, 0);
    }

    if (g_lt != nullptr) {
        for (int io = 0; io < 2; io++) {
            bool f16io = (io == 0);
            int engine = f16io ? 0 : 3;
            cublasComputeType_t ct = f16io ? CUBLAS_COMPUTE_32F : CUBLAS_COMPUTE_32F_FAST_16F;
            cudaDataType_t abt = f16io ? CUDA_R_16F : CUDA_R_32F;
            LtCtx ctx(f16io, M, N, K);
            if (!ctx.ok) continue;
            const void* Aptr = f16io ? (const void*)whc.data_ptr()
                                     : (const void*)wc.data_ptr<float>();
            const void* Bptr = f16io ? (const void*)xh.data_ptr()
                                     : (const void*)xc.data_ptr<float>();

            auto time_algo = [&](cublasLtMatmulAlgo_t* algo, size_t aws) {
                if (aws > kMaxWs) return;
                void* wsp = ensure_ws(aws);
                if (aws > 0 && wsp == nullptr) return;
                auto fn = [&]() -> cublasStatus_t {
                    return cublasLtMatmul(g_lt, ctx.desc, &alpha,
                                          Aptr, ctx.Adesc, Bptr, ctx.Bdesc, &beta0,
                                          out.data_ptr<float>(), ctx.Cdesc,
                                          out.data_ptr<float>(), ctx.Cdesc,
                                          algo, wsp, aws, stream);
                };
                float t = time_op(fn, 1, 3, stream, &ok);
                if (ok) consider(t, f16io, engine, algo, aws);
            };

            // (a) heuristic-picked algorithms
            cublasLtMatmulPreference_t pref;
            if (cublasLtMatmulPreferenceCreate(&pref) == CUBLAS_STATUS_SUCCESS) {
                cublasLtMatmulPreferenceSetAttribute(
                    pref, CUBLASLT_MATMUL_PREF_MAX_WORKSPACE_BYTES, &kMaxWs, sizeof(kMaxWs));
                cublasLtMatmulHeuristicResult_t hres[32];
                int hcount = 0;
                cublasLtMatmulAlgoGetHeuristic(g_lt, ctx.desc, ctx.Adesc, ctx.Bdesc,
                                               ctx.Cdesc, ctx.Cdesc, pref, 32, hres, &hcount);
                for (int i = 0; i < hcount; i++) {
                    if (hres[i].state == CUBLAS_STATUS_SUCCESS)
                        time_algo(&hres[i].algo, hres[i].workspaceSize);
                }
                cublasLtMatmulPreferenceDestroy(pref);
            }

            // (b) exhaustive sweep over all algorithm ids for these types
            std::vector<int> ids(128);
            int nids = 0;
            if (cublasLtMatmulAlgoGetIds(g_lt, ct, CUDA_R_32F, abt, abt,
                                         CUDA_R_32F, CUDA_R_32F,
                                         (int)ids.size(), ids.data(), &nids) == CUBLAS_STATUS_SUCCESS) {
                for (int i = 0; i < nids; i++) {
                    cublasLtMatmulAlgo_t algo;
                    if (cublasLtMatmulAlgoInit(g_lt, ct, CUDA_R_32F, abt, abt,
                                               CUDA_R_32F, CUDA_R_32F,
                                               ids[i], &algo) != CUBLAS_STATUS_SUCCESS)
                        continue;
                    cublasLtMatmulHeuristicResult_t res;
                    if (cublasLtMatmulAlgoCheck(g_lt, ctx.desc, ctx.Adesc, ctx.Bdesc,
                                                ctx.Cdesc, ctx.Cdesc, &algo, &res) != CUBLAS_STATUS_SUCCESS)
                        continue;
                    if (res.state != CUBLAS_STATUS_SUCCESS) continue;
                    time_algo(&algo, res.workspaceSize);
                }
            }
        }
    }

    g_plan = best;
    g_plan.M = M; g_plan.N = N; g_plan.K = K;
    g_plan.valid = true;

    // Final run with the winning algorithm so `out` is deterministic.
    cublasStatus_t st = run_plan(g_plan, blas,
                                 wc.data_ptr<float>(), xc.data_ptr<float>(),
                                 whc.data_ptr(), xh.data_ptr(),
                                 out.data_ptr<float>(), stream);
    TORCH_CHECK(st == CUBLAS_STATUS_SUCCESS, "GEMM failed: ", (int)st);
    return out;
}

// -------------------- fused epilogue: bias + GN + swish --------------------
// Vectorized kernel for group_size == 32 and num_groups % 4 == 0.
__global__ void fused_bias_gn_swish_mul_swish_vec4_kernel(
    const float* __restrict__ x,
    const float* __restrict__ bias,
    const float* __restrict__ gn_w,
    const float* __restrict__ gn_b,
    const float* __restrict__ mul_w,
    float* __restrict__ out,
    int num_groups,
    float eps,
    long total_quads)
{
    long warp_id = (long)blockIdx.x * (blockDim.x >> 5) + (threadIdx.x >> 5);
    if (warp_id >= total_quads) return;
    int lane = threadIdx.x & 31;

    long base = warp_id << 7;  // 128 floats per warp
    int ch0 = (int)(((warp_id << 2) % num_groups) << 5);
    int off = ch0 + (lane << 2);

    float4 v = *reinterpret_cast<const float4*>(x + base + (lane << 2));
    float4 b = *reinterpret_cast<const float4*>(bias + off);
    v.x += b.x; v.y += b.y; v.z += b.z; v.w += b.w;

    float s = (v.x + v.y) + (v.z + v.w);
    #pragma unroll
    for (int o = 4; o > 0; o >>= 1) s += __shfl_xor_sync(FULL_MASK, s, o, 8);
    float mean = s * (1.0f / 32.0f);

    float d0 = v.x - mean, d1 = v.y - mean, d2 = v.z - mean, d3 = v.w - mean;
    float ss = fmaf(d0, d0, fmaf(d1, d1, fmaf(d2, d2, d3 * d3)));
    #pragma unroll
    for (int o = 4; o > 0; o >>= 1) ss += __shfl_xor_sync(FULL_MASK, ss, o, 8);
    float rstd = rsqrtf(ss * (1.0f / 32.0f) + eps);

    float4 w = *reinterpret_cast<const float4*>(gn_w + off);
    float4 g = *reinterpret_cast<const float4*>(gn_b + off);
    float4 m = *reinterpret_cast<const float4*>(mul_w + off);

    #define EPILOGUE(val, ww, gg, mm)                      \
        {                                                  \
            float sc = rstd * (ww);                        \
            float y = fmaf(val, sc, (gg) - mean * sc);     \
            y = y / (1.0f + __expf(-y));                   \
            y = y * (mm);                                  \
            y = y / (1.0f + __expf(-y));                   \
            val = y;                                       \
        }
    EPILOGUE(v.x, w.x, g.x, m.x);
    EPILOGUE(v.y, w.y, g.y, m.y);
    EPILOGUE(v.z, w.z, g.z, m.z);
    EPILOGUE(v.w, w.w, g.w, m.w);
    #undef EPILOGUE

    *reinterpret_cast<float4*>(out + base + (lane << 2)) = v;
}

// Generic fallback: one warp per group, arbitrary group_size.
__global__ void fused_bias_gn_swish_mul_swish_generic_kernel(
    const float* __restrict__ x,
    const float* __restrict__ bias,
    const float* __restrict__ gn_w,
    const float* __restrict__ gn_b,
    const float* __restrict__ mul_w,
    float* __restrict__ out,
    int num_groups,
    int group_size,
    float eps,
    long total_groups)
{
    long warp_id = (long)blockIdx.x * (blockDim.x >> 5) + (threadIdx.x >> 5);
    if (warp_id >= total_groups) return;
    int lane = threadIdx.x & 31;

    long base = warp_id * (long)group_size;
    int ch0 = (int)(warp_id % num_groups) * group_size;

    float sum = 0.f, sumsq = 0.f;
    for (int i = lane; i < group_size; i += 32) {
        float v = x[base + i] + bias[ch0 + i];
        sum += v;
        sumsq = fmaf(v, v, sumsq);
    }
    #pragma unroll
    for (int o = 16; o > 0; o >>= 1) {
        sum += __shfl_xor_sync(FULL_MASK, sum, o);
        sumsq += __shfl_xor_sync(FULL_MASK, sumsq, o);
    }
    float inv_n = 1.0f / (float)group_size;
    float mean = sum * inv_n;
    float var = fmaf(-mean, mean, sumsq * inv_n);
    float rstd = rsqrtf(var + eps);

    for (int i = lane; i < group_size; i += 32) {
        int c = ch0 + i;
        float v = x[base + i] + bias[c];
        float sc = rstd * gn_w[c];
        float y = fmaf(v, sc, gn_b[c] - mean * sc);
        y = y / (1.0f + __expf(-y));
        y = y * mul_w[c];
        y = y / (1.0f + __expf(-y));
        out[base + i] = y;
    }
}

torch::Tensor fused_bias_gn_swish_mul_swish_cuda(
    torch::Tensor x,
    torch::Tensor bias,
    torch::Tensor gn_w,
    torch::Tensor gn_b,
    torch::Tensor mul_w,
    int64_t num_groups,
    double eps)
{
    TORCH_CHECK(x.is_cuda(), "x must be a CUDA tensor");
    TORCH_CHECK(x.scalar_type() == at::kFloat, "x must be float32");

    auto xc = x.contiguous();
    auto bc = bias.contiguous();
    auto wc = gn_w.contiguous();
    auto gc = gn_b.contiguous();
    auto mc = mul_w.contiguous();

    auto out = at::empty_like(xc);

    int64_t channels = xc.size(-1);
    int64_t batch = xc.numel() / channels;
    int64_t group_size = channels / num_groups;
    int64_t total_groups = batch * num_groups;

    const int threads = 256;
    const int warps_per_block = threads / 32;
    cudaStream_t stream = at::cuda::getCurrentCUDAStream();

    if (group_size == 32 && (num_groups % 4) == 0) {
        int64_t quads = total_groups / 4;
        int64_t blocks = (quads + warps_per_block - 1) / warps_per_block;
        fused_bias_gn_swish_mul_swish_vec4_kernel<<<blocks, threads, 0, stream>>>(
            xc.data_ptr<float>(), bc.data_ptr<float>(),
            wc.data_ptr<float>(), gc.data_ptr<float>(), mc.data_ptr<float>(),
            out.data_ptr<float>(),
            (int)num_groups, (float)eps, quads);
    } else {
        int64_t blocks = (total_groups + warps_per_block - 1) / warps_per_block;
        fused_bias_gn_swish_mul_swish_generic_kernel<<<blocks, threads, 0, stream>>>(
            xc.data_ptr<float>(), bc.data_ptr<float>(),
            wc.data_ptr<float>(), gc.data_ptr<float>(), mc.data_ptr<float>(),
            out.data_ptr<float>(),
            (int)num_groups, (int)group_size, (float)eps, total_groups);
    }

    return out;
}
"""

fused_cpp_source = (
    "torch::Tensor gemm_autotune_cuda(torch::Tensor x, torch::Tensor w, torch::Tensor w_half);\n"
    "torch::Tensor fused_bias_gn_swish_mul_swish_cuda("
    "torch::Tensor x, torch::Tensor bias, torch::Tensor gn_w, "
    "torch::Tensor gn_b, torch::Tensor mul_w, int64_t num_groups, double eps);"
)

fused_ops = load_inline(
    name="fused_gemm_gn_swish_v4",
    cpp_sources=fused_cpp_source,
    cuda_sources=fused_cuda_source,
    functions=["gemm_autotune_cuda", "fused_bias_gn_swish_mul_swish_cuda"],
    verbose=False,
    extra_cflags=["-O3"],
    extra_cuda_cflags=["-O3", "--use_fast_math"],
    extra_ldflags=["-lcublas", "-lcublasLt"],
)


class ModelNew(nn.Module):
    """
    Optimized model:
      * GEMM via one-time exhaustively autotuned cuBLASLt/cuBLAS tensor-core
        path. FP16-input algorithms (FP32 accumulate / FP32 output) are
        compared against native FAST_16F FP32-input algorithms with the cost
        of the explicit FP32->FP16 cast included, so the truly fastest
        end-to-end variant wins.
      * Cached half-precision weight.
      * Fused bias + GroupNorm + Swish + Multiply + Swish epilogue kernel.
    """
    def __init__(self, in_features, out_features, num_groups, multiply_weight_shape):
        super(ModelNew, self).__init__()
        self.gemm = nn.Linear(in_features, out_features)
        self.group_norm = nn.GroupNorm(num_groups, out_features)
        self.multiply_weight = nn.Parameter(torch.randn(multiply_weight_shape))
        self.num_groups = num_groups
        self.eps = self.group_norm.eps
        self._w_half = None
        self._w_half_key = None
        self._gemm_fn = fused_ops.gemm_autotune_cuda
        self._epilogue_fn = fused_ops.fused_bias_gn_swish_mul_swish_cuda

    def _half_weight(self):
        w = self.gemm.weight
        key = (w._version, w.data_ptr(), w.device)
        if self._w_half is None or self._w_half_key != key:
            self._w_half = w.detach().half().contiguous()
            self._w_half_key = key
        return self._w_half

    def forward(self, x):
        if x.is_cuda and x.dtype == torch.float32:
            two_d = x.dim() == 2
            x2 = x if two_d else x.reshape(-1, x.shape[-1])
            h = self._gemm_fn(x2, self.gemm.weight, self._half_weight())
            y = self._epilogue_fn(
                h,
                self.gemm.bias,
                self.group_norm.weight,
                self.group_norm.bias,
                self.multiply_weight,
                self.num_groups,
                self.eps,
            )
            return y if two_d else y.view(*x.shape[:-1], y.shape[-1])
        # Fallback for CPU / non-float32 inputs
        x = self.gemm(x)
        x = self.group_norm(x)
        x = x * torch.sigmoid(x)
        x = x * self.multiply_weight
        x = x * torch.sigmoid(x)
        return x
```