File size: 3,521 Bytes
1dc29e9 f5a110e 1dc29e9 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 |
#pragma once
#include "torch/csrc/cuda/Stream.h"
#include "torch/all.h"
#include <ATen/cuda/CUDAContext.h>
#include <cstdio>
#include <cuda_fp16.h>
#include <cuda_runtime.h>
#include <iostream>
// Generates a conflict with CUDA 12.6 between nvtx 2 and 3. Does not
// seem to be used anyway?
//
// #include <nvToolsExt.h>
#include <torch/custom_class.h>
#include <torch/script.h>
#include <vector>
#define TORCH_CHECK_DTYPE(__x, __dtype) TORCH_CHECK((__x).device().is_meta() || (__x).dtype() == torch::__dtype, #__x " is incorrect datatype, must be " #__dtype)
#define TORCH_CHECK_SHAPES(__x, __dim_x, __y, __dim_y, __scale_y) TORCH_CHECK((__x).device().is_meta() || (__x).size(__dim_x) == (__y).size(__dim_y) * __scale_y, #__x " and " #__y " have incompatible shapes")
#define TORCH_CHECK_BUFFER_SIZE(__buffer, __minimum_size) TORCH_CHECK((__buffer).numel() >= __minimum_size, #__buffer " is too small")
#define CHECK_TYPE(x, st) TORCH_CHECK(x.scalar_type() == st, "Inconsistency of Tensor type: " #x)
#define CHECK_TH_CUDA(x) TORCH_CHECK(x.is_cuda(), #x " must be a CUDA tensor")
#define CHECK_CPU(x) TORCH_CHECK(!x.is_cuda(), #x " must be a CPU tensor")
#define CHECK_CONTIGUOUS(x) TORCH_CHECK(x.is_contiguous(), #x " must be contiguous")
#define CHECK_INPUT(x, st) \
CHECK_TH_CUDA(x); \
CHECK_CONTIGUOUS(x); \
CHECK_TYPE(x, st)
#define CHECK_CPU_INPUT(x, st) \
CHECK_CPU(x); \
CHECK_CONTIGUOUS(x); \
CHECK_TYPE(x, st)
#define CHECK_OPTIONAL_INPUT(x, st) \
if (x.has_value()) { \
CHECK_INPUT(x.value(), st); \
}
#define CHECK_OPTIONAL_CPU_INPUT(x, st) \
if (x.has_value()) { \
CHECK_CPU_INPUT(x.value(), st); \
}
#define PRINT_TENSOR(x) std::cout << #x << ":\n" << x << std::endl
#define PRINT_TENSOR_SIZE(x) std::cout << "size of " << #x << ": " << x.sizes() << std::endl
namespace fastertransformer {
template<typename T>
inline T* get_ptr(torch::Tensor& t)
{
return reinterpret_cast<T*>(t.data_ptr());
}
std::vector<size_t> convert_shape(torch::Tensor tensor);
size_t sizeBytes(torch::Tensor tensor);
QuantType get_ft_quant_type(torch::ScalarType quant_type)
{
if (quant_type == torch::kInt8) {
return QuantType::INT8_WEIGHT_ONLY;
}
else if (quant_type == at::ScalarType::QUInt4x2) {
return QuantType::PACKED_INT4_WEIGHT_ONLY;
}
else {
TORCH_CHECK(false, "Invalid quantization type");
}
}
} // namespace fastertransformer
|