已合并
【feature】新增适配QuantMatmulGelu融合算子，支持A4W4/A8W8输入 #4186
Sunshine_Youngster创建于 2月3日
【feature】新增适配QuantMatmulGelu融合算子，支持A4W4/A8W8输入 #4186
已合并
共 11 个文件变更+1286-1
| @@ -4476,6 +4476,168 @@ else: | |||
| 4476 | """ | 4476 | """ |
| 4477 | ) | 4477 | ) |
| 4478 | 4478 | ||
| 4479 | +_add_torch_npu_docstr( | ||
| 4480 | + "npu_quant_matmul_gelu", | ||
| 4481 | + """ | ||
| 4482 | +功能描述: | ||
| 4483 | +完成量化矩阵乘和GELU激活函数的融合计算, 支持A8W8和A4W4量化. 该接口融合了量化矩阵乘和GELU激活, 减少内存访问, 提升性能. | ||
| 4484 | + | ||
| 4485 | +接口原型: | ||
| 4486 | +torch_npu.npu_quant_matmul_gelu(Tensor x1, Tensor x2, Tensor x1_scale, Tensor x2_scale, *, Tensor? bias=None, str? approximate="gelu_erf") -> Tensor | ||
| 4487 | + | ||
| 4488 | +参数说明: | ||
| 4489 | +x1: Tensor类型, 输入激活值, 数据格式支持ND, shape需要在2-6维范围. | ||
| 4490 | +数据类型支持int8(A8W8量化)、int32(A4W4量化, 每个int32数据存放8个int4数据)、quint4x2(A4W4量化, 直接INT4类型). | ||
| 4491 | +x2: Tensor类型(权重), 数据格式支持ND或NZ, shape需要在2-6维范围. | ||
| 4492 | +数据类型与x1的数据类型保持一致. 支持昇腾亲和的NZ数据排布格式, 可通过torch_npu.npu_format_cast转换为NZ格式以提升性能(仅A8W8场景). | ||
| 4493 | +x1_scale: Tensor类型, x1的量化scale参数, 数据格式支持ND. | ||
| 4494 | +数据类型支持float32. shape需要是1维(m,), 其中m与x1的m一致. 采用per-token量化方式, 每个token(行)有一个独立的scale值. | ||
| 4495 | +x2_scale: Tensor类型, x2的量化scale参数, 数据格式支持ND. | ||
| 4496 | +数据类型支持float32或bfloat16. shape需要是1维(n,)或(1,), 其中n与x2的n一致. 采用per-channel量化方式, 每个输出通道有一个独立的scale值, 或使用per-tensor量化(shape为(1,)). | ||
| 4497 | +bias: Tensor类型, 可选参数, 默认值为None, 偏置项, 数据格式支持ND. | ||
| 4498 | +数据类型支持int32、float32、bfloat16、float16. | ||
| 4499 | +A4W4量化场景下: shape仅支持1维(n,), n与x2的n一致. | ||
| 4500 | +A8W8量化场景下: shape支持1维(n,)或3维(batch, 1, n), n与x2的n一致. | ||
| 4501 | +approximate: str类型, 可选参数, 默认值为"gelu_erf". 指定GELU激活函数的类型. | ||
| 4502 | +支持"gelu_tanh"(GELU的tanh近似版本)和"gelu_erf"(GELU的erf精确版本). | ||
| 4503 | + | ||
| 4504 | +输出说明: | ||
| 4505 | +result: Tensor类型, 代表量化矩阵乘融合GELU激活的计算结果. | ||
| 4506 | +如果x2_scale的数据类型为float32, 输出的数据类型为float16. | ||
| 4507 | +如果x2_scale的数据类型为bfloat16, 输出的数据类型为bfloat16. | ||
| 4508 | +输出shape为(batch, m, n), 其中batch根据x1和x2的batch维度广播得到. | ||
| 4509 | + | ||
| 4510 | +约束说明: | ||
| 4511 | +该接口支持推理场景下使用. | ||
| 4512 | +x1、x2、x1_scale、x2_scale不能为空. | ||
| 4513 | +x1、x2的数据类型和数据格式需要在支持的范围之内. | ||
| 4514 | +x1、x2最后一维的shape大小不能超过65535. | ||
| 4515 | +approximate必须为"gelu_tanh"或"gelu_erf". | ||
| 4516 | +对于A4W4量化(INT4/INT32类型输入): | ||
| 4517 | + A4W4量化场景支持两种输入类型: quint4x2(直接INT4类型)和int32(打包存储, 每个int32数据存放8个int4数据). | ||
| 4518 | + x1和x2的内轴(k轴)必须为偶数. | ||
| 4519 | + 当x2为int32类型时, x2的shape为(k, n//8), n必须是8的倍数. | ||
| 4520 | + 当x2为quint4x2类型时, x2的shape为(k, n), n必须是8的倍数. | ||
| 4521 | + A4W4量化仅支持ND格式, 不支持NZ格式. | ||
| 4522 | + 转置信息由算子内部根据tensor的stride自动推导, 无需手动指定. | ||
| 4523 | +对于A8W8量化: | ||
| 4524 | + 支持ND格式和NZ格式. | ||
| 4525 | + 当x2为ND格式时, 自动调用aclnnFusedQuantMatmul接口. | ||
| 4526 | + 当x2为NZ格式时, 自动调用aclnnFusedQuantMatmulWeightNz接口. | ||
| 4527 | + 如果需要使用NZ格式以提升性能, 可以手动调用torch_npu.npu_format_cast完成输入x2(weight)的NZ格式转换. | ||
| 4528 | + 转置信息由算子内部根据tensor的stride自动推导, 无需手动指定. | ||
| 4529 | + | ||
| 4530 | +支持的PyTorch版本 | ||
| 4531 | +PyTorch 2.10 | ||
| 4532 | +PyTorch 2.9 | ||
| 4533 | +PyTorch 2.8 | ||
| 4534 | +PyTorch 2.7 | ||
| 4535 | +PyTorch 2.6 | ||
| 4536 | + | ||
| 4537 | +支持的型号: | ||
| 4538 | +Atlas A2 训练系列产品/Atlas 800I A2 推理产品 | ||
| 4539 | + | ||
| 4540 | +调用示例: | ||
| 4541 | +单算子调用(A8W8量化, ND格式, gelu_tanh激活) | ||
| 4542 | +import torch | ||
| 4543 | +import torch_npu | ||
| 4544 | + | ||
| 4545 | +# 准备输入数据 | ||
| 4546 | +m, k, n = 128, 256, 512 | ||
| 4547 | +x1 = torch.randint(-5, 5, (m, k), dtype=torch.int8).npu() | ||
| 4548 | +x2 = torch.randint(-5, 5, (k, n), dtype=torch.int8).npu() | ||
| 4549 | +x1_scale = torch.randn(m, dtype=torch.float32).abs().npu() * 0.01 | ||
| 4550 | +x2_scale = torch.randn(n, dtype=torch.float32).abs().npu() * 0.01 | ||
| 4551 | + | ||
| 4552 | +# 调用融合算子 | ||
| 4553 | +output = torch_npu.npu_quant_matmul_gelu(x1, x2, x1_scale, x2_scale, approximate="gelu_tanh") | ||
| 4554 | +print(output.shape) # torch.Size([128, 512]) | ||
| 4555 | +print(output.dtype) # torch.float16 | ||
| 4556 | + | ||
| 4557 | +单算子调用(A8W8量化, ND格式, gelu_erf激活, 带bias) | ||
| 4558 | +import torch | ||
| 4559 | +import torch_npu | ||
| 4560 | + | ||
| 4561 | +m, k, n = 128, 256, 512 | ||
| 4562 | +x1 = torch.randint(-5, 5, (m, k), dtype=torch.int8).npu() | ||
| 4563 | +x2 = torch.randint(-5, 5, (k, n), dtype=torch.int8).npu() | ||
| 4564 | +x1_scale = torch.randn(m, dtype=torch.float32).abs().npu() * 0.01 | ||
| 4565 | +x2_scale = torch.randn(n, dtype=torch.float32).abs().npu() * 0.01 | ||
| 4566 | +bias = torch.randn(n, dtype=torch.float32).npu() * 0.1 | ||
| 4567 | + | ||
| 4568 | +# 使用gelu_erf激活并添加bias | ||
| 4569 | +output = torch_npu.npu_quant_matmul_gelu( | ||
| 4570 | + x1, x2, x1_scale, x2_scale, bias=bias, approximate="gelu_erf" | ||
| 4571 | +) | ||
| 4572 | + | ||
| 4573 | +单算子调用(A8W8量化, NZ格式, gelu_tanh激活) | ||
| 4574 | +import torch | ||
| 4575 | +import torch_npu | ||
| 4576 | + | ||
| 4577 | +m, k, n = 128, 256, 512 | ||
| 4578 | +x1 = torch.randint(-5, 5, (m, k), dtype=torch.int8).npu() | ||
| 4579 | +x2 = torch.randint(-5, 5, (k, n), dtype=torch.int8).npu() | ||
| 4580 | + | ||
| 4581 | +# 将x2转换为NZ格式以提升性能 | ||
| 4582 | +x2_nz = torch_npu.npu_format_cast(x2.contiguous(), 29) # 29为ACL_FORMAT_FRACTAL_NZ | ||
| 4583 | + | ||
| 4584 | +x1_scale = torch.randn(m, dtype=torch.float32).abs().npu() * 0.01 | ||
| 4585 | +x2_scale = torch.randn(n, dtype=torch.float32).abs().npu() * 0.01 | ||
| 4586 | + | ||
| 4587 | +# 自动识别NZ格式并调用对应接口 | ||
| 4588 | +output = torch_npu.npu_quant_matmul_gelu(x1, x2_nz, x1_scale, x2_scale, approximate="gelu_tanh") | ||
| 4589 | + | ||
| 4590 | +单算子调用(A8W8量化, BF16输出) | ||
| 4591 | +import torch | ||
| 4592 | +import torch_npu | ||
| 4593 | + | ||
| 4594 | +m, k, n = 64, 128, 256 | ||
| 4595 | +x1 = torch.randint(-5, 5, (m, k), dtype=torch.int8).npu() | ||
| 4596 | +x2 = torch.randint(-5, 5, (k, n), dtype=torch.int8).npu() | ||
| 4597 | +x1_scale = torch.randn(m, dtype=torch.float32).abs().npu() * 0.01 | ||
| 4598 | +x2_scale = torch.randn(n, dtype=torch.bfloat16).abs().npu() * 0.01 # BF16 scale | ||
| 4599 | + | ||
| 4600 | +# 输出数据类型由x2_scale的类型决定, 此处输出为bfloat16 | ||
| 4601 | +output = torch_npu.npu_quant_matmul_gelu(x1, x2, x1_scale, x2_scale, approximate="gelu_tanh") | ||
| 4602 | +print(output.dtype) # torch.bfloat16 | ||
| 4603 | + | ||
| 4604 | +单算子调用(A4W4量化) | ||
| 4605 | +import torch | ||
| 4606 | +import torch_npu | ||
| 4607 | + | ||
| 4608 | +m, k, n = 128, 256, 512 | ||
| 4609 | +# 生成INT4数据(以INT32格式存储) | ||
| 4610 | +x1_fp = torch.randn(m, k, dtype=torch.float32).npu() | ||
| 4611 | +x2_fp = torch.randn(k, n, dtype=torch.float32).npu() | ||
| 4612 | + | ||
| 4613 | +# 量化为INT4 | ||
| 4614 | +scale_tmp = torch.ones(1, dtype=torch.float32).npu() | ||
| 4615 | +x1 = torch_npu.npu_quantize(x1_fp, scale_tmp, None, torch.quint4x2, -1, False) | ||
| 4616 | +x2 = torch_npu.npu_quantize(x2_fp, scale_tmp, None, torch.quint4x2, -1, False) | ||
| 4617 | + | ||
| 4618 | +x1_scale = torch.randn(m, dtype=torch.float32).abs().npu() * 0.01 | ||
| 4619 | +x2_scale = torch.randn(n, dtype=torch.float32).abs().npu() * 0.01 | ||
| 4620 | + | ||
| 4621 | +# A4W4量化仅支持ND格式, 不支持NZ格式 | ||
| 4622 | +output = torch_npu.npu_quant_matmul_gelu(x1, x2, x1_scale, x2_scale, approximate="gelu_tanh") | ||
| 4623 | + | ||
| 4624 | +单算子调用(使用默认approximate="gelu_erf") | ||
| 4625 | +import torch | ||
| 4626 | +import torch_npu | ||
| 4627 | + | ||
| 4628 | +m, k, n = 64, 128, 256 | ||
| 4629 | +x1 = torch.randint(-5, 5, (m, k), dtype=torch.int8).npu() | ||
| 4630 | +x2 = torch.randint(-5, 5, (k, n), dtype=torch.int8).npu() | ||
| 4631 | +x1_scale = torch.randn(m, dtype=torch.float32).abs().npu() * 0.01 | ||
| 4632 | +x2_scale = torch.randn(n, dtype=torch.float32).abs().npu() * 0.01 | ||
| 4633 | + | ||
| 4634 | +# 不指定approximate参数, 使用默认值"gelu_erf" | ||
| 4635 | +output = torch_npu.npu_quant_matmul_gelu(x1, x2, x1_scale, x2_scale) | ||
| 4636 | +print(output.dtype) # torch.float16 | ||
| 4637 | + | ||
| 4638 | +""" | ||
| 4639 | +) | ||
| 4640 | + | ||
| 4479 | _add_torch_npu_docstr( | 4641 | _add_torch_npu_docstr( |
| 4480 | "npu_weight_quant_batchmatmul", | 4642 | "npu_weight_quant_batchmatmul", |
| 4481 | """ | 4643 | """ |
| @@ -0,0 +1,293 @@ | |||
| 1 | +# torch_npu.npu_quant_matmul_gelu | ||
| 2 | + | ||
| 3 | +## 产品支持情况 | ||
| 4 | + | ||
| 5 | +| 产品 | 是否支持 | | ||
| 6 | +|-----------------------------------------------|:----:| | ||
| 7 | +| <term>Atlas A3 训练系列产品/Atlas A3 推理系列产品</term> | √ | | ||
| 8 | +| <term>Atlas A2 训练系列产品/Atlas A2 推理系列产品</term> | √ | | ||
| 9 | + | ||
| 10 | +## 功能说明 | ||
| 11 | + | ||
| 12 | +- API功能：完成量化的矩阵乘和GELU激活函数的融合计算，支持A8W8和A4W4量化。该接口融合了量化矩阵乘和GELU激活，减少内存访问，提升性能。 | ||
| 13 | + | ||
| 14 | +- 计算公式： | ||
| 15 | + | ||
| 16 | + - x1Scale， bias int32（此场景无offset）： | ||
| 17 | + | ||
| 18 | + $$ | ||
| 19 | + qbmmout = (x1@x2 + bias) * x2Scale * x1Scale | ||
| 20 | + $$ | ||
| 21 | + | ||
| 22 | + - x1Scale， bias BFLOAT16/FLOAT16/FLOAT32（此场景无offset）： | ||
| 23 | + | ||
| 24 | + $$ | ||
| 25 | + qbmmout = x1@x2 * x2scale * x1Scale + bias | ||
| 26 | + $$ | ||
| 27 | + | ||
| 28 | + - x1Scale无bias： | ||
| 29 | + | ||
| 30 | + $$ | ||
| 31 | + qbmmout = x1@x2 * x2Scale * x1Scale | ||
| 32 | + $$ | ||
| 33 | + | ||
| 34 | + - gelu类型由approximate输入指定，默认为"gelu_erf"，支持如下： | ||
| 35 | + | ||
| 36 | + - gelu_tanh运算： | ||
| 37 | + | ||
| 38 | + $$ | ||
| 39 | + out = gelu\_tanh(qbmmout) | ||
| 40 | + $$ | ||
| 41 | + | ||
| 42 | + - gelu_erf运算： | ||
| 43 | + | ||
| 44 | + $$ | ||
| 45 | + out = gelu\_erf(qbmmout) | ||
| 46 | + $$ | ||
| 47 | + | ||
| 48 | +## 函数原型 | ||
| 49 | + | ||
| 50 | +``` | ||
| 51 | +npu_quant_matmul_gelu(x1, x2, x1_scale, x2_scale, *, bias=None, approximate="gelu_erf") -> Tensor | ||
| 52 | +``` | ||
| 53 | + | ||
| 54 | +## 参数说明 | ||
| 55 | + | ||
| 56 | +- **x1** (`Tensor`)：必选参数，输入张量，表示矩阵乘法中的左矩阵（激活值），数据格式支持$ND$，shape需要在2-6维范围。 | ||
| 57 | + - 数据类型支持`int8`（A8W8量化）、`int32`（A4W4量化，每个`int32`数据存放8个`int4`数据）和`int4`（A4W4量化，直接int4类型）。 | ||
| 58 | + | ||
| 59 | +- **x2** (`Tensor`)：必选参数，输入张量，表示矩阵乘法中的右矩阵（权重），其与`x1`的数据类型须保持一致。数据格式支持$ND$或$NZ$（昇腾亲和排布格式），shape需要在2-6维范围。 | ||
| 60 | + - 数据类型支持`int8`（A8W8量化）、`int32`（A4W4量化，每个`int32`数据存放8个`int4`数据）和`int4`（A4W4量化，直接int4类型）。 | ||
| 61 | + - 支持昇腾亲和的NZ数据排布格式，可通过`torch_npu.npu_format_cast`转换为NZ格式以提升性能（仅A8W8场景）。 | ||
| 62 | + | ||
| 63 | +- **x1_scale** (`Tensor`)：必选参数，`x1`的量化缩放因子，数据格式支持$ND$。 | ||
| 64 | + - 数据类型支持`float32`。 | ||
| 65 | + - shape需要是1维$(m,)$，其中$m$与`x1`的$m$一致。采用per-token量化方式，每个token（行）有一个独立的scale值。 | ||
| 66 | + | ||
| 67 | +- **x2_scale** (`Tensor`)：必选参数，`x2`的量化缩放因子，数据格式支持$ND$。 | ||
| 68 | + - 数据类型支持`float32`、`bfloat16`。 | ||
| 69 | + - shape需要是1维$(n,)$或$(1,)$，其中$n$与`x2`的$n$一致。采用per-channel量化方式，每个输出通道有一个独立的scale值，或使用per-tensor量化（shape为$(1,)$）。 | ||
| 70 | + | ||
| 71 | +- <strong>*</strong>：必选参数，代表其之前的变量是位置相关的，必须按照顺序输入；之后的变量是可选参数，位置无关，需要使用键值对赋值，不赋值会使用默认值。 | ||
| 72 | + | ||
| 73 | +- **bias** (`Tensor`)：可选参数，默认值为`None`，偏置项，数据格式支持$ND$。 | ||
| 74 | + - 数据类型支持`int32`、`float32`、`bfloat16`、`float16`。 | ||
| 75 | + - A4W4量化场景下：shape仅支持1维$(n,)$，$n$与`x2`的$n$一致。 | ||
| 76 | + - A8W8量化场景下：shape支持1维$(n,)$或3维$(batch, 1, n)$，$n$与`x2`的$n$一致，同时$batch$值需要等于`x1`和`x2` broadcast后推导出的$batch$值。 | ||
| 77 | + | ||
| 78 | +- **approximate** (`str`)：可选参数，默认值为`"gelu_erf"`。指定GELU激活函数的类型。 | ||
| 79 | + - 支持`"gelu_tanh"`（GELU的tanh近似版本）和`"gelu_erf"`（GELU的erf精确版本）。 | ||
| 80 | + | ||
| 81 | +## 返回值说明 | ||
| 82 | + | ||
| 83 | +`Tensor` | ||
| 84 | + | ||
| 85 | +代表量化矩阵乘融合GELU激活的计算结果。 | ||
| 86 | +- 输出数据类型的确定规则： | ||
| 87 | + - 如果`x2_scale`的数据类型为`float32`，输出的数据类型为`float16`。 | ||
| 88 | + - 如果`x2_scale`的数据类型为`bfloat16`，输出的数据类型为`bfloat16`。 | ||
| 89 | + - 如果`bias`的数据类型为`bfloat16`，输出的数据类型强制为`bfloat16`（优先级高于`x2_scale`）。 | ||
| 90 | +- 输出shape为$(batch, m, n)$，其中$batch$根据`x1`和`x2`的batch通过broadcast得到。 | ||
| 91 | + | ||
| 92 | +## 约束说明 | ||
| 93 | + | ||
| 94 | +- 该接口支持推理场景下使用。 | ||
| 95 | +- 该接口支持图模式。 | ||
| 96 | +- 传入的`x1`、`x2`、`x1_scale`、`x2_scale`不能是空。 | ||
| 97 | +- `x1`、`x2`、`bias`、`x1_scale`、`x2_scale`的数据类型和数据格式需要在支持的范围之内。 | ||
| 98 | +- `x1`与`x2`最后一维的shape大小不能超过65535。 | ||
| 99 | +- `approximate`必须为`"gelu_tanh"`或`"gelu_erf"`。 | ||
| 100 | + | ||
| 101 | +- **A4W4量化（int4/int32类型输入）的额外约束**： | ||
| 102 | + | ||
| 103 | + A4W4量化场景支持两种输入类型： | ||
| 104 | + - **int4类型**：直接使用int4数据类型 | ||
| 105 | + - **int32类型**：每个`int32`数据存放8个`int4`数据（打包存储） | ||
| 106 | + | ||
| 107 | + 当使用`int32`类型时，输入的`int32` shape需要将数据原本`int4`类型时shape的最后一维缩小8倍。`int4`数据的shape最后一维应为8的倍数。 | ||
| 108 | + | ||
| 109 | + - `x1`和`x2`的内轴（k轴）必须为偶数。 | ||
| 110 | + - 当`x2`为`int32`类型时，`x2`的shape为$(k, n//8)$，$n$必须是8的倍数。 | ||
| 111 | + - 当`x2`为`int4`类型时，`x2`的shape为$(k, n)$，$n$必须是8的倍数。 | ||
| 112 | + - A4W4量化仅支持ND格式，不支持NZ格式。 | ||
| 113 | + - 转置信息由算子内部根据tensor的stride自动推导，无需手动指定。 | ||
| 114 | + | ||
| 115 | +- **A8W8量化的约束**： | ||
| 116 | + | ||
| 117 | + - 支持ND格式和NZ格式。 | ||
| 118 | + - 如果需要使用NZ格式以提升性能，可以手动调用`torch_npu.npu_format_cast`完成输入`x2`（weight）的NZ格式转换。 | ||
| 119 | + - 转置信息由算子内部根据tensor的stride自动推导，无需手动指定。 | ||
| 120 | + | ||
| 121 | +- 输入参数间支持的数据类型组合情况如下： | ||
| 122 | + | ||
| 123 | + **表 1** <term>Atlas A2 训练系列产品/Atlas A2 推理系列产品</term>、<term>Atlas A3 训练系列产品/Atlas A3 推理系列产品</term> | ||
| 124 | + | ||
| 125 | + | x1 | x2 | x1_scale | x2_scale | bias | 输出数据类型 | | ||
| 126 | + |-------|-------|----------|-----------|-------------------------------------|-----------| | ||
| 127 | + | int8 | int8 | float32 | float32 | int32/float32/bfloat16/float16/None | float16 | | ||
| 128 | + | int8 | int8 | float32 | bfloat16 | int32/float32/bfloat16/float16/None | bfloat16 | | ||
| 129 | + | int32 | int32 | float32 | float32 | int32/None | float16 | | ||
| 130 | + | int32 | int32 | float32 | bfloat16 | int32/None | bfloat16 | | ||
| 131 | + | int4 | int4 | float32 | float32 | int32/None | float16 | | ||
| 132 | + | int4 | int4 | float32 | bfloat16 | int32/None | bfloat16 | | ||
| 133 | + | ||
| 134 | +## 调用示例 | ||
| 135 | + | ||
| 136 | +- 单算子调用（A8W8量化，ND格式，gelu_tanh激活） | ||
| 137 | + | ||
| 138 | + ```python | ||
| 139 | + >>> import torch | ||
| 140 | + >>> import torch_npu | ||
| 141 | + >>> | ||
| 142 | + >>> m, k, n = 128, 256, 512 | ||
| 143 | + >>> x1 = torch.randint(-5, 5, (m, k), dtype=torch.int8).npu() | ||
| 144 | + >>> x2 = torch.randint(-5, 5, (k, n), dtype=torch.int8).npu() | ||
| 145 | + >>> x1_scale = torch.randn(m, dtype=torch.float32).abs().npu() * 0.01 | ||
| 146 | + >>> x2_scale = torch.randn(n, dtype=torch.float32).abs().npu() * 0.01 | ||
| 147 | + >>> | ||
| 148 | + >>> output = torch_npu.npu_quant_matmul_gelu(x1, x2, x1_scale, x2_scale, approximate="gelu_tanh") | ||
| 149 | + >>> print(output.shape) # torch.Size([128, 512]) | ||
| 150 | + >>> print(output.dtype) # torch.float16 | ||
| 151 | + ``` | ||
| 152 | + | ||
| 153 | +- 单算子调用（A8W8量化，ND格式，gelu_erf激活，带bias） | ||
| 154 | + | ||
| 155 | + ```python | ||
| 156 | + >>> import torch | ||
| 157 | + >>> import torch_npu | ||
| 158 | + >>> | ||
| 159 | + >>> m, k, n = 128, 256, 512 | ||
| 160 | + >>> x1 = torch.randint(-5, 5, (m, k), dtype=torch.int8).npu() | ||
| 161 | + >>> x2 = torch.randint(-5, 5, (k, n), dtype=torch.int8).npu() | ||
| 162 | + >>> x1_scale = torch.randn(m, dtype=torch.float32).abs().npu() * 0.01 | ||
| 163 | + >>> x2_scale = torch.randn(n, dtype=torch.float32).abs().npu() * 0.01 | ||
| 164 | + >>> bias = torch.randn(n, dtype=torch.float32).npu() * 0.1 | ||
| 165 | + >>> | ||
| 166 | + >>> # 使用gelu_erf激活并添加bias | ||
| 167 | + >>> output = torch_npu.npu_quant_matmul_gelu( | ||
| 168 | + ... x1, x2, x1_scale, x2_scale, bias=bias, approximate="gelu_erf" | ||
| 169 | + ... ) | ||
| 170 | + >>> print(output.shape) # torch.Size([128, 512]) | ||
| 171 | + >>> print(output.dtype) # torch.float16 | ||
| 172 | + ``` | ||
| 173 | + | ||
| 174 | +- 单算子调用（A8W8量化，NZ格式，gelu_tanh激活） | ||
| 175 | + | ||
| 176 | + ```python | ||
| 177 | + >>> import torch | ||
| 178 | + >>> import torch_npu | ||
| 179 | + >>> | ||
| 180 | + >>> m, k, n = 128, 256, 512 | ||
| 181 | + >>> x1 = torch.randint(-5, 5, (m, k), dtype=torch.int8).npu() | ||
| 182 | + >>> x2 = torch.randint(-5, 5, (k, n), dtype=torch.int8).npu() | ||
| 183 | + >>> | ||
| 184 | + >>> # 将x2转换为NZ格式以提升性能 | ||
| 185 | + >>> x2_nz = torch_npu.npu_format_cast(x2.contiguous(), 29) # 29为ACL_FORMAT_FRACTAL_NZ | ||
| 186 | + >>> | ||
| 187 | + >>> x1_scale = torch.randn(m, dtype=torch.float32).abs().npu() * 0.01 | ||
| 188 | + >>> x2_scale = torch.randn(n, dtype=torch.float32).abs().npu() * 0.01 | ||
| 189 | + >>> | ||
| 190 | + >>> # 自动识别NZ格式并调用对应接口 | ||
| 191 | + >>> output = torch_npu.npu_quant_matmul_gelu(x1, x2_nz, x1_scale, x2_scale, approximate="gelu_tanh") | ||
| 192 | + >>> print(output.shape) # torch.Size([128, 512]) | ||
| 193 | + >>> print(output.dtype) # torch.float16 | ||
| 194 | + ``` | ||
| 195 | + | ||
| 196 | +- 单算子调用（A8W8量化，bfloat16输出） | ||
| 197 | + | ||
| 198 | + ```python | ||
| 199 | + >>> import torch | ||
| 200 | + >>> import torch_npu | ||
| 201 | + >>> | ||
| 202 | + >>> m, k, n = 64, 128, 256 | ||
| 203 | + >>> x1 = torch.randint(-5, 5, (m, k), dtype=torch.int8).npu() | ||
| 204 | + >>> x2 = torch.randint(-5, 5, (k, n), dtype=torch.int8).npu() | ||
| 205 | + >>> x1_scale = torch.randn(m, dtype=torch.float32).abs().npu() * 0.01 | ||
| 206 | + >>> x2_scale = torch.randn(n, dtype=torch.bfloat16).abs().npu() * 0.01 # bfloat16 scale | ||
| 207 | + >>> | ||
| 208 | + >>> # 输出数据类型由x2_scale的类型决定，此处输出为bfloat16 | ||
| 209 | + >>> output = torch_npu.npu_quant_matmul_gelu(x1, x2, x1_scale, x2_scale, approximate="gelu_tanh") | ||
| 210 | + >>> print(output.dtype) # torch.bfloat16 | ||
| 211 | + ``` | ||
| 212 | + | ||
| 213 | +- 单算子调用（A4W4量化） | ||
| 214 | + | ||
| 215 | + ```python | ||
| 216 | + >>> import torch | ||
| 217 | + >>> import torch_npu | ||
| 218 | + >>> | ||
| 219 | + >>> m, k, n = 128, 256, 512 | ||
| 220 | + >>> # 生成int4数据（以int32格式存储） | ||
| 221 | + >>> # 注意：实际使用时需要通过量化接口将float32数据量化为int4并打包为int32 | ||
| 222 | + >>> x1 = torch.randint(-8, 8, (m, k // 8), dtype=torch.int32).npu() | ||
| 223 | + >>> x2 = torch.randint(-8, 8, (k, n // 8), dtype=torch.int32).npu() | ||
| 224 | + >>> | ||
| 225 | + >>> x1_scale = torch.randn(m, dtype=torch.float32).abs().npu() * 0.01 | ||
| 226 | + >>> x2_scale = torch.randn(n, dtype=torch.float32).abs().npu() * 0.01 | ||
| 227 | + >>> | ||
| 228 | + >>> # A4W4量化仅支持ND格式，不支持NZ格式 | ||
| 229 | + >>> # 转置信息由算子内部根据tensor的stride自动推导 | ||
| 230 | + >>> output = torch_npu.npu_quant_matmul_gelu( | ||
| 231 | + ... x1, x2, x1_scale, x2_scale, | ||
| 232 | + ... approximate="gelu_tanh" | ||
| 233 | + ... ) | ||
| 234 | + >>> print(output.shape) # torch.Size([128, 512]) | ||
| 235 | + >>> print(output.dtype) # torch.float16 | ||
| 236 | + ``` | ||
| 237 | + | ||
| 238 | +- 单算子调用（使用默认approximate="gelu_erf"） | ||
| 239 | + | ||
| 240 | + ```python | ||
| 241 | + >>> import torch | ||
| 242 | + >>> import torch_npu | ||
| 243 | + >>> | ||
| 244 | + >>> m, k, n = 64, 128, 256 | ||
| 245 | + >>> x1 = torch.randint(-5, 5, (m, k), dtype=torch.int8).npu() | ||
| 246 | + >>> x2 = torch.randint(-5, 5, (k, n), dtype=torch.int8).npu() | ||
| 247 | + >>> x1_scale = torch.randn(m, dtype=torch.float32).abs().npu() * 0.01 | ||
| 248 | + >>> x2_scale = torch.randn(n, dtype=torch.float32).abs().npu() * 0.01 | ||
| 249 | + >>> | ||
| 250 | + >>> # 不指定approximate参数，使用默认值"gelu_erf" | ||
| 251 | + >>> output = torch_npu.npu_quant_matmul_gelu(x1, x2, x1_scale, x2_scale) | ||
| 252 | + >>> print(output.dtype) # torch.float16 | ||
| 253 | + ``` | ||
| 254 | + | ||
| 255 | +- 图模式调用（A8W8量化，ND格式） | ||
| 256 | + | ||
| 257 | + ```python | ||
| 258 | + >>> import torch | ||
| 259 | + >>> import torch_npu | ||
| 260 | + >>> import torchair as tng | ||
| 261 | + >>> from torchair.configs.compiler_config import CompilerConfig | ||
| 262 | + >>> | ||
| 263 | + >>> os.environ["ENABLE_ACLNN"] = "true" | ||
| 264 | + >>> config = CompilerConfig() | ||
| 265 | + >>> npu_backend = tng.get_npu_backend(compiler_config=config) | ||
| 266 | + >>> | ||
| 267 | + >>> class MyModel(torch.nn.Module): | ||
| 268 | + ... def __init__(self): | ||
| 269 | + ... super().__init__() | ||
| 270 | + ... | ||
| 271 | + ... def forward(self, x1, x2, x1_scale, x2_scale, bias): | ||
| 272 | + ... return torch_npu.npu_quant_matmul_gelu( | ||
| 273 | + ... x1, x2, x1_scale, x2_scale, | ||
| 274 | + ... bias=bias, approximate="gelu_erf" | ||
| 275 | + ... ) | ||
| 276 | + >>> | ||
| 277 | + >>> cpu_model = MyModel() | ||
| 278 | + >>> model = cpu_model.npu() | ||
| 279 | + >>> cpu_x1 = torch.randint(-1, 1, (15, 1, 512), dtype=torch.int8) | ||
| 280 | + >>> cpu_x2 = torch.randint(-1, 1, (15, 512, 128), dtype=torch.int8) | ||
| 281 | + >>> x1_scale = torch.randn(15, dtype=torch.float32).abs() * 0.01 | ||
| 282 | + >>> x2_scale = torch.randn(128, dtype=torch.float32).abs() * 0.01 | ||
| 283 | + >>> bias = torch.randint(-1, 1, (15, 1, 128), dtype=torch.int32) | ||
| 284 | + >>> | ||
| 285 | + >>> model = torch.compile(model, backend=npu_backend, dynamic=True) | ||
| 286 | + >>> npu_out = model( | ||
| 287 | + ... cpu_x1.npu(), cpu_x2.npu(), | ||
| 288 | + ... x1_scale.npu(), x2_scale.npu(), | ||
| 289 | + ... bias.npu() | ||
| 290 | + ... ) | ||
| 291 | + >>> print(npu_out.shape) # torch.Size([15, 1, 128]) | ||
| 292 | + >>> print(npu_out.dtype) # torch.float16 | ||
| 293 | + ``` | ||
| @@ -531,6 +531,11 @@ | |||
| 531 | <td class="cellrowborder" valign="top" width="61.39%" headers="mcps1.2.3.1.2 "><p id="zh-cn_topic_0000001814195101_p156512056161014"><a name="zh-cn_topic_0000001814195101_p156512056161014"></a><a name="zh-cn_topic_0000001814195101_p156512056161014"></a>完成量化的矩阵乘计算，最小支持输入维度为2维，最大支持输入维度为6维。</p> | 531 | <td class="cellrowborder" valign="top" width="61.39%" headers="mcps1.2.3.1.2 "><p id="zh-cn_topic_0000001814195101_p156512056161014"><a name="zh-cn_topic_0000001814195101_p156512056161014"></a><a name="zh-cn_topic_0000001814195101_p156512056161014"></a>完成量化的矩阵乘计算，最小支持输入维度为2维，最大支持输入维度为6维。</p> |
| 532 | </td> | 532 | </td> |
| 533 | </tr> | 533 | </tr> |
| 534 | +<tr id="row9719124019219"><td class="cellrowborder" valign="top" width="38.61%" headers="mcps1.2.3.1.1 "><p id="p104326218132920"><a name="p104326218132920"></a><a name="p104326218132920"></a><a href="torch_npu-npu_quant_matmul_gelu.md">torch_npu.npu_quant_matmul_gelu</a></p> | ||
| 535 | +</td> | ||
| 536 | +<td class="cellrowborder" valign="top" width="61.39%" headers="mcps1.2.3.1.2 "><p id="p2026022511100607"><a name="p2026022511100607"></a><a name="p2026022511100607"></a>完成量化的矩阵乘和GELU激活函数的融合计算，支持A8W8和A4W4量化。</p> | ||
| 537 | +</td> | ||
| 538 | +</tr> | ||
| 534 | <tr id="row202508121056216"><td class="cellrowborder" valign="top" width="38.61%" headers="mcps1.2.3.1.1 "><p id="p202508121056216"><a name="p202508121056216"></a><a name="p202508121056216"></a><a href="torch_npu-npu_quant_matmul_reduce_sum.md">torch_npu.npu_quant_matmul_reduce_sum</a></p> | 539 | <tr id="row202508121056216"><td class="cellrowborder" valign="top" width="38.61%" headers="mcps1.2.3.1.1 "><p id="p202508121056216"><a name="p202508121056216"></a><a name="p202508121056216"></a><a href="torch_npu-npu_quant_matmul_reduce_sum.md">torch_npu.npu_quant_matmul_reduce_sum</a></p> |
| 535 | </td> | 540 | </td> |
| 536 | <td class="cellrowborder" valign="top" width="61.39%" headers="mcps1.2.3.1.2 "><p id="p2025081210578767"><a name="p2025081210578767"></a><a name="p2025081210578767"></a>完成量化的分组矩阵计算，然后所有组的矩阵计算结果相加后输出。</p> | 541 | <td class="cellrowborder" valign="top" width="61.39%" headers="mcps1.2.3.1.2 "><p id="p2025081210578767"><a name="p2025081210578767"></a><a name="p2025081210578767"></a>完成量化的分组矩阵计算，然后所有组的矩阵计算结果相加后输出。</p> |
| @@ -6445,6 +6445,11 @@ custom: | |||
| 6445 | op_api: [v2.1, newest] | 6445 | op_api: [v2.1, newest] |
| 6446 | internal_format_opapi: [v2.1, newest] | 6446 | internal_format_opapi: [v2.1, newest] |
| 6447 | 6447 | ||
| 6448 | + - func: npu_quant_matmul_gelu(Tensor x1, Tensor x2, Tensor x1_scale, Tensor x2_scale, *, Tensor? bias=None, str? approximate='gelu_erf') -> Tensor | ||
| 6449 | + op_api: all_version | ||
| 6450 | + internal_format_opapi: all_version | ||
| 6451 | + exposed: all_version | ||
| 6452 | + | ||
| 6448 | - func: npu_fused_matmul(Tensor x1, Tensor x2, *, Tensor? bias=None, Tensor? x3=None, str fused_op_type='') -> Tensor | 6453 | - func: npu_fused_matmul(Tensor x1, Tensor x2, *, Tensor? bias=None, Tensor? x3=None, str fused_op_type='') -> Tensor |
| 6449 | op_api: all_version | 6454 | op_api: all_version |
| 6450 | 6455 | ||
| @@ -0,0 +1,138 @@ | |||
| 1 | +// Copyright (c) 2026 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 | +// | ||
| 7 | +// Unless required by applicable law or agreed to in writing, software | ||
| 8 | +// distributed under the License is distributed on an "AS IS" BASIS, | ||
| 9 | +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. | ||
| 10 | +// See the License for the specific language governing permissions and | ||
| 11 | +// limitations under the License. | ||
| 12 | + | ||
| 13 | + | ||
| 14 | + | ||
| 15 | + | ||
| 16 | + | ||
| 17 | + | ||
| 18 | +namespace op_api { | ||
| 19 | + | ||
| 20 | +// 常量定义 | ||
| 21 | +constexpr size_t LAST_SECOND_DIM_INDEX = 2; | ||
| 22 | +constexpr int64_t INT4_NUMS_IN_INT32 = 8; | ||
| 23 | +constexpr size_t FUSED_TYPE_ARRAY_SIZE = 100; | ||
| 24 | +using npu_preparation = at_npu::native::OpPreparation; | ||
| 25 | + | ||
| 26 | +// 检查权重是否为 NZ 格式 | ||
| 27 | +static bool is_nz_format(const at::Tensor& x2) | ||
| 28 | +{ | ||
| 29 | + const torch_npu::NPUStorageDesc &tensor_desc = torch_npu::NPUBridge::GetNpuStorageImpl(x2)->npu_desc_; | ||
| 30 | + | ||
| 31 | + return (tensor_desc.npu_format_ == ACL_FORMAT_FRACTAL_NZ || | ||
| 32 | + tensor_desc.npu_format_ == ACL_FORMAT_FRACTAL_NZ_C0_4); | ||
| 33 | +} | ||
| 34 | + | ||
| 35 | +// 推导输出 batch shape | ||
| 36 | +static uint64_t infer_out_batch_shape_gelu(const at::Tensor &x1, const at::Tensor &x2, std::vector<uint64_t> &batch_record) | ||
| 37 | +{ | ||
| 38 | + TORCH_CHECK(at_npu::native::FormatHelper::IsBaseFormatType(x2) || is_nz_format(x2), | ||
| 39 | + "x2 should be in the original format or nz format, but it is ", | ||
| 40 | + npu_preparation::get_tensor_npu_format(x2), OPS_ERROR(ErrCode::PARAM)); | ||
| 41 | + | ||
| 42 | + uint64_t batch_val = 1; | ||
| 43 | + auto x1_dim_num = x1.dim(); | ||
| 44 | + auto x2_dim_num = x2.dim(); | ||
| 45 | + auto out_dim_num = std::max(x1_dim_num, x2_dim_num); | ||
| 46 | + auto &shape_long = x1_dim_num > x2_dim_num ? x1 : x2; | ||
| 47 | + auto &shape_short = x1_dim_num > x2_dim_num ? x2 : x1; | ||
| 48 | + int64_t valid_offset = out_dim_num - std::min(x1_dim_num, x2_dim_num); | ||
| 49 | + | ||
| 50 | + for (int64_t i = 0; i < out_dim_num - LAST_SECOND_DIM_INDEX; i++) { | ||
| 51 | + auto short_dim = i < valid_offset ? 1 : shape_short.size(i - valid_offset); | ||
| 52 | + auto long_dim = shape_long.size(i); | ||
| 53 | + TORCH_CHECK(!(short_dim > 1 && long_dim > 1 && short_dim != long_dim), | ||
| 54 | + "the x1 shape and x2 shape not supported for broadcast, the short_dim is ", | ||
| 55 | + short_dim, " and the long_dim is ", long_dim, OPS_ERROR(ErrCode::PARAM)); | ||
| 56 | + uint64_t cur_batch_value = static_cast<uint64_t>(std::max(short_dim, long_dim)); | ||
| 57 | + batch_val = batch_val * cur_batch_value; | ||
| 58 | + batch_record.push_back(cur_batch_value); | ||
| 59 | + } | ||
| 60 | + | ||
| 61 | + return batch_val; | ||
| 62 | +} | ||
| 63 | + | ||
| 64 | +// 算子接口 | ||
| 65 | +at::Tensor npu_quant_matmul_gelu( | ||
| 66 | + const at::Tensor &x1, | ||
| 67 | + const at::Tensor &x2, | ||
| 68 | + const at::Tensor &x1_scale, | ||
| 69 | + const at::Tensor &x2_scale, | ||
| 70 | + const c10::optional<at::Tensor> &bias, | ||
| 71 | + const c10::optional<c10::string_view> approximate) | ||
| 72 | +{ | ||
| 73 | + // 1. 校验输入场景，确定支持的量化场景（A4W4（INT4/INT32）或A8W8） | ||
| 74 | + bool is_a4w4 = (x1.dtype() == at::ScalarType::QUInt4x2 && x2.dtype() == at::ScalarType::QUInt4x2); | ||
| 75 | + bool is_a4w4_int32 = (x1.dtype() == at::kInt && x2.dtype() == at::kInt); | ||
| 76 | + bool is_a8w8 = (x1.dtype() == at::kChar && x2.dtype() == at::kChar); | ||
| 77 | + TORCH_CHECK(is_a4w4 || is_a4w4_int32 || is_a8w8, | ||
| 78 | + "Only A4W4 (int4/int32) or A8W8 (int8) quantization is supported, " | ||
| 79 | + "but got x1.dtype=", x1.dtype(), ", x2.dtype=", x2.dtype(), | ||
| 80 | + OPS_ERROR(ErrCode::TYPE)); | ||
| 81 | + | ||
| 82 | + // 2. 处理并校验 approximate 参数 | ||
| 83 | + c10::string_view approximate_value = approximate.value_or("gelu_erf"); // 默认值为 "gelu_erf" | ||
| 84 | + TORCH_CHECK(approximate_value == "gelu_tanh" || approximate_value == "gelu_erf", | ||
| 85 | + "approximate must be 'gelu_tanh' or 'gelu_erf', but got: ", | ||
| 86 | + approximate_value, OPS_ERROR(ErrCode::PARAM)); | ||
| 87 | + | ||
| 88 | + // 3. 推导输出 size | ||
| 89 | + int64_t x1_m_dim = x1.size(x1.dim() - LAST_SECOND_DIM_INDEX); | ||
| 90 | + int64_t x1_k_dim = x1.size(x1.dim() - 1); | ||
| 91 | + int64_t x2_k_dim = x2.size(x2.dim() - LAST_SECOND_DIM_INDEX); | ||
| 92 | + int64_t x2_n_dim = x2.size(x2.dim() - 1); | ||
| 93 | + // A4W4场景：当 x1: (m, k1 // 8); x2: (k2, n // 8) 时，需要对推导的 n * 8（INT32存储打包的8个INT4） | ||
| 94 | + if (x1_k_dim * INT4_NUMS_IN_INT32 == x2_k_dim) { | ||
| 95 | + x2_n_dim = x2_n_dim * INT4_NUMS_IN_INT32; | ||
| 96 | + } | ||
| 97 | + | ||
| 98 | + // 推导输出 shape | ||
| 99 | + std::vector<uint64_t> batch_record; | ||
| 100 | + infer_out_batch_shape_gelu(x1, x2, batch_record); | ||
| 101 | + const at::Tensor long_tensor = x1.dim() > x2.dim() ? x1 : x2; | ||
| 102 | + auto output_size = op_infer::array_to_small_vector(long_tensor.sizes()); | ||
| 103 | + output_size[long_tensor.dim() - LAST_SECOND_DIM_INDEX] = x1_m_dim; | ||
| 104 | + output_size[long_tensor.dim() - 1] = x2_n_dim; | ||
| 105 | + for (int64_t i = 0; i < long_tensor.dim() - LAST_SECOND_DIM_INDEX; i++) { | ||
| 106 | + output_size[i] = static_cast<int64_t>(batch_record[i]); | ||
| 107 | + } | ||
| 108 | + | ||
| 109 | + // 4. 分配输出张量 | ||
| 110 | + at::ScalarType output_dtype = (x2_scale.dtype() == at::kBFloat16) ? at::kBFloat16 : at::kHalf; | ||
| 111 | + const at::Tensor &bias_real = bias.value_or(at::Tensor()); | ||
| 112 | + if (bias_real.dtype() == at::kBFloat16) { | ||
| 113 | + output_dtype = at::kBFloat16; | ||
| 114 | + } | ||
| 115 | + c10::TensorOptions options = x1.options().dtype(output_dtype); | ||
| 116 | + at::Tensor result = npu_preparation::apply_tensor_without_format(output_size, options); | ||
| 117 | + | ||
| 118 | + // 5. 调用 aclnn 接口 | ||
| 119 | + // 默认参数处理 | ||
| 120 | + int64_t group_size = 0; | ||
| 121 | + char *approximate_str_ptr = const_cast<char *>(approximate_value.data()); | ||
| 122 | + const at::Tensor empty_tensor = at::Tensor(); | ||
| 123 | + | ||
| 124 | + // 根据 x2 格式选择接口 | ||
| 125 | + if (is_nz_format(x2)) { | ||
| 126 | + EXEC_NPU_CMD(aclnnFusedQuantMatmulWeightNz, x1, x2, x1_scale, x2_scale, | ||
| 127 | + empty_tensor, empty_tensor, empty_tensor, empty_tensor, bias_real, empty_tensor, | ||
| 128 | + approximate_str_ptr, group_size, result); | ||
| 129 | + } else { | ||
| 130 | + EXEC_NPU_CMD(aclnnFusedQuantMatmul, x1, x2, x1_scale, x2_scale, | ||
| 131 | + empty_tensor, empty_tensor, empty_tensor, empty_tensor, bias_real, empty_tensor, | ||
| 132 | + approximate_str_ptr, group_size, result); | ||
| 133 | + } | ||
| 134 | + | ||
| 135 | + return result; | ||
| 136 | +} | ||
| 137 | + | ||
| 138 | +} // namespace op_api | ||
| @@ -3533,6 +3533,158 @@ def npu_quant_matmul_reduce_sum_meta(x1, x2, *, x1_scale=None, x2_scale=None): | |||
| 3533 | return torch.empty(dst_shape, dtype=torch.bfloat16, device=x1.device) | 3533 | return torch.empty(dst_shape, dtype=torch.bfloat16, device=x1.device) |
| 3534 | 3534 | ||
| 3535 | 3535 | ||
| 3536 | + | ||
| 3537 | +def npu_quant_matmul_gelu_meta(x1, x2, x1_scale, x2_scale, *, bias=None, approximate="gelu_erf"): | ||
| 3538 | + INT4_IN_INT32 = 8 | ||
| 3539 | + LAST_SECOND_DIM_INDEX = 2 | ||
| 3540 | + | ||
| 3541 | + # 校验approximate参数 | ||
| 3542 | + torch._check( | ||
| 3543 | + approximate in ["gelu_tanh", "gelu_erf"], | ||
| 3544 | + lambda: f"approximate must be 'gelu_tanh' or 'gelu_erf', but got {approximate} {ops_error(ErrCode.PARAM)}", | ||
| 3545 | + ) | ||
| 3546 | + | ||
| 3547 | + # 校验量化场景（A4W4或A8W8） | ||
| 3548 | + is_a4w4 = ((x1.dtype == torch.int32 or x1.dtype == torch.quint4x2) and | ||
| 3549 | + (x2.dtype == torch.int32 or x2.dtype == torch.quint4x2)) | ||
| 3550 | + is_a8w8 = (x1.dtype == torch.int8 and x2.dtype == torch.int8) | ||
| 3551 | + | ||
| 3552 | + torch._check( | ||
| 3553 | + is_a4w4 or is_a8w8, | ||
| 3554 | + lambda: f"Only A4W4 (int4/int32) or A8W8 (int8) quantization is supported, " | ||
| 3555 | + f"but got x1.dtype={x1.dtype}, x2.dtype={x2.dtype} {ops_error(ErrCode.TYPE)}", | ||
| 3556 | + ) | ||
| 3557 | + | ||
| 3558 | + # 校验x1_scale和x2_scale | ||
| 3559 | + torch._check(x1_scale is not None, lambda: f"x1_scale should not be None.") | ||
| 3560 | + torch._check(x2_scale is not None, lambda: f"x2_scale should not be None.") | ||
| 3561 | + torch._check(x1_scale.dim() == 1, lambda: f"x1_scale dim must be 1, but got {x1_scale.dim()}.") | ||
| 3562 | + torch._check(x2_scale.dim() == 1, lambda: f"x2_scale dim must be 1, but got {x2_scale.dim()}.") | ||
| 3563 | + | ||
| 3564 | + # 推导输出shape | ||
| 3565 | + x1_dim_num = x1.dim() | ||
| 3566 | + x2_dim_num = x2.dim() | ||
| 3567 | + out_dim_num = max(x1_dim_num, x2_dim_num) | ||
| 3568 | + shape_long = x1 if x1_dim_num > x2_dim_num else x2 | ||
| 3569 | + shape_short = x2 if x1_dim_num > x2_dim_num else x1 | ||
| 3570 | + valid_offset = out_dim_num - min(x1_dim_num, x2_dim_num) | ||
| 3571 | + | ||
| 3572 | + # 计算batch维度 | ||
| 3573 | + batch_val = 1 | ||
| 3574 | + dim_list = [] | ||
| 3575 | + for i in range(0, out_dim_num - LAST_SECOND_DIM_INDEX): | ||
| 3576 | + short_dim = 1 if i < valid_offset else shape_short.size(i - valid_offset) | ||
| 3577 | + long_dim = shape_long.size(i) | ||
| 3578 | + torch._check( | ||
| 3579 | + not (short_dim > 1 and long_dim > 1 and short_dim != long_dim), | ||
| 3580 | + lambda: "the batch shape cannot be broadcast" + ops_error(ErrCode.VALUE), | ||
| 3581 | + ) | ||
| 3582 | + cur_batch_val = max(short_dim, long_dim) | ||
| 3583 | + batch_val = batch_val * cur_batch_val | ||
| 3584 | + dim_list.append(cur_batch_val) | ||
| 3585 | + | ||
| 3586 | + # 计算x1的m维度和k维度 | ||
| 3587 | + x1_m_dim = x1.size(x1_dim_num - LAST_SECOND_DIM_INDEX) | ||
| 3588 | + x1_k_dim = x1.size(x1_dim_num - 1) | ||
| 3589 | + | ||
| 3590 | + # 计算x2的k维度和n维度（考虑INT4打包） | ||
| 3591 | + x2_k_dim = x2.size(x2_dim_num - LAST_SECOND_DIM_INDEX) | ||
| 3592 | + x2_n_dim = x2.size(x2_dim_num - 1) | ||
| 3593 | + | ||
| 3594 | + # 校验k维度匹配 | ||
| 3595 | + if is_a4w4: | ||
| 3596 | + # A4W4场景：k维度匹配检查 | ||
| 3597 | + if x1.dtype == torch.int32 and x2.dtype == torch.int32: | ||
| 3598 | + # 两个都是int32（打包INT4），k维度应该相等 | ||
| 3599 | + torch._check( | ||
| 3600 | + x1_k_dim == x2_k_dim, | ||
| 3601 | + lambda: f"A4W4 (int32): k dim of x1 ({x1_k_dim}) must equal k dim of x2 ({x2_k_dim}) {ops_error(ErrCode.VALUE)}", | ||
| 3602 | + ) | ||
| 3603 | + elif x1.dtype == torch.quint4x2 and x2.dtype == torch.quint4x2: | ||
| 3604 | + # 两个都是quint4x2（直接INT4），k维度应该相等 | ||
| 3605 | + torch._check( | ||
| 3606 | + x1_k_dim == x2_k_dim, | ||
| 3607 | + lambda: f"A4W4 (quint4x2): k dim of x1 ({x1_k_dim}) must equal k dim of x2 ({x2_k_dim}) {ops_error(ErrCode.VALUE)}", | ||
| 3608 | + ) | ||
| 3609 | + elif x1.dtype == torch.int32 and x2.dtype == torch.quint4x2: | ||
| 3610 | + # x1是int32（打包），x2是quint4x2（直接），x1的k维度 * 8应该等于x2的k维度 | ||
| 3611 | + torch._check( | ||
| 3612 | + x1_k_dim * INT4_IN_INT32 == x2_k_dim, | ||
| 3613 | + lambda: f"A4W4 (int32/quint4x2): k dim of x1 ({x1_k_dim}) * 8 must equal k dim of x2 ({x2_k_dim}) {ops_error(ErrCode.VALUE)}", | ||
| 3614 | + ) | ||
| 3615 | + elif x1.dtype == torch.quint4x2 and x2.dtype == torch.int32: | ||
| 3616 | + # x1是quint4x2（直接），x2是int32（打包），x1的k维度应该等于x2的k维度 * 8 | ||
| 3617 | + torch._check( | ||
| 3618 | + x1_k_dim == x2_k_dim * INT4_IN_INT32, | ||
| 3619 | + lambda: f"A4W4 (quint4x2/int32): k dim of x1 ({x1_k_dim}) must equal k dim of x2 ({x2_k_dim}) * 8 {ops_error(ErrCode.VALUE)}", | ||
| 3620 | + ) | ||
| 3621 | + | ||
| 3622 | + # A4W4场景：当x2为int32类型时，恢复实际的n维度（INT4打包） | ||
| 3623 | + if x2.dtype == torch.int32: | ||
| 3624 | + x2_n_dim = x2_n_dim * INT4_IN_INT32 | ||
| 3625 | + else: | ||
| 3626 | + # A8W8场景：k维度应该相等 | ||
| 3627 | + torch._check( | ||
| 3628 | + x1_k_dim == x2_k_dim, | ||
| 3629 | + lambda: f"A8W8: k dim of x1 ({x1_k_dim}) must equal k dim of x2 ({x2_k_dim}) {ops_error(ErrCode.VALUE)}", | ||
| 3630 | + ) | ||
| 3631 | + | ||
| 3632 | + dim_list.append(x1_m_dim) | ||
| 3633 | + dim_list.append(x2_n_dim) | ||
| 3634 | + | ||
| 3635 | + # 校验x1_scale和x2_scale的shape | ||
| 3636 | + torch._check( | ||
| 3637 | + x1_scale.size(0) == x1_m_dim, | ||
| 3638 | + lambda: f"x1_scale size(0) must equal to x1's m dimension ({x1_m_dim}), but got {x1_scale.size(0)} {ops_error(ErrCode.VALUE)}", | ||
| 3639 | + ) | ||
| 3640 | + torch._check( | ||
| 3641 | + x2_scale.size(0) == 1 or x2_scale.size(0) == x2_n_dim, | ||
| 3642 | + lambda: f"x2_scale size(0) must be 1 or equal to x2's n dimension ({x2_n_dim}), but got {x2_scale.size(0)} {ops_error(ErrCode.VALUE)}", | ||
| 3643 | + ) | ||
| 3644 | + | ||
| 3645 | + # 校验bias shape（如果提供） | ||
| 3646 | + if bias is not None: | ||
| 3647 | + if is_a4w4: | ||
| 3648 | + # A4W4场景：bias shape仅支持1维(n,) | ||
| 3649 | + torch._check( | ||
| 3650 | + bias.dim() == 1, | ||
| 3651 | + lambda: f"A4W4 quantization only supports 1D bias, but got bias.dim()={bias.dim()} {ops_error(ErrCode.VALUE)}", | ||
| 3652 | + ) | ||
| 3653 | + torch._check( | ||
| 3654 | + bias.size(0) == x2_n_dim, | ||
| 3655 | + lambda: f"bias size(0) must equal to x2's n dimension ({x2_n_dim}), but got {bias.size(0)} {ops_error(ErrCode.VALUE)}", | ||
| 3656 | + ) | ||
| 3657 | + else: | ||
| 3658 | + # A8W8场景：bias shape支持1维(n,)或3维(batch, 1, n) | ||
| 3659 | + torch._check( | ||
| 3660 | + bias.dim() == 1 or bias.dim() == 3, | ||
| 3661 | + lambda: f"A8W8 quantization supports 1D or 3D bias, but got bias.dim()={bias.dim()} {ops_error(ErrCode.VALUE)}", | ||
| 3662 | + ) | ||
| 3663 | + if bias.dim() == 1: | ||
| 3664 | + torch._check( | ||
| 3665 | + bias.size(0) == x2_n_dim, | ||
| 3666 | + lambda: f"bias size(0) must equal to x2's n dimension ({x2_n_dim}), but got {bias.size(0)} {ops_error(ErrCode.VALUE)}", | ||
| 3667 | + ) | ||
| 3668 | + else: # bias.dim() == 3 | ||
| 3669 | + torch._check( | ||
| 3670 | + len(dim_list) == 3, | ||
| 3671 | + lambda: f"when bias dim is 3, out dim need to be 3 {ops_error(ErrCode.TYPE)}", | ||
| 3672 | + ) | ||
| 3673 | + torch._check( | ||
| 3674 | + bias.size(0) == batch_val and bias.size(1) == 1 and bias.size(2) == x2_n_dim, | ||
| 3675 | + lambda: f"bias shape must be ({batch_val}, 1, {x2_n_dim}), but got {tuple(bias.shape)} {ops_error(ErrCode.VALUE)}", | ||
| 3676 | + ) | ||
| 3677 | + | ||
| 3678 | + # 确定输出数据类型（根据x2_scale的数据类型） | ||
| 3679 | + output_dtype = torch.bfloat16 if x2_scale.dtype == torch.bfloat16 else torch.float16 | ||
| 3680 | + | ||
| 3681 | + # 在FakeTensor模式下，使用meta设备；否则使用x1的设备 | ||
| 3682 | + # 确保设备一致性，避免设备传播错误 | ||
| 3683 | + output_device = 'meta' if x1.device.type == 'meta' else x1.device | ||
| 3684 | + | ||
| 3685 | + return torch.empty(tuple(dim_list), dtype=output_dtype, device=output_device) | ||
| 3686 | + | ||
| 3687 | + | ||
| 3536 | 3688 | ||
| 3537 | def npu_quant_grouped_matmul_dequant_meta(x, quantized_weight, weight_scale, group_list, *, | 3689 | def npu_quant_grouped_matmul_dequant_meta(x, quantized_weight, weight_scale, group_list, *, |
| 3538 | bias=None, x_scale=None, x_offset=None, smooth_scale=None, quant_mode="pertoken"): | 3690 | bias=None, x_scale=None, x_offset=None, smooth_scale=None, quant_mode="pertoken"): |
| @@ -143,7 +143,8 @@ | |||
| 143 | "npu_confusion_transpose", | 143 | "npu_confusion_transpose", |
| 144 | "npu_dense_lightning_indexer_grad_kl_loss", | 144 | "npu_dense_lightning_indexer_grad_kl_loss", |
| 145 | "npu_dense_lightning_indexer_softmax_lse", | 145 | "npu_dense_lightning_indexer_softmax_lse", |
| 146 | - "npu_one_hot" | 146 | + "npu_one_hot", |
| 147 | + "npu_quant_matmul_gelu" | ||
| 147 | ] | 148 | ] |
| 148 | } | 149 | } |
| 149 | 150 | ||
| @@ -2731,6 +2731,67 @@ class TestQuantMatmulReduceSum(TestCase): | |||
| 2731 | self.assertTrue(y.dtype == torch.bfloat16) | 2731 | self.assertTrue(y.dtype == torch.bfloat16) |
| 2732 | 2732 | ||
| 2733 | 2733 | ||
| 2734 | +class TestQuantMatmulGelu(TestCase): | ||
| 2735 | + | ||
| 2736 | + def test_npu_quant_matmul_gelu_meta(self): | ||
| 2737 | + with FakeTensorMode(): | ||
| 2738 | + # Test A8W8 scenario | ||
| 2739 | + m, k, n = 128, 256, 512 | ||
| 2740 | + x1 = torch.randint(-5, 5, (m, k), dtype=torch.int8).npu() | ||
| 2741 | + x2 = torch.randint(-5, 5, (k, n), dtype=torch.int8).npu() | ||
| 2742 | + x1_scale = torch.randn(m, dtype=torch.float32).abs() * 0.01 | ||
| 2743 | + x2_scale = torch.randn(n, dtype=torch.float32).abs() * 0.01 | ||
| 2744 | + | ||
| 2745 | + expect_ret = torch.randn((m, n), dtype=torch.float16).npu() | ||
| 2746 | + res = torch_npu.npu_quant_matmul_gelu(x1, x2, x1_scale.npu(), x2_scale.npu(), approximate="gelu_tanh") | ||
| 2747 | + self.assertTrue(expect_ret.shape == res.shape) | ||
| 2748 | + self.assertTrue(expect_ret.dtype == res.dtype) | ||
| 2749 | + | ||
| 2750 | + # Test A8W8 with bias | ||
| 2751 | + bias = torch.randn(n, dtype=torch.float32) * 0.1 | ||
| 2752 | + res_bias = torch_npu.npu_quant_matmul_gelu(x1, x2, x1_scale.npu(), x2_scale.npu(), bias=bias.npu(), approximate="gelu_erf") | ||
| 2753 | + self.assertTrue(expect_ret.shape == res_bias.shape) | ||
| 2754 | + self.assertTrue(expect_ret.dtype == res_bias.dtype) | ||
| 2755 | + | ||
| 2756 | + # Test A8W8 with BF16 output | ||
| 2757 | + x2_scale_bf16 = torch.randn(n, dtype=torch.bfloat16).abs() * 0.01 | ||
| 2758 | + expect_ret_bf16 = torch.randn((m, n), dtype=torch.bfloat16).npu() | ||
| 2759 | + res_bf16 = torch_npu.npu_quant_matmul_gelu(x1, x2, x1_scale.npu(), x2_scale_bf16.npu(), approximate="gelu_tanh") | ||
| 2760 | + self.assertTrue(expect_ret_bf16.shape == res_bf16.shape) | ||
| 2761 | + self.assertTrue(expect_ret_bf16.dtype == res_bf16.dtype) | ||
| 2762 | + | ||
| 2763 | + # Test A8W8 with batch dimensions | ||
| 2764 | + batch, m, k, n = 4, 64, 128, 256 | ||
| 2765 | + x1_batch = torch.randint(-5, 5, (batch, m, k), dtype=torch.int8).npu() | ||
| 2766 | + x2_batch = torch.randint(-5, 5, (batch, k, n), dtype=torch.int8).npu() | ||
| 2767 | + x1_scale_batch = torch.randn(m, dtype=torch.float32).abs() * 0.01 | ||
| 2768 | + x2_scale_batch = torch.randn(n, dtype=torch.float32).abs() * 0.01 | ||
| 2769 | + expect_ret_batch = torch.randn((batch, m, n), dtype=torch.float16).npu() | ||
| 2770 | + res_batch = torch_npu.npu_quant_matmul_gelu(x1_batch, x2_batch, x1_scale_batch.npu(), x2_scale_batch.npu()) | ||
| 2771 | + self.assertTrue(expect_ret_batch.shape == res_batch.shape) | ||
| 2772 | + self.assertTrue(expect_ret_batch.dtype == res_batch.dtype) | ||
| 2773 | + | ||
| 2774 | + # Test A4W4 with int32 (packed INT4) | ||
| 2775 | + m, k, n = 64, 128, 256 | ||
| 2776 | + k_packed = k // 8 | ||
| 2777 | + n_packed = n // 8 | ||
| 2778 | + x1_int32 = torch.randint(-8, 8, (m, k_packed), dtype=torch.int32).npu() | ||
| 2779 | + x2_int32 = torch.randint(-8, 8, (k_packed, n_packed), dtype=torch.int32).npu() | ||
| 2780 | + x1_scale_int32 = torch.randn(m, dtype=torch.float32).abs() * 0.01 | ||
| 2781 | + x2_scale_int32 = torch.randn(n, dtype=torch.float32).abs() * 0.01 | ||
| 2782 | + expect_ret_int32 = torch.randn((m, n), dtype=torch.float16).npu() | ||
| 2783 | + res_int32 = torch_npu.npu_quant_matmul_gelu(x1_int32, x2_int32, x1_scale_int32.npu(), x2_scale_int32.npu(), approximate="gelu_tanh") | ||
| 2784 | + self.assertTrue(expect_ret_int32.shape == res_int32.shape) | ||
| 2785 | + self.assertTrue(expect_ret_int32.dtype == res_int32.dtype) | ||
| 2786 | + | ||
| 2787 | + # Test A4W4 with int32 and bias | ||
| 2788 | + bias_int32 = torch.randint(-5, 5, (n,), dtype=torch.int32) | ||
| 2789 | + res_int32_bias = torch_npu.npu_quant_matmul_gelu(x1_int32, x2_int32, x1_scale_int32.npu(), x2_scale_int32.npu(), | ||
| 2790 | + bias=bias_int32.npu(), approximate="gelu_erf") | ||
| 2791 | + self.assertTrue(expect_ret_int32.shape == res_int32_bias.shape) | ||
| 2792 | + self.assertTrue(expect_ret_int32.dtype == res_int32_bias.dtype) | ||
| 2793 | + | ||
| 2794 | + | ||
| 2734 | class TestRecurrentGatedDeltaRule(TestCase): | 2795 | class TestRecurrentGatedDeltaRule(TestCase): |
| 2735 | def test_recurrent_gated_delta_rule(self): | 2796 | def test_recurrent_gated_delta_rule(self): |
| 2736 | with FakeTensorMode(): | 2797 | with FakeTensorMode(): |
| @@ -308,6 +308,9 @@ | |||
| 308 | "op_api: torch_npu.npu_quant_matmul(*args, **kwargs)": { | 308 | "op_api: torch_npu.npu_quant_matmul(*args, **kwargs)": { |
| 309 | "version": ["v2.1", "v2.5", "v2.6", "v2.7"] | 309 | "version": ["v2.1", "v2.5", "v2.6", "v2.7"] |
| 310 | }, | 310 | }, |
| 311 | + "op_api: torch_npu.npu_quant_matmul_gelu(*args, **kwargs)": { | ||
| 312 | + "version": ["v2.1", "v2.5", "v2.6", "v2.7"] | ||
| 313 | + }, | ||
| 311 | "op_api: torch_npu.npu_all_gather_base_mm(*args, **kwargs)": { | 314 | "op_api: torch_npu.npu_all_gather_base_mm(*args, **kwargs)": { |
| 312 | "version": ["v2.1", "v2.5", "v2.6", "v2.7"] | 315 | "version": ["v2.1", "v2.5", "v2.6", "v2.7"] |
| 313 | }, | 316 | }, |
| @@ -493,6 +496,9 @@ | |||
| 493 | "op_api: torch_npu.npu_quant_matmul(*args: _P.args, **kwargs: _P.kwargs) -> ~_T": { | 496 | "op_api: torch_npu.npu_quant_matmul(*args: _P.args, **kwargs: _P.kwargs) -> ~_T": { |
| 494 | "version": ["v2.8", "newest"] | 497 | "version": ["v2.8", "newest"] |
| 495 | }, | 498 | }, |
| 499 | + "op_api: torch_npu.npu_quant_matmul_gelu(*args: _P.args, **kwargs: _P.kwargs) -> ~_T": { | ||
| 500 | + "version": ["v2.8", "newest"] | ||
| 501 | + }, | ||
| 496 | "op_api: torch_npu.npu_quant_matmul_reduce_sum(*args, **kwargs)": { | 502 | "op_api: torch_npu.npu_quant_matmul_reduce_sum(*args, **kwargs)": { |
| 497 | "version": ["v2.1", "v2.5", "v2.6", "v2.7"] | 503 | "version": ["v2.1", "v2.5", "v2.6", "v2.7"] |
| 498 | }, | 504 | }, |
| @@ -1477,6 +1483,9 @@ | |||
| 1477 | "func: npu_quant_matmul(Tensor x1, Tensor x2, Tensor scale, *, Tensor? offset=None, Tensor? pertoken_scale=None, Tensor? bias=None, int? output_dtype=None, int? x1_dtype=None, int? x2_dtype=None, int? pertoken_scale_dtype=None, int? scale_dtype=None, int[]? group_sizes=None, Tensor? y_scale=None) -> Tensor": { | 1483 | "func: npu_quant_matmul(Tensor x1, Tensor x2, Tensor scale, *, Tensor? offset=None, Tensor? pertoken_scale=None, Tensor? bias=None, int? output_dtype=None, int? x1_dtype=None, int? x2_dtype=None, int? pertoken_scale_dtype=None, int? scale_dtype=None, int[]? group_sizes=None, Tensor? y_scale=None) -> Tensor": { |
| 1478 | "version": ["all_version"] | 1484 | "version": ["all_version"] |
| 1479 | }, | 1485 | }, |
| 1486 | + "func: npu_quant_matmul_gelu(Tensor x1, Tensor x2, Tensor x1_scale, Tensor x2_scale, *, Tensor? bias=None, str? approximate='gelu_erf') -> Tensor": { | ||
| 1487 | + "version": ["all_version"] | ||
| 1488 | + }, | ||
| 1480 | "func: npu_dual_level_quant_matmul(Tensor x1, Tensor x2, Tensor x1_level0_scale, Tensor x2_level0_scale, Tensor x1_level1_scale, Tensor x2_level1_scale, *, Tensor? bias=None, int output_dtype=15) -> Tensor": { | 1489 | "func: npu_dual_level_quant_matmul(Tensor x1, Tensor x2, Tensor x1_level0_scale, Tensor x2_level0_scale, Tensor x1_level1_scale, Tensor x2_level1_scale, *, Tensor? bias=None, int output_dtype=15) -> Tensor": { |
| 1481 | "version": ["v2.1", "newest"] | 1490 | "version": ["v2.1", "newest"] |
| 1482 | }, | 1491 | }, |
| @@ -0,0 +1,458 @@ | |||
| 1 | +import math | ||
| 2 | +import unittest | ||
| 3 | +import numpy as np | ||
| 4 | +import torch | ||
| 5 | +import torch_npu | ||
| 6 | +import torch.nn as nn | ||
| 7 | +from torch_npu.testing.testcase import TestCase, run_tests | ||
| 8 | +from torch_npu.testing.common_utils import SupportedDevices | ||
| 9 | + | ||
| 10 | + | ||
| 11 | +class TestQuantMatmulGelu(TestCase): | ||
| 12 | + | ||
| 13 | + def gelu_tanh(self, x): | ||
| 14 | + """GELU tanh approximation""" | ||
| 15 | + return 0.5 * x * (1.0 + torch.tanh(math.sqrt(2.0 / math.pi) * (x + 0.044715 * torch.pow(x, 3.0)))) | ||
| 16 | + | ||
| 17 | + def gelu_erf(self, x): | ||
| 18 | + """GELU erf""" | ||
| 19 | + m = nn.GELU() | ||
| 20 | + return m(x) | ||
| 21 | + | ||
| 22 | + def cpu_quant_matmul_gelu(self, x1, x2, x1_scale, x2_scale, bias=None, approximate="gelu_erf"): | ||
| 23 | + """CPU reference implementation""" | ||
| 24 | + # Handle different input types | ||
| 25 | + if x1.dtype == torch.quint4x2: | ||
| 26 | + # For quint4x2 (direct INT4), convert to int32 for computation | ||
| 27 | + # quint4x2 stores 2 INT4 values per uint8, so we view as uint8 then to int32 | ||
| 28 | + x1_int32 = x1.view(torch.uint8).to(torch.int32) | ||
| 29 | + x2_int32 = x2.view(torch.uint8).to(torch.int32) | ||
| 30 | + elif x1.dtype == torch.int32: | ||
| 31 | + # For int32 (packed INT4), each int32 stores 8 int4 values | ||
| 32 | + # We need to expand the dimensions: (m, k//8) -> (m, k) and (k, n//8) -> (k, n) | ||
| 33 | + # For simplicity, we'll use the packed dimensions directly | ||
| 34 | + # Note: This is a simplified CPU reference, actual unpacking would be more complex | ||
| 35 | + x1_int32 = x1.to(torch.int32) | ||
| 36 | + x2_int32 = x2.to(torch.int32) | ||
| 37 | + # For packed INT4 in INT32, we need to handle dimension expansion | ||
| 38 | + # Since CPU reference is simplified, we'll use the packed shape directly | ||
| 39 | + # The actual unpacking logic would expand each int32 to 8 int4 values | ||
| 40 | + else: | ||
| 41 | + # A8W8 case: int8 | ||
| 42 | + x1_int32 = x1.to(torch.int32) | ||
| 43 | + x2_int32 = x2.to(torch.int32) | ||
| 44 | + | ||
| 45 | + # Compute quantized matmul: x1 @ x2 * x1_scale * x2_scale | ||
| 46 | + result = torch.matmul(x1_int32, x2_int32).float() | ||
| 47 | + result = result * x1_scale.view(-1, 1) * x2_scale.view(1, -1) | ||
| 48 | + | ||
| 49 | + # Add bias if provided | ||
| 50 | + if bias is not None: | ||
| 51 | + result = result + bias.view(1, -1) | ||
| 52 | + | ||
| 53 | + # Apply GELU | ||
| 54 | + if approximate == "gelu_tanh": | ||
| 55 | + result = self.gelu_tanh(result) | ||
| 56 | + elif approximate == "gelu_erf": | ||
| 57 | + result = self.gelu_erf(result) | ||
| 58 | + else: | ||
| 59 | + raise ValueError(f"Unsupported approximate: {approximate}") | ||
| 60 | + | ||
| 61 | + return result | ||
| 62 | + | ||
| 63 | + | ||
| 64 | + | ||
 | |||
