Fix: Add GPU synchronization for cached tensor transfers

[Pritiks23]

• Add torch.cuda.synchronize() after non-blocking tensor transfers to GPU to ensure cached tensors are fully transferred before use
• Prevents stale data issues where async transfers haven't completed
• Only synchronizes on CUDA devices (no overhead on CPU)
• Add comprehensive test_cache_gpu_synchronization() test to verify fix

Fixes issue where cache.add() with device=cuda could return tensors that are still mid-transfer, causing stale data to be read by subsequent forward passes.

[JadenFiotto-Kaufman]

Thanks for the PR! Before we merge I want to make sure we're solving a real problem, because the fix as written has a significant performance cost and I'm not sure the bug it's describing can actually occur on the current code path.

On non_blocking=True and stream ordering

tensor.to(device, non_blocking=True) enqueues the copy on the current CUDA stream and returns immediately, but CUDA guarantees that any subsequent op submitted to the same stream waits for the copy before running. Host-side access (.cpu(), .item(), .numpy()) also triggers an implicit sync.

The two ways non_blocking=True can actually misbehave are:

• H2D copies from non-pinned host memory. The cache path is D2D in normal usage, so this doesn't apply.
• A consumer running on a different CUDA stream than the copy. This one is worth a closer look in our codebase, but I think we're already covered: Mediator explicitly captures the caller's stream and re-sets it inside the worker thread (src/nnsight/intervention/interleaver.py:960-977) precisely so worker-thread CUDA ops stay on the same stream as the main forward. That means the .to(...) inside Cache.add runs on the same stream as the model's forward kernels, and stream ordering gives us correctness for free.

Do you have a repro that fails on main without this patch? If there's a real ordering bug I want to understand exactly where the streams diverge so we can fix it precisely — but given the stream propagation above, I'd be surprised.

On the cost of the current fix

torch.cuda.synchronize(device) is a full-device sync — it blocks the CPU until every outstanding kernel on that device finishes, not just this copy. Since Cache.add runs from a forward hook on every cached module, on something like a 32-layer model with all layers cached you'd get 32 full-device stalls per forward. That defeats non_blocking=True entirely and is actually worse than a plain blocking copy (which only waits on its own stream).

The test also doesn't quite exercise the claimed failure mode: torch.allclose(cached.cpu(), fresh.cpu()) calls .cpu() which already syncs, and cached_output.sum() runs on the same stream as the copy, so stream ordering already covers it. I'd expect the test to pass without the patch.

Again — appreciate the contribution. Happy to dig in further if you can share a repro.

[khaiwang]

I share @JadenFiotto-Kaufman's concern. We explicitly propagate the caller's CUDA stream into the mediator's worker thread, so ordering between the .to(cuda, non_blocking=True) and any consumer on that stream is already guaranteed — I don't think this race applies to the current codebase.

I ran this PR's test_cache_gpu_synchronization without the sync and it still passes, which is expected: the gpt2 fixture is already on CUDA, so .to(cuda, non_blocking=True) on the layer output is a same-device identity op (same data_ptr). There's no transfer for non_blocking to affect, so the test isn't actually exercising the path it claims to. I also tried scenarios that do involve a real transfer (H2D and D2H) and all of them pass without the sync patch as well.

The performance concern is also real — this PR adds ~40% end-to-end latency on a minimal gpt2 forward pass.

@Pritiks23 we'd appreciate the contribution if you can share a reproducible failure case on main.

[Pritiks23]

Thank you both for the detailed review the stream ordering explanation and the performance concerns make sense @JadenFiotto-Kaufman @khaiwang

To investigate further, I added a focused repro harness to this branch to check whether there’s an actual stream-ordering divergence on main without the synchronization change, and if so, where it occurs.

Repro script: repro_cache_stream_ordering.py
Commit: 71d90e8

What it checks:

Logs the CUDA stream pointer at each Cache.add call.
Verifies whether NULL stream appears in the cache path.
Repeats runs and compares cached vs fresh values without torch.cuda.synchronize.
Includes a different-stream consumer stress path (no explicit stream wait) to probe divergence.

Given the points raised about stream propagation and the cost of full-device synchronization, I agree we should avoid the current fix unless there’s a reproducible correctness issue on main.

If you get a chance to run the repro on a CUDA machine, I’d appreciate any feedback on whether it exposes a real divergence case.