```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 os
import sys
import glob
import site

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


# ---------------------------------------------------------------------------
# Fused epilogue: + combined bias (linear bias + add_value), Swish, Tanh, GELU.
# Note: Tanh maps into (-1, 1) and GELU on (-1, 1) stays within [-0.17, 0.85],
# so the final Hardtanh([-1, 1]) is mathematically an identity and is omitted.
# Fast-but-accurate math: __expf-based sigmoid/tanh (inf-safe), and a Taylor
# polynomial for erf valid on |z| <= 0.7072 (max abs err ~1.2e-7).
# ---------------------------------------------------------------------------
epilogue_source = r"""
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include <cuda_runtime.h>

__device__ __forceinline__ float act_chain(float x) {
    // Swish: x * sigmoid(x)
    float sig = 1.0f / (1.0f + __expf(-x));
    float s = x * sig;
    // Tanh: 1 - 2 / (exp(2s) + 1)  (safe for exp overflow -> inf)
    float u = __expf(2.0f * s);
    float t = 1.0f - __fdividef(2.0f, u + 1.0f);
    // GELU: 0.5 * t * (1 + erf(t / sqrt(2))), |t| <= 1
    float z = t * 0.70710678118654752f;
    float z2 = z * z;
    float p = 0.000120553394f;
    p = fmaf(p, z2, -0.000854832702f);
    p = fmaf(p, z2, 0.00522397763f);
    p = fmaf(p, z2, -0.02686617065f);
    p = fmaf(p, z2, 0.11283791671f);
    p = fmaf(p, z2, -0.37612638903f);
    p = fmaf(p, z2, 1.12837916710f);
    float e = z * p;
    return 0.5f * t * (1.0f + e);
}

template<int THREADS, int UNROLL>
__global__ void fused_bias_act_vec4(const float4* __restrict__ x,
                                    const float4* __restrict__ b,
                                    float4* __restrict__ out,
                                    int n4, int row4) {
    int base = blockIdx.x * (THREADS * UNROLL) + threadIdx.x;
    bool pow2 = (row4 & (row4 - 1)) == 0;
    #pragma unroll
    for (int u = 0; u < UNROLL; ++u) {
        int i = base + u * THREADS;
        if (i < n4) {
            int j = pow2 ? (i & (row4 - 1)) : (i % row4);
            float4 v = __ldcs(&x[i]);
            float4 q = __ldg(&b[j]);
            v.x = act_chain(v.x + q.x);
            v.y = act_chain(v.y + q.y);
            v.z = act_chain(v.z + q.z);
            v.w = act_chain(v.w + q.w);
            __stcs(&out[i], v);
        }
    }
}

__global__ void fused_bias_act_scalar(const float* __restrict__ x,
                                      const float* __restrict__ b,
                                      float* __restrict__ out,
                                      int n, int row) {
    int i = blockIdx.x * blockDim.x + threadIdx.x;
    if (i < n) {
        out[i] = act_chain(__ldcs(&x[i]) + __ldg(&b[i % row]));
    }
}

torch::Tensor fused_bias_act(torch::Tensor x, torch::Tensor b, int64_t variant) {
    TORCH_CHECK(x.is_cuda() && b.is_cuda(), "Inputs must be CUDA tensors");
    TORCH_CHECK(x.scalar_type() == at::kFloat, "Input must be float32");
    auto xc = x.contiguous();
    auto bc = b.contiguous();
    auto out = torch::empty_like(xc);

    int64_t n = xc.numel();
    int64_t row = bc.numel();
    if (n == 0) return out;

    cudaStream_t stream = at::cuda::getCurrentCUDAStream();
    if ((row % 4 == 0) && (n % 4 == 0)) {
        int64_t n4 = n / 4;
        int64_t row4 = row / 4;
        auto launch = [&](auto threads_c, auto unroll_c) {
            constexpr int T = decltype(threads_c)::value;
            constexpr int U = decltype(unroll_c)::value;
            int64_t per_block = (int64_t)T * U;
            int blocks = (int)((n4 + per_block - 1) / per_block);
            fused_bias_act_vec4<T, U><<<blocks, T, 0, stream>>>(
                reinterpret_cast<const float4*>(xc.data_ptr<float>()),
                reinterpret_cast<const float4*>(bc.data_ptr<float>()),
                reinterpret_cast<float4*>(out.data_ptr<float>()),
                (int)n4, (int)row4);
        };
        using i128 = std::integral_constant<int, 128>;
        using i256 = std::integral_constant<int, 256>;
        using i512 = std::integral_constant<int, 512>;
        using u2 = std::integral_constant<int, 2>;
        using u4 = std::integral_constant<int, 4>;
        using u8 = std::integral_constant<int, 8>;
        switch (variant) {
            case 1: launch(i256{}, u8{}); break;
            case 2: launch(i512{}, u4{}); break;
            case 3: launch(i128{}, u8{}); break;
            case 4: launch(i512{}, u2{}); break;
            case 5: launch(i128{}, u4{}); break;
            default: launch(i256{}, u4{}); break;
        }
    } else {
        const int threads = 256;
        int blocks = (int)((n + threads - 1) / threads);
        fused_bias_act_scalar<<<blocks, threads, 0, stream>>>(
            xc.data_ptr<float>(),
            bc.data_ptr<float>(),
            out.data_ptr<float>(),
            (int)n, (int)row);
    }
    return out;
}
"""

epilogue_cpp_source = "torch::Tensor fused_bias_act(torch::Tensor x, torch::Tensor b, int64_t variant);"

epilogue_ext = None
try:
    epilogue_ext = load_inline(
        name="fused_bias_act_v8",
        cpp_sources=epilogue_cpp_source,
        cuda_sources=epilogue_source,
        functions=["fused_bias_act"],
        verbose=False,
        extra_cflags=["-O3"],
        extra_cuda_cflags=["-O3"],
    )
except Exception:
    epilogue_ext = None


# ---------------------------------------------------------------------------
# cuBLASLt GEMM wrapper with deep per-shape autotuning:
#   * FP32 compute with TF32 tensor cores (CUBLAS_COMPUTE_32F_FAST_TF32),
#     explicitly permitted for this workload.
#   * Both weight layouts: TN (original W with OP_T) and NN (pre-transposed Wt
#     with OP_N).
#   * One large workspace per layout so split-K algorithms are available.
#   * Besides the heuristic algorithm list, manual overrides of
#     CUBLASLT_ALGO_CONFIG_SPLITK_NUM / REDUCTION_SCHEME / CTA_SWIZZLING are
#     swept for the most promising algorithms (helps wave quantization for
#     shapes with a small CTA grid like M=1024, N=K=8192).
#
# Row-major goal: Out[M,N] = X[M,K] * W[N,K]^T.
# Column-major view: C'[N,M] = A'[N,K] * B'[K,M] with
#   B' = X^T  -> x buffer, desc (K, M, ld=K), OP_N
#   C' = Out^T-> out buffer, desc (N, M, ld=N)
#   A' = W    -> use_wt=false: W buffer, desc (K, N, ld=K), OP_T
#             -> use_wt=true : Wt buffer ([K,N] row-major), desc (N, K, ld=N), OP_N
# ---------------------------------------------------------------------------
lt_source = r"""
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include <cuda_runtime.h>
#include <cublasLt.h>
#include <vector>
#include <map>
#include <mutex>
#include <memory>
#include <string>

namespace {

struct LtPlan {
    cublasLtMatmulDesc_t opDesc = nullptr;
    cublasLtMatrixLayout_t Adesc = nullptr, Bdesc = nullptr, Cdesc = nullptr;
    std::vector<cublasLtMatmulHeuristicResult_t> heur;
    void* workspace = nullptr;
    size_t ws_bytes = 0;
    ~LtPlan() {
        if (opDesc) cublasLtMatmulDescDestroy(opDesc);
        if (Adesc) cublasLtMatrixLayoutDestroy(Adesc);
        if (Bdesc) cublasLtMatrixLayoutDestroy(Bdesc);
        if (Cdesc) cublasLtMatrixLayoutDestroy(Cdesc);
        if (workspace) cudaFree(workspace);
    }
};

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

std::map<std::string, std::shared_ptr<LtPlan>>& plan_cache() {
    static std::map<std::string, std::shared_ptr<LtPlan>> c;
    return c;
}
std::mutex& cache_mutex() { static std::mutex m; return m; }

std::string make_key(int64_t M, int64_t N, int64_t K, bool use_wt) {
    return std::to_string(M) + "_" + std::to_string(N) + "_" + std::to_string(K)
           + (use_wt ? "_t" : "_n");
}

std::shared_ptr<LtPlan> get_plan(int64_t M, int64_t N, int64_t K, bool use_wt) {
    std::string key = make_key(M, N, K, use_wt);
    std::lock_guard<std::mutex> lock(cache_mutex());
    auto it = plan_cache().find(key);
    if (it != plan_cache().end()) return it->second;

    auto plan = std::make_shared<LtPlan>();
    if (cublasLtMatmulDescCreate(&plan->opDesc, CUBLAS_COMPUTE_32F_FAST_TF32, CUDA_R_32F) != CUBLAS_STATUS_SUCCESS) {
        plan_cache()[key] = plan;
        return plan;
    }
    cublasOperation_t opA = use_wt ? CUBLAS_OP_N : CUBLAS_OP_T;
    cublasOperation_t opB = CUBLAS_OP_N;
    cublasLtMatmulDescSetAttribute(plan->opDesc, CUBLASLT_MATMUL_DESC_TRANSA, &opA, sizeof(opA));
    cublasLtMatmulDescSetAttribute(plan->opDesc, CUBLASLT_MATMUL_DESC_TRANSB, &opB, sizeof(opB));

    if (use_wt) {
        cublasLtMatrixLayoutCreate(&plan->Adesc, CUDA_R_32F, N, K, N);
    } else {
        cublasLtMatrixLayoutCreate(&plan->Adesc, CUDA_R_32F, K, N, K);
    }
    cublasLtMatrixLayoutCreate(&plan->Bdesc, CUDA_R_32F, K, M, K);
    cublasLtMatrixLayoutCreate(&plan->Cdesc, CUDA_R_32F, N, M, N);

    const size_t ws_cands[] = {
        (size_t)512 * 1024 * 1024, (size_t)256 * 1024 * 1024,
        (size_t)64 * 1024 * 1024, (size_t)16 * 1024 * 1024,
        (size_t)1024 * 1024, (size_t)0
    };
    for (size_t cand : ws_cands) {
        if (cand == 0) {
            plan->workspace = nullptr;
            plan->ws_bytes = 0;
            break;
        }
        if (cudaMalloc(&plan->workspace, cand) == cudaSuccess) {
            plan->ws_bytes = cand;
            break;
        }
        cudaGetLastError();
    }

    cublasLtMatmulPreference_t pref = nullptr;
    if (cublasLtMatmulPreferenceCreate(&pref) == CUBLAS_STATUS_SUCCESS) {
        size_t maxws = plan->ws_bytes;
        cublasLtMatmulPreferenceSetAttribute(pref, CUBLASLT_MATMUL_PREF_MAX_WORKSPACE_BYTES,
                                             &maxws, sizeof(maxws));
        plan->heur.resize(128);
        int returned = 0;
        cublasStatus_t st = cublasLtMatmulAlgoGetHeuristic(
            lt_handle(), plan->opDesc, plan->Adesc, plan->Bdesc, plan->Cdesc, plan->Cdesc,
            pref, 128, plan->heur.data(), &returned);
        cublasLtMatmulPreferenceDestroy(pref);
        if (st == CUBLAS_STATUS_SUCCESS && returned > 0) {
            plan->heur.resize(returned);
            std::vector<cublasLtMatmulHeuristicResult_t> good;
            good.reserve(plan->heur.size());
            for (auto& h : plan->heur) {
                if (h.state == CUBLAS_STATUS_SUCCESS && h.workspaceSize <= plan->ws_bytes) {
                    good.push_back(h);
                }
            }
            plan->heur.swap(good);
        } else {
            plan->heur.clear();
        }
    }

    plan_cache()[key] = plan;
    return plan;
}

} // namespace

int64_t lt_num_algos(int64_t M, int64_t N, int64_t K, bool use_wt) {
    return (int64_t)get_plan(M, N, K, use_wt)->heur.size();
}

int64_t lt_gemm_cfg(torch::Tensor x, torch::Tensor w, torch::Tensor out,
                    bool use_wt, int64_t algo_idx, int64_t splitk, int64_t swz) {
    int64_t M = x.size(0);
    int64_t K = x.size(1);
    int64_t N = out.size(1);
    auto plan = get_plan(M, N, K, use_wt);
    if (algo_idx < 0 || algo_idx >= (int64_t)plan->heur.size()) return 999;

    cublasLtMatmulAlgo_t algo = plan->heur[algo_idx].algo;  // mutable copy
    if (splitk > 0) {
        int sk = (int)splitk;
        cublasLtReductionScheme_t rs =
            (sk > 1) ? CUBLASLT_REDUCTION_SCHEME_INPLACE : CUBLASLT_REDUCTION_SCHEME_NONE;
        cublasLtMatmulAlgoConfigSetAttribute(&algo, CUBLASLT_ALGO_CONFIG_SPLITK_NUM, &sk, sizeof(sk));
        cublasLtMatmulAlgoConfigSetAttribute(&algo, CUBLASLT_ALGO_CONFIG_REDUCTION_SCHEME, &rs, sizeof(rs));
    }
    if (swz >= 0) {
        uint32_t s = (uint32_t)swz;
        cublasLtMatmulAlgoConfigSetAttribute(&algo, CUBLASLT_ALGO_CONFIG_CTA_SWIZZLING, &s, sizeof(s));
    }

    float alpha = 1.0f;
    float beta = 0.0f;
    cudaStream_t stream = at::cuda::getCurrentCUDAStream();
    cublasStatus_t st = cublasLtMatmul(lt_handle(), plan->opDesc, &alpha,
                                       w.data_ptr<float>(), plan->Adesc,
                                       x.data_ptr<float>(), plan->Bdesc,
                                       &beta,
                                       out.data_ptr<float>(), plan->Cdesc,
                                       out.data_ptr<float>(), plan->Cdesc,
                                       &algo, plan->workspace, plan->ws_bytes,
                                       stream);
    return (int64_t)st;
}

void lt_prune(int64_t M, int64_t N, int64_t K, int64_t keep_layout) {
    std::string prefix = std::to_string(M) + "_" + std::to_string(N) + "_" + std::to_string(K) + "_";
    std::string keep = (keep_layout == 1) ? make_key(M, N, K, true)
                     : (keep_layout == 0) ? make_key(M, N, K, false)
                                          : std::string("\x01", 1);
    std::lock_guard<std::mutex> lock(cache_mutex());
    auto& c = plan_cache();
    for (auto it = c.begin(); it != c.end();) {
        if (it->first.compare(0, prefix.size(), prefix) == 0 && it->first != keep) {
            it = c.erase(it);
        } else {
            ++it;
        }
    }
}
"""

lt_cpp_source = (
    "int64_t lt_num_algos(int64_t M, int64_t N, int64_t K, bool use_wt);\n"
    "int64_t lt_gemm_cfg(torch::Tensor x, torch::Tensor w, torch::Tensor out, bool use_wt, int64_t algo_idx, int64_t splitk, int64_t swz);\n"
    "void lt_prune(int64_t M, int64_t N, int64_t K, int64_t keep_layout);"
)


def _cublaslt_paths():
    inc_dirs, lib_dirs = [], []
    roots = []
    try:
        roots += site.getsitepackages()
    except Exception:
        pass
    try:
        roots.append(os.path.dirname(os.path.dirname(os.path.abspath(torch.__file__))))
    except Exception:
        pass
    roots.append(sys.prefix)
    roots.append(os.environ.get("CUDA_HOME", "/usr/local/cuda"))
    roots.append("/usr/local/cuda")
    for r in dict.fromkeys(roots):
        for sub in ("lib64", "lib",
                    os.path.join("nvidia", "cublas", "lib"),
                    os.path.join("torch", "lib")):
            p = os.path.join(r, sub)
            if os.path.isdir(p):
                lib_dirs.append(p)
        for sub in ("include", os.path.join("nvidia", "cublas", "include")):
            p = os.path.join(r, sub)
            if os.path.isdir(p):
                inc_dirs.append(p)
    lib_file = None
    for d in lib_dirs:
        for pat in ("libcublasLt.so", "libcublasLt.so.*"):
            hits = sorted(glob.glob(os.path.join(d, pat)))
            hits = [h for h in hits if os.path.isfile(h)]
            if hits:
                lib_file = hits[0]
                break
        if lib_file:
            break
    if lib_file:
        ldflags = [lib_file, "-Wl,-rpath," + os.path.dirname(lib_file)]
    else:
        ldflags = ["-lcublasLt"]
    return inc_dirs, ldflags


lt_ext = None
try:
    _inc, _ld = _cublaslt_paths()
    lt_ext = load_inline(
        name="cublaslt_gemm_autotune_v6",
        cpp_sources=lt_cpp_source,
        cuda_sources=lt_source,
        functions=["lt_num_algos", "lt_gemm_cfg", "lt_prune"],
        verbose=False,
        extra_cflags=["-O3"],
        extra_cuda_cflags=["-O3"],
        extra_include_paths=_inc,
        extra_ldflags=_ld,
    )
except Exception:
    lt_ext = None


class ModelNew(nn.Module):
    """
    Optimized model:
      * GEMM (FP32 with TF32 tensor cores) via cuBLASLt, autotuned per shape over
        heuristic algorithms, weight layouts (TN with the original weight, NN
        with a pre-transposed copy), and manual overrides of split-K factor /
        reduction scheme / CTA swizzling on the most promising algorithms
        (improves CTA-grid wave quantization). Every candidate family is
        numerically validated against the PyTorch cuBLAS result before it can
        win. Falls back to the PyTorch cuBLAS path when cublasLt is
        unavailable.
      * Single fused CUDA epilogue kernel: combined bias (linear bias + add_value)
        + Swish + Tanh + GELU. The trailing Hardtanh([-1, 1]) is provably an
        identity after Tanh -> GELU and is dropped.
      * Whole forward is captured into a CUDA graph (keyed by input/parameter
        buffer signatures) and replayed on subsequent calls to remove host and
        launch overhead; any mismatch or failure falls back to the eager path.
    """
    def __init__(self, in_features, out_features, add_value_shape):
        super(ModelNew, self).__init__()
        # TF32 tensor-core math is explicitly permitted for FP32 workloads.
        torch.backends.cuda.matmul.allow_tf32 = True
        torch.backends.cudnn.allow_tf32 = True

        self.matmul = nn.Linear(in_features, out_features)
        self.add_value = nn.Parameter(torch.randn(add_value_shape))

        self._gemm_choice = {}
        self._wt = None
        self._wt_ver = None
        self._comb = None
        self._comb_ver = None
        self._epi_variant = None

        self._graphs = {}
        self._counts = {}
        self._bad = set()
        self._max_graphs = 16
        self._graphs_ok = hasattr(torch.cuda, "CUDAGraph") and hasattr(torch.cuda, "graph")

    # ----- cached derived buffers -------------------------------------------
    def _get_wt(self):
        w = self.matmul.weight
        ver = (w.data_ptr(), w._version)
        if self._wt is None or self._wt_ver != ver:
            with torch.no_grad():
                self._wt = w.t().contiguous()
            self._wt_ver = ver
        return self._wt

    def _get_comb(self):
        b = self.matmul.bias
        av = self.add_value
        ver = (av.data_ptr(), av._version,
               0 if b is None else b.data_ptr(), -1 if b is None else b._version)
        if self._comb is None or self._comb_ver != ver:
            with torch.no_grad():
                self._comb = av.detach().clone() if b is None else (av + b).contiguous()
            self._comb_ver = ver
        return self._comb

    # ----- reference / eager paths ------------------------------------------
    def _reference(self, x):
        y = F.linear(x, self.matmul.weight, self.matmul.bias)
        y = y + self.add_value
        y = torch.sigmoid(y) * y
        y = torch.tanh(y)
        y = F.gelu(y)
        y = F.hardtanh(y, min_val=-1, max_val=1)
        return y

    @torch.no_grad()
    def _time_fn(self, fn, dev, warmup, iters):
        fn()
        for _ in range(max(0, warmup - 1)):
            fn()
        torch.cuda.synchronize(dev)
        s = torch.cuda.Event(enable_timing=True)
        e = torch.cuda.Event(enable_timing=True)
        s.record()
        for _ in range(iters):
            fn()
        e.record()
        torch.cuda.synchronize(dev)
        return s.elapsed_time(e)

    @torch.no_grad()
    def _pick_epilogue_variant(self, y, comb):
        if epilogue_ext is None:
            return 0
        dev = y.device
        best, best_t = 0, float("inf")
        ref = None
        for v in range(6):
            try:
                r = epilogue_ext.fused_bias_act(y, comb, v)
                if ref is None:
                    ref = r
                elif not torch.allclose(r, ref, rtol=1e-4, atol=1e-5):
                    continue
                t = self._time_fn(lambda vv=v: epilogue_ext.fused_bias_act(y, comb, vv),
                                  dev, 2, 6)
                if t < best_t:
                    best_t, best = t, v
            except Exception:
                continue
        return best

    @torch.no_grad()
    def _autotune(self, x):
        M, K = x.shape
        N = self.matmul.weight.shape[0]
        dev = x.device
        w = self.matmul.weight
        wt = self._get_wt()
        xc = x if x.is_contiguous() else x.contiguous()
        out = torch.empty(M, N, device=dev, dtype=torch.float32)
        ref = torch.mm(x, wt)

        def ok():
            return bool(torch.allclose(out, ref, rtol=5e-2, atol=1.0))

        fns = {}

        def make_lt(use_wt, wbuf, i, sk, sw):
            def f():
                st = lt_ext.lt_gemm_cfg(xc, wbuf, out, use_wt, i, sk, sw)
                if st != 0:
                    raise RuntimeError("cublasLt status %d" % int(st))
            return f

        fns[("torch_tn",)] = lambda: torch.matmul(x, w.t(), out=out)
        fns[("torch_nn",)] = lambda: torch.mm(x, wt, out=out)

        layouts = []
        if lt_ext is not None:
            for use_wt, wbuf in ((False, w), (True, wt)):
                try:
                    n_alg = int(lt_ext.lt_num_algos(M, N, K, use_wt))
                except Exception:
                    n_alg = 0
                if n_alg <= 0:
                    continue
                try:
                    st = lt_ext.lt_gemm_cfg(xc, wbuf, out, use_wt, 0, 0, -1)
                except Exception:
                    st = -1
                if st != 0 or not ok():
                    continue
                layouts.append((use_wt, wbuf, n_alg))
                for i in range(min(n_alg, 12)):
                    fns[("lt", use_wt, i, 0, -1)] = make_lt(use_wt, wbuf, i, 0, -1)

        # Stage 1: coarse timing of base candidates.
        scored = []
        for tag, fn in fns.items():
            try:
                t = self._time_fn(fn, dev, 1, 2)
                scored.append((t, tag))
            except Exception:
                pass
        if not scored:
            return ("torch_tn",)
        scored.sort(key=lambda z: z[0])

        # Stage 2: split-K / swizzle sweep on the best LT algos of each layout.
        sweep = []
        for use_wt, wbuf, n_alg in layouts:
            cnt = 0
            for t0, tag in scored:
                if tag[0] == "lt" and tag[1] == use_wt:
                    i = tag[2]
                    for sk in (1, 2, 3, 4, 6, 8):
                        for sw in (0, 1):
                            ntag = ("lt", use_wt, i, sk, sw)
                            if ntag not in fns:
                                fns[ntag] = make_lt(use_wt, wbuf, i, sk, sw)
                                sweep.append(ntag)
                    cnt += 1
                    if cnt >= 3:
                        break
        for tag in sweep:
            try:
                t = self._time_fn(fns[tag], dev, 1, 2)
                scored.append((t, tag))
            except Exception:
                pass
        scored.sort(key=lambda z: z[0])

        # Stage 3: refine the top few with validation and longer timing.
        best, best_t = ("torch_tn",), float("inf")
        checked = 0
        for t0, tag in scored:
            if checked >= 8:
                break
            fn = fns.get(tag)
            if fn is None:
                continue
            try:
                fn()
                if tag[0] == "lt" and not ok():
                    continue
                t = self._time_fn(fn, dev, 2, 10)
            except Exception:
                continue
            checked += 1
            if t < best_t:
                best_t, best = t, tag

        # Free losing cublasLt plans (and their workspaces) for this shape.
        if lt_ext is not None:
            try:
                kl = -1
                if best[0] == "lt":
                    kl = 1 if best[1] else 0
                lt_ext.lt_prune(M, N, K, kl)
            except Exception:
                pass
        return best

    def _run_gemm(self, x, choice):
        w = self.matmul.weight
        N = w.shape[0]
        M = x.shape[0]
        tag = choice[0]
        if tag == "lt" and lt_ext is not None:
            try:
                _, use_wt, idx, sk, sw = choice
                xc = x if x.is_contiguous() else x.contiguous()
                wbuf = self._get_wt() if use_wt else w
                out = torch.empty(M, N, device=x.device, dtype=torch.float32)
                st = lt_ext.lt_gemm_cfg(xc, wbuf, out, use_wt, idx, sk, sw)
                if st == 0:
                    return out
            except Exception:
                pass
            tag = "torch_nn"  # recorded winner stopped working: safe fallback
        if tag == "torch_nn":
            return torch.mm(x, self._get_wt())
        return torch.matmul(x, w.t())

    def _eager(self, x):
        key = (tuple(x.shape), x.device.index)
        choice = self._gemm_choice.get(key)
        if choice is None:
            choice = self._autotune(x)
            self._gemm_choice[key] = choice
        y = self._run_gemm(x, choice)
        if epilogue_ext is not None:
            comb = self._get_comb()
            v = self._epi_variant
            if v is None:
                v = self._pick_epilogue_variant(y, comb)
                self._epi_variant = v
            return epilogue_ext.fused_bias_act(y, comb, v)
        y = y + self._get_comb()
        y = torch.sigmoid(y) * y
        y = torch.tanh(y)
        y = F.gelu(y)
        return y

    def _capture(self, x, key):
        # Warm up on a side stream (materializes workspaces, caches, etc.).
        s = torch.cuda.Stream(x.device)
        s.wait_stream(torch.cuda.current_stream(x.device))
        with torch.cuda.stream(s):
            for _ in range(3):
                self._eager(x)
        torch.cuda.current_stream(x.device).wait_stream(s)
        torch.cuda.synchronize(x.device)

        g = torch.cuda.CUDAGraph()
        with torch.cuda.graph(g):
            static_out = self._eager(x)
        self._graphs[key] = (g, static_out)

    def forward(self, x):
        if (self._graphs_ok and x.is_cuda and x.dim() == 2
                and x.dtype == torch.float32 and self.matmul.weight.is_cuda):
            w = self.matmul.weight
            b = self.matmul.bias
            key = (x.data_ptr(), tuple(x.shape), tuple(x.stride()), x.device.index,
                   w.data_ptr(), w._version,
                   0 if b is None else b.data_ptr(), -1 if b is None else b._version,
                   self.add_value.data_ptr(), self.add_value._version)
            ent = self._graphs.get(key)
            if ent is not None:
                try:
                    ent[0].replay()
                    return ent[1]
                except Exception:
                    self._graphs.pop(key, None)
                    self._bad.add(key)
            elif key not in self._bad:
                n = self._counts.get(key, 0) + 1
                self._counts[key] = n
                if n >= 3 and len(self._graphs) < self._max_graphs:
                    try:
                        self._capture(x, key)
                        ent = self._graphs.get(key)
                        if ent is not None:
                            ent[0].replay()
                            return ent[1]
                    except Exception:
                        self._graphs.pop(key, None)
                        self._bad.add(key)
            return self._eager(x)
        # Non-CUDA fallback: replicate the reference computation exactly.
        return self._reference(x)
```