#if HAVE_CUBLASLT_HEADERS
typedef cublasStatus_t (*cublasLtCreate_fn)(cublasLtHandle_t *);
typedef cublasStatus_t (*cublasLtDestroy_fn)(cublasLtHandle_t);
typedef cublasStatus_t (*cublasLtMatmulDescCreate_fn)(cublasLtMatmulDesc_t *,
                                                      cublasComputeType_t,
                                                      cudaDataType_t);
typedef cublasStatus_t (*cublasLtMatmulDescDestroy_fn)(cublasLtMatmulDesc_t);
typedef cublasStatus_t (*cublasLtMatmulDescSetAttribute_fn)(cublasLtMatmulDesc_t,
                                                            cublasLtMatmulDescAttributes_t,
                                                            const void *,
                                                            size_t);
typedef cublasStatus_t (*cublasLtMatrixLayoutCreate_fn)(cublasLtMatrixLayout_t *,
                                                        cudaDataType_t,
                                                        uint64_t,
                                                        uint64_t,
                                                        int64_t);
typedef cublasStatus_t (*cublasLtMatrixLayoutDestroy_fn)(cublasLtMatrixLayout_t);
typedef cublasStatus_t (*cublasLtMatmulPreferenceCreate_fn)(cublasLtMatmulPreference_t *);
typedef cublasStatus_t (*cublasLtMatmulPreferenceDestroy_fn)(cublasLtMatmulPreference_t);
typedef cublasStatus_t (*cublasLtMatmulPreferenceSetAttribute_fn)(cublasLtMatmulPreference_t,
                                                                  cublasLtMatmulPreferenceAttributes_t,
                                                                  const void *,
                                                                  size_t);
typedef cublasStatus_t (*cublasLtMatmulAlgoGetHeuristic_fn)(
    cublasLtHandle_t,
    cublasLtMatmulDesc_t,
    cublasLtMatrixLayout_t,
    cublasLtMatrixLayout_t,
    cublasLtMatrixLayout_t,
    cublasLtMatrixLayout_t,
    cublasLtMatmulPreference_t,
    int,
    cublasLtMatmulHeuristicResult_t *,
    int *);
typedef cublasStatus_t (*cublasLtMatmul_fn)(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);

struct cublaslt_api {
    void *lib;
    cublasLtCreate_fn cublasLtCreate;
    cublasLtDestroy_fn cublasLtDestroy;
    cublasLtMatmulDescCreate_fn cublasLtMatmulDescCreate;
    cublasLtMatmulDescDestroy_fn cublasLtMatmulDescDestroy;
    cublasLtMatmulDescSetAttribute_fn cublasLtMatmulDescSetAttribute;
    cublasLtMatrixLayoutCreate_fn cublasLtMatrixLayoutCreate;
    cublasLtMatrixLayoutDestroy_fn cublasLtMatrixLayoutDestroy;
    cublasLtMatmulPreferenceCreate_fn cublasLtMatmulPreferenceCreate;
    cublasLtMatmulPreferenceDestroy_fn cublasLtMatmulPreferenceDestroy;
    cublasLtMatmulPreferenceSetAttribute_fn cublasLtMatmulPreferenceSetAttribute;
    cublasLtMatmulAlgoGetHeuristic_fn cublasLtMatmulAlgoGetHeuristic;
    cublasLtMatmul_fn cublasLtMatmul;
};

struct profile_desc {
    const char *name;
    const char *block_label;
    int min_cc;
    int enabled;
    int needs_scalar_scale;
    int needs_block_scale;
    int min_multiple;
    cudaDataType_t a_type;
    cudaDataType_t b_type;
    cudaDataType_t c_type;
    cudaDataType_t d_type;
    cublasComputeType_t compute_type;
};

struct prepared_profile {
    struct profile_desc desc;
    CUstream stream;
    cublasLtMatmulDesc_t op_desc;
    cublasLtMatrixLayout_t a_layout;
    cublasLtMatrixLayout_t b_layout;
    cublasLtMatrixLayout_t c_layout;
    cublasLtMatrixLayout_t d_layout;
    cublasLtMatmulPreference_t preference;
    cublasLtMatmulHeuristicResult_t heuristic;
    CUdeviceptr a_dev;
    CUdeviceptr b_dev;
    CUdeviceptr c_dev;
    CUdeviceptr d_dev;
    CUdeviceptr a_scale_dev;
    CUdeviceptr b_scale_dev;
    CUdeviceptr workspace_dev;
    size_t workspace_size;
    uint64_t m;
    uint64_t n;
    uint64_t k;
    unsigned long iterations;
    int ready;
};

