torch_npu.npu_quant_matmul

功能描述

完成量化的矩阵乘计算，最小支持输入维度为2维，最大支持输入维度为3维。

接口原型

npu_quant_matmul(Tensor x1, Tensor x2, Tensor scale, Tensor? offset=None, Tensor? bias=None) -> Tensor

参数说明

• x1（计算输入）：Tensor类型，数据类型支持INT8。数据格式支持ND，shape最少是2维，最多是3维。
• x2（计算输入）：Tensor类型，数据类型支持INT8。数据格式支持ND，shape最少是2维，最多是3维。
• scale（计算输入）：Tensor类型，数据类型支持Float32。数据格式支持ND，shape是1维（t，），t = 1或n，其中n与x2的n一致。
• offset（ 计算输入）：可选参数。Tensor类型，数据类型支持Float32，数据格式支持ND，shape是1维（t，），t = 1或n，其中n与x2的n一致。
• bias（ 计算输入）：可选参数。Tensor类型，数据类型支持INT32，数据格式支持ND，shape支持1维（n，）或3维（batch,1,n），n与x2的n一致, batch与x1，x2中的batch一致。

输出说明

一个Tensor类型的输出，代表quant_matmul的计算结果。

约束说明

• 传入的x1、x2、bias、scale或out不能是空指针。
• x1、x2、bias、scale、offset或out的数据类型和数据格式需要在支持的范围之内。
• x1、x2的shape需要在在2-3维范围。
• offset、scale的shape需要为1维（t，），t = 1或n。
• bias的shape支持1维（n，）或3维（batch,1,n），n与x2的n一致, batch与x1，x2中的batch一致。

支持的PyTorch版本

• PyTorch 2.1
• PyTorch 2.0
• PyTorch 1.11.0

支持的型号

Atlas A2 训练系列产品

调用示例

单算子调用：
import torch
import torch_npu
import logging
import os

cpu_x1 = torch.randint(-5, 5, (1, 1, 256, 768), dtype=torch.int8)
cpu_x2 = torch.randint(-5, 5, (2, 31, 768, 16), dtype=torch.int8)
scale = torch.randn(16, dtype=torch.float32)
offset = torch.randn(16, dtype=torch.float32)
bias = torch.randint(-5, 5, (31 * 2, 1, 16), dtype=torch.int32)
npu_out = torch_npu.npu_quant_matmul(cpu_x1.npu(), cpu_x2.npu(), scale.npu(), offset.npu(), bias.npu())

图模式：
import torch
import torch_npu
import torchair as tng
from torchair.ge_concrete_graph import ge_apis as ge
from torchair.configs.compiler_config import CompilerConfig
import logging
from torchair.core.utils import logger
logger.setLevel(logging.DEBUG)
import os
import numpy as np
os.environ["ENABLE_ACLNN"] = "true"
config = CompilerConfig()
model = cpu_model.npu()
npu_backend = tng.get_npu_backend(compiler_config=config)

class MyModel(torch.nn.Module):
    def __init__(self):
        super().__init__()
    def forward(self, x1, x2, scale, offset, bias):
        scale_1 = torch_npu.npu_trans_quant_param(scale, offset)
        return torch_npu.npu_quant_matmul(x1, x2, scale_1, offset, bias)
cpu_model = MyModel()
cpu_x1 = torch.randint(-1, 1, (15, 1, 512), dtype=torch.int8)
cpu_x2 = torch.randint(-1, 1, (15, 512, 128), dtype=torch.int8)
scale = torch.randn(1, dtype=torch.float32)
offset = torch.randn(1, dtype=torch.float32)
bias = torch.randint(-1,1, (15, 1, 128), dtype=torch.int32)
model = cpu_model.npu()
model = torch.compile(cpu_model, backend=npu_backend, dynamic=True)
npu_out = model(cpu_x1.npu(), cpu_x2.npu(), scale.npu(), offset.npu(), bias.npu())