Thanks for sending an requirement! Please fill in the following template to help quickly solve your problem.

Backgroud（背景信息）

KdaInputProj是推理场景下Recurrent KDA的前处理算子，负责qkv，beta，gate和g的投影操作；通过算子融合，实现计算的CV并行，提升KDA前置处理的性能

Benefit / Necessity （价值/作用）

提升KIMI 模型KDA前置投影操作的性能。其中Matmul (beta/gate/g) 所有的 BlockMmad 任务会被所有的AIC单元并行处理，以减少计算轮次，提高硬件利用率

Design（设计方案）

Stage 1

在Stage 1，KdaInputProj算子分别利用AIC和AIV并行执行 Matmul (beta/gate/g) 和 DynamicMxQuant 两个计算任务。其中，Matmul (beta/gate/g) 所有的 BlockMmad 任务会被所有的AIC单元并行处理，以减少计算轮次，提高硬件利用率。计算公式如下：

$$\begin{bmatrix} \mathbf{beta}\\ \mathbf{gate}\\ \mathbf{g}\\ \end{bmatrix} = \mathbf{X} \begin{bmatrix} W_{beta}\\ W_{gate}\\ W_{g} \end{bmatrix}$$

Stage 2

在Stage 2，KdaInputProj算子分别利用AIC和AIV并行执行 QuantMatmul (qkv) 和 Sigmoid 两个计算任务。其中 QuantMatmul 的激活为Stage 1中的 DynamicMxQuant 任务的输出，qkv_weight 为离线量化好的权重。