static const struct profile_desc k_profiles[] = {
    {
        "fp64",
        "fp64",
        80,
        1,
        0,
        0,
        8,
        CUDA_R_64F,
        CUDA_R_64F,
        CUDA_R_64F,
        CUDA_R_64F,
        CUBLAS_COMPUTE_64F,
    },
    {
        "fp32_tf32",
        "fp32",
        80,
        1,
        0,
        0,
        128,
        CUDA_R_32F,
        CUDA_R_32F,
        CUDA_R_32F,
        CUDA_R_32F,
        CUBLAS_COMPUTE_32F_FAST_TF32,
    },
    {
        "fp16_tensor",
        "fp16",
        80,
        1,
        0,
        0,
        128,
        CUDA_R_16F,
        CUDA_R_16F,
        CUDA_R_16F,
        CUDA_R_16F,
        CUBLAS_COMPUTE_32F_FAST_16F,
    },
    {
        "int8_tensor",
        "int8",
        75,
        1,
        0,
        0,
        128,
        CUDA_R_8I,
        CUDA_R_8I,
        CUDA_R_32I,
        CUDA_R_32I,
        CUBLAS_COMPUTE_32I,
    },
    {
        "fp8_e4m3",
        "fp8",
        89,
        1,
        1,
        0,
        128,
        CUDA_R_8F_E4M3,
        CUDA_R_8F_E4M3,
        CUDA_R_16BF,
        CUDA_R_16BF,
        CUBLAS_COMPUTE_32F,
    },
    {
        "fp8_e5m2",
        "fp8",
        89,
        1,
        1,
        0,
        128,
        CUDA_R_8F_E5M2,
        CUDA_R_8F_E5M2,
        CUDA_R_16BF,
        CUDA_R_16BF,
        CUBLAS_COMPUTE_32F,
    },
#if defined(CUDA_R_4F_E2M1) && defined(CUBLASLT_MATMUL_MATRIX_SCALE_VEC16_UE4M3)
    {
        "fp4_e2m1",
        "fp4",
        100,
        1,
        0,
        1,
        128,
        CUDA_R_4F_E2M1,
        CUDA_R_4F_E2M1,
        CUDA_R_16BF,
        CUDA_R_16BF,
        CUBLAS_COMPUTE_32F,
    },
#endif
};

#define PROFILE_COUNT ((int)(sizeof(k_profiles) / sizeof(k_profiles[0])))

static int profile_allowed_for_run(const struct profile_desc *desc, int cc, const char *precision_filter) {
    if (!(desc->enabled && cc >= desc->min_cc)) {
        return 0;
    }
    if (precision_filter != NULL) {
        return strcmp(desc->block_label, precision_filter) == 0;
    }
    /* Mixed/all phases intentionally exclude fp64/fp4 for now: both paths are
     * unstable on the current benchmark fleet and can abort the whole mixed
     * pass after earlier phases already collected useful telemetry. */
    return strcmp(desc->block_label, "fp64") != 0 && strcmp(desc->block_label, "fp4") != 0;
}

static int load_cublaslt(struct cublaslt_api *api) {
    memset(api, 0, sizeof(*api));
    api->lib = dlopen("libcublasLt.so.13", RTLD_NOW | RTLD_LOCAL);
    if (!api->lib) {
        api->lib = dlopen("libcublasLt.so", RTLD_NOW | RTLD_LOCAL);
    }
    if (!api->lib) {
        return 0;
    }
    return
        load_symbol(api->lib, "cublasLtCreate", (void **)&api->cublasLtCreate) &&
        load_symbol(api->lib, "cublasLtDestroy", (void **)&api->cublasLtDestroy) &&
        load_symbol(api->lib, "cublasLtMatmulDescCreate", (void **)&api->cublasLtMatmulDescCreate) &&
        load_symbol(api->lib, "cublasLtMatmulDescDestroy", (void **)&api->cublasLtMatmulDescDestroy) &&
        load_symbol(api->lib,
                    "cublasLtMatmulDescSetAttribute",
                    (void **)&api->cublasLtMatmulDescSetAttribute) &&
        load_symbol(api->lib, "cublasLtMatrixLayoutCreate", (void **)&api->cublasLtMatrixLayoutCreate) &&
        load_symbol(api->lib, "cublasLtMatrixLayoutDestroy", (void **)&api->cublasLtMatrixLayoutDestroy) &&
        load_symbol(api->lib,
                    "cublasLtMatmulPreferenceCreate",
                    (void **)&api->cublasLtMatmulPreferenceCreate) &&
        load_symbol(api->lib,
                    "cublasLtMatmulPreferenceDestroy",
                    (void **)&api->cublasLtMatmulPreferenceDestroy) &&
        load_symbol(api->lib,
                    "cublasLtMatmulPreferenceSetAttribute",
                    (void **)&api->cublasLtMatmulPreferenceSetAttribute) &&
        load_symbol(api->lib,
                    "cublasLtMatmulAlgoGetHeuristic",
                    (void **)&api->cublasLtMatmulAlgoGetHeuristic) &&
        load_symbol(api->lib, "cublasLtMatmul", (void **)&api->cublasLtMatmul);
}

