PR types

Others

PR changes

Others

Description

主要工作：

1. 在fused_rotary_position_embedding算子层面支持 MQA 和 GQA，具体的方法是若 q 的 num_heads 与 k, v 的不相等，则开启 MQA 和 GQA 的功能，分别调用三次 CUDA kernel。

2. 修改 fused_rotary_position_embedding 算子的单测，并增加 MQA 和 GQA 的单测。发现在当前 develop 分支下，fused_rotary_position_embedding 的 fp32 单测存在微小精度误差，在 num_heads >= 4 时 （如 q.shape=[2, 8, 4, 16], batch_size, seqlen, num_heads, head_dim） 相对误差约为 1.1436 * 1e-5（正常情况下 fp32 单测的 rtol 应为 1e-05）。当前的规避办法是修改单测 TestFusedRotaryPositionEmbedding 在 GQA 策略下的输入 shape，使单测的相对误差满足 1e-05 的阈值

3. k, v 都是可选的。目前已支持以下几种输入场景都正确：

• q，k，v
• q，k
• q，v
• q

Pcard-79849