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

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