已合并
add npu_batch_gather_matmul inplace #2274
JialiangShen1创建于 2025年3月14日
add npu_batch_gather_matmul inplace #2274
已合并
从refs/pull/2274/head合入到master
共 5 个文件变更+90-9
| @@ -217,6 +217,7 @@ v (Tensor) - 变量方差。 | |||
| 217 | """ | 217 | """ |
| 218 | ) | 218 | ) |
| 219 | 219 | ||
| 220 | + | ||
| 220 | _add_torch_npu_docstr( | 221 | _add_torch_npu_docstr( |
| 221 | "npu_batch_gather_matmul", | 222 | "npu_batch_gather_matmul", |
| 222 | """ | 223 | """ |
| @@ -246,6 +247,40 @@ y_slice_size (Int, 默认值为-1) - 可选值，表示y更新时的范围，数 | |||
| 246 | """ | 247 | """ |
| 247 | ) | 248 | ) |
| 248 | 249 | ||
| 250 | + | ||
| 251 | +_add_torch_npu_docstr( | ||
| 252 | + "npu_batch_gather_matmul_", | ||
| 253 | + """ | ||
| 254 | +torch_npu.npu_batch_gather_matmul_(y, x, weight_b, indices, weight_a=None, | ||
| 255 | + layer_idx=0, scale=1e-3, y_offset=0, y_slice_size=-1) -> Tensor(a!) | ||
| 256 | + | ||
| 257 | +功能描述 | ||
| 258 | +npu_batch_gather_matmul的inplace版本. 将输入x根据输入索引indices, 分别和对应的weight_a, weight_b 相乘， 然后将结果累加到输入y并输出。 | ||
| 259 | + | ||
| 260 | +参数说明 | ||
| 261 | +y (Tensor) - 必填值，输入tensor，表示待进行累加更新的张量，数据类型Float16，输入示例：[batch_size, y_column]。 | ||
| 262 | +x (Tensor) - 必填值，输入tensor，表示分组前的输入张量，数据类型Float16，输入示例：[batch_size, H1]。 | ||
| 263 | +weight_b (Tensor) - 必填值，输入tensor，表示进行矩阵乘的第二个权重矩阵，数据类型Float16。输入示例：[W, L, H2, R]。 | ||
| 264 | +indices (Tensor) - 必填值，标识输入x的分组索引，数据类型Int32。输入示例：[batch_size]。 | ||
| 265 | +weight_a (Tensor) - 可选值，输入tensor，表示进行矩阵乘的第一个权重矩阵，数据类型Float16。为空时会跳过第一个矩阵乘， 输入示范：[W, L, R, H1]。 | ||
| 266 | +layer_idx (Int, 默认值为0) - 可选值，表示weight的层数索引，数据类型Int。 | ||
| 267 | +scale (Float, 默认值为1e-3) - 可选值，表示matmul结果的缩放系数，数据类型Float。 | ||
| 268 | +y_offset (Int, 默认值为0) - 可选值，表示y更新的偏移值，数据类型Int。 | ||
| 269 | +y_slice_size (Int, 默认值为-1) - 可选值，表示y更新时的范围，数据类型Int。当为-1的时候，会按照y_column的值传入；当非-1 时，以传入的值做更新范围。 | ||
| 270 | + | ||
| 271 | +输出说明 | ||
| 272 | +out：Device侧的Tensor类型，计算输出，复用y输入地址；数据类型和shape与self一致。 | ||
| 273 | +示例 | ||
| 274 | +y = torch.randn(1, 128).half().npu() | ||
| 275 | +x = torch.randn(1, 16).half().npu() | ||
| 276 | +weightA = torch.randn(2, 1, 16, 16).half().npu() | ||
| 277 | +indices = torch.randint(0, 1, (1,)).to(torch.int32).npu() | ||
| 278 | +weightB = torch.randn(2, 1, 128, 16).half().npu() | ||
| 279 | +out = torch_npu.npu_batch_gather_matmul_(y, x, weightB, indices, weightA, y_offset=0, y_slice_size=128, layer_idx=0, scale=2) | ||
| 280 | +out | ||
| 281 | +""" | ||
| 282 | +) | ||
| 283 | + | ||
| 249 | _add_torch_npu_docstr( | 284 | _add_torch_npu_docstr( |
| 250 | "npu_batch_nms", | 285 | "npu_batch_nms", |
| 251 | """ | 286 | """ |
| @@ -6511,7 +6511,10 @@ custom: | |||
| 6511 | - func: npu_apply_rotary_pos_emb(Tensor query, Tensor key, Tensor cos, Tensor sin, str layout='BSH') -> (Tensor, Tensor) | 6511 | - func: npu_apply_rotary_pos_emb(Tensor query, Tensor key, Tensor cos, Tensor sin, str layout='BSH') -> (Tensor, Tensor) |
| 6512 | op_api: all_version | 6512 | op_api: all_version |
| 6513 | 6513 | ||
| 6514 | - - func: npu_batch_gather_matmul(Tensor y, Tensor x, Tensor weight_b, Tensor indices, Tensor? weight_a=None, int layer_idx=0, float scale=1e-3, int y_offset=0, int y_slice_size=-1) -> Tensor | 6514 | + - func: npu_batch_gather_matmul(Tensor self, Tensor x, Tensor weight_b, Tensor indices, Tensor? weight_a=None, int layer_idx=0, float scale=1e-3, int y_offset=0, int y_slice_size=-1) -> Tensor |
| 6515 | + op_api: all_version | ||
| 6516 | + | ||
| 6517 | + - func: npu_batch_gather_matmul_(Tensor(a!) self, Tensor x, Tensor weight_b, Tensor indices, Tensor? weight_a=None, int layer_idx=0, float scale=1e-3, int y_offset=0, int y_slice_size=-1) -> Tensor(a!) | ||
| 6515 | op_api: all_version | 6518 | op_api: all_version |
| 6516 | 6519 | ||
| 6517 | - func: npu_batch_nms(Tensor self, Tensor scores, float score_threshold, float iou_threshold, int max_size_per_class, int max_total_size, bool change_coordinate_frame=False, bool transpose_box=False) -> (Tensor, Tensor, Tensor, Tensor) | 6520 | - func: npu_batch_nms(Tensor self, Tensor scores, float score_threshold, float iou_threshold, int max_size_per_class, int max_total_size, bool change_coordinate_frame=False, bool transpose_box=False) -> (Tensor, Tensor, Tensor, Tensor) |
Rop_plugin/ops/opapi/AddLoraKernelNpuOpApi.cpp→op_plugin/ops/opapi/BatchGatherMatmulKernelNpuOpApi.cpp+27-6
| @@ -35,7 +35,7 @@ void Infer_shape_check(const at::Tensor &y, const at::Tensor &x, const at::Tenso | |||
| 35 | } | 35 | } |
| 36 | 36 | ||
| 37 | at::Tensor npu_batch_gather_matmul( | 37 | at::Tensor npu_batch_gather_matmul( |
| 38 | - const at::Tensor& y, | 38 | + const at::Tensor& self, |
| 39 | const at::Tensor& x, | 39 | const at::Tensor& x, |
| 40 | const at::Tensor& weight_b, | 40 | const at::Tensor& weight_b, |
| 41 | const at::Tensor& indices, | 41 | const at::Tensor& indices, |
| @@ -45,14 +45,35 @@ at::Tensor npu_batch_gather_matmul( | |||
| 45 | int64_t y_offset, | 45 | int64_t y_offset, |
| 46 | int64_t y_slice_size) | 46 | int64_t y_slice_size) |
| 47 | { | 47 | { |
| 48 | - Infer_shape_check(y, x, weight_b, indices, weight_a); | 48 | + Infer_shape_check(self, x, weight_b, indices, weight_a); |
| 49 | 49 | ||
| 50 | if (y_slice_size == -1) { | 50 | if (y_slice_size == -1) { |
| 51 | - y_slice_size = y.size(1); | 51 | + y_slice_size = self.size(1); |
| 52 | } | 52 | } |
| 53 | - at::Tensor result = npu_preparation::apply_tensor_without_format(y); | 53 | + at::Tensor result = npu_preparation::apply_tensor_without_format(self); |
| 54 | 54 | ||
| 55 | - EXEC_NPU_CMD(aclnnAddLora, y, x, weight_b, indices, weight_a, layer_idx, scale, y_offset, y_slice_size, result); | 55 | + EXEC_NPU_CMD(aclnnAddLora, self, x, weight_b, indices, weight_a, layer_idx, scale, y_offset, y_slice_size, result); |
| 56 | - return y; | 56 | + return self; |
| 57 | +} | ||
| 58 | + | ||
| 59 | +at::Tensor &npu_batch_gather_matmul_( | ||
| 60 | + at::Tensor& self, | ||
| 61 | + const at::Tensor& x, | ||
| 62 | + const at::Tensor& weight_b, | ||
| 63 | + const at::Tensor& indices, | ||
| 64 | + const c10::optional<at::Tensor> &weight_a, | ||
| 65 | + int64_t layer_idx, | ||
| 66 | + double scale, | ||
| 67 | + int64_t y_offset, | ||
| 68 | + int64_t y_slice_size) | ||
| 69 | +{ | ||
| 70 | + Infer_shape_check(self, x, weight_b, indices, weight_a); | ||
| 71 | + | ||
| 72 | + if (y_slice_size == -1) { | ||
| 73 | + y_slice_size = self.size(1); | ||
| 74 | + } | ||
| 75 | + | ||
| 76 | + EXEC_NPU_CMD(aclnnAddLora, self, x, weight_b, indices, weight_a, layer_idx, scale, y_offset, y_slice_size, self); | ||
| 77 | + return self; | ||
| 57 | } | 78 | } |
| 58 | } | 79 | } |
| @@ -1236,7 +1236,12 @@ def npu_dequant_bias_meta(x, weight_scale, activation_scale, bias, output_dtype= | |||
| 1236 | 1236 | ||
| 1237 | 1237 | ||
| 1238 | 1238 | ||
| 1239 | -def npu_batch_gather_matmul_meta(y, x, weight_b, indices, weight_a=None, | 1239 | +def npu_batch_gather_matmul_meta(self, x, weight_b, indices, weight_a=None, |
| 1240 | layer_idx=0, scale=1e-3, y_offset=0, y_slice_size=-1): | 1240 | layer_idx=0, scale=1e-3, y_offset=0, y_slice_size=-1): |
| 1241 | - return torch.empty_like(y, dtype=y.dtype) | 1241 | + return torch.empty_like(self, dtype=self.dtype) |
| 1242 | 1242 | ||
| 1243 | + | ||
| 1244 | + | ||
| 1245 | +def npu_batch_gather_matmul__meta(self, x, weight_b, indices, weight_a=None, | ||
| 1246 | + layer_idx=0, scale=1e-3, y_offset=0, y_slice_size=-1): | ||
| 1247 | + return self | ||
| @@ -37,6 +37,23 @@ class TestBatchGatherMatmul(TestCase): | |||
| 37 | 37 | ||
| 38 | self.assertEqual(y.cpu(), output_npu) | 38 | self.assertEqual(y.cpu(), output_npu) |
| 39 | 39 | ||
| 40 | + | ||
| 41 | + | ||
| 42 | + def test_batch_gather_matmul_(self): | ||
| 43 | + torch.manual_seed(12) | ||
| 44 | + y = torch.randn(10, 128).half().npu() | ||
| 45 | + x = torch.randn(10, 128).half().npu() | ||
| 46 | + weightA = torch.randn(2, 1, 16, 128).half().npu() | ||
| 47 | + indices = torch.randint(0, 2, (10,)).to(torch.int32).npu() | ||
| 48 | + weightB = torch.randn(2, 1, 128, 16).half().npu() | ||
| 49 | + | ||
| 50 | + output_npu = self.npu_batch_gather_matmul(y, x, weightB, indices, weightA, y_offset=0, y_slice_size=128, | ||
| 51 | + layer_idx=0, scale=2) | ||
| 52 | + y_out = torch_npu.npu_batch_gather_matmul_(y, x, weightB, indices, weightA, | ||
| 53 | + y_offset=0, y_slice_size=128, layer_idx=0, scale=2) | ||
| 54 | + | ||
| 55 | + self.assertEqual(y_out.cpu(), output_npu) | ||
| 56 | + | ||
| 40 | 57 | ||
| 41 | if __name__ == "__main__": | 58 | if __name__ == "__main__": |
| 42 | run_tests() | 59 | run_tests() |