//
//  ConvFpAIntBExecution.cpp
//  MNN
//
//  Created by MNN on 2024/03/11.
//  Copyright © 2018, Alibaba Group Holding Limited
//

// precision = low/normal 时使用运行时离线反量化，precision = high 时使用在线反量化
#include "ConvFpAIntBExecution.hpp"
#include "../Raster.cuh"
#include "../ConvBaseKernel.cuh"
#include <float.h>
#include <cublas_v2.h>

//#define DEBUG

namespace MNN {
namespace CUDA {

const int TILE_DIM = 16; // GEMM_FpAIntB
const int GEMV_TILE = 64; // GEMV_FpAIntB (legacy, kept for INT8)
const int GEMV_TILE_V2 = 256; // P1: Optimized GEMV tile size
const int GEMV_OC_PER_BLOCK = 4; // P1: Output channels per block for GEMV
const int SUB_GEMM_TILE_DIM = 16; // GEMM_Int
const int BLOCK_SIZE = 256; // Quant / SumBq / Bias
const int DEQUANT_TILE_DIM = 32; // Dequant
const int WARP_SIZE = 32;

template<typename T, typename dT>
__global__ void DequantizeInt8Weight(
    const int8_t* quantized_kernel,
    dT* dequantized_kernel,
    const T* scale,
    const T* offset,
    const int oc, const int ic, const int ic_p, const int quanC
) {
    // 每个线程负责反量化一个权重值
    const int col = blockIdx.x * blockDim.x + threadIdx.x; // input channel
    const int row = blockIdx.y * blockDim.y + threadIdx.y; // output channel

    if (row < oc && col < ic) {
        // 计算量化参数的索引
        const int num_quan_groups_per_channel = (quanC > 0) ? (quanC / oc) : 1;
        const int ic_per_group = (num_quan_groups_per_channel > 0) ? (ic / num_quan_groups_per_channel) : ic;
        const int group_idx = col / ic_per_group;
        const int quan_param_index = row * num_quan_groups_per_channel + group_idx;

        const float x_scale  = (float)scale[quan_param_index];
        const float x_offset = (float)offset[quan_param_index];

        // 读取量化的权重
        const float quantized_val = quantized_kernel[row * ic_p + col];

        // 反量化并写入目标 Tensor
        dequantized_kernel[row * ic_p + col] = (dT)(quantized_val * x_scale + x_offset);
    }
}

// Optimized INT4 weight dequantization: each thread unpacks 8 bytes → 16 FP16 values
// Uses uint2 vectorized load for weights and float4 (8 half) vectorized writes
template<typename T, typename dT>
__global__ void DequantizeInt4Weight(
    const uint8_t* __restrict__ packed_weight,
    dT* __restrict__ output,
    const T* __restrict__ scale,
    const T* __restrict__ offset,
    const int oc, const int ic, const int ic_p, const int quanC
) {
    // Each thread handles 8 packed bytes = 16 int4 values
    const int tid = blockIdx.y * blockDim.x + threadIdx.x;
    const int row = blockIdx.x;  // OC index

    if (row >= oc) return;

    const int half_ic = ic_p / 2;  // bytes per row
    const int col_byte = tid * 8;  // starting byte within row
    if (col_byte >= half_ic) return;

    const int num_qg = (quanC > 0) ? (quanC / oc) : 1;
    const int ic_per_group = max(1, (num_qg > 0) ? (ic / num_qg) : ic);
    const int qparam_base = row * num_qg;

    // Vectorized 8-byte weight load (two uint32s)
    const uint32_t* src = reinterpret_cast<const uint32_t*>(packed_weight + row * half_ic + col_byte);
    const uint32_t packed0 = __ldg(src);
    const uint32_t packed1 = (col_byte + 4 < half_ic) ? __ldg(src + 1) : 0;

    const int col = col_byte * 2;  // starting ic index
    const int out_base = row * ic_p + col;

    // Process 8 bytes = 16 int4 values, write as half2 pairs
    half result[16];
    #pragma unroll
    for (int j = 0; j < 4; j++) {
        const int kk = col + j * 2;
        if (kk >= ic) break;
        const int gidx = min(kk / ic_per_group, num_qg - 1);
        const float s = (float)__ldg(&scale[qparam_base + gidx]);
        const float o = (float)__ldg(&offset[qparam_base + gidx]);
        const uint8_t byte_val = (packed0 >> (j * 8)) & 0xFF;
        result[j * 2] = (half)(((float)(byte_val >> 4) - 8.0f) * s + o);
        result[j * 2 + 1] = (half)(((float)(byte_val & 0x0F) - 8.0f) * s + o);
    }
    #pragma unroll
    for (int j = 0; j < 4; j++) {
        const int kk = col + 8 + j * 2;
        if (kk >= ic) break;
        const int gidx = min(kk / ic_per_group, num_qg - 1);
        const float s = (float)__ldg(&scale[qparam_base + gidx]);
        const float o = (float)__ldg(&offset[qparam_base + gidx]);
        const uint8_t byte_val = (packed1 >> (j * 8)) & 0xFF;
        result[8 + j * 2] = (half)(((float)(byte_val >> 4) - 8.0f) * s + o);
        result[8 + j * 2 + 1] = (half)(((float)(byte_val & 0x0F) - 8.0f) * s + o);
    }

    // Vectorized write: 16 halves = 32 bytes = float4 + float4
    if (col + 8 <= ic) {
        *reinterpret_cast<float4*>(output + out_base) = *reinterpret_cast<float4*>(result);
    } else {
        for (int i = 0; i < 8 && col + i < ic; i++) output[out_base + i] = (dT)result[i];
    }
    if (col + 16 <= ic) {
        *reinterpret_cast<float4*>(output + out_base + 8) = *reinterpret_cast<float4*>(result + 8);
    } else if (col + 8 < ic) {
        for (int i = 8; i < 16 && col + i < ic; i++) output[out_base + i] = (dT)result[i];
    }
}

// R6: Pre-compute GEMV params: scale and adj_offset as float2 array
// Eliminates scattered offset loads and runtime computation in GEMV
// Grid: (ceil(total/256)), Block: 256
template<typename T>
__global__ void PrecomputeGemvParams(
    const T* __restrict__ scale,
    const T* __restrict__ offset,
    float2* __restrict__ params,   // [oc * num_qg], .x=scale, .y=adj_offset
    const int total                // oc * num_qg
) {
    const int idx = blockIdx.x * blockDim.x + threadIdx.x;
    if (idx >= total) return;
    const float s = (float)scale[idx];
    const float o = (float)offset[idx];
    params[idx] = make_float2(s, o - 8.0f * s);  // adj_offset = offset - 8*scale
}

// 动态量化 A、计算 sum_A
template<typename T>
__global__ void QuantA(
    const T* A_sub_fp, int8_t* A_sub_q,
    T* scale_A_out, T* offset_A_out, int32_t* sum_A_q_out,
    const int M, const int K_i, const int lda
) {
    // 使用 union 复用共享内存，减少总占用量
    __shared__ union {
        float min_max_vals[2][BLOCK_SIZE / WARP_SIZE]; // 用于 min/max 的 warp 间规约
        int32_t sum_vals[BLOCK_SIZE / WARP_SIZE];      // 用于 sum 的 warp 间规约
        float scale_offset[2];                        // 用于向块内所有线程广播 scale 和 offset
    } smem;

    const int m = blockIdx.x;
    if (m >= M) return;

    const int tid = threadIdx.x;
    const int warp_id = tid / WARP_SIZE;
    const int lane_id = tid % WARP_SIZE;

    float my_min = FLT_MAX;
    float my_max = -FLT_MAX;

    for (int k = tid; k < K_i; k += BLOCK_SIZE) {
        float val = (float)A_sub_fp[m * lda + k];
        my_min = min(my_min, val);
        my_max = max(my_max, val);
    }

    // Warp 内规约：使用 __shfl_down_sync 在 warp 内无锁计算 min/max
    for (int offset = WARP_SIZE / 2; offset > 0; offset >>= 1) {
        my_min = min(my_min, __shfl_down_sync(0xFFFFFFFF, my_min, offset));
        my_max = max(my_max, __shfl_down_sync(0xFFFFFFFF, my_max, offset));
    }

    // Warp 间规约：每个 warp 的 0 号线程将 warp 结果写入共享内存
    if (lane_id == 0) {
        smem.min_max_vals[0][warp_id] = my_min;
        smem.min_max_vals[1][warp_id] = my_max;
    }
    __syncthreads();

    if (warp_id == 0) {
        // 从共享内存加载各 warp 的结果
        if (lane_id < (BLOCK_SIZE / WARP_SIZE)) {
            my_min = smem.min_max_vals[0][lane_id];
            my_max = smem.min_max_vals[1][lane_id];
        } else {
            my_min = FLT_MAX;
            my_max = -FLT_MAX;
        }

        // 在第一个 warp 内部完成最终规约
        for (int offset = (BLOCK_SIZE / WARP_SIZE) / 2; offset > 0; offset >>= 1) {
            my_min = min(my_min, __shfl_down_sync(0xFFFFFFFF, my_min, offset));
            my_max = max(my_max, __shfl_down_sync(0xFFFFFFFF, my_max, offset));
        }
    }
    // 线程 0 计算 scale/offset 并写入共享内存进行广播
    if (tid == 0) {
        float scale = (my_max - my_min) / 255.0f;
        float offset;
        if (abs(scale) > 1e-5) {
            offset = my_max - scale * 127.0f; // my_max 映射到 127, my_min 映射到 -128
        } else {
            scale = 1.0f;
            offset = my_max;
        }

        scale_A_out[m] = (T)scale;
        offset_A_out[m] = (T)offset;
        
        smem.scale_offset[0] = scale;
        smem.scale_offset[1] = offset;
    }
    __syncthreads();

    // 所有线程从共享内存读取 scale/offset, 并进行量化和求和
    const float s_scale_A = smem.scale_offset[0];
    const float s_offset_A = smem.scale_offset[1];

    int32_t my_sum_q = 0;
    for (int k = tid; k < K_i; k += BLOCK_SIZE) {
        float val_fp = (float)A_sub_fp[m * lda + k];
        int32_t val_q = roundf((val_fp - s_offset_A) / s_scale_A);
        int8_t a_q = (int8_t)max(-128, min(127, val_q));
        A_sub_q[m * K_i + k] = a_q;
        my_sum_q += a_q;
    }

    for (int offset = WARP_SIZE / 2; offset > 0; offset >>= 1) {
        my_sum_q += __shfl_down_sync(0xFFFFFFFF, my_sum_q, offset);
    }

    if (lane_id == 0) smem.sum_vals[warp_id] = my_sum_q;
    __syncthreads();

    if (warp_id == 0) {
        my_sum_q = (lane_id < (BLOCK_SIZE / WARP_SIZE)) ? smem.sum_vals[lane_id] : 0;
        for (int offset = (BLOCK_SIZE / WARP_SIZE) / 2; offset > 0; offset >>= 1) {
             my_sum_q += __shfl_down_sync(0xFFFFFFFF, my_sum_q, offset);
        }
    }

    // 线程 0 将最终的和写入全局内存
    if (tid == 0) sum_A_q_out[m] = my_sum_q;
}

__global__ void GEMM_Int8(
    const int8_t* A_q, const int8_t* B_q, int32_t* C_q,
    const int M, const int N, const int K_i,
    const int lda_q, const int ldb, const int ldc
) {
    __shared__ int8_t A_tile_s8[SUB_GEMM_TILE_DIM][SUB_GEMM_TILE_DIM];
    __shared__ int8_t B_tile_s8[SUB_GEMM_TILE_DIM][SUB_GEMM_TILE_DIM];

    const int block_row = blockIdx.y;
    const int block_col = blockIdx.x;
    const int ty = threadIdx.y;
    const int tx = threadIdx.x;

    const int row = block_row * SUB_GEMM_TILE_DIM + ty;
    const int col = block_col * SUB_GEMM_TILE_DIM + tx;

    int32_t acc = 0;
    const int num_k_tiles = UP_DIV(K_i, SUB_GEMM_TILE_DIM);

    for (int k_tile = 0; k_tile < num_k_tiles; ++k_tile) {
        const int k_tile_base = k_tile * SUB_GEMM_TILE_DIM;

        A_tile_s8[ty][tx] = (row < M && (k_tile_base + tx) < K_i) ? A_q[row * lda_q + k_tile_base + tx] : 0;
        B_tile_s8[ty][tx] = (col < N && (k_tile_base + ty) < K_i) ? B_q[col * ldb + k_tile_base + ty] : 0;
        __syncthreads();

        #pragma unroll
        for (int k = 0; k < SUB_GEMM_TILE_DIM; ++k) {
            acc += (int32_t)A_tile_s8[ty][k] * (int32_t)B_tile_s8[k][tx];
        }
        __syncthreads();
    }

    if (row < M && col < N) C_q[row * ldc + col] = acc;
}

template<typename T>
__global__ void DequantAndAcc(
    const int32_t* C_q, T* C_fp_final,
    const T* scale_A_in, const T* offset_A_in,
    const T* base_scale_B, const T* base_offset_B,
    const int32_t* base_sum_B_q,
    const int group_idx, const int num_oc_groups,
    const int32_t* sum_A_q_in,
    const int M, const int N, const int K_i, const int ldc
) {
    // 使用共享内存缓存行和列的量化参数
    __shared__ float smem_scale_A[DEQUANT_TILE_DIM];
    __shared__ float smem_offset_A[DEQUANT_TILE_DIM];
    __shared__ int32_t smem_sum_A_q[DEQUANT_TILE_DIM];

    __shared__ float smem_scale_B[DEQUANT_TILE_DIM];
    __shared__ float smem_offset_B[DEQUANT_TILE_DIM];
    __shared__ int32_t smem_sum_B_q[DEQUANT_TILE_DIM];

    const int tx = threadIdx.x;
    const int ty = threadIdx.y;

    const int block_row_start = blockIdx.y * DEQUANT_TILE_DIM;
    const int block_col_start = blockIdx.x * DEQUANT_TILE_DIM;

    // 每个线程加载一个 A 的参数
    int m_load_idx = block_row_start + ty;
    if (tx == 0 && m_load_idx < M) {
        smem_scale_A[ty]   = (float)scale_A_in[m_load_idx];
        smem_offset_A[ty]  = (float)offset_A_in[m_load_idx];
        smem_sum_A_q[ty]   = sum_A_q_in[m_load_idx];
    }

    // 每个线程加载一个 B 的参数
    int n_load_idx = block_col_start + tx;
    if (ty == 0 && n_load_idx < N) {
        const size_t b_param_idx = n_load_idx * num_oc_groups + group_idx;
        smem_scale_B[tx]  = (float)base_scale_B[b_param_idx];
        smem_offset_B[tx] = (float)base_offset_B[b_param_idx];
        smem_sum_B_q[tx]  = base_sum_B_q[b_param_idx];
    }

    __syncthreads();

    // 计算每个线程负责的输出点
    const int m = block_row_start + ty;
    const int n = block_col_start + tx;

    if (m < M && n < N) {
        // 从共享内存中读取参数
        const float scale_A = smem_scale_A[ty];
        const float offset_A = smem_offset_A[ty];
        const float sum_A_q = (float)smem_sum_A_q[ty];

        const float scale_B = smem_scale_B[tx];
        const float offset_B = smem_offset_B[tx];
        const float sum_B_q = (float)smem_sum_B_q[tx];
        
        const float c_q_val = (float)C_q[m * ldc + n];

        float term1 = scale_A * (c_q_val * scale_B + sum_A_q * offset_B);
        float term2 = offset_A * (sum_B_q * scale_B + K_i * offset_B);
        float final_val = term1 + term2;

        #if (defined(__CUDACC__) && (!defined(__CUDA_ARCH__) || (__CUDA_ARCH__ >= 700))) || defined(_NVHPC_CUDA)
            if constexpr (std::is_same_v<T, float>) {
                atomicAdd(&C_fp_final[m * ldc + n], final_val);
            } else if constexpr (std::is_same_v<T, half>) {
                // printf("[final_val]%.4f %.4f\n", __half2float(C_fp_final[m * ldc + n]), final_val);
                atomicAdd(&C_fp_final[m * ldc + n], __float2half(final_val));
                // printf("[C_fp_final]%.4f\n", __half2float(C_fp_final[m * ldc + n]));
            }
        #else
            C_fp_final[m * ldc + n] += final_val;
        #endif
    }
}

// 预计算权重 B 的修正项 (sum_B_q)
__global__ void Precompute_SumBq(
    const int8_t* B_q, int32_t* sum_B_q_out,
    const int num_groups, const int ic_per_group,
    const int oc, const int ic_p
) {
    // 每个线程处理一个 (output_channel, group) 的修正项
    for (size_t index = blockIdx.x * blockDim.x + threadIdx.x; index < (size_t)oc * num_groups; index += blockDim.x * gridDim.x) {
        
        const int i = index % num_groups; // group_idx
        const int n = index / num_groups; // output_channel
        
        const int k_start = i * ic_per_group;
        
        int32_t sum = 0;
        for (int k_offset = 0; k_offset < ic_per_group; ++k_offset) {
            // 访问 B 矩阵的布局是 [oc][ic_p]
            sum += (int32_t)B_q[n * ic_p + k_start + k_offset];
        }
        sum_B_q_out[index] = sum;
    }
}

template<typename T>
__global__ void BiasAndActivation(
    T* data, const T* bias,
    const float minV, const float maxV,
    const int M, const int N, const int ldc
) {
    for (size_t index = blockIdx.x * blockDim.x + threadIdx.x; index < (size_t)M * N; index += blockDim.x * gridDim.x) {
        const int m = index / N;
        const int n = index % N;

        const size_t buffer_idx = m * ldc + n;
        float val = (float)data[buffer_idx];
        val += (float)bias[n];
        val = max(val, minV);
        val = min(val, maxV);
        data[buffer_idx] = (T)val;
    }
}

template <typename T>
__global__ void print_tensor_kernel(const T* d_tensor, int rows, int cols) {
    if (blockIdx.x == 0 && threadIdx.x == 0) {
        printf("----------- Tensor in Kernel -----------\n");
        printf("Rows: %d, Cols: %d\n", rows, cols);
        for (int i = 0; i < rows; ++i) {
            for (int j = 0; j < cols; ++j) {
                T value = d_tensor[i * cols + j];

                if constexpr (std::is_same_v<T, half>) {
                    printf("%.4f ", __half2float(value));
                } else {
                    printf("%.4f ", value);
                }
            }
            printf("\n");
        }
        printf("---------------------------------------\n");
    }
}

template<typename T>
__global__ void GEMM_FpAInt8B(
    const T* input,
    const int8_t* kernel,
    const T* scale, const T* offset, const T* bias,
    T* output,
    const float maxV, const float minV,
    const int ic, const int ic_p, const int oc, const int oc_p,
    const int batch, const int quanC
) {
    __shared__ T A_tile[TILE_DIM][TILE_DIM]; // [batch, ic]
    __shared__ int8_t B_tile_s8[TILE_DIM][TILE_DIM]; // [ic, oc] kernel 本身是 B^T [oc, ic]
    __shared__ T B_tile_fp[TILE_DIM][TILE_DIM];

    const int tx = threadIdx.x;
    const int ty = threadIdx.y;

    const int block_row = blockIdx.y; // batch
    const int block_col = blockIdx.x; // output channel

    // 每个线程负责计算输出Tile中的一个元素
    const int out_row = block_row * TILE_DIM + ty; // M
    const int out_col = block_col * TILE_DIM + tx; // N

    float acc = 0.0f;
    // 沿K维度（输入通道ic）分块循环
    const int num_k_tiles = UP_DIV(ic, TILE_DIM);

    const int num_quan_groups_per_channel = (quanC > 0) ? (quanC / oc) : 1;
    const int ic_per_group = (num_quan_groups_per_channel > 0) ? (ic / num_quan_groups_per_channel) : ic;

    for (int k_tile = 0; k_tile < num_k_tiles; ++k_tile) {
        const int k_tile_base = k_tile * TILE_DIM;

        // 合并访问加载 A_tile (input)，线程 (ty, tx) 加载 A_tile[ty][tx]
        int a_col_idx = k_tile_base + tx;
        A_tile[ty][tx] = (out_row < batch && a_col_idx < ic) ? input[out_row * ic_p + a_col_idx] : (T)0.0f;
        
        // 合并访问加载 B_tile (kernel)，线程 (ty, tx) 加载 B_tile[ty][tx]
        // kernel 布局为 [oc, ic]，需要 B(k,n)，即 kernel(n,k)
        int b_load_row = block_col * TILE_DIM + ty;
        int b_col_idx = k_tile_base + tx;
        B_tile_s8[ty][tx] = (b_load_row < oc && b_col_idx < ic) ? kernel[b_load_row * ic_p + b_col_idx] : 0;

        __syncthreads();

        // 反量化 + 转置
        const int K = ty;
        const int N = tx;

        const int global_k = k_tile_base + K;
        const int global_n = block_col * TILE_DIM + N;

        if (global_n < oc && global_k < ic) {
            const int group_idx = global_k / ic_per_group;
            const int quan_param_index = global_n * num_quan_groups_per_channel + group_idx;

            const float x_scale  = (float)scale[quan_param_index];
            const float x_offset = (float)offset[quan_param_index];
            
            // B_tile_s8(n,k), thread(n,k) -> (tx, ty)
            // So we need to read from B_tile_s8[n_dim_in_tile][k_dim_in_tile] -> B_tile_s8[tx][ty]
            const float b_quant = (float)B_tile_s8[N][K];

            // B_tile_fp[k][n] -> B_tile_fp[ty][tx]
            B_tile_fp[K][N] = (T)(b_quant * x_scale + x_offset);
        }
        
        __syncthreads();

        // 在 SMEM 中进行子矩阵乘法
        if (out_col < oc) {
            #pragma unroll
            for (int k = 0; k < TILE_DIM; ++k) {
                acc += (float)A_tile[ty][k] * (float)B_tile_fp[k][tx];
            }
        }

        __syncthreads();
    }

    // 写回全局内存
    if (out_row < batch && out_col < oc) {
        acc += (float)bias[out_col];
        acc = max(acc, minV);
        acc = min(acc, maxV);
        output[out_row * oc_p + out_col] = (T)acc;
    }
}

template<typename T>
__global__ void GEMV_FpAInt8B(
    const T* input,
    const int8_t* kernel,
    const T* scale, const T* offset, const T* bias,
    T* output,
    const float maxV, const float minV,
    const int batch, const int ic, const int ic_p,
    const int oc, const int oc_p, const int quanC
) {
    __shared__ float partial_sums[GEMV_TILE]; // 每个线程块计算一个输出通道

    const int oz = blockIdx.x;
    const int b  = blockIdx.y;

    // if (oz >= oc || b >= batch) return;

    const int tid = threadIdx.x; // 块内线程ID
    
    // 准备量化参数
    const int num_quan_groups_per_channel = (quanC > 0) ? (quanC / oc) : 1;
    const int ic_per_group = (num_quan_groups_per_channel > 0) ? (ic / num_quan_groups_per_channel) : ic;
    const int quan_param_index_base = oz * num_quan_groups_per_channel;

    // 每个线程计算自己的部分和
    float my_sum = 0.0f;
    for (int k = tid; k < ic; k += blockDim.x) {
        const float inp0 = input[b * ic_p + k];
        const int8_t ker0 = kernel[oz * ic_p + k];
        
        const int group_idx = k / ic_per_group;
        const int quan_param_index = quan_param_index_base + group_idx;
        const float x_scale  = scale[quan_param_index];
        const float x_offset = offset[quan_param_index];
        
        my_sum += inp0 * ((float)ker0 * x_scale + x_offset);
    }

    partial_sums[tid] = my_sum;
    __syncthreads();

    // 在 SMEM 中执行并行规约（蝶式交换）
    for (unsigned int s = blockDim.x / 2; s > 0; s >>= 1) {
        if (tid < s) partial_sums[tid] += partial_sums[tid + s];
        __syncthreads();
    }

    if (tid == 0) { // 写回全局内存
        float final_val = partial_sums[0] + (float)bias[oz];
        final_val = max(final_val, minV);
        final_val = min(final_val, maxV);
        output[b * oc_p + oz] = (T)final_val;
    }
}

template<typename T>
__global__ void GEMM_FpAInt4B(
    const T* input,
    const uint8_t* kernel, // kernel 是打包的 int4
    const T* scale, const T* offset, const T* bias,
    T* output,
    const float maxV, const float minV,
    const int ic, const int ic_p, const int oc, const int oc_p,
    const int batch, const int quanC
) {
    __shared__ T A_tile[TILE_DIM][TILE_DIM]; // [batch, ic]
    __shared__ uint8_t B_tile_u8[TILE_DIM][TILE_DIM / 2]; // [ic, oc] kernel 本身是 B^T [oc, ic]
    __shared__ T B_tile_fp[TILE_DIM][TILE_DIM];

    const int tx = threadIdx.x;
    const int ty = threadIdx.y;

    const int block_row = blockIdx.y; // batch
    const int block_col = blockIdx.x; // output channel

    const int out_row = block_row * TILE_DIM + ty; // M
    const int out_col = block_col * TILE_DIM + tx; // N

    float acc = 0.0f;
    const int num_k_tiles = UP_DIV(ic, TILE_DIM);
    
    const int num_quan_groups_per_channel = (quanC > 0) ? (quanC / oc) : 1;
    const int ic_per_group = (num_quan_groups_per_channel > 0) ? (ic / num_quan_groups_per_channel) : ic;

    for (int k_tile = 0; k_tile < num_k_tiles; ++k_tile) {
        const int k_tile_base = k_tile * TILE_DIM;

        // 合并访问加载 A_tile (input)
        int a_col_idx = k_tile_base + tx;
        A_tile[ty][tx] = (out_row < batch && a_col_idx < ic) ? input[out_row * ic_p + a_col_idx] : (T)0.0f;
        
        // 合并访问加载 B_tile (kernel)，每个线程加载一个 byte，但 tile 宽度减半
        if (tx < TILE_DIM / 2) {
            int b_load_row = block_col * TILE_DIM + ty;
            // kernel 布局为 [oc, ic/2]，所以列索引要除以2
            int b_col_idx = (k_tile_base / 2) + tx;
            B_tile_u8[ty][tx] = (b_load_row < oc && (k_tile_base + tx*2) < ic) ? kernel[b_load_row * (ic_p / 2) + b_col_idx] : 0;
        }

        __syncthreads();

        // 反量化 + 转置
        const int K = ty;
        const int N = tx;

        const int global_k = k_tile_base + K;
        const int global_n = block_col * TILE_DIM + N;

        if (global_n < oc && global_k < ic) {
            const int group_idx = global_k / ic_per_group;
            const int quan_param_index = global_n * num_quan_groups_per_channel + group_idx;

            const float x_scale  = (float)scale[quan_param_index];
            const float x_offset = (float)offset[quan_param_index];
            
            // B_tile_u8(n,k) -> B_tile_u8[tx][ty/2] (转置后)
            const uint8_t b_packed = B_tile_u8[N][K / 2];
            // 根据 K 的奇偶性解包 int4
            const int8_t b_quant_s8 = (K % 2 == 0) ? ((b_packed >> 4) - 8) : ((b_packed & 0x0F) - 8);
            const float b_quant = (float)b_quant_s8;
            
            // B_tile_fp[k][n] -> B_tile_fp[ty][tx]
            B_tile_fp[K][N] = (T)(b_quant * x_scale + x_offset);
        }
        
        __syncthreads();

        // 在 SMEM 中进行子矩阵乘法
        if (out_col < oc) {
            #pragma unroll
            for (int k = 0; k < TILE_DIM; ++k) {
                acc += (float)A_tile[ty][k] * (float)B_tile_fp[k][tx];
            }
        }

        __syncthreads();
    }

    // 写回全局内存
    if (out_row < batch && out_col < oc) {
        acc += (float)bias[out_col];
        acc = max(acc, minV);
        acc = min(acc, maxV);
        output[out_row * oc_p + out_col] = (T)acc;
    }
}

template<typename T>
__global__ void GEMV_FpAInt4B(
    const T* input,
    const uint8_t* kernel, // kernel 是打包的 int4
    const T* scale, const T* offset, const T* bias,
    T* output,
    const float maxV, const float minV,
    const int batch, const int ic, const int ic_p,
    const int oc, const int oc_p, const int quanC
) {
    extern __shared__ uint8_t smem_buffer[];
    T* smem_input = reinterpret_cast<T*>(smem_buffer);
    float* partial_sums = reinterpret_cast<float*>(smem_buffer + ic_p * sizeof(T));
    
    const int oz = blockIdx.x; // 当前线程块负责计算的输出通道索引
    const int b  = blockIdx.y;

    const int tid = threadIdx.x;
    
    // if (oz >= oc || b >= batch) return;
    
    // 加载 input 到共享内存
    for (int i = tid; i < ic; i += blockDim.x) smem_input[i] = input[b * ic_p + i];
    for (int i = ic + tid; i < ic_p; i += blockDim.x) smem_input[i] = (T)0.0f;
    __syncthreads();

    const int num_quan_groups_per_channel = (quanC > 0) ? (quanC / oc) : 1;
    const int ic_per_group = (num_quan_groups_per_channel > 0) ? (ic / num_quan_groups_per_channel) : ic;
    const int quan_param_index_base = oz * num_quan_groups_per_channel;

    float my_sum = 0.0f;
    for (int k = tid * 2; k < ic; k += blockDim.x * 2) {
        // 加载打包的权重并解包
        const uint8_t ker_packed = kernel[oz * (ic_p / 2) + k / 2];
        const int8_t ker0_s8 = (ker_packed >> 4) - 8;
        const int8_t ker1_s8 = (ker_packed & 0x0F) - 8;

        const int group_idx0 = k / ic_per_group;
        const int quan_param_index0 = quan_param_index_base + group_idx0;
        const float x_scale0  = scale[quan_param_index0];
        const float x_offset0 = offset[quan_param_index0];
        my_sum += (float)smem_input[k] * ((float)ker0_s8 * x_scale0 + x_offset0);
        my_sum += (float)smem_input[k + 1] * ((float)ker1_s8 * x_scale0 + x_offset0);
        
        // if (k + 1 < ic) {
        //     const float inp1 = (float)smem_input[k + 1];
        //     const int group_idx1 = (k + 1) / ic_per_group;

        //     if (group_idx0 == group_idx1) {
        //         my_sum += inp1 * ((float)ker1_s8 * x_scale0 + x_offset0);
        //     } else {
        //         const int quan_param_index1 = quan_param_index_base + group_idx1;
        //         const float x_scale1  = (float)scale[quan_param_index1];
        //         const float x_offset1 = (float)offset[quan_param_index1];
        //         my_sum += inp1 * ((float)ker1_s8 * x_scale1 + x_offset1);
        //     }
        // }
    }

    partial_sums[tid] = my_sum;
    __syncthreads();

    // 并行规约
    for (unsigned int s = blockDim.x / 2; s > 0; s >>= 1) {
        if (tid < s) partial_sums[tid] += partial_sums[tid + s];
        __syncthreads();
    }

    if (tid == 0) {
        float final_val = partial_sums[0] + (float)bias[oz];
        final_val = max(final_val, minV);
        final_val = min(final_val, maxV);
        output[b * oc_p + oz] = (T)final_val;
    }
}


// =====================================================================
// P1-V5: 128-thread GEMV — all 128 threads cooperate on 1 OC
// - Grid: (oc, batch), Block: 128
// =====================================================================
template<typename T>
__global__ void GEMV_FpAInt4B_V5(
    const T* __restrict__ input,
    const uint8_t* __restrict__ kernel,  // packed int4 weights [oc, ic_p/2]
    const T* __restrict__ scale, const T* __restrict__ offset, const T* __restrict__ bias,
    T* __restrict__ output,
    const float maxV, const float minV,
    const int batch, const int ic, const int ic_p,
    const int oc, const int oc_p, const int quanC
) {
    const int oz = blockIdx.x;  // 1 OC per block
    const int b = blockIdx.y;
    if (oz >= oc || b >= batch) return;

    const int tid = threadIdx.x;  // 0..127
    const int warp_id = tid / 32;
    const int lane_id = tid % 32;

    const int num_qg = (quanC > 0) ? (quanC / oc) : 1;
    const int ic_per_group = (num_qg > 0) ? (ic / num_qg) : ic;
    const int qparam_base = oz * num_qg;

    const int row_bytes = ic_p / 2;
    const uint8_t* w_row = kernel + oz * row_bytes;
    const T* in_ptr = input + b * ic_p;

    float acc = 0.0f;

    // 128 threads × 8 elements = 1024 per iteration
    for (int k = tid * 8; k < ic; k += 1024) {
        uint32_t w32 = __ldg(reinterpret_cast<const uint32_t*>(w_row + k / 2));

        const int gidx = k / ic_per_group;
        const float s = (float)__ldg(&scale[qparam_base + gidx]);
        const float o = (float)__ldg(&offset[qparam_base + gidx]);
        const float adj_off = o - 8.0f * s;

        #pragma unroll
        for (int j = 0; j < 4; j++) {
            const int kk = k + j * 2;
            if (kk >= ic) break;
            const uint8_t packed = (w32 >> (j * 8)) & 0xFF;
            const float w0 = (float)(packed >> 4) * s + adj_off;
            const float w1 = (float)(packed & 0x0F) * s + adj_off;
            acc += (float)in_ptr[kk] * w0;
            if (kk + 1 < ic) acc += (float)in_ptr[kk + 1] * w1;
        }
    }

    // Warp-level reduction
    #pragma unroll
    for (int off = 16; off > 0; off >>= 1) {
        acc += __shfl_xor_sync(0xffffffff, acc, off);
    }

    // Cross-warp reduction via shared memory (4 warps → 1 result)
    __shared__ float smem[4];
    if (lane_id == 0) smem[warp_id] = acc;
    __syncthreads();

    if (tid == 0) {
        float result = smem[0] + smem[1] + smem[2] + smem[3];
        result += (float)bias[oz];
        result = fmaxf(result, minV);
        result = fminf(result, maxV);
        output[b * oc_p + oz] = (T)result;
    }
}

// =====================================================================
// V9: Multi-OC GEMV — each block processes OC_PER_BLK output channels
// 128 threads cooperatively reduce each OC in sequence (fewer blocks, better latency hiding)
// Grid: (ceil(oc/OC_PER_BLK), batch), Block: 128
// =====================================================================
template<typename T, int OC_PER_BLK>
__global__ void GEMV_FpAInt4B_V9(
    const T* __restrict__ input,
    const uint8_t* __restrict__ kernel,  // packed int4 weights [oc, ic_p/2]
    const T* __restrict__ scale, const T* __restrict__ offset, const T* __restrict__ bias,
    T* __restrict__ output,
    const float maxV, const float minV,
    const int batch, const int ic, const int ic_p,
    const int oc, const int oc_p, const int quanC
) {
    const int oc_base = blockIdx.x * OC_PER_BLK;
    const int b = blockIdx.y;
    if (oc_base >= oc || b >= batch) return;

    const int tid = threadIdx.x;
    const int warp_id = tid / 32;
    const int lane_id = tid % 32;

    const int num_qg = (quanC > 0) ? (quanC / oc) : 1;
    const int ic_per_group = (num_qg > 0) ? (ic / num_qg) : ic;

    const int row_bytes = ic_p / 2;
    const T* in_ptr = input + b * ic_p;

    // blockDim.x threads, (blockDim.x/32) warps
    const int nwarps = blockDim.x / 32;
    __shared__ float smem[8];  // up to 8 warps (256 threads)

    #pragma unroll
    for (int oc_off = 0; oc_off < OC_PER_BLK; oc_off++) {
        const int oz = oc_base + oc_off;
        if (oz >= oc) break;

        const int qparam_base = oz * num_qg;
        const uint8_t* w_row = kernel + oz * row_bytes;

        float acc = 0.0f;

        // Each thread handles 8 IC elements per iteration, stride = nthreads * 8
        const int stride = blockDim.x * 8;
        for (int k = tid * 8; k < ic; k += stride) {
            uint32_t w32 = __ldg(reinterpret_cast<const uint32_t*>(w_row + k / 2));

            const int gidx = k / ic_per_group;
            const float s = (float)__ldg(&scale[qparam_base + gidx]);
            const float o = (float)__ldg(&offset[qparam_base + gidx]);
            const float adj_off = o - 8.0f * s;

            #pragma unroll
            for (int j = 0; j < 4; j++) {
                const int kk = k + j * 2;
                if (kk >= ic) break;
                const uint8_t packed = (w32 >> (j * 8)) & 0xFF;
                const float w0 = (float)(packed >> 4) * s + adj_off;
                const float w1 = (float)(packed & 0x0F) * s + adj_off;
                acc += (float)in_ptr[kk] * w0;
                if (kk + 1 < ic) acc += (float)in_ptr[kk + 1] * w1;
            }
        }

        // Warp-level reduction
        #pragma unroll
        for (int off = 16; off > 0; off >>= 1)
            acc += __shfl_xor_sync(0xffffffff, acc, off);

        if (lane_id == 0) smem[warp_id] = acc;
        __syncthreads();

        if (tid == 0) {
            float result = 0.0f;
            for (int w = 0; w < nwarps; w++) result += smem[w];
            result += (float)bias[oz];
            result = fmaxf(result, minV);
            result = fminf(result, maxV);
            output[b * oc_p + oz] = (T)result;
        }
        __syncthreads();
    }
}




// =====================================================================
// V14: Factored-dequant multi-OC GEMV — reduces ALU per weight byte.
// Key insight: acc += in * (s * nibble + adj) = s * sum(in*nibble) + adj * sum(in)
// By factoring out scale/adj, we save 6 FMA per OC per 8-IC iteration.
// Also: explicit load-compute separation for better instruction scheduling.
// Grid: (ceil(oc/OC_PER_BLK), batch), Block: 128
// =====================================================================
template<typename T, int OC_PER_BLK>
__global__ void __launch_bounds__(128, 12) GEMV_FpAInt4B_V14(
    const T* __restrict__ input,
    const uint8_t* __restrict__ kernel,  // packed int4 weights [oc, ic_p/2]
    const float2* __restrict__ gemv_params,  // [oc * num_qg], .x=scale, .y=adj_offset
    const T* __restrict__ bias,
    T* __restrict__ output,
    const float maxV, const float minV,
    const int batch, const int ic, const int ic_p,
    const int oc, const int oc_p, const int num_qg
) {
    // Grid layout: (oc/OC_PER_BLK, batch)
    const int oc_base = blockIdx.x * OC_PER_BLK;
    const int b = blockIdx.y;
    if (oc_base >= oc || b >= batch) return;

    const int tid = threadIdx.x;
    const int warp_id = tid / 32;
    const int lane_id = tid % 32;
    const int nwarps = blockDim.x / 32;

    const int ic_per_group = ic / num_qg;
    const int row_bytes = ic_p / 2;
    const T* in_ptr = input + b * ic_p;

    // Compute per-OC base pointers (on-the-fly, no pointer arrays to save registers)
    float accs[OC_PER_BLK];
    #pragma unroll
    for (int i = 0; i < OC_PER_BLK; i++)
        accs[i] = 0.0f;

    // Main IC loop — all OCs processed simultaneously
    const int stride = blockDim.x * 8;
    for (int k = tid * 8; k < ic; k += stride) {
        // Phase 1: Load ALL input values
        float in_vals[8];
        #pragma unroll
        for (int j = 0; j < 8; j++)
            in_vals[j] = (float)in_ptr[k + j];

        // Pre-compute input sum for factored adj_offset term (shared across all OCs)
        float input_sum = in_vals[0] + in_vals[1] + in_vals[2] + in_vals[3]
                        + in_vals[4] + in_vals[5] + in_vals[6] + in_vals[7];

        const int gidx = k / ic_per_group;

        // Phase 2: Load ALL weights and params before any compute
        uint32_t w_data[OC_PER_BLK];
        float s_data[OC_PER_BLK], adj_data[OC_PER_BLK];
        #pragma unroll
        for (int oi = 0; oi < OC_PER_BLK; oi++) {
            const int oc_idx = oc_base + oi;
            if (oc_idx < oc) {
                w_data[oi] = __ldg(reinterpret_cast<const uint32_t*>(
                    kernel + oc_idx * row_bytes + k / 2));
                const float2 p = __ldg(gemv_params + oc_idx * num_qg + gidx);
                s_data[oi] = p.x;
                adj_data[oi] = p.y;
            }
        }

        // Phase 3: Compute raw dot products (no scale/adj in the hot loop)
        #pragma unroll
        for (int oi = 0; oi < OC_PER_BLK; oi++) {
            if (oc_base + oi >= oc) break;
            float raw_dot = 0.0f;
            #pragma unroll
            for (int j = 0; j < 4; j++) {
                const uint32_t packed = (w_data[oi] >> (j * 8)) & 0xFF;
                raw_dot += in_vals[j * 2] * (float)(packed >> 4);
                raw_dot += in_vals[j * 2 + 1] * (float)(packed & 0x0F);
            }
            // Apply scale and adj_offset: acc += s * raw_dot + adj * input_sum
            accs[oi] = fmaf(adj_data[oi], input_sum, fmaf(s_data[oi], raw_dot, accs[oi]));
        }
    }

    // Warp-level reduction for all accumulators
    #pragma unroll
    for (int oi = 0; oi < OC_PER_BLK; oi++) {
        #pragma unroll
        for (int off = 16; off > 0; off >>= 1)
            accs[oi] += __shfl_xor_sync(0xffffffff, accs[oi], off);
    }

    // Single-sync parallel cross-warp reduction for ALL OCs at once
    __shared__ float smem[OC_PER_BLK][8];  // [oc_idx][warp_idx]

    if (lane_id == 0) {
        #pragma unroll
        for (int oi = 0; oi < OC_PER_BLK; oi++)
            smem[oi][warp_id] = accs[oi];
    }
    __syncthreads();

    // First OC_PER_BLK threads each reduce one OC
    if (tid < OC_PER_BLK && oc_base + tid < oc) {
        float result = 0.0f;
        for (int w = 0; w < nwarps; w++) result += smem[tid][w];
        result += (float)__ldg(&bias[oc_base + tid]);
        result = fmaxf(result, minV);
        result = fminf(result, maxV);
        output[b * oc_p + oc_base + tid] = (T)result;
    }
}

// V14_MB: Multi-batch variant — one block processes ALL batch elements for the same OC group.
// Weights are loaded ONCE and reused across batch elements, reducing memory bandwidth by ~batch×.
template<typename T, int OC_PER_BLK, int MAX_BATCH>
__global__ void __launch_bounds__(128, 8) GEMV_FpAInt4B_V14_MB(
    const T* __restrict__ input,
    const uint8_t* __restrict__ kernel,
    const float2* __restrict__ gemv_params,
    const T* __restrict__ bias,
    T* __restrict__ output,
    const float maxV, const float minV,
    const int batch, const int ic, const int ic_p,
    const int oc, const int oc_p, const int num_qg
) {
    const int oc_base = blockIdx.x * OC_PER_BLK;
    if (oc_base >= oc) return;

    const int tid = threadIdx.x;
    const int warp_id = tid / 32;
    const int lane_id = tid % 32;
    const int nwarps = blockDim.x / 32;

    const int ic_per_group = ic / num_qg;
    const int row_bytes = ic_p / 2;

    float accs[MAX_BATCH * OC_PER_BLK];
    #pragma unroll
    for (int i = 0; i < MAX_BATCH * OC_PER_BLK; i++)
        accs[i] = 0.0f;

    const int stride = blockDim.x * 8;
    for (int k = tid * 8; k < ic; k += stride) {
        const int gidx = k / ic_per_group;

        // Load weights ONCE for all batch elements
        uint32_t w_data[OC_PER_BLK];
        float s_data[OC_PER_BLK], adj_data[OC_PER_BLK];
        #pragma unroll
        for (int oi = 0; oi < OC_PER_BLK; oi++) {
            const int oc_idx = oc_base + oi;
            if (oc_idx < oc) {
                w_data[oi] = __ldg(reinterpret_cast<const uint32_t*>(
                    kernel + oc_idx * row_bytes + k / 2));
                const float2 p = __ldg(gemv_params + oc_idx * num_qg + gidx);
                s_data[oi] = p.x;
                adj_data[oi] = p.y;
            }
        }

        // Process each batch element with the cached weights
        #pragma unroll
        for (int b = 0; b < MAX_BATCH; b++) {
            if (b >= batch) break;
            const T* in_ptr = input + b * ic_p;
            float in_vals[8];
            #pragma unroll
            for (int j = 0; j < 8; j++)
                in_vals[j] = (float)in_ptr[k + j];

            float input_sum = in_vals[0] + in_vals[1] + in_vals[2] + in_vals[3]
                            + in_vals[4] + in_vals[5] + in_vals[6] + in_vals[7];

            #pragma unroll
            for (int oi = 0; oi < OC_PER_BLK; oi++) {
                if (oc_base + oi >= oc) break;
                float raw_dot = 0.0f;
                #pragma unroll
                for (int j = 0; j < 4; j++) {
                    const uint32_t packed = (w_data[oi] >> (j * 8)) & 0xFF;
                    raw_dot += in_vals[j * 2] * (float)(packed >> 4);
                    raw_dot += in_vals[j * 2 + 1] * (float)(packed & 0x0F);
                }
                accs[b * OC_PER_BLK + oi] = fmaf(adj_data[oi], input_sum, fmaf(s_data[oi], raw_dot, accs[b * OC_PER_BLK + oi]));
            }
        }
    }

    // Warp-level reduction for all accumulators
    #pragma unroll
    for (int i = 0; i < MAX_BATCH * OC_PER_BLK; i++) {
        #pragma unroll
        for (int off = 16; off > 0; off >>= 1)
            accs[i] += __shfl_xor_sync(0xffffffff, accs[i], off);
    }

    // Cross-warp reduction using shared memory
    __shared__ float smem[MAX_BATCH * OC_PER_BLK][4]; // [batch*oc][warp_id], nwarps<=4

    if (lane_id == 0) {
        #pragma unroll
        for (int i = 0; i < MAX_BATCH * OC_PER_BLK; i++)
            smem[i][warp_id] = accs[i];
    }
    __syncthreads();

    // First batch*OC_PER_BLK threads each reduce one accumulator
    const int reduce_count = batch * OC_PER_BLK;
    if (tid < reduce_count) {
        const int b = tid / OC_PER_BLK;
        const int oi = tid % OC_PER_BLK;
        if (oc_base + oi < oc) {
            float result = 0.0f;
            for (int w = 0; w < nwarps; w++) result += smem[tid][w];
            result += (float)__ldg(&bias[oc_base + oi]);
            result = fmaxf(result, minV);
            result = fminf(result, maxV);
            output[b * oc_p + oc_base + oi] = (T)result;
        }
    }
}










// P1: Optimized INT8 GEMV with warp shuffle reduction
template<typename T>
__global__ void GEMV_FpAInt8B_V2(
    const T* __restrict__ input,
    const int8_t* __restrict__ kernel,
    const T* __restrict__ scale, const T* __restrict__ offset, const T* __restrict__ bias,
    T* __restrict__ output,
    const float maxV, const float minV,
    const int batch, const int ic, const int ic_p,
    const int oc, const int oc_p, const int quanC
) {
    const int b = blockIdx.y;
    const int oz_base = blockIdx.x * GEMV_OC_PER_BLOCK;
    const int warp_id = threadIdx.y;
    const int lane_id = threadIdx.x;
    const int oz = oz_base + warp_id;

    if (oz >= oc || b >= batch) return;

    const int num_quan_groups_per_channel = (quanC > 0) ? (quanC / oc) : 1;
    const int ic_per_group = (num_quan_groups_per_channel > 0) ? (ic / num_quan_groups_per_channel) : ic;
    const int quan_param_index_base = oz * num_quan_groups_per_channel;

    float my_sum = 0.0f;
    for (int k = lane_id; k < ic; k += WARP_SIZE) {
        const float inp0 = (float)input[b * ic_p + k];
        const int8_t ker0 = __ldg(&kernel[oz * ic_p + k]);

        const int group_idx = k / ic_per_group;
        const int quan_param_index = quan_param_index_base + group_idx;
        const float x_scale = (float)__ldg(&scale[quan_param_index]);
        const float x_offset = (float)__ldg(&offset[quan_param_index]);

        my_sum += inp0 * ((float)ker0 * x_scale + x_offset);
    }

    // Warp shuffle reduction
    #pragma unroll
    for (int offset_val = WARP_SIZE / 2; offset_val > 0; offset_val >>= 1) {
        my_sum += __shfl_xor_sync(0xffffffff, my_sum, offset_val);
    }

    if (lane_id == 0) {
        float final_val = my_sum + (float)bias[oz];
        final_val = max(final_val, minV);
        final_val = min(final_val, maxV);
        output[b * oc_p + oz] = (T)final_val;
    }
}

template<typename T>
__global__ void CONV_FpAInt8B(const T* input,
    const int8_t* kernel,
    const T* scale, const T* offset, const T* bias,
    T *output,
    const float maxV, const float minV,
    const int ic,  const int ic_p, const int iw, const int ih,
    const int c, const int c_p, const int ow, const int oh,
    const int kw, const int kh, const int dw, const int dh,
    const int sw, const int sh, const int pw, const int ph,
    const int total, const int quanC,
    DivModFast d_oc, DivModFast d_ow, DivModFast d_oh
) {

    for (size_t index = blockIdx.x * blockDim.x + threadIdx.x; index < total; index += blockDim.x * gridDim.x) {
        int oz_2, tmp2, oy, ox, tmp1, ob;
        d_oc.divmod(index, tmp1, oz_2);
        d_ow.divmod(tmp1, tmp2, ox);
        d_oh.divmod(tmp2, ob, oy);

        int oz = oz_2;
        int ix = ox * sw - pw;
        int iy = oy * sh - ph;
        float color0 = bias[oz];
        
        const int num_quan_groups_per_channel = (c > 0 && quanC > 0) ? (quanC / c) : 1;
        const int ic_per_group = (num_quan_groups_per_channel > 0) ? (ic / num_quan_groups_per_channel) : ic;
        const int quan_param_index_base = oz * num_quan_groups_per_channel;

        int fxSta = max(0, (UP_DIV(-ix, dw)));
        int fySta = max(0, (UP_DIV(-iy, dh)));
        int fxEnd = min(kw, UP_DIV(iw - ix, dw));
        int fyEnd = min(kh, UP_DIV(ih - iy, dh));
        int fx, fy, fz;
        for (fy=fySta; fy<fyEnd; ++fy) {
            int sy = fy*dh + iy;
            for (fx=fxSta; fx<fxEnd; ++fx) {
                int sx = fx*dw + ix;
                for (int group_idx = 0; group_idx < num_quan_groups_per_channel; ++group_idx) {
                    const int quan_param_index = quan_param_index_base + group_idx;
                    const float x_scale = scale[quan_param_index];
                    const float x_offset = offset[quan_param_index];

                    const int sz_start = group_idx * ic_per_group;
                    const int sz_end = sz_start + ic_per_group;

                    for (int sz = sz_start; sz < sz_end && sz < ic_p; ++ sz) {
                        int src_offset = ((ob * ih + sy) * iw + sx) * ic_p + sz;
                        float inp0 = input[src_offset];
                        //[Cop, KhKw, Cip]
                        int8_t ker0 = kernel[((oz * kh + fy) * kw + fx) * ic_p + sz];

                        color0 = color0 + inp0 * ((float)ker0 * x_scale + x_offset);
                    }
                }
            }
        }
        color0 = max(color0, minV);
        color0 = min(color0, maxV);

        int dst_offset = ((ob * oh + oy) * ow + ox) * c_p + oz;

        output[dst_offset] = color0;
    }
}

template<typename T>
__global__ void CONV_FpAInt4B(const T* input,
    const uint8_t* kernel,
    const T* scale, const T* offset, const T* bias,
    T *output,
    const float maxV, const float minV,
    const int ic, const int ic_p, const int iw, const int ih,
    const int c, const int c_p, const int ow, const int oh,
    const int kw, const int kh, const int dw, const int dh,
    const int sw, const int sh, const int pw, const int ph,
    const int total, const int quanC,
    DivModFast d_oc, DivModFast d_ow, DivModFast d_oh
) {

    for (size_t index = blockIdx.x * blockDim.x + threadIdx.x; index < total; index += blockDim.x * gridDim.x) {
        int oz_2, tmp2, oy, ox, tmp1, ob;
        d_oc.divmod(index, tmp1, oz_2);
        d_ow.divmod(tmp1, tmp2, ox);
        d_oh.divmod(tmp2, ob, oy);

        int oz = oz_2;
        int ix = ox * sw - pw;
        int iy = oy * sh - ph;
        float color0 = bias[oz];

        const int num_quan_groups_per_channel = (c > 0 && quanC > 0) ? (quanC / c) : 1;
        const int ic_per_group = (num_quan_groups_per_channel > 0) ? (ic / num_quan_groups_per_channel) : ic;
        const int quan_param_index_base = oz * num_quan_groups_per_channel;

        int fxSta = max(0, (UP_DIV(-ix, dw)));
        int fySta = max(0, (UP_DIV(-iy, dh)));
        int fxEnd = min(kw, UP_DIV(iw - ix, dw));
        int fyEnd = min(kh, UP_DIV(ih - iy, dh));
        int fx, fy, fz;
        for (fy=fySta; fy<fyEnd; ++fy) {
            int sy = fy*dh + iy;
            for (fx=fxSta; fx<fxEnd; ++fx) {
                int sx = fx*dw + ix;
                for (int group_idx = 0; group_idx < num_quan_groups_per_channel; ++group_idx) {
                    const int quan_param_index = quan_param_index_base + group_idx;
                    const float x_scale = scale[quan_param_index];
                    const float x_offset = offset[quan_param_index];

                    const int sz_start = group_idx * ic_per_group / 2;
                    const int sz_end = sz_start + ic_per_group / 2;

                    for (int sz = sz_start; sz < sz_end && sz * 2 < ic_p; ++ sz) {
                        int src_offset = ((ob * ih + sy) * iw + sx) * ic_p + 2 * sz;
                        float inp0 = input[src_offset];
                        float inp1 = input[src_offset+1];

                        //[Cop, KhKw, Cip]
                        uint8_t ker = kernel[((oz * kh + fy) * kw + fx) * ic_p / 2 + sz];
                        int8_t ker0 = (ker >> 4) - 8;
                        int8_t ker1 = (ker & 15) - 8;
                        color0 = color0 + inp0 * ((float)ker0 * x_scale + x_offset);
                        color0 = color0 + inp1 * ((float)ker1 * x_scale + x_offset);
                    }
                }
            }
        }
        color0 = max(color0, minV);
        color0 = min(color0, maxV);

        int dst_offset = ((ob * oh + oy) * ow + ox) * c_p + oz;

        output[dst_offset] = color0;
    }
}

__global__ void Rearrange_Packed_Weight_Int4(const uint8_t* param,
    uint8_t* output,
    const int khw,
    const size_t maxCount,
    const int oc,
    const int ic,
    DivModFast d_khw,
    DivModFast d_icp2
) {
    for (size_t index = blockIdx.x * blockDim.x + threadIdx.x; index < maxCount; index += blockDim.x * gridDim.x) {
        int icp2Index, temp, ocpIndex, khwIndex;
        d_icp2.divmod(index, temp, icp2Index);
        d_khw.divmod(temp, ocpIndex, khwIndex);
        if(ocpIndex >= oc || 2 * icp2Index >= ic) {
            output[index] = 0;
            continue;
        }

        int src_index = (ocpIndex * UP_DIV(ic, 2) + icp2Index) * khw + khwIndex;
        output[index] = param[src_index];
    }
}

__global__ void Rearrange_Weight_Int4(const int8_t* param,
    uint8_t* output,
    const int khw,
    const size_t maxCount,
    const int oc,
    const int ic,
    DivModFast d_khw,
    DivModFast d_icp2
) {
    for (size_t index = blockIdx.x * blockDim.x + threadIdx.x; index < maxCount; index += blockDim.x * gridDim.x) {
        int icp2Index, temp, ocpIndex, khwIndex;
        d_icp2.divmod(index, temp, icp2Index);
        d_khw.divmod(temp, ocpIndex, khwIndex);
        if(2*icp2Index >= ic || ocpIndex >= oc) {
            output[index] = 0;
            continue;
        }
        // [Co, Ci, KhKw] -> [Cop, KhKw, Cip/2], Ci available for vectorize
        output[index] = ((uint8_t)(param[(ocpIndex * ic + 2*icp2Index+0) * khw + khwIndex] + 8 )) * 16;
        if(2*icp2Index+1 < ic) {
            output[index] += ((uint8_t)(param[(ocpIndex * ic + 2*icp2Index+1) * khw + khwIndex] + 8 ));
        }
    }
}

__global__ void Rearrange_Weight_Int8(const int8_t* param,
    int8_t* output,
    const int khw,
    const size_t maxCount,
    const int oc,
    const int ic,
    DivModFast d_khw,
    DivModFast d_icp
) {
    for (size_t index = blockIdx.x * blockDim.x + threadIdx.x; index < maxCount; index += blockDim.x * gridDim.x) {
        int icpIndex, temp, ocpIndex, khwIndex;
        d_icp.divmod(index, temp, icpIndex);
        d_khw.divmod(temp, ocpIndex, khwIndex);
        if(icpIndex >= ic || ocpIndex >= oc) {
            output[index] = 0;
            continue;
        }
        // [Co, Ci, KhKw] -> [Cop, KhKw, Cip], Ci available for vectorize
        output[index] = param[(ocpIndex * ic + icpIndex) * khw + khwIndex];
    }
}

bool ConvFpAIntBExecution::isValid(const Convolution2D* conv, Backend* backend) {
    return true;
}


ConvFpAIntBExecution::Resource::Resource(Backend* bn, const MNN::Op* op) {
    mBackend = bn;
    auto runtime = static_cast<CUDABackend*>(bn)->getCUDARuntime();

    auto conv       = op->main_as_Convolution2D();
    auto common     = conv->common();

    //weight host->device
    std::shared_ptr<ConvolutionCommon::Int8Common> quanCommon = ConvolutionCommon::load(op, mBackend, false, true);

    auto oc = common->outputCount();
    if (quanCommon->asymmetric) {
        mQuanC = quanCommon->alpha.size() / 2;
    } else {
        mQuanC = quanCommon->alpha.size();
    }
    auto weightSize = quanCommon->weight.size();
    int l = weightSize / oc;
    int h = oc;
    int ic = common->inputCount();
    if(ic == 0) {
        ic = l / common->kernelX() / common->kernelY();
    }

    int lp = UP_DIV(l, 8) * 8;
    int hp = UP_DIV(h, 8) * 8;
    int quanHp = UP_DIV(mQuanC, 8) * 8;

    // set dequant scale/offset
    {
        float * dequantAlpha = quanCommon->alpha.get();
        std::vector<float> dequantScale(quanHp, 0.0);
        std::vector<float> dequantOffset(quanHp, 0.0);

        for (int o = 0; o < mQuanC; o++) {
            float min = 0.0f;
            float alpha = 0.0f;
            if (quanCommon->asymmetric) {
                min = dequantAlpha[2*o];
                alpha = dequantAlpha[2*o+1];
            } else {
                alpha = dequantAlpha[o];
            }
            dequantScale[o] = alpha;
            dequantOffset[o] = min;
        }

        if(static_cast<CUDABackend*>(bn)->useFp16()) {
            auto tempScaleStorage = static_cast<CUDABackend*>(bn)->getStaticBufferPool()->alloc(quanHp*sizeof(float));
            auto scaleTemp = (float*)((uint8_t*)tempScaleStorage.first + tempScaleStorage.second);
            cuda_check(cudaMemcpy(scaleTemp, dequantScale.data(), quanHp*sizeof(float), cudaMemcpyHostToDevice));

            scaleTensor.reset(Tensor::createDevice<int16_t>({quanHp}));
            bn->onAcquireBuffer(scaleTensor.get(), Backend::STATIC);
            mScale = (void *)scaleTensor.get()->buffer().device;
            callFloat2Half((const void*)scaleTemp, (void*)mScale, quanHp, runtime);

            // Reuse scaleTemp buffer
            cuda_check(cudaMemcpy(scaleTemp, dequantOffset.data(), quanHp*sizeof(float), cudaMemcpyHostToDevice));

            offsetTensor.reset(Tensor::createDevice<int16_t>({quanHp}));
            bn->onAcquireBuffer(offsetTensor.get(), Backend::STATIC);
            mOffset = (void *)offsetTensor.get()->buffer().device;
            callFloat2Half((const void*)scaleTemp, (void*)mOffset, quanHp, runtime);

            static_cast<CUDABackend*>(bn)->getStaticBufferPool()->free(tempScaleStorage);
        } else {
            scaleTensor.reset(Tensor::createDevice<int32_t>({quanHp}));
            bn->onAcquireBuffer(scaleTensor.get(), Backend::STATIC);
            mScale = (void *)scaleTensor.get()->buffer().device;
            cuda_check(cudaMemcpy(mScale, dequantScale.data(), quanHp*sizeof(float), cudaMemcpyHostToDevice));
            
            offsetTensor.reset(Tensor::createDevice<int32_t>({quanHp}));
            bn->onAcquireBuffer(offsetTensor.get(), Backend::STATIC);
            mOffset = (void *)offsetTensor.get()->buffer().device;
            cuda_check(cudaMemcpy(mOffset, dequantOffset.data(), quanHp*sizeof(float), cudaMemcpyHostToDevice));
        }
    }

    // Reorder weight
    {
        int khw = common->kernelX() * common->kernelY();
        int icp = UP_DIV(ic, 8) * 8;
        DivModFast khwD(khw);
        DivModFast icp2D(icp/2);
        DivModFast icpD(icp);

        if(quanCommon->canUseInt4) {
            mIsWeightInt4 = true;

            auto tempCacheBuffer = static_cast<CUDABackend*>(bn)->getStaticBufferPool()->alloc(weightSize * sizeof(int8_t));
            float* cacheWeight = (float*)((uint8_t*)tempCacheBuffer.first + tempCacheBuffer.second);
            runtime->memcpy(cacheWeight, quanCommon->weight.get(), weightSize * sizeof(int8_t), MNNMemcpyHostToDevice);

            weightTensor.reset(Tensor::createDevice<uint8_t>({khw * icp/2 * hp}));
            bn->onAcquireBuffer(weightTensor.get(), Backend::STATIC);
            mFilter = (void *)weightTensor.get()->buffer().device;

            //[Co, Ci, KhKw] -> [Cop, KhKw, Cip/2]
            int block_num = runtime->blocks_num(khw*icp/2*hp);
            int block_size = runtime->threads_num();
            Rearrange_Packed_Weight_Int4<<<block_num, block_size>>>((const uint8_t*)cacheWeight, (uint8_t*)mFilter, khw, khw*icp/2*hp, oc, ic, khwD, icp2D);
            checkKernelErrors;

            static_cast<CUDABackend*>(bn)->getStaticBufferPool()->free(tempCacheBuffer);

            // R6: Pre-compute GEMV params (scale, adj_offset) as float2 array for 1x1 conv
            if (khw == 1) {
                mNumQg = (mQuanC > 0) ? (mQuanC / oc) : 1;
                const size_t param_count = (size_t)oc * mNumQg;
                cudaError_t err = cudaMalloc(&mGemvParams, param_count * sizeof(float2));
                if (err == cudaSuccess) {
                    dim3 pg_block(256);
                    dim3 pg_grid((param_count + 255) / 256);
                    if (static_cast<CUDABackend*>(bn)->useFp16()) {
                        PrecomputeGemvParams<half><<<pg_grid, pg_block>>>(
                            (const half*)mScale, (const half*)mOffset,
                            mGemvParams, (int)param_count);
                    } else {
                        PrecomputeGemvParams<float><<<pg_grid, pg_block>>>(
                            (const float*)mScale, (const float*)mOffset,
                            mGemvParams, (int)param_count);
                    }
                    checkKernelErrors;
                }
            }
        } else {
            auto tempCacheBuffer = static_cast<CUDABackend*>(bn)->getStaticBufferPool()->alloc(weightSize * sizeof(int8_t));
            float* cacheWeight = (float*)((uint8_t*)tempCacheBuffer.first + tempCacheBuffer.second);
            runtime->memcpy(cacheWeight, quanCommon->weight.get(), weightSize * sizeof(int8_t), MNNMemcpyHostToDevice);

            weightTensor.reset(Tensor::createDevice<int8_t>({khw * icp * hp}));
            bn->onAcquireBuffer(weightTensor.get(), Backend::STATIC);
            mFilter = (void *)weightTensor.get()->buffer().device;

            //[Co, Ci, KhKw] -> [Cop, KhKw, Cip]
            int block_num = runtime->blocks_num(khw*icp*hp);
            int block_size = runtime->threads_num();
            Rearrange_Weight_Int8<<<block_num, block_size>>>((const int8_t*)cacheWeight, (int8_t*)mFilter, khw, khw*icp*hp, oc, ic, khwD, icpD);
            checkKernelErrors;
            static_cast<CUDABackend*>(bn)->getStaticBufferPool()->free(tempCacheBuffer);
        }
    }

    // Copy Bias
    {
        if(static_cast<CUDABackend*>(bn)->useFp16()) {
            int biasSize = conv->bias()->size();
            int hp = UP_DIV(biasSize, 8) * 8;

            auto tempBiasStorage = static_cast<CUDABackend*>(bn)->getStaticBufferPool()->alloc(hp*sizeof(float));
            auto biasTemp = (float*)((uint8_t*)tempBiasStorage.first + tempBiasStorage.second);
            runtime->memset(biasTemp, 0, hp * sizeof(int32_t));
            cuda_check(cudaMemcpy(biasTemp, conv->bias()->data(), conv->bias()->size()*sizeof(float), cudaMemcpyHostToDevice));

            biasTensor.reset(Tensor::createDevice<int16_t>({hp}));
            bn->onAcquireBuffer(biasTensor.get(), Backend::STATIC);
            mBias = (void *)biasTensor.get()->buffer().device;
            callFloat2Half((const void*)biasTemp, (void*)mBias, hp, runtime);

            static_cast<CUDABackend*>(bn)->getStaticBufferPool()->free(tempBiasStorage);
        } else {
            int biasSize = conv->bias()->size();
            int hp = UP_DIV(biasSize, 8) * 8;
            biasTensor.reset(Tensor::createDevice<int32_t>({hp}));
            bn->onAcquireBuffer(biasTensor.get(), Backend::STATIC);
            mBias = (void *)biasTensor.get()->buffer().device;
            runtime->memset(mBias, 0, hp * sizeof(int32_t));
            cuda_check(cudaMemcpy(mBias, conv->bias()->data(), conv->bias()->size()*sizeof(float), cudaMemcpyHostToDevice));
        }
    }

    // High + Int8：预计算 B 的修正项 sum_B_q
    // if(!quanCommon->canUseInt4) {
    //     const int icp = UP_DIV(ic, 8) * 8;
    //     const int num_groups = (mQuanC > 0) ? (mQuanC / oc) : 1;
    //     const int ic_per_group = (num_groups > 0) ? (ic / num_groups) : ic;

    //     const int param_count = oc * num_groups;
    //     mSumBQTensor.reset(Tensor::createDevice<int32_t>({param_count}));
    //     bn->onAcquireBuffer(mSumBQTensor.get(), Backend::STATIC);
    //     mSumBQ = (void*)mSumBQTensor.get()->buffer().device;

    //     const int block_size = BLOCK_SIZE;
    //     const int num_blocks = (param_count + block_size - 1) / block_size;

    //     Precompute_SumBq<<<num_blocks, block_size>>>((const int8_t*)mFilter, (int32_t*)mSumBQ, num_groups, ic_per_group, oc, icp);
    //     checkKernelErrors;
    // }
}

ConvFpAIntBExecution::Resource::~Resource() {
    if (mGemvParams) cudaFree(mGemvParams);
}
ConvFpAIntBExecution::ConvFpAIntBExecution(Backend* backend, const MNN::Op* op, std::shared_ptr<Resource> res) : CutlassConvCommonExecution(backend) {
    mOp = op;
    mResource = res;
    auto runtime = static_cast<CUDABackend*>(backend)->getCUDARuntime();
    mPrecisonLevel = static_cast<CUDABackend*>(backend)->getPrecision();
    mFp16Infer = (mPrecisonLevel == 2);
    mFp32Infer = (mPrecisonLevel == 1);
    mFp16Fp32MixInfer = (mPrecisonLevel == 0);
    mBf16Infer = (mPrecisonLevel == 3);
}

ConvFpAIntBExecution::~ConvFpAIntBExecution() {
    // DYNAMIC dequant buffers are managed by the memory planner (alloc+release in onResize)
    // Only STATIC buffers need explicit release in destructor
    if (mDequantFilterTensor && mDequantIsStatic) {
        backend()->onReleaseBuffer(mDequantFilterTensor.get(), Backend::STATIC);
    }
}
bool ConvFpAIntBExecution::onClone(Backend* bn, const Op* op, Execution** dst) {
    if (!mValid) {
        return false;
    }
    if (nullptr == dst) {
        return true;
    }
    auto dstExe = new ConvFpAIntBExecution(bn, op, mResource);
    *dst = dstExe;
    return true;
}


ErrorCode ConvFpAIntBExecution::onResize(const std::vector<Tensor*> &inputs, const std::vector<Tensor*> &outputs) {
    auto runtime = static_cast<CUDABackend*>(backend())->getCUDARuntime();
    auto input = inputs[0], output = outputs[0];
    const int UNIT = PACK_NUMBER;
    auto convCommon = mOp->main_as_Convolution2D()->common();
    auto pads = ConvolutionCommon::convolutionPadFull(input, output, mOp->main_as_Convolution2D()->common());
    int ic = input->channel();
    const int icp = UP_DIV(ic, 8) * 8;
    auto icDiv = UP_DIV(ic, UNIT);

    mIm2ColParamter.dilateX         = convCommon->dilateX();
    mIm2ColParamter.dilateY         = convCommon->dilateY();
    mIm2ColParamter.strideX         = convCommon->strideX();
    mIm2ColParamter.strideY         = convCommon->strideY();
    mIm2ColParamter.icDiv4          = icDiv;
    mIm2ColParamter.ic              = ic;
    mIm2ColParamter.kernelX         = convCommon->kernelX();
    mIm2ColParamter.kernelY         = convCommon->kernelY();
    mIm2ColParamter.padX = std::get<0>(pads);
    mIm2ColParamter.padY = std::get<1>(pads);

    mIm2ColParamter.ih = input->height();
    mIm2ColParamter.iw = input->width();
    mIm2ColParamter.oh = output->height();
    mIm2ColParamter.ow = output->width();
    mIm2ColParamter.srcZStep = input->height() * input->width() * UNIT * input->batch();
    mIm2ColParamter.srcYStep = input->width() * UNIT;
    mIm2ColParamter.packCUnit = UNIT;

    mActivationType = convCommon->relu() ? 1 : convCommon->relu6() ? 2 : 0;

    //MNN_PRINT("conv size:%d-%d, %d-%d-%d, %d-%d-%d\n", mIm2ColParamter.kernelX, mIm2ColParamter.strideX, input->height(), input->width(), input->channel(), output->height(), output->width(), output->channel());
    int e = output->height() * output->width() * output->batch();
    int l = ic * mIm2ColParamter.kernelX * mIm2ColParamter.kernelY;
    int h = output->channel();
    mGemmInfo.elh[0] = e;
    mGemmInfo.elh[1] = l;
    mGemmInfo.elh[2] = h;
    mGemmInfo.elhPad[0] = UP_DIV(e, 8) * 8;
    mGemmInfo.elhPad[1] = UP_DIV(l, 8) * 8;
    mGemmInfo.elhPad[2] = UP_DIV(h, 8) * 8;
    const int oc = mGemmInfo.elh[2];
    const int ocp = mGemmInfo.elhPad[2];

    mIsConv1x1S1D1P0 = (mIm2ColParamter.kernelX == 1 && mIm2ColParamter.kernelY == 1 &&
        mIm2ColParamter.strideX == 1 && mIm2ColParamter.strideY == 1 &&
        mIm2ColParamter.dilateX == 1 && mIm2ColParamter.dilateY == 1 &&
        mIm2ColParamter.padX == 0 && mIm2ColParamter.padY == 0);

    mNeedIm2Col = !(mIsConv1x1S1D1P0 && (mFp16Infer || mFp32Infer));

    auto pool = static_cast<CUDABackend*>(backend())->getBufferPool();
    if(mNeedIm2Col) {
        size_t im2colBytes = 2;
        // Only when fp32 Im2Col convert to fp32, Fp16Fp32Mix Im2Col convert to fp16
        if(mFp32Infer) {
            im2colBytes = 4;
        }
        auto buffer = pool->alloc(im2colBytes * (size_t)mGemmInfo.elh[0] * (size_t)mGemmInfo.elhPad[1]);
        mIm2ColBuffer = (void*)((uint8_t*)buffer.first + buffer.second);
        pool->free(buffer);
    }

    if (mDequantFilterTensor) {
        if (mDequantIsStatic) {
            backend()->onReleaseBuffer(mDequantFilterTensor.get(), Backend::STATIC);
        }
        // DYNAMIC buffers were already released in previous onResize, no action needed
        mDequantFilterTensor = nullptr;
    }
    mDequantFilter = nullptr;
    mNeedRuntimeDequant = false;
    mDequantIsStatic = false;

    // Weight dequantization: use DYNAMIC buffer with runtime dequant in onExecute.
    // DYNAMIC buffers are shared across all Conv layers (only one active at a time),
    // saving massive GPU memory vs STATIC (which would need ~30GB for 27B models).
    if (!mFp32Infer) {
        if (mIsConv1x1S1D1P0) {
            std::vector<int> dequantShape = {mGemmInfo.elhPad[2], mGemmInfo.elhPad[1]}; // Shape: [oc_p, ic_p]
            if (mFp16Infer || mFp16Fp32MixInfer) {
                mDequantFilterTensor.reset(Tensor::createDevice<int16_t>(dequantShape));
            } else {
                mDequantFilterTensor.reset(Tensor::createDevice<float>(dequantShape));
            }
            backend()->onAcquireBuffer(mDequantFilterTensor.get(), Backend::DYNAMIC);
            backend()->onReleaseBuffer(mDequantFilterTensor.get(), Backend::DYNAMIC);
            mNeedRuntimeDequant = true;
            mDequantIsStatic = false;

            mBiasAddr   = mResource->mBias;
            mBackendPtr = mResource->mBackend;
            mDequantFilter = (void*)mDequantFilterTensor->buffer().device;
            mFilterAddr = mDequantFilter;

            if (mFp32Infer) {
                return callCutlassGemmCudaCoreFloat32(inputs, outputs);
            }
            mGpuComputeCap = static_cast<CUDABackend*>(backend())->getCUDARuntime()->compute_capability();
            if (mGpuComputeCap < 70) {
                return callCutlassGemmCudaCoreFloat16(inputs, outputs);
            } else if (mGpuComputeCap < 75) {
                return callCutlassGemmTensorCore884(inputs, outputs);
            }
            return callCutlassGemmTensorCore(inputs, outputs);
        }
    }

    return NO_ERROR;
}

ErrorCode ConvFpAIntBExecution::onExecute(const std::vector<Tensor*> &inputs, const std::vector<Tensor*> &outputs) {
    MNN_ASSERT(inputs.size() == 1);
    MNN_ASSERT(outputs.size() == 1);
    auto input = inputs[0];
    auto output = outputs[0];

    //printf("convcutlass:%p %p\n", input->deviceId(), output->deviceId());
    //MNN_PRINT("cutlass hw:%d-%d\n", input->height(), input->width());
    auto runtime = static_cast<CUDABackend*>(backend())->getCUDARuntime();
    const void *input_addr = (const void*)inputs[0]->deviceId();
    const void *filter_addr = mResource->mFilter;
    const void *bias_addr = mResource->mBias;
    auto bn = backend();
    void *output_addr = (void*)outputs[0]->deviceId();

    const int sw = mIm2ColParamter.strideX;
    const int sh = mIm2ColParamter.strideY;
    const int kw = mIm2ColParamter.kernelX;
    const int kh = mIm2ColParamter.kernelY;
    const int dw = mIm2ColParamter.dilateX;
    const int dh = mIm2ColParamter.dilateY;
    const int pw = mIm2ColParamter.padX;
    const int ph = mIm2ColParamter.padY;
    const int ic = mIm2ColParamter.ic;
    const int icp = UP_DIV(ic, 8) * 8;
    const int iw = mIm2ColParamter.iw;
    const int ih = mIm2ColParamter.ih;

    const int oc = mGemmInfo.elh[2];
    const int ocp = mGemmInfo.elhPad[2];
    const int ow = mIm2ColParamter.ow;
    const int oh = mIm2ColParamter.oh;

    float maxV = FLT_MAX;
    float minV = -FLT_MAX;
    if (mActivationType == 1) {
        minV = 0.0f;
    }
    if (mActivationType == 2) {
        minV = 0.0f;
        maxV = 6.0f;
    }

    auto total = outputs[0]->batch() * oh * ow * ocp;
    auto& prop = runtime->prop();
    int limitThreads = UP_DIV(total, prop.multiProcessorCount);
    int threadNum = ALIMIN(prop.maxThreadsPerBlock/2, limitThreads);
    int blockNum = prop.multiProcessorCount;

    DivModFast d_oc(ocp);
    DivModFast d_ow(ow);
    DivModFast d_oh(oh);

    const int batch = inputs[0]->batch();

    // Weight dequantization now happens once in onResize (STATIC buffer).
    // GEMV path (batch <= 4) reads packed INT4 weights directly.


    if(mResource->mIsWeightInt4) {
        if (mIsConv1x1S1D1P0) {
            if (mFp32Infer) { // High + Int4
                if (batch <= 16) {
                    // P1-V5: 128-thread GEMV, 1 OC per block
                    dim3 v5_grid(oc, batch);
                    dim3 v5_block(128);
                    GEMV_FpAInt4B_V5<<<v5_grid, v5_block>>>(
                        (const float*)input_addr, (const uint8_t*)mResource->mFilter,
                        (const float*)mResource->mScale, (const float*)mResource->mOffset,
                        (const float*)bias_addr, (float*)output_addr,
                        maxV, minV, batch, ic, icp, oc, ocp, mResource->mQuanC);
                } else {
                    dim3 threads(TILE_DIM, TILE_DIM);
                    dim3 blocks(UP_DIV(ocp, TILE_DIM), UP_DIV(batch, TILE_DIM));

                    GEMM_FpAInt4B<<<blocks, threads>>>(
                        (const float*)input_addr, (const uint8_t*)mResource->mFilter,
                        (const float*)mResource->mScale,  (const float*)mResource->mOffset,
                        (const float*)bias_addr, (float*)output_addr,
                        maxV, minV, ic, icp, oc, ocp, batch, mResource->mQuanC
                    );
                }
            } else {
                if (batch <= 4) {
                    const int num_qg = (mResource->mQuanC > 0) ? (mResource->mQuanC / oc) : 1;
                    if (mResource->mGemvParams) {
                        constexpr int V14_OC = 4;
                        if (batch > 1) {
                            // V14_MB: One block processes ALL batch elements, weights loaded once
                            dim3 mb_grid((oc + V14_OC - 1) / V14_OC);
                            dim3 mb_block(128);
                            if (mFp16Infer) {
                                GEMV_FpAInt4B_V14_MB<half, V14_OC, 4><<<mb_grid, mb_block>>>(
                                    (const half*)input_addr, (const uint8_t*)mResource->mFilter,
                                    mResource->mGemvParams,
                                    (const half*)bias_addr, (half*)output_addr,
                                    maxV, minV, batch, ic, icp, oc, ocp, num_qg);
                            } else if (mFp16Fp32MixInfer) {
                                GEMV_FpAInt4B_V14_MB<float, V14_OC, 4><<<mb_grid, mb_block>>>(
                                    (const float*)input_addr, (const uint8_t*)mResource->mFilter,
                                    mResource->mGemvParams,
                                    (const float*)bias_addr, (float*)output_addr,
                                    maxV, minV, batch, ic, icp, oc, ocp, num_qg);
                            }
                        } else {
                            // batch=1: original V14
                            dim3 v14_grid((oc + V14_OC - 1) / V14_OC, batch);
                            dim3 v14_block(128);
                            if (mFp16Infer) {
                                GEMV_FpAInt4B_V14<half, V14_OC><<<v14_grid, v14_block>>>(
                                    (const half*)input_addr, (const uint8_t*)mResource->mFilter,
                                    mResource->mGemvParams,
                                    (const half*)bias_addr, (half*)output_addr,
                                    maxV, minV, batch, ic, icp, oc, ocp, num_qg);
                            } else if (mFp16Fp32MixInfer) {
                                GEMV_FpAInt4B_V14<float, V14_OC><<<v14_grid, v14_block>>>(
                                    (const float*)input_addr, (const uint8_t*)mResource->mFilter,
                                    mResource->mGemvParams,
                                    (const float*)bias_addr, (float*)output_addr,
                                    maxV, minV, batch, ic, icp, oc, ocp, num_qg);
                            }
                        }
                    } else {
                        // Fallback to V9/V5
                        dim3 v5_grid(oc, batch);
                        dim3 v5_block(128);
                        if (mFp16Infer) {
                            constexpr int OC_PER_BLK = 2;
                            dim3 v9_grid((oc + OC_PER_BLK - 1) / OC_PER_BLK, batch);
                            GEMV_FpAInt4B_V9<half, OC_PER_BLK><<<v9_grid, v5_block>>>(
                                (const half*)input_addr, (const uint8_t*)mResource->mFilter,
                                (const half*)mResource->mScale, (const half*)mResource->mOffset,
                                (const half*)bias_addr, (half*)output_addr,
                                maxV, minV, batch, ic, icp, oc, ocp, mResource->mQuanC);
                        } else if (mFp16Fp32MixInfer) {
                            GEMV_FpAInt4B_V5<<<v5_grid, v5_block>>>(
                                (const float*)input_addr, (const uint8_t*)mResource->mFilter,
                                (const float*)mResource->mScale, (const float*)mResource->mOffset,
                                (const float*)bias_addr, (float*)output_addr,
                                maxV, minV, batch, ic, icp, oc, ocp, mResource->mQuanC);
                        }
                    }
                } else {
                    // Prefill/batched GEMM path: CUTLASS with dequantized FP16 weights
                    if (mNeedRuntimeDequant) {
                        // Runtime dequantization: dequant INT4→FP16 into DYNAMIC buffer before CUTLASS
                        if (mResource->mIsWeightInt4) {
                            int threads_per_row = UP_DIV(icp, 16);
                            dim3 dq_block(min(threads_per_row, 256));
                            dim3 dq_grid(oc, UP_DIV(threads_per_row, (int)dq_block.x));
                            if (mFp16Infer) {
                                DequantizeInt4Weight<half, half><<<dq_grid, dq_block>>>(
                                    (const uint8_t*)mResource->mFilter, (half*)mDequantFilter,
                                    (const half*)mResource->mScale, (const half*)mResource->mOffset,
                                    oc, ic, icp, mResource->mQuanC);
                            } else if (mFp16Fp32MixInfer) {
                                DequantizeInt4Weight<float, half><<<dq_grid, dq_block>>>(
                                    (const uint8_t*)mResource->mFilter, (half*)mDequantFilter,
                                    (const float*)mResource->mScale, (const float*)mResource->mOffset,
                                    oc, ic, icp, mResource->mQuanC);
                            }
                        } else {
                            dim3 threads(32, 32);
                            dim3 blocks(UP_DIV(ic, threads.x), UP_DIV(oc, threads.y));
                            if (mFp16Infer) {
                                DequantizeInt8Weight<half, half><<<blocks, threads>>>(
                                    (const int8_t*)mResource->mFilter, (half*)mDequantFilter,
                                    (const half*)mResource->mScale, (const half*)mResource->mOffset,
                                    oc, ic, icp, mResource->mQuanC);
                            } else if (mFp16Fp32MixInfer) {
                                DequantizeInt8Weight<float, half><<<blocks, threads>>>(
                                    (const int8_t*)mResource->mFilter, (half*)mDequantFilter,
                                    (const float*)mResource->mScale, (const float*)mResource->mOffset,
                                    oc, ic, icp, mResource->mQuanC);
                            }
                        }
                    }
                    if (mFp16Fp32MixInfer) {
                        size_t maxCount = mGemmInfo.elh[0] * mGemmInfo.elhPad[1];
                        callFloat2Half(input_addr, mIm2ColBuffer, maxCount, runtime);
                    }
                    runCutlassGemmFunc();
                }
            }
        } else {
            if(mFp16Infer) {
                CONV_FpAInt4B<<<blockNum, threadNum>>>((const half*)input_addr, (const uint8_t*)mResource->mFilter,
                    (const half*)mResource->mScale,  (const half*)mResource->mOffset, (const half*)bias_addr, (half*)output_addr,
                    maxV, minV, ic, icp, iw, ih, oc, ocp, ow, oh, kw, kh, dw, dh, sw, sh, pw, ph, total, mResource -> mQuanC,
                    d_oc, d_ow, d_oh);
            } else {
                CONV_FpAInt4B<<<blockNum, threadNum>>>((const float*)input_addr, (const uint8_t*)mResource->mFilter,
                    (const float*)mResource->mScale,  (const float*)mResource->mOffset, (const float*)bias_addr, (float*)output_addr,
                    maxV, minV, ic, icp, iw, ih, oc, ocp, ow, oh, kw, kh, dw, dh, sw, sh, pw, ph, total, mResource -> mQuanC,
                    d_oc, d_ow, d_oh);
            }
        }
    } else { // int8
        if (mIsConv1x1S1D1P0) {
            if (mFp32Infer) { // High + Int8
                if (batch <= 16) {
                    // P1: Use optimized GEMV with warp shuffle reduction
                    dim3 threads(WARP_SIZE, GEMV_OC_PER_BLOCK);
                    dim3 blocks(UP_DIV(oc, GEMV_OC_PER_BLOCK), batch);

                    if (mFp16Infer) {
                        GEMV_FpAInt8B_V2<<<blocks, threads>>>(
                            (const half*)input_addr, (const int8_t*)mResource->mFilter,
                            (const half*)mResource->mScale,  (const half*)mResource->mOffset,
                            (const half*)bias_addr, (half*)output_addr,
                            maxV, minV, batch, ic, icp, oc, ocp, mResource->mQuanC
                        );
                    } else {
                        GEMV_FpAInt8B_V2<<<blocks, threads>>>(
                            (const float*)input_addr, (const int8_t*)mResource->mFilter,
                            (const float*)mResource->mScale,  (const float*)mResource->mOffset,
                            (const float*)bias_addr, (float*)output_addr,
                            maxV, minV, batch, ic, icp, oc, ocp, mResource->mQuanC
                        );
                    }
                } else {
                    dim3 threads(TILE_DIM, TILE_DIM);
                    dim3 blocks(UP_DIV(ocp, TILE_DIM), UP_DIV(batch, TILE_DIM));

                    GEMM_FpAInt8B<<<blocks, threads>>>(
                        (const float*)input_addr, (const int8_t*)mResource->mFilter,
                        (const float*)mResource->mScale,  (const float*)mResource->mOffset,
                        (const float*)bias_addr, (float*)output_addr,
                        maxV, minV, ic, icp, oc, ocp, batch, mResource->mQuanC
                    );
                }
            } else {
                if (batch <= 16) {
                    // P1: FP16 direct GEMV for INT8 small batch
                    dim3 threads(WARP_SIZE, GEMV_OC_PER_BLOCK);
                    dim3 blocks(UP_DIV(oc, GEMV_OC_PER_BLOCK), batch);
                    if (mFp16Infer) {
                        GEMV_FpAInt8B_V2<<<blocks, threads>>>(
                            (const half*)input_addr, (const int8_t*)mResource->mFilter,
                            (const half*)mResource->mScale, (const half*)mResource->mOffset,
                            (const half*)bias_addr, (half*)output_addr,
                            maxV, minV, batch, ic, icp, oc, ocp, mResource->mQuanC
                        );
                    } else if (mFp16Fp32MixInfer) {
                        GEMV_FpAInt8B_V2<<<blocks, threads>>>(
                            (const float*)input_addr, (const int8_t*)mResource->mFilter,
                            (const float*)mResource->mScale, (const float*)mResource->mOffset,
                            (const float*)bias_addr, (float*)output_addr,
                            maxV, minV, batch, ic, icp, oc, ocp, mResource->mQuanC
                        );
                    }
                } else {
                    // Prefill path: CUTLASS GEMM with dequantized FP16 weights
                    if (mFp16Fp32MixInfer) {
                        size_t maxCount = mGemmInfo.elh[0] * mGemmInfo.elhPad[1];
                        callFloat2Half(input_addr, mIm2ColBuffer, maxCount, runtime);
                    }
                    runCutlassGemmFunc();
                }
            }
        } else {
            if(mFp16Infer) {
                CONV_FpAInt8B<<<blockNum, threadNum>>>((const half*)input_addr, (const int8_t*)mResource->mFilter,
                    (const half*)mResource->mScale,  (const half*)mResource->mOffset, (const half*)bias_addr, (half*)output_addr,
                    maxV, minV, ic, icp, iw, ih, oc, ocp, ow, oh, kw, kh, dw, dh, sw, sh, pw, ph, total, mResource -> mQuanC,
                    d_oc, d_ow, d_oh);
            } else {
                CONV_FpAInt8B<<<blockNum, threadNum>>>((const float*)input_addr, (const int8_t*)mResource->mFilter,
                        (const float*)mResource->mScale,  (const float*)mResource->mOffset, (const float*)bias_addr, (float*)output_addr,
                    maxV, minV, ic, icp, iw, ih, oc, ocp, ow, oh, kw, kh, dw, dh, sw, sh, pw, ph, total, mResource -> mQuanC,
                    d_oc, d_ow, d_oh);
            }
        }
    }
    
    checkKernelErrors;
    return NO_ERROR;
}


}// namespace CUDA
}// namespace MNN