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

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

# The GEMM dominates runtime. With M=128, N=K=16384 and the default
# 256x128 tile, only 64 CTAs are launched on a 108-SM A100, so the GEMM is
# tensor-core-bound at ~59% SM utilization. Splitting K (or using smaller
# tiles) engages all SMs and approaches the DRAM floor (~512 MiB of fp16
# weight traffic). cuBLASLt's default heuristic does not pick such configs,
# so we EXPLICITLY enumerate manual (tile, split-K) algorithm configs via
# cublasLtMatmulAlgoInit/ConfigSetAttribute/Check, plus the heuristic list,
# time every candidate on the real problem shape, numerically validate each
# against the default algorithm's output, and cache the fastest.
torch.backends.cuda.matmul.allow_tf32 = True
torch.backends.cudnn.allow_tf32 = True
torch.backends.cuda.matmul.allow_fp16_reduced_precision_reduction = False

cuda_src = r"""
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include <c10/cuda/CUDAGuard.h>
#include <cuda_runtime.h>
#include <cuda_fp16.h>

namespace {

__global__ void f32_to_f16_vec_kernel(const float* __restrict__ in,
                                      __half* __restrict__ out, long n) {
    long base = ((long)blockIdx.x * blockDim.x + threadIdx.x) * 8;
    if (base >= n) return;
    float4 a = *reinterpret_cast<const float4*>(in + base);
    float4 b = *reinterpret_cast<const float4*>(in + base + 4);
    __half2* o = reinterpret_cast<__half2*>(out + base);
    o[0] = __floats2half2_rn(a.x, a.y);
    o[1] = __floats2half2_rn(a.z, a.w);
    o[2] = __floats2half2_rn(b.x, b.y);
    o[3] = __floats2half2_rn(b.z, b.w);
}

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

__global__ void bias_min0_vec_kernel(const __half* __restrict__ y,
                                     const float* __restrict__ bias,
                                     const float* __restrict__ cp,
                                     float* __restrict__ out,
                                     long n, long cols) {
    long i = ((long)blockIdx.x * blockDim.x + threadIdx.x) * 8;
    if (i >= n) return;
    float c = *cp;
    int j = (int)(i % cols);
    float4 bLo = *reinterpret_cast<const float4*>(bias + j);
    float4 bHi = *reinterpret_cast<const float4*>(bias + j + 4);
    float4 hv = *reinterpret_cast<const float4*>(y + i);
    const __half2* h2 = reinterpret_cast<const __half2*>(&hv);
    float2 f0 = __half22float2(h2[0]);
    float2 f1 = __half22float2(h2[1]);
    float2 f2 = __half22float2(h2[2]);
    float2 f3 = __half22float2(h2[3]);
    float4 oLo = make_float4(
        fminf(f0.x + (bLo.x - c), 0.f),
        fminf(f0.y + (bLo.y - c), 0.f),
        fminf(f1.x + (bLo.z - c), 0.f),
        fminf(f1.y + (bLo.w - c), 0.f));
    float4 oHi = make_float4(
        fminf(f2.x + (bHi.x - c), 0.f),
        fminf(f2.y + (bHi.y - c), 0.f),
        fminf(f3.x + (bHi.z - c), 0.f),
        fminf(f3.y + (bHi.w - c), 0.f));
    reinterpret_cast<float4*>(out + i)[0] = oLo;
    reinterpret_cast<float4*>(out + i)[1] = oHi;
}

__global__ void bias_min0_scalar_kernel(const __half* __restrict__ y,
                                        const float* __restrict__ bias,
                                        const float* __restrict__ cp,
                                        float* __restrict__ out,
                                        long n, long cols) {
    long i = (long)blockIdx.x * blockDim.x + threadIdx.x;
    if (i < n) {
        float c = *cp;
        int j = (int)(i % cols);
        out[i] = fminf(__half2float(y[i]) + (bias[j] - c), 0.f);
    }
}

__device__ __forceinline__ void atomicMaxNonNeg(float* addr, float v) {
    // Valid because v >= 0: for non-negative floats the integer bit pattern
    // ordering matches the float ordering (NaN maps to a huge int and is
    // therefore treated as "reject" by the host-side threshold check).
    atomicMax(reinterpret_cast<int*>(addr), __float_as_int(v));
}

__global__ void diff_stats_kernel(const __half* __restrict__ a,
                                  const __half* __restrict__ b, long n,
                                  float* __restrict__ out2) {
    __shared__ float s_d, s_r;
    if (threadIdx.x == 0) {
        s_d = 0.f;
        s_r = 0.f;
    }
    __syncthreads();
    long stride = (long)gridDim.x * blockDim.x;
    float md = 0.f, mr = 0.f;
    for (long i = (long)blockIdx.x * blockDim.x + threadIdx.x; i < n;
         i += stride) {
        float fa = __half2float(a[i]);
        float fb = __half2float(b[i]);
        md = fmaxf(md, fabsf(fa - fb));
        mr = fmaxf(mr, fabsf(fb));
    }
    atomicMaxNonNeg(&s_d, md);
    atomicMaxNonNeg(&s_r, mr);
    __syncthreads();
    if (threadIdx.x == 0) {
        atomicMaxNonNeg(&out2[0], s_d);
        atomicMaxNonNeg(&out2[1], s_r);
    }
}

}  // namespace

void launch_f32_to_f16(const void* in, void* out, long n, cudaStream_t stream) {
    if (n <= 0) return;
    const float* ip = reinterpret_cast<const float*>(in);
    __half* op = reinterpret_cast<__half*>(out);
    const int block = 256;
    if ((n & 7L) == 0) {
        long items = n >> 3;
        long blocks = (items + block - 1) / block;
        f32_to_f16_vec_kernel<<<blocks, block, 0, stream>>>(ip, op, n);
    } else {
        long blocks = (n + block - 1) / block;
        f32_to_f16_scalar_kernel<<<blocks, block, 0, stream>>>(ip, op, n);
    }
}

void launch_bias_min0(const void* y, const float* bias, const float* cp,
                      float* out, long n, long cols, cudaStream_t stream) {
    if (n <= 0) return;
    const __half* yp = reinterpret_cast<const __half*>(y);
    const int block = 256;
    if ((cols & 7L) == 0) {
        long items = n >> 3;
        long blocks = (items + block - 1) / block;
        bias_min0_vec_kernel<<<blocks, block, 0, stream>>>(yp, bias, cp, out,
                                                           n, cols);
    } else {
        long blocks = (n + block - 1) / block;
        bias_min0_scalar_kernel<<<blocks, block, 0, stream>>>(yp, bias, cp, out,
                                                              n, cols);
    }
}

void launch_diff_stats(const void* a, const void* b, long n, float* out2,
                       cudaStream_t stream) {
    if (n <= 0) return;
    int blocks = (int)std::min<long>(512, (n + 255) / 256);
    diff_stats_kernel<<<blocks, 256, 0, stream>>>(
        reinterpret_cast<const __half*>(a),
        reinterpret_cast<const __half*>(b), n, out2);
}

torch::Tensor bias_min0_epilogue(torch::Tensor y, torch::Tensor bias,
                                 torch::Tensor c) {
    TORCH_CHECK(y.is_cuda() && bias.is_cuda() && c.is_cuda(),
                "CUDA tensors required");
    TORCH_CHECK(y.scalar_type() == at::kHalf, "y must be float16");
    TORCH_CHECK(bias.scalar_type() == at::kFloat, "bias must be float32");
    TORCH_CHECK(c.scalar_type() == at::kFloat, "c must be float32");

    const at::cuda::OptionalCUDAGuard device_guard(at::device_of(y));
    auto yc = y.contiguous();
    auto bc = bias.contiguous();
    auto cc = c.contiguous();
    auto out = torch::empty(yc.sizes(), yc.options().dtype(at::kFloat));

    long n = yc.numel();
    if (n == 0) {
        return out;
    }
    long cols = yc.size(-1);
    cudaStream_t stream = at::cuda::getCurrentCUDAStream().stream();
    launch_bias_min0(yc.data_ptr<at::Half>(), bc.data_ptr<float>(),
                     cc.data_ptr<float>(), out.data_ptr<float>(), n, cols,
                     stream);
    return out;
}
"""

