Star1388
2139
Code
Code
Issues837
Pull Requests675
Actions
Discussion
Wiki
Analysis
Star1388
2139
Merged
test(rms_norm_quant): 补充 Ascend 950 golden 测试资产 #11974
raoliang_sacAdd on 24 days ago
test(rms_norm_quant): 补充 Ascend 950 golden 测试资产 #11974
Merged
raoliang_sacAdd on 24 days ago
Merged fromraoliang_sac/ops-nntest/rms_norm_quant-golden-assetsintocann/ops-nnmaster
Total: 5 files changed+983-1
Anorm/rms_norm_quant/tests/assets/golden.py+217-0
@@ -0,0 +1,217 @@
1+#!/usr/bin/env python3
2+# -*- coding: UTF-8 -*-
3+# Copyright (c) 2026 Huawei Technologies Co., Ltd.
4+# This program is free software, you can redistribute it and/or modify it under the terms and conditions of
5+# CANN Open Software License Agreement Version 2.0 (the "License").
6+# Please refer to the License for details. You may not use this file except in compliance with the License.
7+# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
8+# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
9+# See LICENSE in the root of the software repository for the full text of the License.
10+ 
11+import numpy as np
12+import torch
13+ 
14+ 
15+__golden__ = {"aclnn": {"aclnnRmsNormQuant": "aclnn_rms_norm_quant_golden"}}
16+ 
17+_DST_SAT_MAX = {
18+ "int8": 127.0,
19+ "int4": 7.0,
20+ "float8_e5m2": 57344.0,
21+ "float8_e4m3fn": 448.0,
22+ "hifloat8": 32768.0,
23+}
24+_DST_SAT_MIN = {
25+ "int8": -128.0,
26+ "int4": -8.0,
27+ "float8_e5m2": -57344.0,
28+ "float8_e4m3fn": -448.0,
29+ "hifloat8": -32768.0,
30+}
31+_INT4_VALUES_PER_BYTE = 2
32+_INT4_VALUES_PER_INT32 = np.dtype(np.int32).itemsize * _INT4_VALUES_PER_BYTE
33+ 
34+ 
35+def _to_numpy(t):
36+ if t is None:
37+ return None
38+ if isinstance(t, np.ndarray):
39+ return t
40+ if hasattr(t, "numpy"):
41+ try:
42+ return t.numpy()
43+ except (TypeError, RuntimeError):
44+ return t.float().numpy()
45+ if hasattr(t, "detach"):
46+ try:
47+ return t.detach().cpu().numpy()
48+ except (TypeError, RuntimeError):
49+ return t.detach().cpu().float().numpy()
50+ return np.array(t)
51+ 
52+ 
53+def _np_to_torch(arr):
54+ if arr is None:
55+ return None
56+ if hasattr(arr, "view") and hasattr(arr, "dim"):
57+ return arr
58+ if isinstance(arr, np.ndarray):
59+ dt = str(arr.dtype)
60+ if np.issubdtype(arr.dtype, np.integer):
61+ return torch.from_numpy(arr)
62+ if "float8_e5m2" in dt:
63+ return torch.from_numpy(arr.view(np.uint8)).view(torch.float8_e5m2)
64+ if "float8_e4m3" in dt:
65+ return torch.from_numpy(arr.view(np.uint8)).view(torch.float8_e4m3fn)
66+ if "hifloat8" in dt:
67+ return torch.from_numpy(arr.view(np.int8)).view(torch.int8)
68+ return torch.from_numpy(arr.astype(np.float32))
69+ return arr
70+ 
71+ 
72+def _quant_to_dtype(y_quant_np, dtype_str):
73+ if dtype_str not in ("int8", "hifloat8", "float8_e5m2", "float8_e4m3fn", "int4"):
74+ return y_quant_np
75+ 
76+ y_quant_np = np.nan_to_num(
77+ y_quant_np,
78+ nan=0.0,
79+ posinf=_DST_SAT_MAX.get(dtype_str, 0.0),
80+ neginf=_DST_SAT_MIN.get(dtype_str, 0.0),
81+ )
82+ 
83+ if dtype_str in ("int8", "int4"):
84+ y_quant_np = np.round(y_quant_np, 0)
85+ 
86+ if dtype_str in _DST_SAT_MAX:
87+ y_quant_np = np.clip(
88+ y_quant_np, _DST_SAT_MIN[dtype_str], _DST_SAT_MAX[dtype_str]
89+ )
90+ 
91+ if dtype_str == "int8":
92+ return y_quant_np.astype(np.int8)
93+ if dtype_str == "int4":
94+ from ttk.utilities.dtypes import numpy_int4
95+ 
96+ y_quant_np = np.clip(y_quant_np, -8, 7)
97+ return y_quant_np.astype(numpy_int4(), copy=False)
98+ if dtype_str == "hifloat8":
99+ from en_dtypes import hifloat8
100+ 
101+ return y_quant_np.astype(hifloat8, copy=False)
102+ if dtype_str == "float8_e5m2":
103+ from ml_dtypes import float8_e5m2
104+ 
105+ return y_quant_np.astype(float8_e5m2, copy=False)
106+ if dtype_str == "float8_e4m3fn":
107+ from ml_dtypes import float8_e4m3fn
108+ 
109+ return y_quant_np.astype(float8_e4m3fn, copy=False)
110+ return y_quant_np
111+ 
112+ 
113+def _pack_int4_to_int32(y_int4):
114+ """Pack the public ACLNN INT4 result into its int32 storage ABI."""
115+ if y_int4.ndim == 0 or y_int4.shape[-1] % _INT4_VALUES_PER_INT32 != 0:
116+ raise ValueError(
117+ f"INT4 last dimension must be divisible by {_INT4_VALUES_PER_INT32}, "
118+ f"but got shape {y_int4.shape}"
119+ )
120+ from ttk.utilities.dtypes import pack_4bits
121+ 
122+ packed_bytes = pack_4bits(np.ascontiguousarray(y_int4))
123+ packed_shape = y_int4.shape[:-1] + (y_int4.shape[-1] // _INT4_VALUES_PER_INT32,)
124+ return packed_bytes.view(np.int32).reshape(packed_shape)
125+ 
126+ 
127+def aclnn_rms_norm_quant_golden(
128+ x, gamma, beta, scale, offset, epsilon=1e-6, y=None, dst_type=2, **kwargs
129+):
130+ x = _to_numpy(x)
131+ gamma = _to_numpy(gamma)
132+ beta = _to_numpy(beta)
133+ scale = _to_numpy(scale)
134+ offset = _to_numpy(offset)
135+ 
136+ if scale is not None and scale.dtype != np.float32:
137+ scale = scale.astype(np.float32)
138+ if offset is not None and offset.dtype != np.float32:
139+ offset = offset.astype(np.float32)
140+ 
141+ x_fp32 = x.astype(np.float32)
142+ 
143+ # ACLNN squeezes a legal [1, H] gamma to [H] before last-axis RMSNorm.
144+ n = len(x.shape) - 1
145+ gamma_flat = gamma.reshape(-1)
146+ x_reshaped = x_fp32.reshape(x_fp32.shape[:n] + gamma_flat.shape)
147+ 
148+ variance = np.mean(np.power(x_reshaped, 2), axis=n, keepdims=True)
149+ std = np.sqrt(variance + epsilon)
150+ rstd = 1 / std
151+ 
152+ result_mid = x_reshaped * rstd
153+ 
154+ gamma_fp32 = gamma.astype(np.float32) if gamma.dtype != np.float32 else gamma
155+ y_array = result_mid * gamma_fp32.reshape(gamma_flat.shape)
156+ 
157+ if beta is not None:
158+ beta_reshaped = beta.reshape(gamma_flat.shape)
159+ y_array = y_array + beta_reshaped.astype(np.float32)
160+ 
161+ y_dtype_str = "int8"
162+ if y is not None:
163+ dt_str = str(y.dtype)
164+ if "int8" in dt_str:
165+ y_dtype_str = "int8"
166+ elif "hifloat8" in dt_str:
167+ y_dtype_str = "hifloat8"
168+ elif "e5m2" in dt_str:
169+ y_dtype_str = "float8_e5m2"
170+ elif "e4m3" in dt_str:
171+ y_dtype_str = "float8_e4m3fn"
172+ elif "int4" in dt_str or "int32" in dt_str:
173+ y_dtype_str = "int4"
174+ else:
175+ DST_TYPE_MAP = {
176+ 2: "int8",
177+ 29: "int4",
178+ 34: "hifloat8",
179+ 35: "float8_e5m2",
180+ 36: "float8_e4m3fn",
181+ }
182+ y_dtype_str = DST_TYPE_MAP.get(dst_type, "int8")
183+ 
184+ y_quant = (
185+ y_array / scale.reshape(-1).reshape(scale.shape)
186+ if scale is not None
187+ else y_array
188+ )
189+ if offset is not None:
190+ y_quant = y_quant + offset.reshape(-1).reshape(offset.shape)
191+ 
192+ y_result = _quant_to_dtype(y_quant, y_dtype_str)
193+ if y_dtype_str == "int4":
194+ y_result = _pack_int4_to_int32(y_result)
195+ return _np_to_torch(y_result)
196+ 
197+ 
198+class AclnnRmsNormQuantTestSpec:
199+ """量化输出按 1 LSB 判（standard='quant'，判据实现见 comparison/quant.py）。
200+ 
201+ y = round(x / scale) 的 round 是阶跃函数：商落在 .5 附近时，浮点末位 1 ulp 的差异
202+ 就足以让结果跳到相邻整数格。而末位差异来自归约次序、FMA 融合、硬件舍入等实现自由度，
203+ 不同实现之间必然存在。要求逐位相等等于要求复现某一个实现的浮点运算次序，
204+ 故按绝对误差 <= 1 判定。
205+ """
206+ 
207+ golden = staticmethod(aclnn_rms_norm_quant_golden)
208+ tolerance = {
209+ "int8": {"standard": "quant"},
210+ "int4": {"standard": "quant"},
211+ "int32": {"standard": "quant", "packed_bits": 4},
212+ "float8_e5m2": {"standard": "requant"},
213+ "float8_e4m3fn": {"standard": "requant"},
214+ }
215+ 
216+ 
217+__spec__ = {"aclnnRmsNormQuant": "AclnnRmsNormQuantTestSpec"}
Mnorm/rms_norm_quant_v2/README.md+1-1
@@ -157,4 +157,4 @@
157| 调用方式 | 样例代码 | 说明 |157| 调用方式 | 样例代码 | 说明 |
158| ---------------- | --------------------------- | --------------------------------------------------- |158| ---------------- | --------------------------- | --------------------------------------------------- |
159| aclnn接口 | [test_aclnn_rms_norm_quant](../rms_norm_quant/examples/test_aclnn_rms_norm_quant.cpp) | 通过[aclnnRmsNormQuant](../rms_norm_quant/docs/aclnnRmsNormQuant.md)接口方式调用RmsNormQuantV2算子。 |159| aclnn接口 | [test_aclnn_rms_norm_quant](../rms_norm_quant/examples/test_aclnn_rms_norm_quant.cpp) | 通过[aclnnRmsNormQuant](../rms_norm_quant/docs/aclnnRmsNormQuant.md)接口方式调用RmsNormQuantV2算子。 |
160-| 图模式 | - | 通过[算子IR](op_graph/rms_norm_quant_v2_proto.h)构图方式调用RmsNormQuantV2算子。 |160+| 图模式 | [test_geir_rms_norm_quant_v2](examples/test_geir_rms_norm_quant_v2.cpp) | 通过[算子IR](op_graph/rms_norm_quant_v2_proto.h)构图方式调用RmsNormQuantV2算子。 |
Anorm/rms_norm_quant_v2/examples/test_geir_rms_norm_quant_v2.cpp+131-0
@@ -0,0 +1,131 @@
1+/**
2+ * Copyright (c) 2026 Huawei Technologies Co., Ltd.
3+ * This program is free software, you can redistribute it and/or modify it under the terms and conditions of
4+ * CANN Open Software License Agreement Version 2.0 (the "License").
5+ * Please refer to the License for details. You may not use this file except in compliance with the License.
6+ * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
7+ * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
8+ * See LICENSE in the root of the software repository for the full text of the License.
9+ */
10+ 
11+/**
12+ * @file test_geir_rms_norm_quant_v2.cpp
13+ * @brief RmsNormQuantV2 GE graph example with two INT8 outputs on Ascend 950.
14+ */
15+ 
16+#include <cstdint>
17+#include <cstdio>
18+#include <map>
19+#include <utility>
20+#include <vector>
21+ 
22+#include "ge_api.h"
23+#include "graph.h"
24+#include "ops_proto_legacy.h"
25+#include "tensor.h"
26+ 
27+#include "../op_graph/rms_norm_quant_v2_proto.h"
28+ 
29+namespace {
30+constexpr int32_t SUCCESS_CODE = 0;
31+constexpr int32_t FAILED_CODE = 1;
32+ 
33+ge::Tensor MakeTensor(const std::vector<int64_t>& shape, const std::vector<float>& values)
34+{
35+ ge::TensorDesc desc(ge::Shape(shape), ge::FORMAT_ND, ge::DT_FLOAT);
36+ desc.SetPlacement(ge::kPlacementHost);
37+ const auto* bytes = static_cast<const uint8_t*>(static_cast<const void*>(values.data()));
38+ return ge::Tensor(desc, bytes, values.size() * sizeof(float));
39+}
40+ 
41+bool CheckOutput(const ge::Tensor& tensor, const std::vector<int64_t>& shape, uint8_t expected)
42+{
43+ const auto& desc = tensor.GetTensorDesc();
44+ if (desc.GetShape().GetDims() != shape || desc.GetDataType() != ge::DT_INT8 || tensor.GetData() == nullptr ||
45+ tensor.GetSize() != static_cast<size_t>(desc.GetShape().GetShapeSize())) {
46+ std::printf("Unexpected output shape, type or data size\n");
47+ return false;
48+ }
49+ for (size_t index = 0; index < tensor.GetSize(); ++index) {
50+ if (tensor.GetData()[index] != expected) {
51+ std::printf("Output mismatch at %zu: got %u, expected %u\n", index,
52+ static_cast<unsigned int>(tensor.GetData()[index]), static_cast<unsigned int>(expected));
53+ return false;
54+ }
55+ }
56+ return true;
57+}
58+ 
59+int32_t RunExample()
60+{
61+ const std::vector<int64_t> xShape = {2, 16};
62+ // x and gamma are all ones. Division by scales1=0.5 and scales2=0.25 gives y1=2 and y2=4.
63+ const std::vector<ge::Tensor> inputs = {MakeTensor(xShape, std::vector<float>(32, 1.0f)),
64+ MakeTensor({16}, std::vector<float>(16, 1.0f)), MakeTensor({1}, {0.5f}),
65+ MakeTensor({1}, {0.25f})};
66+ std::vector<ge::Operator> dataNodes;
67+ const char* const inputNames[] = {"x", "gamma", "scales1", "scales2"};
68+ ge::Graph graph("rms_norm_quant_v2_example");
69+ for (size_t index = 0; index < inputs.size(); ++index) {
70+ auto data = ge::op::Data(inputNames[index]).set_attr_index(static_cast<int64_t>(index));
71+ const auto& desc = inputs[index].GetTensorDesc();
72+ if (data.update_input_desc_x(desc) != ge::GRAPH_SUCCESS ||
73+ data.update_output_desc_y(desc) != ge::GRAPH_SUCCESS || graph.AddOp(data) != ge::GRAPH_SUCCESS) {
74+ std::printf("Failed to create input %s\n", inputNames[index]);
75+ return FAILED_CODE;
76+ }
77+ dataNodes.push_back(data);
78+ }
79+ 
80+ auto quant = ge::op::RmsNormQuantV2("rms_norm_quant_v2");
81+ quant.set_input_x(dataNodes[0])
82+ .set_input_gamma(dataNodes[1])
83+ .set_input_scales1(dataNodes[2])
84+ .set_input_scales2(dataNodes[3])
85+ .set_attr_epsilon(1.0e-6f)
86+ .set_attr_div_mode(true)
87+ .set_attr_dst_type(ge::DT_INT8);
88+ // zero_points1, zero_points2 and beta are omitted. Both outputs are valid because scales2 is present.
89+ const ge::TensorDesc outputDesc(ge::Shape(xShape), ge::FORMAT_ND, ge::DT_INT8);
90+ if (quant.update_output_desc_y1(outputDesc) != ge::GRAPH_SUCCESS ||
91+ quant.update_output_desc_y2(outputDesc) != ge::GRAPH_SUCCESS || graph.AddOp(quant) != ge::GRAPH_SUCCESS) {
92+ std::printf("Failed to create RmsNormQuantV2 node\n");
93+ return FAILED_CODE;
94+ }
95+ graph.SetInputs(dataNodes).SetOutputs(std::vector<std::pair<ge::Operator, std::vector<size_t>>>{{quant, {0, 1}}});
96+ 
97+ const std::map<ge::AscendString, ge::AscendString> sessionOptions;
98+ ge::Session session(sessionOptions);
99+ if (session.AddGraph(0, graph) != ge::SUCCESS) {
100+ std::printf("AddGraph failed\n");
101+ return FAILED_CODE;
102+ }
103+ std::vector<ge::Tensor> outputs;
104+ if (session.RunGraph(0, inputs, outputs) != ge::SUCCESS) {
105+ std::printf("RunGraph failed\n");
106+ return FAILED_CODE;
107+ }
108+ if (outputs.size() != 2 || !CheckOutput(outputs[0], xShape, 2) || !CheckOutput(outputs[1], xShape, 4)) {
109+ std::printf("RmsNormQuantV2 output verification failed\n");
110+ return FAILED_CODE;
111+ }
112+ std::printf("RmsNormQuantV2 PASS: y1=[2,...], y2=[4,...], both shapes=[2,16]\n");
113+ return SUCCESS_CODE;
114+}
115+} // namespace
116+ 
117+int main()
118+{
119+ const std::map<ge::AscendString, ge::AscendString> options = {{"ge.exec.deviceId", "0"}, {"ge.graphRunMode", "1"}};
120+ if (ge::GEInitialize(options) != ge::SUCCESS) {
121+ std::printf("GEInitialize failed\n");
122+ return FAILED_CODE;
123+ }
124+ const int32_t result = RunExample();
125+ // RunExample destroys the session before the global GE resources are released.
126+ if (ge::GEFinalize() != ge::SUCCESS) {
127+ std::printf("GEFinalize failed\n");
128+ return FAILED_CODE;
129+ }
130+ return result;
131+}
Anorm/rms_norm_quant_v2/tests/assets/golden.py+199-0
@@ -0,0 +1,199 @@
1+#!/usr/bin/env python3
2+# -*- coding: UTF-8 -*-
3+# Copyright (c) 2026 Huawei Technologies Co., Ltd.
4+# This program is free software, you can redistribute it and/or modify it under the terms and conditions of
5+# CANN Open Software License Agreement Version 2.0 (the "License").
6+# Please refer to the License for details. You may not use this file except in compliance with the License.
7+# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
8+# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
9+# See LICENSE in the root of the software repository for the full text of the License.
10+ 
11+import numpy as np
12+ 
13+ 
14+_DST_SAT_MAX = {
15+ "int8": 127.0,
16+ "int4": 7.0,
17+ "float8_e5m2": 57344.0,
18+ "float8_e4m3fn": 448.0,
19+ "hifloat8": 32768.0,
20+}
21+_DST_SAT_MIN = {
22+ "int8": -128.0,
23+ "int4": -8.0,
24+ "float8_e5m2": -57344.0,
25+ "float8_e4m3fn": -448.0,
26+ "hifloat8": -32768.0,
27+}
28+ 
29+ 
30+__golden__ = {"kernel": {"rms_norm_quant_v2": "rms_norm_quant_v2"}}
31+__spec__ = {"rms_norm_quant_v2": "RmsNormQuantV2TestSpec"}
32+ 
33+ 
34+def rms_norm_quant_v2(
35+ input_x,
36+ input_gamma,
37+ input_scales1,
38+ input_scales2=None,
39+ input_zero_points1=None,
40+ input_zero_points2=None,
41+ input_beta=None,
42+ epsilon=1e-6,
43+ div_mode=True,
44+ output_dtypes=None,
45+ dst_type=2,
46+ **kwargs,
47+):
48+ DST_TYPE_MAP = {
49+ 2: "int8",
50+ 29: "int4",
51+ 34: "hifloat8",
52+ 35: "float8_e5m2",
53+ 36: "float8_e4m3fn",
54+ }
55+ if output_dtypes is None:
56+ quant_dtype = DST_TYPE_MAP.get(dst_type, "int8")
57+ output_dtypes = (quant_dtype, quant_dtype)
58+ y1_dtype = str(output_dtypes[0]) if output_dtypes is not None else "int8"
59+ y2_dtype = (
60+ str(output_dtypes[1])
61+ if (output_dtypes is not None and len(output_dtypes) > 1)
62+ else None
63+ )
64+ 
65+ # Flatten only the normalization axes. Quantization parameters describe
66+ # their own suffix of x, which need not have the same rank as gamma.
67+ original_x_shape = input_x.shape
68+ norm_dim_num = max(1, input_gamma.ndim)
69+ n = input_x.ndim - norm_dim_num
70+ input_gamma = input_gamma.reshape(-1)
71+ x_shape = input_x.shape[0:n] + input_gamma.shape
72+ input_x = input_x.reshape(x_shape)
73+ 
74+ len_shape_x = len(input_x.shape)
75+ len_shape_gamma = len(input_gamma.shape)
76+ axis = len_shape_x - len_shape_gamma
77+ input_x_dtype = input_x.dtype
78+ input_scales_dtype = input_scales1.dtype
79+ input_zp_dtype = None
80+ if input_zero_points1 is not None:
81+ input_zp_dtype = input_zero_points1.dtype
82+ if input_zero_points2 is not None:
83+ input_zp_dtype = input_zero_points2.dtype
84+ 
85+ if input_scales_dtype is not np.float32:
86+ input_scales1 = input_scales1.astype(np.float32)
87+ if input_scales2 is not None:
88+ if input_scales_dtype is not np.float32:
89+ input_scales2 = input_scales2.astype(np.float32)
90+ if input_zero_points1 is not None:
91+ if input_zp_dtype == np.int32:
92+ input_zero_points1 = input_zero_points1.astype(np.float32)
93+ elif input_scales_dtype is not np.float32:
94+ input_zero_points1 = input_zero_points1.astype(np.float32)
95+ if input_zero_points2 is not None:
96+ if input_zp_dtype == np.int32:
97+ input_zero_points2 = input_zero_points2.astype(np.float32)
98+ elif input_scales_dtype is not np.float32:
99+ input_zero_points2 = input_zero_points2.astype(np.float32)
100+ 
101+ x_fp32 = input_x.astype(np.float32)
102+ variance = np.mean(np.power(x_fp32, 2), axis=axis, keepdims=True)
103+ std = np.sqrt(variance + epsilon)
104+ rstd = 1 / std
105+ result_mid = x_fp32 * rstd
106+ 
107+ if input_x_dtype == np.float32:
108+ y_array = result_mid * input_gamma
109+ elif input_x_dtype == np.float16:
110+ input_gamma_fp32 = input_gamma.astype(np.float32)
111+ y_array = result_mid * input_gamma_fp32
112+ else:
113+ input_gamma_fp32 = input_gamma.astype(np.float32)
114+ y_array = result_mid * input_gamma_fp32
115+ 
116+ if input_beta is not None:
117+ beta_reshaped = input_beta.reshape(input_gamma.shape)
118+ y_array = y_array + beta_reshaped.astype(np.float32)
119+ 
120+ # Restore x's axes before broadcasting scalar or multidimensional scales
121+ # and zero points; never index shape[0] on a scalar quantization parameter.
122+ y_array = y_array.reshape(original_x_shape)
123+ 
124+ if div_mode:
125+ input_scales1 = 1.0 / input_scales1
126+ if input_scales2 is not None:
127+ input_scales2 = 1.0 / input_scales2
128+ 
129+ y1_quant = y_array * input_scales1
130+ y2_quant = None
131+ if input_zero_points1 is not None:
132+ y1_quant = y1_quant + input_zero_points1
133+ 
134+ if input_scales2 is not None:
135+ y2_quant = y_array * input_scales2
136+ if input_zero_points2 is not None:
137+ y2_quant = y2_quant + input_zero_points2
138+ 
139+ def quant_process(y_quant, q_dtype):
140+ y_quant = np.nan_to_num(
141+ y_quant,
142+ nan=0.0,
143+ posinf=_DST_SAT_MAX.get(q_dtype, 0.0),
144+ neginf=_DST_SAT_MIN.get(q_dtype, 0.0),
145+ )
146+ if q_dtype in _DST_SAT_MAX:
147+ y_quant = np.clip(y_quant, _DST_SAT_MIN[q_dtype], _DST_SAT_MAX[q_dtype])
148+ if q_dtype in ("int8", "int4"):
149+ y_quant = y_quant.astype(np.float16).astype(np.float32)
150+ y_quant = np.round(y_quant, 0)
151+ elif q_dtype not in ("hifloat8", "float8_e5m2", "float8_e4m3fn"):
152+ y_quant = np.round(y_quant, 0)
153+ if q_dtype == "int8":
154+ y_quant = y_quant.clip(-128, 127)
155+ y_quant = y_quant.astype("int8", copy=False)
156+ elif q_dtype == "int4":
157+ from ttk.utilities.dtypes import numpy_int4
158+ 
159+ y_quant = np.clip(y_quant, -8, 7)
160+ y_quant = y_quant.astype(numpy_int4(), copy=False)
161+ elif q_dtype == "int32":
162+ y_quant = np.nan_to_num(y_quant, nan=0)
163+ y_quant = np.clip(y_quant, -2147483648, 2147483647)
164+ y_quant = y_quant.astype("int32", copy=False)
165+ elif q_dtype == "float8_e5m2":
166+ from ml_dtypes import float8_e5m2 as numpy_float8_e5m2
167+ 
168+ y_quant = y_quant.astype(numpy_float8_e5m2, copy=False)
169+ elif q_dtype == "float8_e4m3fn":
170+ from ml_dtypes import float8_e4m3fn as numpy_float8_e4m3fn
171+ 
172+ y_quant = y_quant.astype(numpy_float8_e4m3fn, copy=False)
173+ elif q_dtype == "hifloat8":
174+ from en_dtypes import hifloat8 as numpy_hifloat8
175+ 
176+ y_quant = y_quant.astype(numpy_hifloat8, copy=False)
177+ return y_quant
178+ 
179+ return quant_process(y1_quant, y1_dtype), quant_process(
180+ y2_quant, y2_dtype
181+ ) if y2_quant is not None else None
182+ 
183+ 
184+class RmsNormQuantV2TestSpec:
185+ """整数量化输出按 1 LSB 判；FP8 输出按相邻量化码判。
186+ 
187+ y = round(x / scale) 的 round 是阶跃函数：商落在 .5 附近时，浮点末位 1 ulp 的差异
188+ 就足以让结果跳到相邻整数格。而末位差异来自归约次序、FMA 融合、硬件舍入等实现自由度，
189+ 不同实现之间必然存在。要求逐位相等等于要求复现某一个实现的浮点运算次序，
190+ 故整数按绝对误差 <= 1 判定，FP8 按 requant 的码距判定。
191+ """
192+ 
193+ golden = staticmethod(rms_norm_quant_v2)
194+ tolerance = {
195+ "int8": {"standard": "quant"},
196+ "int4": {"standard": "quant"},
197+ "float8_e5m2": {"standard": "requant"},
198+ "float8_e4m3fn": {"standard": "requant"},
199+ }
Anorm/rms_norm_quant_v3/tests/assets/golden.py+435-0
@@ -0,0 +1,435 @@
1+#!/usr/bin/env python3
2+# -*- coding: UTF-8 -*-
3+# Copyright (c) 2026 Huawei Technologies Co., Ltd.
4+# This program is free software, you can redistribute it and/or modify it under the terms and conditions of
5+# CANN Open Software License Agreement Version 2.0 (the "License").
6+# Please refer to the License for details. You may not use this file except in compliance with the License.
7+# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
8+# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
9+# See LICENSE in the root of the software repository for the full text of the License.
10+ 
11+import numpy as np
12+import torch
13+ 
14+ 
15+_DST_SAT_MAX = {
16+ "int8": 127.0,
17+ "int4": 7.0,
18+ "float8_e5m2": 57344.0,
19+ "float8_e4m3fn": 448.0,
20+ "hifloat8": 32768.0,
21+}
22+_DST_SAT_MIN = {
23+ "int8": -128.0,
24+ "int4": -8.0,
25+ "float8_e5m2": -57344.0,
26+ "float8_e4m3fn": -448.0,
27+ "hifloat8": -32768.0,
28+}
29+ 
30+ 
31+__golden__ = {
32+ "kernel": {"rms_norm_quant_v3": "rms_norm_quant_v3"},
33+ "aclnn": {"aclnnRmsNormQuantV3": "aclnn_rms_norm_quant_v3_golden"},
34+}
35+__spec__ = {
36+ "rms_norm_quant_v3": "RmsNormQuantV3TestSpec",
37+ "aclnnRmsNormQuantV3": "AclnnRmsNormQuantV3TestSpec",
38+}
39+ 
40+ 
41+def rms_norm_quant_v3(
42+ input_x,
43+ input_gamma,
44+ input_scales1,
45+ input_scales2=None,
46+ input_zero_points1=None,
47+ input_zero_points2=None,
48+ input_beta=None,
49+ epsilon=1e-6,
50+ div_mode=True,
51+ output_dtypes=None,
52+ dst_type=2,
53+ output_rstd=False,
54+ **kwargs,
55+):
56+ DST_TYPE_MAP = {
57+ 2: "int8",
58+ 29: "int4",
59+ 34: "hifloat8",
60+ 35: "float8_e5m2",
61+ 36: "float8_e4m3fn",
62+ }
63+ if output_dtypes is None:
64+ quant_dtype = DST_TYPE_MAP.get(dst_type, "int8")
65+ output_dtypes = (quant_dtype, quant_dtype, "float32")
66+ y1_dtype = str(output_dtypes[0]) if output_dtypes is not None else "int8"
67+ y2_dtype = (
68+ str(output_dtypes[1])
69+ if (output_dtypes is not None and len(output_dtypes) > 1)
70+ else None
71+ )
72+ 
73+ # Flatten only the normalization axes. Quantization parameters describe
74+ # their own suffix of x, which need not have the same rank as gamma.
75+ original_x_shape = input_x.shape
76+ norm_dim_num = max(1, input_gamma.ndim)
77+ n = input_x.ndim - norm_dim_num
78+ input_gamma = input_gamma.reshape(-1)
79+ x_shape = input_x.shape[0:n] + input_gamma.shape
80+ input_x = input_x.reshape(x_shape)
81+ 
82+ len_shape_x = len(input_x.shape)
83+ len_shape_gamma = len(input_gamma.shape)
84+ axis = len_shape_x - len_shape_gamma
85+ input_x_dtype = input_x.dtype
86+ input_scales_dtype = input_scales1.dtype
87+ input_zp_dtype = None
88+ if input_zero_points1 is not None:
89+ input_zp_dtype = input_zero_points1.dtype
90+ if input_zero_points2 is not None:
91+ input_zp_dtype = input_zero_points2.dtype
92+ 
93+ if input_scales_dtype is not np.float32:
94+ input_scales1 = input_scales1.astype(np.float32)
95+ if input_scales2 is not None:
96+ if input_scales_dtype is not np.float32:
97+ input_scales2 = input_scales2.astype(np.float32)
98+ if input_zero_points1 is not None:
99+ if input_zp_dtype == np.int32:
100+ input_zero_points1 = input_zero_points1.astype(np.float32)
101+ elif input_scales_dtype is not np.float32:
102+ input_zero_points1 = input_zero_points1.astype(np.float32)
103+ if input_zero_points2 is not None:
104+ if input_zp_dtype == np.int32:
105+ input_zero_points2 = input_zero_points2.astype(np.float32)
106+ elif input_scales_dtype is not np.float32:
107+ input_zero_points2 = input_zero_points2.astype(np.float32)
108+ 
109+ x_fp32 = input_x.astype(np.float32)
110+ variance = np.mean(np.power(x_fp32, 2), axis=axis, keepdims=True)
111+ std = np.sqrt(variance + epsilon)
112+ rstd = 1 / std
113+ result_mid = x_fp32 * rstd
114+ 
115+ if input_x_dtype == np.float32:
116+ y_array = result_mid * input_gamma
117+ elif input_x_dtype == np.float16:
118+ input_gamma_fp32 = input_gamma.astype(np.float32)
119+ y_array = result_mid * input_gamma_fp32
120+ else:
121+ input_gamma_fp32 = input_gamma.astype(np.float32)
122+ y_array = result_mid * input_gamma_fp32
123+ 
124+ if input_beta is not None:
125+ beta_reshaped = input_beta.reshape(input_gamma.shape)
126+ y_array = y_array + beta_reshaped.astype(np.float32)
127+ 
128+ # Restore x's axes before broadcasting scalar or multidimensional scales
129+ # and zero points; never index shape[0] on a scalar quantization parameter.
130+ y_array = y_array.reshape(original_x_shape)
131+ 
132+ if div_mode:
133+ input_scales1 = 1.0 / input_scales1
134+ if input_scales2 is not None:
135+ input_scales2 = 1.0 / input_scales2
136+ 
137+ y1_quant = y_array * input_scales1
138+ y2_quant = None
139+ if input_zero_points1 is not None:
140+ y1_quant = y1_quant + input_zero_points1
141+ 
142+ if input_scales2 is not None:
143+ y2_quant = y_array * input_scales2
144+ if input_zero_points2 is not None:
145+ y2_quant = y2_quant + input_zero_points2
146+ 
147+ def quant_process(y_quant, q_dtype):
148+ y_quant = np.nan_to_num(
149+ y_quant,
150+ nan=0.0,
151+ posinf=_DST_SAT_MAX.get(q_dtype, 0.0),
152+ neginf=_DST_SAT_MIN.get(q_dtype, 0.0),
153+ )
154+ if q_dtype in _DST_SAT_MAX:
155+ y_quant = np.clip(y_quant, _DST_SAT_MIN[q_dtype], _DST_SAT_MAX[q_dtype])
156+ if q_dtype in ("int8", "int4"):
157+ y_quant = y_quant.astype(np.float16).astype(np.float32)
158+ y_quant = np.round(y_quant, 0)
159+ elif q_dtype not in ("hifloat8", "float8_e5m2", "float8_e4m3fn"):
160+ y_quant = np.round(y_quant, 0)
161+ if q_dtype == "int8":
162+ y_quant = y_quant.clip(-128, 127)
163+ y_quant = y_quant.astype("int8", copy=False)
164+ elif q_dtype == "int4":
165+ from ttk.utilities.dtypes import numpy_int4
166+ 
167+ y_quant = np.clip(y_quant, -8, 7)
168+ y_quant = y_quant.astype(numpy_int4(), copy=False)
169+ elif q_dtype == "int32":
170+ y_quant = np.nan_to_num(y_quant, nan=0)
171+ y_quant = np.clip(y_quant, -2147483648, 2147483647)
172+ y_quant = y_quant.astype("int32", copy=False)
173+ elif q_dtype == "float8_e5m2":
174+ from ml_dtypes import float8_e5m2 as numpy_float8_e5m2
175+ 
176+ y_quant = y_quant.astype(numpy_float8_e5m2, copy=False)
177+ elif q_dtype == "float8_e4m3fn":
178+ from ml_dtypes import float8_e4m3fn as numpy_float8_e4m3fn
179+ 
180+ y_quant = y_quant.astype(numpy_float8_e4m3fn, copy=False)
181+ elif q_dtype == "hifloat8":
182+ from en_dtypes import hifloat8 as numpy_hifloat8
183+ 
184+ y_quant = y_quant.astype(numpy_hifloat8, copy=False)
185+ return y_quant
186+ 
187+ rstd_shape = original_x_shape[:n] + (1,) * norm_dim_num
188+ rstd_out = (
189+ rstd.reshape(rstd_shape).astype(np.float32, copy=False) if output_rstd else None
190+ )
191+ return (
192+ quant_process(y1_quant, y1_dtype),
193+ quant_process(y2_quant, y2_dtype) if y2_quant is not None else None,
194+ rstd_out,
195+ )
196+ 
197+ 
198+class RmsNormQuantV3TestSpec:
199+ """整数量化输出按 1 LSB 判；FP8 输出按相邻量化码判。
200+ 
201+ y = round(x / scale) 的 round 是阶跃函数：商落在 .5 附近时，浮点末位 1 ulp 的差异
202+ 就足以让结果跳到相邻整数格。而末位差异来自归约次序、FMA 融合、硬件舍入等实现自由度，
203+ 不同实现之间必然存在。要求逐位相等等于要求复现某一个实现的浮点运算次序，
204+ 故整数按绝对误差 <= 1 判定，FP8 按 requant 的码距判定。
205+ """
206+ 
207+ golden = staticmethod(rms_norm_quant_v3)
208+ tolerance = {
209+ "int8": {"standard": "quant"},
210+ "int4": {"standard": "quant"},
211+ "float8_e5m2": {"standard": "requant"},
212+ "float8_e4m3fn": {"standard": "requant"},
213+ "float32": {"standard": "stat_rel_err"},
214+ }
215+ 
216+ 
217+_INT4_VALUES_PER_BYTE = 2
218+_INT4_VALUES_PER_INT32 = np.dtype(np.int32).itemsize * _INT4_VALUES_PER_BYTE
219+ 
220+ 
221+def _to_numpy(t):
222+ if t is None:
223+ return None
224+ if isinstance(t, np.ndarray):
225+ return t
226+ if hasattr(t, "numpy"):
227+ try:
228+ return t.numpy()
229+ except (TypeError, RuntimeError):
230+ return t.float().numpy()
231+ if hasattr(t, "detach"):
232+ try:
233+ return t.detach().cpu().numpy()
234+ except (TypeError, RuntimeError):
235+ return t.detach().cpu().float().numpy()
236+ return np.array(t)
237+ 
238+ 
239+def _np_to_torch(arr):
240+ if arr is None:
241+ return None
242+ if hasattr(arr, "view") and hasattr(arr, "dim"):
243+ return arr
244+ if isinstance(arr, np.ndarray):
245+ dt = str(arr.dtype)
246+ if np.issubdtype(arr.dtype, np.integer):
247+ return torch.from_numpy(arr)
248+ if "float8_e5m2" in dt:
249+ return torch.from_numpy(arr.view(np.uint8)).view(torch.float8_e5m2)
250+ if "float8_e4m3" in dt:
251+ return torch.from_numpy(arr.view(np.uint8)).view(torch.float8_e4m3fn)
252+ if "hifloat8" in dt:
253+ return torch.from_numpy(arr.view(np.int8)).view(torch.int8)
254+ return torch.from_numpy(arr.astype(np.float32))
255+ return arr
256+ 
257+ 
258+def _quant_to_dtype(y_quant_np, dtype_str):
259+ if dtype_str not in ("int8", "hifloat8", "float8_e5m2", "float8_e4m3fn", "int4"):
260+ return y_quant_np
261+ 
262+ y_quant_np = np.nan_to_num(
263+ y_quant_np,
264+ nan=0.0,
265+ posinf=_DST_SAT_MAX.get(dtype_str, 0.0),
266+ neginf=_DST_SAT_MIN.get(dtype_str, 0.0),
267+ )
268+ 
269+ if dtype_str in ("int8", "int4"):
270+ y_quant_np = np.round(y_quant_np, 0)
271+ 
272+ if dtype_str in _DST_SAT_MAX:
273+ y_quant_np = np.clip(
274+ y_quant_np, _DST_SAT_MIN[dtype_str], _DST_SAT_MAX[dtype_str]
275+ )
276+ 
277+ if dtype_str == "int8":
278+ return y_quant_np.astype(np.int8)
279+ if dtype_str == "int4":
280+ from ttk.utilities.dtypes import numpy_int4
281+ 
282+ y_quant_np = np.clip(y_quant_np, -8, 7)
283+ return y_quant_np.astype(numpy_int4(), copy=False)
284+ if dtype_str == "hifloat8":
285+ from en_dtypes import hifloat8
286+ 
287+ return y_quant_np.astype(hifloat8, copy=False)
288+ if dtype_str == "float8_e5m2":
289+ from ml_dtypes import float8_e5m2
290+ 
291+ return y_quant_np.astype(float8_e5m2, copy=False)
292+ if dtype_str == "float8_e4m3fn":
293+ from ml_dtypes import float8_e4m3fn
294+ 
295+ return y_quant_np.astype(float8_e4m3fn, copy=False)
296+ return y_quant_np
297+ 
298+ 
299+def _pack_int4_to_int32(y_int4):
300+ """Pack the public ACLNN INT4 result into its int32 storage ABI."""
301+ if y_int4.ndim == 0 or y_int4.shape[-1] % _INT4_VALUES_PER_INT32 != 0:
302+ raise ValueError(
303+ f"INT4 last dimension must be divisible by {_INT4_VALUES_PER_INT32}, "
304+ f"but got shape {y_int4.shape}"
305+ )
306+ from ttk.utilities.dtypes import pack_4bits
307+ 
308+ packed_bytes = pack_4bits(np.ascontiguousarray(y_int4))
309+ packed_shape = y_int4.shape[:-1] + (y_int4.shape[-1] // _INT4_VALUES_PER_INT32,)
310+ return packed_bytes.view(np.int32).reshape(packed_shape)
311+ 
312+ 
313+# 形参表必须逐位对齐 aclnnRmsNormQuantV3 的原型：属性（epsilon/divMode/outputRstd）
314+# 排在输出 tensor（y/rstd）**前面**。把 y/rstd 提前会让 epsilon 收到一个 tensor，
315+# 典型报错 'float' object has no attribute 'dtype' / 精度整片错。
316+# dst_type 不在原型里，只能挂末尾走 extra_attrs。
317+def aclnn_rms_norm_quant_v3_golden(
318+ x,
319+ gamma,
320+ scale,
321+ offset,
322+ beta,
323+ epsilon=1e-6,
324+ div_mode=True,
325+ output_rstd=False,
326+ y=None,
327+ rstd=None,
328+ dst_type=2,
329+ **kwargs,
330+):
331+ divMode = div_mode
332+ outputRstd = output_rstd
333+ x = _to_numpy(x)
334+ gamma = _to_numpy(gamma)
335+ scale = _to_numpy(scale)
336+ offset = _to_numpy(offset)
337+ beta = _to_numpy(beta)
338+ 
339+ if scale is not None and scale.dtype != np.float32:
340+ scale = scale.astype(np.float32)
341+ if offset is not None and offset.dtype != np.float32:
342+ offset = offset.astype(np.float32)
343+ 
344+ x_fp32 = x.astype(np.float32)
345+ 
346+ # ACLNN squeezes a legal [1, H] gamma to [H] before last-axis RMSNorm.
347+ n = len(x.shape) - 1
348+ gamma_flat = gamma.reshape(-1)
349+ x_reshaped = x_fp32.reshape(x_fp32.shape[:n] + gamma_flat.shape)
350+ 
351+ variance = np.mean(np.power(x_reshaped, 2), axis=n, keepdims=True)
352+ std = np.sqrt(variance + epsilon)
353+ rstd_val = 1 / std
354+ 
355+ result_mid = x_reshaped * rstd_val
356+ 
357+ gamma_fp32 = gamma.astype(np.float32) if gamma.dtype != np.float32 else gamma
358+ y_array = result_mid * gamma_fp32.reshape(gamma_flat.shape)
359+ 
360+ if beta is not None:
361+ beta_reshaped = beta.reshape(gamma_flat.shape)
362+ y_array = y_array + beta_reshaped.astype(np.float32)
363+ 
364+ y_dtype_str = "int8"
365+ if y is not None:
366+ dt_str = str(y.dtype)
367+ if "int8" in dt_str:
368+ y_dtype_str = "int8"
369+ elif "hifloat8" in dt_str:
370+ y_dtype_str = "hifloat8"
371+ elif "e5m2" in dt_str:
372+ y_dtype_str = "float8_e5m2"
373+ elif "e4m3" in dt_str:
374+ y_dtype_str = "float8_e4m3fn"
375+ elif "int4" in dt_str or "int32" in dt_str:
376+ y_dtype_str = "int4"
377+ else:
378+ DST_TYPE_MAP = {
379+ 2: "int8",
380+ 29: "int4",
381+ 34: "hifloat8",
382+ 35: "float8_e5m2",
383+ 36: "float8_e4m3fn",
384+ }
385+ y_dtype_str = DST_TYPE_MAP.get(dst_type, "int8")
386+ 
387+ if divMode:
388+ y_quant = (
389+ y_array / scale.reshape(-1).reshape(scale.shape)
390+ if scale is not None
391+ else y_array
392+ )
393+ else:
394+ y_quant = (
395+ y_array * scale.reshape(-1).reshape(scale.shape)
396+ if scale is not None
397+ else y_array
398+ )
399+ if offset is not None:
400+ y_quant = y_quant + offset.reshape(-1).reshape(offset.shape)
401+ 
402+ y_result = _quant_to_dtype(y_quant, y_dtype_str)
403+ if y_dtype_str == "int4":
404+ y_result = _pack_int4_to_int32(y_result)
405+ y_t = _np_to_torch(y_result)
406+ 
407+ if outputRstd:
408+ rstd_np = rstd_val.astype(np.float32)
409+ rstd_t = _np_to_torch(rstd_np)
410+ return [y_t, rstd_t]
411+ 
412+ return [y_t, None]
413+ 
414+ 
415+class AclnnRmsNormQuantV3TestSpec:
416+ """量化输出按 1 LSB 判（standard='quant'，判据实现见 comparison/quant.py）。
417+ 
418+ y = round(x / scale) 的 round 是阶跃函数：商落在 .5 附近时，浮点末位 1 ulp 的差异
419+ 就足以让结果跳到相邻整数格。而末位差异来自归约次序、FMA 融合、硬件舍入等实现自由度，
420+ 不同实现之间必然存在。要求逐位相等等于要求复现某一个实现的浮点运算次序，
421+ 故按绝对误差 <= 1 判定。
422+ """
423+ 
424+ golden = staticmethod(aclnn_rms_norm_quant_v3_golden)
425+ tolerance = {
426+ "int8": {"standard": "quant"},
427+ "int4": {"standard": "quant"},
428+ "int32": {"standard": "quant", "packed_bits": 4},
429+ "float8_e5m2": {"standard": "requant"},
430+ "float8_e4m3fn": {"standard": "requant"},
431+ # rstd is already produced as float32 by the golden. Avoid whole-call
432+ # dtype promotion, which cannot cross the Torch boundary for a
433+ # HIFLOAT8 sibling output.
434+ "float32": {"standard": "stat_rel_err"},
435+ }