Skip to content

Hierarchical Caching for SGLang#2693

Merged
Ying1123 merged 67 commits intomainfrom
xiezhq-hierarchical
Feb 24, 2025
Merged

Hierarchical Caching for SGLang#2693
Ying1123 merged 67 commits intomainfrom
xiezhq-hierarchical

Conversation

@xiezhq-hermann
Copy link
Copy Markdown
Collaborator

@xiezhq-hermann xiezhq-hermann commented Jan 1, 2025 •
edited
Loading

Motivation

While RadixTree-based context caching provides significant performance benefits, these gains are not always fully realized. A key bottleneck is the capacity limit of GPU memory. Currently, SGLang stores historical KV caches exclusively in GPU memory; whenever more memory is required for batch execution, existing caches are discarded.

To address this issue, we propose a hierarchical caching mechanism for LLM serving, treating GPU memory as an L1 cache, host memory as an L2 cache, and disk as an L3 cache (future). This PR introduces such a mechanism in SGLang through a separate host memory pool that backs up KV caches, allowing them to be reloaded into GPU memory when needed.

Modifications

  • A HiRadixCache that extends RadixCache with host memory addresses and synchronization mechanisms.
  • A host memory pool that synchronizes with the device memory pool of KV caches.
  • A memory controller that implements efficient data transfer between host and device, and handles various cache write policies for hierarchical caching.

Todo:

  • Update benchmark results.
  • Remove deprecated design and implementation.

Checklist

  • Format your code according to the Contributor Guide.
  • Add unit tests as outlined in the Contributor Guide.
  • Update documentation as needed, including docstrings or example tutorials.

@zhyncs
Copy link
Copy Markdown
Collaborator

zhyncs commented Jan 1, 2025

It's amazing! Happy new year!

@zhyncs zhyncs added the enhancement New feature or request label Jan 1, 2025
merrymercy
merrymercy reviewed Jan 2, 2025
View reviewed changes
Comment thread python/sglang/srt/mem_cache/radix_cache.py Outdated
zhaohaidao
zhaohaidao reviewed Jan 4, 2025
View reviewed changes
Comment thread python/sglang/srt/mem_cache/radix_cache.py Outdated
Comment thread python/sglang/srt/mem_cache/radix_cache.py Outdated
@Ying1123 Ying1123 merged commit 6c7a152 into main Feb 24, 2025
@Ying1123 Ying1123 deleted the xiezhq-hierarchical branch February 24, 2025 05:56
@lambert0312
Copy link
Copy Markdown
Contributor

lambert0312 commented Feb 25, 2025