static const char *cublas_status_text(cublasStatus_t status) {
    switch (status) {
        case CUBLAS_STATUS_SUCCESS:
            return "CUBLAS_STATUS_SUCCESS";
        case CUBLAS_STATUS_NOT_INITIALIZED:
            return "CUBLAS_STATUS_NOT_INITIALIZED";
        case CUBLAS_STATUS_ALLOC_FAILED:
            return "CUBLAS_STATUS_ALLOC_FAILED";
        case CUBLAS_STATUS_INVALID_VALUE:
            return "CUBLAS_STATUS_INVALID_VALUE";
        case CUBLAS_STATUS_ARCH_MISMATCH:
            return "CUBLAS_STATUS_ARCH_MISMATCH";
        case CUBLAS_STATUS_MAPPING_ERROR:
            return "CUBLAS_STATUS_MAPPING_ERROR";
        case CUBLAS_STATUS_EXECUTION_FAILED:
            return "CUBLAS_STATUS_EXECUTION_FAILED";
        case CUBLAS_STATUS_INTERNAL_ERROR:
            return "CUBLAS_STATUS_INTERNAL_ERROR";
        case CUBLAS_STATUS_NOT_SUPPORTED:
            return "CUBLAS_STATUS_NOT_SUPPORTED";
        default:
            return "CUBLAS_STATUS_UNKNOWN";
    }
}

static int check_cublas(const char *step, cublasStatus_t status) {
    if (status == CUBLAS_STATUS_SUCCESS) {
        return 1;
    }
    fprintf(stderr, "%s failed: %s (%d)\n", step, cublas_status_text(status), (int)status);
    return 0;
}

static size_t bytes_for_elements(cudaDataType_t type, uint64_t elements) {
    switch (type) {
        case CUDA_R_32F:
        case CUDA_R_32I:
            return (size_t)(elements * 4u);
        case CUDA_R_16F:
        case CUDA_R_16BF:
            return (size_t)(elements * 2u);
        case CUDA_R_8I:
        case CUDA_R_8F_E4M3:
        case CUDA_R_8F_E5M2:
            return (size_t)(elements);
#if defined(CUDA_R_4F_E2M1)
        case CUDA_R_4F_E2M1:
            return (size_t)((elements + 1u) / 2u);
#endif
        default:
            return (size_t)(elements * 4u);
    }
}

static cudaDataType_t matmul_scale_type(const struct profile_desc *desc) {
    if (desc->compute_type == CUBLAS_COMPUTE_32I) {
        return CUDA_R_32I;
    }
    if (desc->compute_type == CUBLAS_COMPUTE_64F) {
        return CUDA_R_64F;
    }
    return CUDA_R_32F;
}

static size_t fp4_scale_bytes(uint64_t rows, uint64_t cols) {
    uint64_t row_tiles = (rows + 127u) / 128u;
    uint64_t col_tiles = (cols + 63u) / 64u;
    return (size_t)(row_tiles * col_tiles * 128u);
}

static uint64_t choose_square_dim(size_t budget_bytes, size_t bytes_per_cell, int multiple) {
    double approx = sqrt((double)budget_bytes / (double)bytes_per_cell);
    uint64_t dim = (uint64_t)approx;
    if (dim < (uint64_t)multiple) {
        dim = (uint64_t)multiple;
    }
    dim = (uint64_t)round_down_size((size_t)dim, (size_t)multiple);
    if (dim < (uint64_t)multiple) {
        dim = (uint64_t)multiple;
    }
    if (dim > 65536u) {
        dim = 65536u;
    }
    return dim;
}

static int device_upload(struct cuda_api *cuda, CUdeviceptr dev, const void *src, size_t bytes) {
    return check_rc(cuda, "cuMemcpyHtoD", cuda->cuMemcpyHtoD(dev, src, bytes));
}

static int alloc_filled(struct cuda_api *cuda, CUdeviceptr *ptr, size_t bytes, unsigned char pattern) {
    if (!check_rc(cuda, "cuMemAlloc", cuda->cuMemAlloc(ptr, bytes))) {
        return 0;
    }
    if (!check_rc(cuda, "cuMemsetD8", cuda->cuMemsetD8(*ptr, pattern, bytes))) {
        cuda->cuMemFree(*ptr);
        *ptr = 0;
        return 0;
    }
    return 1;
}

