Add fused_transpose_quant op

[lshpku]

PR types

Performance optimization

PR changes

APIs

Description

新增fused_transpose_quant算子，定义如下：

/**
 * Doing quantization on dim[-2] of X, then transpose dim[-1] and dim[-2] of X.
 *
 * Inputs:
 *   X    : [*, M, K], bfloat16
 *
 * Outputs:
 *   out  : [*, K, M], float8_e4m3fn
 *   scale: [*, M/128, K], float32
 *
 * Requirements:
 *   1) batch_size <= 65535
 *   2) M <= 65535 * 128 and M % 128 == 0
 */
std::vector<paddle::Tensor> fused_transpose_quant(const paddle::Tensor& X);

实现亮点

1. 使用128x128的分块进行transpose，实现高并行度
2. 由于M维保证128对齐，因此总是使用4x向量化写回out
3. 由于K维无对齐保证，因此读X和写scale时区分了1x/2x/4x向量化实例

性能测试

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

输入x.shape	用时(ns)	带宽(GBps)	带宽利用率	说明
[4, 7168, 4096]	278,176	1280	82.3%	4x向量化
[4, 7168, 4096 + 2]	294,663	1208	77.7%	2x向量化
[4, 7168, 4096 + 1]	404,087	881	56.7%	无向量化

Pcard-85711

[paddle-bot]

Thanks for your contribution!