已合并
Add obfuscation_initialize obfucation_finialize obfuscation_calculate #3043
jiang_xue_bo创建于 2025年8月21日
Add obfuscation_initialize obfucation_finialize obfuscation_calculate #3043
已合并
从refs/pull/3043/head合入到master
共 14 个文件变更+569-0
| @@ -9918,3 +9918,135 @@ probs = torch.ones_like(indices) / 2 | |||
| 9918 | unpermuted_tokens = torch_npu.npu_moe_token_unpermute(permuted_tokens, sorted_indices, probs=probs) | 9918 | unpermuted_tokens = torch_npu.npu_moe_token_unpermute(permuted_tokens, sorted_indices, probs=probs) |
| 9919 | """ | 9919 | """ |
| 9920 | ) | 9920 | ) |
| 9921 | + | ||
| 9922 | + | ||
| 9923 | +_add_torch_npu_docstr( | ||
| 9924 | + "obfuscation_initialize", | ||
| 9925 | + """ | ||
| 9926 | +功能描述: | ||
| 9927 | +该接口用于完成PMCC模型混淆引擎的资源初始化。 | ||
| 9928 | + | ||
| 9929 | +接口原型: | ||
| 9930 | +torch_npu.obfuscation_initialize(int hidden_size, int tp_rank, int cmd, int data_type, int model_obf_seed_id, int data_obf_seed_id, int thread_num, float obf_coefficient) -> Tensor | ||
| 9931 | + | ||
| 9932 | +参数说明: | ||
| 9933 | +- hidden_size(Int, 计算输入): 必选参数,隐藏层的维度,数据类型为int32,支持输入范围为1-10000,仅在cmd设置为1或2时需要填写有效值,否则填0。 | ||
| 9934 | +- tp_rank(Int, 计算输入): 必选参数, 张量并行TP Rank,数据类型为int32,支持输入范围为0-1024,仅在cmd设置为1或2时需要填写有效值,否则填0。 | ||
| 9935 | +- cmd(Int, 计算输入): 必选参数, setup指令编号,在{1, 2, 3}中选择,设置为1时进行浮点推理模式资源初始化、为2时进行量化推理模式资源初始化,设置为3时进行资源释放。数据类型为int32。 | ||
| 9936 | +- data_type(Int, 计算输入): 可选参数, 代表Tensor数据类型的编号,数据类型为INT32,仅在cmd设置为1或2时需要填写有效值,否则填0。 | ||
| 9937 | + 昇腾310P AI处理器:在{0, 1}中选择,0表示ACL_FLOAT、1表示ACL_FLOAT16 | ||
| 9938 | + 昇腾910B AI处理器:在{0, 1, 27}中选择,0表示ACL_FLOAT、1表示ACL_FLOAT16、27表示ACL_BF16 | ||
| 9939 | +- model_obf_seed_id(Int, 计算输入): 可选参数, 模型混淆因子id,用于TA从TEE KMC查询模型混淆因子,数据类型为int32,仅在cmd设置为1或2时需要填写有效值,否则填0。 | ||
| 9940 | +- data_obf_seed_id(Int, 计算输入): 必选参数, 数据混淆因子id,用于TA从TEE KMC查询数据混淆因子,数据类型为int32,仅在cmd设置为1或2时需要填写有效值,否则填0。 | ||
| 9941 | +- thread_num(Int, 计算输入): 可选参数, CA/TA进行混淆处理使用的线程数。在{1, 2, 3, 4, 5, 6}中选择,数据类型为int32,仅在cmd设置为1或2时需要填写有效值,否则填0。 | ||
| 9942 | +- obf_coefficient(Float, 计算输入): 可选参数,混淆系数,数据类型为float,支持输入范围为0-1,默认值1.0。 | ||
| 9943 | +- fd(Tensor, 计算输出): socket连接符,1D,shape为(1),int32。 | ||
| 9944 | + | ||
| 9945 | +支持的芯片型号: | ||
| 9946 | +Atlas A2 训练系列产品/Atlas 800I A2 推理产品/A200I A2 Box 异构组件 | ||
| 9947 | +Atlas 推理系列产品 | ||
| 9948 | + | ||
| 9949 | +调用示例: | ||
| 9950 | +# 单算子调用 | ||
| 9951 | +import torch | ||
| 9952 | +import torch_npu | ||
| 9953 | + | ||
| 9954 | +device = "npu:0" | ||
| 9955 | +hidden_size = int(3584) | ||
| 9956 | +cmd = 1 | ||
| 9957 | +data_type = torch.bfloat16 | ||
| 9958 | +model_obf_seed = 0 | ||
| 9959 | +data_obf_seed = 0 | ||
| 9960 | +thread_num = 4 | ||
| 9961 | +tp_rank = 0 | ||
| 9962 | +i = 0 | ||
| 9963 | +hidden_states = torch.randn((1024,3584), dtype=torch.bfloat16, device=device) | ||
| 9964 | +obf_cft = 1.0 | ||
| 9965 | +fd = torch_npu.obfuscation_initialize(hidden_size, tp_rank, cmd, data_type=data_type, thread_num= thread_num, obf_coefficient=obf_cft) | ||
| 9966 | +""" | ||
| 9967 | +) | ||
| 9968 | + | ||
| 9969 | + | ||
| 9970 | +_add_torch_npu_docstr( | ||
| 9971 | + "obfuscation_finalize", | ||
| 9972 | + """ | ||
| 9973 | +功能描述: | ||
| 9974 | +该接口用于完成PMCC模型混淆引擎的资源释放。 | ||
| 9975 | + | ||
| 9976 | +接口原型: | ||
| 9977 | +torch_npu.obfuscation_finalize(Tensor fd_to_close) -> Tensor | ||
| 9978 | + | ||
| 9979 | +参数说明: | ||
| 9980 | +- fd_to_close(Tensor, 计算输入): 填写本算子初始化时返回的fd。数据类型为int32。 | ||
| 9981 | + | ||
| 9982 | +支持的芯片型号: | ||
| 9983 | +Atlas A2 训练系列产品/Atlas 800I A2 推理产品/A200I A2 Box 异构组件 | ||
| 9984 | +Atlas 推理系列产品 | ||
| 9985 | + | ||
| 9986 | +调用示例: | ||
| 9987 | +# 单算子调用 | ||
| 9988 | +import torch | ||
| 9989 | +import torch_npu | ||
| 9990 | + | ||
| 9991 | +device = "npu:0" | ||
| 9992 | +hidden_size = int(3584) | ||
| 9993 | +cmd = 1 | ||
| 9994 | +data_type = torch.bfloat16 | ||
| 9995 | +model_obf_seed = 0 | ||
| 9996 | +data_obf_seed = 0 | ||
| 9997 | +thread_num = 4 | ||
| 9998 | +tp_rank = 0 | ||
| 9999 | +i = 0 | ||
| 10000 | +hidden_states = torch.randn((1024,3584), dtype=torch.bfloat16, device=device) | ||
| 10001 | +obf_cft = 1.0 | ||
| 10002 | +fd = torch_npu.obfuscation_initialize(hidden_size, tp_rank, cmd, data_type=data_type, thread_num= thread_num, obf_coefficient=obf_cft) | ||
| 10003 | +torch_npu.obfuscation_finalize(fd) | ||
| 10004 | +""" | ||
| 10005 | +) | ||
| 10006 | + | ||
| 10007 | + | ||
| 10008 | +_add_torch_npu_docstr( | ||
| 10009 | + "obfuscation_calculate", | ||
| 10010 | + """ | ||
| 10011 | +功能描述: | ||
| 10012 | +该接口用于实现矩阵乘计算输入和输出的transpose操作。 | ||
| 10013 | + | ||
| 10014 | +接口原型: | ||
| 10015 | +torch_npu.obfuscation_calculate(Tensor fd, Tensor x, Tensor param, float obf_coefficient) -> Tensor | ||
| 10016 | + | ||
| 10017 | +参数说明: | ||
| 10018 | +- fd(Tensor, 计算输入): 必选参数,待关闭的socket连接符,cmd为3时,填写本算子在cmd为1时返回的fd,否则填0,数据类型为int32。 | ||
| 10019 | +- x(Tensor, 计算输入): 必选参数,待混淆处理的Tensor输入,数据类型如下,对Tensor维度不作限制,Shape为( , *, ... , hiddenSize),即最后一维的size是hiddenSize。数据格式支持ND。 | ||
| 10020 | + 昇腾310P AI处理器: Tensor数据类型支持ACL_FLOAT、 ACL_FLOAT16 torch.float16 / torch.float32 / torch.bfloat16 | ||
| 10021 | + 昇腾910B AI处理器: Tensor数据类型支持ACL_FLOAT、 ACL_FLOAT16、ACL_BF16 | ||
| 10022 | +- param(Tensor, 计算输入): 必选参数,预留的参数字段,Tensor数据类型为int32。 | ||
| 10023 | +- obf_coefficient(Float, 计算输入): 可选参数,混淆系数,数据类型为float,支持输入范围为0-1,默认值1.0。 | ||
| 10024 | +- y(Tensor, 计算输出): 混淆处理后的张量,输出数据类型及Shape与x相同。 | ||
| 10025 | + | ||
| 10026 | + | ||
| 10027 | +支持的芯片型号: | ||
| 10028 | +Atlas A2 训练系列产品/Atlas 800I A2 推理产品/A200I A2 Box 异构组件 | ||
| 10029 | +Atlas 推理系列产品 | ||
| 10030 | + | ||
| 10031 | +调用示例: | ||
| 10032 | +# 单算子调用 | ||
| 10033 | +import torch | ||
| 10034 | +import torch_npu | ||
| 10035 | + | ||
| 10036 | +device = "npu:0" | ||
| 10037 | +hidden_size = int(3584) | ||
| 10038 | +cmd = 1 | ||
| 10039 | +data_type = torch.bfloat16 | ||
| 10040 | +model_obf_seed = 0 | ||
| 10041 | +data_obf_seed = 0 | ||
| 10042 | +thread_num = 4 | ||
| 10043 | +tp_rank = 0 | ||
| 10044 | +i = 0 | ||
| 10045 | +hidden_states = torch.randn((1024,3584), dtype=torch.bfloat16, device=device) | ||
| 10046 | +obf_cft = 1.0 | ||
| 10047 | +fd = torch_npu.obfuscation_initialize(hidden_size, tp_rank, cmd, data_type=data_type, thread_num= thread_num, obf_coefficient=obf_cft) | ||
| 10048 | +param = torch.tensor([3584], device=device) | ||
| 10049 | +x_obf_out = torch_npu.obfuscation_calculate(fd, hidden_states, param, obf_coefficient=obf_cft) | ||
| 10050 | + | ||
| 10051 | +""" | ||
| 10052 | +) | ||
| @@ -211,6 +211,9 @@ | |||
| 211 | - [torch_npu.npu.graph_task_group_end](./context/torch_npu-npu-graph_task_group_end.md) | 211 | - [torch_npu.npu.graph_task_group_end](./context/torch_npu-npu-graph_task_group_end.md) |
| 212 | - [torch_npu.npu.graph_task_update_begin](./context/torch_npu-npu-graph_task_update_begin.md) | 212 | - [torch_npu.npu.graph_task_update_begin](./context/torch_npu-npu-graph_task_update_begin.md) |
| 213 | - [torch_npu.npu.graph_task_update_end](./context/torch_npu-npu-graph_task_update_end.md) | 213 | - [torch_npu.npu.graph_task_update_end](./context/torch_npu-npu-graph_task_update_end.md) |
| 214 | + - [torch_npu.obfuscation_initialize](./context/torch_npu-obfuscation_initialize.md) | ||
| 215 | + - [torch_npu.obfuscation_finalize](./context/torch_npu-obfuscation_finalize.md) | ||
| 216 | + - [torch_npu.obfuscation_calculate](./context/torch_npu-obfuscation_calculate.md) | ||
| 214 | 217 | ||
| 215 | - [amp](./context/amp.md) | 218 | - [amp](./context/amp.md) |
| 216 | - [（beta）torch_npu.npu.get_amp_supported_dtype](./context/（beta）torch_npu-npu-get_amp_supported_dtype.md) | 219 | - [（beta）torch_npu.npu.get_amp_supported_dtype](./context/（beta）torch_npu-npu-get_amp_supported_dtype.md) |
| @@ -0,0 +1,54 @@ | |||
| 1 | + | ||
| 2 | + | ||
| 3 | +# torch_npu.npu.obfuscation_calculate | ||
| 4 | + | ||
| 5 | +## 产品支持情况 | ||
| 6 | + | ||
| 7 | +| 产品 | 是否支持 | | ||
| 8 | +| ------------------------------------------------------------ | :------: | | ||
| 9 | +|<term>Atlas A2 训练系列产品/Atlas 800I A2 推理产品/A200I A2 Box 异构组件</term> | √ | | ||
| 10 | +|<term>Atlas 推理系列产品 </term> | √ | | ||
| 11 | + | ||
| 12 | +## 功能说明 | ||
| 13 | + | ||
| 14 | +该接口用于实现矩阵乘计算输入和输出的transpose操作。 | ||
| 15 | + | ||
| 16 | +## 函数原型 | ||
| 17 | + | ||
| 18 | +``` | ||
| 19 | +torch_npu.obfuscation_calculate(Tensor fd, Tensor x, Tensor param, float obf_coefficient) -> Tensor | ||
| 20 | +``` | ||
| 21 | + | ||
| 22 | +## 参数说明 | ||
| 23 | + | ||
| 24 | +- **x**（`Tensor`）：必选参数,待混淆处理的Tensor输入,数据类型如下,对Tensor维度不作限制,Shape为( , *, ... , hiddenSize),即最后一维的size是hiddenSize。数据格式支持ND。 | ||
| 25 | + * <term>Atlas 推理系列产品</term>: Tensor数据类型支持torch.float16 / torch.float32 / torch.int8 | ||
| 26 | + * <term>Atlas A2 训练系列产品/Atlas 800I A2 推理产品/A200I A2 Box 异构组件</term>: Tensor数据类型支持torch.float16 / torch.float32 / torch.bfloat16 / torch.int8 | ||
| 27 | +- **param**（`Tensor`）：必选参数,预留的参数字段,Tensor数据类型为int32。 | ||
| 28 | +- **obf_coefficient**（`Float`）：可选参数,混淆系数,数据类型为float,支持输入范围为0-1,默认值1.0。 | ||
| 29 | + | ||
| 30 | +## 返回值 | ||
| 31 | + | ||
| 32 | +- **y**（`Tensor`）：混淆处理后的张量,输出数据类型及Shape与x相同。 | ||
| 33 | + | ||
| 34 | +## 调用示例 | ||
| 35 | + | ||
| 36 | +```python | ||
| 37 | +import torch | ||
| 38 | +import torch_npu | ||
| 39 | + | ||
| 40 | +device = "npu:0" | ||
| 41 | +hidden_size = int(3584) | ||
| 42 | +cmd = 1 | ||
| 43 | +data_type = torch.bfloat16 | ||
| 44 | +model_obf_seed = 0 | ||
| 45 | +data_obf_seed = 0 | ||
| 46 | +thread_num = 4 | ||
| 47 | +tp_rank = 0 | ||
| 48 | +i = 0 | ||
| 49 | +hidden_states = torch.randn((1024,3584), dtype=torch.bfloat16, device=device) | ||
| 50 | +obf_cft = 1.0 | ||
| 51 | +fd = torch_npu.obfuscation_initialize(hidden_size, tp_rank, cmd, data_type=data_type, thread_num= thread_num, obf_coefficient=obf_cft) | ||
| 52 | +param = torch.tensor([3584], device=device) | ||
| 53 | +x_obf_out = torch_npu.obfuscation_calculate(fd, hidden_states, param, obf_coefficient=obf_cft) | ||
| 54 | +``` | ||
| @@ -0,0 +1,45 @@ | |||
| 1 | + | ||
| 2 | + | ||
| 3 | +# torch_npu.npu.obfuscation_finalize | ||
| 4 | + | ||
| 5 | +## 产品支持情况 | ||
| 6 | + | ||
| 7 | +| 产品 | 是否支持 | | ||
| 8 | +| ------------------------------------------------------------ | :------: | | ||
| 9 | +|<term>Atlas A2 训练系列产品/Atlas 800I A2 推理产品/A200I A2 Box 异构组件</term> | √ | | ||
| 10 | +|<term>Atlas 推理系列产品 </term> | √ | | ||
| 11 | + | ||
| 12 | +## 功能说明 | ||
| 13 | + | ||
| 14 | +该接口用于完成PMCC模型混淆引擎的资源释放。 | ||
| 15 | + | ||
| 16 | +## 函数原型 | ||
| 17 | + | ||
| 18 | +``` | ||
| 19 | +torch_npu.obfuscation_finalize(Tensor fd_to_close) -> Tensor | ||
| 20 | +``` | ||
| 21 | + | ||
| 22 | +## 参数说明 | ||
| 23 | + | ||
| 24 | +- **fd_to_close**（`Tensor`）：填写调用obfuscation_initialize返回的fd。数据类型为int32。 | ||
| 25 | + | ||
| 26 | +## 调用示例 | ||
| 27 | + | ||
| 28 | +```python | ||
| 29 | +import torch | ||
| 30 | +import torch_npu | ||
| 31 | + | ||
| 32 | +device = "npu:0" | ||
| 33 | +hidden_size = int(3584) | ||
| 34 | +cmd = 1 | ||
| 35 | +data_type = torch.bfloat16 | ||
| 36 | +model_obf_seed = 0 | ||
| 37 | +data_obf_seed = 0 | ||
| 38 | +thread_num = 4 | ||
| 39 | +tp_rank = 0 | ||
| 40 | +i = 0 | ||
| 41 | +hidden_states = torch.randn((1024,3584), dtype=torch.bfloat16, device=device) | ||
| 42 | +obf_cft = 1.0 | ||
| 43 | +fd = torch_npu.obfuscation_initialize(hidden_size, tp_rank, cmd, data_type=data_type, thread_num= thread_num, obf_coefficient=obf_cft) | ||
| 44 | +torch_npu.obfuscation_finalize(fd) | ||
| 45 | +``` | ||
| @@ -0,0 +1,57 @@ | |||
| 1 | + | ||
| 2 | + | ||
| 3 | +# torch_npu.npu.obfuscation_initialize | ||
| 4 | + | ||
| 5 | +## 产品支持情况 | ||
| 6 | + | ||
| 7 | +| 产品 | 是否支持 | | ||
| 8 | +| ------------------------------------------------------------ | :------: | | ||
| 9 | +|<term>Atlas A2 训练系列产品/Atlas 800I A2 推理产品/A200I A2 Box 异构组件</term> | √ | | ||
| 10 | +|<term>Atlas 推理系列产品 </term> | √ | | ||
| 11 | + | ||
| 12 | +## 功能说明 | ||
| 13 | + | ||
| 14 | +该接口用于完成PMCC模型混淆引擎的资源初始化。 | ||
| 15 | + | ||
| 16 | +## 函数原型 | ||
| 17 | + | ||
| 18 | +``` | ||
| 19 | +torch_npu.obfuscation_initialize(int hidden_size, int tp_rank, int cmd, int data_type, int model_obf_seed_id, int data_obf_seed_id, int thread_num, float obf_coefficient) -> Tensor | ||
| 20 | +``` | ||
| 21 | + | ||
| 22 | +## 参数说明 | ||
| 23 | + | ||
| 24 | +- **hidden_size**（`Int`）：必选参数,隐藏层的维度,数据类型为int32,支持输入范围为1-10000,仅在cmd设置为1或2时需要填写有效值,否则填0。 | ||
| 25 | +- **tp_rank**（`Int`）：必选参数, 张量并行TP Rank,数据类型为int32,支持输入范围为0-1024,仅在cmd设置为1或2时需要填写有效值,否则填0。 | ||
| 26 | +- **cmd**（`Int`）：必选参数, setup指令编号,在{1, 2, 3}中选择,设置为1时进行浮点推理模式资源初始化、为2时进行量化推理模式资源初始化,设置为3时进行资源释放。数据类型为int32。 | ||
| 27 | +- **data_type**（`Int`）：可选参数, 代表Tensor数据类型的编号,数据类型为INT32,仅在cmd设置为1或2时需要填写有效值,否则填0。 | ||
| 28 | + * <term>Atlas 推理系列产品</term>: Tensor数据类型支持torch.float16 / torch.float32 / torch.int8 | ||
| 29 | + * <term>Atlas A2 训练系列产品/Atlas 800I A2 推理产品/A200I A2 Box 异构组件</term>: Tensor数据类型支持torch.float16 / torch.float32 / torch.bfloat16 / torch.int8 | ||
| 30 | +- **model_obf_seed_id**（`Int`）：可选参数, 模型混淆因子id,用于TA从TEE KMC查询模型混淆因子,数据类型为int32,仅在cmd设置为1或2时需要填写有效值,否则填0。 | ||
| 31 | +- **data_obf_seed_id**（`Int`）：必选参数, 数据混淆因子id,用于TA从TEE KMC查询数据混淆因子,数据类型为int32,仅在cmd设置为1或2时需要填写有效值,否则填0。 | ||
| 32 | +- **thread_num**（`Int`）：可选参数, CA/TA进行混淆处理使用的线程数。在{1, 2, 3, 4, 5, 6}中选择,数据类型为int32,仅在cmd设置为1或2时需要填写有效值,否则填0。 | ||
| 33 | +- **obf_coefficient**（`Float`）：可选参数,混淆系数,数据类型为float,支持输入范围为0-1,默认值1.0。 | ||
| 34 | + | ||
| 35 | +## 返回值 | ||
| 36 | + | ||
| 37 | +- **fd**（`Tensor`）：socket连接符,1D,shape为(1),int32。 | ||
| 38 | + | ||
| 39 | +## 调用示例 | ||
| 40 | + | ||
| 41 | +```python | ||
| 42 | +import torch | ||
| 43 | +import torch_npu | ||
| 44 | + | ||
| 45 | +device = "npu:0" | ||
| 46 | +hidden_size = int(3584) | ||
| 47 | +cmd = 1 | ||
| 48 | +data_type = torch.bfloat16 | ||
| 49 | +model_obf_seed = 0 | ||
| 50 | +data_obf_seed = 0 | ||
| 51 | +thread_num = 4 | ||
| 52 | +tp_rank = 0 | ||
| 53 | +i = 0 | ||
| 54 | +hidden_states = torch.randn((1024,3584), dtype=torch.bfloat16, device=device) | ||
| 55 | +obf_cft = 1.0 | ||
| 56 | +fd = torch_npu.obfuscation_initialize(hidden_size, tp_rank, cmd, data_type=data_type, thread_num= thread_num, obf_coefficient=obf_cft) | ||
| 57 | +``` | ||
| @@ -7581,6 +7581,18 @@ custom: | |||
| 7581 | acl_op: all_version | 7581 | acl_op: all_version |
| 7582 | exposed: all_version | 7582 | exposed: all_version |
| 7583 | 7583 | ||
| 7584 | + - func: obfuscation_initialize(int hidden_size, int tp_rank, int cmd, *, ScalarType? data_type=None, int? model_obf_seed_id=0, int? data_obf_seed_id=0, int? thread_num=0, float? obf_coefficient=1.0) -> Tensor | ||
| 7585 | + op_api: all_version | ||
| 7586 | + internal_format_opapi: all_version | ||
| 7587 | + dispatch: | ||
| 7588 | + CompositeExplicitAutograd: obfuscation_initialize | ||
| 7589 | + | ||
| 7590 | + - func: obfuscation_finalize(Tensor fd_to_close) -> Tensor | ||
| 7591 | + op_api: all_version | ||
| 7592 | + | ||
| 7593 | + - func: obfuscation_calculate(Tensor fd, Tensor x, Tensor param, *, float? obf_coefficient=1.0) -> Tensor | ||
| 7594 | + op_api: all_version | ||
| 7595 | + | ||
| 7584 | - func: one_(Tensor(a!) self) -> Tensor(a!) | 7596 | - func: one_(Tensor(a!) self) -> Tensor(a!) |
| 7585 | acl_op: all_version | 7597 | acl_op: all_version |
| 7586 | op_api: all_version | 7598 | op_api: all_version |
| @@ -0,0 +1,56 @@ | |||
| 1 | +// Copyright (c) 2025 Huawei Technologies Co., Ltd | ||
| 2 | +// All rights reserved. | ||
| 3 | +// | ||
| 4 | +// Licensed under the BSD 3-Clause License (the "License"); | ||
| 5 | +// you may not use this file except in compliance with the License. | ||
| 6 | +// You may obtain a copy of the License at | ||
| 7 | +// | ||
| 8 | +// https://opensource.org/licenses/BSD-3-Clause | ||
| 9 | +// | ||
| 10 | +// Unless required by applicable law or agreed to in writing, software | ||
| 11 | +// distributed under the License is distributed on an "AS IS" BASIS, | ||
| 12 | +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. | ||
| 13 | +// See the License for the specific language governing permissions and | ||
| 14 | +// limitations under the License. | ||
| 15 | + | ||
| 16 | + | ||
| 17 | + | ||
| 18 | + | ||
| 19 | +namespace op_api { | ||
| 20 | +using npu_preparation = at_npu::native::OpPreparation; | ||
| 21 | + | ||
| 22 | +at::Tensor obfuscation_calculate( | ||
| 23 | + const at::Tensor &fd, const at::Tensor &x, | ||
| 24 | + const at::Tensor ¶m, c10::optional<double> obf_coefficient | ||
| 25 | + ) | ||
| 26 | +{ | ||
| 27 | + int32_t fd_real; | ||
| 28 | + int32_t param_real; | ||
| 29 | + if (!fd.defined()) { | ||
| 30 | + throw std::runtime_error("fd cannot be empty"); | ||
| 31 | + } | ||
| 32 | + try { | ||
| 33 | + auto fd_scalar = fd[0].item(); | ||
| 34 | + fd_real = fd_scalar.to<int32_t>(); | ||
| 35 | + } catch (const std::exception& e) { | ||
| 36 | + throw std::runtime_error("Failed to process fd: " + std::string(e.what())); | ||
| 37 | + } | ||
| 38 | + if (!param.defined()) { | ||
| 39 | + throw std::runtime_error("param cannot be empty"); | ||
| 40 | + } | ||
| 41 | + try { | ||
| 42 | + auto param_scalar = param[0].item(); | ||
| 43 | + param_real = param_scalar.to<int32_t>(); | ||
| 44 | + } catch (const std::exception& e) { | ||
| 45 | + throw std::runtime_error("Failed to process param: " + std::string(e.what())); | ||
| 46 | + } | ||
| 47 | + auto cmd_real = static_cast<int32_t>(1); | ||
| 48 | + auto obf_coefficient_real = static_cast<float>(obf_coefficient.value_or(1)); | ||
| 49 | + auto out_size = op_infer::array_to_small_vector(x.sizes()); | ||
| 50 | + auto out_type = x.scalar_type(); | ||
| 51 | + c10::TensorOptions options = x.options().dtype(out_type); | ||
| 52 | + at::Tensor y = npu_preparation::apply_tensor_without_format(out_size, options); | ||
| 53 | + EXEC_NPU_CMD(aclnnObfuscationCalculateV2, fd_real, x, param_real, cmd_real, obf_coefficient_real, y); | ||
| 54 | + return y; | ||
| 55 | +} | ||
| 56 | +} | ||
| @@ -0,0 +1,42 @@ | |||
| 1 | +// Copyright (c) 2025 Huawei Technologies Co., Ltd | ||
| 2 | +// All rights reserved. | ||
| 3 | +// | ||
| 4 | +// Licensed under the BSD 3-Clause License (the "License"); | ||
| 5 | +// you may not use this file except in compliance with the License. | ||
| 6 | +// You may obtain a copy of the License at | ||
| 7 | +// | ||
| 8 | +// https://opensource.org/licenses/BSD-3-Clause | ||
| 9 | +// | ||
| 10 | +// Unless required by applicable law or agreed to in writing, software | ||
| 11 | +// distributed under the License is distributed on an "AS IS" BASIS, | ||
| 12 | +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. | ||
| 13 | +// See the License for the specific language governing permissions and | ||
| 14 | +// limitations under the License. | ||
| 15 | + | ||
| 16 | + | ||
| 17 | + | ||
| 18 | + | ||
| 19 | +namespace op_api { | ||
| 20 | +using npu_preparation = at_npu::native::OpPreparation; | ||
| 21 | + | ||
| 22 | +at::Tensor obfuscation_finalize(const at::Tensor &fd_to_close) | ||
| 23 | +{ | ||
| 24 | + auto out_type = fd_to_close.scalar_type(); | ||
| 25 | + auto out_size = op_infer::array_to_small_vector(fd_to_close.sizes()); | ||
| 26 | + c10::TensorOptions options = fd_to_close.options().dtype(out_type); | ||
| 27 | + at::Tensor fd = npu_preparation::apply_tensor_without_format(out_size, options); | ||
| 28 | + auto fd_to_close_scalar = fd_to_close[0].item(); | ||
| 29 | + auto fd_to_close_real = fd_to_close_scalar.to<int32_t>(); | ||
| 30 | + int32_t data_type = 0; | ||
| 31 | + int32_t hidden_size = 0; | ||
| 32 | + int32_t tp_rank = 0; | ||
| 33 | + int32_t model_obf_seed_id = 0; | ||
| 34 | + int32_t data_obf_seed_id = 0; | ||
| 35 | + int32_t mode = 16; | ||
| 36 | + int32_t thread_num = 0; | ||
| 37 | + float obf_cft = 1.0; | ||
| 38 | + EXEC_NPU_CMD(aclnnObfuscationSetupV2, fd_to_close_real, data_type, hidden_size, tp_rank, model_obf_seed_id, | ||
| 39 | + data_obf_seed_id, mode, thread_num, obf_cft, fd); | ||
| 40 | + return fd; | ||
| 41 | +} | ||
| 42 | +} | ||
| @@ -0,0 +1,75 @@ | |||
| 1 | +// Copyright (c) 2025 Huawei Technologies Co., Ltd | ||
| 2 | +// All rights reserved. | ||
| 3 | +// | ||
| 4 | +// Licensed under the BSD 3-Clause License (the "License"); | ||
| 5 | +// you may not use this file except in compliance with the License. | ||
| 6 | +// You may obtain a copy of the License at | ||
| 7 | +// | ||
| 8 | +// https://opensource.org/licenses/BSD-3-Clause | ||
| 9 | +// | ||
| 10 | +// Unless required by applicable law or agreed to in writing, software | ||
| 11 | +// distributed under the License is distributed on an "AS IS" BASIS, | ||
| 12 | +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. | ||
| 13 | +// See the License for the specific language governing permissions and | ||
| 14 | +// limitations under the License. | ||
| 15 | + | ||
| 16 | + | ||
| 17 | + | ||
| 18 | + | ||
| 19 | +namespace op_api { | ||
| 20 | +using npu_preparation = at_npu::native::OpPreparation; | ||
| 21 | +enum class ObfDataType : int32_t { | ||
| 22 | + FLOAT = 0, | ||
| 23 | + HALF = 1, | ||
| 24 | + CHAR = 2, | ||
| 25 | + BFLOAT = 27, | ||
| 26 | + UNDEFINED = -1 | ||
| 27 | +}; | ||
| 28 | + | ||
| 29 | +at::Tensor obfuscation_initialize( | ||
| 30 | + int64_t hidden_size, int64_t tp_rank, | ||
| 31 | + int64_t cmd, c10::optional<c10::ScalarType> data_type, | ||
| 32 | + c10::optional<int64_t> model_obf_seed_id, | ||
| 33 | + c10::optional<int64_t> data_obf_seed_id, | ||
| 34 | + c10::optional<int64_t> thread_num, | ||
| 35 | + c10::optional<double> obf_coefficient | ||
| 36 | + ) | ||
| 37 | +{ | ||
| 38 | + auto hidden_size_real = static_cast<int32_t>(hidden_size); | ||
| 39 | + auto tp_rank_real = static_cast<int32_t>(tp_rank); | ||
| 40 | + auto cmd_real = static_cast<int32_t>(cmd); | ||
| 41 | + int32_t data_type_real; | ||
| 42 | + if (!data_type.has_value()) { | ||
| 43 | + throw std::runtime_error("data_type cannot be empty"); | ||
| 44 | + } | ||
| 45 | + | ||
| 46 | + switch (data_type.value()) { | ||
| 47 | + case at::ScalarType::Half: | ||
| 48 | + data_type_real = static_cast<int32_t>(ObfDataType::HALF); | ||
| 49 | + break; | ||
| 50 | + case at::ScalarType::Float: | ||
| 51 | + data_type_real = static_cast<int32_t>(ObfDataType::FLOAT); | ||
| 52 | + break; | ||
| 53 | + case at::ScalarType::Char: | ||
| 54 | + data_type_real = static_cast<int32_t>(ObfDataType::CHAR); | ||
| 55 | + break; | ||
| 56 | + case at::ScalarType::BFloat16: | ||
| 57 | + data_type_real = static_cast<int32_t>(ObfDataType::BFLOAT); | ||
| 58 | + break; | ||
| 59 | + default: | ||
| 60 | + data_type_real = static_cast<int32_t>(ObfDataType::UNDEFINED); | ||
| 61 | + } | ||
| 62 | + auto model_obf_seed_id_real = static_cast<int32_t>(model_obf_seed_id.value_or(0)); | ||
| 63 | + auto data_obf_seed_id_real = static_cast<int32_t>(data_obf_seed_id.value_or(0)); | ||
| 64 | + auto thread_num_real = static_cast<int32_t>(thread_num.value_or(4)); | ||
| 65 | + auto obf_coefficient_real = static_cast<float>(obf_coefficient.value_or(1)); | ||
| 66 | + at::ScalarType out_type = at::ScalarType::Int; | ||
| 67 | + c10::SmallVector<int64_t> out_size = {1}; | ||
| 68 | + auto options = c10::TensorOptions().device(c10::DeviceType::PrivateUse1).dtype(out_type); | ||
| 69 | + int32_t fd_to_close = 0; | ||
| 70 | + at::Tensor fd = npu_preparation::apply_tensor_without_format(out_size, options); | ||
| 71 | + EXEC_NPU_CMD(aclnnObfuscationSetupV2, fd_to_close, data_type_real, hidden_size_real, tp_rank_real, model_obf_seed_id_real, | ||
| 72 | + data_obf_seed_id_real, cmd_real, thread_num_real, obf_coefficient_real, fd); | ||
| 73 | + return fd; | ||
| 74 | +} | ||
| 75 | +} | ||
| @@ -1801,6 +1801,16 @@ def quant_matmul_scale_offset_out_check(scale, offset, pertoken_scale, output_dt | |||
| 1801 | ) | 1801 | ) |
| 1802 | 1802 | ||
| 1803 | 1803 | ||
| 1804 | + | ||
| 1805 | +def obfuscation_calculate_meta(fd, x, param, cmd): | ||
| 1806 | + return torch.empty_like(x) | ||
| 1807 | + | ||
| 1808 | + | ||
| 1809 | + | ||
| 1810 | +def obfuscation_finalize_meta(fd_to_close): | ||
| 1811 | + return torch.empty_like(fd_to_close) | ||
| 1812 | + | ||
| 1813 | + | ||
| 1804 | 1814 | ||
| 1805 | def npu_quant_matmul_meta(x1, x2, scale, *, offset=None, pertoken_scale=None, bias=None, output_dtype=None, group_sizes=None): | 1815 | def npu_quant_matmul_meta(x1, x2, scale, *, offset=None, pertoken_scale=None, bias=None, output_dtype=None, group_sizes=None): |
| 1806 | INT4_IN_INT32 = 8 | 1816 | INT4_IN_INT32 = 8 |
| @@ -2903,6 +2903,15 @@ | |||
| 2903 | "func: normal_(Tensor(a!) self, float mean=0, float std=1, *, Generator? generator=None) -> Tensor(a!)": { | 2903 | "func: normal_(Tensor(a!) self, float mean=0, float std=1, *, Generator? generator=None) -> Tensor(a!)": { |
| 2904 | "version": ["all_version"] | 2904 | "version": ["all_version"] |
| 2905 | }, | 2905 | }, |
| 2906 | + "func: obfuscation_initialize(int hidden_size, int tp_rank, int cmd, *, ScalarType? data_type=None, int? model_obf_seed_id=0, int? data_obf_seed_id=0, int? thread_num=0, float? obf_coefficient=1.0) -> Tensor": { | ||
| 2907 | + "version": ["all_version"] | ||
| 2908 | + }, | ||
| 2909 | + "func: obfuscation_finalize(Tensor fd_to_close) -> Tensor": { | ||
| 2910 | + "version": ["all_version"] | ||
| 2911 | + }, | ||
| 2912 | + "func: obfuscation_calculate(Tensor fd, Tensor x, Tensor param, *, float? obf_coefficient=1.0) -> Tensor": { | ||
| 2913 | + "version": ["all_version"] | ||
| 2914 | + }, | ||
| 2906 | "func: one_hot(Tensor self, int num_classes=-1) -> Tensor": { | 2915 | "func: one_hot(Tensor self, int num_classes=-1) -> Tensor": { |
| 2907 | "version": ["all_version"] | 2916 | "version": ["all_version"] |
| 2908 | }, | 2917 | }, |
| @@ -0,0 +1,25 @@ | |||
| 1 | +import unittest | ||
| 2 | +from unittest.mock import patch | ||
| 3 | + | ||
| 4 | +import torch | ||
| 5 | + | ||
| 6 | +import torch_npu | ||
| 7 | +from torch_npu.testing.testcase import TestCase, run_tests | ||
| 8 | +from torch_npu.testing.common_utils import SupportedDevices | ||
| 9 | + | ||
| 10 | + | ||
| 11 | +class TestObfuscationCalculate(TestCase): | ||
| 12 | + | ||
| 13 | + def test_obfuscation_calculate(self, device="npu"): | ||
| 14 | + x = torch.randn(1024, 3584, device=device, dtype=torch.bfloat16) | ||
| 15 | + fd = torch.tensor([35], device=device, dtype=torch.int32) | ||
| 16 | + param = torch.tensor([3584], device=device, dtype=torch.int32) | ||
| 17 | + obf_cft = 1.0 | ||
| 18 | + mock_output = torch.randn(1024, 3584, device=device, dtype=torch.bfloat16) | ||
| 19 | + with patch('torch_npu.obfuscation_calculate', return_value=mock_output): | ||
| 20 | + result = torch_npu.obfuscation_calculate(fd, x, param, obf_coefficient=obf_cft) | ||
| 21 | + self.assertEqual(result, mock_output) | ||
| 22 | + | ||
| 23 | + | ||
| 24 | +if __name__ == "__main__": | ||
| 25 | + run_tests() | ||
| @@ -0,0 +1,22 @@ | |||
| 1 | +import unittest | ||
| 2 | +from unittest.mock import patch | ||
| 3 | + | ||
| 4 | +import torch | ||
| 5 | + | ||
| 6 | +import torch_npu | ||
| 7 | +from torch_npu.testing.testcase import TestCase, run_tests | ||
| 8 | +from torch_npu.testing.common_utils import SupportedDevices | ||
| 9 | + | ||
| 10 | + | ||
| 11 | +class TestObfuscationFinalize(TestCase): | ||
| 12 | + | ||
| 13 | + def test_obfuscation_finalize(self, device="npu"): | ||
| 14 | + fd_to_close = torch.tensor([35], device=device, dtype=torch.int32) | ||
| 15 | + mock_output = torch.tensor([35], device=device, dtype=torch.int32) | ||
| 16 | + with patch('torch_npu.obfuscation_finalize', return_value=mock_output): | ||
| 17 | + result = torch_npu.obfuscation_finalize(fd_to_close) | ||
| 18 | + self.assertEqual(result, mock_output) | ||
| 19 | + | ||
| 20 | + | ||
| 21 | +if __name__ == "__main__": | ||
| 22 | + run_tests() | ||
| @@ -0,0 +1,27 @@ | |||
| 1 | +import unittest | ||
| 2 | +from unittest.mock import patch | ||
| 3 | + | ||
| 4 | +import torch | ||
| 5 | + | ||
| 6 | +import torch_npu | ||
| 7 | +from torch_npu.testing.testcase import TestCase, run_tests | ||
| 8 | +from torch_npu.testing.common_utils import SupportedDevices | ||
| 9 | + | ||
| 10 | + | ||
| 11 | +class TestObfuscationInitialize(TestCase): | ||
| 12 | + | ||
| 13 | + def test_obfuscation_initialize(self, device="npu"): | ||
| 14 | + hidden_size = 3584 | ||
| 15 | + tp_rank = 4 | ||
| 16 | + cmd = 1 | ||
| 17 | + data_type = torch.bfloat16 | ||
| 18 | + thread_num = 4 | ||
| 19 | + obf_cft = 1.0 | ||
| 20 | + mock_output = torch.tensor([35], device=device, dtype=torch.int32) | ||
| 21 | + with patch('torch_npu.obfuscation_initialize', return_value=mock_output): | ||
| 22 | + result = torch_npu.obfuscation_initialize(hidden_size, tp_rank, cmd, data_type=data_type, thread_num=thread_num, obf_coefficient=obf_cft) | ||
| 23 | + self.assertEqual(result, mock_output) | ||
| 24 | + | ||
| 25 | + | ||
| 26 | +if __name__ == "__main__": | ||
| 27 | + run_tests() | ||