已合并
torch_extension API文档新增产品型号标签 #8233
chenjiao创建于 12 天前
torch_extension API文档新增产品型号标签 #8233
已合并
共 36 个文件变更+224-168
| @@ -1,34 +1,46 @@ | |||
| 1 | # cann_ops_nn.swiglu_group | 1 | # swiglu_group |
| 2 | 2 | ||
| 3 | ## 产品支持情况 | 3 | ## 产品支持情况 |
| 4 | 4 | ||
| 5 | <!-- npu="950" id1 --> | ||
| 5 | - <term>Ascend 950PR/Ascend 950DT</term>：支持 | 6 | - <term>Ascend 950PR/Ascend 950DT</term>：支持 |
| 7 | <!-- end id1 --> | ||
| 8 | <!-- npu="A3" id2 --> | ||
| 6 | - <term>Atlas A3 训练系列产品/Atlas A3 推理系列产品</term>：不支持 | 9 | - <term>Atlas A3 训练系列产品/Atlas A3 推理系列产品</term>：不支持 |
| 10 | <!-- end id2 --> | ||
| 11 | <!-- npu="910b" id3 --> | ||
| 7 | - <term>Atlas A2 训练系列产品/Atlas A2 推理系列产品</term>：不支持 | 12 | - <term>Atlas A2 训练系列产品/Atlas A2 推理系列产品</term>：不支持 |
| 13 | <!-- end id3 --> | ||
| 14 | <!-- npu="310b" id4 --> | ||
| 8 | - <term>Atlas 200I/500 A2 推理产品</term>：不支持 | 15 | - <term>Atlas 200I/500 A2 推理产品</term>：不支持 |
| 16 | <!-- end id4 --> | ||
| 17 | <!-- npu="310p" id5 --> | ||
| 9 | - <term>Atlas 推理系列产品</term>：不支持 | 18 | - <term>Atlas 推理系列产品</term>：不支持 |
| 19 | <!-- end id5 --> | ||
| 20 | <!-- npu="910" id6 --> | ||
| 10 | - <term>Atlas 训练系列产品</term>：不支持 | 21 | - <term>Atlas 训练系列产品</term>：不支持 |
| 22 | <!-- end id6 --> | ||
| 11 | 23 | ||
| 12 | ## 功能说明 | 24 | ## 功能说明 |
| 13 | 25 | ||
| 14 | - 接口功能： | 26 | - 接口功能： |
| 15 | 27 | ||
| 16 | 对输入 `x` 执行 SwiGLU 激活，支持可选的逐 token 权重和 count 模式分组。底层封装 `aclnnSwigluGroup`。 | 28 | 对输入`x`执行SwiGLU激活，支持可选的逐token权重和count模式分组。底层封装`aclnnSwigluGroup`。 |
| 17 | 29 | ||
| 18 | - 计算公式： | 30 | - 计算公式： |
| 19 | 31 | ||
| 20 | 令输入 `x` 最后一维大小为 `D`，`A = x[..., :D/2]`，`B = x[..., D/2:]`。 | 32 | 令输入`x`最后一维大小为`D`，`A = x[..., :D/2]`，`B = x[..., D/2:]`。 |
| 21 | 33 | ||
| 22 | $$ | 34 | $$ |
| 23 | y = silu(A) \times B = \frac{A}{1 + e^{-A}} \times B | 35 | y = silu(A) \times B = \frac{A}{1 + e^{-A}} \times B |
| 24 | $$ | 36 | $$ |
| 25 | 37 | ||
| 26 | 当 `clamp_limit > 0` 时，先对 `A` 和 `B` 做截断；当 `weight` 非空时，对结果逐 token 乘以 `weight`。 | 38 | 当`clamp_limit > 0`时，先对`A`和`B`做截断；当`weight`非空时，对结果逐token乘以 `weight`。 |
| 27 | 39 | ||
| 28 | ## 函数原型 | 40 | ## 函数原型 |
| 29 | 41 | ||
| 30 | ```python | 42 | ```python |
| 31 | torch.ops.cann_ops_nn.swiglu_group( | 43 | cann_ops_nn.swiglu_group( |
 | |||
| 32 | x, | 44 | x, |
| 33 | *, | 45 | *, |
| 34 | weight=None, | 46 | weight=None, |
| @@ -41,24 +53,24 @@ torch.ops.cann_ops_nn.swiglu_group( | |||
| 41 | 53 | ||
| 42 | | 参数名 | 参数类型 | 可选/必选 | 描述 | 数据类型 | 维度(shape) | | 54 | | 参数名 | 参数类型 | 可选/必选 | 描述 | 数据类型 | 维度(shape) | |
| 43 | | --- | --- | --- | --- | --- | --- | | 55 | | --- | --- | --- | --- | --- | --- | |
| 44 | | `x` | Tensor | 必选 | SwiGLU 输入，最后一维会被均分为两部分。 | `torch.float16`、`torch.bfloat16`、`torch.float32` | 1-8维，最后一维为偶数 | | 56 | | `x` | Tensor | 必选 | SwiGLU输入，最后一维会被均分为两部分。 | `torch.float16`、`torch.bfloat16`、`torch.float32` | 1-8维，最后一维为偶数 | |
| 45 | | `weight` | Tensor | 可选 | 逐 token 权重，非空时乘到 SwiGLU 结果上。 | `torch.float32` | 元素个数等于 `x` 除最后一维外的元素个数 | | 57 | | `weight` | Tensor | 可选 | 逐token权重，非空时乘到SwiGLU结果上。 | `torch.float32` | 元素个数等于`x`除最后一维外的元素个数 | |
| 46 | | `group_index` | Tensor | 可选 | count 模式分组 token 数。 | `torch.int64` | 1维 | | 58 | | `group_index` | Tensor | 可选 | count模式分组token 数。 | `torch.int64` | 1维 | |
| 47 | | `clamp_limit` | float | 可选 | 激活前截断阈值，默认 `-1.0` 表示不启用截断。 | - | - | | 59 | | `clamp_limit` | float | 可选 | 激活前截断阈值，默认`-1.0`表示不启用截断。 | - | - | |
| 48 | 60 | ||
| 49 | ## 返回值说明 | 61 | ## 返回值说明 |
| 50 | 62 | ||
| 51 | | 参数名 | 参数类型 | 描述 | 数据类型 | 维度(shape) | | 63 | | 参数名 | 参数类型 | 描述 | 数据类型 | 维度(shape) | |
| 52 | | --- | --- | --- | --- | --- | | 64 | | --- | --- | --- | --- | --- | |
| 53 | | `y` | Tensor | SwiGLU 激活结果。 | 与 `x` 相同 | 与 `x` 相同，但最后一维为 `x.shape[-1] // 2` | | 65 | | `y` | Tensor | SwiGLU激活结果。 | 与`x`相同 | 与`x`相同，但最后一维为`x.shape[-1] // 2` | |
| 54 | 66 | ||
| 55 | ## 约束说明 | 67 | ## 约束说明 |
| 56 | 68 | ||
| 57 | - 该接口支持单算子模式和 TorchAir 图模式调用。 | 69 | - 该接口支持单算子模式和TorchAir图模式调用。 |
| 58 | - `x`、`weight`、`group_index` 均需为 NPU Tensor；可选 Tensor 可以传 `None`。 | 70 | - `x`、`weight`、`group_index`均需为NPU Tensor；可选Tensor可以传 `None`。 |
| 59 | - `x` 的最后一维必须为偶数。 | 71 | - `x`的最后一维必须为偶数。 |
| 60 | - 不支持空 Tensor 和非连续 Tensor。 | 72 | - 不支持空Tensor和非连续Tensor。 |
| 61 | - `group_index` 表示每个 group 的 token 数，所有元素需大于等于 0。 | 73 | - `group_index`表示每个group的token数，所有元素需大于等于0。 |
| 62 | 74 | ||
| 63 | ## 确定性计算 | 75 | ## 确定性计算 |
| 64 | 76 | ||
| @@ -1,31 +1,43 @@ | |||
| 1 | # cann_ops_nn.swiglu_group_quant | 1 | # swiglu_group_quant |
| 2 | 2 | ||
| 3 | ## 产品支持情况 | 3 | ## 产品支持情况 |
| 4 | 4 | ||
| 5 | <!-- npu="950" id1 --> | ||
| 5 | - <term>Ascend 950PR/Ascend 950DT</term>：支持 | 6 | - <term>Ascend 950PR/Ascend 950DT</term>：支持 |
| 7 | <!-- end id1 --> | ||
| 8 | <!-- npu="A3" id2 --> | ||
| 6 | - <term>Atlas A3 训练系列产品/Atlas A3 推理系列产品</term>：不支持 | 9 | - <term>Atlas A3 训练系列产品/Atlas A3 推理系列产品</term>：不支持 |
| 10 | <!-- end id2 --> | ||
| 11 | <!-- npu="910b" id3 --> | ||
| 7 | - <term>Atlas A2 训练系列产品/Atlas A2 推理系列产品</term>：不支持 | 12 | - <term>Atlas A2 训练系列产品/Atlas A2 推理系列产品</term>：不支持 |
| 13 | <!-- end id3 --> | ||
| 14 | <!-- npu="310b" id4 --> | ||
| 8 | - <term>Atlas 200I/500 A2 推理产品</term>：不支持 | 15 | - <term>Atlas 200I/500 A2 推理产品</term>：不支持 |
| 16 | <!-- end id4 --> | ||
| 17 | <!-- npu="310p" id5 --> | ||
| 9 | - <term>Atlas 推理系列产品</term>：不支持 | 18 | - <term>Atlas 推理系列产品</term>：不支持 |
| 19 | <!-- end id5 --> | ||
| 20 | <!-- npu="910" id6 --> | ||
| 10 | - <term>Atlas 训练系列产品</term>：不支持 | 21 | - <term>Atlas 训练系列产品</term>：不支持 |
| 22 | <!-- end id6 --> | ||
| 11 | 23 | ||
| 12 | ## 功能说明 | 24 | ## 功能说明 |
| 13 | 25 | ||
| 14 | - 接口功能： | 26 | - 接口功能： |
| 15 | 27 | ||
| 16 | 对输入 `x` 执行 SwiGLU 激活后进行分组低比特量化，支持 Block FP8、MX FP8、MX FP4、HiFloat8 静态量化和 HiFloat8 动态量化。底层封装 `aclnnSwigluGroupQuant`。 | 28 | 对输入`x`执行SwiGLU激活后进行分组低比特量化，支持Block FP8、MX FP8、MX FP4、HiFloat8静态量化和HiFloat8动态量化。底层封装 `aclnnSwigluGroupQuant`。 |
| 17 | 29 | ||
| 18 | - 计算流程： | 30 | - 计算流程： |
| 19 | 31 | ||
| 20 | 1. 将 `x` 的最后一维均分为 `A`、`B` 两部分，计算 `y_origin = silu(A) * B`。 | 32 | 1. 将`x`的最后一维均分为`A`、`B`两部分，计算`y_origin = silu(A) * B`。 |
| 21 | 2. 当 `weight` 非空时，对 `y_origin` 逐 token 乘以 `weight`。 | 33 | 2. 当`weight`非空时，对`y_origin`逐token乘以`weight`。 |
| 22 | 3. 按 `quant_mode` 对 `y_origin` 量化，输出 `y` 和 `y_scale`。 | 34 | 3. 按`quant_mode`对`y_origin`量化，输出`y`和`y_scale`。 |
| 23 | 4. 当 `output_origin=True` 时，额外返回量化前的 `y_origin`；否则返回 shape 为 `[0]` 的占位 Tensor。 | 35 | 4. 当`output_origin=True`时，额外返回量化前的`y_origin`；否则返回shape为`[0]`的占位Tensor。 |
| 24 | 36 | ||
| 25 | ## 函数原型 | 37 | ## 函数原型 |
| 26 | 38 | ||
| 27 | ```python | 39 | ```python |
| 28 | torch.ops.cann_ops_nn.swiglu_group_quant( | 40 | cann_ops_nn.swiglu_group_quant( |
| 29 | x, | 41 | x, |
| 30 | *, | 42 | *, |
| 31 | weight=None, | 43 | weight=None, |
| @@ -45,68 +57,68 @@ torch.ops.cann_ops_nn.swiglu_group_quant( | |||
| 45 | 57 | ||
| 46 | | 参数名 | 参数类型 | 可选/必选 | 描述 | 数据类型 | 维度(shape) | | 58 | | 参数名 | 参数类型 | 可选/必选 | 描述 | 数据类型 | 维度(shape) | |
| 47 | | --- | --- | --- | --- | --- | --- | | 59 | | --- | --- | --- | --- | --- | --- | |
| 48 | | `x` | Tensor | 必选 | SwiGLU 输入，最后一维会被均分为两部分。 | `torch.float16`、`torch.bfloat16`、`torch.float32` | 2-8维 | | 60 | | `x` | Tensor | 必选 | SwiGLU输入，最后一维会被均分为两部分。 | `torch.float16`、`torch.bfloat16`、`torch.float32` | 1-7维 | |
| 49 | | `weight` | Tensor | 可选 | 逐 token 权重，非空时乘到量化前结果上。 | `torch.float32` | 1-8维，元素个数等于 `x` 除最后一维外的元素个数 | | 61 | | `weight` | Tensor | 可选 | 逐token权重，非空时乘到量化前结果上。 | `torch.float32` | 1维，元素个数等于 `x` 除最后一维外的元素个数 | |
| 50 | | `group_index` | Tensor | 可选 | count 模式分组 token 数。 | `torch.int64` | 1维 | | 62 | | `group_index` | Tensor | 可选 | count模式分组token数。 | `torch.int64` | 1维 | |
| 51 | | `scale` | Tensor | 可选 | HiFloat8 静态量化使用的 scale。 | `torch.float32` | 1维 | | 63 | | `scale` | Tensor | 可选 | HiFloat8静态量化使用的scale。 | `torch.float32` | 1维 | |
| 52 | | `dst_type` | int | 可选 | 目标量化类型的 torch dtype 编码，默认 `291`。 | - | - | | 64 | | `dst_type` | int | 可选 | 目标量化类型的torch dtype编码，默认`291`。 | - | - | |
| 53 | | `quant_mode` | int | 可选 | 量化模式，支持 `0`、`1`、`2`、`3`。 | - | - | | 65 | | `quant_mode` | int | 可选 | 量化模式，支持`0`、`1`、`2`、`3`。 | - | - | |
| 54 | | `block_size` | int | 可选 | 量化块大小，`0` 表示使用模式默认值。 | - | - | | 66 | | `block_size` | int | 可选 | 量化块大小，`0`表示使用模式默认值。 | - | - | |
| 55 | | `round_scale` | bool | 可选 | MX 量化是否将 scale 舍入为 2 的幂。 | - | - | | 67 | | `round_scale` | bool | 可选 | MX量化是否将scale舍入为2的幂。 | - | - | |
| 56 | | `clamp_limit` | float | 可选 | SwiGLU 计算前的截断阈值，默认 `-1.0` 表示不启用截断。 | - | - | | 68 | | `clamp_limit` | float | 可选 | SwiGLU计算前的截断阈值，默认`-1.0`表示不启用截断。 | - | - | |
| 57 | | `dst_type_max` | float | 可选 | HiFloat8 动态量化计算 scale 时使用的最大有限值。 | - | - | | 69 | | `dst_type_max` | float | 可选 | HiFloat8动态量化计算 scale时使用的最大有限值。 | - | - | |
| 58 | | `output_origin` | bool | 可选 | 是否返回量化前的 SwiGLU 结果。 | - | - | | 70 | | `output_origin` | bool | 可选 | 是否返回量化前的SwiGLU结果。 | - | - | |
| 59 | 71 | ||
| 60 | ### quant_mode 与 dst_type | 72 | ### quant_mode 与 dst_type |
| 61 | 73 | ||
| 62 | | `quant_mode` | 含义 | `dst_type` 支持值 | `y` 的 torch dtype | `y_scale` 的 torch dtype | | 74 | | `quant_mode` | 含义 | `dst_type` 支持值 | `y` 的 torch dtype | `y_scale` 的 torch dtype | |
| 63 | | --- | --- | --- | --- | --- | | 75 | | --- | --- | --- | --- | --- | |
| 64 | | `0` | Block FP8 | `23`/`291` 表示 `float8_e5m2`；`24`/`292` 表示 `float8_e4m3fn` | `torch.float8_e5m2` 或 `torch.float8_e4m3fn` | `torch.float32` | | 76 | | `0` | Block FP8 | `23`/`291`表示`float8_e5m2`；`24`/`292`表示`float8_e4m3fn` | `torch.float8_e5m2`或`torch.float8_e4m3fn` | `torch.float32` | |
| 65 | | `1` | MX FP8 / MX FP4 | `23`、`24`、`291`、`292`、`296`、`297` | FP8 使用 torch FP8 dtype；FP4 在 eager/ACL Graph 路径使用 `torch.uint8` 打包存储。TorchAir GE 图模式中，`296` 输出类型为 `torch.float4_e2m1fn_x2`，`297` 输出类型为 `torch.uint8` | `torch.float8_e8m0fnu` | | 77 | | `1` | MX FP8 / MX FP4 | `23`、`24`、`291`、`292`、`296`、`297` | FP8使用torch FP8 dtype；FP4在eager/图模式（aclgraph）路径使用`torch.uint8`打包存储。TorchAir GE图模式中`296`输出类型为`torch.float4_e2m1fn_x2`，`297`输出类型为`torch.uint8` | `torch.float8_e8m0fnu` | |
| 66 | | `2` | HiFloat8 静态量化 | 建议使用 `290` 表示 HiFloat8；当前 torch 接口实际下发 HIFLOAT8 | `torch.uint8` | `torch.float32`，shape 为 `[0]` | | 78 | | `2` | HiFloat8 静态量化 | `290`表示HiFloat8 | `torch.uint8` | `torch.float32`，shape为`[0]` | |
| 67 | | `3` | HiFloat8 动态量化 | 建议使用 `290` 表示 HiFloat8；当前 torch 接口实际下发 HIFLOAT8 | `torch.uint8` | `torch.float32` | | 79 | | `3` | HiFloat8 动态量化 | `290`表示HiFloat8 | `torch.uint8` | `torch.float32` | |
| 68 | 80 | ||
| 69 | #### dst_type 编码说明 | 81 | #### dst_type 编码说明 |
| 70 | 82 | ||
| 71 | | `dst_type` | 对应类型 | 来源 | 下发到 GE/ACL 的类型 | 适用 `quant_mode` | | 83 | | `dst_type` | 对应类型 | 来源 | 下发到 GE/ACL 的类型 | 适用 `quant_mode` | |
| 72 | | --- | --- | --- | --- | --- | | 84 | | --- | --- | --- | --- | --- | |
| 73 | | `23` | `torch.float8_e5m2` | PyTorch 原生 dtype int 值 | `DT_FLOAT8_E5M2` / `ACL_FLOAT8_E5M2` | `0`、`1` | | 85 | | `23` | `torch.float8_e5m2` | PyTorch原生dtype int值 | `DT_FLOAT8_E5M2` / `ACL_FLOAT8_E5M2` | `0`、`1` | |
| 74 | | `24` | `torch.float8_e4m3fn` | PyTorch 原生 dtype int 值 | `DT_FLOAT8_E4M3FN` / `ACL_FLOAT8_E4M3FN` | `0`、`1` | | 86 | | `24` | `torch.float8_e4m3fn` | PyTorch原生dtype int值 | `DT_FLOAT8_E4M3FN` / `ACL_FLOAT8_E4M3FN` | `0`、`1` | |
| 75 | | `291` | `torch_npu.float8_e5m2`，语义同 `torch.float8_e5m2` | torch_npu 扩展 dtype 编码 | `DT_FLOAT8_E5M2` / `ACL_FLOAT8_E5M2` | `0`、`1` | | 87 | | `291` | `torch_npu.float8_e5m2`，语义同`torch.float8_e5m2` | torch_npu扩展dtype编码 | `DT_FLOAT8_E5M2` / `ACL_FLOAT8_E5M2` | `0`、`1` | |
| 76 | | `292` | `torch_npu.float8_e4m3fn`，语义同 `torch.float8_e4m3fn` | torch_npu 扩展 dtype 编码 | `DT_FLOAT8_E4M3FN` / `ACL_FLOAT8_E4M3FN` | `0`、`1` | | 88 | | `292` | `torch_npu.float8_e4m3fn`，语义同`torch.float8_e4m3fn` | torch_npu扩展dtype编码 | `DT_FLOAT8_E4M3FN` / `ACL_FLOAT8_E4M3FN` | `0`、`1` | |
| 77 | | `290` | `torch_npu.hifloat8` | torch_npu 扩展 dtype 编码 | `DT_HIFLOAT8` / `ACL_HIFLOAT8` | `2`、`3` | | 89 | | `290` | `torch_npu.hifloat8` | torch_npu扩展dtype编码 | `DT_HIFLOAT8` / `ACL_HIFLOAT8` | `2`、`3` | |
| 78 | | `296` | `torch_npu.float4_e2m1fn_x2` | torch_npu 扩展 dtype 编码 | `DT_FLOAT4_E2M1` / `ACL_FLOAT4_E2M1` | `1` | | 90 | | `296` | `torch_npu.float4_e2m1fn_x2` | torch_npu扩展dtype编码 | `DT_FLOAT4_E2M1` / `ACL_FLOAT4_E2M1` | `1` | |
| 79 | | `297` | `torch_npu.float4_e1m2fn_x2` | torch_npu 扩展 dtype 编码 | `DT_FLOAT4_E1M2` / `ACL_FLOAT4_E1M2` | `1` | | 91 | | `297` | `torch_npu.float4_e1m2fn_x2` | torch_npu扩展dtype编码 | `DT_FLOAT4_E1M2` / `ACL_FLOAT4_E1M2` | `1` | |
| 80 | 92 | ||
| 81 | 说明：`quant_mode=2/3` 为 HiFloat8 模式，实际下发为 `DT_HIFLOAT8` / `ACL_HIFLOAT8`。graph_convert 会把上表中的 torch dtype 编码转换为 GE dtype attr，ACLNN 路径会把对应编码转换为 `aclDataType` 后调用底层接口。 | 93 | 说明：`quant_mode=2/3`为HiFloat8模式，实际下发为`DT_HIFLOAT8`/`ACL_HIFLOAT8`。graph_convert会把上表中的torch dtype编码转换为GE dtype attr，ACLNN路径会把对应编码转换为`aclDataType`后调用底层接口。 |
| 82 | 94 | ||
| 83 | ## 返回值说明 | 95 | ## 返回值说明 |
| 84 | 96 | ||
| 85 | | 参数名 | 参数类型 | 描述 | 数据类型 | 维度(shape) | | 97 | | 参数名 | 参数类型 | 描述 | 数据类型 | 维度(shape) | |
| 86 | | --- | --- | --- | --- | --- | | 98 | | --- | --- | --- | --- | --- | |
| 87 | | `y` | Tensor | 量化输出。 | 见 `quant_mode 与 dst_type` | FP8/HiFloat8 为 `x.shape[:-1] + [D/2]`；FP4 为 `x.shape[:-1] + [ceil((D/2)/2)]`，在 `D` 可被 256 整除时等价于 `x.shape[:-1] + [D/4]` | | 99 | | `y` | Tensor | 量化输出。 | 参见`quant_mode 与 dst_type` | FP8/HiFloat8为`x.shape[:-1] + [D/2]`；FP4为`x.shape[:-1] + [ceil((D/2)/2)]`，在`D`可被256整除时等价于`x.shape[:-1] + [D/4]` | |
| 88 | | `y_scale` | Tensor | 量化 scale 输出。 | 见 `quant_mode 与 dst_type` | `quant_mode=0` 为 `x.shape[:-1] + [ceil((D/2)/128)]`；`quant_mode=1` 为 `x.shape[:-1] + [ceil(ceil((D/2)/32)/2), 2]`；`quant_mode=2` 为 `[0]`；`quant_mode=3` 为 `group_index.shape` 或 `[1]` | | 100 | | `y_scale` | Tensor | 量化scale输出。 | 参见`quant_mode 与 dst_type` | `quant_mode=0`为`x.shape[:-1] + [ceil((D/2)/128)]`；`quant_mode=1`为`x.shape[:-1] + [ceil(ceil((D/2)/32)/2), 2]`；`quant_mode=2`为`[0]`；`quant_mode=3`为`group_index.shape`或`[1]` | |
| 89 | | `y_origin` | Tensor | 量化前 SwiGLU 结果或占位 Tensor。 | 与 `x` 相同 | `output_origin=True` 时为 `x.shape[:-1] + [D/2]`，否则为 `[0]` | | 101 | | `y_origin` | Tensor | 量化前SwiGLU结果或占位Tensor。 | 与`x`相同 | `output_origin=True`时为`x.shape[:-1] + [D/2]`，否则为`[0]` | |
| 90 | 102 | ||
| 91 | 其中 `D = x.shape[-1]`。 | 103 | 其中 `D = x.shape[-1]`。 |
| 92 | 104 | ||
| 93 | ## 约束说明 | 105 | ## 约束说明 |
| 94 | 106 | ||
| 95 | - 该接口支持单算子模式和 TorchAir 图模式调用。 | 107 | - 该接口支持单算子模式和TorchAir图模式调用。 |
| 96 | - `x`、`weight`、`group_index`、`scale` 均需为 NPU Tensor；可选 Tensor 可以传 `None`。 | 108 | - `x`、`weight`、`group_index`、`scale`均需为NPU Tensor；可选Tensor可以传 `None`。 |
| 97 | - 输入 `x` 的 rank 必须大于 0，最后一维 `D` 必须大于等于 256 且能被 256 整除。 | 109 | - 输入`x`的rank 必须大于0，最后一维`D`必须大于等于256且能被256整除。 |
| 98 | - `dst_type` 支持 FP8、FP4 和 HiFloat8 对应的 torch dtype 编码，详见 `dst_type 编码说明`。 | 110 | - `dst_type`支持FP8、FP4和HiFloat8对应的torch dtype编码，详见`dst_type 编码说明`。 |
| 99 | - `quant_mode=0` 时仅支持 FP8 输出，`dst_type` 支持 `23`、`24`、`291`、`292`，`block_size` 支持 `0` 或 `128`。 | 111 | - `quant_mode=0`时仅支持FP8输出，`dst_type`支持`23`、`24`、`291`、`292`，`block_size`支持`0`或`128`。 |
| 100 | - `quant_mode=1` 时支持 FP8/FP4 输出，`dst_type` 支持 `23`、`24`、`291`、`292`、`296`、`297`，`block_size` 支持 `0` 或 `32`，`round_scale` 必须为 `True`。 | 112 | - `quant_mode=1`时支持FP8/FP4 输出，`dst_type`支持`23`、`24`、`291`、`292`、`296`、`297`，`block_size`支持`0`或`32`，`round_scale`必须为`True`。 |
| 101 | - `quant_mode=2` 时支持 HiFloat8 静态量化输出，需传入 `scale`，`dst_type`、`block_size` 和 `round_scale` 不生效，实际下发 HiFloat8。 | 113 | - `quant_mode=2`时支持HiFloat8静态量化输出，需传入`scale`，`dst_type`、`block_size`和`round_scale`不生效，实际下发HiFloat8。 |
| 102 | - `quant_mode=3` 时支持 HiFloat8 动态量化输出，不使用 `scale`，`dst_type`、`block_size` 和 `round_scale` 不生效，实际下发 HiFloat8。 | 114 | - `quant_mode=3`时支持HiFloat8动态量化输出，不使用`scale`，`dst_type`、`block_size`和`round_scale`不生效，实际下发 HiFloat8。 |
| 103 | - `dst_type` 为 `296` 或 `297`，即 `FLOAT4_E2M1` 或 `FLOAT4_E1M2` 时，必须使用 `quant_mode=1`。 | 115 | - `dst_type`为`296`或`297`，即`FLOAT4_E2M1`或`FLOAT4_E1M2`时，必须使用`quant_mode=1`。 |
| 104 | - `y_scale` 的数据类型必须与 `quant_mode` 匹配：Block FP8 为 `torch.float32`，MX 为 `torch.float8_e8m0fnu`，HiFloat8 为 `torch.float32`。 | 116 | - `y_scale`的数据类型必须与`quant_mode`匹配：Block FP8为`torch.float32`，MX为`torch.float8_e8m0fnu`，HiFloat8为`torch.float32`。 |
| 105 | - `quant_mode=3` 时，`group_index` 可用于 MoE 场景的分组动态量化；`y_scale` 的 shape 为 `group_index.shape`，未传 `group_index` 时为 `[1]`。 | 117 | - `quant_mode=3`时，`group_index`可用于MoE场景的分组动态量化；`y_scale`的shape为`group_index.shape`，未传`group_index`时为`[1]`。 |
| 106 | - `clamp_limit` 不启用时使用默认占位值 `-1.0`；启用时必须大于 0。 | 118 | - `clamp_limit`不启用时使用默认占位值`-1.0`；启用时必须大于0。 |
| 107 | - quant_mode=0和quant_mode=1时，output_origin仅支持False。 | 119 | - quant_mode=0和quant_mode=1时，output_origin仅支持False。 |
| 108 | - group_index中的元素值须大于等于0。 | 120 | - group_index中的元素值须大于等于0。 |
| 109 | - 不支持空 Tensor 和非连续 Tensor。 | 121 | - 不支持空Tensor和非连续Tensor。 |
| 110 | 122 | ||
| 111 | ## 确定性计算 | 123 | ## 确定性计算 |
| 112 | 124 | ||
| @@ -20,32 +20,32 @@ CANN算子量化是指对神经网络中Matmul等矩阵（cube）类算子的输 | |||
| 20 | 20 | ||
| 21 | 假设左矩阵shape为(m, k)，右矩阵shape为(k, n)，k为reduce轴，生成量化参数的shape为(1, )。 | 21 | 假设左矩阵shape为(m, k)，右矩阵shape为(k, n)，k为reduce轴，生成量化参数的shape为(1, )。 |
| 22 | 22 | ||
| 23 |  | 23 |  |
| 24 | 24 | ||
| 25 | - perchannel量化（简称C量化）：量化对象是右矩阵，每个channel分别使用独立的量化参数。 | 25 | - perchannel量化（简称C量化）：量化对象是右矩阵，每个channel分别使用独立的量化参数。 |
| 26 | 26 | ||
| 27 | 假设右矩阵shape为(k, n)，k为reduce轴，生成量化参数的shape为(n, )。 | 27 | 假设右矩阵shape为(k, n)，k为reduce轴，生成量化参数的shape为(n, )。 |
| 28 | 28 | ||
| 29 |  | 29 |  |
| 30 | 30 | ||
| 31 | - pertoken量化（简称K量化）：量化对象是左矩阵，每个token分别使用独立的量化参数。 | 31 | - pertoken量化（简称K量化）：量化对象是左矩阵，每个token分别使用独立的量化参数。 |
| 32 | 32 | ||
| 33 | 假设左矩阵shape为(m, k)，k为reduce轴，生成量化参数的shape为(m, )。 | 33 | 假设左矩阵shape为(m, k)，k为reduce轴，生成量化参数的shape为(m, )。 |
| 34 | 34 | ||
| 35 |  | 35 |  |
| 36 | 36 | ||
| 37 | - pergroup量化（简称G量化）：量化对象既可以是左矩阵，也可以是右矩阵，在reduce轴上对数据分组，每组使用独立的量化参数。 | 37 | - pergroup量化（简称G量化）：量化对象既可以是左矩阵，也可以是右矩阵，在reduce轴上对数据分组，每组使用独立的量化参数。 |
| 38 | - 假设左矩阵shape为(m, k)，k为reduce轴，在k轴上分组，group size为gs，生成量化参数的shape为(m, k/gs)。 | 38 | - 假设左矩阵shape为(m, k)，k为reduce轴，在k轴上分组，group size为gs，生成量化参数的shape为(m, k/gs)。 |
| 39 | - 假设右矩阵shape为(k, n)，k为reduce轴，在k轴上分组，group size为gs，生成量化参数的shape为(k/gs, n)。 | 39 | - 假设右矩阵shape为(k, n)，k为reduce轴，在k轴上分组，group size为gs，生成量化参数的shape为(k/gs, n)。 |
| 40 | 40 | ||
| 41 |  | 41 |  |
| 42 | 42 | ||
| 43 | - perblock量化（简称B量化）：量化对象既可以是左矩阵，也可以是右矩阵，在所有轴上对数据分块，每块使用独立的量化参数。 | 43 | - perblock量化（简称B量化）：量化对象既可以是左矩阵，也可以是右矩阵，在所有轴上对数据分块，每块使用独立的量化参数。 |
| 44 | 44 | ||
| 45 | - 假设左矩阵shape为(m, k)，k为reduce轴，在m、k轴上分别按(bs, bs)块对数据分组，bs为block size，生成量化参数的shape为(m/bs, k/bs)。 | 45 | - 假设左矩阵shape为(m, k)，k为reduce轴，在m、k轴上分别按(bs, bs)块对数据分组，bs为block size，生成量化参数的shape为(m/bs, k/bs)。 |
| 46 | - 假设右矩阵shape为(k, n)，k为reduce轴，在k、n轴上分别按(bs, bs)块对数据分组，bs为block size，生成量化参数的shape为(k/bs, n/bs)。 | 46 | - 假设右矩阵shape为(k, n)，k为reduce轴，在k、n轴上分别按(bs, bs)块对数据分组，bs为block size，生成量化参数的shape为(k/bs, n/bs)。 |
| 47 | 47 | ||
| 48 |  | 48 |  |
| 49 | 49 | ||
| 50 | ## 常见组合量化 | 50 | ## 常见组合量化 |
| 51 | 51 | ||
文件重命名但无更改。
文件重命名但无更改。
文件重命名但无更改。
文件重命名但无更改。
文件重命名但无更改。
| @@ -1,23 +1,35 @@ | |||||||
| 1 | # cann_ops_nn.quant_matmul_activation_quant | 1 | # quant_matmul_activation_quant | ||||
| 2 | 2 | ||||||
| 3 | ## 产品支持情况 | 3 | ## 产品支持情况 | ||||
| 4 | 4 | ||||||
| 5 | <!-- npu="950" id1 --> | ||||||
| 5 | - <term>Ascend 950PR/Ascend 950DT</term>：支持 | 6 | - <term>Ascend 950PR/Ascend 950DT</term>：支持 | ||||
| 7 | <!-- end id1 --> | ||||||
| 8 | <!-- npu="A3" id2 --> | ||||||
| 6 | - <term>Atlas A3 训练系列产品/Atlas A3 推理系列产品</term>：不支持 | 9 | - <term>Atlas A3 训练系列产品/Atlas A3 推理系列产品</term>：不支持 | ||||
| 10 | <!-- end id2 --> | ||||||
| 11 | <!-- npu="910b" id3 --> | ||||||
| 7 | - <term>Atlas A2 训练系列产品/Atlas A2 推理系列产品</term>：不支持 | 12 | - <term>Atlas A2 训练系列产品/Atlas A2 推理系列产品</term>：不支持 | ||||
| 13 | <!-- end id3 --> | ||||||
| 14 | <!-- npu="310b" id4 --> | ||||||
| 8 | - <term>Atlas 200I/500 A2 推理产品</term>：不支持 | 15 | - <term>Atlas 200I/500 A2 推理产品</term>：不支持 | ||||
| 16 | <!-- end id4 --> | ||||||
| 17 | <!-- npu="310p" id5 --> | ||||||
| 9 | - <term>Atlas 推理系列产品</term>：不支持 | 18 | - <term>Atlas 推理系列产品</term>：不支持 | ||||
| 19 | <!-- end id5 --> | ||||||
| 20 | <!-- npu="910" id6 --> | ||||||
| 10 | - <term>Atlas 训练系列产品</term>：不支持 | 21 | - <term>Atlas 训练系列产品</term>：不支持 | ||||
| 22 | <!-- end id6 --> | ||||||
| 11 | 23 | ||||||
| 12 | ## 功能说明 | 24 | ## 功能说明 | ||||
| 13 | 25 | ||||||
| 14 | - 接口功能： | 26 | - 接口功能： | ||||
| 15 | 27 | ||||||
| 16 | 融合量化的矩阵乘、激活以及动态量化，封装 `aclnnQuantMatmulActivationQuantWeightNz`。当前支持激活为 gelu、MX 量化模式。输入 `x1`、`x2` 为 FP8 量化矩阵，必选输入 `x2_scale`、可选输入 `x1_scale` 为 MX 量化缩放因子，`bias` 为偏置项；矩阵乘结果经激活函数后做动态量化，输出量化结果 `y` 和量化尺度 `y_scale`。M/N/K 维度及转置标志由 `x1`、`x2` 最后两维自动匹配推导，无需显式传入。 | 28 | 融合量化的矩阵乘、激活以及动态量化，封装`aclnnQuantMatmulActivationQuantWeightNz`。当前支持激活为gelu、MX量化模式。输入`x1`、`x2`为P8量化矩阵，必选输入`x2_scale`、可选输入`x1_scale`为MX量化缩放因子，`bias`为偏置项；矩阵乘结果经激活函数后做动态量化，输出量化结果`y`和量化尺度`y_scale`。M/N/K维度及转置标志由`x1`、`x2`最后两维自动匹配推导，无需显式传入。 | ||||
🟡 Medium Priority 在功能说明段落中，原文"FP8量化矩阵"（Float Point 8）被误改为"P8量化矩阵"。该文档其余各处（第74、75、82、88、90、127行）均一致使用"FP8"，此处"P8"为排版清理过程中意外丢失"F"导致的拼写错误。用户阅读时会误解为某种不存在的"P8"数据类型。 建议：将"P8"修正为"FP8"，恢复丢失的字符"F"。 改动建议
  | |||||||
| 17 | 29 | ||||||
| 18 | - 计算公式： | 30 | - 计算公式： | ||||
| 19 | 31 | ||||||
| 20 | - QuantMatmul MX 量化模式： | 32 | - QuantMatmul MX量化模式： | ||||
| 21 | 33 | ||||||
| 22 | $$ | 34 | $$ | ||||
| 23 | matmulOut[m,n] = \sum_{j=0}^{kLoops-1} ((\sum_{k=0}^{gsK-1} (x1Slice * x2Slice))* (x1Scale[m/gsM, j] * x2Scale[j, n/gsN]))+bias[n] | 35 | matmulOut[m,n] = \sum_{j=0}^{kLoops-1} ((\sum_{k=0}^{gsK-1} (x1Slice * x2Slice))* (x1Scale[m/gsM, j] * x2Scale[j, n/gsN]))+bias[n] | ||||
| @@ -39,7 +51,7 @@ | |||||||
| 39 | 51 | ||||||
| 40 | - 动态量化计算公式： | 52 | - 动态量化计算公式： | ||||
| 41 | 53 | ||||||
| 42 | - 场景 1，当 scale_alg 为 0 时： | 54 | - **场景1，当scale_alg为0时**： | ||||
| 43 | - 将输入 activationOut 在尾轴上按 $k = 32$ 个数分组，一组 k 个数 $\{\{V_i\}_{i=1}^{k}\}$ 动态量化为 $\{mxscale1, \{P_i\}_{i=1}^{k}\}， k = 32$ | 55 | - 将输入 activationOut 在尾轴上按 $k = 32$ 个数分组，一组 k 个数 $\{\{V_i\}_{i=1}^{k}\}$ 动态量化为 $\{mxscale1, \{P_i\}_{i=1}^{k}\}， k = 32$ | ||||
| 44 | 56 | ||||||
| 45 | $$ | 57 | $$ | ||||
| @@ -59,7 +71,7 @@ | |||||||
| 59 | | FLOAT8_E4M3FN | 8 | | 71 | | FLOAT8_E4M3FN | 8 | | ||||
| 60 | | FLOAT8_E5M2 | 15 | | 72 | | FLOAT8_E5M2 | 15 | | ||||
| 61 | 73 | ||||||
| 62 | - 场景 2，当 scale_alg 为 1 时，只涉及 FP8 类型： | 74 | - **场景2，当scale_alg为1时，只涉及FP8类型**： | ||||
| 63 | - 将输入activationOut在尾轴上按$k = 32$个数分块，对每块单独计算一个块缩放因子$S_{fp32}^b$，再把块内所有元素用同一个$S_{fp32}^b$映射到目标低精度类型FP8。如果最后一块不足$k = 32$个元素，把缺失值视为0，按照完整块处理。 | 75 | - 将输入activationOut在尾轴上按$k = 32$个数分块，对每块单独计算一个块缩放因子$S_{fp32}^b$，再把块内所有元素用同一个$S_{fp32}^b$映射到目标低精度类型FP8。如果最后一块不足$k = 32$个元素，把缺失值视为0，按照完整块处理。 | ||||
| 64 | - 找到该块中数值的最大绝对值： | 76 | - 找到该块中数值的最大绝对值： | ||||
| 65 | 77 | ||||||
| @@ -88,7 +100,7 @@ | |||||||
| 88 | ## 函数原型 | 100 | ## 函数原型 | ||||
| 89 | 101 | ||||||
| 90 | ```python | 102 | ```python | ||||
| 91 | torch.ops.cann_ops_nn.quant_matmul_activation_quant(x1, x2, x2_scale, *, x1_scale=None, bias=None, | 103 | cann_ops_nn.quant_matmul_activation_quant(x1, x2, x2_scale, *, x1_scale=None, bias=None, | ||||
| 92 | output_dtype=None, x1_dtype=None, x2_dtype=None, x1scale_dtype=None, x2scale_dtype=None, | 104 | output_dtype=None, x1_dtype=None, x2_dtype=None, x1scale_dtype=None, x2scale_dtype=None, | ||||
| 93 | group_sizes=None, activation_type="gelu_tanh", quant_mode="mx", round_mode="rint", | 105 | group_sizes=None, activation_type="gelu_tanh", quant_mode="mx", round_mode="rint", | ||||
| 94 | scale_alg=0, dst_type_max=0.0) -> (Tensor y, Tensor y_scale) | 106 | scale_alg=0, dst_type_max=0.0) -> (Tensor y, Tensor y_scale) | ||||
| @@ -98,51 +110,49 @@ torch.ops.cann_ops_nn.quant_matmul_activation_quant(x1, x2, x2_scale, *, x1_scal | |||||||
| 98 | 110 | ||||||
| 99 | | 参数名 | 参数类型 | 可选/必选 | 描述 | 数据类型 | 维度(shape) | | 111 | | 参数名 | 参数类型 | 可选/必选 | 描述 | 数据类型 | 维度(shape) | | ||||
| 100 | | --- | --- | --- | --- | --- | --- | | 112 | | --- | --- | --- | --- | --- | --- | | ||||
| 101 | | `x1` | Tensor | 必选 | 矩阵乘运算中的左矩阵。数据格式为 ND。最后两维为 `(M, K)` 或 `(K, M)`，由与 `x2` 的维度匹配自动推导。 | torch.float8_e4m3fn、torch.float8_e5m2 | 2-6 维，`(..., M, K)` 或 `(..., K, M)` | | 113 | | `x1` | Tensor | 必选 | 矩阵乘运算中的左矩阵。数据格式为ND。最后两维为`(M, K)`或`(K, M)`，由与`x2`的维度匹配自动推导。 | torch.float8_e4m3fn、torch.float8_e5m2 | 2-6维，`(..., M, K)`或`(..., K, M)` | | ||||
| 102 | | `x2` | Tensor | 必选 | 矩阵乘运算中的右矩阵。数据格式为 FRACTAL_NZ。最后两维为 `(K, N)` 或 `(N, K)`，由与 `x1` 的维度匹配自动推导。 | torch.float8_e4m3fn | 2-6 维，`(..., K, N)` 或 `(..., N, K)` | | 114 | | `x2` | Tensor | 必选 | 矩阵乘运算中的右矩阵。数据格式为FRACTAL_NZ。最后两维为`(K, N)`或`(N, K)`，由与`x1`的维度匹配自动推导。 | torch.float8_e4m3fn | 2-6维，`(..., K, N)`或`(..., N, K)` | | ||||
| 103 | | `x2_scale` | Tensor | 必选 | 矩阵乘计算时 x2 的 MX 量化缩放因子。数据格式为 ND。batch 维须与 `x2` 一致。 | torch.float8_e8m0fnu | `(..., K//64, N, 2)` 或 `(..., N, K//64, 2)`（随 `x2` 方向） | | 115 | | `x2_scale` | Tensor | 必选 | 矩阵乘计算时x2的MX量化缩放因子。数据格式为ND。batch维须与`x2`一致。 | torch.float8_e8m0fnu | `(..., K//64, N, 2)`或`(..., N, K//64, 2)`（随`x2 方向） | | ||||
🟡 Medium Priority
建议：在 改动建议
| |||||||
| 104 | | `x1_scale` | Tensor | 可选 | 矩阵乘计算时 x1 的 MX 量化缩放因子。数据格式为 ND。batch 维须与 `x1` 一致。 | torch.float8_e8m0fnu | `(..., M, K//64, 2)` 或 `(..., K//64, M, 2)`（随 `x1` 方向） | | 116 | | `x1_scale` | Tensor | 可选 | 矩阵乘计算时x1的MX量化缩放因子。数据格式为ND。batch维须与`x1`一致。 | torch.float8_e8m0fnu | `(..., M, K//64, 2)`或`(..., K//64, M, 2)`（随`x1`方向） | | ||||
| 105 | | `bias` | Tensor | 可选 | 矩阵乘运算后累加的偏置。数据格式为 ND。 | float32 | `(N,)` | | 117 | | `bias` | Tensor | 可选 | 矩阵乘运算后累加的偏置。数据格式为ND。 | float32 | `(N,)` | | ||||
| 106 | | `output_dtype` | int | 可选 | 输出 `y` 的数据类型枚举值。支持 torch.float8_e4m3fn、torch.float8_e5m2 等。默认值 None（等价于 0，表示与 `x1` 同类型）。 | int | - | | 118 | | `output_dtype` | int | 可选 | 输出`y`的数据类型枚举值。支持torch.float8_e4m3fn、torch.float8_e5m2等。默认值None（等价于 0，表示与`x1`同类型）。 | int | - | | ||||
| 107 | | `x1_dtype` | int | 可选 | `x1` 的数据类型枚举值。不传入时根据 `x1` 的 scalar_type 自动推导。 | int | - | | 119 | | `x1_dtype` | int | 可选 | `x1`的数据类型枚举值。不传入时根据`x1`的scalar_type自动推导。 | int | - | | ||||
| 108 | | `x2_dtype` | int | 可选 | `x2` 的数据类型枚举值。不传入时根据 `x2` 的 scalar_type 自动推导。 | int | - | | 120 | | `x2_dtype` | int | 可选 | `x2`的数据类型枚举值。不传入时根据`x2`的scalar_type自动推导。 | int | - | | ||||
| 109 | | `x1scale_dtype` | int | 可选 | `x1_scale` 的数据类型枚举值。不传入时根据 `x1_scale` 的 scalar_type 自动推导。 | int | - | | 121 | | `x1scale_dtype` | int | 可选 | `x1_scale`的数据类型枚举值。不传入时根据`x1_scale`的scalar_type自动推导。 | int | - | | ||||
| 110 | | `x2scale_dtype` | int | 可选 | `x2_scale` 的数据类型枚举值。不传入时根据 `x2_scale` 的 scalar_type 自动推导。 | int | - | | 122 | | `x2scale_dtype` | int | 可选 | `x2_scale`的数据类型枚举值。不传入时根据`x2_scale`的scalar_type自动推导。 | int | - | | ||||
| 111 | | `group_sizes` | List[int] | 可选 | 分组量化大小 `[groupSizeM, groupSizeN, groupSizeK]`，每个元素取值范围为 [0, 65535]。 | list | `(3,)` | | 123 | | `group_sizes` | List[int] | 可选 | 分组量化大小 `[groupSizeM, groupSizeN, groupSizeK]`，每个元素取值范围为[0, 65535]。 | list | `(3,)` | | ||||
| 112 | | `activation_type` | str | 可选 | 激活函数类型，支持 `"gelu_tanh"`、`"gelu_erf"`，默认值 `"gelu_tanh"`。 | string | - | | 124 | | `activation_type` | str | 可选 | 激活函数类型，支持`"gelu_tanh"`、`"gelu_erf"`，默认值`"gelu_tanh"`。 | string | - | | ||||
| 113 | | `quant_mode` | str | 可选 | 量化模式，当前支持 `"mx"`，默认值 `"mx"`。 | string | - | | 125 | | `quant_mode` | str | 可选 | 量化模式，当前支持`"mx"`，默认值`"mx"`。 | string | - | | ||||
| 114 | | `round_mode` | str | 可选 | 舍入模式，支持 `"rint"`、`"floor"`、`"round"`，当前仅支持 `"rint"`。 | string | - | | 126 | | `round_mode` | str | 可选 | 舍入模式，支持`"rint"`、`"floor"`、`"round"`，当前仅支持`"rint"`。 | string | - | | ||||
| 115 | | `scale_alg` | int | 可选 | 缩放算法。0 表示 MX 动态量化（场景 1），1 表示 FP8 块缩放量化（场景 2），默认值 0。 | int | - | | 127 | | `scale_alg` | int | 可选 | 缩放算法。0 表示MX动态量化（场景 1），1表示FP8块缩放量化（场景 2），默认值0。 | int | - | | ||||
| 116 | | `dst_type_max` | float | 可选 | 目标数据类型最大值，用于量化范围控制，默认值 0.0（表示使用目标类型的默认最大值）。 | float32 | - | | 128 | | `dst_type_max` | float | 可选 | 目标数据类型最大值，用于量化范围控制，默认值0.0（表示使用目标类型的默认最大值）。 | float32 | - | | ||||
| 117 | 129 | ||||||
| 118 | ## 返回值说明 | 130 | ## 返回值说明 | ||||
| 119 | 131 | ||||||
| 120 | | 输出名 | 输出类型 | 可选/必选 | 描述 | 数据类型 | 维度(shape) | | 132 | | 输出名 | 输出类型 | 可选/必选 | 描述 | 数据类型 | 维度(shape) | | ||||
| 121 | | --- | --- | --- | --- | --- | --- | | 133 | | --- | --- | --- | --- | --- | --- | | ||||
| 122 | | `y` | Tensor | 必选 | 动态量化后的矩阵乘及激活计算结果。 | torch.float8_e4m3fn、torch.float8_e5m2 | `(..., M, N)`； | | 134 | | `y` | Tensor | 必选 | 动态量化后的矩阵乘及激活计算结果。 | torch.float8_e4m3fn、torch.float8_e5m2 | `(..., M, N)`； | | ||||
| 123 | | `y_scale` | Tensor | 必选 | 动态量化后每个分组对应的量化尺度，最后一维固定为 2。 | torch.float8_e8m0fnu | `(..., M, CeilDiv(N, 64), 2)` | | 135 | | `y_scale` | Tensor | 必选 | 动态量化后每个分组对应的量化尺度，最后一维固定为2。 | torch.float8_e8m0fnu | `(..., M, CeilDiv(N, 64), 2)` | | ||||
| 124 | 136 | ||||||
| 125 | ## 约束说明 | 137 | ## 约束说明 | ||||
| 126 | 138 | ||||||
| 127 | - 该接口支持训练、推理场景下使用。 | 139 | - 该接口支持训练、推理场景下使用。 | ||||
| 128 | - 该接口支持单算子模式调用。 | 140 | - 该接口支持单算子模式调用。 | ||||
| 129 | - 不支持空 Tensor。 | 141 | - 不支持空Tensor。 | ||||
| 130 | - 支持连续 Tensor，非连续 Tensor 仅支持最后两根轴转置场景。 | 142 | - 支持连续Tensor，非连续Tensor仅支持最后两根轴转置场景。 | ||||
| 131 | - `x1` 支持 2-6 维，`x2`仅支持NZ 4-8 维。 | 143 | - `x1`支持2-6维，`x2`仅支持NZ 4-8 维。 | ||||
| 132 | - 当 `K` 或 `N` 为 1 时，无法使用weightNz特性，本接口不支持此种场景。 | 144 | - 当`K`或`N`为1时，无法使用weightNz特性，本接口不支持此种场景。 | ||||
| 133 | - M/N/K 维度及 `transpose_x1`、`transpose_x2` 由 `x1`、`x2` 最后两维自动匹配推导：取 `x1` 最后两维和 `x2` 最后两维共四个值中相等的一对作为 K，`x1` 中剩余的为 M，`x2` 中剩余的为 N。若四组组合中无相等维度则报错。 | 145 | - M/N/K维度及`transpose_x1`、`transpose_x2`由`x1`、`x2`最后两维自动匹配推导：取`x1`最后两维和`x2`最后两维共四个值中相等的一对作为K，`x1`中剩余的为M，`x2`中剩余的为N。若四组组合中无相等维度则报错。 | ||||
| 134 | - `x1`、`x2` 的 batch 维度（除最后两维外的维度）支持广播（右对齐），如 `x1=(1,M,K)`、`x2=(8,K,N)` 输出 `(8,M,N)`。 | 146 | - `x1`、`x2`的batch维度（除最后两维外的维度）支持广播（右对齐），如`x1=(1,M,K)`、`x2=(8,K,N)`输出`(8,M,N)`。 | ||||
| 135 | - `x1_scale`、`x2_scale` 若传入，其 batch 维度（除最后三维外的维度）的数量和每一维的值必须与对应的 `x1`、`x2` 完全一致；若 `x1` 无 batch 维度（2D），则 `x1_scale`、`x2_scale` 须为 3D。 | 147 | - `x1_scale`、`x2_scale`若传入，其batch维度（除最后三维外的维度）的数量和每一维的值必须与对应的`x1`、`x2`完全一致；若`x1`无 batch维度（2D），则`x1_scale`、`x2_scale`须为3D。 | ||||
| 136 | - `x1_scale`、`x2_scale` 最后一维必须为 2。 | 148 | - `x1_scale`、`x2_scale`最后一维必须为2。 | ||||
| 137 | - `group_sizes` 若传入，必须包含三个元素 `[groupSizeM, groupSizeN, groupSizeK]`，每个元素取值范围为 [0, 65535]，当前mx场景仅支持[1, 1, 32]。 | 149 | - `group_sizes`若传入，必须包含三个元素`[groupSizeM, groupSizeN, groupSizeK]`，每个元素取值范围为[0, 65535]，当前MX场景仅支持[1, 1, 32]。 | ||||
| 138 | - 输入和输出支持以下数据类型组合： | 150 | - 输入和输出支持以下数据类型组合： | ||||
| 139 | 151 | ||||||
| 140 | - <term>Ascend 950PR/Ascend 950DT</term>： | 152 | | x1 | x2 | x1_scale | x2_scale | bias | y | y_scale | | ||||
| 141 | 153 | |---------------|---------------|-------------|-------------|------------------|---------------------------|-------------| | |||||
| 142 | | x1 | x2 | x1_scale | x2_scale | bias | y | y_scale | | 154 | | torch.float8_e4m3fn | torch.float8_e4m3fn | torch.float8_e8m0fnu | torch.float8_e8m0fnu | None/torch.float32 | torch.float8_e4m3fn2 | torch.float8_e8m0fnu | | ||||
| 143 | |---------------|---------------|-------------|-------------|------------------|---------------------------|-------------| | 155 | | torch.float8_e5m2 | torch.float8_e4m3fn | torch.float8_e8m0fnu | torch.float8_e8m0fnu | None/torch.float32 | torch.float8_e5m2 | torch.float8_e8m0fnu | | ||||
| 144 | | torch.float8_e4m3fn | torch.float8_e4m3fn | torch.float8_e8m0fnu | torch.float8_e8m0fnu | None/torch.float32 | torch.float8_e4m3fn2 | torch.float8_e8m0fnu | | ||||||
| 145 | | torch.float8_e5m2 | torch.float8_e4m3fn | torch.float8_e8m0fnu | torch.float8_e8m0fnu | None/torch.float32 | torch.float8_e5m2 | torch.float8_e8m0fnu | | ||||||
| 146 | 156 | ||||||
| 147 | ## 确定性计算 | 157 | ## 确定性计算 | ||||
| 148 | 158 | ||||||
| @@ -25,7 +25,7 @@ | |||
| 25 | 25 | ||
| 26 | - 接口功能： | 26 | - 接口功能： |
| 27 | 27 | ||
| 28 | 融合 RMS Normalization 与对称动态量化（INT8），封装 aclnnRmsNormDynamicQuant。将输入 x 先做 RMSNorm 归一化（可选加 smooth_scales 缩放和 beta 偏置），再按行做动态 INT8 量化，输出 INT8 量化结果和 FLOAT32 的 scale 因子。 | 28 | 融合RMS Normalization与对称动态量化（INT8），封装aclnnRmsNormDynamicQuant。将输入x先做RMSNorm归一化（可选加smooth_scales缩放和beta偏置），再按行做动态INT8量化，输出INT8量化结果和FLOAT32的scale因子。 |
| 29 | 29 | ||
| 30 | - 计算公式： | 30 | - 计算公式： |
| 31 | 31 | ||
| @@ -91,7 +91,7 @@ cann_ops_nn.rms_norm_dynamic_quant(x, gamma, smooth_scales=None, beta=None, *, e | |||
| 91 | <td>x</td> | 91 | <td>x</td> |
| 92 | <td>Tensor</td> | 92 | <td>Tensor</td> |
| 93 | <td>必选</td> | 93 | <td>必选</td> |
| 94 | <td>标准化过程中的源数据张量，公式中的 x。</td> | 94 | <td>标准化过程中的源数据张量，公式中的x。</td> |
| 95 | <td>float16、bfloat16</td> | 95 | <td>float16、bfloat16</td> |
| 96 | <td>(..., D)</td> | 96 | <td>(..., D)</td> |
| 97 | </tr> | 97 | </tr> |
| @@ -99,7 +99,7 @@ cann_ops_nn.rms_norm_dynamic_quant(x, gamma, smooth_scales=None, beta=None, *, e | |||
| 99 | <td>gamma</td> | 99 | <td>gamma</td> |
| 100 | <td>Tensor</td> | 100 | <td>Tensor</td> |
| 101 | <td>必选</td> | 101 | <td>必选</td> |
| 102 | <td>标准化过程中的权重张量，公式中的 gamma。shape 需与 x 最后一维一致。</td> | 102 | <td>标准化过程中的权重张量，公式中的gamma。shape需与x最后一维一致。</td> |
| 103 | <td>float16、bfloat16</td> | 103 | <td>float16、bfloat16</td> |
| 104 | <td>(D,)</td> | 104 | <td>(D,)</td> |
| 105 | </tr> | 105 | </tr> |
| @@ -107,7 +107,7 @@ cann_ops_nn.rms_norm_dynamic_quant(x, gamma, smooth_scales=None, beta=None, *, e | |||
| 107 | <td>smooth_scales</td> | 107 | <td>smooth_scales</td> |
| 108 | <td>Tensor</td> | 108 | <td>Tensor</td> |
| 109 | <td>可选</td> | 109 | <td>可选</td> |
| 110 | <td>量化过程中使用的 smoothScale 张量，公式中的 smooth_scales。shape 和 dtype 需与 gamma 一致。</td> | 110 | <td>量化过程中使用的smoothScale张量，公式中的smooth_scales。shape和dtype需与gamma一致。</td> |
| 111 | <td>float16、bfloat16</td> | 111 | <td>float16、bfloat16</td> |
| 112 | <td>(D,)</td> | 112 | <td>(D,)</td> |
| 113 | </tr> | 113 | </tr> |
| @@ -115,7 +115,7 @@ cann_ops_nn.rms_norm_dynamic_quant(x, gamma, smooth_scales=None, beta=None, *, e | |||
| 115 | <td>beta</td> | 115 | <td>beta</td> |
| 116 | <td>Tensor</td> | 116 | <td>Tensor</td> |
| 117 | <td>可选</td> | 117 | <td>可选</td> |
| 118 | <td>标准化过程中的偏置项，公式中的 beta。shape 和 dtype 需与 gamma 一致。</td> | 118 | <td>标准化过程中的偏置项，公式中的beta。shape和dtype需gamma一致。</td> |
| 119 | <td>float16、bfloat16</td> | 119 | <td>float16、bfloat16</td> |
| 120 | <td>(D,)</td> | 120 | <td>(D,)</td> |
| 121 | </tr> | 121 | </tr> |
| @@ -123,7 +123,7 @@ cann_ops_nn.rms_norm_dynamic_quant(x, gamma, smooth_scales=None, beta=None, *, e | |||
| 123 | <td>epsilon</td> | 123 | <td>epsilon</td> |
| 124 | <td>float</td> | 124 | <td>float</td> |
| 125 | <td>可选</td> | 125 | <td>可选</td> |
| 126 | <td>用于防止除 0 错误，公式中的 epsilon，必须大于零。默认值 1e-6。</td> | 126 | <td>用于防止除0错误，公式中的epsilon，必须大于零。默认值1e-6。</td> |
| 127 | <td>float32</td> | 127 | <td>float32</td> |
| 128 | <td>-</td> | 128 | <td>-</td> |
| 129 | </tr> | 129 | </tr> |
| @@ -131,7 +131,7 @@ cann_ops_nn.rms_norm_dynamic_quant(x, gamma, smooth_scales=None, beta=None, *, e | |||
| 131 | <td>dst_type</td> | 131 | <td>dst_type</td> |
| 132 | <td>int</td> | 132 | <td>int</td> |
| 133 | <td>可选</td> | 133 | <td>可选</td> |
| 134 | <td>输出 y 的数据类型枚举值。当前仅支持 iNT8 量化，默认值 2。</td> | 134 | <td>输出 y 的数据类型枚举值。当前仅支持INT8量化，默认值2。</td> |
| 135 | <td>int</td> | 135 | <td>int</td> |
| 136 | <td>-</td> | 136 | <td>-</td> |
| 137 | </tr> | 137 | </tr> |
| @@ -163,7 +163,7 @@ cann_ops_nn.rms_norm_dynamic_quant(x, gamma, smooth_scales=None, beta=None, *, e | |||
| 163 | <td>y</td> | 163 | <td>y</td> |
| 164 | <td>Tensor</td> | 164 | <td>Tensor</td> |
| 165 | <td>必选</td> | 165 | <td>必选</td> |
| 166 | <td>量化输出，公式中的 yOut，shape 与 x 一致。</td> | 166 | <td>量化输出，公式中的yOut，shape与x一致。</td> |
| 167 | <td>int8</td> | 167 | <td>int8</td> |
| 168 | <td>(..., D)</td> | 168 | <td>(..., D)</td> |
| 169 | </tr> | 169 | </tr> |
| @@ -171,7 +171,7 @@ cann_ops_nn.rms_norm_dynamic_quant(x, gamma, smooth_scales=None, beta=None, *, e | |||
| 171 | <td>scale</td> | 171 | <td>scale</td> |
| 172 | <td>Tensor</td> | 172 | <td>Tensor</td> |
| 173 | <td>必选</td> | 173 | <td>必选</td> |
| 174 | <td>量化 scale 因子，公式中的 scaleOut，shape 为 x 去掉最后一维。</td> | 174 | <td>量化 scale 因子，公式中的scaleOut，shape为x去掉最后一维。</td> |
| 175 | <td>float32</td> | 175 | <td>float32</td> |
| 176 | <td>(...,)</td> | 176 | <td>(...,)</td> |
| 177 | </tr> | 177 | </tr> |
| @@ -182,7 +182,7 @@ cann_ops_nn.rms_norm_dynamic_quant(x, gamma, smooth_scales=None, beta=None, *, e | |||
| 182 | 182 | ||
| 183 | - 该接口支持训练、推理场景下使用。 | 183 | - 该接口支持训练、推理场景下使用。 |
| 184 | - 该接口支持单算子模式和TorchAir图模式调用。 | 184 | - 该接口支持单算子模式和TorchAir图模式调用。 |
| 185 | - 输入`x`的最后一维必须小于等于 8192，否则可能会出现精度问题。 | 185 | - 输入`x`的最后一维必须小于等于8192，否则可能会出现精度问题。 |
| 186 | - `x`不支持全0输入。 | 186 | - `x`不支持全0输入。 |
| 187 | 187 | ||
| 188 | ## 确定性计算 | 188 | ## 确定性计算 |
| @@ -25,7 +25,7 @@ | |||
| 25 | 25 | ||
| 26 | - 接口功能： | 26 | - 接口功能： |
| 27 | 27 | ||
| 28 | 为矩阵x依次进行两次克罗内克积小矩阵乘法，然后针对矩阵乘的结果进行量化处理。底层封装 `aclnnFlatQuantV3`。 | 28 | 为矩阵x依次进行两次克罗内克积小矩阵乘法，然后针对矩阵乘的结果进行量化处理。底层封装`aclnnFlatQuantV3`。 |
| 29 | 29 | ||
| 30 | - 计算公式： | 30 | - 计算公式： |
| 31 | 31 | ||
| @@ -182,20 +182,19 @@ cann_ops_nn.flat_quant(x, kronecker_p1, kronecker_p2, clip_ratio=1.0, dst_dtype= | |||
| 182 | 182 | ||
| 183 | ## 约束说明 | 183 | ## 约束说明 |
| 184 | 184 | ||
| 185 | - 该接口支持训练、推理场景下使用。 | ||
| 186 | - 该接口支持单算子模式和TorchAir图模式调用。 | ||
| 185 | - 输入x的N2维度必须满足以下条件： | 187 | - 输入x的N2维度必须满足以下条件： |
| 186 | - dst_dtype = torch.quint4x2时，N2必须是8的整数倍。 | 188 | - dst_dtype = torch.quint4x2时，N2必须是8的整数倍。 |
| 187 | - dst_dtype = torch_npu.float4_e2m1fn_x2时，N2必须是偶数。 | 189 | - dst_dtype = torch_npu.float4_e2m1fn_x2时，N2必须是偶数。 |
| 188 | - 输入kronecker_p1和kronecker_p2的数据类型必须与x一致。 | 190 | - 输入kronecker_p1和kronecker_p2的数据类型必须与x一致。 |
| 189 | - clip_ratio范围为(0, 1]。 | ||
| 190 | - dst_dtype支持torch.quint4x2（默认）、torch_npu.float4_e2m1fn_x2，输出说明如下： | 191 | - dst_dtype支持torch.quint4x2（默认）、torch_npu.float4_e2m1fn_x2，输出说明如下： |
| 191 | - 如果dtype为torch.quint4x2时，输出out类型为int32，由8个int4拼接，查看具体值需自行解包，输出quant_scale类型为float32。 | 192 | - 如果dtype为torch.quint4x2时，输出out类型为int32，由8个int4拼接，查看具体值需自行解包，输出quant_scale类型为float32。 |
| 192 | - 如果dtype为torch_npu.float4_e2m1fn_x2时，输出out类型为uint8，由2个float4_e2m1fn_x2拼接，查看具体值需自行解包，输出quant_scale类型为uint8，查看实际值需自行转换成float8_e8m0fnu。 | 193 | - 如果dtype为torch_npu.float4_e2m1fn_x2时，输出out类型为uint8，由2个float4_e2m1fn_x2拼接，查看具体值需自行解包，输出quant_scale类型为uint8，查看实际值需自行转换成float8_e8m0fnu。 |
| 193 | - dst_type_max只能为0或[6, 12]范围内的数。 | ||
| 194 | - group_list_type取值范围为0-2，当group_list_type为0或1时，group_list的shape为(G,)，当group_list_type为2时，group_list的shape为(G, 2)，G表示分组数，G需要小于等于1024。 | 194 | - group_list_type取值范围为0-2，当group_list_type为0或1时，group_list的shape为(G,)，当group_list_type为2时，group_list的shape为(G, 2)，G表示分组数，G需要小于等于1024。 |
| 195 | - group_list需要满足以下条件： | 195 | - group_list需要满足以下条件： |
| 196 | <!-- npu="950" id7 --> | 196 | <!-- npu="950" id7 --> |
| 197 | - <term>Ascend 950PR/Ascend 950DT</term>： | 197 | - <term>Ascend 950PR/Ascend 950DT</term>：group_list仅支持None输入。 |
| 198 | - group_list仅支持None输入。 | ||
| 199 | <!-- end id7 --> | 198 | <!-- end id7 --> |
| 200 | <!-- npu="A3,910b" id8 --> | 199 | <!-- npu="A3,910b" id8 --> |
| 201 | - <term>Atlas A3 训练系列产品/Atlas A3 推理系列产品</term>、<term>Atlas A2 训练系列产品/Atlas A2 推理系列产品</term>： | 200 | - <term>Atlas A3 训练系列产品/Atlas A3 推理系列产品</term>、<term>Atlas A2 训练系列产品/Atlas A2 推理系列产品</term>： |
| @@ -204,8 +203,6 @@ cann_ops_nn.flat_quant(x, kronecker_p1, kronecker_p2, clip_ratio=1.0, dst_dtype= | |||
| 204 | - 当group_list_type为1时，group_list必须为非负数列，表示分组后每组大小，数值的总和应小于等于x中tensor的第一维。 | 203 | - 当group_list_type为1时，group_list必须为非负数列，表示分组后每组大小，数值的总和应小于等于x中tensor的第一维。 |
| 205 | - 当group_list_type为2时，group_list必须为非负数列，数据排布为[[groupIdx0, groupSize0], [groupIdx1, groupSize1]...]，其中groupSize为分组后每组大小，第二列数值的总和应小于等于x中tensor的第一维。 | 204 | - 当group_list_type为2时，group_list必须为非负数列，数据排布为[[groupIdx0, groupSize0], [groupIdx1, groupSize1]...]，其中groupSize为分组后每组大小，第二列数值的总和应小于等于x中tensor的第一维。 |
| 206 | <!-- end id8 --> | 205 | <!-- end id8 --> |
| 207 | - 该接口支持训练、推理场景下使用。 | ||
| 208 | - 该接口支持单算子模式和TorchAir图模式调用。 | ||
| 209 | 206 | ||
| 210 | ## 确定性计算 | 207 | ## 确定性计算 |
| 211 | 208 | ||
Mquant/grouped_dynamic_mx_quant_with_dual_axis/docs/torchapi_grouped_dynamic_mx_quant_with_dual_axis.md+18-6
| @@ -1,23 +1,35 @@ | |||
| 1 | # cann_ops_nn.grouped_dynamic_mx_quant_with_dual_axis | 1 | # grouped_dynamic_mx_quant_with_dual_axis |
| 2 | 2 | ||
| 3 | ## 产品支持情况 | 3 | ## 产品支持情况 |
| 4 | 4 | ||
| 5 | <!-- npu="950" id1 --> | ||
| 5 | - <term>Ascend 950PR/Ascend 950DT</term>：支持 | 6 | - <term>Ascend 950PR/Ascend 950DT</term>：支持 |
| 7 | <!-- end id1 --> | ||
| 8 | <!-- npu="A3" id2 --> | ||
| 6 | - <term>Atlas A3 训练系列产品/Atlas A3 推理系列产品</term>：不支持 | 9 | - <term>Atlas A3 训练系列产品/Atlas A3 推理系列产品</term>：不支持 |
| 10 | <!-- end id2 --> | ||
| 11 | <!-- npu="910b" id3 --> | ||
| 7 | - <term>Atlas A2 训练系列产品/Atlas A2 推理系列产品</term>：不支持 | 12 | - <term>Atlas A2 训练系列产品/Atlas A2 推理系列产品</term>：不支持 |
| 13 | <!-- end id3 --> | ||
| 14 | <!-- npu="310b" id4 --> | ||
| 8 | - <term>Atlas 200I/500 A2 推理产品</term>：不支持 | 15 | - <term>Atlas 200I/500 A2 推理产品</term>：不支持 |
| 16 | <!-- end id4 --> | ||
| 17 | <!-- npu="310p" id5 --> | ||
| 9 | - <term>Atlas 推理系列产品</term>：不支持 | 18 | - <term>Atlas 推理系列产品</term>：不支持 |
| 19 | <!-- end id5 --> | ||
| 20 | <!-- npu="910" id6 --> | ||
| 10 | - <term>Atlas 训练系列产品</term>：不支持 | 21 | - <term>Atlas 训练系列产品</term>：不支持 |
| 22 | <!-- end id6 --> | ||
| 11 | 23 | ||
| 12 | ## 功能说明 | 24 | ## 功能说明 |
| 13 | 25 | ||
| 14 | - 接口功能： | 26 | - 接口功能： |
| 15 | 27 | ||
| 16 | 根据 x 和 group_index 一次完成双轴 MX 量化，输出 y1、mxscale1、y2、mxscale2。其中 y1、mxscale1 为最后一维方向量化结果，y2、mxscale2 为倒数第二维方向量化结果。 | 28 | 根据x和group_index一次完成双轴MX量化，输出y1、mxscale1、y2、mxscale2。其中y1、mxscale1为最后一维方向量化结果，y2、mxscale2为倒数第二维方向量化结果。 |
| 17 | 29 | ||
| 18 | - 计算说明： | 30 | - 计算说明： |
| 19 | 31 | ||
| 20 | 量化以 32 个元素为一个 MX block。scale_alg=1 时，每个 block 生成 FLOAT8_E8M0 类型的 scale，并使用该 scale 的倒数将 block 内元素转换为 dst_type 指定的数据类型。 | 32 | 量化以32个元素为一个MX block。scale_alg=1时，每个block生成FLOAT8_E8M0类型的scale，并使用该scale的倒数将block内元素转换为 dst_type指定的数据类型。 |
| 21 | 33 | ||
| 22 | - 计算公式： | 34 | - 计算公式： |
| 23 | 35 | ||
| @@ -78,7 +90,7 @@ | |||
| 78 | ## 函数原型 | 90 | ## 函数原型 |
| 79 | 91 | ||
| 80 | ```python | 92 | ```python |
| 81 | torch.ops.cann_ops_nn.grouped_dynamic_mx_quant_with_dual_axis( | 93 | cann_ops_nn.grouped_dynamic_mx_quant_with_dual_axis( |
🟡 Medium Priority 变更将函数原型从
建议：统一示例代码中的 API 调用形式，将第265行的 | |||
| 82 | x, group_index, *, round_mode="rint", scale_alg=1, dst_type=24, dst_type_max=0) | 94 | x, group_index, *, round_mode="rint", scale_alg=1, dst_type=24, dst_type_max=0) |
| 83 | -> (Tensor y1, Tensor mxscale1, Tensor y2, Tensor mxscale2) | 95 | -> (Tensor y1, Tensor mxscale1, Tensor y2, Tensor mxscale2) |
| 84 | ``` | 96 | ``` |
| @@ -116,7 +128,7 @@ torch.ops.cann_ops_nn.grouped_dynamic_mx_quant_with_dual_axis( | |||
| 116 | <td>group_index</td> | 128 | <td>group_index</td> |
| 117 | <td>Tensor</td> | 129 | <td>Tensor</td> |
| 118 | <td>必选</td> | 130 | <td>必选</td> |
| 119 | <td>量化分组索引，采用cumsum形式描述各group边界。G表示分组数；每个元素值需非负且非递减，最后一个元素需等于M。</td> | 131 | <td>量化分组索引，采用cumsum形式描述各group边界。G表示分组数；每个元素值非负且非递减，最后一个元素需等于M。</td> |
| 120 | <td>torch.int64</td> | 132 | <td>torch.int64</td> |
| 121 | <td>(G,)</td> | 133 | <td>(G,)</td> |
| 122 | </tr> | 134 | </tr> |
| @@ -226,7 +238,7 @@ torch.ops.cann_ops_nn.grouped_dynamic_mx_quant_with_dual_axis( | |||
| 226 | - `y1`、`y2`的shape需与`x`一致，均为(M, N)。 | 238 | - `y1`、`y2`的shape需与`x`一致，均为(M, N)。 |
| 227 | - `mxscale1`的shape为(M, ceil(N/64), 2)。 | 239 | - `mxscale1`的shape为(M, ceil(N/64), 2)。 |
| 228 | - `mxscale2`的shape为(floor(M/64)+G, N, 2)，其中G为group_index的元素个数。 | 240 | - `mxscale2`的shape为(floor(M/64)+G, N, 2)，其中G为group_index的元素个数。 |
| 229 | - x 的第0维 M 允许为0；N 和 group_index 的长度仍需大于0。 | 241 | - x的第0维M允许为0；N和group_index的长度仍需大于0。 |
| 230 | 242 | ||
| 231 | ## 确定性计算 | 243 | ## 确定性计算 |
| 232 | 244 | ||
| @@ -1,4 +1,4 @@ | |||
| 1 | # cann_ops_nn.mx_to_block_mx_quant | 1 | # mx_to_block_mx_quant |
| 2 | 2 | ||
| 3 | ## 产品支持情况 | 3 | ## 产品支持情况 |
| 4 | 4 | ||
| @@ -23,8 +23,7 @@ | |||
| 23 | 23 | ||
| 24 | ## 功能说明 | 24 | ## 功能说明 |
| 25 | 25 | ||
| 26 | - 接口功能： | 26 | - 接口功能：将调用npu_dynamic_mx_quant量化得到的FLOAT4的Tensor结合FLOAT8_E8M0缩放系数，转换为FLOAT8分块量化格式，同时输出-1轴和-2轴方向的量化尺度。 |
| 27 | 将调用 npu_dynamic_mx_quant 量化得到的 FLOAT4 的 Tensor 结合 FLOAT8_E8M0 缩放系数，转换为 FLOAT8 分块量化格式，同时输出 -1 轴和 -2 轴方向的量化尺度。 | ||
| 28 | 27 | ||
| 29 | - 计算公式： | 28 | - 计算公式： |
| 30 | 29 | ||
| @@ -36,7 +35,7 @@ | |||
| 36 | \end{aligned} | 35 | \end{aligned} |
| 37 | $$ | 36 | $$ |
| 38 | 37 | ||
| 39 | - 其中 $mxscale_{fp4\_block}$ 是输入 mxscale 提供的 FP8_E8M0 缩放系数；$MAX\_OFFSET$ 是输入和输出数据类型之间的最大偏移量；$x_{fp4}$ 是量化得到的 FLOAT4 张量；$x_{fp8}$ 是转换得到的 FLOAT8 张量。 | 38 | - 其中 $mxscale_{fp4\_block}$ 是输入mxscale提供的FP8_E8M0缩放系数；$MAX\_OFFSET$ 是输入和输出数据类型之间的最大偏移量；$x_{fp4}$ 是量化得到的FLOAT4张量；$x_{fp8}$ 是转换得到的FLOAT8张量。 |
| 40 | 39 | ||
| 41 | - MAX_OFFSET 对照表： | 40 | - MAX_OFFSET 对照表： |
| 42 | 41 | ||
| @@ -50,7 +49,7 @@ | |||
| 50 | ## 函数原型 | 49 | ## 函数原型 |
| 51 | 50 | ||
| 52 | ```python | 51 | ```python |
| 53 | torch.ops.cann_ops_nn.mx_to_block_mx_quant(x, mxscale, *, dst_type=292, x_type=296) | 52 | cann_ops_nn.mx_to_block_mx_quant(x, mxscale, *, dst_type=292, x_type=296) |
| 54 | -> (Tensor, Tensor, Tensor) | 53 | -> (Tensor, Tensor, Tensor) |
| 55 | ``` | 54 | ``` |
| 56 | 55 | ||
| @@ -79,7 +78,7 @@ torch.ops.cann_ops_nn.mx_to_block_mx_quant(x, mxscale, *, dst_type=292, x_type=2 | |||
| 79 | <td>x</td> | 78 | <td>x</td> |
| 80 | <td>Tensor</td> | 79 | <td>Tensor</td> |
| 81 | <td>必选</td> | 80 | <td>必选</td> |
| 82 | <td>算子输入的量化 Tensor，公式中的 x_fp4。</td> | 81 | <td>算子输入的量化Tensor，公式中的x_fp4。</td> |
| 83 | <td>FLOAT4_E2M1、FLOAT4_E1M2</td> | 82 | <td>FLOAT4_E2M1、FLOAT4_E1M2</td> |
| 84 | <td>2-3</td> | 83 | <td>2-3</td> |
| 85 | </tr> | 84 | </tr> |
| @@ -87,7 +86,7 @@ torch.ops.cann_ops_nn.mx_to_block_mx_quant(x, mxscale, *, dst_type=292, x_type=2 | |||
| 87 | <td>mxscale </td> | 86 | <td>mxscale </td> |
| 88 | <td>Tensor</td> | 87 | <td>Tensor</td> |
| 89 | <td>必选</td> | 88 | <td>必选</td> |
| 90 | <td>调用 DynamicMxQuant 计算得到的量化尺度，公式中的 mxscale_fp4_block。shape 需满足约束说明中的公式。</td> | 89 | <td>调用DynamicMxQuant计算得到的量化尺度，公式中的mxscale_fp4_block。shape需满足约束说明中的公式。</td> |
| 91 | <td>FLOAT8_E8M0</td> | 90 | <td>FLOAT8_E8M0</td> |
| 92 | <td>3-4</td> | 91 | <td>3-4</td> |
| 93 | </tr> | 92 | </tr> |
| @@ -95,7 +94,7 @@ torch.ops.cann_ops_nn.mx_to_block_mx_quant(x, mxscale, *, dst_type=292, x_type=2 | |||
| 95 | <td>dst_type</td> | 94 | <td>dst_type</td> |
| 96 | <td>int</td> | 95 | <td>int</td> |
| 97 | <td>可选</td> | 96 | <td>可选</td> |
| 98 | <td>指定输出 y 的数据类型。当前仅支持输入范围为{291, 292}，分别对应输出 y 的数据类型为 {291: FLOAT8_E5M2, 292: FLOAT8_E4M3FN}。</td> | 97 | <td>指定输出y的数据类型。当前仅支持输入范围为{291, 292}，分别对应输出y的数据类型为{291: FLOAT8_E5M2, 292: FLOAT8_E4M3FN}。</td> |
| 99 | <td>INT</td> | 98 | <td>INT</td> |
| 100 | <td>-</td> | 99 | <td>-</td> |
| 101 | </tr> | 100 | </tr> |
| @@ -103,7 +102,7 @@ torch.ops.cann_ops_nn.mx_to_block_mx_quant(x, mxscale, *, dst_type=292, x_type=2 | |||
| 103 | <td>x_type</td> | 102 | <td>x_type</td> |
| 104 | <td>int</td> | 103 | <td>int</td> |
| 105 | <td>可选</td> | 104 | <td>可选</td> |
| 106 | <td>表示输入 x 的真实数据类型。当前仅支持输入范围为{296, 297}，分别对应输入 x 的数据类型为 {296: FLOAT4_E2M1, 297: FLOAT4_E1M2}。</td> | 105 | <td>表示输入x的真实数据类型。当前仅支持输入范围为{296, 297}，分别对应输入x的数据类型为{296: FLOAT4_E2M1, 297: FLOAT4_E1M2}。</td> |
| 107 | <td>INT</td> | 106 | <td>INT</td> |
| 108 | <td>-</td> | 107 | <td>-</td> |
| 109 | </tr> | 108 | </tr> |
| @@ -135,7 +134,7 @@ torch.ops.cann_ops_nn.mx_to_block_mx_quant(x, mxscale, *, dst_type=292, x_type=2 | |||
| 135 | <td>y</td> | 134 | <td>y</td> |
| 136 | <td>Tensor</td> | 135 | <td>Tensor</td> |
| 137 | <td>必选</td> | 136 | <td>必选</td> |
| 138 | <td>量化输出，公式中的 x_fp8，shape 与 x 一致，数据类型由 dst_type 决定。</td> | 137 | <td>量化输出，公式中的x_fp8，shape与x一致，数据类型由dst_type决定。</td> |
| 139 | <td>FLOAT8_E5M2、FLOAT8_E4M3FN</td> | 138 | <td>FLOAT8_E5M2、FLOAT8_E4M3FN</td> |
| 140 | <td>2-3</td> | 139 | <td>2-3</td> |
| 141 | </tr> | 140 | </tr> |
| @@ -143,7 +142,7 @@ torch.ops.cann_ops_nn.mx_to_block_mx_quant(x, mxscale, *, dst_type=292, x_type=2 | |||
| 143 | <td>scale1</td> | 142 | <td>scale1</td> |
| 144 | <td>Tensor</td> | 143 | <td>Tensor</td> |
| 145 | <td>必选</td> | 144 | <td>必选</td> |
| 146 | <td>表示 -1 轴每个分组对应的量化尺度。shape 需满足约束说明中的公式。</td> | 145 | <td>表示-1轴每个分组对应的量化尺度。shape需满足约束说明中的公式。</td> |
| 147 | <td>FLOAT8_E8M0</td> | 146 | <td>FLOAT8_E8M0</td> |
| 148 | <td>3-4</td> | 147 | <td>3-4</td> |
| 149 | </tr> | 148 | </tr> |
| @@ -151,7 +150,7 @@ torch.ops.cann_ops_nn.mx_to_block_mx_quant(x, mxscale, *, dst_type=292, x_type=2 | |||
| 151 | <td>scale2</td> | 150 | <td>scale2</td> |
| 152 | <td>Tensor</td> | 151 | <td>Tensor</td> |
| 153 | <td>必选</td> | 152 | <td>必选</td> |
| 154 | <td>表示 -2 轴每个分组对应的量化尺度，输出需要对每两行数据进行交织处理。shape 需满足约束说明中的公式。</td> | 153 | <td>表示-2轴每个分组对应的量化尺度，输出需要对每两行数据进行交织处理。shape需满足约束说明中的公式。</td> |
| 155 | <td>FLOAT8_E8M0</td> | 154 | <td>FLOAT8_E8M0</td> |
| 156 | <td>3-4</td> | 155 | <td>3-4</td> |
| 157 | </tr> | 156 | </tr> |
| @@ -160,23 +159,23 @@ torch.ops.cann_ops_nn.mx_to_block_mx_quant(x, mxscale, *, dst_type=292, x_type=2 | |||
| 160 | 159 | ||
| 161 | ## 约束说明 | 160 | ## 约束说明 |
| 162 | 161 | ||
| 163 | - x 只支持 2 维或 3 维输入，且 -2 轴是 64 的倍数, -1 轴是 2 的倍数。 | 162 | - `x`只支持2维或3维输入，且-2轴是64的倍数，-1轴是2的倍数。 |
| 164 | - 关于 x、mxscale、scale1、scale2 的 shape 约束说明如下： | 163 | - 关于`x`、`mxscale`、`scale1`、`scale2`的shape约束说明如下： |
| 165 | - rank(mxscale) = rank(x) + 1。 | 164 | - rank(mxscale) = rank(x) + 1。 |
| 166 | - mxscale.shape[-2] = (Ceil(x.shape[-1] / 32) + 2 - 1) / 2。 | 165 | - mxscale.shape[-2] = (Ceil(x.shape[-1], 32) + 2 - 1) / 2。 |
| 167 | - mxscale.shape[-1] = 2。 | 166 | - mxscale.shape[-1] = 2。 |
| 168 | - 其它维度与输入 x 一致。 | 167 | - 其它维度与输入`x`一致。 |
| 169 | - 关于输出 scale1 的 shape 约束说明如下： | 168 | - 关于输出`scale1`的shape约束说明如下： |
| 170 | - rank(scale1) = rank(x) + 1。 | 169 | - rank(scale1) = rank(x) + 1。 |
| 171 | - scale1.shape[-2] = (Ceil(x.shape[-1] / 32) + 2 - 1) / 2。 | 170 | - scale1.shape[-2] = (Ceil(x.shape[-1], 32) + 2 - 1) / 2。 |
| 172 | - scale1.shape[-1] = 2。 | 171 | - scale1.shape[-1] = 2。 |
| 173 | - 其它维度和输入 x 保持一致。 | 172 | - 其它维度和输入`x`保持一致。 |
| 174 | - 关于输出 scale2 的 shape 约束说明如下： | 173 | - 关于输出`scale2`的shape约束说明如下： |
| 175 | - rank(scale2) = rank(x) + 1。 | 174 | - rank(scale2) = rank(x) + 1。 |
| 176 | - scale2.shape[-3] = ((Ceil(x.shape[-2] / 32) + 2 - 1) / 2) * 2 / 2。 | 175 | - scale2.shape[-3] = ((Ceil(x.shape[-2], 32) + 2 - 1) / 2) * 2 / 2。 |
🟡 Medium Priority 在 shape 约束公式中，原文 受影响位置：
建议：将所有 | |||
| 177 | - scale2.shape[-2] = x.shape[-1]。 | 176 | - scale2.shape[-2] = x.shape[-1]。 |
| 178 | - scale2.shape[-1] = 2。 | 177 | - scale2.shape[-1] = 2。 |
| 179 | - 其它维度和输入 x 保持一致。 | 178 | - 其它维度和输入`x`保持一致。 |
| 180 | - 该接口支持推理场景下使用。 | 179 | - 该接口支持推理场景下使用。 |
| 181 | - 该接口支持单算子模式和TorchAir图模式调用。 | 180 | - 该接口支持单算子模式和TorchAir图模式调用。 |
| 182 | 181 | ||
| @@ -1,23 +1,35 @@ | |||
| 1 | # cann_ops_nn.swiglu_group_backward | 1 | # swiglu_group_backward |
| 2 | 2 | ||
| 3 | ## 产品支持情况 | 3 | ## 产品支持情况 |
| 4 | 4 | ||
| 5 | <!-- npu="950" id1 --> | ||
| 5 | - <term>Ascend 950PR/Ascend 950DT</term>：支持 | 6 | - <term>Ascend 950PR/Ascend 950DT</term>：支持 |
| 7 | <!-- end id1 --> | ||
| 8 | <!-- npu="A3" id2 --> | ||
| 6 | - <term>Atlas A3 训练系列产品/Atlas A3 推理系列产品</term>：不支持 | 9 | - <term>Atlas A3 训练系列产品/Atlas A3 推理系列产品</term>：不支持 |
| 10 | <!-- end id2 --> | ||
| 11 | <!-- npu="910b" id3 --> | ||
| 7 | - <term>Atlas A2 训练系列产品/Atlas A2 推理系列产品</term>：不支持 | 12 | - <term>Atlas A2 训练系列产品/Atlas A2 推理系列产品</term>：不支持 |
| 13 | <!-- end id3 --> | ||
| 14 | <!-- npu="310b" id4 --> | ||
| 8 | - <term>Atlas 200I/500 A2 推理产品</term>：不支持 | 15 | - <term>Atlas 200I/500 A2 推理产品</term>：不支持 |
| 16 | <!-- end id4 --> | ||
| 17 | <!-- npu="310p" id5 --> | ||
| 9 | - <term>Atlas 推理系列产品</term>：不支持 | 18 | - <term>Atlas 推理系列产品</term>：不支持 |
| 19 | <!-- end id5 --> | ||
| 20 | <!-- npu="910" id6 --> | ||
| 10 | - <term>Atlas 训练系列产品</term>：不支持 | 21 | - <term>Atlas 训练系列产品</term>：不支持 |
| 22 | <!-- end id6 --> | ||
| 11 | 23 | ||
| 12 | ## 功能说明 | 24 | ## 功能说明 |
| 13 | 25 | ||
| 14 | - 接口功能： | 26 | - 接口功能： |
| 15 | 27 | ||
| 16 | 完成 ClampedSwiglu 激活函数的反向梯度计算。从上游梯度 `grad_output` 和前向输入 `x` 重算 clamp mask 与 sigmoid，输出 `grad_x` 与可选 `grad_weight`。底层调用 `aclnnSwigluGroupGrad`。 | 28 | 完成ClampedSwiglu激活函数的反向梯度计算。从上游梯度`grad_output`和前向输入`x`重算clamp mask与sigmoid，输出`grad_x`与可选 `grad_weight`。底层调用`aclnnSwigluGroupGrad`。 |
| 17 | 29 | ||
| 18 | - 计算公式： | 30 | - 计算公式： |
| 19 | 31 | ||
| 20 | 前向分解：`x` 按 hidden 维劈半得到 gate (`g`) 和 up (`u`)；可选 clamp 产生 `g̃ = min(c, g)`、`ũ = clip(u, −c, c)`；`SiLU(g̃) = g̃·σ(g̃)`；`y = SiLU(g̃)·ũ·w_t`。 | 32 | 前向分解：`x`按hidden维劈半得到gate(`g`)和up(`u`)；可选clamp产生`g̃ = min(c, g)`、`ũ = clip(u, −c, c)`；`SiLU(g̃) = g̃·σ(g̃)`；`y = SiLU(g̃)·ũ·w_t`。 |
| 21 | 33 | ||
| 22 | $$ | 34 | $$ |
| 23 | silu'(g̃) = s + f − f·s | 35 | silu'(g̃) = s + f − f·s |
| @@ -35,7 +47,7 @@ | |||
| 35 | grad\_weight = \Sigma(grad\_output \cdot y\_origin) \text{ along hidden dim} | 47 | grad\_weight = \Sigma(grad\_output \cdot y\_origin) \text{ along hidden dim} |
| 36 | $$ | 48 | $$ |
| 37 | 49 | ||
| 38 | 其中 `I` 为开区间指示函数（边界值时 mask=0），`m_r` 为 group_index mask，`w_t` 为 weight 的 broadcast。 | 50 | 其中`I`为开区间指示函数（边界值时 mask=0），`m_r`为group_index mask，`w_t`为weight的broadcast。 |
| 39 | 51 | ||
| 40 | ## 函数原型 | 52 | ## 函数原型 |
| 41 | 53 | ||
| @@ -55,29 +67,29 @@ torch.ops.cann_ops_nn.swiglu_group_backward( | |||
| 55 | 67 | ||
| 56 | | 参数名 | 参数类型 | 可选/必选 | 描述 | 数据类型 | 维度(shape) | | 68 | | 参数名 | 参数类型 | 可选/必选 | 描述 | 数据类型 | 维度(shape) | |
| 57 | | --- | --- | --- | --- | --- | --- | | 69 | | --- | --- | --- | --- | --- | --- | |
| 58 | | `grad_output` | Tensor | 必选 | 上游梯度。 | `torch.float16`、`torch.bfloat16`、`torch.float32` | 2维或3维 | | 70 | | `grad_output`| Tensor | 必选 | 上游梯度。 | `torch.float16`、`torch.bfloat16`、`torch.float32`| 2维或3维 | |
| 59 | | `x` | Tensor | 必选 | 前向输入，包含 gate 和 up 分支。 | `torch.float16`、`torch.bfloat16`、`torch.float32` | 2维或3维，最后一维为 `2 * grad_output.shape[-1]` | | 71 | | `x`| Tensor | 必选 | 前向输入，包含 gate 和 up 分支。 | `torch.float16`、`torch.bfloat16`、`torch.float32`| 2维或3维，最后一维为 `2 * grad_output.shape[-1]`| |
| 60 | | `weight` | Tensor | 可选 | MoE top-k 路由权重，缺省视作全1。 | `torch.float32` | 与 `grad_output` 前导维度一致，最后一维为1 | | 72 | | `weight`| Tensor | 可选 | MoE top-k路由权重，缺省视作全1。 | `torch.float32`| 与`grad_output`前导维度一致，最后一维为1 | |
| 61 | | `y_origin` | Tensor | 可选 | 前向输出 y，weight 存在时 y 已乘该权重。 | 与 `grad_output` 相同 | 与 `grad_output` 相同 | | 73 | | `y_origin`| Tensor | 可选 | 前向输出 y，weight存在时y已乘该权重。 | 与`grad_output`相同 | 与`grad_output`相同 | |
| 62 | | `group_index` | Tensor | 可选 | 各分组 token 数量索引，缺省视作全部行有效。 | `torch.int64` | 1维 | | 74 | | `group_index`| Tensor | 可选 | 各分组token数量索引，缺省视作全部行有效。 | `torch.int64`| 1维 | |
| 63 | | `clamp_limit` | float | 可选 | 截断门限，默认 `0.0` 表示不 clamp。 | - | - | | 75 | | `clamp_limit`| float | 可选 | 截断门限，默认`0.0`表示不clamp。 | - | - | |
| 64 | 76 | ||
| 65 | ## 返回值说明 | 77 | ## 返回值说明 |
| 66 | 78 | ||
| 67 | | 参数名 | 参数类型 | 描述 | 数据类型 | 维度(shape) | | 79 | | 参数名 | 参数类型 | 描述 | 数据类型 | 维度(shape) | |
| 68 | | --- | --- | --- | --- | --- | | 80 | | --- | --- | --- | --- | --- | |
| 69 | | `grad_x` | Tensor | x 的梯度。 | 与 `grad_output` 相同 | 与 `x` 相同 | | 81 | |`grad_x`| Tensor | x的梯度。 | 与`grad_output`相同 | 与`x`相同 | |
| 70 | | `grad_weight` | Tensor? | weight 的梯度，仅 weight 和 y_origin 同时提供时返回。 | `torch.float32` | 与 `weight` 相同 | | 82 | | `grad_weight`| Tensor | weight的梯度，仅weight和y_origin同时提供时返回。 | `torch.float32`| 与`weight`相同 | |
| 71 | 83 | ||
| 72 | ## 约束说明 | 84 | ## 约束说明 |
| 73 | 85 | ||
| 74 | - 该接口支持单算子模式和 TorchAir 图模式调用。 | 86 | - 该接口支持单算子模式和TorchAir图模式调用。 |
| 75 | - `grad_output`、`x`、`weight`、`y_origin`、`group_index` 均需为 NPU Tensor；可选 Tensor 可以传 `None`。 | 87 | - `grad_output`、`x`、`weight`、`y_origin`、`group_index`均需为NPU Tensor；可选Tensor可以传`None`。 |
| 76 | - `weight` 和 `y_origin` 必须同时提供或同时为空。 | 88 | - `weight`和`y_origin`必须同时提供或同时为空。 |
| 77 | - `grad_output` 与 `x` 的前导维度必须一致，且二者均为 2 维或 3 维 Tensor。 | 89 | - `grad_output`与`x`的前导维度必须一致，且二者均为2维或3维Tensor。 |
| 78 | - `x.shape[-1]` 必须等于 `2 * grad_output.shape[-1]`。 | 90 | - `x.shape[-1]`必须等于`2 * grad_output.shape[-1]`。 |
| 79 | - `group_index` 非空时必须是一维非空 Tensor（G > 0）。 | 91 | - `group_index`非空时必须是一维非空Tensor（G > 0）。 |
| 80 | - `clamp_limit` 必须 >= 0.0。 | 92 | - `clamp_limit`必须>= 0.0。 |
| 81 | 93 | ||
| 82 | ## 确定性计算 | 94 | ## 确定性计算 |
| 83 | 95 | ||
🟡 Medium Priority
变更将函数原型从
torch.ops.cann_ops_nn.swiglu_group(改为cann_ops_nn.swiglu_group(（第43行），但下方示例代码（第92行、第107行）仍然使用torch.ops.cann_ops_nn.swiglu_group(。建议：统一示例代码中的 API 调用形式，将第92行和第107行的
torch.ops.cann_ops_nn.swiglu_group(改为cann_ops_nn.swiglu_group(，或将函数原型恢复为torch.ops.cann_ops_nn.swiglu_group(。