已合并
MoeComputeExpertTokens文档更新 #5728
wuxiyuan创建于 20 天前
MoeComputeExpertTokens文档更新 #5728
已合并
共 2 个文件变更+17-15
| @@ -4003,7 +4003,7 @@ _add_torch_npu_docstr( | |||
| 4003 | "npu_mhc_pre", | 4003 | "npu_mhc_pre", |
| 4004 | """ | 4004 | """ |
| 4005 | 接口原型： | 4005 | 接口原型： |
| 4006 | -torch_npu.npu_mhc_pre(Tensor x, Tensor phi, Tensor alpha, Tensor bias, *, Tensor? gamma=None, float norm_eps=1e-6, float hc_eps=1e-6, int out_flag=0, int inner_precise=0) -> (Tensor, Tensor, Tensor, Tensor, Tensor, Tensor) | 4006 | +torch_npu.npu_mhc_pre(Tensor x, Tensor phi, Tensor alpha, Tensor bias, *, Tensor? gamma=None, float norm_eps=1e-6, float hc_eps=1e-6, int out_flag=0, int inner_precise=0) -> (Tensor, Tensor, Tensor, Tensor, Tensor, Tensor) |
| 4007 | 4007 | ||
| 4008 | 功能描述： | 4008 | 功能描述： |
| 4009 | 通过一系列计算，可得到 MHC (流形约束超连接)架构中 hidden 层对应的投影矩阵 Hres 和 Hpost，以及作为 Atten 或 MLP 层输入的矩阵 Hin。 | 4009 | 通过一系列计算，可得到 MHC (流形约束超连接)架构中 hidden 层对应的投影矩阵 Hres 和 Hpost，以及作为 Atten 或 MLP 层输入的矩阵 Hin。 |
| @@ -4017,7 +4017,7 @@ gamma: Tensor类型，可选输入，表示进行 RMSNorm 计算时的缩放因 | |||
| 4017 | norm_eps: Float类型，可选输入，RMSNorm 的防除零参数。 | 4017 | norm_eps: Float类型，可选输入，RMSNorm 的防除零参数。 |
| 4018 | hc_eps: Float类型，可选输入，H_pre 经过 sigmoid 运算后的 eps 参数。 | 4018 | hc_eps: Float类型，可选输入，H_pre 经过 sigmoid 运算后的 eps 参数。 |
| 4019 | out_flag: Int类型，可选输入，表示是否输出中间结果标识，默认值为0(仅输出最终变换结果)。 | 4019 | out_flag: Int类型，可选输入，表示是否输出中间结果标识，默认值为0(仅输出最终变换结果)。 |
| 4020 | -inner_precise: Int类型，可选输入，指定MhcPre算子的计算模式，默认值为0。0表示在Cube中使用FP32模式计算；1表示在Cube中使用HF32模式计算。 | 4020 | +inner_precise: Int类型，可选输入，指定MhcPre算子的计算模式，默认值为0。0表示在Cube中使用FP32模式计算；1表示在Cube中使用HF32模式计算。 |
| 4021 | n：shape 中的 n 常取 4、6、8。 | 4021 | n：shape 中的 n 常取 4、6、8。 |
| 4022 | 4022 | ||
| 4023 | 输出说明: | 4023 | 输出说明: |
| @@ -10539,14 +10539,14 @@ torch_npu.npu_moe_compute_expert_tokens(Tensor sorted_expert_for_source_row, int | |||
| 10539 | 功能描述 | 10539 | 功能描述 |
| 10540 | 算子功能: MoE(Mixture of Experts, 混合专家模型)计算中, 通过二分查找的方式查找每个专家处理的最后一行的位置. | 10540 | 算子功能: MoE(Mixture of Experts, 混合专家模型)计算中, 通过二分查找的方式查找每个专家处理的最后一行的位置. |
| 10541 | 计算公式: | 10541 | 计算公式: |
| 10542 | -expertTokens_{i}=BinaerSearch(sortedExpertForSourceRow,numExpert) | 10542 | +expertTokens_{i}=BinarySearch(sortedExpertForSourceRow,i), i in [0,numExpert) |
| 10543 | 10543 | ||
| 10544 | 参数说明 | 10544 | 参数说明 |
| 10545 | -sorted_expert_for_source_row: Tensor类型, 必选参数, 经过专家处理过的结果, 要求是一个1D的Tensor, 数据类型支持int32, 数据格式要求为ND. shape大小需要小于2147483647. | 10545 | +sorted_expert_for_source_row: Tensor类型, 必选参数, 每个source row经过排序后对应的专家索引, 要求是一个1D的Tensor, 数据类型支持int32, 数据格式要求为ND. 输入值取值范围为[0, num_expert-1], 且需按非递减顺序排列; Tensor元素个数需小于2^24. |
| 10546 | -num_expert: int类型, 必选参数, 总专家数. | 10546 | +num_expert: int类型, 必选参数, 总专家数, 取值范围为(0, 2048]. |
| 10547 | 10547 | ||
| 10548 | 输出说明 | 10548 | 输出说明 |
| 10549 | -expertTokens: Tensor类型, 公式中的输出, 要求的是一个1D的Tensor, 数据类型与sorted_expert_for_source_row保持一致. | 10549 | +expertTokens: Tensor类型, 公式中的输出, 要求的是一个1D的Tensor, shape为[num_expert], 数据类型与sorted_expert_for_source_row保持一致. |
| 10550 | 10550 | ||
| 10551 | 约束说明 | 10551 | 约束说明 |
| 10552 | 该接口支持推理场景下使用. | 10552 | 该接口支持推理场景下使用. |
| @@ -10560,14 +10560,15 @@ PyTorch 2.0 | |||
| 10560 | PyTorch 1.11.0 | 10560 | PyTorch 1.11.0 |
| 10561 | 10561 | ||
| 10562 | 支持的型号 | 10562 | 支持的型号 |
| 10563 | +Ascend 950PR/Ascend 950DT | ||
| 10563 | Atlas A2 训练系列产品/Atlas 800I A2 推理产品 | 10564 | Atlas A2 训练系列产品/Atlas 800I A2 推理产品 |
| 10564 | 10565 | ||
| 10565 | 调用示例 | 10566 | 调用示例 |
| 10566 | 单算子模式调用 | 10567 | 单算子模式调用 |
| 10567 | import torch | 10568 | import torch |
| 10568 | import torch_npu | 10569 | import torch_npu |
| 10569 | -sorted_experts = torch.tensor([3,3,4,5,6,7], dtype=torch.int32) | 10570 | +sorted_experts = torch.tensor([0,0,0,1,1,2], dtype=torch.int32) |
| 10570 | -num_experts = 5 | 10571 | +num_experts = 3 |
| 10571 | output = torch_npu.npu_moe_compute_expert_tokens(sorted_experts.npu(), num_experts) | 10572 | output = torch_npu.npu_moe_compute_expert_tokens(sorted_experts.npu(), num_experts) |
| 10572 | 图模式调用 | 10573 | 图模式调用 |
| 10573 | import torch | 10574 | import torch |
| @@ -10584,8 +10585,8 @@ class GMMModel(nn.Module): | |||
| 10584 | def forward(self, sorted_experts, num_experts): | 10585 | def forward(self, sorted_experts, num_experts): |
| 10585 | return torch_npu.npu_moe_compute_expert_tokens(sorted_experts, num_experts) | 10586 | return torch_npu.npu_moe_compute_expert_tokens(sorted_experts, num_experts) |
| 10586 | def main(): | 10587 | def main(): |
| 10587 | - sorted_experts = torch.tensor([3,3,4,5,6,7], dtype=torch.int32) | 10588 | + sorted_experts = torch.tensor([0,0,0,1,1,2], dtype=torch.int32) |
| 10588 | - num_experts = 5 | 10589 | + num_experts = 3 |
| 10589 | model = GMMModel().npu() | 10590 | model = GMMModel().npu() |
| 10590 | model = torch.compile(model, backend=npu_backend, dynamic=False) | 10591 | model = torch.compile(model, backend=npu_backend, dynamic=False) |
| 10591 | custom_output = model(sorted_experts, num_experts) | 10592 | custom_output = model(sorted_experts, num_experts) |
| @@ -14790,7 +14791,7 @@ out(Tensor)：表示公式中的f(x)，即原地更新后的input张量。 | |||
| 14790 | PyTorch 2.6及更高版本 | 14791 | PyTorch 2.6及更高版本 |
| 14791 | 14792 | ||
| 14792 | 支持的型号: | 14793 | 支持的型号: |
| 14793 | -Ascend 950PR/Ascend 950DT | 14794 | +Ascend 950PR/Ascend 950DT |
| 14794 | Atlas A2训练系列产品 | 14795 | Atlas A2训练系列产品 |
| 14795 | Atlas A3训练系列产品 | 14796 | Atlas A3训练系列产品 |
| 14796 | 14797 | ||
| @@ -4,6 +4,7 @@ | |||
| 4 | 4 | ||
| 5 | | 产品 | 是否支持 | | 5 | | 产品 | 是否支持 | |
| 6 | | ------------------------------------------------------------ | :------: | | 6 | | ------------------------------------------------------------ | :------: | |
| 7 | +|<term>Ascend 950PR/Ascend 950DT</term> | √ | | ||
| 7 | |<term>Atlas A2 训练系列产品/Atlas A2 推理系列产品</term> | √ | | 8 | |<term>Atlas A2 训练系列产品/Atlas A2 推理系列产品</term> | √ | |
| 8 | 9 | ||
| 9 | ## 功能说明 | 10 | ## 功能说明 |
| @@ -12,7 +13,7 @@ | |||
| 12 | - 计算公式： | 13 | - 计算公式： |
| 13 | 14 | ||
| 14 | $$ | 15 | $$ |
| 15 | - expertTokens_i = BinarySearch(sortedExpertForSourceRow,numExpert) | 16 | + expertTokens_i = BinarySearch(sortedExpertForSourceRow,i), \quad i \in [0,numExpert) |
| 16 | $$ | 17 | $$ |
| 17 | 18 | ||
| 18 | ## 函数原型 | 19 | ## 函数原型 |
| @@ -23,15 +24,15 @@ torch_npu.npu_moe_compute_expert_tokens(sorted_expert_for_source_row, num_expert | |||
| 23 | 24 | ||
| 24 | ## 参数说明 | 25 | ## 参数说明 |
| 25 | 26 | ||
| 26 | -- **sorted_expert_for_source_row** (`Tensor`)：必选参数，每个source row经过排序后对应的专家索引，对应公式中的$sortedExpertForSourceRow$，要求是一个1维变量，数据类型支持int32，数据格式要求为$ND$。shape需小于2147483647。 | 27 | +- **sorted_expert_for_source_row** (`Tensor`)：必选参数，每个source row经过排序后对应的专家索引，对应公式中的$sortedExpertForSourceRow$，要求是一个1维变量，数据类型支持int32，数据格式要求为$ND$。输入值取值范围为[0, `num_expert`-1]，且需按非递减顺序排列；Tensor元素个数需小于$2^{24}$。 |
| 27 | 28 | ||
| 28 | -- **num_expert** (`int`)：必选参数，表示总专家数。对应公式中的$numExpert$。 | 29 | +- **num_expert** (`int`)：必选参数，表示总专家数。对应公式中的$numExpert$，取值范围为(0, 2048]。 |
| 29 | 30 | ||
| 30 | ## 返回值说明 | 31 | ## 返回值说明 |
| 31 | 32 | ||
| 32 | `Tensor` | 33 | `Tensor` |
| 33 | 34 | ||
| 34 | - 对应公式中的$expertTokens$，要求的是一个1维张量，数据类型与`sorted_expert_for_source_row`保持一致。 | 35 | + 对应公式中的$expertTokens$，要求的是一个1维张量，shape为[`num_expert`]，数据类型与`sorted_expert_for_source_row`保持一致。 |
| 35 | 36 | ||
| 36 | ## 约束说明 | 37 | ## 约束说明 |
| 37 | 38 | ||