Star1092
443
代码
代码
Issues71
Pull Requests93
Actions
项目讨论
Wiki
分析
Star1092
443
已合并
适配mxfp4图模式 #3473
wuziyu创建于 7月30日
适配mxfp4图模式 #3473
已合并
wuziyu创建于 7月30日
从wuziyu/torchairgraph合入到Ascend/torchairmaster
共 2 个文件变更+133-14
Mpython/torchair/_ge_concrete_graph/auto_generated_ge_raw_ops.py+3-3
@@ -88653,14 +88653,14 @@ def AllGatherMatmulV2( # pylint: disable=dangerous-default-value
88653):88653):
88654 """88654 """
88655 REG_OP(AllGatherMatmulV2)\n88655 REG_OP(AllGatherMatmulV2)\n
88656- .INPUT(x1, TensorType({DT_FLOAT16, DT_BF16, DT_FLOAT8_E4M3FN, DT_FLOAT8_E5M2, DT_HIFLOAT8}))\n88656+ .INPUT(x1, TensorType({DT_FLOAT16, DT_BF16, DT_FLOAT8_E4M3FN, DT_FLOAT8_E5M2, DT_HIFLOAT8, DT_FLOAT4_E2M1}))\n
88657- .INPUT(x2, TensorType({DT_FLOAT16, DT_BF16, DT_FLOAT8_E4M3FN, DT_FLOAT8_E5M2, DT_HIFLOAT8}))\n88657+ .INPUT(x2, TensorType({DT_FLOAT16, DT_BF16, DT_FLOAT8_E4M3FN, DT_FLOAT8_E5M2, DT_HIFLOAT8, DT_FLOAT4_E2M1}))\n
88658 .OPTIONAL_INPUT(bias, TensorType({DT_FLOAT16, DT_BF16, DT_FLOAT}))\n88658 .OPTIONAL_INPUT(bias, TensorType({DT_FLOAT16, DT_BF16, DT_FLOAT}))\n
88659 .OPTIONAL_INPUT(x1_scale, TensorType({DT_FLOAT, DT_FLOAT8_E8M0}))\n88659 .OPTIONAL_INPUT(x1_scale, TensorType({DT_FLOAT, DT_FLOAT8_E8M0}))\n
88660 .OPTIONAL_INPUT(x2_scale, TensorType({DT_FLOAT, DT_FLOAT8_E8M0}))\n88660 .OPTIONAL_INPUT(x2_scale, TensorType({DT_FLOAT, DT_FLOAT8_E8M0}))\n
88661 .OPTIONAL_INPUT(quant_scale, TensorType({DT_FLOAT}))\n88661 .OPTIONAL_INPUT(quant_scale, TensorType({DT_FLOAT}))\n
88662 .OUTPUT(y, TensorType({DT_FLOAT16, DT_BF16, DT_FLOAT}))\n88662 .OUTPUT(y, TensorType({DT_FLOAT16, DT_BF16, DT_FLOAT}))\n
88663- .OUTPUT(gather_out, TensorType({DT_FLOAT16, DT_BF16, DT_FLOAT8_E4M3FN, DT_FLOAT8_E5M2, DT_HIFLOAT8}))\n88663+ .OUTPUT(gather_out, TensorType({DT_FLOAT16, DT_BF16, DT_FLOAT8_E4M3FN, DT_FLOAT8_E5M2, DT_HIFLOAT8, DT_FLOAT4_E2M1}))\n
88664 .OUTPUT(amax_out, TensorType({DT_FLOAT}))\n88664 .OUTPUT(amax_out, TensorType({DT_FLOAT}))\n
88665 .REQUIRED_ATTR(group, String)\n88665 .REQUIRED_ATTR(group, String)\n
88666 .ATTR(is_trans_a, Bool, false)\n88666 .ATTR(is_trans_a, Bool, false)\n
Mpython/torchair/_ge_concrete_graph/ge_converter/custom/all_gather_matmul_v2.py+130-11
@@ -45,8 +45,9 @@ from torchair._ge_concrete_graph.ge_converter.converter_utils import * # pylint
45 45 
46 46 
47X1_X2_SAME_TYPES = {DataType.DT_FLOAT16, DataType.DT_BF16, DataType.DT_HIFLOAT8}47X1_X2_SAME_TYPES = {DataType.DT_FLOAT16, DataType.DT_BF16, DataType.DT_HIFLOAT8}
48- 
49FP8_TYPES = {DataType.DT_HIFLOAT8, DataType.DT_FLOAT8_E4M3FN, DataType.DT_FLOAT8_E5M2}48FP8_TYPES = {DataType.DT_HIFLOAT8, DataType.DT_FLOAT8_E4M3FN, DataType.DT_FLOAT8_E5M2}
49+GROUP_SIZE_MAX_VALUE = 65535
50+WORLD_SIZE_SUPPORT_LIST = [2, 4, 8, 16, 32, 64]
50 51 
51 52 
52@declare_supported(53@declare_supported(
@@ -108,9 +109,19 @@ def convert_npu_all_gather_quant_mm(
108 Tensor? quant_scale=None, int block_size=0, int gather_index=0, bool gather_output=True, int comm_turn=0,109 Tensor? quant_scale=None, int block_size=0, int gather_index=0, bool gather_output=True, int comm_turn=0,
109 int[]? group_sizes=None, bool amax_output=False, int? y_dtype=None,110 int[]? group_sizes=None, bool amax_output=False, int? y_dtype=None,
110 int? x1_dtype=None, int? x2_dtype=None, int? x1_scale_dtype=None,111 int? x1_dtype=None, int? x2_dtype=None, int? x1_scale_dtype=None,
111- int? x2_scale_dtype=None) -> (Tensor, Tensor, Tensor)112+ int? x2_scale_dtype=None, str? comm_mode=None) -> (Tensor, Tensor, Tensor)
112 '''113 '''
113- check_dtype(self, x2, y_dtype=y_dtype)114+ check_dtype(
115+ self,
116+ x2,
117+ x1_dtype,
118+ x2_dtype,
119+ x1_scale_dtype,
120+ x2_scale_dtype,
121+ y_dtype=y_dtype,
122+ )
123+ if world_size not in WORLD_SIZE_SUPPORT_LIST:
124+ raise RuntimeError(f"world_size only supports value in {WORLD_SIZE_SUPPORT_LIST}, but got {world_size}.")
114 group_max = GROUP_SIZE_MAX_VALUE # 65535是指group_size中的值最大不能超过16位可表示的范围125 group_max = GROUP_SIZE_MAX_VALUE # 65535是指group_size中的值最大不能超过16位可表示的范围
115 group_size = 0126 group_size = 0
116 if group_sizes is not None and isinstance(group_sizes, List):127 if group_sizes is not None and isinstance(group_sizes, List):
@@ -134,12 +145,42 @@ def convert_npu_all_gather_quant_mm(
134 )145 )
135 if y_dtype is None:146 if y_dtype is None:
136 y_dtype = 1147 y_dtype = 1
137- if x1_dtype is not None:148+ if self.rank < 2:
138- self = ge.Bitcast(self, type=torch_dtype_value_to_ge_type(x1_dtype))149+ raise RuntimeError("x1 dimension cannot be less than 2, but actual x1 dimension is: " + str(self.rank) + ".")
139- self.desc.dtype = torch_dtype_value_to_ge_proto_type(x1_dtype)150+ if x2.rank < 2:
140- if x2_dtype is not None:151+ raise RuntimeError("x2 dimension cannot be less than 2, but actual x2 dimension is: " + str(x2.rank) + ".")
141- x2 = ge.Bitcast(x2, type=torch_dtype_value_to_ge_type(x2_dtype))152+ if self.dtype in [DataType.DT_UINT8] and x2.dtype in [DataType.DT_UINT8]:
142- x2.desc.dtype = torch_dtype_value_to_ge_proto_type(x2_dtype)153+ is_mxfp4 = torch_npu.float4_e2m1fn_x2 in (x1_dtype, x2_dtype)
154+ trans_x2 = self.symsize[-1] == x2.symsize[-2]
155+ x1_ge_proto_dtype = torch_dtype_value_to_ge_proto_type(x1_dtype)
156+ x2_ge_proto_dtype = torch_dtype_value_to_ge_proto_type(x2_dtype)
157+ if is_mxfp4:
158+ factor = 2
159+ self_const = ge.Const([1] * (self.rank - 1) + [factor])
160+ self_shape = ge.Shape(self)
161+ self_new_shape = ge.Mul(self_shape, self_const)
162+ self = ge.Bitcast(self, type=torch_dtype_value_to_ge_type(x1_dtype))
163+ self.desc.dtype = x1_ge_proto_dtype
164+ self = ge.Reshape(self, self_new_shape)
165+ 
166+ if trans_x2:
167+ perm = list(range(x2.rank))
168+ perm[-2], perm[-1] = perm[-1], perm[-2]
169+ x2_const = ge.Const([1] * (x2.rank - 1) + [factor])
170+ if trans_x2:
171+ x2 = ge.Transpose(x2, perm)
172+ x2_shape = ge.Shape(x2)
173+ x2_new_shape = ge.Mul(x2_shape, x2_const)
174+ x2 = ge.Bitcast(x2, type=torch_dtype_value_to_ge_type(x2_dtype))
175+ x2.desc.dtype = x2_ge_proto_dtype
176+ x2 = ge.Reshape(x2, x2_new_shape)
177+ if trans_x2:
178+ x2 = ge.Transpose(x2, perm)
179+ else:
180+ self = ge.Bitcast(self, type=torch_dtype_value_to_ge_type(x1_dtype))
181+ self.desc.dtype = x1_ge_proto_dtype
182+ x2 = ge.Bitcast(x2, type=torch_dtype_value_to_ge_type(x2_dtype))
183+ x2.desc.dtype = x2_ge_proto_dtype
143 if x1_scale_dtype is not None:184 if x1_scale_dtype is not None:
144 x1_scale = ge.Bitcast(x1_scale, type=torch_dtype_value_to_ge_type(x1_scale_dtype))185 x1_scale = ge.Bitcast(x1_scale, type=torch_dtype_value_to_ge_type(x1_scale_dtype))
145 x1_scale.desc.dtype = torch_dtype_value_to_ge_proto_type(x1_scale_dtype)186 x1_scale.desc.dtype = torch_dtype_value_to_ge_proto_type(x1_scale_dtype)
@@ -168,13 +209,91 @@ def convert_npu_all_gather_quant_mm(
168 )209 )
169 # 对于非原生的torch数据类型需要做类型标注210 # 对于非原生的torch数据类型需要做类型标注
170 out.desc.dtype = torch_dtype_value_to_ge_proto_type(y_dtype)211 out.desc.dtype = torch_dtype_value_to_ge_proto_type(y_dtype)
171- gather_out_dtype = DataType.DT_HIFLOAT8 if x1_dtype == torch_npu.hifloat8 else self.dtype212+ gather_out_dtype = self.dtype
213+ if x1_dtype is not None:
214+ gather_out_dtype = torch_dtype_value_to_ge_type(x1_dtype)
172 gather_out.desc.dtype = _ge_dtype_to_ge_proto_dtype(gather_out_dtype)215 gather_out.desc.dtype = _ge_dtype_to_ge_proto_dtype(gather_out_dtype)
216+ if gather_output is True:
217+ if gather_out_dtype in (DataType.DT_FLOAT4_E2M1, DataType.DT_FLOAT4_E1M2):
218+ bit_shape = [1, 2]
219+ div_x2 = ge.Cast(ge.Const(bit_shape), dst_type=DataType.DT_INT32)
220+ y_shape_int4 = ge.Shape(gather_out)
221+ y_shape_uint8 = ge.Div(y_shape_int4, div_x2)
222+ y_shape_int4_2bit = ge.ConcatV2(
223+ [y_shape_uint8, ge.Cast(ge.Const([2]), dst_type=DataType.DT_INT32)], concat_dim=0, N=2
224+ )
225+ gather_out = ge.Bitcast(ge.Reshape(gather_out, y_shape_int4_2bit), type=DataType.DT_UINT8)
226+ gather_out = ge.Reshape(gather_out, y_shape_uint8)
227+ gather_out.desc.dtype = torch_dtype_value_to_ge_proto_type(DataType.DT_UINT8)
atomgit-bot
atomgit-botatomgit-bot7月30日

