[A5 TQuant]: Adding support for 2D DN MXFP4 Quantization-pto-isa-AtomGit

已合并

[A5 TQuant]: Adding support for 2D DN MXFP4 Quantization #1187

omarzohir创建于 6月25日从omarzohir/pto-isa： tquant_dn_mxfp4_2 合入到 cann/pto-isa： master

[A5 TQuant]: Adding support for 2D DN MXFP4 Quantization #1187

已合并

omarzohir创建于 6月25日

从omarzohir/pto-isa tquant_dn_mxfp4_2合入到cann/pto-isa master

omarzohir

6月25日

描述

[A5 TQuant]: Adding support for 2D DN MXFP4 Quantization

关联的Issue

https://gitcode.com/cann/pto-isa/issues/429

测试

All st cases pass on 9.1.0 CANN, by running on board

文档更新

类型标签

Bug修复
新特性
性能优化
文档更新
其他，请描述：

Pull Request已成功合入, 合并人@CANN-robot

（感谢 omarzohir 的贡献）

CANN-robot成员

6月25日添加了label：cann-cla/yes

CANN-robot成员

6月25日评论：

CLA Signature Pass

omarzohir, thanks for your pull request. All authors of the commits have signed the CLA. 👍

CANN-robot成员

6月25日评论：

Thanks for your pull-request.
The full list of commands accepted by me can be found at here。
You can get sig-info at here

PR Approval Progress

✅ Congratulations! All modules have met the lgtm and approve requirements.

Module Approval Details

module	lgtm status	approve status
repo-cann/pto-isa	✅ Zhanghaijian, csjlchen (2/2)	✅ Zhanghaijian (1/1)

💡 Tip:

Committer can comment /approve or /lgtm

Commenting /approve implies both code review (lgtm) and intent to merge (approve)

CLA Signature Pass

omarzohir, thanks for your pull request. All authors of the commits have signed the CLA. 👍

omarzohir

6月25日评论：

/compile

atomgit-bot

6月25日评论：

变更摘要

此 PR 为 A5 TQuant 模块新增了 2D DN（Decomposed interleaved N）模式下的 MXFP4 E2M1 量化支持，覆盖 BF16 和 FP16 两种输入数据类型。变更涉及类型系统扩展、指令接口统一、量化核心管线实现（含 ND 和 DN 两种分组轴模式），以及完整的系统测试（包含 golden 数据生成、内核实现及测试用例）。

主要改动

类型系统扩展：在 QuantType 枚举中新增 MXFP4_E2M1，在 MxQuantAlg 统一枚举中新增 OcpMxFp4E2M1 和 NvMxFp4E2M1，为 MXFP4 量化提供类型标签和算法分派基础。
统一指令入口：在 pto_instr.hpp 中新增 TQuant 模板重载（接受 grp_axis 和 mx_alg 参数），作为 MXFP8/MXFP4 在 ND/DN 模式下的统一调用入口。
核心量化管线：在 TQuant.hpp 中新增 MXFP4 全链路实现，包括 1D 平坦量化 TQuant_MXFP4_E2M1_B16、2D 分块量化 TQuant_MXFP4_E2M1_B16_2D、DN 模式三阶段管线 TQuant_MXFP4_E2M1_DN（含 BF16 和 FP16 各自的 stage-3 量化函数），以及对应的 _Impl / _Impl_DN 分派实现；同时更新 CheckTQuantMxTypes 对 MXFP4 的输入输出类型进行编译期校验。
DN 模式测试覆盖：新增 gen_data.py 中 BF16/FP16 的 MXFP4 DN golden 数据生成逻辑（基于 ml_dtypes.float4_e2m1fn 独立验证），新增 tquant_dn_kernel.cpp 中 runTQuantDN_MXFP4 设备内核及 CompactFp4PackedRows 辅助函数，并在 main.cpp 中添加 6 个测试用例（BF16 和 FP16 各 3 种 shape：64×128、128×128、64×256）。

