Skip to content

fix(ascend): synchronize PyTorch operators with NPU streams - #974

Merged
voltjia merged 2 commits into
masterfrom
fix/ascend-torch-stream
Sep 4, 2026
Merged

fix(ascend): synchronize PyTorch operators with NPU streams#974
voltjia merged 2 commits into
masterfrom
fix/ascend-torch-stream

Conversation

@baominghelly

Copy link
Copy Markdown
Contributor

Summary

  • Add an Ascend C10 stream adapter in src/torch/ascend/c10.h using c10_npu::NPUStreamGuard and c10_npu::getStreamFromExternal.
  • Enable generated PyTorch slot 8 operators to execute on the caller-provided Ascend stream, with an explicit default-stream fallback when the runtime stream is null.
  • Discover and link the installed torch_npu C++ headers and library for Ascend PyTorch builds, and add regression coverage for code generation and NPU stream ordering.

Motivation

Ascend was not included in the generated PyTorch backend's C10 stream-guard path. Consequently, slot 8 operators could execute on a different NPU stream from the InfiniOps runtime stream, which risks incorrect ordering and results when surrounding work is asynchronous.

This change gives Ascend the same stream-guard behavior as the other PyTorch-backed accelerator platforms. It also fails at configure time with a clear diagnostic when the installed torch_npu does not provide the required external-stream API.

Type of Change

  • feat — new feature / new operator / new platform
  • fix — bug fix
  • perf — performance improvement (no behavioral change)
  • refactor — code restructuring without behavior change
  • test — adding or fixing tests only
  • docs — documentation only
  • build / ci — build system or CI configuration
  • chore — tooling, formatting, or other non-code changes
  • Breaking change (requires a ! in the Conventional Commits prefix or a BREAKING CHANGE: footer)

Platforms Affected

  • CPU (WITH_CPU)
  • NVIDIA (WITH_NVIDIA)
  • Iluvatar (WITH_ILUVATAR)
  • MetaX (WITH_METAX)
  • Cambricon (WITH_CAMBRICON)
  • Moore (WITH_MOORE)
  • Ascend (WITH_ASCEND)
  • PyTorch C++ bindings (WITH_TORCH)
  • Build system / CMake / CI
  • Python bindings / user-facing API

Smoke Test Result

Ascend 910C, CANN 9.0.0, Python 3.11, PyTorch 2.10.0, and torch_npu 2.10.0.post4:

python -m pytest tests/test_generate_torch_ops.py -q
21 passed in 1.13s

cmake -S /workspace/InfiniOps-ascend-torch-stream-20260902 \
  -B /workspace/ascend-torch-stream-test-20260902/build-6b256e0 \
  -DCMAKE_INSTALL_PREFIX=/workspace/ascend-torch-stream-test-20260902/python3/infini \
  -DPython_EXECUTABLE=/workspace/torch-npu-wheel-audit-20260902/venv-post4/bin/python
cmake --build /workspace/ascend-torch-stream-test-20260902/build-6b256e0 -j16
[100%] Built target ops

python -m pytest tests/test_abs.py --devices ascend -q
20 passed in 1.14s

python -m pytest tests/test_torch_ops.py -k 'tril or triu' --devices ascend -q
18 passed, 9 deselected in 1.08s

The CMake cache used WITH_ASCEND=ON, WITH_CPU=ON, WITH_TORCH=ON, INFINI_RT_ROOT=/workspace/InfiniRT-6b256e0-prefix, and generated abs, tril, and triu PyTorch operators.

A CPU-only regression build using WITH_CPU=ON, WITH_ASCEND=OFF, and WITH_TORCH=ON also completed successfully:

cmake -S /workspace/InfiniOps-ascend-torch-stream-20260902 \
  -B /workspace/ascend-torch-stream-test-20260902/build-cpu-regression-v2
cmake --build /workspace/ascend-torch-stream-test-20260902/build-cpu-regression-v2 -j16
[100%] Built target ops

Test Results on Supported Platforms

Platform Affected Build / Smoke Result Full Result / Notes
NVIDIA No N/A - not affected N/A - generated include remains guarded by WITH_NVIDIA
Iluvatar No N/A - not affected N/A
MetaX No N/A - not affected N/A
Cambricon No N/A - not affected N/A
Moore No N/A - not affected N/A
Ascend Yes Build passed; targeted smoke passed Generator: 21 passed; abs: 20 passed; tril/triu: 18 passed
Additional integration results
Llama-3.2-3B-Instruct: inference passed
ChatGLM3-6B: inference passed
GLM-4-9B-0414: inference passed

Trace output confirmed:
"operator_name": "Triu", "device_type": "ascend", "implementation": 8
"operator_name": "Tril", "device_type": "ascend", "implementation": 8

Benchmark / Performance Impact

N/A. This is a correctness and stream-ordering fix; no performance claim is made.

Notes for Reviewers

  • The torch_npu integration is enabled only when both WITH_ASCEND and generated PyTorch sources are present. Non-Ascend builds do not discover or link torch_npu.
  • torch_npu 2.10.0.post4 was validated. An older 2.10.0.post2 wheel does not expose c10_npu::getStreamFromExternal; CMake now reports this incompatibility explicitly.
  • The model-level tests were performed in a separate integration checkout that also contained the existing Ascend rotary-embedding fix. The slot trace verifies that tril and triu used implementation 8.
  • No native Ascend slot 0 implementation for tril or triu is included in this PR.

Comment thread src/torch/ascend/c10.h
Comment on lines +20 to +24
if (stream == nullptr) {
return c10_npu::getDefaultNPUStream(
static_cast<c10::DeviceIndex>(device_index));
}

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

这个感觉不是必要的,按理说应该都是相当于转发到 C10 的 getStreamFromExternal 就行,我看别的平台都没加 nullptr 的处理,可以看看昇腾是不是这样。

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

确认过了,昇腾这里和其他平台的行为不同。当前验证的 torch_npu 2.10.0.post4 对应源码提交 5dd8ef3f9b375b5ae4a83538d5785754148c3302,其 c10_npu::getStreamFromExternal 会显式检查 stream != nullptr,否则直接报 External NPU stream does not support nullptr.

https://github.com/Ascend/pytorch/blob/5dd8ef3f9b375b5ae4a83538d5785754148c3302/torch_npu/csrc/core/npu/NPUStream.cpp

InfiniOps 的 stream=0 会传入空指针来表示默认流,因此这里不能无条件转发;非空流仍然直接调用 getStreamFromExternal。我补充了英文注释说明这个平台差异,并重新验证了 Ascend 构建和 tests/test_abs.py --devices ascend,结果为 20 passed

@voltjia
voltjia marked this pull request as ready for review September 4, 2026 02:37
@voltjia
voltjia requested a review from a team September 4, 2026 02:37
@voltjia
voltjia merged commit 9576444 into master Sep 4, 2026
12 of 20 checks passed
@voltjia
voltjia deleted the fix/ascend-torch-stream branch September 4, 2026 02:38
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants