Skip to content

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

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

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

Conversation

@voltjia

@voltjia voltjia commented Aug 4, 2026

Copy link
Copy Markdown
Collaborator

Summary

  • Uses the target device's default PyTorch stream when an InfiniOps Handle has no stream, and imports a non-null handle stream through the vendor's external-stream API before constructing the native PyTorch stream guard.
  • Declares the shared C10<kDev> template in src/torch/c10.h and keeps each backend specialization in src/torch/<platform>/c10.h, following the existing platform-oriented C10 design without a separate stream bridge.
  • Exposes only the C10-aligned default-stream lookup, external-stream import, and native StreamGuard operations from each specialization, while keeping the null-versus-non-null selection explicit at generated, linked, and hand-written call sites.
  • Implements the Hygon integration with its PyTorch fork's HIP C10 API and keeps the HIP-only implementation isolated from InfiniRT's DTK CUDA compatibility headers.
  • Adds generator assertions and runtime regressions for default-stream and explicit-handle ordering, including a warmed non-default-stream test that excludes first-call operator construction overhead.

Motivation

Generated ATen backends received the InfiniOps handle stream through stream_, but they invoked PyTorch without installing the selected stream through a native PyTorch guard. PyTorch work could therefore run on a different stream from the caller-selected InfiniRT stream and violate the caller's ordering assumptions.

The previous revision also introduced C10<kDev>::GetStream, which combined InfiniOps policy with vendor C10 integration and changed a null handle from the InfiniOps default stream to the current PyTorch stream. The specializations now expose separate operations corresponding to the vendor C10 APIs, and each consumer preserves the existing handle contract explicitly.

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).
  • Hygon (WITH_HYGON).
  • 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

Current validation SHA: a074719146fb7375a9bff681d20d0b0b5104180e.

Ruff 0.15.22 check: passed.
Ruff 0.15.22 format --check: passed.
Clang Format 21.1.8: passed.
Generator and architecture tests: 45 passed.
NVIDIA smoke wheel build with current InfiniRT master: passed.
NVIDIA smoke tests: 121 passed, 28 skipped in 22.82s.
NVIDIA focused non-default handle-stream test: 1 passed, 1 skipped.

The main and shadow platform CI workflows passed for every configured accelerator backend on the current SHA.

Test Results on Supported Platforms

Platform Affected Build / Smoke Result Full Result / Notes
CPU Yes Covered by the NVIDIA smoke wheel and smoke suite. The CPU specialization directly defines a no-op StreamGuard and exposes matching no-op default/external stream operations.
NVIDIA Yes The current SHA built and passed 121 smoke tests in accelerator-dev/nvidia:latest; its main and shadow CI jobs also passed. Both default-stream and explicit-handle ordering regressions passed on an A100.
Iluvatar Yes Current-SHA main and shadow CI passed. The CUDA-compatible specialization uses InfiniRT's native stream type and the PyTorch fork's CUDA C10 stream API.
Hygon Yes Current-SHA hardware validation is unavailable. The isolated implementation now uses c10::hip::getDefaultHIPStream and c10::hip::getStreamFromExternal.
MetaX Yes Current-SHA main and shadow CI passed. The CUDA-compatible specialization uses the PyTorch fork's CUDA C10 stream API.
Cambricon Yes Current-SHA main and shadow CI passed. The specialization uses torch_mlu::getDefaultMLUStream and torch_mlu::getStreamFromExternal.
Moore Yes Current-SHA main and shadow CI passed. The specialization uses the corresponding MUSA C10 stream APIs.
Ascend Yes Current-SHA main and shadow CI passed. Older torch_npu releases can use the default stream but reject explicit non-null external streams clearly.

Benchmark / Performance Impact

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

Notes for Reviewers

InfiniRT remains the source of truth for each backend's native stream type through Runtime<kDev>::Stream. Vendor PyTorch APIs are used only by the corresponding C10<kDev> specialization to obtain the default stream or translate an external native stream into the vendor's C10 stream type.

src/torch/c10.h contains only the template declaration. Each platform owns its specialization under src/torch/<platform>/c10.h, matching the original C10 organization. The old src/linked/torch/c10.h, backend-local linked C10 headers, and centralized src/torch/stream_.h bridge have been removed.

Generated, linked, and hand-written call sites select GetDefaultStream when stream_ == nullptr and GetStreamFromExternal otherwise, then construct the platform's native StreamGuard. This preserves the InfiniOps handle contract without hiding selection policy inside the C10 trait.

The Ascend external-stream API was introduced together with NPUStreamUtils.h, so that header remains the compile-time capability signal. Older torch_npu releases still support null handles through getDefaultNPUStream, while a non-null handle produces a clear error instead of silently running on another stream.

Hygon is not a no-op. Its PyTorch fork exports HIP C10 stream symbols, but directly combining the HIP headers with InfiniRT's DTK CUDA compatibility headers causes runtime type redefinitions. src/torch/hygon/c10.cc therefore performs only the HIP stream conversion behind generic c10::Stream declarations, and the normal c10::StreamGuard remains visible in the platform specialization.

@voltjia
voltjia requested a review from a team August 4, 2026 11:36
@voltjia
voltjia force-pushed the fix/torch-backend-stream branch 11 times, most recently from f6e6f27 to bbc7b58 Compare August 10, 2026 06:42
@voltjia
voltjia force-pushed the fix/torch-backend-stream branch 2 times, most recently from ce71ffc to 5d4f0dc Compare August 10, 2026 08:11
@voltjia
voltjia force-pushed the fix/torch-backend-stream branch from 5d4f0dc to a074719 Compare August 11, 2026 05:41
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