🟠 High Priority

第 233 行：gather_out.desc.dtype = torch_dtype_value_to_ge_proto_type(DataType.DT_UINT8) 中，DataType.DT_UINT8 是 GE 侧的枚举值（其整数为 4），而 torch_dtype_value_to_ge_proto_type 期望接收的是 torch dtype 整数值（torch.uint8 = 0）。传入 4 会命中 torch_dtype_value_to_ge_type 中 dtype_value == 4 的分支，返回 ProtoDataType.DT_INT64（对应 torch.int64），而非预期的 ProtoDataType.DT_UINT8。

失败模式：当 mxfp4 量化场景且 gather_output=True 时（第 222-233 行路径），gather_out 被 reshape 为 uint8 存储，但其 .desc.dtype 被错误标注为 INT64。这会导致下游图编译或执行时类型不匹配，可能引发 GE 图编译失败或静默数据错误。

正确的调用应为 _ge_dtype_to_ge_proto_dtype(DataType.DT_UINT8)，该函数接收 GE DataType 并返回对应的 ProtoDataType（已确认 _ge_dtype_to_ge_proto_dtype(DataType.DT_UINT8) → ProtoDataType.DT_UINT8，第 376-377 行）。

建议：将 torch_dtype_value_to_ge_proto_type(DataType.DT_UINT8) 替换为 _ge_dtype_to_ge_proto_dtype(DataType.DT_UINT8)，因为 DataType.DT_UINT8 是 GE DataType 枚举，应使用 GE→Proto 的转换函数而非 torch→Proto 的转换函数。