static size_t profile_scale_bytes(const struct profile_desc *desc, uint64_t m, uint64_t n, uint64_t k) {
    size_t bytes = 0;
    if (desc->needs_scalar_scale) {
        bytes += 2u * sizeof(float);
    }
#if defined(CUBLASLT_MATMUL_MATRIX_SCALE_VEC16_UE4M3)
    if (desc->needs_block_scale) {
        bytes += fp4_scale_bytes(k, m);
        bytes += fp4_scale_bytes(k, n);
    }
#else
    (void)m;
    (void)n;
    (void)k;
#endif
    return bytes;
}

static void destroy_profile(struct cublaslt_api *cublas, struct cuda_api *cuda, struct prepared_profile *profile) {
    if (profile->workspace_dev) {
        cuda->cuMemFree(profile->workspace_dev);
    }
    if (profile->a_scale_dev) {
        cuda->cuMemFree(profile->a_scale_dev);
    }
    if (profile->b_scale_dev) {
        cuda->cuMemFree(profile->b_scale_dev);
    }
    if (profile->d_dev) {
        cuda->cuMemFree(profile->d_dev);
    }
    if (profile->c_dev) {
        cuda->cuMemFree(profile->c_dev);
    }
    if (profile->b_dev) {
        cuda->cuMemFree(profile->b_dev);
    }
    if (profile->a_dev) {
        cuda->cuMemFree(profile->a_dev);
    }
    if (profile->preference) {
        cublas->cublasLtMatmulPreferenceDestroy(profile->preference);
    }
    if (profile->d_layout) {
        cublas->cublasLtMatrixLayoutDestroy(profile->d_layout);
    }
    if (profile->c_layout) {
        cublas->cublasLtMatrixLayoutDestroy(profile->c_layout);
    }
    if (profile->b_layout) {
        cublas->cublasLtMatrixLayoutDestroy(profile->b_layout);
    }
    if (profile->a_layout) {
        cublas->cublasLtMatrixLayoutDestroy(profile->a_layout);
    }
    if (profile->op_desc) {
        cublas->cublasLtMatmulDescDestroy(profile->op_desc);
    }
    memset(profile, 0, sizeof(*profile));
}

