From e4d12621e350c4c3caa1d693686fb8004b1b8957 Mon Sep 17 00:00:00 2001 From: Panadestein Date: Thu, 29 Oct 2026 06:58:04 +0000 Subject: [PATCH] =?UTF-8?q?fix(group):=20=E2=9A=A1=EF=B8=8F=20run=20each?= =?UTF-8?q?=20rank=20on=20a=20non-blocking=20stream?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The ranks of a DeviceGroup ran on CuPy's legacy null stream, where a peer copy into one GPU serializes with the copy in the other direction. Every allgather over PCIe got half the bandwidth: 19.5/11.0/6.6 GB/s per receiver at 2/4/8 GPUs, against 34.8/15.3/11.3 on non-blocking streams. Each rank now runs on a non-blocking stream of its device, one per device and reused across calls so that CuPy's pool keeps reusing its cached blocks. The stream waits for the caller's current stream on entry and is synchronized on exit. On 8x RTX PRO 6000 (PCIe, complex64, D_M = 4000, chi = 2000) the 8-GPU sweep drops from 169.3 s to 145.5 s, 63% to 73% parallel efficiency, with unchanged outputs. Assisted-by: Claude Code:claude-opus-5-5 --- docs/content/docs/features/large-problems.mdx | 6 +++- src/src_method/_group.py | 5 +-- src/src_method/utils/__init__.py | 2 ++ src/src_method/utils/_backend.py | 36 +++++++++++++++++++ tests/test_backend.py | 6 ++++ tests/test_gpu_backend.py | 33 +++++++++++++++++ 6 files changed, 85 insertions(+), 3 deletions(-) diff --git a/docs/content/docs/features/large-problems.mdx b/docs/content/docs/features/large-problems.mdx index 11a6804..dc6cca2 100644 --- a/docs/content/docs/features/large-problems.mdx +++ b/docs/content/docs/features/large-problems.mdx @@ -97,7 +97,11 @@ out = src( the first one gathers the sketches and runs the QR. Everything runs in the calling process, one thread per GPU, so nothing changes for the caller: the call returns the output cores as with one GPU, errors are raised as usual, and the inputs are -read once and shared. +read once and shared. Each thread runs on a non-blocking CUDA stream of its own GPU, +not on CuPy's legacy default stream, where a copy between two GPUs waits for the +copy in the other direction and the exchanges between GPUs would get half the +bandwidth. The stream first waits for the work already queued on the caller's +current stream, so inputs given as CuPy arrays may still be in flight. The sketch index is split: each GPU owns a contiguous block of the `chi_out` sketch columns and computes only those columns of every environment and of every diff --git a/src/src_method/_group.py b/src/src_method/_group.py index 20c6602..413a11f 100644 --- a/src/src_method/_group.py +++ b/src/src_method/_group.py @@ -28,6 +28,7 @@ enable_peer_access, get_xp, use_device, + use_stream, ) if TYPE_CHECKING: @@ -125,14 +126,14 @@ def run(self, fn: Callable[[int], T]) -> list[T]: Exception: The first error raised by a rank, after every rank stopped. """ if self.size == 1: - with use_device(self.xp, self.devices[0]): + with use_device(self.xp, self.devices[0]), use_stream(self.xp): return [fn(0)] first: list[BaseException] = [] lock = threading.Lock() def worker(rank: int) -> T: try: - with use_device(self.xp, self.devices[rank]): + with use_device(self.xp, self.devices[rank]), use_stream(self.xp): return fn(rank) except BaseException as err: secondary = isinstance( diff --git a/src/src_method/utils/__init__.py b/src/src_method/utils/__init__.py index 91b2111..bed085c 100644 --- a/src/src_method/utils/__init__.py +++ b/src/src_method/utils/__init__.py @@ -27,6 +27,7 @@ to_host_async, to_numpy, use_device, + use_stream, ) from .linalg import truncated_qr @@ -56,4 +57,5 @@ "to_numpy", "truncated_qr", "use_device", + "use_stream", ] diff --git a/src/src_method/utils/_backend.py b/src/src_method/utils/_backend.py index ff78a37..3124d57 100644 --- a/src/src_method/utils/_backend.py +++ b/src/src_method/utils/_backend.py @@ -213,6 +213,42 @@ def use_device(xp: ModuleType, device: int) -> Iterator[None]: yield +# One per device, reused across calls: CuPy's pool hands a freed block only to the +# stream that allocated it, so a new stream per call would strand the blocks the +# previous one cached. +_streams: dict[int, Any] = {} + + +@contextmanager +def use_stream(xp: ModuleType) -> Iterator[None]: + """Run the block on a non-blocking stream of the current device. + + On the legacy null stream, a peer copy into one device serializes with the copy + in the other direction, which halves the bandwidth of every collective over + PCIe. The stream first waits for the work already queued on the previous + current stream, so inputs written there are ready, and is synchronized on a + normal exit, so whatever the block returns is ready on any stream. On the host + this is a no-op. + + Args: + xp: The array namespace. + + Yields: + Nothing. + """ + if is_host(xp): + yield + return + device = xp.cuda.Device().id + stream = _streams.get(device) or _streams.setdefault( + device, xp.cuda.Stream(non_blocking=True) + ) + stream.wait_event(xp.cuda.get_current_stream().record()) + with stream: + yield + stream.synchronize() + + # cudaErrorPeerAccessAlreadyEnabled: another library, or an earlier call, did it. _PEER_ALREADY_ENABLED = 704 diff --git a/tests/test_backend.py b/tests/test_backend.py index 62a1ca1..4e18e60 100644 --- a/tests/test_backend.py +++ b/tests/test_backend.py @@ -21,6 +21,7 @@ to_device_async, to_host_async, use_device, + use_stream, ) @@ -90,6 +91,11 @@ def test_host_devices_are_simulated(): assert current_device(np) == 0 +def test_host_use_stream_is_a_no_op(): + with use_stream(np): + assert isinstance(current_stream(np), NullStream) + + def test_copy_into_on_host(): dst = np.empty((2, 3), dtype=complex) diff --git a/tests/test_gpu_backend.py b/tests/test_gpu_backend.py index 0e851ba..9788162 100644 --- a/tests/test_gpu_backend.py +++ b/tests/test_gpu_backend.py @@ -13,6 +13,7 @@ import src_method._sweep as sweep_module from src_method import Resources, apply, compress, src +from src_method._group import DeviceGroup from src_method._kernels import SiteKernels from src_method._plan import ( GPU_MARGIN_FRACTION, @@ -475,3 +476,35 @@ def test_too_many_devices_raise() -> None: with pytest.raises(ValueError, match="visible"): src(*stack, chi_out=4, device="gpu", resources=Resources(devices=count + 1)) + + +def test_group_runs_on_a_non_blocking_stream() -> None: + """Ranks leave the null stream, keep one stream per device, and see prior work.""" + group = DeviceGroup("gpu", [0]) + data = cupy.zeros(1 << 24, dtype=cupy.float32) + data += 1 # queued on the null stream before the run + + def fn(_rank: int) -> tuple[int, float]: + stream = cupy.cuda.get_current_stream() + assert stream != cupy.cuda.Stream.null + return stream.ptr, float((data * 2).sum()) + + ((first, total),) = group.run(fn) + ((second, _),) = group.run(fn) + + assert total == 2 * data.size + assert first == second + assert cupy.cuda.get_current_stream() == cupy.cuda.Stream.null + + +@pytest.mark.usefixtures("two_gpus") +def test_two_gpus_run_on_their_own_streams() -> None: + """Each rank of a group runs on a non-blocking stream of its own device.""" + group = DeviceGroup("gpu", 2) + + def fn(rank: int) -> tuple[int, int]: + stream = cupy.cuda.get_current_stream() + assert stream != cupy.cuda.Stream.null + return rank, stream.device_id + + assert group.run(fn) == [(0, 0), (1, 1)]