[CUDA graph] Drain the device after the last warmup, before capture

[jiaqiang-dot-liu]

Problem

Both CUDA-graph backends run two warmup iterations before capture, with a
synchronize() + barrier() at the top of each iteration:

for warmup_step in range(2):
    self._device_module.synchronize()
    self._tp_group.barrier()
    with self._precarve.measure():
        output = forward_fn()
    ...
    if post_warmup_hook is not None:
        post_warmup_hook()

# ... nothing drains the last warmup's own async work ...
graph = torch.cuda.CUDAGraph()

Those syncs order work issued before each warmup. Nothing drains the final
warmup's own async work.

Impact

Lazy / JIT kernel compilation kicked off by a first-seen shape can still be in flight
when capture starts. A JIT module that finishes loading mid-capture issues illegal
driver calls (cuModuleLoadData and friends) on the capturing stream, surfacing as
CUDA_ERROR_ILLEGAL_ADDRESS or cudaErrorStreamCaptureUnsupported.

This is intermittent by nature — it depends on whether the last warmup happened to
trigger a compile and whether that compile lands before or after capture begins.

Fix

Drain the device and re-align ranks once more, immediately before the capture.

Applied to both backends. sgl-project#33795 fixed
FullCudaGraphBackend; BreakableCudaGraphBackend has the identical pattern and was
left unfixed.

Verification

Reasoning-only. The ordering gap is visible in the source and the fix is the same
drain sgl-project#33795 already established as correct for the sibling backend. I have not
constructed a reproducer — the failure is a race and reproducing it reliably would
mean forcing a JIT compile on the last warmup shape.

Cost is two extra synchronization points per captured graph, at capture time only.
No effect on the replay path.

---

Provenance: originally authored by an automated kernel-optimization agent
(hyperloom session 20260806T082051Z), rebased onto current main and reviewed by
hand before submission.

---

CI States

Latest PR Test (Base): ❌ Run #34923945812
Latest PR Test (Extra): ❌ Run #34923945548
Latest PR Test (AMD ROCm 10): ❌ Run #34923945923