static int prepare_profile(struct cublaslt_api *cublas,
                           cublasLtHandle_t handle,
                           struct cuda_api *cuda,
                           const struct profile_desc *desc,
                           CUstream stream,
                           size_t profile_budget_bytes,
                           struct prepared_profile *out) {
    size_t bytes_per_cell = 0;
    size_t attempt_budget = profile_budget_bytes;

    bytes_per_cell += bytes_for_elements(desc->a_type, 1);
    bytes_per_cell += bytes_for_elements(desc->b_type, 1);
    bytes_per_cell += bytes_for_elements(desc->c_type, 1);
    bytes_per_cell += bytes_for_elements(desc->d_type, 1);
    if (bytes_per_cell == 0) {
        return 0;
    }

    while (attempt_budget >= MIN_PROFILE_BUDGET_BYTES) {
        memset(out, 0, sizeof(*out));
        out->desc = *desc;
        out->stream = stream;

        uint64_t dim = choose_square_dim(attempt_budget, bytes_per_cell, desc->min_multiple);
        out->m = dim;
        out->n = dim;
        out->k = dim;

        size_t desired_workspace = attempt_budget / 8u;
        if (desired_workspace > 32u * 1024u * 1024u) {
            desired_workspace = 32u * 1024u * 1024u;
        }
        desired_workspace = round_down_size(desired_workspace, 256u);

        size_t a_bytes = 0;
        size_t b_bytes = 0;
        size_t c_bytes = 0;
        size_t d_bytes = 0;
        size_t scale_bytes = 0;
        while (1) {
            a_bytes = bytes_for_elements(desc->a_type, out->k * out->m);
            b_bytes = bytes_for_elements(desc->b_type, out->k * out->n);
            c_bytes = bytes_for_elements(desc->c_type, out->m * out->n);
            d_bytes = bytes_for_elements(desc->d_type, out->m * out->n);
            scale_bytes = profile_scale_bytes(desc, out->m, out->n, out->k);

            size_t matrix_bytes = a_bytes + b_bytes + c_bytes + d_bytes + scale_bytes;
            if (matrix_bytes <= attempt_budget) {
                size_t remaining = attempt_budget - matrix_bytes;
                out->workspace_size = desired_workspace;
                if (out->workspace_size > remaining) {
                    out->workspace_size = round_down_size(remaining, 256u);
                }
                break;
            }

            if (out->m <= (uint64_t)desc->min_multiple) {
                break;
            }
            out->m -= (uint64_t)desc->min_multiple;
            out->n = out->m;
            out->k = out->m;
        }
        if (out->m < (uint64_t)desc->min_multiple) {
            attempt_budget /= 2u;
            continue;
        }

        if (!alloc_filled(cuda, &out->a_dev, a_bytes, 0x11) ||
            !alloc_filled(cuda, &out->b_dev, b_bytes, 0x11) ||
            !alloc_filled(cuda, &out->c_dev, c_bytes, 0x00) ||
            !alloc_filled(cuda, &out->d_dev, d_bytes, 0x00)) {
            destroy_profile(cublas, cuda, out);
            return 0;
        }

        cudaDataType_t scale_type = matmul_scale_type(desc);
        if (!check_cublas("cublasLtMatmulDescCreate",
                          cublas->cublasLtMatmulDescCreate(&out->op_desc, desc->compute_type, scale_type))) {
            destroy_profile(cublas, cuda, out);
            return 0;
        }

        cublasOperation_t transa = CUBLAS_OP_T;
        cublasOperation_t transb = CUBLAS_OP_N;
        if (!check_cublas("set TRANSA",
                          cublas->cublasLtMatmulDescSetAttribute(out->op_desc,
                                                                 CUBLASLT_MATMUL_DESC_TRANSA,
                                                                 &transa,
                                                                 sizeof(transa))) ||
            !check_cublas("set TRANSB",
                          cublas->cublasLtMatmulDescSetAttribute(out->op_desc,
                                                                 CUBLASLT_MATMUL_DESC_TRANSB,
                                                                 &transb,
                                                                 sizeof(transb)))) {
            destroy_profile(cublas, cuda, out);
            return 0;
        }

        if (desc->needs_scalar_scale) {
            float one = 1.0f;
            if (!alloc_filled(cuda, &out->a_scale_dev, sizeof(one), 0x00) ||
                !alloc_filled(cuda, &out->b_scale_dev, sizeof(one), 0x00)) {
                destroy_profile(cublas, cuda, out);
                return 0;
            }
            if (!device_upload(cuda, out->a_scale_dev, &one, sizeof(one)) ||
                !device_upload(cuda, out->b_scale_dev, &one, sizeof(one))) {
                destroy_profile(cublas, cuda, out);
                return 0;
            }
            void *a_scale_ptr = (void *)(uintptr_t)out->a_scale_dev;
            void *b_scale_ptr = (void *)(uintptr_t)out->b_scale_dev;
            if (!check_cublas("set A scale ptr",
                              cublas->cublasLtMatmulDescSetAttribute(out->op_desc,
                                                                     CUBLASLT_MATMUL_DESC_A_SCALE_POINTER,
                                                                     &a_scale_ptr,
                                                                     sizeof(a_scale_ptr))) ||
                !check_cublas("set B scale ptr",
                              cublas->cublasLtMatmulDescSetAttribute(out->op_desc,
                                                                     CUBLASLT_MATMUL_DESC_B_SCALE_POINTER,
                                                                     &b_scale_ptr,
                                                                     sizeof(b_scale_ptr)))) {
                destroy_profile(cublas, cuda, out);
                return 0;
            }
        }

#if defined(CUBLASLT_MATMUL_MATRIX_SCALE_VEC16_UE4M3)
    if (desc->needs_block_scale) {
        size_t a_scale_bytes = fp4_scale_bytes(out->k, out->m);
        size_t b_scale_bytes = fp4_scale_bytes(out->k, out->n);
        if (!alloc_filled(cuda, &out->a_scale_dev, a_scale_bytes, 0x11) ||
            !alloc_filled(cuda, &out->b_scale_dev, b_scale_bytes, 0x11)) {
            destroy_profile(cublas, cuda, out);
            return 0;
        }
        cublasLtMatmulMatrixScale_t scale_mode = CUBLASLT_MATMUL_MATRIX_SCALE_VEC16_UE4M3;
        void *a_scale_ptr = (void *)(uintptr_t)out->a_scale_dev;
        void *b_scale_ptr = (void *)(uintptr_t)out->b_scale_dev;
        if (!check_cublas("set A scale mode",
                          cublas->cublasLtMatmulDescSetAttribute(out->op_desc,
                                                                 CUBLASLT_MATMUL_DESC_A_SCALE_MODE,
                                                                 &scale_mode,
                                                                 sizeof(scale_mode))) ||
            !check_cublas("set B scale mode",
                          cublas->cublasLtMatmulDescSetAttribute(out->op_desc,
                                                                 CUBLASLT_MATMUL_DESC_B_SCALE_MODE,
                                                                 &scale_mode,
                                                                 sizeof(scale_mode))) ||
            !check_cublas("set A block scale ptr",
                          cublas->cublasLtMatmulDescSetAttribute(out->op_desc,
                                                                 CUBLASLT_MATMUL_DESC_A_SCALE_POINTER,
                                                                 &a_scale_ptr,
                                                                 sizeof(a_scale_ptr))) ||
            !check_cublas("set B block scale ptr",
                          cublas->cublasLtMatmulDescSetAttribute(out->op_desc,
                                                                 CUBLASLT_MATMUL_DESC_B_SCALE_POINTER,
                                                                 &b_scale_ptr,
                                                                 sizeof(b_scale_ptr)))) {
            destroy_profile(cublas, cuda, out);
            return 0;
        }
    }
#endif

        if (!check_cublas("create A layout",
                          cublas->cublasLtMatrixLayoutCreate(&out->a_layout, desc->a_type, out->k, out->m, out->k)) ||
            !check_cublas("create B layout",
                          cublas->cublasLtMatrixLayoutCreate(&out->b_layout, desc->b_type, out->k, out->n, out->k)) ||
            !check_cublas("create C layout",
                          cublas->cublasLtMatrixLayoutCreate(&out->c_layout, desc->c_type, out->m, out->n, out->m)) ||
            !check_cublas("create D layout",
                          cublas->cublasLtMatrixLayoutCreate(&out->d_layout, desc->d_type, out->m, out->n, out->m))) {
            destroy_profile(cublas, cuda, out);
            return 0;
        }

        if (!check_cublas("create preference", cublas->cublasLtMatmulPreferenceCreate(&out->preference))) {
            destroy_profile(cublas, cuda, out);
            return 0;
        }

        if (out->workspace_size > 0) {
            if (!alloc_filled(cuda, &out->workspace_dev, out->workspace_size, 0x00)) {
                destroy_profile(cublas, cuda, out);
                return 0;
            }
        }

        if (!check_cublas("set workspace",
                          cublas->cublasLtMatmulPreferenceSetAttribute(
                              out->preference,
                              CUBLASLT_MATMUL_PREF_MAX_WORKSPACE_BYTES,
                              &out->workspace_size,
                              sizeof(out->workspace_size)))) {
            destroy_profile(cublas, cuda, out);
            return 0;
        }

        int found = 0;
        if (check_cublas("heuristic",
                         cublas->cublasLtMatmulAlgoGetHeuristic(handle,
                                                                out->op_desc,
                                                                out->a_layout,
                                                                out->b_layout,
                                                                out->c_layout,
                                                                out->d_layout,
                                                                out->preference,
                                                                1,
                                                                &out->heuristic,
                                                                &found)) &&
            found > 0) {
            out->ready = 1;
            return 1;
        }

        destroy_profile(cublas, cuda, out);
        attempt_budget = round_down_size(attempt_budget * 3u / 4u, 256u);
        if (attempt_budget < MIN_PROFILE_BUDGET_BYTES) {
            break;
        }
    }

    return 0;
}

