wehub-resource-sync
137 lines
6.2 KiB
C++
137 lines
6.2 KiB
C++
// Copyright (c) Microsoft Corporation.
|
|
// SPDX-License-Identifier: Apache-2.0
|
|
|
|
// DeepSpeed Team
|
|
|
|
#pragma once
|
|
|
|
#include <cuda.h>
|
|
#include <stdint.h>
|
|
|
|
#include <cuda_fp16.h>
|
|
// Note: BF16 support on AMD but we have to exclude here cuda_bf16.h (which turn to
|
|
// <hip/hip_bfloat16.h> after hipifying), because this header is pulled into .cpp translation units
|
|
// that are compiled by a host-only compiler, which triggers build errors. Added forward declaration
|
|
// instead, see code block below
|
|
#if defined(BF16_AVAILABLE)
|
|
#if !defined(__HIP_PLATFORM_AMD__)
|
|
#include <cuda_bf16.h>
|
|
#else
|
|
struct __hip_bfloat16;
|
|
#endif
|
|
#endif
|
|
|
|
#include <cuda_runtime_api.h>
|
|
#include <stdio.h>
|
|
|
|
#define QUANT_SWITCH(Q_BITS, ...) \
|
|
[&] { \
|
|
if (12 == Q_BITS) { \
|
|
constexpr int CONST_STOCHASTIC_ROUNDING = 0; \
|
|
constexpr int CONST_Q_BITS = 8; \
|
|
constexpr int CONST_Q_MANTISA_BITS = 3; \
|
|
__VA_ARGS__(); \
|
|
} else if (13 == Q_BITS) { \
|
|
constexpr int CONST_STOCHASTIC_ROUNDING = 1; \
|
|
constexpr int CONST_Q_BITS = 8; \
|
|
constexpr int CONST_Q_MANTISA_BITS = 3; \
|
|
__VA_ARGS__(); \
|
|
} else if (10 == Q_BITS) { \
|
|
constexpr int CONST_STOCHASTIC_ROUNDING = 0; \
|
|
constexpr int CONST_Q_BITS = 8; \
|
|
constexpr int CONST_Q_MANTISA_BITS = 2; \
|
|
__VA_ARGS__(); \
|
|
} else if (11 == Q_BITS) { \
|
|
constexpr int CONST_STOCHASTIC_ROUNDING = 1; \
|
|
constexpr int CONST_Q_BITS = 8; \
|
|
constexpr int CONST_Q_MANTISA_BITS = 2; \
|
|
__VA_ARGS__(); \
|
|
} else if (28 == Q_BITS) { \
|
|
constexpr int CONST_STOCHASTIC_ROUNDING = 0; \
|
|
constexpr int CONST_Q_BITS = 12; \
|
|
constexpr int CONST_Q_MANTISA_BITS = 7; \
|
|
__VA_ARGS__(); \
|
|
} else if (29 == Q_BITS) { \
|
|
constexpr int CONST_STOCHASTIC_ROUNDING = 1; \
|
|
constexpr int CONST_Q_BITS = 12; \
|
|
constexpr int CONST_Q_MANTISA_BITS = 7; \
|
|
__VA_ARGS__(); \
|
|
} else if (6 == Q_BITS) { \
|
|
constexpr int CONST_STOCHASTIC_ROUNDING = 0; \
|
|
constexpr int CONST_Q_BITS = 6; \
|
|
constexpr int CONST_Q_MANTISA_BITS = 2; \
|
|
__VA_ARGS__(); \
|
|
} else if (7 == Q_BITS) { \
|
|
constexpr int CONST_STOCHASTIC_ROUNDING = 1; \
|
|
constexpr int CONST_Q_BITS = 6; \
|
|
constexpr int CONST_Q_MANTISA_BITS = 2; \
|
|
__VA_ARGS__(); \
|
|
} else if (2 == Q_BITS) { \
|
|
constexpr int CONST_STOCHASTIC_ROUNDING = 0; \
|
|
constexpr int CONST_Q_BITS = 4; \
|
|
constexpr int CONST_Q_MANTISA_BITS = 1; \
|
|
__VA_ARGS__(); \
|
|
} else { \
|
|
constexpr int CONST_STOCHASTIC_ROUNDING = 1; \
|
|
constexpr int CONST_Q_BITS = 4; \
|
|
constexpr int CONST_Q_MANTISA_BITS = 1; \
|
|
__VA_ARGS__(); \
|
|
} \
|
|
}()
|
|
|
|
#define DEQUANT_SWITCH(Q_MANTISA_EXPONENT_BITS, ...) \
|
|
[&] { \
|
|
if (12 == Q_MANTISA_EXPONENT_BITS) { \
|
|
constexpr int CONST_Q_MANTISA_BITS = 3; \
|
|
constexpr int CONST_Q_EXPONENT_BITS = 4; \
|
|
__VA_ARGS__(); \
|
|
} else if (10 == Q_MANTISA_EXPONENT_BITS) { \
|
|
constexpr int CONST_Q_MANTISA_BITS = 2; \
|
|
constexpr int CONST_Q_EXPONENT_BITS = 5; \
|
|
__VA_ARGS__(); \
|
|
} else if (28 == Q_MANTISA_EXPONENT_BITS) { \
|
|
constexpr int CONST_Q_MANTISA_BITS = 7; \
|
|
constexpr int CONST_Q_EXPONENT_BITS = 4; \
|
|
__VA_ARGS__(); \
|
|
} else if (6 == Q_MANTISA_EXPONENT_BITS) { \
|
|
constexpr int CONST_Q_MANTISA_BITS = 2; \
|
|
constexpr int CONST_Q_EXPONENT_BITS = 3; \
|
|
__VA_ARGS__(); \
|
|
} else { \
|
|
constexpr int CONST_Q_MANTISA_BITS = 1; \
|
|
constexpr int CONST_Q_EXPONENT_BITS = 2; \
|
|
__VA_ARGS__(); \
|
|
} \
|
|
}()
|
|
|
|
template <typename T, int mantisa, int exponent>
|
|
void launch_quantization(T* val,
|
|
uint8_t* q_val,
|
|
int num_groups,
|
|
int group_size,
|
|
cudaStream_t stream,
|
|
float q_range,
|
|
int q_bits,
|
|
int q_mantisa_bits,
|
|
int stochastic_rounding);
|
|
|
|
template <typename T, int mantisa>
|
|
void launch_dequantization(uint8_t* val,
|
|
T* q_val,
|
|
int num_groups,
|
|
int group_size,
|
|
int q_mantisa_bits,
|
|
int q_exponent_bits,
|
|
cudaStream_t stream);
|
|
|
|
template <typename T, int mantisa>
|
|
void launch_selective_dequantization(uint8_t* val,
|
|
T* q_val,
|
|
int32_t* indexes,
|
|
int num_groups,
|
|
int group_size,
|
|
int num_indexes,
|
|
int q_mantisa_bits,
|
|
int q_exponent_bits,
|
|
cudaStream_t stream);
|