改动建议
227
- gather_out.desc.dtype = torch_dtype_value_to_ge_proto_type(DataType.DT_UINT8)
227
+ gather_out.desc.dtype = _ge_dtype_to_ge_proto_dtype(DataType.DT_UINT8)
应用建议
likedislike
不准确？
173 return (out, gather_out, amax_out)228 return (out, gather_out, amax_out)
174 229 
175 230 
176-def check_dtype(x1: Tensor, x2: Tensor, y_dtype: int = None):231+_GE_DTYPE_TO_NAME = {
232+ DataType.DT_FLOAT: "float32",
233+ DataType.DT_FLOAT16: "float16",
234+ DataType.DT_INT8: "int8",
235+ DataType.DT_INT16: "int16",
236+ DataType.DT_UINT16: "uint16",
237+ DataType.DT_UINT8: "uint8",
238+ DataType.DT_INT32: "int32",
239+ DataType.DT_INT64: "int64",
240+ DataType.DT_UINT32: "uint32",
241+ DataType.DT_UINT64: "uint64",
242+ DataType.DT_BOOL: "bool",
243+ DataType.DT_DOUBLE: "float64",
244+ DataType.DT_COMPLEX32: "complex32",
245+ DataType.DT_COMPLEX64: "complex64",
246+ DataType.DT_COMPLEX128: "complex128",
247+ DataType.DT_QINT8: "qint8",
248+ DataType.DT_QINT32: "qint32",
249+ DataType.DT_QUINT8: "quint8",
250+ DataType.DT_BF16: "bfloat16",
251+ DataType.DT_INT4: "int4",
252+ DataType.DT_HIFLOAT8: "hifloat8",
253+ DataType.DT_FLOAT8_E5M2: "float8_e5m2",
254+ DataType.DT_FLOAT8_E4M3FN: "float8_e4m3fn",
255+ DataType.DT_FLOAT8_E8M0: "float8_e8m0",
256+ DataType.DT_FLOAT4_E2M1: "float4_e2m1",
257+ DataType.DT_FLOAT4_E1M2: "float4_e1m2",
258+}
259+ 
260+ 
261+def check_dtype(
262+ x1: Tensor,
263+ x2: Tensor,
264+ x1_dtype: int = None,
265+ x2_dtype: int = None,
266+ x1_scale_dtype: int = None,
267+ x2_scale_dtype: int = None,
268+ y_dtype: int = None,
269+):
177 if (x1.dtype != x2.dtype) and x1.dtype in X1_X2_SAME_TYPES:270 if (x1.dtype != x2.dtype) and x1.dtype in X1_X2_SAME_TYPES:
178 raise AssertionError(f"Type of x1:{x1.dtype} and x2:{x2.dtype} must be same.")271 raise AssertionError(f"Type of x1:{x1.dtype} and x2:{x2.dtype} must be same.")
179 if x1.dtype in FP8_TYPES and y_dtype is None:272 if x1.dtype in FP8_TYPES and y_dtype is None:
180 raise AssertionError(f"When type of x1 is {x1.dtype} should input y_dtype.")273 raise AssertionError(f"When type of x1 is {x1.dtype} should input y_dtype.")
274+ if x1.dtype in [DataType.DT_UINT8] and x2.dtype in [DataType.DT_UINT8]:
275+ if x1_dtype is None:
276+ raise AssertionError("In this quant scene: x1_dtype cannot be none.")
277+ if x2_dtype is None:
278+ raise AssertionError("In this quant scene: x2_dtype cannot be none.")
279+ x1_ge_dtype = torch_dtype_value_to_ge_type(x1_dtype)
280+ x2_ge_dtype = torch_dtype_value_to_ge_type(x2_dtype)
281+ is_mxfp4 = DataType.DT_FLOAT4_E2M1 in (x1_ge_dtype, x2_ge_dtype)
282+ if is_mxfp4:
283+ if x1_scale_dtype is None:
284+ raise AssertionError("In mxfp4 quant scene: x1_scale_dtype cannot be none.")
285+ if x2_scale_dtype is None:
286+ raise AssertionError("In mxfp4 quant scene: x2_scale_dtype cannot be none.")
287+ x1scale_ge_dtype = torch_dtype_value_to_ge_type(x1_scale_dtype)
288+ x2scale_ge_dtype = torch_dtype_value_to_ge_type(x2_scale_dtype)
289+ if x1_ge_dtype != DataType.DT_FLOAT4_E2M1 or x2_ge_dtype != DataType.DT_FLOAT4_E2M1:
290+ raise AssertionError(
291+ f"In mxfp4 quant scene: the data types of x1 and x2 must both be float4_e2m1, "
292+ f"but got x1_dtype:{_GE_DTYPE_TO_NAME.get(x1_ge_dtype)} and x2_dtype:{_GE_DTYPE_TO_NAME.get(x2_ge_dtype)}."
293+ )
294+ if x1scale_ge_dtype != DataType.DT_FLOAT8_E8M0 or x2scale_ge_dtype != DataType.DT_FLOAT8_E8M0:
295+ raise AssertionError(
296+ f"In mxfp4 quant scene: the data types of x1scale and x2scale must both be float8_e8m0, "
297+ f"but got x1_scale_dtype:{_GE_DTYPE_TO_NAME.get(x1scale_ge_dtype)} "
298+ f"and x2_scale_dtype:{_GE_DTYPE_TO_NAME.get(x2scale_ge_dtype)}."
299+ )