static int run_cublas_profile(cublasLtHandle_t handle,
                              struct cublaslt_api *cublas,
                              struct prepared_profile *profile) {
    int32_t alpha_i32 = 1;
    int32_t beta_i32 = 0;
    double alpha_f64 = 1.0;
    double beta_f64 = 0.0;
    float alpha = 1.0f;
    float beta = 0.0f;
    const void *alpha_ptr = &alpha;
    const void *beta_ptr = &beta;
    if (profile->desc.compute_type == CUBLAS_COMPUTE_32I) {
        alpha_ptr = &alpha_i32;
        beta_ptr = &beta_i32;
    } else if (profile->desc.compute_type == CUBLAS_COMPUTE_64F) {
        alpha_ptr = &alpha_f64;
        beta_ptr = &beta_f64;
    }
    return check_cublas(profile->desc.name,
                        cublas->cublasLtMatmul(handle,
                                               profile->op_desc,
                                               alpha_ptr,
                                               (const void *)(uintptr_t)profile->a_dev,
                                               profile->a_layout,
                                               (const void *)(uintptr_t)profile->b_dev,
                                               profile->b_layout,
                                               beta_ptr,
                                               (const void *)(uintptr_t)profile->c_dev,
                                               profile->c_layout,
                                               (void *)(uintptr_t)profile->d_dev,
                                               profile->d_layout,
                                               &profile->heuristic.algo,
                                               (void *)(uintptr_t)profile->workspace_dev,
                                               profile->workspace_size,
                                               profile->stream));
}

