# npu_quant

**Repository Path**: libarry123/npu_quant

## Basic Information

- **Project Name**: npu_quant
- **Description**: No description available
- **Primary Language**: Unknown
- **License**: Not specified
- **Default Branch**: master
- **Homepage**: None
- **GVP Project**: No

## Statistics

- **Stars**: 0
- **Forks**: 0
- **Created**: 2025-06-08
- **Last Updated**: 2025-06-19

## Categories & Tags

**Categories**: Uncategorized

**Tags**: None

## README

# NPU 量化算子功能说明

本文档旨在说明 `quant_op.py` 文件中主要 NPU 量化算子的功能和区别。

## 核心矩阵乘法算子对比

`npu_quant_matmul` 和 `npu_weight_quant_batchmatmul` 是两个核心的量化矩阵乘法算子，其主要区别在于对输入张量的量化要求不同，导致了它们在功能和使用场景上的差异。

### `npu_quant_matmul` vs `npu_weight_quant_batchmatmul`

| 特性 | `npu_quant_matmul` | `npu_weight_quant_batchmatmul` |
| :--- | :--- | :--- |
| **核心功能** | 两个**量化**矩阵相乘 | **浮点**矩阵与**量化**矩阵相乘 |
| **输入 `x` (激活)** | 量化类型 (INT8) | 浮点类型 (FP16/BF16) |
| **输入 `weight` (权重)** | 量化类型 (INT8) | 量化类型 (INT8) |
| **内部操作** | 量化乘法 -> (反)量化输出 | 权重反量化 -> 浮点乘法 -> (可选)量化输出 |
| **主要场景** | 全量化模型 (激活和权重均量化) | 仅权重(Weight-Only)量化模型 |

**简单总结**:

*   如果激活值**是**量化的，用 `npu_quant_matmul`。
*   如果激活值**不是**量化的（即浮点），而权重是量化的，用 `npu_weight_quant_batchmatmul`。

---

## 各算子功能详解

### 1. `npu_quant_matmul`

*   **主要功能**: 执行两个**量化后**的矩阵（激活和权重）之间的乘法运算。
*   **详细说明**: 这是一个纯粹的"量化矩阵乘"算子，执行的是 `INT8 x INT8` 的计算。它接收已经量化好的 `x1` 和 `x2` 张量，利用 `scale`、`offset` 等参数在计算后进行反量化缩放。适用于模型的激活和权重都已被量化的**全量化**场景。

### 2. `npu_weight_quant_batchmatmul`

*   **主要功能**: 执行**浮点**激活矩阵和**量化**权重矩阵之间的乘法运算。
*   **详细说明**: 这是一个混合精度计算的融合算子。它内部会先对量化的 `weight` 张量进行反量化（恢复到浮点），然后与浮点类型的 `x` 张量进行浮点矩阵乘法。主要用于**仅权重(Weight-Only)量化**的场景，在保持激活值精度的同时，通过量化权重来减小模型体积和内存占用。

### 3. `npu_dynamic_quant`

*   **主要功能**: 对输入张量进行**对称**动态量化。
*   **详细说明**: 该算子自动计算输入张量（通常是激活值）每个 "token" 的最大绝对值，并以此为依据计算出缩放系数 `scale`，然后将浮点数据量化为整数类型。因对称量化的零点固定为0，故只输出 `scale`。

### 4. `npu_dynamic_quant_asymmetric`

*   **主要功能**: 对输入张量进行**非对称**动态量化。
*   **详细说明**: 与对称量化不同，此算子会为每个 "token" 计算最大值和最小值，从而确定 `scale` 和 `offset`（零点）两个参数。这使得量化能更好地适应数据分布不均匀的情况。它会同时输出量化后的张量、`scale` 和 `offset`。

### 5. `npu_anti_quant`

*   **主要功能**: 执行**反量化**（Dequantization）。
*   **详细说明**: 量化操作的逆过程。它接收一个量化后的整数张量，并利用给定的 `scale` 和 `offset` 将其恢复为浮点张量。

### 6. `npu_quantize`

*   **主要功能**: 通用的静态量化处理函数。
*   **详细说明**: 与 `dynamic` 系列算子不同，这个算子不会自动计算量化参数。用户必须**显式地提供** `scales` 和 `zero_points`。它根据这些给定的参数来量化输入张量。

### 7. `npu_trans_quant_param`

*   **主要功能**: 转换量化参数的数据类型。
*   **详细说明**: 这是一个辅助函数，主要作用是将 `float32` 类型的量化参数 `scale` 和 `offset` 转换为 `int64` 等特定类型。在某些图模式下的高性能计算场景中，传入转换后的参数可以获得更好的性能。

### 8. `npu_quant_scatter` 和 `npu_quant_scatter_`