cpp_src = r"""
#include <torch/extension.h>
#include <c10/cuda/CUDAGuard.h>
#include <c10/cuda/CUDAStream.h>
#include <cuda_runtime.h>
#include <cublasLt.h>
#include <dlfcn.h>
#include <string.h>
#include <algorithm>
#include <cmath>
#include <map>
#include <mutex>
#include <string>
#include <vector>

// Implemented in cuda.cu
void launch_f32_to_f16(const void* in, void* out, long n, cudaStream_t stream);
void launch_bias_min0(const void* y, const float* bias, const float* cp,
                      float* out, long n, long cols, cudaStream_t stream);
void launch_diff_stats(const void* a, const void* b, long n, float* out2,
                       cudaStream_t stream);
torch::Tensor bias_min0_epilogue(torch::Tensor y, torch::Tensor bias,
                                 torch::Tensor c);

namespace {

// ----------------- runtime-resolved cuBLASLt entry points -----------------

struct LtApi {
    cublasStatus_t (*create)(cublasLtHandle_t*) = nullptr;
    cublasStatus_t (*desc_create)(cublasLtMatmulDesc_t*, cublasComputeType_t,
                                  cudaDataType_t) = nullptr;
    cublasStatus_t (*desc_set)(cublasLtMatmulDesc_t,
                               cublasLtMatmulDescAttributes_t, const void*,
                               size_t) = nullptr;
    cublasStatus_t (*layout_create)(cublasLtMatrixLayout_t*, cudaDataType_t,
                                    uint64_t, uint64_t, int64_t) = nullptr;
    cublasStatus_t (*pref_create)(cublasLtMatmulPreference_t*) = nullptr;
    cublasStatus_t (*pref_set)(cublasLtMatmulPreference_t,
                               cublasLtMatmulPreferenceAttributes_t,
                               const void*, size_t) = nullptr;
    cublasStatus_t (*get_heuristic)(cublasLtHandle_t, cublasLtMatmulDesc_t,
                                    cublasLtMatrixLayout_t,
                                    cublasLtMatrixLayout_t,
                                    cublasLtMatrixLayout_t,
                                    cublasLtMatrixLayout_t,
                                    cublasLtMatmulPreference_t, int,
                                    cublasLtMatmulHeuristicResult_t*,
                                    int*) = nullptr;
    cublasStatus_t (*matmul)(cublasLtHandle_t, cublasLtMatmulDesc_t,
                             const void*, const void*, cublasLtMatrixLayout_t,
                             const void*, cublasLtMatrixLayout_t, const void*,
                             const void*, cublasLtMatrixLayout_t, void*,
                             cublasLtMatrixLayout_t,
                             const cublasLtMatmulAlgo_t*, void*, size_t,
                             cudaStream_t) = nullptr;
    // Optional (used for explicit tile / split-K enumeration).
    cublasStatus_t (*algo_get_ids)(cublasLtHandle_t, cublasComputeType_t,
                                   cudaDataType_t, cudaDataType_t,
                                   cudaDataType_t, cudaDataType_t,
                                   cudaDataType_t, int, int*, int*) = nullptr;
    cublasStatus_t (*algo_init)(cublasLtHandle_t, cublasComputeType_t,
                                cudaDataType_t, cudaDataType_t, cudaDataType_t,
                                cudaDataType_t, cudaDataType_t, int,
                                cublasLtMatmulAlgo_t*) = nullptr;
    cublasStatus_t (*algo_cfg_set)(cublasLtMatmulAlgo_t*,
                                   cublasLtMatmulAlgoConfigAttributes_t,
                                   const void*, size_t) = nullptr;
    cublasStatus_t (*algo_check)(cublasLtHandle_t, cublasLtMatmulDesc_t,
                                 cublasLtMatrixLayout_t, cublasLtMatrixLayout_t,
                                 cublasLtMatrixLayout_t, cublasLtMatrixLayout_t,
                                 const cublasLtMatmulAlgo_t*,
                                 cublasLtMatmulHeuristicResult_t*) = nullptr;
};

LtApi g_api;
void* g_dlhandle = nullptr;
bool g_resolved = false;
bool g_ok = false;
std::mutex g_init_mtx;
cublasLtHandle_t g_lt = nullptr;

void* get_sym(const char* name) {
    void* s = dlsym(RTLD_DEFAULT, name);
    if (s == nullptr && g_dlhandle != nullptr) {
        s = dlsym(g_dlhandle, name);
    }
    return s;
}

// ------------------------------- GEMM plans --------------------------------

struct Plan {
    cublasLtMatmulDesc_t opDesc = nullptr;
    cublasLtMatrixLayout_t Adesc = nullptr;
    cublasLtMatrixLayout_t Bdesc = nullptr;
    cublasLtMatrixLayout_t Ddesc = nullptr;
    cublasLtMatmulAlgo_t algo{};
    size_t workspace_size = 0;
    bool tuned = false;
};

struct PlanKey {
    long m, n, k;
    bool operator<(const PlanKey& o) const {
        if (m != o.m) return m < o.m;
        if (n != o.n) return n < o.n;
        return k < o.k;
    }
};

struct Cand {
    cublasLtMatmulAlgo_t algo;
    size_t ws;
};

std::map<PlanKey, Plan> g_plans;
std::mutex g_plan_mtx;
void* g_ws = nullptr;
size_t g_ws_cap = 0;
void* g_refbuf = nullptr;
size_t g_refbuf_cap = 0;
float* g_stats = nullptr;

void* ensure_buffer(void** buf, size_t* cap, size_t bytes) {
    if (bytes == 0) return nullptr;
    if (bytes <= *cap) return *buf;
    if (*buf != nullptr) {
        cudaFree(*buf);
        *buf = nullptr;
        *cap = 0;
    }
    if (cudaMalloc(buf, bytes) != cudaSuccess) {
        cudaGetLastError();
        *buf = nullptr;
        *cap = 0;
        return nullptr;
    }
    *cap = bytes;
    return *buf;
}

// Row-major y (m x n) = row-major x (m x k) @ row-major w (n x k)^T.
// Column-major view used by cuBLASLt:
//   A = w viewed as (k x n) col-major, ld = k, op = T
//   B = x viewed as (k x m) col-major, ld = k, op = N
//   D = y viewed as (n x m) col-major, ld = n
const Plan& get_plan(long m, long n, long k, const void* A, const void* B,
                     void* D, cudaStream_t stream) {
    PlanKey key{m, n, k};
    std::lock_guard<std::mutex> lock(g_plan_mtx);
    auto it = g_plans.find(key);
    if (it != g_plans.end()) {
        return it->second;
    }

    Plan plan;
    cublasStatus_t st;
    st = g_api.desc_create(&plan.opDesc, CUBLAS_COMPUTE_32F, CUDA_R_32F);
    TORCH_CHECK(st == CUBLAS_STATUS_SUCCESS,
                "cublasLtMatmulDescCreate failed: ", (int)st);
    cublasOperation_t opT = CUBLAS_OP_T;
    cublasOperation_t opN = CUBLAS_OP_N;
    g_api.desc_set(plan.opDesc, CUBLASLT_MATMUL_DESC_TRANSA, &opT,
                   sizeof(opT));
    g_api.desc_set(plan.opDesc, CUBLASLT_MATMUL_DESC_TRANSB, &opN,
                   sizeof(opN));
    g_api.layout_create(&plan.Adesc, CUDA_R_16F, (uint64_t)k, (uint64_t)n,
                        (int64_t)k);
    g_api.layout_create(&plan.Bdesc, CUDA_R_16F, (uint64_t)k, (uint64_t)m,
                        (int64_t)k);
    g_api.layout_create(&plan.Ddesc, CUDA_R_16F, (uint64_t)n, (uint64_t)m,
                        (int64_t)n);

    const float alpha = 1.f, beta = 0.f;
    std::vector<Cand> cands;

    // ------------------------- heuristic candidates -------------------------
    cublasLtMatmulPreference_t pref = nullptr;
    if (g_api.pref_create(&pref) == CUBLAS_STATUS_SUCCESS) {
        size_t max_ws = 128ull * 1024 * 1024;
        g_api.pref_set(pref, CUBLASLT_MATMUL_PREF_MAX_WORKSPACE_BYTES, &max_ws,
                       sizeof(max_ws));
        cublasLtReductionScheme_t mask = CUBLASLT_REDUCTION_SCHEME_MASK;
        g_api.pref_set(pref, CUBLASLT_MATMUL_PREF_REDUCTION_SCHEME_MASK, &mask,
                       sizeof(mask));
        std::vector<cublasLtMatmulHeuristicResult_t> results(32);
        int returned = 0;
        st = g_api.get_heuristic(g_lt, plan.opDesc, plan.Adesc, plan.Bdesc,
                                 plan.Ddesc, plan.Ddesc, pref, (int)results.size(),
                                 results.data(), &returned);
        if (st == CUBLAS_STATUS_SUCCESS) {
            for (int i = 0; i < returned; ++i) {
                if (results[i].state != CUBLAS_STATUS_SUCCESS) continue;
                if (results[i].workspaceSize > 128ull * 1024 * 1024) continue;
                Cand c;
                c.algo = results[i].algo;
                c.ws = results[i].workspaceSize;
                cands.push_back(c);
            }
        }
    }

    // ------------------- explicit tile / split-K candidates -----------------
    if (g_api.algo_get_ids != nullptr && g_api.algo_init != nullptr &&
        g_api.algo_cfg_set != nullptr && g_api.algo_check != nullptr) {
        int ids[16];
        int nIds = 0;
        memset(ids, 0, sizeof(ids));
        st = g_api.algo_get_ids(g_lt, CUBLAS_COMPUTE_32F, CUDA_R_32F,
                                CUDA_R_16F, CUDA_R_16F, CUDA_R_16F, CUDA_R_16F,
                                16, ids, &nIds);
        if (st == CUBLAS_STATUS_SUCCESS && nIds > 0) {
            if (nIds > 8) nIds = 8;
            const cublasLtMatmulTile_t tiles[] = {
                CUBLASLT_MATMUL_TILE_256x128, CUBLASLT_MATMUL_TILE_128x128,
                CUBLASLT_MATMUL_TILE_128x64,  CUBLASLT_MATMUL_TILE_64x128,
                CUBLASLT_MATMUL_TILE_64x64,   CUBLASLT_MATMUL_TILE_256x64,
                CUBLASLT_MATMUL_TILE_128x256};
            const int splitks[] = {3, 2, 4, 6};
            for (size_t ti = 0; ti < sizeof(tiles) / sizeof(tiles[0]); ++ti) {
                for (size_t si = 0; si < sizeof(splitks) / sizeof(splitks[0]);
                     ++si) {
                    for (int j = 0; j < nIds; ++j) {
                        if (cands.size() >= 72) break;
                        cublasLtMatmulAlgo_t algo;
                        memset(&algo, 0, sizeof(algo));
                        if (g_api.algo_init(g_lt, CUBLAS_COMPUTE_32F,
                                            CUDA_R_32F, CUDA_R_16F, CUDA_R_16F,
                                            CUDA_R_16F, CUDA_R_16F, ids[j],
                                            &algo) != CUBLAS_STATUS_SUCCESS) {
                            continue;
                        }
                        cublasLtMatmulTile_t tile = tiles[ti];
                        int32_t sk = splitks[si];
                        g_api.algo_cfg_set(&algo, CUBLASLT_ALGO_CONFIG_TILE_ID,
                                           &tile, sizeof(tile));
                        g_api.algo_cfg_set(&algo,
                                           CUBLASLT_ALGO_CONFIG_SPLITK_NUM,
                                           &sk, sizeof(sk));
                        cublasLtReductionScheme_t rs =
                            CUBLASLT_REDUCTION_SCHEME_INPLACE;
                        g_api.algo_cfg_set(
                            &algo, CUBLASLT_ALGO_CONFIG_REDUCTION_SCHEME, &rs,
                            sizeof(rs));
                        cublasLtMatmulHeuristicResult_t r;
                        memset(&r, 0, sizeof(r));
                        cublasStatus_t cst = g_api.algo_check(
                            g_lt, plan.opDesc, plan.Adesc, plan.Bdesc,
                            plan.Ddesc, plan.Ddesc, &algo, &r);
                        cudaGetLastError();
                        if (cst != CUBLAS_STATUS_SUCCESS ||
                            r.state != CUBLAS_STATUS_SUCCESS) {
                            continue;
                        }
                        if (r.workspaceSize > 128ull * 1024 * 1024) continue;
                        bool dup = false;
                        for (const auto& c : cands) {
                            if (c.ws == r.workspaceSize &&
                                memcmp(&c.algo, &algo, sizeof(algo)) == 0) {
                                dup = true;
                                break;
                            }
                        }
                        if (dup) continue;
                        Cand c;
                        c.algo = algo;
                        c.ws = r.workspaceSize;
                        cands.push_back(c);
                    }
                }
            }
        }
    }

    // ------------------------------ tuning ----------------------------------
    size_t need = 0;
    for (const auto& c : cands) need = std::max(need, c.ws);
    void* ws = ensure_buffer(&g_ws, &g_ws_cap, need);
    const size_t out_bytes = (size_t)m * (size_t)n * 2;
    void* ref = ensure_buffer(&g_refbuf, &g_refbuf_cap, out_bytes);
    if (g_stats == nullptr) {
        if (cudaMalloc(&g_stats, 2 * sizeof(float)) != cudaSuccess) {
            cudaGetLastError();
            g_stats = nullptr;
        }
    }

    cudaEvent_t e0 = nullptr, e1 = nullptr;
    bool events_ok = (cudaEventCreate(&e0) == cudaSuccess) &&
                     (cudaEventCreate(&e1) == cudaSuccess);

    auto run_one = [&](const cublasLtMatmulAlgo_t* algo, void* wsp,
                       size_t wsz) -> bool {
        cublasStatus_t s = g_api.matmul(
            g_lt, plan.opDesc, &alpha, A, plan.Adesc, B, plan.Bdesc, &beta, D,
            plan.Ddesc, D, plan.Ddesc, algo, wsp, wsz, stream);
        if (s != CUBLAS_STATUS_SUCCESS) {
            cudaGetLastError();
            return false;
        }
        return true;
    };

    auto time_cand = [&](const cublasLtMatmulAlgo_t* algo, void* wsp,
                         size_t wsz, float* avg) -> bool {
        if (!run_one(algo, wsp, wsz)) return false;
        if (!events_ok) {
            *avg = 1e30f;
            return true;
        }
        const int reps = 4;
        cudaEventRecord(e0, stream);
        for (int r = 0; r < reps; ++r) {
            g_api.matmul(g_lt, plan.opDesc, &alpha, A, plan.Adesc, B,
                         plan.Bdesc, &beta, D, plan.Ddesc, D, plan.Ddesc, algo,
                         wsp, wsz, stream);
        }
        cudaEventRecord(e1, stream);
        cudaEventSynchronize(e1);
        float ms = 0.f;
        cudaEventElapsedTime(&ms, e0, e1);
        *avg = ms / (float)reps;
        return true;
    };

    auto validate = [&](const cublasLtMatmulAlgo_t* algo, void* wsp,
                        size_t wsz) -> bool {
        if (ref == nullptr || g_stats == nullptr) return true;  // can't check
        if (!run_one(algo, wsp, wsz)) return false;
        cudaMemsetAsync(g_stats, 0, 2 * sizeof(float), stream);
        launch_diff_stats(D, ref, (long)m * (long)n, g_stats, stream);
        float h[2] = {0.f, 0.f};
        cudaMemcpyAsync(h, g_stats, 2 * sizeof(float), cudaMemcpyDeviceToHost,
                        stream);
        cudaStreamSynchronize(stream);
        float maxdiff = h[0], maxref = h[1];
        if (!std::isfinite(maxdiff) || !std::isfinite(maxref)) return false;
        return maxdiff <= 0.5f + 0.02f * maxref;
    };

    // Reference output from the default algorithm.
    bool def_ok = run_one(nullptr, nullptr, 0);
    TORCH_CHECK(def_ok, "default cublasLtMatmul failed");
    if (ref != nullptr) {
        cudaMemcpyAsync(ref, D, out_bytes, cudaMemcpyDeviceToDevice, stream);
        cudaStreamSynchronize(stream);
    }

    float best_ms = 1e30f;
    float def_ms = 1e30f;
    if (time_cand(nullptr, nullptr, 0, &def_ms)) {
        best_ms = def_ms;
    }
    int best = -1;
    for (size_t i = 0; i < cands.size(); ++i) {
        auto& c = cands[i];
        if (c.ws > 0 && (ws == nullptr || c.ws > g_ws_cap)) continue;
        void* wsp = c.ws ? ws : nullptr;
        float ms = 0.f;
        if (!time_cand(&c.algo, wsp, c.ws, &ms)) continue;
        if (ms >= best_ms * 0.995f) continue;  // must beat current best
        if (!validate(&c.algo, wsp, c.ws)) continue;
        best_ms = ms;
        best = (int)i;
    }
    if (events_ok) {
        cudaEventDestroy(e0);
        cudaEventDestroy(e1);
    }
    if (pref != nullptr) {
        // no destroy symbol resolved; preference is tiny and process-lifetime
    }

    if (best >= 0 && best_ms <= def_ms * 0.995f) {
        plan.algo = cands[best].algo;
        plan.workspace_size = cands[best].ws;
        plan.tuned = true;
    }

    auto res = g_plans.emplace(key, plan);
    return res.first->second;
}

}  // namespace

bool init_cublaslt(std::string path) {
    std::lock_guard<std::mutex> lock(g_init_mtx);
    if (g_resolved) {
        return g_ok;
    }
    g_resolved = true;
    g_ok = false;

    if (dlsym(RTLD_DEFAULT, "cublasLtCreate") == nullptr) {
        if (!path.empty()) {
            g_dlhandle = dlopen(path.c_str(), RTLD_NOW | RTLD_GLOBAL);
        }
        if (g_dlhandle == nullptr) {
            const char* fallbacks[] = {"libcublasLt.so.13",
                                       "libcublasLt.so.12",
                                       "libcublasLt.so.11",
                                       "libcublasLt.so"};
            for (const char* f : fallbacks) {
                g_dlhandle = dlopen(f, RTLD_NOW | RTLD_GLOBAL);
                if (g_dlhandle != nullptr) break;
            }
        }
    }

    bool ok = true;
    auto need = [&](void** dst, const char* name) {
        *dst = get_sym(name);
        if (*dst == nullptr) ok = false;
    };
    auto want = [&](void** dst, const char* name) {
        *dst = get_sym(name);
    };
    need(reinterpret_cast<void**>(&g_api.create), "cublasLtCreate");
    need(reinterpret_cast<void**>(&g_api.desc_create),
         "cublasLtMatmulDescCreate");
    need(reinterpret_cast<void**>(&g_api.desc_set),
         "cublasLtMatmulDescSetAttribute");
    need(reinterpret_cast<void**>(&g_api.layout_create),
         "cublasLtMatrixLayoutCreate");
    need(reinterpret_cast<void**>(&g_api.pref_create),
         "cublasLtMatmulPreferenceCreate");
    need(reinterpret_cast<void**>(&g_api.pref_set),
         "cublasLtMatmulPreferenceSetAttribute");
    need(reinterpret_cast<void**>(&g_api.get_heuristic),
         "cublasLtMatmulAlgoGetHeuristic");
    need(reinterpret_cast<void**>(&g_api.matmul), "cublasLtMatmul");
    if (!ok) {
        return false;
    }
    want(reinterpret_cast<void**>(&g_api.algo_get_ids),
         "cublasLtMatmulAlgoGetIds");
    want(reinterpret_cast<void**>(&g_api.algo_init), "cublasLtMatmulAlgoInit");
    want(reinterpret_cast<void**>(&g_api.algo_cfg_set),
         "cublasLtMatmulAlgoConfigSetAttribute");
    want(reinterpret_cast<void**>(&g_api.algo_check),
         "cublasLtMatmulAlgoCheck");
    if (g_api.create(&g_lt) != CUBLAS_STATUS_SUCCESS || g_lt == nullptr) {
        return false;
    }
    g_ok = true;
    return true;
}

torch::Tensor fused_linear_min_sub(torch::Tensor x, torch::Tensor w_half,
                                   torch::Tensor bias,
                                   torch::Tensor constant) {
    TORCH_CHECK(g_ok, "cuBLASLt is not available");
    TORCH_CHECK(x.is_cuda() && w_half.is_cuda() && bias.is_cuda() &&
                    constant.is_cuda(),
                "CUDA tensors required");
    TORCH_CHECK(x.scalar_type() == at::kFloat, "x must be fp32");
    TORCH_CHECK(w_half.scalar_type() == at::kHalf, "w_half must be fp16");
    TORCH_CHECK(bias.scalar_type() == at::kFloat, "bias must be fp32");
    TORCH_CHECK(constant.scalar_type() == at::kFloat,
                "constant must be fp32");
    TORCH_CHECK(x.dim() == 2, "x must be 2D");
    TORCH_CHECK(constant.numel() == 1, "constant must be scalar");

    const c10::cuda::OptionalCUDAGuard device_guard(x.device());
    auto xc = x.contiguous();
    auto wc = w_half.contiguous();
    auto bc = bias.contiguous();
    auto cc = constant.contiguous();

    const long m = xc.size(0);
    const long k = xc.size(1);
    const long n = wc.size(0);
    TORCH_CHECK(wc.dim() == 2 && wc.size(1) == k, "weight shape mismatch");
    TORCH_CHECK(bc.numel() == n, "bias shape mismatch");

    cudaStream_t stream = c10::cuda::getCurrentCUDAStream().stream();

    auto out = torch::empty({m, n}, xc.options());
    const long nx = m * k;
    const long ny = m * n;
    if (ny == 0) {
        return out;
    }

    const long x_pad = (nx + 127) & ~127L;  // keep y 256B aligned
    auto buf = torch::empty({x_pad + ny}, xc.options().dtype(at::kHalf));
    void* xh = buf.data_ptr();
    void* yh = reinterpret_cast<char*>(buf.data_ptr()) + x_pad * 2;

    launch_f32_to_f16(xc.data_ptr<float>(), xh, nx, stream);

    const Plan& plan = get_plan(m, n, k, wc.data_ptr(), xh, yh, stream);
    const float alpha = 1.f, beta = 0.f;
    cublasStatus_t st;
    if (plan.tuned) {
        void* ws = ensure_buffer(&g_ws, &g_ws_cap, plan.workspace_size);
        st = g_api.matmul(g_lt, plan.opDesc, &alpha, wc.data_ptr(), plan.Adesc,
                          xh, plan.Bdesc, &beta, yh, plan.Ddesc, yh,
                          plan.Ddesc, &plan.algo, ws, plan.workspace_size,
                          stream);
    } else {
        st = g_api.matmul(g_lt, plan.opDesc, &alpha, wc.data_ptr(), plan.Adesc,
                          xh, plan.Bdesc, &beta, yh, plan.Ddesc, yh,
                          plan.Ddesc, nullptr, nullptr, 0, stream);
    }
    TORCH_CHECK(st == CUBLAS_STATUS_SUCCESS, "cublasLtMatmul failed: ",
                (int)st);

    launch_bias_min0(yh, bc.data_ptr<float>(), cc.data_ptr<float>(),
                     out.data_ptr<float>(), ny, n, stream);
    return out;
}
"""

_ext = load_inline(
    name="fused_linear_min_sub_v5",
    cpp_sources=cpp_src,
    cuda_sources=cuda_src,
    functions=["init_cublaslt", "fused_linear_min_sub", "bias_min0_epilogue"],
    verbose=False,
    extra_cflags=["-O3"],
    extra_cuda_cflags=["-O3"],
    extra_ldflags=["-ldl"],
)


def _find_cublaslt():
    here = os.path.dirname(torch.__file__)
    sp = os.path.dirname(here)
    patterns = (
        os.path.join(sp, "nvidia", "cublas", "lib", "libcublasLt.so*"),
        os.path.join(sp, "nvidia", "cu13", "lib", "libcublasLt.so*"),
        os.path.join(sp, "nvidia", "cu13", "lib64", "libcublasLt.so*"),
        os.path.join(sp, "nvidia", "cu12", "lib", "libcublasLt.so*"),
        os.path.join