[UnifiedTree]: Fix Unified HiCache tombstone lock release replay

[hzh0425]

Motivation

Record component node ids skipped during lock acquire and pass them back
to release. This prevents a temporary admission lock from
decrementing SWA or Mamba locks that were acquired later after HiCache
load-back revived a tombstoned device value.

The main race is _lock_node -> init_load_back -> release: acquire sees
a component tombstone and skips it, load-back restores it, then release
must not consume the newer load-back/request lock.

Modifications

Accuracy Tests

Speed Tests and Profiling

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.

[hzh0425]

/rerun-test test/registered/unit/mem_cache/test_unified_radix_cache_unittest.py test/registered/radix_cache/test_unified_radix_cache_kl.py

[github-actions]

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

cd test/ && python3 registered/unit/mem_cache/test_unified_radix_cache_unittest.py

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

cd test/ && python3 registered/radix_cache/test_unified_radix_cache_kl.py

[gemini-code-assist]

Code Review

This pull request fixes a bug in the unified cache locking mechanism where temporary locks could incorrectly release nodes that were restored from tombstones while the lock was held. It introduces skip_lock_node_ids in the lock parameters to track nodes skipped during acquisition, ensuring that the release operation accurately mirrors the acquisition state. The fix is implemented for both SWA and Mamba components, and new unit tests have been added to verify that restored tombstones are not prematurely released. I have no feedback to provide as there were no review comments.

[baoskee]

I've been running write_through_selective on this branch for last 8 hours. Will let you know if it is stable. Thank you!

[baoskee]

Hi, I am still getting the same error:

         ^^^^^^^^^^^^^^^^^^
  File "/home/baoskee/.venv/lib/python3.12/site-packages/sglang/srt/managers/schedule_policy.py", line 878, in add_one_req
    new_indices, req.last_node = self.tree_cache.init_load_back(
                                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/baoskee/.venv/lib/python3.12/site-packages/sglang/srt/mem_cache/unified_radix_cache.py", line 1390, in init_load_back
    loading_values = self.load_back(last_node, mem_quota, req=req)
                     ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/baoskee/.venv/lib/python3.12/site-packages/sglang/srt/mem_cache/unified_radix_cache.py", line 1240, in load_back
    t = comp.build_hicache_transfers(
        ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/baoskee/.venv/lib/python3.12/site-packages/sglang/srt/mem_cache/unified_cache_components/swa_component.py", line 448, in build_hicache_transfers
    assert cd.host_value is not None or cd.value is not None
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
AssertionError

[2026-05-12 06:06:13 TP1] Scheduler hit an exception: Traceback (most recent call last):
  File "/home/baoskee/.venv/lib/python3.12/site-packages/sglang/srt/managers/scheduler.py", line 4001, in run_scheduler_process
    scheduler.run_event_loop()
  File "/home/baoskee/.venv/lib/python3.12/site-packages/sglang/srt/managers/scheduler.py", line 1521, in run_event_loop
    dispatch_event_loop(self)
  File "/home/baoskee/.venv/lib/python3.12/site-packages/sglang/srt/managers/scheduler.py", line 3871, in dispatch_event_loop
    scheduler.event_loop_overlap()
  File "/home/baoskee/.venv/lib/python3.12/site-packages/torch/utils/_contextlib.py", line 124, in decorate_context
    return func(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^
  File "/home/baoskee/.venv/lib/python3.12/site-packages/sglang/srt/managers/scheduler.py", line 1570, in event_loop_overlap
    batch = self.get_next_batch_to_run()
            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/baoskee/.venv/lib/python3.12/site-packages/sglang/srt/managers/scheduler.py", line 2538, in get_next_batch_to_run
    new_batch = self.get_new_batch_prefill()
                ^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/baoskee/.venv/lib/python3.12/site-packages/sglang/srt/managers/scheduler.py", line 2592, in get_new_batch_prefill
    ret = self._get_new_batch_prefill_raw(
          ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/baoskee/.venv/lib/python3.12/site-packages/sglang/srt/managers/scheduler.py", line 2723, in _get_new_batch_prefill_raw
    res = adder.add_one_req(
          ^^^^^^^^^^^^^^^^^^
  File "/home/baoskee/.venv/lib/python3.12/site-packages/sglang/srt/managers/schedule_policy.py", line 878, in add_one_req
    new_indices, req.last_node = self.tree_cache.init_load_back(
                                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/baoskee/.venv/lib/python3.12/site-packages/sglang/srt/mem_cache/unified_radix_cache.py", line 1390, in init_load_back
    loading_values = self.load_back(last_node, mem_quota, req=req)
                     ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/baoskee/.venv/lib/python3.12/site-packages/sglang/srt/mem_cache/unified_radix_cache.py", line 1240, in load_back
    t = comp.build_hicache_transfers(
        ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/baoskee/.venv/lib/python3.12/site-packages/sglang/srt/mem_cache/unified_cache_components/swa_component.py", line 448, in build_hicache_transfers
    assert cd.host_value is not None or cd.value is not None
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
AssertionError

[2026-05-12 06:06:13] SIGQUIT received. signum=None, frame=None. It usually means one child failed.

It looks like all 4 workers failed, but over time rather than all at once (see TFTT spike and then no requests being served).

Same configuration:

CUDA_VISIBLE_DEVICES=0,1 \
SGLANG_ENABLE_UNIFIED_RADIX_TREE=1 \
sglang serve \
  --model-path google/gemma-4-31B-it \
  --host 0.0.0.0 --port 30001 \
  --tp 2 \
  --context-length 32768 \
  --mem-fraction-static 0.7 \
  --kv-cache-dtype fp8_e4m3 \
  ...
  --enable-cache-report \
  --enable-metrics \
  --trust-remote-code \
  --page-size 64 \
  --enable-hierarchical-cache \
  --hicache-ratio 1 \
  --hicache-write-policy write_through_selective \
  --hicache-storage-prefetch-policy wait_complete \
  --hicache-io-backend kernel \
  --hicache-mem-layout page_first \
  --trust-remote-code \
  --log-level info

[baoskee]

Our full dashboard if this helps.

[hzh0425]

**baoskee **

@baoskee how can I reach you in slack?

BTW, I noticed that in your setup, hicache-ratio=1, which means the host memory size is equal to the device memory size.
Usually, we set it to at least 2.

[baoskee]

I sent an invitation on Slack to hzh0425@apache.org. I had to set to hicache-ratio=1 because our node did not have enough System RAM.

For some reason, the workers starting process crash (won't even run) if I set it to hicache ratio = 2.

The node are B200s on 1800GiB RAM machine. You would think that's enough RAM....