feat: add ExternalStream class with unsupported-scenario docs-pytorch-AtomGit

Mtest/npu/test_stream.py+69-0

	@@ -1,3 +1,6 @@
		1	+import contextlib
		2	+import ctypes
		3	+
1	import torch	4	import torch
2		5
3	import torch_npu	6	import torch_npu
	@@ -47,5 +50,71 @@ class TestNpuStream(TestCase):
47	self.assertTrue((s.stream_id >> 5) == 4)	50	self.assertTrue((s.stream_id >> 5) == 4)
48		51
49		52
		53	+class TestExternalStream(TestCase):
		54	+
		55	+ @contextlib.contextmanager
		56	+ def _get_external_stream(self, device=None):
		57	+ rt = torch_npu.npu.npurt()
		58	+ stream = ctypes.c_void_p(0)
		59	+ with torch_npu.npu.device(device):
		60	+ try:
		61	+ ret = rt.npuStreamCreate(ctypes.addressof(stream))
		62	+ self.assertEqual(int(ret), 0)
		63	+ self.assertNotEqual(stream.value, 0)
		64	+ yield stream.value
		65	+ finally:
		66	+ if stream.value:
		67	+ ret = rt.npuStreamDestroy(stream.value)
		68	+ self.assertEqual(int(ret), 0)
		69	+
		70	+ def test_external_stream_creation(self):
		71	+ with self._get_external_stream() as stream_v:
		72	+ ext_stream = torch_npu.npu.ExternalStream(stream_v)
		73	+ self.assertEqual(stream_v, ext_stream.npu_stream)
		74	+ self.assertEqual(ext_stream.device.index, torch_npu.npu.current_device())
		75	+
		76	+ def test_external_stream_as_current(self):
		77	+ with self._get_external_stream() as stream_v:
		78	+ ext_stream = torch_npu.npu.ExternalStream(stream_v)
		79	+ with torch_npu.npu.stream(ext_stream):
		80	+ self.assertEqual(
		81	+ torch_npu.npu.current_stream().npu_stream,
		82	+ ext_stream.npu_stream)
		83	+
		84	+ def test_external_stream_op(self):
		85	+ with self._get_external_stream() as stream_v:
		86	+ ext_stream = torch_npu.npu.ExternalStream(stream_v)
		87	+ with torch_npu.npu.stream(ext_stream):
		88	+ x = torch.randn(10, device='npu')
		89	+ y = x + 1
		90	+ torch_npu.npu.synchronize()
		91	+ self.assertEqual(y.sum().item(), x.sum().item() + x.numel())
		92	+
		93	+ def test_external_stream_same_ptr(self):
		94	+ with self._get_external_stream() as stream_v:
		95	+ ext1 = torch_npu.npu.ExternalStream(stream_v)
		96	+ ext2 = torch_npu.npu.ExternalStream(stream_v)
		97	+ self.assertEqual(ext1.npu_stream, ext2.npu_stream)
		98	+ self.assertEqual(ext1, ext2)
		99	+
		100	+ def test_external_stream_isinstance(self):
		101	+ with self._get_external_stream() as stream_v:
		102	+ ext_stream = torch_npu.npu.ExternalStream(stream_v)
		103	+ self.assertIsInstance(ext_stream, torch_npu.npu.Stream)
		104	+ self.assertIsInstance(ext_stream, torch_npu.npu.ExternalStream)
		105	+
		106	+ def test_external_stream_synchronize_restriction(self):
		107	+ with self._get_external_stream() as stream_v:
		108	+ ext_stream = torch_npu.npu.ExternalStream(stream_v)
		109	+ with self.assertRaisesRegex(RuntimeError, "NPUStream::synchronize"):
		110	+ ext_stream.synchronize()
		111	+
		112	+ def test_external_stream_query_restriction(self):
		113	+ with self._get_external_stream() as stream_v:
		114	+ ext_stream = torch_npu.npu.ExternalStream(stream_v)
		115	+ with self.assertRaisesRegex(RuntimeError, "Cannot query"):
		116	+ ext_stream.query()
		117	+
		118	+