不准确？

CANN-robot成员

6月25日添加了label：ci-pipeline-running

CANN-robot成员

6月25日评论：

流水线任务触发成功
任务链接 [8814a05c45474da196de35739d0cdbe1][流水线指导]

任务名称	状态	日志	下载链接
codecheck	✅ SUCCESS	>>>>>
SCA	✅ SUCCESS	>>>>>
antipoison	✅ SUCCESS	>>>>>
Check_Pr	✅ SUCCESS	>>>>>
Compile_Ascend_X86	✅ SUCCESS	>>>>>	>>>>>
Compile_Ascend_ARM	✅ SUCCESS	>>>>>	>>>>>
pto-isa_check	✅ SUCCESS
pre_comment	✅ SUCCESS	>>>>>
Compile_Ascend_X86_ubuntu24	✅ SUCCESS	>>>>>	>>>>>
cocheck_codestyle	✅ SUCCESS
precommit	⚠️ WARNING	>>>>>	>>>>>
Compile_Ascend_ARM_A5	❌ FAILED	>>>>>
Compile_Ascend_X86_A5_ubuntu24	❌ FAILED	>>>>>
UT_Test	✅ SUCCESS

CANN-robot成员

6月25日删除了label：ci-pipeline-running

CANN-robot成员

6月25日添加了label：ci-pipeline-failed

CANN-robot成员

6月25日评论：

流水线 PR-pipeline_cann_pto-isa#4149 [ commitID：6121f0d5 ] 运行失败

atomgit-bot

6月25日评论：

代码审查

我已经完成了对所有6个变更文件的审查。以下是我的发现摘要：

审查总结

P2 问题：1 个（gen_data.py — NaN→+Inf 结果被丢弃）
P3 问题：1 个（gen_data.py — 误导性 docstring）
整体风险：低 — 核心逻辑正确，唯一的中等问题影响 fp16 golden 数据生成器在 NaN 输入下的正确性（当前测试数据不触发）。

各文件审查结果

文件	审查结论
`include/pto/common/pto_instr.hpp`	无问题 — 新增的 `TQuant` 包装函数正确转发至 `TQUANT_IMPL`
`include/pto/common/type.hpp`	无问题 — 新增的 `MxQuantAlg` 枚举正确编码了格式+算法组合
`include/pto/npu/a5/TQuant.hpp`	无问题 — FP4 DN 管线（reduce→exponent/scaling→quantize）正确实现了三阶段流水，并正确复用已有的 OCP 指数提取和 NaN 饱和逻辑
`tests/npu/a5/src/st/testcase/tquant_dn/gen_data.py`	1 个 P2 + 1 个 P3 — 见上
`tests/npu/a5/src/st/testcase/tquant_dn/main.cpp`	无问题 — MXFP4 测试用例正确分配 host/device 内存、加载 golden 输入、启动 kernel、与 golden 对比
`tests/npu/a5/src/st/testcase/tquant_dn/tquant_dn_kernel.cpp`	无问题 — MXFP4 DN kernel 正确设置了 UB 布局、tile 类型，并通过 `CompactFp4PackedRows` 桥接 flat FP4 tile 与 2D TSTORE tile

类型	数量
🔴 阻塞	0
🟡 建议	2

💬 仅评论

不准确？

atomgit-bot6月25日进行代码检视1

tests/npu/a5/src/st/testcase/tquant_dn/gen_data.py

381+def _gen_src_fp16_safe(m, n_pad):

382+ """fp16-safe variant of _gen_src: per-group max in [0.25, 16] (no *10000), so

383+ all values stay well within fp16 normal range. Uses the same RNG stream as

384+ _gen_src so bf16 tests (which use _gen_src) keep their existing goldens."""

atomgit-bot6月25日评论：

🔵 Low Priority