static int run_cublaslt_stress(struct cuda_api *cuda,
                               CUdevice dev,
                               const char *device_name,
                               int cc_major,
                               int cc_minor,
                               int seconds,
                               int size_mb,
                               const char *precision_filter,
                               struct stress_report *report) {
    struct cublaslt_api cublas;
    struct prepared_profile prepared[MAX_STRESS_STREAMS * PROFILE_COUNT];
    cublasLtHandle_t handle = NULL;
    CUcontext ctx = NULL;
    CUstream streams[MAX_STRESS_STREAMS] = {0};
    uint16_t sample[256];
    int cc = cc_major * 10 + cc_minor;
    int planned = 0;
    int active = 0;
    int mp_count = 0;
    int stream_count = 1;
    int profile_count = PROFILE_COUNT;
    int prepared_count = 0;
    size_t requested_budget = 0;
    size_t total_budget = 0;
    size_t per_profile_budget = 0;
    int budget_profiles = 0;

    memset(report, 0, sizeof(*report));
    snprintf(report->backend, sizeof(report->backend), "cublasLt");
    snprintf(report->device, sizeof(report->device), "%s", device_name);
    report->cc_major = cc_major;
    report->cc_minor = cc_minor;
    report->buffer_mb = size_mb;

    if (!load_cublaslt(&cublas)) {
        snprintf(report->details, sizeof(report->details), "cublasLt=unavailable\n");
        return 0;
    }
    if (!check_rc(cuda, "cuCtxCreate", cuda->cuCtxCreate(&ctx, 0, dev))) {
        return 0;
    }
    if (!check_cublas("cublasLtCreate", cublas.cublasLtCreate(&handle))) {
        cuda->cuCtxDestroy(ctx);
        return 0;
    }

    /* Count profiles matching the filter (for deciding what to run). */
    for (size_t i = 0; i < sizeof(k_profiles) / sizeof(k_profiles[0]); i++) {
        if (profile_allowed_for_run(&k_profiles[i], cc, precision_filter)) {
            planned++;
        }
    }
    if (planned <= 0) {
        snprintf(report->details, sizeof(report->details), "cublasLt_profiles=unsupported\n");
        cublas.cublasLtDestroy(handle);
        cuda->cuCtxDestroy(ctx);
        return 0;
    }

    /* Count all profiles active on this GPU regardless of filter.
     * Mixed phases still divide budget across the full precision set, while
     * single-precision benchmark phases dedicate budget only to active
     * profiles matching precision_filter. */
    int planned_total = 0;
    for (size_t i = 0; i < sizeof(k_profiles) / sizeof(k_profiles[0]); i++) {
        if (profile_allowed_for_run(&k_profiles[i], cc, precision_filter)) {
            planned_total++;
        }
    }
    if (planned_total < planned) {
        planned_total = planned;
    }
    budget_profiles = planned_total;
    if (precision_filter != NULL) {
        budget_profiles = planned;
    }
    if (budget_profiles <= 0) {
        budget_profiles = planned_total;
    }

    requested_budget = (size_t)size_mb * 1024u * 1024u;
    if (requested_budget < (size_t)budget_profiles * MIN_PROFILE_BUDGET_BYTES) {
        requested_budget = (size_t)budget_profiles * MIN_PROFILE_BUDGET_BYTES;
    }
    total_budget = clamp_budget_to_free_memory(cuda, requested_budget);
    if (total_budget < (size_t)budget_profiles * MIN_PROFILE_BUDGET_BYTES) {
        total_budget = (size_t)budget_profiles * MIN_PROFILE_BUDGET_BYTES;
    }
    if (query_multiprocessor_count(cuda, dev, &mp_count) &&
        cuda->cuStreamCreate &&
        cuda->cuStreamDestroy) {
        stream_count = choose_stream_count(mp_count, budget_profiles, total_budget, 1);
    }
    if (precision_filter != NULL && stream_count > MAX_SINGLE_PRECISION_STREAMS) {
        stream_count = MAX_SINGLE_PRECISION_STREAMS;
    }
    if (stream_count > 1) {
        int created = 0;
        for (; created < stream_count; created++) {
            if (!check_rc(cuda, "cuStreamCreate", cuda->cuStreamCreate(&streams[created], 0))) {
                destroy_streams(cuda, streams, created);
                stream_count = 1;
                break;
            }
        }
    }
    report->stream_count = stream_count;
    per_profile_budget = total_budget / ((size_t)budget_profiles * (size_t)stream_count);
    if (per_profile_budget < MIN_PROFILE_BUDGET_BYTES) {
        per_profile_budget = MIN_PROFILE_BUDGET_BYTES;
    }
    if (precision_filter != NULL) {
        per_profile_budget = clamp_single_precision_profile_budget(per_profile_budget);
    }
    report->buffer_mb = (int)(total_budget / (1024u * 1024u));
    append_detail(report->details,
                  sizeof(report->details),
                  "requested_mb=%d actual_mb=%d streams=%d mp_count=%d budget_profiles=%d per_worker_mb=%zu\n",
                  size_mb,
                  report->buffer_mb,
                  report->stream_count,
                  mp_count,
                  budget_profiles,
                  per_profile_budget / (1024u * 1024u));

    for (int i = 0; i < profile_count; i++) {
        const struct profile_desc *desc = &k_profiles[i];
        if (!(desc->enabled && cc >= desc->min_cc)) {
            append_detail(report->details,
                          sizeof(report->details),
                          "%s=SKIPPED cc<%d\n",
                          desc->name,
                          desc->min_cc);
            continue;
        }
        if (!profile_allowed_for_run(desc, cc, precision_filter)) {
            append_detail(report->details,
                          sizeof(report->details),
                          "%s=SKIPPED benchmark_disabled\n",
                          desc->name);
            continue;
        }
        for (int lane = 0; lane < stream_count; lane++) {
            CUstream stream = streams[lane];
            if (prepared_count >= (int)(sizeof(prepared) / sizeof(prepared[0]))) {
                break;
            }
            if (prepare_profile(&cublas, handle, cuda, desc, stream, per_profile_budget, &prepared[prepared_count])) {
                active++;
                append_detail(report->details,
                              sizeof(report->details),
                              "%s[%d]=READY dim=%llux%llux%llu block=%s stream=%d\n",
                              desc->name,
                              lane,
                              (unsigned long long)prepared[prepared_count].m,
                              (unsigned long long)prepared[prepared_count].n,
                              (unsigned long long)prepared[prepared_count].k,
                              desc->block_label,
                              lane);
                prepared_count++;
            } else {
                append_detail(report->details,
                              sizeof(report->details),
                              "%s[%d]=SKIPPED unsupported\n",
                              desc->name,
                              lane);
            }
        }
    }

    if (active <= 0) {
        cublas.cublasLtDestroy(handle);
        destroy_streams(cuda, streams, stream_count);
        cuda->cuCtxDestroy(ctx);
        return 0;
    }

    /* Keep the GPU queue continuously full by submitting kernels without
     * synchronizing after every wave.  A sync barrier after each small batch
     * creates CPU-to-GPU ping-pong gaps that prevent full TDP utilisation,
     * especially when individual kernels are short.  Instead we sync at most
     * once per second (for error detection) and once at the very end. */
    double deadline = now_seconds() + (double)seconds;
    double next_sync = now_seconds() + 1.0;
    while (now_seconds() < deadline) {
        int launched = 0;
        for (int i = 0; i < prepared_count; i++) {
            if (!prepared[i].ready) {
                continue;
            }
            if (!run_cublas_profile(handle, &cublas, &prepared[i])) {
                append_detail(report->details,
                              sizeof(report->details),
                              "%s=FAILED runtime\n",
                              prepared[i].desc.name);
                for (int j = 0; j < prepared_count; j++) {
                    destroy_profile(&cublas, cuda, &prepared[j]);
                }
                cublas.cublasLtDestroy(handle);
                destroy_streams(cuda, streams, stream_count);
                cuda->cuCtxDestroy(ctx);
                return 0;
            }
            prepared[i].iterations++;
            report->iterations++;
            launched++;
        }
        if (launched <= 0) {
            break;
        }
        double now = now_seconds();
        if (now >= next_sync || now >= deadline) {
            if (!check_rc(cuda, "cuCtxSynchronize", cuda->cuCtxSynchronize())) {
                for (int i = 0; i < prepared_count; i++) {
                    destroy_profile(&cublas, cuda, &prepared[i]);
                }
                cublas.cublasLtDestroy(handle);
                destroy_streams(cuda, streams, stream_count);
                cuda->cuCtxDestroy(ctx);
                return 0;
            }
            next_sync = now + 1.0;
        }
    }
    /* Final drain: ensure all queued work finishes before we read results. */
    cuda->cuCtxSynchronize();

    for (int i = 0; i < prepared_count; i++) {
        if (!prepared[i].ready) {
            continue;
        }
        append_detail(report->details,
                      sizeof(report->details),
                      "%s_iterations=%lu\n",
                      prepared[i].desc.name,
                      prepared[i].iterations);
    }

    for (int i = 0; i < prepared_count; i++) {
        if (prepared[i].ready) {
            if (check_rc(cuda, "cuMemcpyDtoH", cuda->cuMemcpyDtoH(sample, prepared[i].d_dev, sizeof(sample)))) {
                for (size_t j = 0; j < sizeof(sample) / sizeof(sample[0]); j++) {
                    report->checksum += sample[j];
                }
            }
            break;
        }
    }

    for (int i = 0; i < prepared_count; i++) {
        destroy_profile(&cublas, cuda, &prepared[i]);
    }
    cublas.cublasLtDestroy(handle);
    destroy_streams(cuda, streams, stream_count);
    cuda->cuCtxDestroy(ctx);
    return 1;
}
#endif