Pull Request已成功合入, 合并人@ascend-robot
（感谢 wuziyu 的贡献）
变更摘要
本次 PR 主要对 convert_npu_all_gather_quant_mm 函数及其辅助函数 check_dtype 进行扩展，以适配 mxfp4（MXFP4，即 float4_e2m1fn_x2）量化图模式。核心改动包括：新增 mxfp4 场景下对输入 x1/x2 的 Bitcast + Reshape（factor=2）处理逻辑、对 gather_out 在 float4 数据类型下的后处理（转换为 DT_UINT8）、增强 check_dtype 的校验规则（校验 x1/x2 dtype 必须为 float4_e2m1、scale dtype 必须为 float8_e8m0），以及新增 world_size 合法范围检查和输入张量 rank 校验。
主要改动
-
新增 mxfp4 输入处理逻辑: 当
x1和x2的 dtype 均为DT_UINT8且目标 dtype 为 mxfp4（float4_e2m1fn_x2）时，对x1/x2执行 Bitcast 后再按 factor=2 的 ratio 进行 Reshape，以正确将 UINT8 存储 reinterpret 为 float4 数据类型；若x2需要转置（trans_x2），则在 reshape 前后分别做 Transpose。 -
新增 gather_out 的 float4 转 UINT8 后处理: 当
gather_output=True且gather_out_dtype为DT_FLOAT4_E2M1或DT_FLOAT4_E1M2时，通过 Div/ConcatV2/Reshape/Bitcast 将输出从 int4 2bit 格式转换为DT_UINT8格式，确保 gather 输出的 dtype 正确。 -
重构 gather_out_dtype 的确定方式: 原逻辑
gather_out_dtype = DT_HIFLOAT8 if x1_dtype == torch_npu.hifloat8 else self.dtype被替换为通用的torch_dtype_value_to_ge_type(x1_dtype)，使 gather_out 的 dtype 推导不再仅硬编码 hifloat8 特例。 -
增强
check_dtype校验逻辑:check_dtype新增x1_dtype/x2_dtype/x1_scale_dtype/x2_scale_dtype参数，针对 mxfp4 量化场景强制要求 x1/x2 dtype 均为float4_e2m1、scale dtype 均为float8_e8m0，并确保这些参数不为 None。 -
新增输入合法性校验: 增加
world_size必须在[2, 4, 8, 16, 32, 64]范围内的检查；增加x1和x2的 rank 必须 ≥ 2 的检查；函数签名新增可选参数comm_mode: str?。
代码审查
审查总结
审查了 1 个变更文件：python/torchair/_ge_concrete_graph/ge_converter/custom/all_gather_matmul_v2.py（+105 / -10 行）。
| 优先级 | 数量 | 说明 |
|---|---|---|
| P1 | 1 | 第 233 行：torch_dtype_value_to_ge_proto_type 被错误传入 GE DataType（DataType.DT_UINT8=4）而非 torch dtype（torch.uint8=0），导致 gather_out.desc.dtype 被错误设为 DT_INT64 而非 DT_UINT8，在 mxfp4 量化 + gather_output=True 路径下会引发类型标注错误 |
| P2 | 2 | 第 159 行和第 257 行：is_mxfp4 判定使用 OR 而非 AND，与 check_dtype 内部的双重校验（第 265 行）及代码库其他并行实现（quantmatmul.py、matmul_reduce_scatter_v2.py）不一致，存在未来重构时的逻辑脆弱性 |
| P3 | 1 | 第 53-75 行：@declare_supported 测试声明中缺少新增的 comm_mode 参数，测试覆盖不完整 |
整体风险评估：中等风险。 P1 问题会影响 mxfp4 量化场景下 gather_output 的类型正确性，可能导致下游 GE 图编译失败或静默错误，建议优先修复。P2 问题当前有安全网保护不会触发，但建议与 P1 一同修复以保持一致性。
| 类型 | 数量 |
|---|---|
| 🔴 阻塞 | 1 |
| 🟡 建议 | 0 |
⛔ 需要修改
Pull Request 已合并或已关闭。
If you want to solve this problem, you can click here to do it in the FAQs.
Pull Request 已合并或已关闭。
If you want to solve this problem, you can click here to do it in the FAQs.
Pull Request 已合并或已关闭。
If you want to solve this problem, you can click here to do it in the FAQs.
Pull Request 已合并或已关闭。
If you want to solve this problem, you can click here to do it in the FAQs.
Pull Request 已合并或已关闭。
If you want to solve this problem, you can click here to do it in the FAQs.
适配mxfp4图模式
https://gitcode.com/Ascend/torchair/issues/617