_gen_src_fp16_safe 的文档字符串声称 "Uses the same RNG stream as _gen_src so bf16 tests (which use _gen_src) keep their existing goldens."，但该函数并未调用 _gen_src 或复用其 RNG 状态，而是独立调用 np.random.uniform 生成随机数。该注释具有误导性。

实际上 bf16 golden 不受影响仅仅是因为 _gen_src_fp16_safe 在 __main__ 中的调用顺序排在所有 bf16 测试之后，而非因为它使用了"相同的 RNG 流"。

这是一个文档/注释准确性问题，不影响代码正确性。

建议：修正 docstring 为更准确的描述，例如："Generates fp16-safe data with per-group max in [0.25, 16] (no *10000 multiplier). Called after all bf16 cases so existing bf16 goldens are unaffected."

不准确？

atomgit-bot6月25日进行代码检视1

tests/npu/a5/src/st/testcase/tquant_dn/gen_data.py

258+ scaled_fp32 = scaled_bf16.astype(np.float32)

259+ with np.errstate(invalid="ignore"):

260+ scaled_fp32 = np.where(np.isnan(scaled_fp32), np.float32(np.inf), scaled_fp32)

261+ fp4_codes = scaled_bf16.astype(float4_e2m1fn).view(np.uint8).astype(np.uint8)

atomgit-bot6月25日评论：

🟡 Medium Priority

在 quant_fp16_to_mxfp4_dn 函数中（文件 gen_data.py），NaN→+Inf 的处理被应用到变量 scaled_fp32 上，但随后转换为 fp4 时使用的是未经 NaN 处理的 scaled_bf16：

硬件路径 calcQuantizedFP4E2M1Values_DN_Fp16_Row 在 vcvt bf16→fp4 之前调用了 SaturateBf16NaNToPosInf。因此若输入包含 NaN，golden 与硬件输出将不一致。

触发条件：源数据中包含 NaN（当前测试数据 _gen_src_fp16_safe 生成的值在正常范围内，NaN 概率极低，故降低 confidence）。

改动建议

+scaled_fp32 = scaled_bf16.astype(np.float32)
+with np.errstate(invalid="ignore"):
+    scaled_fp32 = np.where(np.isnan(scaled_fp32), np.float32(np.inf), scaled_fp32)
-    fp4_codes = scaled_bf16.astype(float4_e2m1fn).view(np.uint8).astype(np.uint8)
+fp4_codes = scaled_fp32.astype(bfloat16).astype(float4_e2m1fn).view(np.uint8).astype(np.uint8)

应用建议

不准确？

omarzohir

6月25日强制推送 1 个提交:408861ba-[A5 TQuant]: Adding support for 2D DN MXFP4 Quantization

omarzohir

6月25日评论：

/compile

CANN-robot成员

6月25日删除了label：ci-pipeline-failed

CANN-robot成员

6月25日添加了label：ci-pipeline-running

CANN-robot成员

6月25日评论：

流水线任务触发成功
任务链接 [0631069243a746c5a24c3cba16b20e3e][流水线指导]

任务名称	状态	日志	下载链接
codecheck	✅ SUCCESS	>>>>>
SCA	✅ SUCCESS	>>>>>
antipoison	✅ SUCCESS	>>>>>
Check_Pr	✅ SUCCESS	>>>>>
Compile_Ascend_X86	✅ SUCCESS	>>>>>	>>>>>
Compile_Ascend_ARM	✅ SUCCESS	>>>>>	>>>>>
pto-isa_check	✅ SUCCESS
pre_comment	✅ SUCCESS	>>>>>
Compile_Ascend_X86_ubuntu24	✅ SUCCESS	>>>>>	>>>>>
cocheck_codestyle	✅ SUCCESS
precommit	⚠️ WARNING	>>>>>	>>>>>
Compile_Ascend_ARM_A5	✅ SUCCESS	>>>>>
Compile_Ascend_X86_A5_ubuntu24	✅ SUCCESS	>>>>>
UT_Test	✅ SUCCESS