fix(group): ⚡️ run each rank on a non-blocking stream (#75)

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

Author: Panadestein

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

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(

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",
 ]

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
 

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)
 

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)]