50	if __name__ == "__main__":	119	if __name__ == "__main__":
51	run_tests()	120	run_tests()

Mtest/torch_npu_schema.json+6-0

	@@ -719,6 +719,9 @@
719	"torch_npu.npu.graphs.make_graphed_callables": {	719	"torch_npu.npu.graphs.make_graphed_callables": {
720	"signature": "(callables, sample_args, num_warmup_iters=3, allow_unused_input=False, pool=None)"	720	"signature": "(callables, sample_args, num_warmup_iters=3, allow_unused_input=False, pool=None)"
721	},	721	},
		722	+ "torch_npu.npu.streams.ExternalStream": {
		723	+ "signature": "(stream_ptr, device=None, **kwargs)"
		724	+ },
722	"torch_npu.npu.streams.ExternalEvent": {	725	"torch_npu.npu.streams.ExternalEvent": {
723	"signature": "()"	726	"signature": "()"
724	},	727	},
	@@ -791,6 +794,9 @@
791	"torch_npu.npu.Event.ipc_handle": {	794	"torch_npu.npu.Event.ipc_handle": {
792	"signature": "(self)"	795	"signature": "(self)"
793	},	796	},
		797	+ "torch_npu.npu.ExternalStream": {
		798	+ "signature": "(stream_ptr, device=None, **kwargs)"
		799	+ },
794	"torch_npu.npu.ExternalEvent": {	800	"torch_npu.npu.ExternalEvent": {
795	"signature": "()"	801	"signature": "()"
796	},	802	},

Mtorch_npu/csrc/npu/Stream.cpp+15-1

@@ -56,12 +56,26 @@ static PyObject *THNPStream_pynew(

56 return nullptr; 56 return nullptr;

57 } 57 }

58 58

59 + if (stream_ptr) {

60 + TORCH_CHECK(

61 + priority == 0,

62 + "Priority was explicitly set for an external stream",

63 + PTA_ERROR(ErrCode::PARAM));

64 + TORCH_CHECK(

65 + is_sync_launch == 0,

66 + "is_sync_launch was explicitly set for an external stream",

67 + PTA_ERROR(ErrCode::PARAM));

68 + }

69 +

59 c10_npu::NPUStream stream = 70 c10_npu::NPUStream stream =

60 (stream_id || device_index || device_type) ? 71 (stream_id || device_index || device_type) ?

61 c10_npu::NPUStream::unpack3( 72 c10_npu::NPUStream::unpack3(

62 stream_id, device_index, static_cast<c10::DeviceType>(device_type)) : 73 stream_id, device_index, static_cast<c10::DeviceType>(device_type)) :

74 + (stream_ptr ?

75 + c10_npu::getStreamFromExternal(

76 + reinterpret_cast<aclrtStream>(stream_ptr), current_device) :

63 (is_sync_launch ? c10_npu::getNPUStreamFromSyncLaunchPool() : 77 (is_sync_launch ? c10_npu::getNPUStreamFromSyncLaunchPool() :

64 - c10_npu::getStreamFromPool(priority)); 78 + c10_npu::getStreamFromPool(priority)));

atomgit-bot7月28日

🟡 Medium Priority

在 THNPStream_pynew 的三元表达式链中（第 70-78 行），(stream_id || device_index || device_type) 的条件判断优先于 stream_ptr。当调用者同时传入非零的 stream_ptr 和 stream_id（或其他解包参数）时，代码会走 unpack3 路径而非 getStreamFromExternal，导致外部流包装失败。

触发条件：从 Python 调用 ExternalStream(ptr, stream_id=456) 时，stream_id 通过 **kwargs 透传到 C++ 层，使 stream_id || device_index || device_type 为真，从而绕过 getStreamFromExternal。

失效模式：外部流指针被当作普通的 stream_id/device_index/device_type 组合来解包，产生错误或未定义的流对象。

建议：将 stream_ptr 分支提升为三元表达式的第一个判断条件，确保外部流指针优先于 unpack 路径。

改动建议

