fp4-gemm / torch-ext /torch_binding.h
liangsu9988's picture
Publish GROOT N1.7 Thor kernel sources
88b8ef2 verified
Raw
History Blame
687 Bytes
// SPDX-License-Identifier: Apache-2.0
#pragma once
#include <torch/all.h>
void fp4_w4a16_linear_bf16(
torch::Tensor const& a_packed,
torch::Tensor const& b_packed,
torch::Tensor const& sfa,
torch::Tensor const& sfb,
torch::Tensor& out,
double alpha,
int64_t variant);
void quantize_fp4_sfa_fp16(
torch::Tensor const& x,
torch::Tensor& packed,
torch::Tensor& sfa,
bool is_sfb);
void quantize_fp4_sfa_bf16(
torch::Tensor const& x,
torch::Tensor& packed,
torch::Tensor& sfa,
bool is_sfb);
void dequantize_fp4_sfa_fp16(
torch::Tensor const& packed,
torch::Tensor const& sfa,
torch::Tensor& out,
bool is_sfb);