Skip to content

fix(torch): honor handle streams in generated operators - #880

Open
voltjia wants to merge 2 commits into
masterfrom
fix/torch-backend-stream
Open

fix(torch): honor handle streams in generated operators#880
voltjia wants to merge 2 commits into
masterfrom
fix/torch-backend-stream

Conversation

@voltjia

@voltjia voltjia commented Aug 4, 2026

Copy link
Copy Markdown
Collaborator

Summary

  • Uses InfiniRT's Runtime<kDev>::Stream aliases to bridge explicit handle streams into generated PyTorch operators on NVIDIA, MetaX, and supported Ascend torch_npu versions.
  • Preserves PyTorch's current stream when Handle carries no stream and restores the previous device and stream after an explicit bridge.
  • Adds vendor PyTorch build integration, generated-wrapper coverage, and smoke regressions for current-stream and explicit-handle execution.

Motivation

Generated ATen backends received the InfiniOps handle stream through stream_, but they invoked PyTorch without installing that stream as the framework's current stream. PyTorch work could therefore run on a different stream from the caller-selected InfiniRT stream and violate the caller's ordering assumptions. The bridge must also remain cross-platform and must not convert a missing handle stream into a vendor default stream.

No issue is linked.

Type of Change

  • feat - New feature, operator, or platform.
  • fix - Bug fix.
  • perf - Performance improvement without a behavioral change.
  • refactor - Code restructuring without a behavioral 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.

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, or CI.
  • Python bindings or user-facing API.

Smoke Test Result

NVIDIA:
$ PYTHONPATH=build-nvidia/install python -m pytest tests -m smoke -q --devices cpu nvidia
117 passed, 25 skipped, 13342 deselected in 15.91s

MetaX:
$ python -m pytest tests -m smoke -q --devices metax
75 passed, 23 skipped, 5924 deselected

Ascend CI:
$ pytest tests/ -m smoke -n 1 -v --tb=short --junitxml=/workspace/results/test-results.xml
103 passed, 28 skipped in 21.68s

Test Results on Supported Platforms

Platform Affected Build / Smoke Result Full Result / Notes
NVIDIA Yes Smoke passed. The generated Torch build passed, and both current-stream and explicit-handle stream regressions passed.
Iluvatar No N/A. No Iluvatar-specific bridge was added.
MetaX Yes Smoke passed. The generated Torch build passed, and both current-stream and explicit-handle stream regressions passed on MetaX's CUDA-compatible PyTorch API.
Cambricon No N/A. No Cambricon-specific bridge was added.
Moore No N/A. No Moore-specific bridge was added.
Ascend Yes Smoke passed. Generated Torch code compiled, linked, installed, and imported. Native external-stream behavior was not executed because the installed stable torch_npu lacks the required APIs.
Full validation output.
Final validation SHA:
f99e11e925cbf42f7d19c6a8cd6e8343bf7190a0

NVIDIA targeted stream regression:
2 passed in 2.54s

MetaX targeted stream regression:
2 passed

Ascend targeted stream regression:
npu-current: skipped because the device module does not expose the required `_sleep` test API.
npu-handle: skipped because the installed `torch_npu` does not expose `ExternalStream`.

Generator tests:
19 passed in 0.94s

Formatting:
ruff check: passed.
ruff format --check: 142 files already formatted.
clang-format 21: passed.

Benchmark / Performance Impact

N/A. This is a stream-ordering correctness fix, and no performance benchmark was run.

Notes for Reviewers

InfiniRT remains the source of truth for each backend's native stream type. Vendor PyTorch APIs are still required to install that native stream into PyTorch's thread-local current-stream state.

The guard is constructed inside every generated C++ ATen operator wrapper, so the fix applies to C++ execution and is not a Python-only adaptation. A null handle stream leaves the framework's existing current stream unchanged.

NVIDIA and MetaX use c10::cuda::getStreamFromExternal. Ascend uses c10_npu::getStreamFromExternal only when the installed headers include NPUStreamUtils.h, which was added with the official external-stream implementation in Ascend/pytorch commit 03027fe. Older torch_npu releases retain the current NPU stream because they cannot safely encode an arbitrary aclrtStream.

Validation used InfiniRT 95c70080f9551e61241110497d163dfcdf9dc7e7 for NVIDIA and MetaX. The Ascend environment used PyTorch 2.9.0+cpu and torch_npu 2.9.0.post1+gitee7ba04; it validates the compatibility fallback but not the newer native external-stream path.

@voltjia
voltjia requested a review from a team August 4, 2026 11:36
@voltjia
voltjia force-pushed the fix/torch-backend-stream branch 4 times, most recently from 8f6a607 to f457ca2 Compare August 5, 2026 10:00
@voltjia
voltjia force-pushed the fix/torch-backend-stream branch from f457ca2 to f99e11e Compare August 5, 2026 10:15
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.

1 participant