| 65 | + def test_npu_quant_matmul_gelu_a8w8_nd_gelu_tanh(self): | ||
| 66 | + """Test A8W8 with ND format and gelu_tanh""" | ||
| 67 | + torch.manual_seed(0) | ||
| 68 | + m, k, n = 128, 256, 512 | ||
| 69 | + | ||
| 70 | + x1 = torch.randint(-5, 5, (m, k), dtype=torch.int8) | ||
| 71 | + x2 = torch.randint(-5, 5, (k, n), dtype=torch.int8) | ||
| 72 | + x1_scale = torch.randn(m, dtype=torch.float32).abs() * 0.01 | ||
| 73 | + x2_scale = torch.randn(n, dtype=torch.float32).abs() * 0.01 | ||
| 74 | + | ||
| 75 | + # CPU reference | ||
| 76 | + cpu_result = self.cpu_quant_matmul_gelu(x1, x2, x1_scale, x2_scale, approximate="gelu_tanh") | ||
| 77 | + | ||
| 78 | + # NPU custom op | ||
| 79 | + npu_result = torch_npu.npu_quant_matmul_gelu( | ||
| 80 | + x1.npu(), x2.npu(), x1_scale.npu(), x2_scale.npu(), approximate="gelu_tanh") | ||
| 81 | + | ||
| 82 | + self.assertRtolEqual(cpu_result.numpy(), npu_result.cpu().float().numpy(), 0.01) | ||
| 83 | + | ||
| 84 | + | ||
| 85 | + | ||
| 86 | + def test_npu_quant_matmul_gelu_a8w8_nd_gelu_erf(self): | ||
| 87 | + """Test A8W8 with ND format and gelu_erf""" | ||
| 88 | + torch.manual_seed(0) | ||
| 89 | + m, k, n = 128, 256, 512 | ||
| 90 | + | ||
| 91 | + x1 = torch.randint(-5, 5, (m, k), dtype=torch.int8) | ||
| 92 | + x2 = torch.randint(-5, 5, (k, n), dtype=torch.int8) | ||
| 93 | + x1_scale = torch.randn(m, dtype=torch.float32).abs() * 0.01 | ||
| 94 | + x2_scale = torch.randn(n, dtype=torch.float32).abs() * 0.01 | ||
| 95 | + | ||
| 96 | + # CPU reference | ||
| 97 | + cpu_result = self.cpu_quant_matmul_gelu(x1, x2, x1_scale, x2_scale, approximate="gelu_erf") | ||
| 98 | + | ||
| 99 | + # NPU custom op | ||
| 100 | + npu_result = torch_npu.npu_quant_matmul_gelu( | ||
| 101 | + x1.npu(), x2.npu(), x1_scale.npu(), x2_scale.npu(), approximate="gelu_erf") | ||
| 102 | + | ||
| 103 | + self.assertRtolEqual(cpu_result.numpy(), npu_result.cpu().float().numpy(), 0.01) | ||
| 104 | + | ||
| 105 | + | ||
| 106 | + | ||
| 107 | + def test_npu_quant_matmul_gelu_a8w8_nd_with_bias(self): | ||
| 108 | + """Test A8W8 with ND format, gelu_tanh and bias""" | ||
| 109 | + torch.manual_seed(0) | ||
| 110 | + m, k, n = 128, 256, 512 | ||
| 111 | + | ||
| 112 | + x1 = torch.randint(-5, 5, (m, k), dtype=torch.int8) | ||
| 113 | + x2 = torch.randint(-5, 5, (k, n), dtype=torch.int8) | ||
| 114 | + x1_scale = torch.randn(m, dtype=torch.float32).abs() * 0.01 | ||
| 115 | + x2_scale = torch.randn(n, dtype=torch.float32).abs() * 0.01 | ||
| 116 | + bias = torch.randn(n, dtype=torch.float32) * 0.1 | ||
| 117 | + | ||
| 118 | + # CPU reference | ||
| 119 | + cpu_result = self.cpu_quant_matmul_gelu(x1, x2, x1_scale, x2_scale, approximate="gelu_tanh", bias=bias) | ||
| 120 | + | ||
| 121 | + # NPU custom op | ||
| 122 | + npu_result = torch_npu.npu_quant_matmul_gelu( | ||
| 123 | + x1.npu(), x2.npu(), x1_scale.npu(), x2_scale.npu(), | ||
| 124 | + bias=bias.npu(), approximate="gelu_tanh") | ||
| 125 | + | ||
| 126 | + self.assertRtolEqual(cpu_result.numpy(), npu_result.cpu().float().numpy(), 0.01) | ||
| 127 | + | ||
| 128 | + | ||
| 129 | + | ||
| 130 | + def test_npu_quant_matmul_gelu_a8w8_nz_gelu_tanh(self): | ||
| 131 | + """Test A8W8 with NZ format and gelu_tanh""" | ||
| 132 | + torch.manual_seed(0) | ||
| 133 | + m, k, n = 128, 256, 512 | ||
| 134 | + | ||
| 135 | + x1 = torch.randint(-1, 1, (m, k), dtype=torch.int8) | ||
| 136 | + x2 = torch.randint(-1, 1, (k, n), dtype=torch.int8) | ||
| 137 | + x1_scale = torch.randn(m, dtype=torch.float32).abs() * 0.01 | ||
| 138 | + x2_scale = torch.randn(n, dtype=torch.float32).abs() * 0.01 | ||
| 139 | + | ||
| 140 | + # CPU reference | ||
| 141 | + cpu_result = self.cpu_quant_matmul_gelu(x1, x2, x1_scale, x2_scale, approximate="gelu_tanh") | ||
| 142 | + | ||
| 143 | + # Convert x2 to NZ format | ||
| 144 | + x2_nz = torch_npu.npu_format_cast(x2.npu().contiguous(), 29) # 29 is ACL_FORMAT_FRACTAL_NZ | ||
| 145 | + | ||
| 146 | + # NPU custom op with NZ format | ||
| 147 | + npu_result = torch_npu.npu_quant_matmul_gelu( | ||
| 148 | + x1.npu(), x2_nz, x1_scale.npu(), x2_scale.npu(), approximate="gelu_tanh") | ||
| 149 | + | ||
| 150 | + self.assertRtolEqual(cpu_result.numpy(), npu_result.cpu().float().numpy(), 0.01) | ||
| 151 | + | ||
| 152 | + | ||
| 153 | + | ||
| 154 | + def test_npu_quant_matmul_gelu_a8w8_nz_gelu_erf(self): | ||
| 155 | + """Test A8W8 with NZ format and gelu_erf""" | ||
| 156 | + torch.manual_seed(0) | ||
| 157 | + m, k, n = 128, 256, 512 | ||
| 158 | + | ||
| 159 | + x1 = torch.randint(-1, 1, (m, k), dtype=torch.int8) | ||
| 160 | + x2 = torch.randint(-1, 1, (k, n), dtype=torch.int8) | ||
| 161 | + x1_scale = torch.randn(m, dtype=torch.float32).abs() * 0.01 | ||
| 162 | + x2_scale = torch.randn(n, dtype=torch.float32).abs() * 0.01 | ||
| 163 | + | ||
| 164 | + # CPU reference | ||
| 165 | + cpu_result = self.cpu_quant_matmul_gelu(x1, x2, x1_scale, x2_scale, approximate="gelu_erf") | ||
| 166 | + | ||
| 167 | + # Convert x2 to NZ format | ||
| 168 | + x2_nz = torch_npu.npu_format_cast(x2.npu().contiguous(), 29) # 29 is ACL_FORMAT_FRACTAL_NZ | ||
| 169 | + | ||
| 170 | + # NPU custom op with NZ format | ||
| 171 | + npu_result = torch_npu.npu_quant_matmul_gelu( | ||
| 172 | + x1.npu(), x2_nz, x1_scale.npu(), x2_scale.npu(), approximate="gelu_erf") | ||
| 173 | + | ||
| 174 | + self.assertRtolEqual(cpu_result.numpy(), npu_result.cpu().float().numpy(), 0.01) | ||
| 175 | + | ||
| 176 | + | ||
| 177 | + | ||
| 178 | + def test_npu_quant_matmul_gelu_a8w8_nz_simple(self): | ||
| 179 | + """Test A8W8 with NZ format - simple case with smaller dimensions""" | ||
| 180 | + torch.manual_seed(0) | ||
| 181 | + m, k, n = 32, 64, 128 | ||
| 182 | + | ||
| 183 | + x1 = torch.randint(-1, 1, (m, k), dtype=torch.int8) | ||
| 184 | + x2 = torch.randint(-1, 1, (k, n), dtype=torch.int8) | ||
| 185 | + x1_scale = torch.randn(m, dtype=torch.float32).abs() * 0.01 | ||
| 186 | + x2_scale = torch.randn(n, dtype=torch.float32).abs() * 0.01 | ||
| 187 | + | ||
| 188 | + # CPU reference | ||
| 189 | + cpu_result = self.cpu_quant_matmul_gelu(x1, x2, x1_scale, x2_scale, approximate="gelu_tanh") | ||
| 190 | + | ||
| 191 | + # Convert x2 to NZ format | ||
| 192 | + x2_nz = torch_npu.npu_format_cast(x2.npu().contiguous(), 29) # 29 is ACL_FORMAT_FRACTAL_NZ | ||
| 193 | + | ||
| 194 | + # NPU custom op with NZ format | ||
| 195 | + npu_result = torch_npu.npu_quant_matmul_gelu( | ||
| 196 | + x1.npu(), x2_nz, x1_scale.npu(), x2_scale.npu(), approximate="gelu_tanh") | ||
| 197 | + | ||
| 198 | + self.assertRtolEqual(cpu_result.numpy(), npu_result.cpu().float().numpy(), 0.01) | ||
| 199 | + | ||
| 200 | + | ||
| 201 | + | ||
| 202 | + def test_npu_quant_matmul_gelu_a4w4_nd_gelu_tanh_int4(self): | ||
| 203 | + """Test A4W4 with ND format, gelu_tanh, using quint4x2 (direct INT4) type""" | ||
| 204 | + torch.manual_seed(0) | ||
| 205 | + m, k, n = 128, 256, 512 | ||
| 206 | + | ||
| 207 | + # Generate INT4 data using quint4x2 type | ||
| 208 | + x1_fp = torch.randn(m, k, dtype=torch.float32) | ||
| 209 | + x2_fp = torch.randn(k, n, dtype=torch.float32) | ||
| 210 | + | ||
| 211 | + # Quantize to INT4 (quint4x2) | ||
| 212 | + scale_tmp = torch.ones(1, dtype=torch.float32).npu() | ||
| 213 | + x1 = torch_npu.npu_quantize(x1_fp.npu(), scale_tmp, None, torch.quint4x2, -1, False) | ||
| 214 | + x2 = torch_npu.npu_quantize(x2_fp.npu(), scale_tmp, None, torch.quint4x2, -1, False) | ||
| 215 | + | ||
| 216 | + x1_scale = torch.randn(m, dtype=torch.float32).abs() * 0.01 | ||
| 217 | + x2_scale = torch.randn(n, dtype=torch.float32).abs() * 0.01 | ||
| 218 | + | ||
| 219 | + # NPU custom op | ||
| 220 | + npu_result = torch_npu.npu_quant_matmul_gelu( | ||
| 221 | + x1, x2, x1_scale.npu(), x2_scale.npu(), approximate="gelu_tanh") | ||
| 222 | + | ||
| 223 | + # For A4W4, we mainly test that the op runs without errors | ||
| 224 | + # Full accuracy comparison requires proper INT4 CPU implementation | ||
| 225 | + self.assertEqual(npu_result.shape, (m, n)) | ||
| 226 | + self.assertTrue(torch.isfinite(npu_result).all()) | ||
| 227 | + self.assertEqual(npu_result.dtype, torch.float16) | ||
| 228 | + | ||
| 229 | + | ||
| 230 | + | ||
| 231 | + def test_npu_quant_matmul_gelu_a4w4_nd_gelu_erf_int4(self): | ||
| 232 | + """Test A4W4 with ND format, gelu_erf, using quint4x2 (direct INT4) type""" | ||
| 233 | + torch.manual_seed(0) | ||
| 234 | + m, k, n = 64, 128, 256 | ||
| 235 | + | ||
| 236 | + # Generate INT4 data using quint4x2 type | ||
| 237 | + x1_fp = torch.randn(m, k, dtype=torch.float32) | ||
| 238 | + x2_fp = torch.randn(k, n, dtype=torch.float32) | ||
| 239 | + | ||
| 240 | + # Quantize to INT4 (quint4x2) | ||
| 241 | + scale_tmp = torch.ones(1, dtype=torch.float32).npu() | ||
| 242 | + x1 = torch_npu.npu_quantize(x1_fp.npu(), scale_tmp, None, torch.quint4x2, -1, False) | ||
| 243 | + x2 = torch_npu.npu_quantize(x2_fp.npu(), scale_tmp, None, torch.quint4x2, -1, False) | ||
| 244 | + | ||
| 245 | + x1_scale = torch.randn(m, dtype=torch.float32).abs() * 0.01 | ||
| 246 | + x2_scale = torch.randn(n, dtype=torch.float32).abs() * 0.01 | ||
| 247 | + | ||
| 248 | + # NPU custom op | ||
| 249 | + npu_result = torch_npu.npu_quant_matmul_gelu( | ||
| 250 | + x1, x2, x1_scale.npu(), x2_scale.npu(), approximate="gelu_erf") | ||
| 251 | + | ||
| 252 | + self.assertEqual(npu_result.shape, (m, n)) | ||
| 253 | + self.assertTrue(torch.isfinite(npu_result).all()) | ||
| 254 | + self.assertEqual(npu_result.dtype, torch.float16) | ||
| 255 | + | ||
| 256 | + | ||
| 257 | + | ||
| 258 | + def test_npu_quant_matmul_gelu_a4w4_nd_with_bias_int4(self): | ||
| 259 | + """Test A4W4 with ND format, gelu_tanh, bias, using quint4x2 (direct INT4) type""" | ||
| 260 | + torch.manual_seed(0) | ||
| 261 | + m, k, n = 64, 128, 256 | ||
| 262 | + | ||
| 263 | + # Generate INT4 data using quint4x2 type | ||
| 264 | + x1_fp = torch.randn(m, k, dtype=torch.float32) | ||
| 265 | + x2_fp = torch.randn(k, n, dtype=torch.float32) | ||
| 266 | + | ||
| 267 | + # Quantize to INT4 (quint4x2) | ||
| 268 | + scale_tmp = torch.ones(1, dtype=torch.float32).npu() | ||
| 269 | + x1 = torch_npu.npu_quantize(x1_fp.npu(), scale_tmp, None, torch.quint4x2, -1, False) | ||
| 270 | + x2 = torch_npu.npu_quantize(x2_fp.npu(), scale_tmp, None, torch.quint4x2, -1, False) | ||
| 271 | + | ||
| 272 | + x1_scale = torch.randn(m, dtype=torch.float32).abs() * 0.01 | ||
| 273 | + x2_scale = torch.randn(n, dtype=torch.float32).abs() * 0.01 | ||
| 274 | + bias = torch.randint(-5, 5, (n,), dtype=torch.int32) | ||
| 275 | + | ||
| 276 | + # NPU custom op | ||
| 277 | + npu_result = torch_npu.npu_quant_matmul_gelu( | ||
| 278 | + x1, x2, x1_scale.npu(), x2_scale.npu(), | ||
| 279 | + bias=bias.npu(), approximate="gelu_tanh") | ||
| 280 | + | ||
| 281 | + self.assertEqual(npu_result.shape, (m, n)) | ||
| 282 | + self.assertTrue(torch.isfinite(npu_result).all()) | ||
| 283 | + self.assertEqual(npu_result.dtype, torch.float16) | ||
| 284 | + | ||
| 285 | + | ||
| 286 | + | ||
| 287 | + def test_npu_quant_matmul_gelu_a4w4_nd_gelu_tanh_int32(self): | ||
| 288 | + """Test A4W4 with ND format, gelu_tanh, using int32 (packed INT4) type""" | ||
| 289 | + torch.manual_seed(0) | ||
| 290 | + m, k, n = 128, 256, 512 | ||
| 291 | + | ||
| 292 | + # Generate INT4 data packed in INT32 format (8 int4 values per int32) | ||
| 293 | + # Shape for int32: (m, k//8) and (k, n//8) | ||
| 294 | + k_packed = k // 8 | ||
| 295 | + n_packed = n // 8 | ||
| 296 | + | ||
| 297 | + x1 = torch.randint(-8, 8, (m, k_packed), dtype=torch.int32) | ||
| 298 | + x2 = torch.randint(-8, 8, (k, n_packed), dtype=torch.int32) | ||
| 299 | + | ||
| 300 | + x1_scale = torch.randn(m, dtype=torch.float32).abs() * 0.01 | ||
| 301 | + x2_scale = torch.randn(n, dtype=torch.float32).abs() * 0.01 | ||
| 302 | + | ||
| 303 | + # NPU custom op | ||
| 304 | + npu_result = torch_npu.npu_quant_matmul_gelu( | ||
| 305 | + x1.npu(), x2.npu(), x1_scale.npu(), x2_scale.npu(), approximate="gelu_tanh") | ||
| 306 | + | ||
| 307 | + # Verify output shape: (m, n) where n is recovered from n_packed * 8 | ||
| 308 | + self.assertEqual(npu_result.shape, (m, n)) | ||
| 309 | + self.assertTrue(torch.isfinite(npu_result).all()) | ||
| 310 | + self.assertEqual(npu_result.dtype, torch.float16) | ||
| 311 | + | ||
| 312 | + | ||
| 313 | + | ||
| 314 | + def test_npu_quant_matmul_gelu_a4w4_nd_gelu_erf_int32(self): | ||
| 315 | + """Test A4W4 with ND format, gelu_erf, using int32 (packed INT4) type""" | ||
| 316 | + torch.manual_seed(0) | ||
| 317 | + m, k, n = 64, 128, 256 | ||
| 318 | + | ||
| 319 | + # Generate INT4 data packed in INT32 format | ||
| 320 | + k_packed = k // 8 | ||
| 321 | + n_packed = n // 8 | ||
| 322 | + | ||
| 323 | + x1 = torch.randint(-8, 8, (m, k_packed), dtype=torch.int32) | ||
| 324 | + x2 = torch.randint(-8, 8, (k, n_packed), dtype=torch.int32) | ||
| 325 | + | ||
| 326 | + x1_scale = torch.randn(m, dtype=torch.float32).abs() * 0.01 | ||
| 327 | + x2_scale = torch.randn(n, dtype=torch.float32).abs() * 0.01 | ||
| 328 | + | ||
| 329 | + # NPU custom op | ||
| 330 | + npu_result = torch_npu.npu_quant_matmul_gelu( | ||
| 331 | + x1.npu(), x2.npu(), x1_scale.npu(), x2_scale.npu(), approximate="gelu_erf") | ||
| 332 | + | ||
| 333 | + self.assertEqual(npu_result.shape, (m, n)) | ||
| 334 | + self.assertTrue(torch.isfinite(npu_result).all()) | ||
| 335 | + self.assertEqual(npu_result.dtype, torch.float16) | ||
| 336 | + | ||
| 337 | + | ||
| 338 | + | ||
| 339 | + def test_npu_quant_matmul_gelu_a4w4_nd_with_bias_int32(self): | ||
| 340 | + """Test A4W4 with ND format, gelu_tanh, bias, using int32 (packed INT4) type""" | ||
| 341 | + torch.manual_seed(0) | ||
| 342 | + m, k, n = 64, 128, 256 | ||
| 343 | + | ||
| 344 | + # Generate INT4 data packed in INT32 format | ||
| 345 | + k_packed = k // 8 | ||
| 346 | + n_packed = n // 8 | ||
| 347 | + | ||
| 348 | + x1 = torch.randint(-8, 8, (m, k_packed), dtype=torch.int32) | ||
| 349 | + x2 = torch.randint(-8, 8, (k, n_packed), dtype=torch.int32) | ||
| 350 | + | ||
| 351 | + x1_scale = torch.randn(m, dtype=torch.float32).abs() * 0.01 | ||
| 352 | + x2_scale = torch.randn(n, dtype=torch.float32).abs() * 0.01 | ||
| 353 | + bias = torch.randint(-5, 5, (n,), dtype=torch.int32) | ||
| 354 | + | ||
| 355 | + # NPU custom op | ||
| 356 | + npu_result = torch_npu.npu_quant_matmul_gelu( | ||
| 357 | + x1.npu(), x2.npu(), x1_scale.npu(), x2_scale.npu(), | ||
| 358 | + bias=bias.npu(), approximate="gelu_tanh") | ||
| 359 | + | ||
| 360 | + self.assertEqual(npu_result.shape, (m, n)) | ||
| 361 | + self.assertTrue(torch.isfinite(npu_result).all()) | ||
| 362 | + self.assertEqual(npu_result.dtype, torch.float16) | ||
| 363 | + | ||
| 364 | + | ||
| 365 | + | ||
| 366 | + def test_npu_quant_matmul_gelu_bf16_output(self): | ||
| 367 | + """Test with BF16 output (x2_scale is BF16)""" | ||
| 368 | + torch.manual_seed(0) | ||
| 369 | + m, k, n = 64, 128, 256 | ||
| 370 | + | ||
| 371 | + x1 = torch.randint(-5, 5, (m, k), dtype=torch.int8) | ||
| 372 | + x2 = torch.randint(-5, 5, (k, n), dtype=torch.int8) | ||
| 373 | + x1_scale = torch.randn(m, dtype=torch.float32).abs() * 0.01 | ||
| 374 | + x2_scale = torch.randn(n, dtype=torch.bfloat16).abs() * 0.01 # BF16 scale | ||
| 375 | + | ||
| 376 | + # NPU custom op | ||
| 377 | + npu_result = torch_npu.npu_quant_matmul_gelu( | ||
| 378 | + x1.npu(), x2.npu(), x1_scale.npu(), x2_scale.npu(), approximate="gelu_tanh") | ||
| 379 | + | ||
| 380 | + # Check output dtype is BF16 | ||
| 381 | + self.assertEqual(npu_result.dtype, torch.bfloat16) | ||
| 382 | + self.assertEqual(npu_result.shape, (m, n)) | ||
| 383 | + self.assertTrue(torch.isfinite(npu_result.float()).all()) | ||
| 384 | + | ||
| 385 | + | ||
| 386 | + | ||
| 387 | + def test_npu_quant_matmul_gelu_batch(self): | ||
| 388 | + """Test with batch dimensions""" | ||
| 389 | + torch.manual_seed(0) | ||
| 390 | + batch, m, k, n = 4, 64, 128, 256 | ||
| 391 | + | ||
| 392 | + x1 = torch.randint(-5, 5, (batch, m, k), dtype=torch.int8) | ||
| 393 | + x2 = torch.randint(-5, 5, (batch, k, n), dtype=torch.int8) | ||
| 394 | + # x1_scale: per-token scale, shape should be (m,) for all batches or (batch*m,) for per-batch-per-token | ||
| 395 | + # For simplicity, use same scale for all batches: (m,) | ||
| 396 | + x1_scale = torch.randn(m, dtype=torch.float32).abs() * 0.01 | ||
| 397 | + x2_scale = torch.randn(n, dtype=torch.float32).abs() * 0.01 | ||
| 398 | + | ||
| 399 | + # Process each batch separately for CPU reference | ||
| 400 | + cpu_results = [] | ||
| 401 | + for i in range(batch): | ||
| 402 | + cpu_result = self.cpu_quant_matmul_gelu( | ||
| 403 | + x1[i], x2[i], x1_scale, x2_scale, approximate="gelu_tanh") | ||
| 404 | + cpu_results.append(cpu_result) | ||
| 405 | + cpu_result = torch.stack(cpu_results, dim=0) | ||
| 406 | + | ||
| 407 | + # NPU custom op should handle batch automatically | ||
| 408 | + npu_result = torch_npu.npu_quant_matmul_gelu( | ||
| 409 | + x1.npu(), x2.npu(), | ||
| 410 | + x1_scale.npu(), # Per-token scale: (m,) | ||
| 411 | + x2_scale.npu(), | ||
| 412 | + approximate="gelu_tanh") | ||
| 413 | + | ||
| 414 | + self.assertEqual(npu_result.shape, (batch, m, n)) | ||
| 415 | + self.assertTrue(torch.isfinite(npu_result.float()).all()) | ||
| 416 | + | ||
| 417 | + | ||
| 418 | + | ||
| 419 | + def test_npu_quant_matmul_gelu_invalid_approximate(self): | ||
| 420 | + """Test with invalid approximate values""" | ||
| 421 | + m, k, n = 64, 128, 256 | ||
| 422 | + | ||
| 423 | + x1 = torch.randint(-5, 5, (m, k), dtype=torch.int8) | ||
| 424 | + x2 = torch.randint(-5, 5, (k, n), dtype=torch.int8) | ||
| 425 | + x1_scale = torch.randn(m, dtype=torch.float32).abs() * 0.01 | ||
| 426 | + x2_scale = torch.randn(n, dtype=torch.float32).abs() * 0.01 | ||
| 427 | + | ||
| 428 | + # Test various invalid approximate values | ||
| 429 | + invalid_values = ["invalid_type", "gelu", "tanh", "erf", "gelu_relu", "", "none"] | ||
| 430 | + for invalid_val in invalid_values: | ||
| 431 | + with self.assertRaisesRegex(RuntimeError, "approximate must be 'gelu_tanh' or 'gelu_erf'"): | ||
| 432 | + torch_npu.npu_quant_matmul_gelu( | ||
| 433 | + x1.npu(), x2.npu(), x1_scale.npu(), x2_scale.npu(), approximate=invalid_val) | ||
| 434 | + | ||
| 435 | + | ||
| 436 | + | ||
| 437 | + def test_npu_quant_matmul_gelu_default_approximate(self): | ||
| 438 | + """Test with default approximate (gelu_erf)""" | ||
| 439 | + torch.manual_seed(0) | ||
| 440 | + m, k, n = 64, 128, 256 | ||
| 441 | + | ||
| 442 | + x1 = torch.randint(-5, 5, (m, k), dtype=torch.int8) | ||
| 443 | + x2 = torch.randint(-5, 5, (k, n), dtype=torch.int8) | ||
| 444 | + x1_scale = torch.randn(m, dtype=torch.float32).abs() * 0.01 | ||
| 445 | + x2_scale = torch.randn(n, dtype=torch.float32).abs() * 0.01 | ||
| 446 | + | ||
| 447 | + # CPU reference with default gelu_erf | ||
| 448 | + cpu_result = self.cpu_quant_matmul_gelu(x1, x2, x1_scale, x2_scale, approximate="gelu_erf") | ||
| 449 | + | ||
| 450 | + # NPU custom op without specifying approximate (should use default "gelu_erf") | ||
| 451 | + npu_result = torch_npu.npu_quant_matmul_gelu( | ||
| 452 | + x1.npu(), x2.npu(), x1_scale.npu(), x2_scale.npu()) | ||
| 453 | + | ||
| 454 | + self.assertRtolEqual(cpu_result.numpy(), npu_result.cpu().float().numpy(), 0.01) | ||
| 455 | + | ||
| 456 | + | ||
| 457 | +if __name__ == "__main__": | ||
| 458 | + run_tests() | ||
增加跳过
@unittest.skip("Skipping due to outdated CANN version; please update CANN to the latest version and remove this skip")