[Bug] Fix Negative Cuda Memory Usage

[yewentao256]

Purpose

2025-09-24T04:15:29.297966219Z (EngineCore_DP4 pid=1673) INFO 09-24 00:15:29 [v1/worker/gpu_model_runner.py:3380] Graph capturing finished in 79 secs, took -9.96 GiB
2025-09-24T04:15:29.298151199Z (EngineCore_DP4 pid=1673) DEBUG 09-24 00:15:29 [v1/worker/gpu_worker.py:401] Free memory on device (176.17/178.35 GiB) on startup. Desired GPU memory utilization is (0.9, 160.52 GiB). Actual usage is 60.79 GiB for weight, 13.33 GiB for peak activation, 10.96 GiB for non-torch memory, and -9.96 GiB for CUDAGraph memory. Replace gpu_memory_utilization config with `--kv-cache-memory=91539781222` (85.25 GiB) to fit into requested memory, or `--kv-cache-memory=108345741312` (100.9 GiB) to fully utilize gpu memory. Current kv cache memory in use is 75.44 GiB.

This PR updates the logic for memory usage capture, hopefully fix the issue.

Test

@smarterclayton Could you try again using this branch?

[gemini-code-assist]

Code Review

This pull request addresses a bug where CUDA graph memory usage was reported as a negative value. The root cause was asynchronous CUDA operations leading to inaccurate memory measurements. The fix correctly moves the memory measurement logic inside the graph_capture context and, crucially, adds a torch.cuda.synchronize() call before measuring the final memory usage. This ensures that all graph capture operations are complete before the memory is queried, providing an accurate measurement. The changes are logical, well-targeted, and should resolve the issue. The implementation is correct.

[yewentao256]

@smarterclayton Hi Clayton, wondering if you could validate this fix if you have time? So that we can get this landed.