tlc-pack / cutlass_fpa_intb_gemm Goto Github PK

void fpA_intB_gemm_forward_cuda(torch::Tensor &input,
                                torch::Tensor &weight,
                                torch::Tensor &scale,
                                torch::Tensor &output,
                                int m, int n, int k)
{
    c10::cuda::CUDAGuard device_guard(input.device());
    const fastertransformer::half *input_ptr = reinterpret_cast<fastertransformer::half *>(input.data_ptr());
    const uint8_t *weight_ptr = reinterpret_cast<const uint8_t *>(weight.data_ptr());
    const fastertransformer::half *scale_ptr = reinterpret_cast<fastertransformer::half *>(scale.data_ptr());
    fastertransformer::half *output_ptr = reinterpret_cast<fastertransformer::half *>(output.data_ptr());

    fastertransformer::gemm_fp16_int_bias_act(
        input_ptr,
        weight_ptr,
        scale_ptr,
        nullptr,
        output_ptr,
        std::nullopt,
        m, n, k,
        0,
        nullptr,
        0,
        0);
}