TBE CPU 自动向量化¶
FP8/16/32 自动向量化实现方法¶
-
template<typename InType = std::uint8_t, typename IndexType = std::int64_t, typename OffsetType = std::int32_t, typename OutType = float>
bool EmbeddingSpMDM_autovec(const std::int64_t block_size, const std::int64_t output_size, const std::int64_t index_size, const std::int64_t data_size, const InType *input, const IndexType *indices, const OffsetType *offsets_or_lengths, const float *weights, bool normalize_by_lengths, OutType *out, bool is_weight_positional = false, bool use_offsets = true, std::int64_t output_stride = -1, std::int64_t input_stride = -1, bool scale_bias_last = true, bool no_bag = false, bool is_bf16_out = false, bool is_bf16_in = false)¶ 方法
EmbeddingSpMDM_ref
的自动向量化版本,用于 FP32 权重类型。- 模板参数:
InType – 输入数据类型(使用
uint8_t
)IndexType – 索引数据类型(使用
int64_t
)OffsetType – 偏移量数据类型(使用
int32_t
)OutType – 输出数据类型(使用
float
)
- 参数:
block_size – 块中的元素数量(
int64_t
)output_size – 输出中的元素数量(
int64_t
)index_size – 索引中的元素数量(
int64_t
)data_size – 数据中的元素数量(
int64_t
)input – 输入的地址(
InType*
)indices – 索引的地址(
IndexType*
)offsets_or_lengths – 偏移量的地址(
OffsetType*
)weights – 求和的权重;可选,对于非加权求和可以为 null(
float*
)normalize_by_lengths – 是否按长度归一化(
bool
)out – 输出的地址(
OutType*
)is_weight_positional – 如果为
true
,则权重是位置性的;对于 FP32 自动向量化实现,设置为false
(bool
)use_offsets – 如果
true
,将使用偏移量而不是长度;对于 FP32 autovec 实现,设置为true
(bool
)output_stride – 如果为 -1,则 output_stride 与 block_size 相同;对于 FP32 autovec 实现,设置为 -1 (
int64_t
)input_stride – 如果为 -1,则 input_stride 与 block_size 相同;对于 FP32 autovec 实现,设置为 -1 (
int64_t
)scale_bias_last – 如果
true
,则 scale 和 bias 出现在每一行的末尾;对于 FP32 autovec 实现,设置为true
(bool
)no_bag – 如果
true
,则没有 embedding bag;对于 FP32 autovec 实现,设置为false
(bool
)is_bf16_out – 如果
true
,则输出为BFLOAT16
类型;对于 FP32 autovec 实现,设置为false
(bool
)is_bf16_in – 如果
true
,则输入为BFLOAT16
类型;对于 FP32 autovec 实现,设置为false
(bool
)
-
template<typename IndexType, typename OffsetType, typename OutType>
bool EmbeddingSpMDMFP8_autovec(const int64_t block_size, const int64_t output_size, const int64_t index_size, const int64_t data_size, const uint8_t *input, const IndexType *indices, const OffsetType *offsets_or_lengths, const float *weights, bool normalize_by_lengths, OutType *out, bool is_weight_positional, bool use_offsets, int64_t output_stride, int64_t input_stride, int exponent_bits, int exponent_bias, bool is_bf16_out)¶ 方法
EmbeddingSpMDM_ref
的自动向量化版本,用于 FP8 权重类型。- 模板参数:
InType – 输入数据类型(使用
uint8_t
)IndexType – 索引数据类型(使用
int64_t
)OffsetType – 偏移量数据类型(使用
int32_t
)OutType – 输出数据类型(使用
float
)
- 参数:
block_size – 块中的元素数量(
int64_t
)output_size – 输出中的元素数量(
int64_t
)index_size – 索引中的元素数量(
int64_t
)data_size – 数据中的元素数量(
int64_t
)input – 输入的地址(
InType*
)indices – 索引的地址(
IndexType*
)offsets_or_lengths – 偏移量的地址(
OffsetType*
)weights – 求和的权重;可选,对于非加权求和可以为 null(
float*
)normalize_by_lengths – 是否按长度归一化(
bool
)out – 输出的地址(
OutType*
)is_weight_positional – 如果
true
,则权重是位置性的;对于 FP8 autovec 实现,设置为false
(bool
)use_offsets – 如果
true
,将使用偏移量而不是长度;对于 FP8 autovec 实现,设置为true
(bool
)output_stride – 如果为 -1,则 output_stride 与 block_size 相同;对于 FP8 autovec 实现,设置为 -1 (
int64_t
)exponent_bits – 指数中使用的位数
exponent_bias – 指数中使用的偏差
is_bf16_out – 如果
true
,则输出为BFLOAT16
类型;对于 FP8 autovec 实现,设置为false
(bool
)