[Panadestein]

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

🤖 AI text below 🤖

What

DeviceGroup.run now runs each rank on a non-blocking CUDA stream of its device
(utils._backend.use_stream) instead of CuPy's legacy null stream.

• One stream per device, reused across calls. CuPy's pool reuses a freed block
  only on the stream that allocated it, so a new stream per call would strand the
  previous call's cached blocks.
• On entry the stream waits for the caller's current stream, so CuPy inputs that
  are still being computed are ordered correctly. On normal exit it is
  synchronized, so the returned cores are ready on any stream.
• The host backend is unchanged: use_stream is a no-op there.

Why

On the legacy null stream, a peer copy into GPU A serializes with the copy from A
into B. Every allgather in the sweep (S_j and the core uploads) therefore got
half the PCIe bandwidth. The scaling runs on 8x RTX PRO 6000 (PCIe only) traced
this, see benches/large/nebius/summary.md on bench/nebius-rtx:

Peer copy, GPU 0 <-> 1 at once	GB/s
null streams	20.7
non-blocking streams	37.3

DeviceGroup.allgather alone, per receiver	G = 2	G = 4	G = 8
before (null stream)	19.5	11.0	6.6
per-thread / non-blocking streams	34.8	15.3	11.3

Results

8x RTX PRO 6000 Blackwell, PCIe 5.0, no NVLink, complex64, the same stacks as the
bench branch. Efficiency is against the 1-GPU time from that branch's step 4 (847 s
and 20.3 s).

D_M = 4000, chi = 2000	G = 2	G = 4	G = 8
before (s)	454.4	258.7	169.3
after (s)	448.7	252.4	145.5
efficiency before / after	93% / 94%	82% / 84%	63% / 73%

D_M = 1000, chi = 500	G = 2	G = 4	G = 8
before (s)	10.5	7.7	9.7
after (s)	10.1	6.8	9.8

The outputs are unchanged. At D_M = 1000 the QR distance of each GPU count to the
1-GPU output is the same as before to every printed digit (8.45e-5, 7.79e-5,
7.84e-5). The before and after times were measured in separate runs, since the
library change cannot run side by side with the old code.

The remaining gap to the bandwidth probe at 4 and 8 GPUs comes from the
collective's events and barriers. It is not addressed here.

Tests

• test_host_use_stream_is_a_no_op.
• test_group_runs_on_a_non_blocking_stream: not the null stream, the same stream
  across runs, and work queued earlier on the null stream is seen.
• test_two_gpus_run_on_their_own_streams.
• uv run prek run --all-files is clean. uv run pytest -m "not slow" passes
  (396 tests), including tests/test_gpu_backend.py on 8 GPUs.

🤖 Generated with Claude Code

[github-actions]

Test Results

    5 files  ±0      5 suites  ±0   4m 32s ⏱️ +23s
  397 tests +3    389 ✅ +2   8 💤 +1  0 ❌ ±0 
1 899 runs  +9  1 882 ✅ +7  17 💤 +2  0 ❌ ±0 

Results for commit e4d1262. ± Comparison against base commit dc6ecf6.