Add release_kv_cache() API

[andakai]

Motivation

Sleep level 2 can release KV cache memory, but it also discards more state than needed for cases where the model weights should stay resident. This PR provides a direct API for the common case where an upper-level framework wants to temporarily free KV cache memory, then later call wake_up() and continue using the same model instance.

Summary

This PR adds release_kv_cache() API that releases KV cache memory while keeping model weights resident.

The change is intentionally scoped to KV cache release only, rather than introducing generic tag-level memory management in the public API such as this pr did #42658.

Changes include:

• Add LLM.release_kv_cache() and async/server plumbing for /release_kv_cache.
• Add executor, engine core, and worker support for releasing only KV cache memory.
• Add allocator-level release_tags() support for CUDA/XPU memory pools.
• Track whether allocator handles are currently mapped to avoid double unmap/map when release_kv_cache(), sleep(), and wake_up() are mixed.
• Add Prometheus sleep-state metric label kv_cache_released.
• Record kv_cache_released only when KV cache release actually happens.
• Add tests for offline LLM.release_kv_cache() and dev server metrics.

Similar work

#42658. This PR adds a tag parameter to sleep, but its semantics are unclear—it is not obvious whether the tag represents discard or offload. This pr also proposed another approach, such as release(tag=..., offload_tag=...), this would introduce significant code changes and overlap with existing sleep levels. Based on my investigation, the current need is limited to releasing KV cache only, so we suggest adding a dedicated interface for this use case. More fine-grained tag management can be abstracted later if needed. The PR has already implemented a release_tag interface in cumem.py, which can be easily extended in the future.

Tests

pytest tests/basic_correctness/test_mem.py -k release_kv_cache -v

Passed.

[claude]

Claude Code Review

This pull request is from a fork — automated review is disabled. A repository maintainer can comment @claude review to run a one-time review.

[mergify]

This pull request has merge conflicts that must be resolved before it can be
merged. Please rebase the PR, @andakai.

https://docs.github.com/en/pull-requests/collaborating-with-pull-requests/working-with-forks/syncing-a-fork

[AlanFokCo]

Hi @andakai — this is closely related to #46438 (CuMemAllocator.discard()), which provides the same "release KV cache while keeping weights resident" capability as a general-purpose allocator primitive:

allocator.discard(tags=("kv_cache",))

The overlap:

• Your release_tags() ≈ our discard(tags) — both call unmap_and_release for specific tags and leave others mapped
• Your mapped-state tracking ≈ the is_asleep flag in AllocationData — both prevent double unmap/map

The difference is scope: discard() is a generic allocator method (works for any tag), while release_kv_cache() adds the full API surface (LLM/AsyncLLM/server endpoint/Prometheus). Both are useful — discard() is the building block, release_kv_cache() is the user-facing convenience.

Would it make sense to build release_kv_cache() on top of discard() once #46438 lands? That way the allocator-level logic is shared and the higher layers just call allocator.discard(("kv_cache",)) instead of reimplementing the unmap + flag tracking.

[andakai]

Would it make sense to build release_kv_cache() on top of discard() once #46438 lands? That way the allocator-level logic is shared and the higher layers just call allocator.discard(("kv_cache",)) instead of reimplementing the unmap + flag tracking.

Sure. I did the allocator-level logic and the higher layers all in a single pr now. I checked we share the same logic betweeen release_tag and discard. I think separating these two can be fine and be more clear to do code review. If your pr is merged, I can change my pr accordingly. Thanks!

[AlanFokCo]

Sounds good — I'll ping you once #46438 is merged so you can rebase on top of it. Happy to help review the integration on your side too.