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

# ---------------------------------------------------------------------------
# Single fused extension:
#   * vectorized FP32 -> FP16 cast (8 floats per thread, streaming loads)
#   * GEMM via cuBLASLt / cuBLAS tensor-core paths on dedicated handles
#   * fused bias + GroupNorm + Swish + per-channel multiply + Swish epilogue
#   * software pipeline: the batch is split into chunks processed on two
#     CUDA streams so the bandwidth-bound casts / epilogues overlap with the
#     compute-bound GEMMs.  One-time autotuning picks (chunk count, GEMM
#     engine, cublasLt algorithm) by timing the whole pipelined forward.
# ---------------------------------------------------------------------------
fused_cuda_source = r"""
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include <c10/cuda/CUDAGuard.h>
#include <cublas_v2.h>
#include <cublasLt.h>
#include <cuda_runtime.h>
#include <cuda_fp16.h>
#include <functional>
#include <algorithm>
#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 = __ldcs(in + 2 * i);
    float4 b = __ldcs(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_tail_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_tail_kernel<<<blocks, threads, 0, stream>>>(
            in, reinterpret_cast<__half*>(out), start, numel);
    }
}

// -------------------- fused epilogue: bias + GN + swish --------------------
__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;
}

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

static void epilogue_launch(const float* h, const float* bias, const float* gnw,
                            const float* gnb, const float* mulw, float* out,
                            long rows, long channels, int num_groups, float eps,
                            cudaStream_t stream) {
    long group_size = channels / num_groups;
    long total_groups = rows * (long)num_groups;
    const int threads = 256;
    const int wpb = threads / 32;
    if (group_size == 32 && (num_groups % 4) == 0) {
        long quads = total_groups / 4;
        long blocks = (quads + wpb - 1) / wpb;
        fused_bias_gn_swish_mul_swish_vec4_kernel<<<blocks, threads, 0, stream>>>(
            h, bias, gnw, gnb, mulw, out, num_groups, eps, quads);
    } else {
        long blocks = (total_groups + wpb - 1) / wpb;
        fused_bias_gn_swish_mul_swish_generic_kernel<<<blocks, threads, 0, stream>>>(
            h, bias, gnw, gnb, mulw, out, num_groups, (int)group_size, eps, total_groups);
    }
}

// ------------------------------ GEMM plumbing ------------------------------
namespace {

struct Plan {
    long M = -1, N = -1, K = -1;
    int engine = 2;   // 0: Lt fp16, 1: gemmEx FAST_16F (fp32 io), 2: gemmEx fp16
    int nchunks = 1;
    long Mc = 0;
    cublasLtMatmulAlgo_t algo{};
    size_t ws = 0;
    cublasLtMatmulDesc_t desc = nullptr;
    cublasLtMatrixLayout_t Adesc = nullptr, Bdesc = nullptr, Cdesc = nullptr;
    bool valid = false;
};

struct IO {
    const float* x;
    const float* w32;
    const __half* w16;
    const float* bias;
    const float* gnw;
    const float* gnb;
    const float* mulw;
    float* out;
    long M, N, K;
    int num_groups;
    float eps;
};

Plan g_plan;
cublasHandle_t g_blas[2] = {nullptr, nullptr};
cublasLtHandle_t g_lt = nullptr;
cudaStream_t g_streams[2] = {nullptr, nullptr};
cudaEvent_t g_ev_start = nullptr;
cudaEvent_t g_ev_done[2] = {nullptr, nullptr};
void* g_ws[2] = {nullptr, nullptr};
size_t g_ws_bytes[2] = {0, 0};
void* g_xh = nullptr;
size_t g_xh_cap = 0;
void* g_h = nullptr;
size_t g_h_cap = 0;
const size_t kLtWsCap = 32u << 20;
const size_t kHandleWs = 32u << 20;

void init_once() {
    static bool done = false;
    if (done) return;
    done = true;
    for (int i = 0; i < 2; i++) {
        TORCH_CHECK(cublasCreate(&g_blas[i]) == CUBLAS_STATUS_SUCCESS,
                    "cublasCreate failed");
        void* p = nullptr;
        if (cudaMalloc(&p, kHandleWs) == cudaSuccess && p != nullptr) {
            cublasSetWorkspace(g_blas[i], p, kHandleWs);
        }
        cudaStreamCreateWithFlags(&g_streams[i], cudaStreamNonBlocking);
        cudaEventCreateWithFlags(&g_ev_done[i], cudaEventDisableTiming);
    }
    cudaEventCreateWithFlags(&g_ev_start, cudaEventDisableTiming);
    if (cublasLtCreate(&g_lt) != CUBLAS_STATUS_SUCCESS) g_lt = nullptr;
}

void* ensure_buf(void** p, size_t* cap, size_t bytes) {
    if (bytes == 0) return nullptr;
    if (*cap < bytes) {
        cudaDeviceSynchronize();
        if (*p) cudaFree(*p);
        *p = nullptr;
        *cap = 0;
        if (cudaMalloc(p, bytes) != cudaSuccess) return nullptr;
        *cap = bytes;
    }
    return *p;
}

void free_lt_objs(Plan& p) {
    if (p.Adesc) { cublasLtMatrixLayoutDestroy(p.Adesc); p.Adesc = nullptr; }
    if (p.Bdesc) { cublasLtMatrixLayoutDestroy(p.Bdesc); p.Bdesc = nullptr; }
    if (p.Cdesc) { cublasLtMatrixLayoutDestroy(p.Cdesc); p.Cdesc = nullptr; }
    if (p.desc)  { cublasLtMatmulDescDestroy(p.desc);    p.desc = nullptr; }
}

void plan_clear() {
    free_lt_objs(g_plan);
    g_plan = Plan();
}

bool make_lt_objs(Plan& p, long Mc, long N, long K) {
    if (cublasLtMatmulDescCreate(&p.desc, CUBLAS_COMPUTE_32F, CUDA_R_32F) != CUBLAS_STATUS_SUCCESS)
        return false;
    cublasOperation_t opT = CUBLAS_OP_T, opN = CUBLAS_OP_N;
    cublasLtMatmulDescSetAttribute(p.desc, CUBLASLT_MATMUL_DESC_TRANSA, &opT, sizeof(opT));
    cublasLtMatmulDescSetAttribute(p.desc, CUBLASLT_MATMUL_DESC_TRANSB, &opN, sizeof(opN));
    bool ok = (cublasLtMatrixLayoutCreate(&p.Adesc, CUDA_R_16F, K, N, K) == CUBLAS_STATUS_SUCCESS &&
               cublasLtMatrixLayoutCreate(&p.Bdesc, CUDA_R_16F, K, Mc, K) == CUBLAS_STATUS_SUCCESS &&
               cublasLtMatrixLayoutCreate(&p.Cdesc, CUDA_R_32F, N, Mc, N) == CUBLAS_STATUS_SUCCESS);
    if (!ok) free_lt_objs(p);
    return ok;
}

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

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

cublasStatus_t gemm_chunk(const Plan& p, int slot, const IO& io,
                          const float* x32, const __half* x16, float* h,
                          long Mc, cudaStream_t stream) {
    const float alpha = 1.0f, beta = 0.0f;
    long N = io.N, K = io.K;
    if (p.engine == 0) {
        return cublasLtMatmul(g_lt, p.desc, &alpha,
                              io.w16, p.Adesc, x16, p.Bdesc, &beta,
                              h, p.Cdesc, h, p.Cdesc,
                              &p.algo, g_ws[slot], p.ws, stream);
    }
    cublasHandle_t hb = g_blas[slot];
    cublasSetStream(hb, stream);
    if (p.engine == 1) {
        return cublasGemmEx(hb, CUBLAS_OP_T, CUBLAS_OP_N,
                            (int)N, (int)Mc, (int)K, &alpha,
                            io.w32, CUDA_R_32F, (int)K,
                            x32, CUDA_R_32F, (int)K,
                            &beta, h, CUDA_R_32F, (int)N,
                            CUBLAS_COMPUTE_32F_FAST_16F, CUBLAS_GEMM_DEFAULT);
    }
    return cublasGemmEx(hb, CUBLAS_OP_T, CUBLAS_OP_N,
                        (int)N, (int)Mc, (int)K, &alpha,
                        io.w16, CUDA_R_16F, (int)K,
                        x16, CUDA_R_16F, (int)K,
                        &beta, h, CUDA_R_32F, (int)N,
                        CUBLAS_COMPUTE_32F, CUBLAS_GEMM_DEFAULT);
}

void run_pipeline(const IO& io, const Plan& p, cudaStream_t main_stream) {
    int nc = p.nchunks;
    if (nc <= 1) {
        if (p.engine != 1) cast_to_half(io.x, g_xh, io.M * io.K, main_stream);
        gemm_chunk(p, 0, io, io.x, (const __half*)g_xh, (float*)g_h, io.M, main_stream);
        epilogue_launch((const float*)g_h, io.bias, io.gnw, io.gnb, io.mulw,
                        io.out, io.M, io.N, io.num_groups, io.eps, main_stream);
        return;
    }
    long Mc = p.Mc;
    cudaEventRecord(g_ev_start, main_stream);
    cudaStreamWaitEvent(g_streams[0], g_ev_start, 0);
    cudaStreamWaitEvent(g_streams[1], g_ev_start, 0);
    for (int i = 0; i < nc; i++) {
        int s = i & 1;
        cudaStream_t st = g_streams[s];
        const float* xc = io.x + (long)i * Mc * io.K;
        const __half* xhc = (const __half*)g_xh + (long)i * Mc * io.K;
        float* hc = (float*)g_h + (long)i * Mc * io.N;
        float* oc = io.out + (long)i * Mc * io.N;
        if (p.engine != 1) cast_to_half(xc, (void*)xhc, Mc * io.K, st);
        gemm_chunk(p, s, io, xc, xhc, hc, Mc, st);
        epilogue_launch(hc, io.bias, io.gnw, io.gnb, io.mulw, oc,
                        Mc, io.N, io.num_groups, io.eps, st);
    }
    cudaEventRecord(g_ev_done[0], g_streams[0]);
    cudaEventRecord(g_ev_done[1], g_streams[1]);
    cudaStreamWaitEvent(main_stream, g_ev_done[0], 0);
    cudaStreamWaitEvent(main_stream, g_ev_done[1], 0);
}

float time_pipeline(const IO& io, const Plan& p, cudaStream_t main_stream,
                    int warmup, int iters) {
    for (int i = 0; i < warmup; i++) run_pipeline(io, p, main_stream);
    cudaEvent_t e0, e1;
    cudaEventCreate(&e0);
    cudaEventCreate(&e1);
    cudaEventRecord(e0, main_stream);
    for (int i = 0; i < iters; i++) run_pipeline(io, p, main_stream);
    cudaEventRecord(e1, main_stream);
    cudaEventSynchronize(e1);
    float ms = 0.f;
    cudaEventElapsedTime(&ms, e0, e1);
    cudaEventDestroy(e0);
    cudaEventDestroy(e1);
    return ms / iters;
}

void autotune(const IO& io, cudaStream_t main_stream) {
    plan_clear();
    Plan best;
    float best_t = 1e30f;
    bool have = false;
    const float alpha = 1.0f, beta = 0.0f;
    const int cand_nc[4] = {1, 2, 4, 8};

    for (int ci = 0; ci < 4; ci++) {
        int nc = cand_nc[ci];
        if ((long)nc > io.M) break;
        if (io.M % nc != 0) continue;
        long Mc = io.M / nc;
        bool ok = false;

        float t_cast = 0.f;
        {
            auto cfn = [&]() -> cublasStatus_t {
                cast_to_half(io.x, g_xh, Mc * io.K, main_stream);
                return CUBLAS_STATUS_SUCCESS;
            };
            t_cast = time_op(cfn, 1, 2, main_stream, &ok);
            if (!ok) t_cast = 0.f;
        }

        bool fp16_ok = false, f16f_ok = false;
        {
            auto fn = [&]() -> cublasStatus_t {
                cublasSetStream(g_blas[0], main_stream);
                return cublasGemmEx(g_blas[0], CUBLAS_OP_T, CUBLAS_OP_N,
                                    (int)io.N, (int)Mc, (int)io.K, &alpha,
                                    io.w16, CUDA_R_16F, (int)io.K,
                                    (const __half*)g_xh, CUDA_R_16F, (int)io.K,
                                    &beta, (float*)g_h, CUDA_R_32F, (int)io.N,
                                    CUBLAS_COMPUTE_32F, CUBLAS_GEMM_DEFAULT);
            };
            float t = time_op(fn, 1, 2, main_stream, &ok);
            if (ok) fp16_ok = true;
        }
        {
            auto fn = [&]() -> cublasStatus_t {
                cublasSetStream(g_blas[0], main_stream);
                return cublasGemmEx(g_blas[0], CUBLAS_OP_T, CUBLAS_OP_N,
                                    (int)io.N, (int)Mc, (int)io.K, &alpha,
                                    io.w32, CUDA_R_32F, (int)io.K,
                                    io.x, CUDA_R_32F, (int)io.K,
                                    &beta, (float*)g_h, CUDA_R_32F, (int)io.N,
                                    CUBLAS_COMPUTE_32F_FAST_16F, CUBLAS_GEMM_DEFAULT);
            };
            float t = time_op(fn, 1, 2, main_stream, &ok);
            if (ok) f16f_ok = true;
        }

        bool lt_ok = false;
        float t_lt = 1e30f;
        cublasLtMatmulAlgo_t lt_algo{};
        size_t lt_ws = 0;
        if (g_lt != nullptr) {
            LtCtx ctx(true, Mc, io.N, io.K);
            if (ctx.ok) {
                cublasLtMatmulPreference_t pref;
                if (cublasLtMatmulPreferenceCreate(&pref) == CUBLAS_STATUS_SUCCESS) {
                    size_t maxws = kLtWsCap;
                    cublasLtMatmulPreferenceSetAttribute(
                        pref, CUBLASLT_MATMUL_PREF_MAX_WORKSPACE_BYTES, &maxws, sizeof(maxws));
                    cublasLtMatmulHeuristicResult_t hres[8];
                    int hc = 0;
                    cublasLtMatmulAlgoGetHeuristic(g_lt, ctx.desc, ctx.Adesc, ctx.Bdesc,
                                                   ctx.Cdesc, ctx.Cdesc, pref, 8, hres, &hc);
                    for (int i = 0; i < hc; i++) {
                        if (hres[i].state != CUBLAS_STATUS_SUCCESS) continue;
                        size_t aws = hres[i].workspaceSize;
                        if (aws > kLtWsCap) continue;
                        if (aws > 0 && ensure_buf(&g_ws[0], &g_ws_bytes[0], aws) == nullptr)
                            continue;
                        cublasLtMatmulAlgo_t algo = hres[i].algo;
                        auto fn = [&]() -> cublasStatus_t {
                            return cublasLtMatmul(g_lt, ctx.desc, &alpha,
                                                  io.w16, ctx.Adesc,
                                                  (const __half*)g_xh, ctx.Bdesc, &beta,
                                                  (float*)g_h, ctx.Cdesc,
                                                  (float*)g_h, ctx.Cdesc,
                                                  &algo, g_ws[0], aws, main_stream);
                        };
                        float t = time_op(fn, 1, 2, main_stream, &ok);
                        if (ok && t < t_lt) {
                            t_lt = t;
                            lt_algo = algo;
                            lt_ws = aws;
                            lt_ok = true;
                        }
                    }
                    cublasLtMatmulPreferenceDestroy(pref);
                }
            }
        }

        struct Cand { int engine; cublasLtMatmulAlgo_t algo; size_t ws; };
        Cand cands[3];
        int ncand = 0;
        if (lt_ok)   cands[ncand++] = Cand{0, lt_algo, lt_ws};
        if (fp16_ok) cands[ncand++] = Cand{2, cublasLtMatmulAlgo_t{}, 0};
        if (f16f_ok) cands[ncand++] = Cand{1, cublasLtMatmulAlgo_t{}, 0};

        for (int k = 0; k < ncand; k++) {
            Plan p;
            p.engine = cands[k].engine;
            p.nchunks = nc;
            p.Mc = Mc;
            p.ws = cands[k].ws;
            p.algo = cands[k].algo;
            if (p.engine == 0) {
                if (!make_lt_objs(p, Mc, io.N, io.K)) continue;
                bool ws_ok = true;
                if (p.ws > 0) {
                    for (int s = 0; s < 2; s++) {
                        if (ensure_buf(&g_ws[s], &g_ws_bytes[s], p.ws) == nullptr)
                            ws_ok = false;
                    }
                }
                if (!ws_ok) { free_lt_objs(p); continue; }
            }
            float t = time_pipeline(io, p, main_stream, 1, 3);
            if (t < best_t) {
                if (have) free_lt_objs(best);
                best = p;
                best_t = t;
                have = true;
            } else {
                free_lt_objs(p);
            }
        }
    }

    if (!have) {
        best = Plan();
        best.engine = 2;
        best.nchunks = 1;
        best.Mc = io.M;
    }
    g_plan = best;
    g_plan.M = io.M;
    g_plan.N = io.N;
    g_plan.K = io.K;
    g_plan.valid = true;
}

}  // namespace

torch::Tensor fused_forward_cuda(torch::Tensor x, torch::Tensor w, torch::Tensor w_half,
                                 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() && w.is_cuda(), "inputs must be CUDA tensors");
    TORCH_CHECK(x.scalar_type() == at::kFloat, "x must be float32");
    const c10::cuda::CUDAGuard device_guard(x.device());
    init_once();

    auto xc = x.contiguous();
    auto wc = w.contiguous();
    auto whc = w_half.contiguous();
    auto bc = bias.contiguous();
    auto gwc = gn_w.contiguous();
    auto gbc = gn_b.contiguous();
    auto mc = mul_w.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 main_stream = at::cuda::getCurrentCUDAStream();

    TORCH_CHECK(ensure_buf(&g_xh, &g_xh_cap, (size_t)M * (size_t)K * 2) != nullptr,
                "half buffer allocation failed");
    TORCH_CHECK(ensure_buf(&g_h, &g_h_cap, (size_t)M * (size_t)N * 4) != nullptr,
                "gemm buffer allocation failed");

    IO io;
    io.x = xc.data_ptr<float>();
    io.w32 = wc.data_ptr<float>();
    io.w16 = reinterpret_cast<const __half*>(whc.data_ptr());
    io.bias = bc.data_ptr<float>();
    io.gnw = gwc.data_ptr<float>();
    io.gnb = gbc.data_ptr<float>();
    io.mulw = mc.data_ptr<float>();
    io.out = out.data_ptr<float>();
    io.M = M;
    io.N = N;
    io.K = K;
    io.num_groups = (int)num_groups;
    io.eps = (float)eps;

    if (!(g_plan.valid && g_plan.M == M && g_plan.N == N && g_plan.K == K)) {
        autotune(io, main_stream);
    }
    run_pipeline(io, g_plan, main_stream);
    return out;
}
"""

fused_cpp_source = (
    "torch::Tensor fused_forward_cuda(torch::Tensor x, torch::Tensor w, "
    "torch::Tensor w_half, 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_v5",
    cpp_sources=fused_cpp_source,
    cuda_sources=fused_cuda_source,
    functions=["fused_forward_cuda"],
    verbose=False,
    extra_cflags=["-O3"],
    extra_cuda_cflags=["-O3", "--use_fast_math"],
    extra_ldflags=["-lcublas", "-lcublasLt"],
)


class ModelNew(nn.Module):
    """
    Optimized model:
      * One fused CUDA entry point for the whole forward pass.
      * GEMM on tensor cores (FP16 inputs, FP32 accumulate/output, or native
        FAST_16F on FP32 inputs) using dedicated cuBLAS handles on two CUDA
        streams.
      * Software pipelining: the batch is split into chunks whose casts and
        epilogues (bandwidth-bound) overlap with the GEMMs (compute-bound).
      * One-time autotuning of (chunk count, GEMM engine, cublasLt algorithm)
        by timing the entire pipelined forward.
      * 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._fwd = fused_ops.fused_forward_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])
            y = self._fwd(
                x2,
                self.gemm.weight,
                self._half_weight(),
                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
```