Add fused_stack_transpose_quant kernel (optional transpose)

[lshpku]

PR types

Function optimization

PR changes

APIs

Description

实现了fused_stack_quant和fused_stack_transpose_quant两个融合算子，定义如下：

/**
 * Stack tensors in X and do quantization on both dim[-1] and dim[-2].
 *
 * Inputs:
 *   X    : N tensors of [M, K], bfloat16
 *
 * Outputs:
 *   out  : [N * M, K], float8_e4m3fn
 *   scale: [N * M / 128, K / 128], float
 *
 * Requirements:
 *   1) N <= 65535
 *   2) M % 128 == 0
 *   3) K % 128 == 0
 */
std::vector<paddle::Tensor> fused_stack_quant(
    const std::vector<paddle::Tensor>& X);

/**
 * Stack tensors in X, transpose dim[-1] and dim[-2], and do quantization 
 * on both dim[-1] and dim[-2].
 *
 * Inputs:
 *   X    : N tensors of [M, K], bfloat16
 *
 * Outputs:
 *   out  : [N * K, M], float8_e4m3fn
 *   scale: [N * K / 128, M / 128], float
 *
 * Requirements:
 *   1) N <= 65535
 *   2) M % 128 == 0
 *   3) K % 128 == 0
 */
std::vector<paddle::Tensor> fused_stack_transpose_quant(
    const std::vector<paddle::Tensor>& X);

（两者唯一区别是，带transpose的版本会在stack后先交换 M、K 维度，再进行后续quant操作）

性能测试

在A100-40G上做了初步测试，由于A100不支持fp8，因此在cast fp32 to fp8的时候用了int8代替，基本可以反映H卡上的性能

测试对象	输入个数与shape	用时(ns)	带宽(GBps)	带宽利用率
fused_stack_quant	4 * [7168, 4096]	258,841	1361	87.5%
fused_stack_transpose_quant	4 * [7168, 4096]	270,972	1300	83.6%

Pcard-85711

[paddle-bot]

Thanks for your contribution!