master

ascend-robot[Bugfix] Fix chunk_kda_bwd precision

2e0108b8创建于 3 天前历史提交

文件	最后提交记录	最后更新时间
tilelang	[Feat] Add tilelang op sinkhorn Co-authored-by: zhuweichen<calvin_zhu0210@outlook.com> # message auto-generated for no-merge-commit merge: !73 merge sinkhorn into master [Feat] Add tilelang op sinkhorn Created-by: zhuweichen Commit-by: zhuweichen Merged-by: ascend-robot Description: ## Related Issue https://gitcode.com/Ascend/MindSpeed-Ops/issues/50 ## What this PR does / why we need it? This PR adds a forward-only TileLang implementation of the HC-Split Sinkhorn operator for Ascend arch32. Main changes: - Add mindspeed_ops.api.tilelang.sinkhorn.sinkhorn as the public Python API. - Add the arch32 TileLang kernel implementation for HC-Split Sinkhorn. - Support hc_mult=4. - Support torch.float32 and torch.bfloat16 inputs/outputs. For BF16, the kernel keeps BF16 IO and uses FP32 UB buffers for internal sigmoid, exp, reduce, and Sinkhorn normalization. - Add input validation for shape, dtype, hc_mult, sinkhorn_iters, and unsupported arch35. - Add unit tests and ATK accuracy cases for TileLang Sinkhorn. - Add operator documentation at docs/tilelang/sinkhorn.md. ## Does this PR introduce any user-facing change? Yes. This PR introduces a new TileLang Sinkhorn API: python from mindspeed_ops.api.tilelang.sinkhorn import sinkhorn Supported usage: - mixes: shape [batch_size, seq_len, (2 + hc_mult) * hc_mult] - hc_scale: shape [3] - hc_base: shape [(2 + hc_mult) * hc_mult] - hc_mult: 4 - dtype: torch.float32 or torch.bfloat16 - platform: Ascend arch32 only The operator is forward-only and does not provide autograd backward. Usage details are documented in docs/tilelang/sinkhorn.md. ## How was this patch tested? Verified on Ascend NPU environment: Ascend910_9382, CANN 9.0.0, Python 3.11, TileLang-Ascend installed. 1. Unit test bash pytest -q tests/unit_tests/tilelang/test_sinkhorn.py Result: 42 passed 2. ATK accuracy test sh cd tests/atk_tests/tilelang/sinkhorn export PYTHONPATH=/home/z30046196/MindSpeed-Ops-pr73:$PYTHONPATH atk node --backend triton --devices 0 node --backend npu --devices 0 task -c result/sinkhorn/json/all_sinkhorn.json --task accuracy -tup ./ -p tilelang_sinkhorn.py ATK accuracy log: sh +-------+----------+------------------+------------------+--------------+----------------------+--------+--------------+ | 名称 | 总用例数 | 执行成功用例个数 | 执行失败用例个数 | 通过用例个数 | 错误信息匹配用例个数 | 通过率 | 精度是否达标 | +-------+----------+------------------+------------------+--------------+----------------------+--------+--------------+ | npu_0 | 36 | 36 | 0 | 36 | 0 | 100.0 | Pass | +-------+----------+------------------+------------------+--------------+----------------------+--------+--------------+ Total Task: 36, success 36, failed 0 Summary info: acc_pass_result:Pass is_acc_dc_pass:None is_e2e_perf_pass:None is_benchmark_perf_pass:None is_device_perf_pass:None is_device_memory_pass:None is_memory_check_pass:None . atk task success! atk task finish. cost time: 9 s 3. ATK device performance test sh cd tests/atk_tests/tilelang/sinkhorn export PYTHONPATH=/home/z30046196/MindSpeed-Ops-pr73:$PYTHONPATH atk node --backend triton --devices 0 node --backend npu --devices 0 task -c result/sinkhorn/json/all_sinkhorn.json --task performance_device -tup ./ -p tilelang_sinkhorn.py ATK performance log: sh +-------+----------+------------------+------------------+--------------------+ | 名称 | 总用例数 | device性能通过率 | 平均device性能比 | device性能是否达标 | +-------+----------+------------------+------------------+--------------------+ | npu_0 | 36 | 94.4444 | 7.2842 | Pass | +-------+----------+------------------+------------------+--------------------+ Total Task: 36, success 36, failed 0 Summary info: acc_pass_result:None is_acc_dc_pass:None is_e2e_perf_pass:None is_benchmark_perf_pass:None is_device_perf_pass:Pass is_device_memory_pass:None is_memory_check_pass:None . save result excel file: /home/z30046196/MindSpeed-Ops-pr73/tests/atk_tests/tilelang/sinkhorn/atk_output/all_sinkhorn_2026-07-09-09-54-25-397923/report/ all_sinkhorn_reports_2026-07-09-09-54-26.xlsx atk task success! atk task finish. cost time: 537 s See merge request: Ascend/MindSpeed-Ops!73	7 天前
triton	[Bugfix] Fix chunk_kda_bwd precision Co-authored-by: zhuweichen<calvin_zhu0210@outlook.com> # message auto-generated for no-merge-commit merge: !105 merge kda_fix into master [Bugfix] Fix chunk_kda_bwd precision Created-by: zhuweichen Commit-by: zhuweichen Merged-by: ascend-robot Description: ## What this PR does / why we need it? Fix chunk_kda_bwd precision. ## Does this PR introduce any user-facing change? No. ## How was this patch tested? ATK & UT   See merge request: Ascend/MindSpeed-Ops!105	3 天前
__init__.py	[Feat] Add tilelang op sinkhorn Co-authored-by: zhuweichen<calvin_zhu0210@outlook.com> # message auto-generated for no-merge-commit merge: !73 merge sinkhorn into master [Feat] Add tilelang op sinkhorn Created-by: zhuweichen Commit-by: zhuweichen Merged-by: ascend-robot Description: ## Related Issue https://gitcode.com/Ascend/MindSpeed-Ops/issues/50 ## What this PR does / why we need it? This PR adds a forward-only TileLang implementation of the HC-Split Sinkhorn operator for Ascend arch32. Main changes: - Add mindspeed_ops.api.tilelang.sinkhorn.sinkhorn as the public Python API. - Add the arch32 TileLang kernel implementation for HC-Split Sinkhorn. - Support hc_mult=4. - Support torch.float32 and torch.bfloat16 inputs/outputs. For BF16, the kernel keeps BF16 IO and uses FP32 UB buffers for internal sigmoid, exp, reduce, and Sinkhorn normalization. - Add input validation for shape, dtype, hc_mult, sinkhorn_iters, and unsupported arch35. - Add unit tests and ATK accuracy cases for TileLang Sinkhorn. - Add operator documentation at docs/tilelang/sinkhorn.md. ## Does this PR introduce any user-facing change? Yes. This PR introduces a new TileLang Sinkhorn API: python from mindspeed_ops.api.tilelang.sinkhorn import sinkhorn Supported usage: - mixes: shape [batch_size, seq_len, (2 + hc_mult) * hc_mult] - hc_scale: shape [3] - hc_base: shape [(2 + hc_mult) * hc_mult] - hc_mult: 4 - dtype: torch.float32 or torch.bfloat16 - platform: Ascend arch32 only The operator is forward-only and does not provide autograd backward. Usage details are documented in docs/tilelang/sinkhorn.md. ## How was this patch tested? Verified on Ascend NPU environment: Ascend910_9382, CANN 9.0.0, Python 3.11, TileLang-Ascend installed. 1. Unit test bash pytest -q tests/unit_tests/tilelang/test_sinkhorn.py Result: 42 passed 2. ATK accuracy test sh cd tests/atk_tests/tilelang/sinkhorn export PYTHONPATH=/home/z30046196/MindSpeed-Ops-pr73:$PYTHONPATH atk node --backend triton --devices 0 node --backend npu --devices 0 task -c result/sinkhorn/json/all_sinkhorn.json --task accuracy -tup ./ -p tilelang_sinkhorn.py ATK accuracy log: sh +-------+----------+------------------+------------------+--------------+----------------------+--------+--------------+ | 名称 | 总用例数 | 执行成功用例个数 | 执行失败用例个数 | 通过用例个数 | 错误信息匹配用例个数 | 通过率 | 精度是否达标 | +-------+----------+------------------+------------------+--------------+----------------------+--------+--------------+ | npu_0 | 36 | 36 | 0 | 36 | 0 | 100.0 | Pass | +-------+----------+------------------+------------------+--------------+----------------------+--------+--------------+ Total Task: 36, success 36, failed 0 Summary info: acc_pass_result:Pass is_acc_dc_pass:None is_e2e_perf_pass:None is_benchmark_perf_pass:None is_device_perf_pass:None is_device_memory_pass:None is_memory_check_pass:None . atk task success! atk task finish. cost time: 9 s 3. ATK device performance test sh cd tests/atk_tests/tilelang/sinkhorn export PYTHONPATH=/home/z30046196/MindSpeed-Ops-pr73:$PYTHONPATH atk node --backend triton --devices 0 node --backend npu --devices 0 task -c result/sinkhorn/json/all_sinkhorn.json --task performance_device -tup ./ -p tilelang_sinkhorn.py ATK performance log: sh +-------+----------+------------------+------------------+--------------------+ | 名称 | 总用例数 | device性能通过率 | 平均device性能比 | device性能是否达标 | +-------+----------+------------------+------------------+--------------------+ | npu_0 | 36 | 94.4444 | 7.2842 | Pass | +-------+----------+------------------+------------------+--------------------+ Total Task: 36, success 36, failed 0 Summary info: acc_pass_result:None is_acc_dc_pass:None is_e2e_perf_pass:None is_benchmark_perf_pass:None is_device_perf_pass:Pass is_device_memory_pass:None is_memory_check_pass:None . save result excel file: /home/z30046196/MindSpeed-Ops-pr73/tests/atk_tests/tilelang/sinkhorn/atk_output/all_sinkhorn_2026-07-09-09-54-25-397923/report/ all_sinkhorn_reports_2026-07-09-09-54-26.xlsx atk task success! atk task finish. cost time: 537 s See merge request: Ascend/MindSpeed-Ops!73	7 天前