+    c10_npu::NPUStream stream =
+        (stream_ptr ?
+        c10_npu::getStreamFromExternal(
+            reinterpret_cast<aclrtStream>(stream_ptr), current_device) :
+        (stream_id || device_index || device_type) ?
+        c10_npu::NPUStream::unpack3(
+            stream_id, device_index, static_cast<c10::DeviceType>(device_type)) :
+        (is_sync_launch ? c10_npu::getNPUStreamFromSyncLaunchPool() :
         c10_npu::getStreamFromPool(priority)));

应用建议

Y

yuht9

7月29日评论：

65 79

66 THNPStream *self = (THNPStream *)ptr.get(); 80 THNPStream *self = (THNPStream *)ptr.get();

67 self->stream_id = static_cast<int64_t>(stream.id()); 81 self->stream_id = static_cast<int64_t>(stream.id());

Mtorch_npu/npu/__init__.py+2-1

	@@ -115,6 +115,7 @@ __all__ = [
115	"is_current_stream_capturing",	115	"is_current_stream_capturing",
116	"make_graphed_callables",	116	"make_graphed_callables",
117	"ExternalEvent",	117	"ExternalEvent",
		118	+ "ExternalStream",
118	"graph_task_group_begin",	119	"graph_task_group_begin",
119	"graph_task_group_end",	120	"graph_task_group_end",
120	"graph_task_update_begin",	121	"graph_task_update_begin",
	@@ -164,7 +165,7 @@ from .utils import (obfuscation_initialize, obfuscation_calculate, obfuscation_f
164	finalize_dump, set_dump, get_npu_overflow_flag, clear_npu_overflow_flag,	165	finalize_dump, set_dump, get_npu_overflow_flag, clear_npu_overflow_flag,
165	check_uce_in_memory, stress_detect, _get_uce_addr, ipc_collect, set_op_timeout_ms)	166	check_uce_in_memory, stress_detect, _get_uce_addr, ipc_collect, set_op_timeout_ms)
166	from ._recovery import restart_device, stop_device	167	from ._recovery import restart_device, stop_device
167	-from .streams import Stream, Event, SyncLaunchStream, ExternalEvent	168	+from .streams import Stream, Event, SyncLaunchStream, ExternalStream, ExternalEvent
168	from .mstx import mstx	169	from .mstx import mstx
169	from .npu_config import * # noqa: F403	170	from .npu_config import * # noqa: F403
170	from .autocast_utils import * # noqa: F403	171	from .autocast_utils import * # noqa: F403

Mtorch_npu/npu/streams.py+67-1

@@ -3,7 +3,7 @@ import ctypes

3 import torch_npu 3 import torch_npu

4 import torch_npu._C 4 import torch_npu._C

5 5

6 -__all__ = ["Stream", "Event", "SyncLaunchStream", "ExternalEvent"] 6 +__all__ = ["Stream", "Event", "SyncLaunchStream", "ExternalStream", "ExternalEvent"]

7 7

8 8

9 class Stream(torch_npu._C._NPUStreamBase): 9 class Stream(torch_npu._C._NPUStreamBase):

@@ -110,6 +110,72 @@ class Stream(torch_npu._C._NPUStreamBase):

110 .format(self.device, self.npu_stream)) 110 .format(self.device, self.npu_stream))

111 111

112 112

113 +class ExternalStream(Stream):

