已合并
fix(swiglu_mx_quant): 修复检视意见 + 新增标杆和aclnn文档 #9292
caorenlei创建于 13 天前
fix(swiglu_mx_quant): 修复检视意见 + 新增标杆和aclnn文档 #9292
已合并
共 3 个文件变更+380-10
| @@ -60,6 +60,54 @@ | |||
| 60 | 60 | ||
| 61 | 其中，x\_glu表示x<sub>i</sub>的偶数索引部分，x\_linear表示x<sub>i</sub>的奇数索引部分。 | 61 | 其中，x\_glu表示x<sub>i</sub>的偶数索引部分，x\_linear表示x<sub>i</sub>的奇数索引部分。 |
| 62 | 62 | ||
| 63 | +- swigluMode为2时的计算公式： | ||
| 64 | + | ||
| 65 | + $$ | ||
| 66 | + x\_glu = x\_glu.clamp(min=None, max=clampLimit) | ||
| 67 | + $$ | ||
| 68 | + | ||
| 69 | + $$ | ||
| 70 | + x\_linear = x\_linear.clamp(min=-clampLimit, max=clampLimit) | ||
| 71 | + $$ | ||
| 72 | + | ||
| 73 | + $$ | ||
| 74 | + out\_glu = x\_glu * sigmoid(gluAlpha * x\_glu) | ||
| 75 | + $$ | ||
| 76 | + | ||
| 77 | + $$ | ||
| 78 | + swigluOut_i = out\_glu * (x\_linear + gluBias) | ||
| 79 | + $$ | ||
| 80 | + | ||
| 81 | + $$ | ||
| 82 | + y, mxscale = DynamicMxQuant(swigluOut_i) | ||
| 83 | + $$ | ||
| 84 | + | ||
| 85 | + 其中，x\_glu和x\_linear由x<sub>i</sub>沿activateDim轴前后分半得到（非交错分块），当activateLeft为true时x\_glu为前半部分，否则为后半部分。此模式下axis必须为-1。 | ||
| 86 | + | ||
| 87 | +- swigluMode为3时的计算公式： | ||
| 88 | + | ||
| 89 | + $$ | ||
| 90 | + x\_glu = x\_glu * sigmoid(x\_glu) | ||
| 91 | + $$ | ||
| 92 | + | ||
| 93 | + $$ | ||
| 94 | + x\_glu = x\_glu.clamp(min=None, max=clampLimit) | ||
| 95 | + $$ | ||
| 96 | + | ||
| 97 | + $$ | ||
| 98 | + x\_linear = x\_linear.clamp(min=-clampLimit, max=clampLimit) | ||
| 99 | + $$ | ||
| 100 | + | ||
| 101 | + $$ | ||
| 102 | + swigluOut_i = x\_glu * x\_linear | ||
| 103 | + $$ | ||
| 104 | + | ||
| 105 | + $$ | ||
| 106 | + y, mxscale = DynamicMxQuant(swigluOut_i) | ||
| 107 | + $$ | ||
| 108 | + | ||
| 109 | + 其中，x\_glu和x\_linear由x<sub>i</sub>沿activateDim轴前后分半得到（非交错分块）。此模式下忽略gluAlpha和gluBias参数（sigmoid固定使用alpha=1，不添加bias），axis必须为-1。 | ||
| 110 | + | ||
| 63 | - 场景1，当scaleAlg为0时，DynamicMxQuant计算逻辑： | 111 | - 场景1，当scaleAlg为0时，DynamicMxQuant计算逻辑： |
| 64 | - 将输入x在axis维度上按k = blocksize = 32个数分组，一组k个数 $\{\{V_i\}_{i=1}^{k}\}$ 动态量化为 $\{mxscale1, \{P_i\}_{i=1}^{k}\}$, k = blocksize | 112 | - 将输入x在axis维度上按k = blocksize = 32个数分组，一组k个数 $\{\{V_i\}_{i=1}^{k}\}$ 动态量化为 $\{mxscale1, \{P_i\}_{i=1}^{k}\}$, k = blocksize |
| 65 | 113 | ||
| @@ -204,7 +252,7 @@ aclnnStatus aclnnSwigluMxQuant( | |||
| 204 | <td>swigluMode（int64_t）</td> | 252 | <td>swigluMode（int64_t）</td> |
| 205 | <td>输入</td> | 253 | <td>输入</td> |
| 206 | <td>表示swiglu的计算模式。</td> | 254 | <td>表示swiglu的计算模式。</td> |
| 207 | - <td><ul><li>取值范围为：[0, 1]。</li><li>0：表示传统swiglu计算方式。</li><li>1：表示swiglu的变种，使用奇偶分块方式，并支持clampLimit、激活系数以及偏差。</li></ul></td> | 255 | + <td><ul><li>取值范围为：[0, 3]。</li><li>0：表示传统swiglu计算方式。</li><li>1：表示swiglu的变种，使用奇偶分块方式，并支持clampLimit、激活系数以及偏差。</li><li>2：表示swiglu的变种，使用前后分半方式，先clamp再sigmoid(alpha)激活，并支持偏差。此模式要求axis=-1。</li><li>3：表示swiglu的变种，使用前后分半方式，先sigmoid激活（alpha固定为1）再clamp，不支持偏差。此模式要求axis=-1。</li></ul></td> |
| 208 | <td>-</td> | 256 | <td>-</td> |
| 209 | <td>-</td> | 257 | <td>-</td> |
| 210 | <td>-</td> | 258 | <td>-</td> |
| @@ -396,6 +444,9 @@ aclnnStatus aclnnSwigluMxQuant( | |||
| 396 | <tr> | 444 | <tr> |
| 397 | <td>scaleAlg不在指定的取值范围内。</td> | 445 | <td>scaleAlg不在指定的取值范围内。</td> |
| 398 | </tr> | 446 | </tr> |
| 447 | + <tr> | ||
| 448 | + <td>swigluMode为2或3时axis不为-1。</td> | ||
| 449 | + </tr> | ||
| 399 | </tbody></table> | 450 | </tbody></table> |
| 400 | 451 | ||
| 401 | ## aclnnSwigluMxQuant | 452 | ## aclnnSwigluMxQuant |
| @@ -446,7 +497,8 @@ aclnnStatus aclnnSwigluMxQuant( | |||
| 446 | - aclnnSwigluMxQuant默认确定性实现。 | 497 | - aclnnSwigluMxQuant默认确定性实现。 |
| 447 | 498 | ||
| 448 | - 输入x对应activateDim的维度需要是2的倍数，且x的维数必须大于1维。 | 499 | - 输入x对应activateDim的维度需要是2的倍数，且x的维数必须大于1维。 |
| 449 | -- activateDim为非last轴，swigluMode必须为0。 | 500 | +- activateDim为非last轴，swigluMode必须为0或1。 |
| 501 | +- swigluMode为2或3时，axis必须为-1（仅支持尾轴量化）。 | ||
| 450 | - 当输出yOut的数据类型为FLOAT4_E2M1、FLOAT4_E1M2时，yOut的最后一维需要是2的倍数。 | 502 | - 当输出yOut的数据类型为FLOAT4_E2M1、FLOAT4_E1M2时，yOut的最后一维需要是2的倍数。 |
| 451 | - 当输出yOut的数据类型为FLOAT4_E2M1、FLOAT4_E1M2时，scaleAlg必须为0。 | 503 | - 当输出yOut的数据类型为FLOAT4_E2M1、FLOAT4_E1M2时，scaleAlg必须为0。 |
| 452 | - groupIndexOptional所有元素之和不能大于输入x除尾轴之外的剩余轴的乘积，groupIndexOptional的每个元素需要大于0。 | 504 | - groupIndexOptional所有元素之和不能大于输入x除尾轴之外的剩余轴的乘积，groupIndexOptional的每个元素需要大于0。 |
| @@ -21,6 +21,7 @@ | |||
| 21 | 21 | ||
| 22 | 22 | ||
| 23 | 23 | ||
| 24 | + | ||
| 24 | 25 | ||
| 25 | 26 | ||
| 26 | using namespace std; | 27 | using namespace std; |
| @@ -145,17 +146,20 @@ ge::graphStatus SwigluMxQuantRegbaseTiling::ParseAttrs() | |||
| 145 | // Get swiglu_mode (int64 type) | 146 | // Get swiglu_mode (int64 type) |
| 146 | auto* attrSwigluMode = attrs->GetAttrPointer<int64_t>(INDEX_ATTR_SWIGLU_MODE); | 147 | auto* attrSwigluMode = attrs->GetAttrPointer<int64_t>(INDEX_ATTR_SWIGLU_MODE); |
| 147 | attrParam_.swigluMode = (attrSwigluMode != nullptr) ? static_cast<int64_t>(*attrSwigluMode) : 0; | 148 | attrParam_.swigluMode = (attrSwigluMode != nullptr) ? static_cast<int64_t>(*attrSwigluMode) : 0; |
| 148 | - OP_CHECK_IF(attrParam_.swigluMode < 0 || attrParam_.swigluMode > 3, | 149 | + OP_CHECK_IF( |
| 149 | - OP_LOGE(context_->GetNodeName(), "swigluMode must be in [0, 3], but is %ld", attrParam_.swigluMode), | 150 | + attrParam_.swigluMode < 0 || attrParam_.swigluMode > 3, |
| 150 | - return ge::GRAPH_FAILED); | 151 | + OP_LOGE_FOR_INVALID_VALUE_WITH_REASON(context_->GetNodeName(), "swiglu_mode", |
| 152 | + std::to_string(attrParam_.swigluMode), "swiglu_mode must be in [0, 3]"), | ||
| 153 | + return ge::GRAPH_FAILED); | ||
| 151 | 154 | ||
| 152 | // Get clamp_limit (float type) | 155 | // Get clamp_limit (float type) |
| 153 | auto* attrClampLimit = attrs->GetAttrPointer<float>(INDEX_ATTR_CLAMP_LIMIT); | 156 | auto* attrClampLimit = attrs->GetAttrPointer<float>(INDEX_ATTR_CLAMP_LIMIT); |
| 154 | attrParam_.clampLimit = (attrClampLimit != nullptr) ? *attrClampLimit : 7.0f; | 157 | attrParam_.clampLimit = (attrClampLimit != nullptr) ? *attrClampLimit : 7.0f; |
| 155 | OP_CHECK_IF((attrParam_.swigluMode == 1 || attrParam_.swigluMode == 2 || attrParam_.swigluMode == 3) && | 158 | OP_CHECK_IF((attrParam_.swigluMode == 1 || attrParam_.swigluMode == 2 || attrParam_.swigluMode == 3) && |
| 156 | (attrParam_.clampLimit <= 0.0f), | 159 | (attrParam_.clampLimit <= 0.0f), |
| 157 | - OP_LOGE(context_->GetNodeName(), "swigluMode == %ld, clampLimit must > 0, but is %f", | 160 | + OP_LOGE_FOR_INVALID_VALUE_WITH_REASON(context_->GetNodeName(), "clamp_limit", |
| 158 | - attrParam_.swigluMode, attrParam_.clampLimit), | 161 | + std::to_string(attrParam_.clampLimit), |
| 162 | + "clamp_limit must be > 0 when swiglu_mode is 1, 2, or 3"), | ||
| 159 | return ge::GRAPH_FAILED); | 163 | return ge::GRAPH_FAILED); |
| 160 | // Get glu_alpha (float type) | 164 | // Get glu_alpha (float type) |
| 161 | auto* attrGluAlpha = attrs->GetAttrPointer<float>(INDEX_ATTR_GLU_ALPHA); | 165 | auto* attrGluAlpha = attrs->GetAttrPointer<float>(INDEX_ATTR_GLU_ALPHA); |
| @@ -168,9 +172,6 @@ ge::graphStatus SwigluMxQuantRegbaseTiling::ParseAttrs() | |||
| 168 | // Get axis (int64 type) | 172 | // Get axis (int64 type) |
| 169 | auto* attrAxis = attrs->GetAttrPointer<int64_t>(INDEX_ATTR_AXIS); | 173 | auto* attrAxis = attrs->GetAttrPointer<int64_t>(INDEX_ATTR_AXIS); |
| 170 | attrParam_.axis = (attrAxis != nullptr) ? static_cast<int64_t>(*attrAxis) : -1; | 174 | attrParam_.axis = (attrAxis != nullptr) ? static_cast<int64_t>(*attrAxis) : -1; |
| 171 | - OP_CHECK_IF((attrParam_.axis != -1) && (attrParam_.swigluMode == 2 || attrParam_.swigluMode == 3), | ||
| 172 | - OP_LOGE(context_->GetNodeName(), "swigluMode 2/3 requires axis=-1, but axis=%ld", attrParam_.axis), | ||
| 173 | - return ge::GRAPH_FAILED); | ||
| 174 | 175 | ||
| 175 | // Get dst_type (int type) | 176 | // Get dst_type (int type) |
| 176 | auto* attrDstType = attrs->GetAttrPointer<int64_t>(INDEX_ATTR_DST_TYPE); | 177 | auto* attrDstType = attrs->GetAttrPointer<int64_t>(INDEX_ATTR_DST_TYPE); |
| @@ -241,6 +242,12 @@ ge::graphStatus SwigluMxQuantRegbaseTiling::ParseAttrs() | |||
| 241 | OP_CHECK_IF((attrParam_.axis != -1 && attrParam_.axis != -2), | 242 | OP_CHECK_IF((attrParam_.axis != -1 && attrParam_.axis != -2), |
| 242 | OP_LOGE(context_->GetNodeName(), "axis must be -1 or -2, but got %ld.", attrParam_.axis), | 243 | OP_LOGE(context_->GetNodeName(), "axis must be -1 or -2, but got %ld.", attrParam_.axis), |
| 243 | return ge::GRAPH_FAILED); | 244 | return ge::GRAPH_FAILED); |
| 245 | + | ||
| 246 | + // swigluMode 2/3 only support axis=-1 (last axis) | ||
| 247 | + OP_CHECK_IF((attrParam_.axis != -1) && (attrParam_.swigluMode == 2 || attrParam_.swigluMode == 3), | ||
| 248 | + OP_LOGE_FOR_INVALID_VALUE_WITH_REASON(context_->GetNodeName(), "axis", std::to_string(attrParam_.axis), | ||
| 249 | + "axis must be -1 when swiglu_mode is 2 or 3"), | ||
| 250 | + return ge::GRAPH_FAILED); | ||
| 244 | return ge::GRAPH_SUCCESS; | 251 | return ge::GRAPH_SUCCESS; |
| 245 | } | 252 | } |
| 246 | 253 | ||
| @@ -0,0 +1,311 @@ | |||
| 1 | +#!/usr/bin/env python3 | ||
| 2 | +# ----------------------------------------------------------------------------------------------------------- | ||
| 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 | +""" | ||
| 12 | +TTK custom golden for swiglu_mx_quant (SwiGLU + DynamicMxQuant fusion operator). | ||
| 13 | + | ||
| 14 | +Inputs (positional, in op-def order): | ||
| 15 | + x : numpy array (fp16/bf16) | ||
| 16 | + group_index : numpy int64 array, OPTIONAL (absent -> None) | ||
| 17 | + | ||
| 18 | +Attributes (via **kwargs, from CSV `attributes`): | ||
| 19 | + activate_dim : int (default -1) SwiGLU split axis | ||
| 20 | + activate_left : bool (default False) True=left half is gate | ||
| 21 | + swiglu_mode : int (default 0) 0=SwiGLU, 1=interleaved clamp, 2=split clamp, 3=split sigmoid-clamp | ||
| 22 | + clamp_limit : float (default 7.0) clamp bound for mode 1/2/3 | ||
| 23 | + glu_alpha : float (default 1.702) sigmoid scale for mode 1/2 | ||
| 24 | + glu_bias : float (default 1.0) bias added to linear path for mode 1/2 | ||
| 25 | + group_mode : int (default 0) | ||
| 26 | + axis : int (default -1) quantization axis | ||
| 27 | + dst_type : int (default 40) 40=FP4_E2M1, 41=FP4_E1M2, 36=FP8_E4M3FN, 35=FP8_E5M2 | ||
| 28 | + round_mode : str (default "rint") | ||
| 29 | + scale_alg : int (default 0) 0=per-blockscale, 1=per-block FP8 | ||
| 30 | + max_dtype_value: float (default 0.0) | ||
| 31 | + block_size : int (fixed 32) | ||
| 32 | + | ||
| 33 | +Outputs: | ||
| 34 | + y : quantized result (dst_type) | ||
| 35 | + mxscale : scale factors (FP8_E8M0) | ||
| 36 | + | ||
| 37 | +Reference (mirrors docs/aclnnSwigluMxQuant.md and kernel ComputeVfSwigluV1-V4): | ||
| 38 | + | ||
| 39 | + mode 0 (SwiGLU): | ||
| 40 | + chunk x along activate_dim into [A, B] | ||
| 41 | + y = silu(A) * B (if activate_left: silu(A)*B, else silu(B)*A) | ||
| 42 | + | ||
| 43 | + mode 1 (interleaved clamp): | ||
| 44 | + A = x[..., ::2], B = x[..., 1::2] | ||
| 45 | + A = clamp(A, max=clamp_limit) | ||
| 46 | + B = clamp(B, -clamp_limit, clamp_limit) | ||
| 47 | + y = A * sigmoid(glu_alpha * A) * (B + glu_bias) | ||
| 48 | + | ||
| 49 | + mode 2 (split clamp): | ||
| 50 | + chunk x along activate_dim into [x_glu, x_linear] | ||
| 51 | + x_glu = clamp(x_glu, max=clamp_limit) | ||
| 52 | + x_linear = clamp(x_linear, -clamp_limit, clamp_limit) | ||
| 53 | + y = x_glu * sigmoid(glu_alpha * x_glu) * (x_linear + glu_bias) | ||
| 54 | + | ||
| 55 | + mode 3 (split sigmoid-then-clamp): | ||
| 56 | + chunk x along activate_dim into [x_glu, x_linear] | ||
| 57 | + x_glu = x_glu * sigmoid(x_glu) # alpha=1, no bias | ||
| 58 | + x_glu = clamp(x_glu, max=clamp_limit) | ||
| 59 | + x_linear = clamp(x_linear, -clamp_limit, clamp_limit) | ||
| 60 | + y = x_glu * x_linear | ||
| 61 | +""" | ||
| 62 | + | ||
| 63 | +import math | ||
| 64 | +import numpy as np | ||
| 65 | + | ||
| 66 | +try: | ||
| 67 | + from ml_dtypes import bfloat16 as _bf16 | ||
| 68 | +except ImportError: | ||
| 69 | + _bf16 = None | ||
| 70 | + | ||
| 71 | +try: | ||
| 72 | + from ml_dtypes import float8_e4m3fn as _fp8_e4m3 | ||
| 73 | + from ml_dtypes import float8_e5m2 as _fp8_e5m2 | ||
| 74 | + from ml_dtypes import float8_e8m0 as _fp8_e8m0 | ||
| 75 | +except ImportError: | ||
| 76 | + _fp8_e4m3 = None | ||
| 77 | + _fp8_e5m2 = None | ||
| 78 | + _fp8_e8m0 = None | ||
| 79 | + | ||
| 80 | +try: | ||
| 81 | + from ml_dtypes import float4_e2m1_fn as _fp4_e2m1 | ||
| 82 | + from ml_dtypes import float4_e1m2_fn as _fp4_e1m2 | ||
| 83 | +except ImportError: | ||
| 84 | + _fp4_e2m1 = None | ||
| 85 | + _fp4_e1m2 = None | ||
| 86 | + | ||
| 87 | + | ||
| 88 | +_DST_TYPE_MAP = { | ||
| 89 | + 40: ("float4_e2m1", _fp4_e2m1, 2), | ||
| 90 | + 41: ("float4_e1m2", _fp4_e1m2, 0), | ||
| 91 | + 36: ("float8_e4m3fn", _fp8_e4m3, 8), | ||
| 92 | + 35: ("float8_e5m2", _fp8_e5m2, 15), | ||
| 93 | +} | ||
| 94 | + | ||
| 95 | + | ||
| 96 | +def _prod(seq): | ||
| 97 | + p = 1 | ||
| 98 | + for v in seq: | ||
| 99 | + p *= int(v) | ||
| 100 | + return p | ||
| 101 | + | ||
| 102 | + | ||
| 103 | +def _sigmoid(x): | ||
| 104 | + with np.errstate(over="ignore", invalid="ignore"): | ||
| 105 | + return 1.0 / (1.0 + np.exp(-x)) | ||
| 106 | + | ||
| 107 | + | ||
| 108 | +def _swiglu( | ||
| 109 | + x_fp32, dim_pos, swiglu_mode, activate_left, clamp_limit, glu_alpha, glu_bias | ||
| 110 | +): | ||
| 111 | + """Compute SwiGLU activation (modes 0-3), return fp32 result.""" | ||
| 112 | + pre = _prod(x_fp32.shape[:dim_pos]) if dim_pos > 0 else 1 | ||
| 113 | + cut = _prod(x_fp32.shape[dim_pos:]) | ||
| 114 | + xf = x_fp32.reshape(pre, cut).astype(np.float32) | ||
| 115 | + h = cut // 2 | ||
| 116 | + | ||
| 117 | + if swiglu_mode == 0: | ||
| 118 | + a = xf[:, :h] | ||
| 119 | + b = xf[:, h:] | ||
| 120 | + if activate_left: | ||
| 121 | + res = _sigmoid(a) * a * b | ||
| 122 | + else: | ||
| 123 | + res = _sigmoid(b) * a * b | ||
| 124 | + elif swiglu_mode == 1: | ||
| 125 | + a = xf[:, 0::2] | ||
| 126 | + b = xf[:, 1::2] | ||
| 127 | + a = np.clip(a, None, clamp_limit) | ||
| 128 | + b = np.clip(b, -clamp_limit, clamp_limit) | ||
| 129 | + res = a * _sigmoid(glu_alpha * a) * (b + glu_bias) | ||
| 130 | + elif swiglu_mode == 2: | ||
| 131 | + a = xf[:, :h] | ||
| 132 | + b = xf[:, h:] | ||
| 133 | + if not activate_left: | ||
| 134 | + a, b = b, a | ||
| 135 | + a = np.clip(a, None, clamp_limit) | ||
| 136 | + b = np.clip(b, -clamp_limit, clamp_limit) | ||
| 137 | + res = a * _sigmoid(glu_alpha * a) * (b + glu_bias) | ||
| 138 | + elif swiglu_mode == 3: | ||
| 139 | + a = xf[:, :h] | ||
| 140 | + b = xf[:, h:] | ||
| 141 | + if not activate_left: | ||
| 142 | + a, b = b, a | ||
| 143 | + a = a * _sigmoid(a) | ||
| 144 | + a = np.clip(a, None, clamp_limit) | ||
| 145 | + b = np.clip(b, -clamp_limit, clamp_limit) | ||
| 146 | + res = a * b | ||
| 147 | + else: | ||
| 148 | + raise ValueError(f"Unsupported swiglu_mode: {swiglu_mode}") | ||
| 149 | + | ||
| 150 | + y = np.zeros((pre, h), dtype=np.float32) | ||
| 151 | + y[:] = res.astype(np.float32) | ||
| 152 | + out_shape = list(x_fp32.shape) | ||
| 153 | + out_shape[dim_pos] = out_shape[dim_pos] // 2 | ||
| 154 | + y = y.reshape(out_shape) | ||
| 155 | + return y | ||
| 156 | + | ||
| 157 | + | ||
| 158 | +def _mx_quantize(data_fp32, axis_pos, dst_type, block_size, round_mode, scale_alg): | ||
| 159 | + """Dynamic MX quantization, returns (quantized_y, mxscale).""" | ||
| 160 | + dst_name, dst_dtype, emax = _DST_TYPE_MAP[dst_type] | ||
| 161 | + | ||
| 162 | + shape = list(data_fp32.shape) | ||
| 163 | + pre_q = _prod(shape[:axis_pos]) if axis_pos > 0 else 1 | ||
| 164 | + q_dim = shape[axis_pos] | ||
| 165 | + post_q = _prod(shape[axis_pos + 1 :]) if axis_pos + 1 < len(shape) else 1 | ||
| 166 | + flat = data_fp32.reshape(pre_q, q_dim, post_q) | ||
| 167 | + | ||
| 168 | + n_blocks = math.ceil(q_dim / block_size) | ||
| 169 | + y_flat = np.zeros((pre_q, q_dim, post_q), dtype=np.float32) | ||
| 170 | + scale_flat = np.zeros((pre_q, n_blocks, post_q), dtype=np.float32) | ||
| 171 | + | ||
| 172 | + for b in range(n_blocks): | ||
| 173 | + start = b * block_size | ||
| 174 | + end = min(start + block_size, q_dim) | ||
| 175 | + chunk = flat[:, start:end, :] | ||
| 176 | + pad_len = block_size - (end - start) | ||
| 177 | + if pad_len > 0: | ||
| 178 | + chunk = np.pad(chunk, ((0, 0), (0, pad_len), (0, 0)), mode="constant") | ||
| 179 | + | ||
| 180 | + abs_max = np.max(np.abs(chunk), axis=1, keepdims=True) | ||
| 181 | + abs_max = np.where(abs_max == 0, 1.0, abs_max) | ||
| 182 | + | ||
| 183 | + shared_exp = np.floor(np.log2(abs_max)) - emax | ||
| 184 | + shared_exp = np.where(shared_exp < 0, 0, shared_exp) | ||
| 185 | + mxscale = np.power(2.0, shared_exp.astype(np.int32).astype(np.float32)) | ||
| 186 | + | ||
| 187 | + scaled = chunk / np.where(mxscale == 0, 1.0, mxscale) | ||
| 188 | + | ||
| 189 | + if round_mode == "floor": | ||
| 190 | + scaled_q = np.floor(scaled) | ||
| 191 | + elif round_mode == "round": | ||
| 192 | + scaled_q = np.round(scaled) | ||
| 193 | + else: | ||
| 194 | + scaled_q = np.rint(scaled) | ||
| 195 | + | ||
| 196 | + actual_len = end - start | ||
| 197 | + y_flat[:, start:end, :] = scaled_q[:, :actual_len, :] | ||
| 198 | + scale_flat[:, b, :] = mxscale[:, 0, :] | ||
| 199 | + | ||
| 200 | + y_out = y_flat.reshape(shape) | ||
| 201 | + scale_shape = list(shape) | ||
| 202 | + scale_shape[axis_pos] = n_blocks | ||
| 203 | + scale_shape.append(2) | ||
| 204 | + scale_out = np.zeros(scale_shape, dtype=np.float32) | ||
| 205 | + | ||
| 206 | + if post_q == 1: | ||
| 207 | + scale_out[:, :, 0, 0] = scale_flat[:, :, 0] | ||
| 208 | + scale_out[:, :, 1, 0] = scale_flat[:, :, 0] | ||
| 209 | + else: | ||
| 210 | + for i in range(2): | ||
| 211 | + scale_out[:, :, :, i] = scale_flat[:, :, :] | ||
 🟡 Medium Priority 变更行：golden.py 第 200-211 行（新增的 _mx_quantize 的 scale 输出构造部分）。 影响行为： 失败模式（对默认/主路径 axis=-1 必然触发）：
由于 axis=-1 是默认属性、且 mode 2/3 仅允许 axis=-1，该 golden 在主要测试路径上直接崩溃，无法产出参考输出。 建议：将 scale_flat 先 reshape 到 scale_shape[:-1] 再写入 scale_out 的两个 trailing 切片，统一处理 post_q==1 与 post_q>1 两种情况，避免按展开形状索引未展开数组。   不准确？ | |||
| 212 | + | ||
| 213 | + return y_out, scale_out, dst_dtype | ||
| 214 | + | ||
| 215 | + | ||
| 216 | +def __golden_swiglu_mx_quant(*input_arrays, **kwargs): | ||
| 217 | + x = np.asarray(input_arrays[0]) | ||
| 218 | + group_index = None | ||
| 219 | + if len(input_arrays) > 1 and input_arrays[1] is not None: | ||
| 220 | + group_index = np.asarray(input_arrays[1]) | ||
| 221 | + | ||
| 222 | + activate_dim = int(kwargs.get("activate_dim", -1)) | ||
| 223 | + activate_left = bool(kwargs.get("activate_left", False)) | ||
| 224 | + swiglu_mode = int(kwargs.get("swiglu_mode", 0)) | ||
| 225 | + clamp_limit = float(kwargs.get("clamp_limit", 7.0)) | ||
| 226 | + glu_alpha = float(kwargs.get("glu_alpha", 1.702)) | ||
| 227 | + glu_bias = float(kwargs.get("glu_bias", 1.0)) | ||
| 228 | + axis = int(kwargs.get("axis", -1)) | ||
| 229 | + dst_type = int(kwargs.get("dst_type", 40)) | ||
| 230 | + round_mode = str(kwargs.get("round_mode", "rint")) | ||
| 231 | + scale_alg = int(kwargs.get("scale_alg", 0)) | ||
| 232 | + block_size = 32 | ||
| 233 | + | ||
| 234 | + output_dtypes = kwargs.get("output_dtypes") | ||
| 235 | + if output_dtypes is not None and len(output_dtypes) > 0: | ||
| 236 | + pass | ||
| 237 | + | ||
| 238 | + ndim = x.ndim | ||
| 239 | + dim_pos = activate_dim % ndim | ||
| 240 | + axis_pos = axis % ndim | ||
| 241 | + | ||
| 242 | + if "bfloat16" in str(x.dtype): | ||
| 243 | + x_fp32 = x.astype(np.float32) | ||
| 244 | + elif "float16" in str(x.dtype): | ||
| 245 | + x_fp32 = x.astype(np.float32) | ||
| 246 | + else: | ||
| 247 | + x_fp32 = x.astype(np.float32) | ||
| 248 | + | ||
| 249 | + swiglu_result = _swiglu( | ||
| 250 | + x_fp32, dim_pos, swiglu_mode, activate_left, clamp_limit, glu_alpha, glu_bias | ||
| 251 | + ) | ||
| 252 | + | ||
| 253 | + if "bfloat16" in str(x.dtype): | ||
| 254 | + swiglu_result = ( | ||
| 255 | + swiglu_result.astype(_bf16) if _bf16 is not None else swiglu_result | ||
| 256 | + ) | ||
| 257 | + elif "float16" in str(x.dtype): | ||
| 258 | + swiglu_result = swiglu_result.astype(np.float16) | ||
| 259 | + | ||
| 260 | + swiglu_fp32 = swiglu_result.astype(np.float32) | ||
| 261 | + | ||
| 262 | + if group_index is not None: | ||
| 263 | + y_shape = list(swiglu_fp32.shape) | ||
| 264 | + scale_shape = list(swiglu_fp32.shape) | ||
| 265 | + scale_shape[axis_pos] = math.ceil(scale_shape[axis_pos] / block_size) | ||
| 266 | + scale_shape.append(2) | ||
| 267 | + | ||
| 268 | + _, dst_dtype, _ = _DST_TYPE_MAP[dst_type] | ||
| 269 | + if dst_dtype is not None: | ||
| 270 | + y = np.zeros(y_shape, dtype=dst_dtype) | ||
| 271 | + else: | ||
| 272 | + y = np.zeros(y_shape, dtype=np.float32) | ||
| 273 | + scale = np.zeros( | ||
| 274 | + scale_shape, dtype=_fp8_e8m0 if _fp8_e8m0 is not None else np.float32 | ||
| 275 | + ) | ||
| 276 | + | ||
| 277 | + start = 0 | ||
| 278 | + for gv in group_index: | ||
| 279 | + gv = int(gv) | ||
| 280 | + y_part, scale_part, dst_dt = _mx_quantize( | ||
| 281 | + swiglu_fp32[start : start + gv], | ||
| 282 | + axis_pos, | ||
| 283 | + dst_type, | ||
| 284 | + block_size, | ||
| 285 | + round_mode, | ||
| 286 | + scale_alg, | ||
| 287 | + ) | ||
| 288 | + if dst_dtype is not None: | ||
| 289 | + y[start : start + gv] = y_part.astype(dst_dtype) | ||
| 290 | + else: | ||
| 291 | + y[start : start + gv] = y_part | ||
| 292 | + scale[start : start + gv] = scale_part.astype( | ||
| 293 | + _fp8_e8m0 if _fp8_e8m0 is not None else np.float32 | ||
| 294 | + ) | ||
| 295 | + start += gv | ||
| 296 | + else: | ||
| 297 | + y_np, scale_np, dst_dtype = _mx_quantize( | ||
| 298 | + swiglu_fp32, axis_pos, dst_type, block_size, round_mode, scale_alg | ||
| 299 | + ) | ||
| 300 | + if dst_dtype is not None: | ||
| 301 | + y = y_np.astype(dst_dtype) | ||
| 302 | + else: | ||
| 303 | + y = y_np | ||
| 304 | + scale = scale_np.astype(_fp8_e8m0 if _fp8_e8m0 is not None else np.float32) | ||
| 305 | + | ||
| 306 | + y = np.nan_to_num(y, nan=0.0) | ||
| 307 | + scale = np.nan_to_num(scale, nan=0.0) | ||
| 308 | + return [y, scale] | ||
| 309 | + | ||
| 310 | + | ||
| 311 | +__golden__ = {"kernel": {"swiglu_mx_quant": "__golden_swiglu_mx_quant"}} | ||
🟠 High Priority
changed line: golden.py 新增的
_mx_quantize中scale_out组装逻辑（206-211 行）。受影响行为/契约：
scale_flat始终是扁平的三维(pre_q, n_blocks, post_q)，而scale_out的形状是shape把axis_pos维替换为n_blocks后再append(2)，即ndim+1维；两者维数/形状根本没有对齐，赋值语句却按固定维数写死。失败模式：
即该 golden 在算子最常用配置（axis=-1，rank 2/3，以及 group_index 分组路径 280-295 行同样调用本函数）下必然崩溃，无法生成参考输出，新增标杆完全不可用。
建议：将 scale_flat 先 reshape 回原始维度（axis_pos 处为 n_blocks）再写入 scale_out 的最后一个维度（两个槽位写同一份 scale），替换整个 if/else 分支。