|
#pragma once |
|
|
|
#include <torch/torch.h> |
|
|
|
bool cutlass_scaled_mm_supports_fp8(int64_t cuda_device_capability); |
|
|
|
void cutlass_scaled_mm(torch::Tensor& out, torch::Tensor const& a, |
|
torch::Tensor const& b, torch::Tensor const& a_scales, |
|
torch::Tensor const& b_scales, |
|
c10::optional<torch::Tensor> const& bias); |
|
|
|
void cutlass_scaled_mm_azp(torch::Tensor& out, torch::Tensor const& a, |
|
torch::Tensor const& b, |
|
torch::Tensor const& a_scales, |
|
torch::Tensor const& b_scales, |
|
torch::Tensor const& azp_adj, |
|
c10::optional<torch::Tensor> const& azp, |
|
c10::optional<torch::Tensor> const& bias); |
|
|