DeepSeek MLA is not supported yet, and an error will be reported when starting the model:

  File "/sgl-workspace/sglang/python/sglang/srt/managers/scheduler.py", line 1849, in run_scheduler_process
    scheduler = Scheduler(server_args, port_args, gpu_id, tp_rank, dp_rank)
  File "/sgl-workspace/sglang/python/sglang/srt/managers/scheduler.py", line 305, in __init__
    HiRadixCache(
  File "/sgl-workspace/sglang/python/sglang/srt/mem_cache/hiradix_cache.py", line 26, in __init__
    self.token_to_kv_pool_host = MLATokenToKVPoolHost(token_to_kv_pool)
  File "/sgl-workspace/sglang/python/sglang/srt/mem_cache/memory_pool.py", line 461, in __init__
    self.head_num = device_pool.head_num
AttributeError: 'MLATokenToKVPool' object has no attribute 'head_num'

@xiezhq-hermann
Copy link
Copy Markdown
Collaborator Author

xiezhq-hermann commented Feb 25, 2025

DeepSeek MLA is not supported yet, and an error will be reported when starting the model:

  File "/sgl-workspace/sglang/python/sglang/srt/managers/scheduler.py", line 1849, in run_scheduler_process
    scheduler = Scheduler(server_args, port_args, gpu_id, tp_rank, dp_rank)
  File "/sgl-workspace/sglang/python/sglang/srt/managers/scheduler.py", line 305, in __init__
    HiRadixCache(
  File "/sgl-workspace/sglang/python/sglang/srt/mem_cache/hiradix_cache.py", line 26, in __init__
    self.token_to_kv_pool_host = MLATokenToKVPoolHost(token_to_kv_pool)
  File "/sgl-workspace/sglang/python/sglang/srt/mem_cache/memory_pool.py", line 461, in __init__
    self.head_num = device_pool.head_num
AttributeError: 'MLATokenToKVPool' object has no attribute 'head_num'

Thank you @lambert0312 for pointing out, yes, this feature is still under meta stage and currently only supported MHA and GQA style memory pool. I will keep you posted once MLA is supported, which should be soon.
For further question about this feature, feel free to reach out to me on SGLang slack for a more prompt reply.

@lambert0312
Copy link
Copy Markdown
Contributor

lambert0312 commented Feb 25, 2025 •
edited
Loading

Thank you @lambert0312 for pointing out, yes, this feature is still under meta stage and currently only supported MHA and GQA style memory pool. I will keep you posted once MLA is supported, which should be soon.
For further question about this feature, feel free to reach out to me on SGLang slack for a more prompt reply.

Thanks @xiezhq-hermann

@zhaochenyang20 zhaochenyang20 mentioned this pull request Mar 3, 2025
Closed
22 tasks
@zhyncs zhyncs mentioned this pull request Mar 4, 2025
Closed
67 tasks
@xiezhq-hermann
Copy link
Copy Markdown
Collaborator Author

xiezhq-hermann commented Mar 4, 2025 •
edited
Loading

Thank you @lambert0312 for pointing out, yes, this feature is still under meta stage and currently only supported MHA and GQA style memory pool. I will keep you posted once MLA is supported, which should be soon.
For further question about this feature, feel free to reach out to me on SGLang slack for a more prompt reply.

Thanks @xiezhq-hermann

@lambert0312 just FYI, there is a PR from the community supporting MLA with hierarchical caching, which will be merged soon but feel free to check it out: #4009

@lambert0312
Copy link
Copy Markdown
Contributor

lambert0312 commented Mar 4, 2025

@lambert0312 just FYI, there is a PR from the community supporting MLA with hierarchical caching, which will be merged soon but feel free to check it out: #4009

@xiezhq-hermann Thanks, but I've encountered a problem. I just experimented with #4009 and found that there is indeed a concurrency problem when TP>1. The program will enter a locked state. There may be a concurrency problem. Please follow up. Thank you!

@shensimeteor
Copy link
Copy Markdown

shensimeteor commented Mar 4, 2025

After code cleaning and basic performance benchmark, this PR is ready to merge. You can add --enable-hierarchical-cache option when starting a SGLang server to turn on this feature. This feature will still be under active development in the future months, your feedback will be greatly welcomed : ) Following is a throughput v.s. median TTFT curve that demonstrates the benefit of hierarchical caching using a synthetic multi-turn benchmark, and you can reproduce it with Qwen/Qwen2.5-14B-Instruct on an A100-80G GPU as explained here:

throughput_latency_curve

Besides --enable-hierarchical-cache, do we also need to set cpu_offload_gb?

@xiezhq-hermann xiezhq-hermann mentioned this pull request Mar 5, 2025
Merged
6 tasks
@xiezhq-hermann
Copy link
Copy Markdown
Collaborator Author

xiezhq-hermann commented Mar 5, 2025

After code cleaning and basic performance benchmark, this PR is ready to merge. You can add --enable-hierarchical-cache option when starting a SGLang server to turn on this feature. This feature will still be under active development in the future months, your feedback will be greatly welcomed : ) Following is a throughput v.s. median TTFT curve that demonstrates the benefit of hierarchical caching using a synthetic multi-turn benchmark, and you can reproduce it with Qwen/Qwen2.5-14B-Instruct on an A100-80G GPU as explained here:
throughput_latency_curve

Besides --enable-hierarchical-cache, do we also need to set cpu_offload_gb?

Right now it allocate a host memory pool which is 4 times of the size of the device memory pool by default, so no need to set other things but more options will be added.

@JensenFire JensenFire mentioned this pull request Apr 3, 2025
Merged
6 tasks
@Fridge003 Fridge003 mentioned this pull request Apr 8, 2025
Closed
This was referenced Jun 17, 2025
Merged
Merged
@wangyibin-gh
Copy link
Copy Markdown

wangyibin-gh commented Jun 24, 2025

Hi I'm wondering - when are you planning to support L3 cache? I think it's reasonable to support pluggable L3 caches, which encourages storage providers to implement their L3 caches according to their product features. What you need to do is to define a bunch of kv cache apis for getting/putting/evicting kv cache chunk/item and give them some demo implentation using something like local SSD.

@msharmavikram
Copy link
Copy Markdown

msharmavikram commented Jun 24, 2025

This is in works @wangyibin-gh !

@wangyibin-gh
Copy link
Copy Markdown

wangyibin-gh commented Jun 25, 2025

This is in works @wangyibin-gh !

when do you expect this feature can be merged? and btw is there any documentation about it, especially w.r.t the APIs.

@xiezhq-hermann xiezhq-hermann mentioned this pull request Jul 11, 2025
Merged
6 tasks
@hubertlu-tw hubertlu-tw mentioned this pull request Jul 21, 2025
Merged
6 tasks
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

@merrymercy merrymercy merrymercy left review comments

@Ying1123 Ying1123 Ying1123 approved these changes

@zhyncs zhyncs Awaiting requested review from zhyncs

@hnyls2002 hnyls2002 Awaiting requested review from hnyls2002 hnyls2002 is a code owner

@ispobock ispobock Awaiting requested review from ispobock ispobock is a code owner

@ByronHsu ByronHsu Awaiting requested review from ByronHsu

+1 more reviewer

@zhaohaidao zhaohaidao zhaohaidao left review comments

Reviewers whose approvals may not affect merge requirements

Assignees

@Ying1123 Ying1123

Labels

enhancement New feature or request

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.