compressor-算子接口（torch_extension）-算子库-API参考-CANN9.2.0-beta.1开发文档-昇腾社区

[object Object][object Object]

[object Object]Ascend 950PR/Ascend 950DT[object Object]：支持
[object Object]Atlas A3 训练系列产品/Atlas A3 推理系列产品[object Object]：支持
[object Object]Atlas A2 训练系列产品/Atlas A2 推理系列产品[object Object]：不支持
[object Object]Atlas 200I/500 A2 推理产品[object Object]：不支持
[object Object]Atlas 推理系列产品[object Object]：不支持
[object Object]Atlas 训练系列产品[object Object]：不支持

[object Object]

接口功能：Compressor是推理场景下SMLA和QLI的前处理算子，用于将每4或128个token的KV cache压缩成一个，然后每个token与这些压缩的KV cache进行DSA计算。在长序列的情况下，Compressor可以有效地减少计算开销。主要计算过程为：
1. 将输入 $X$ 与 $W^{KV}$ 做Matmul运算得到 $kv\_state$ ，将输入 $X$ 与 $W^{Gate}$ 做Matmul运算后再与 $Ape$ 做Add运算得到 $score\_state$ ， $kv\_state$ 与 $score\_state$ 根据输入的start_pos及cu_seqlens完成更新。
2. 在coff为2的情况下对 $kv\_state$ 和 $score\_state$ 进行数据重排。
3. 对 $score\_state$ 进行softmax运算将softmax结果与 $kv\_state$ 做Mul计算，后进行ReduceSum运算。
计算公式：
1. 计算矩阵乘法：
$C4A：\left[kv\_state^a, score\_state^a\right] = X @ \left[W^{aKV}, W^{aGate}\right], \left[kv\_state^b, score\_state^b\right] = X @ \left[W^{bKV}, W^{bGate}\right];$ $C128A：\left[kv\_state, score\_state\right] = X @ \left[W^{KV}, W^{Gate}\right]$
1. 计算分组加法：
$C4A：score\_state_i^\prime = \left[score\_state_{\left[4(i-1)+1:4i,:\right]}^a; score\_state_{\left[4i+1:4(i+1),:\right]}^b\right] + Ape,~i=1,2,\cdots, \frac{s}{4};$ $C128A：score\_state_i^\prime = score\_state_{\left[128(i-1)+1:128i,:\right]} + Ape,~i=1,2,\cdots, \frac{s}{128};$
1. 计算分组Softmax：
$C4A：S_i^\prime = softmax(score\_state_i^\prime),~i=1,2,\cdots, \frac{s}{4};$ $C128A：S_i^\prime = softmax(score\_state_i^\prime),~i=1,2,\cdots, \frac{s}{128};$
1. 计算Hadamard乘积：
$C4A：(S_H)_i = S_i^\prime \odot \left[kv\_state^a_{\left[4(i-1)+1:4i,:\right]} ;kv\_state^b_{\left[4i+1:4(i+1),:\right]}\right],~i=1,2,\cdots, \frac{s}{4};$ $C128A：S_H = S_i^\prime \odot kv\_state;$
1. 沿着压缩轴分组求和：
$C4A：C_{i}^{\text{Comp}} = \left[1\right]_{1\times8} @ (S_H)_i, ~i=1,2,\cdots, \frac{s}{4};$ $C128A：C_{i}^{\text{Comp}} = \left[1\right]_{1\times128} @ (S_H)_i, ~i=1,2,\cdots, \frac{s}{128};$

[object Object]

[object Object][object Object]undefined

[object Object]Atlas A3 训练系列产品/Atlas A3 推理系列产品[object Object]： cache_mode不支持输入2，且不支持0轴非连续。

[object Object][object Object]undefined

[object Object]

该接口支持推理场景下使用。
该接口支持单算子模式和TorchAir图模式调用。
x参数维度含义：B（Batch Size）表示输入样本批量大小、S（Sequence Length）表示输入样本序列长度、H（Head Size）表示hidden层的大小、D（Head Dim）表示hidden层的最小单元大小、T表示所有Batch输入样本序列长度的累加和。
该接口支持B、S泛化，且存在如下场景限制：
- 只支持B、S为0。
- 部分长序列场景下，如果计算量过大可能会导致出现超过NPU内存的报错，注：这里计算量会受x输入shape的影响，值越大计算量越大。典型的长序列（即B、S的乘积或T较大）场景包括但不限于：
[object Object]
支持D为128/512。
支持H为1K~10K，512对齐。
支持block_size为1~1024。
支持cmp_ratio为2/4/8/16/32/64/128。支持如下三种典型组合场景：
- C4A: D=512, coff=2, cmp_ratio=4;
- C4Li: D=128, coff=2, cmp_ratio=4;
- C128A: D=512, coff=1, cmp_ratio=128。
该接口支持aclgraph模式。
该接口不支持GE图模式。

[object Object]

默认支持确定性计算。

[object Object]

单算子模式调用：

[object Object]
TorchAir图模式调用：

[object Object]