*   **主要功能**: 融合了量化和更新（scatter）操作。
*   **详细说明**: 此算子会先对 `updates` 张量进行量化，然后根据 `indices` 索引将量化后的值更新到 `input` 张量中。`npu_quant_scatter_` 是其 in-place 版本，会直接修改 `input` 张量。

### 9. `npu_kronecker_quant`

*   **主要功能**: 融合了两次矩阵乘法和量化。
*   **详细说明**: 这是一个专用的融合算子，它对输入 `x` 依次进行两次矩阵乘法（与 `kronecker_p1` 和 `kronecker_p2`），然后对最终结果执行量化。用于特定的计算模式优化。

### 10. `npu_convert_weight_to_int4pack`

*   **主要功能**: 将权重张量打包为 INT4 格式。
*   **详细说明**: 这是一个数据格式转换算子，它将 `int32` 类型的输入张量打包为 `int4` 存储格式。每8个 `int4` 数据通过一个 `int32` 数据承载，并进行交叠排放。这种打包格式可以有效减少内存占用，同时为后续的量化矩阵运算提供高效的数据排布。主要用于 **W4A4（4位权重4位激活）量化** 场景中的权重预处理。

### 11. `npu_grouped_matmul_swiglu_quant`

*   **主要功能**: 为 `deepseek` 模型定制的**超级融合算子**。
*   **详细说明**: 它将分组矩阵乘（Grouped Matmul）、反量化、SwiGLU激活函数和最终的量化等多个步骤融合在了一个算子中，旨在最大化特定大语言模型在 NPU 上的运行性能。

---

## W8A8 和 W4A4 量化线性层实现

本项目还包含了两种主要的量化线性层实现，它们基于上述 NPU 量化算子构建，为实际的模型量化应用提供了高层封装。

### W8A8 量化线性层 (`w8a8_linear.py`)

**W8A8** 表示 **8位权重 + 8位激活** 的量化方案，这是一种相对保守但稳定的量化策略。

#### 核心特性

*   **量化精度**: 权重和激活值均使用 8 位整数表示
*   **NPU 实现**: `NPUQuantLinear` - 使用 `npu_quant_matmul` 算子
*   **参考实现**: `GoldenQuantLinear` - CPU 上的纯 PyTorch 实现
*   **适用场景**: 对精度要求较高，但仍需要内存和计算优化的场景

#### 实现细节

*   激活值和权重在计算前都需要进行量化
*   支持非对称量化（包含 offset/zero_point）
*   自动处理类型约束，确保 NPU 算子的输入类型要求
*   提供完整的前向传播流程：浮点输入 → 量化 → 矩阵乘 → 反量化 → 浮点输出

### W4A4 量化线性层 (`w4a4_linear.py`)

**W4A4** 表示 **4位权重 + 4位激活** 的极致量化方案，追求最大的内存和计算效率。

#### 核心特性

*   **量化精度**: 权重和激活值均使用 4 位整数表示
*   **NPU 实现**: `NPUQuantLinearW4A4` - 使用 `npu_kronecker_quant` 和 `npu_anti_quant` 算子
*   **参考实现**: `GoldenQuantLinearW4A4` - CPU 上的模拟实现
*   **数据打包**: 使用 `npu_convert_weight_to_int4pack` 进行权重预处理
*   **适用场景**: 资源极度受限的边缘设备或需要最大吞吐量的推理场景

#### 实现细节

*   采用创新的 **Kronecker 量化** 方法，通过两个小矩阵的乘积来近似量化过程
*   权重使用 INT4 打包格式（8个 INT4 值存储在1个 INT32 中）
*   支持 `clip_ratio` 参数来控制量化的激进程度
*   动态计算激活值的量化参数，无需预先校准

### 实现对比

| 特性 | W8A8 | W4A4 |
|------|------|------|
| **内存占用** | 中等 (8位) | 极低 (4位) |
| **计算精度** | 较高 | 较低 |
| **实现复杂度** | 简单 | 复杂 |
| **NPU 算子** | `npu_quant_matmul` | `npu_kronecker_quant` + `npu_anti_quant` |
| **权重格式** | INT8 直接存储 | INT4 打包存储 |
| **量化方式** | 传统对称/非对称量化 | Kronecker 量化 |
| **适用模型** | 通用神经网络 | 对精度要求相对宽松的大模型 |

### 使用建议

1. **精度优先场景**: 选择 W8A8，它在精度和效率之间提供了良好的平衡
2. **效率优先场景**: 选择 W4A4，特别适合大规模语言模型的推理加速  
3. **开发调试**: 两种实现都提供了 Golden 参考版本，便于验证 NPU 实现的正确性
4. **类型安全**: 所有实现都包含自动类型转换机制，确保 NPU 算子的类型要求得到满足