114 + r"""Wrapper around an externally allocated NPU stream.

115 +

116 + This class is used to wrap streams allocated in other libraries in order

117 + to facilitate data exchange and multi-library interactions.

118 +

119 + .. note:: This class doesn't manage the stream life-cycle, it is the user

120 + responsibility to keep the referenced stream alive while this class is

121 + being used.

122 +

123 + .. note:: ``priority`` and ``is_sync_launch`` must keep their defaults

124 + (``0``) when wrapping an external stream; passing non-default values

125 + raises a ``RuntimeError``.

126 +

127 + Args:

128 + stream_ptr(int): Integer representation of the `aclrtStream` value.

129 + allocated externally.

130 + device(torch.device or int, optional): the device where the stream

131 + was originally allocated. If device is specified incorrectly,

132 + subsequent launches using this stream may fail.

133 +

134 + Unsupported scenarios (raise ``RuntimeError``):

135 +

136 + .. list-table::

137 + :widths: 35 65

138 + :header-rows: 1

139 +

140 + * - API

141 + - Description

142 + * - ``synchronize()``

143 + - Caller must synchronize the external stream externally.

144 + * - ``query()``

145 + - Caller must track stream completion externally.

146 + * - ``record_event()`` / ``Event.record(ext)`` /

147 + ``ExternalEvent.record(ext)``

148 + - Event recording on an external stream is not supported.

149 + * - ``wait_event(event)`` / ``wait_stream(stream)`` /

150 + ``Event.wait(ext)`` / ``ExternalEvent.wait(ext)``

151 + - Event waiting on an external stream is not supported.

152 + * - ``Tensor.record_stream(ext)``

153 + - Caching allocator stream tracking is not supported.

154 + * - ``NPUGraph.capture_begin()`` / ``capture_end()``

155 + - Graph capture on an external stream is not supported.

156 + * - ``graph_task_group_begin`` / ``end`` /

157 + ``graph_task_update_begin`` / ``end``

158 + - Graph task group APIs are not supported.

159 + * - ``super_kernel_scope_begin`` / ``end``

160 + - Super kernel scope is not supported.

161 + * - ``launch_callback``

162 + - Host callback launch is not supported.

163 + * - ``subscribe_report`` / ``unsubscribe_report``

164 + - Task report subscribe/unsubscribe is not supported.

165 + * - ``LaunchRecordEventTask`` / ``LaunchWaitEventTask``

166 + - Async task queue event ops are not supported.

167 +

168 + The root cause: the NPU backend routes the above APIs through torch_npu's

169 + internal ``AsyncTaskQueue`` (taskqueue), which only tracks streams it

170 + created itself. An external stream is outside the taskqueue's tracking

171 + scope. For full capabilities, use a torch_npu-managed ``Stream`` instead.

172 + """

173 +

174 + def __new__(cls, stream_ptr, device=None, **kwargs):

175 + with torch_npu.npu.device(device):

176 + return super().__new__(cls, stream_ptr=stream_ptr, **kwargs)

atomgit-bot7月28日

🟡 Medium Priority

changed line: torch_npu/npu/streams.py 第174行 ExternalStream.__new__ 不校验 stream_ptr 是否为 0。

affected behavior: 当用户错误地传入 stream_ptr=0（或从 ctypes 获取空指针值）时，ExternalStream.__new__ 经 Stream.__new__ → C++ THNPStream_pynew，由于 stream_ptr=0 在 C++ 中被视为 falsy，三元表达式会跳过 getStreamFromExternal 分支，最终 fallthrough 到 getStreamFromPool(priority=0)，静默创建一个 pool stream 而非 external stream。

failure mode: 用户认为自己包装了外部流，实际却与 torch_npu 内部 pool stream 共享同一个流对象，可能导致并发冲突、数据竞争或难以排查的异常行为。getStreamFromExternal 内部已有 TORCH_CHECK(stream != nullptr) 保护，但因 C++ 的三元优先级，该检查永远不会在 stream_ptr=0 时被触发。

建议：在 ExternalStream.__new__ 中增加 stream_ptr 非零校验，例如：if not stream_ptr: raise ValueError("stream_ptr must be a non-zero aclrtStream pointer")。

改动建议

+    def __new__(cls, stream_ptr, device=None, **kwargs):
+        if not stream_ptr:
+            raise ValueError("stream_ptr must be a non-zero aclrtStream pointer")
+        with torch_npu.npu.device(device):
             return super().__new__(cls, stream_ptr=stream_ptr, **kwargs)

应用建议

Y

yuht9

7月29日评论：

177 +

178 +

113 class Event(torch_npu._C._NPUEventBase): 179 class Event(torch_npu._C._NPUEventBase):

114 r"""Wrapper around a NPU event. 180 r"""Wrapper around a NPU event.

115 181