[HiCache] Fence full prefill CUDA-graph replay behind pending load-back

[rodamani]

Motivation

HiCache load-back gates device-state reads with Python-level per-layer waits (layer_transfer_counter.wait_until in the KV / Mamba pool accessors). A captured CUDA graph never executes those waits on replay. The breakable backend runs attention as eager graph breaks, so the waits still fire, and tc_piecewise is already disabled with hierarchical cache. The full prefill graph backend replays the whole transformer body as one graph, so a prefill batch with a pending load-back could read KV / Mamba state before the H2D copy lands. This is the open "disable cuda graph execution if hicache loading triggered" TODO in Scheduler.get_new_batch_prefill.

Modifications

• In prefill_cuda_graph_runner.py, under the full backend, enqueue one wait on the pending load op's final event on the forward stream before replay. The graph is kept; only that batch loses layerwise H2D/compute overlap.
• test/registered/unit/model_executor/test_full_prefill_graph_hicache_load_fence.py.

Accuracy Tests

CPU: 3 passed. With the runner change reverted: 3 failed. CUDA-graph replay itself was not exercised on GPU for this branch.

Speed Tests and Profiling

Adds one stream wait only for full-backend prefill batches that have a pending load-back; other batches are unchanged. Not separately benchmarked.

Checklist

• Format your code according to the Format code with pre-commit.
• Add unit tests according to the Run and add unit tests.
• Update documentation according to Write documentations.
• Provide accuracy and speed benchmark results according to Test the accuracy and Benchmark the speed.
• Follow the SGLang code style guidance.

Review and Merge Process

1. Ping Merge Oncalls to start the process. See the PR Merge Process.
2. Get approvals from CODEOWNERS and other reviewers.
3. Trigger CI tests with comments or contact authorized users to do so.
   • Common commands include /tag-and-rerun-ci, /tag-run-ci-label, /rerun-failed-ci
4. After green CI and required approvals, ask Merge Oncalls or people with Write permission to merge the PR.

---

CI States

Latest PR Test (Base): ❌ Run #36767706030
Latest PR Test (Extra): ❌ Run #36767705632
Latest PR Test (AMD ROCm 10): ❌ Run #36767706183

[rodamani]

/rerun-test -c test_full_cuda_graph_prefill.py test_hicache_variants.py

[github-actions]

Results for /rerun-test -c test_full_cuda_graph_prefill.py test_hicache_variants.py:

🚀 4-gpu-b200 (1 test): ✅ View workflow run

cd test/ && python3 registered/cuda_graph/full_prefill/test_full_cuda_graph_prefill.py

🚀 1-gpu-h100 (1 test): ✅ View workflow run

cd test/ && python3 registered/hicache/test_hicache_variants.py

⛔ test/registered/unit/model_executor/test_full_prefill_graph_hicache_load_fence.py: File not found: test/registered/unit/model_executor/test_full_prefill_graph_hicache_load_fence.py

[rodamani]

/tag-and-rerun-ci