chenchenghao/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

[Performance] Decode Offload improves the long texts performance 100% through dynamic block offload. (#17216)

Browse Source 
Co-authored-by: zhangheng <hzh0425@apache.org>
This commit is contained in:
yuyu5333
2026-03-08 17:16:53 +08:00
committed by GitHub
parent 97a2a9be0f
commit 230fb55899
8 changed files with 290 additions and 33 deletions
Download Patch File Download Diff File Expand all files Collapse all files

119
python/sglang/srt/disaggregation/decode_kvcache_offload_manager.py
View File

@@ -7,6 +7,8 @@ from typing import TYPE_CHECKING
import torch
from sglang.srt.disaggregation.kv_events import OffloadedState
from sglang.srt.environ import envs
from sglang.srt.managers.cache_controller import HiCacheController
from sglang.srt.mem_cache.allocator import BaseTokenToKVPoolAllocator
from sglang.srt.mem_cache.base_prefix_cache import BasePrefixCache
@@ -45,6 +47,13 @@ class DecodeKVCacheOffloadManager:
self.server_args = server_args
self.request_counter = 0
self.tree_cache = tree_cache
env_stride = envs.SGLANG_HICACHE_DECODE_OFFLOAD_STRIDE.get()
if env_stride is None or env_stride <= 0:
self.offload_stride = self.page_size
else:
self.offload_stride = max(
self.page_size, (env_stride // self.page_size) * self.page_size
)
kv_cache = self.token_to_kv_pool_allocator.get_kvcache()
if isinstance(kv_cache, MHATokenToKVPool):
self.decode_host_mem_pool = MHATokenToKVPoolHost(
@@ -82,6 +91,7 @@ class DecodeKVCacheOffloadManager:
self.ongoing_offload = {}
self.ongoing_backup = {}
self.offloaded_state = {}
logger.info("Enable offload kv cache for decode side")
def offload_kv_cache(self, req) -> bool:
@@ -102,23 +112,38 @@ class DecodeKVCacheOffloadManager:
prefill_offloaded_len = (
len(req.origin_input_ids) // self.page_size * self.page_size
)
incremental_len = len(all_tokens) - prefill_offloaded_len
incremental_aligned_len = incremental_len // self.page_size * self.page_size
state = self.offloaded_state.get(req.rid)
if state is None:
prefill_hashes = self._compute_prefix_hash(
req.origin_input_ids[:prefill_offloaded_len]
)
last_prefill_hash = (
prefill_hashes[-1] if prefill_offloaded_len > 0 else None
)
state = OffloadedState(
prefill_len=prefill_offloaded_len,
inc_len=0,
last_hash=last_prefill_hash,
)
self.offloaded_state[req.rid] = state
incremental_total = len(all_tokens) - state.prefill_len
incremental_new = incremental_total - state.inc_len
incremental_aligned_len = (
incremental_new // self.offload_stride * self.offload_stride
)
if incremental_aligned_len == 0:
return False
# Extract incremental tokens and indices
start, end = (
prefill_offloaded_len,
prefill_offloaded_len + incremental_aligned_len,
)
# Extract incremental tokens and indices for the newly available chunk
start = state.prefill_len + state.inc_len
end = start + incremental_aligned_len
incremental_tokens = all_tokens[start:end]
incremental_indices = token_indices[start:end]
# Early free prefill-offloaded GPU memory
if prefill_offloaded_len > 0:
self.token_to_kv_pool_allocator.free(token_indices[:prefill_offloaded_len])
if state.prefill_len > 0 and state.inc_len == 0:
self.token_to_kv_pool_allocator.free(token_indices[: state.prefill_len])
# Asynchronously offload incremental KV cache from device to host
self.request_counter += 1
@@ -136,8 +161,10 @@ class DecodeKVCacheOffloadManager:
host_indices,
incremental_tokens,
time.time(),
prefill_offloaded_len,
start,
end,
)
state.inc_len += incremental_aligned_len
return True
def check_offload_progress(self):
@@ -171,26 +198,36 @@ class DecodeKVCacheOffloadManager:
host_indices,
incremental_tokens,
start_time,
prefill_offloaded_len,
start,
end,
) = self.ongoing_offload.pop(ack_id)
self._release_finished_req(req, prefill_offloaded_len)
self._trigger_backup(
req,
host_indices,
incremental_tokens,
start_time,
prefill_offloaded_len,
if req.finished():
self._release_finished_req(req, start)
else:
kv_indices = self.req_to_token_pool.req_to_token[
req.req_pool_idx, start:end
]
self.token_to_kv_pool_allocator.free(kv_indices)
prior_hash = (
self.offloaded_state[req.rid].last_hash
if req.rid in self.offloaded_state
else None
)
last_hash = self._trigger_backup(
req, host_indices, incremental_tokens, start_time, prior_hash
)
if req.rid in self.offloaded_state:
self.offloaded_state[req.rid].last_hash = last_hash
finish_count -= 1
def _release_finished_req(self, req: Req, prefill_offloaded_len: int):
def _release_finished_req(self, req: Req, start_offset: int):
kv_committed_len = req.pop_committed_kv_cache()
kv_indices = self.req_to_token_pool.req_to_token[
req.req_pool_idx, prefill_offloaded_len:kv_committed_len
]
# Free the incremental part of the request
start = start_offset
end = kv_committed_len
# Free the incremental part of the request (NSA-aware)
kv_indices = self.req_to_token_pool.req_to_token[req.req_pool_idx, start:end]
self.token_to_kv_pool_allocator.free(kv_indices)
# Free over-allocated KV cache slots (e.g. from speculative decoding v2).
@@ -206,6 +243,8 @@ class DecodeKVCacheOffloadManager:
self.req_to_token_pool.free(req)
self.tree_cache.protected_size_ -= len(req.prefix_indices)
if req.rid in self.offloaded_state:
del self.offloaded_state[req.rid]
def _check_backup_progress(self, finish_count):
"""Check the progress of backup from host to storage."""
@@ -222,21 +261,17 @@ class DecodeKVCacheOffloadManager:
)
def _trigger_backup(
self, req, host_indices, incremental_tokens, start_time, prefill_offloaded_len
self, req, host_indices, incremental_tokens, start_time, prior_hash
):
"""Trigger async backup from host to storage."""
prefill_hashes = self._compute_prefix_hash(
req.origin_input_ids[:prefill_offloaded_len]
)
last_prefill_hash = prefill_hashes[-1] if prefill_offloaded_len > 0 else ""
page_hashes = self._compute_prefix_hash(incremental_tokens, last_prefill_hash)
page_hashes = self._compute_prefix_hash(incremental_tokens, prior_hash)
ack_id = self.cache_controller.write_storage(
host_indices,
incremental_tokens,
hash_value=page_hashes,
)
self.ongoing_backup[ack_id] = (req.rid, host_indices, start_time)
return page_hashes[-1] if len(page_hashes) > 0 else prior_hash
def _compute_prefix_hash(self, tokens, prior_hash=""):
page_hashes = []
@@ -246,3 +281,25 @@ class DecodeKVCacheOffloadManager:
last_hash = self.cache_controller.get_hash_str(page_tokens, last_hash)
page_hashes.append(last_hash)
return page_hashes
def finalize_release_on_finish(self, req: Req):
"""Free any remaining tail KV that was not offloaded due to non-aligned length."""
if req.req_pool_idx == -1:
return
state = self.offloaded_state.get(req.rid)
if state is None:
prefill_len = len(req.origin_input_ids) // self.page_size * self.page_size
inc_len = 0
else:
prefill_len = state.prefill_len
inc_len = state.inc_len
# If no incremental offload ever happened, the prefill-aligned part was never freed.
# Free the prefill portion on request finish to avoid leaks.
if prefill_len > 0 and inc_len == 0:
token_indices = self.req_to_token_pool.req_to_token[req.req_pool_idx]
self.token_to_kv_pool_allocator.free(token_indices[:prefill_len])
logger.info(
f"Finalize release: freed prefill-aligned KV for req {req.rid}, len:{prefill_len}"
)
start_offset = prefill_len + inc_len
self._release_finished_req(req, start_offset)

17
python/sglang/srt/disaggregation/kv_events.py
View File

@@ -61,6 +61,23 @@ MEDIUM_GPU = "GPU"
MEDIUM_CPU = "CPU_PINNED"
class OffloadedState:
"""
OffloadedState represents the state of a KV cache block offloaded to the hicache.
- prefill_len (int): The length of the prefill part of the KV cache block.
- inc_len (int): The length of the incremental part of the KV cache block.
- last_hash (Optional[str]): The hash of the last token in the KV cache block.
"""
def __init__(
self, prefill_len: int, inc_len: int = 0, last_hash: Optional[str] = None
):
self.prefill_len = prefill_len
self.inc_len = inc_len
self.last_hash = last_hash
class BlockStored(KVCacheEvent):
block_hashes: list[int]
parent_block_hash: Optional[int]

1
python/sglang/srt/environ.py
View File

@@ -276,6 +276,7 @@ class Envs:
# Hi-Cache
SGLANG_HICACHE_HF3FS_CONFIG_PATH = EnvStr(None)
SGLANG_HICACHE_DECODE_OFFLOAD_STRIDE = EnvInt(None)
SGLANG_HICACHE_FILE_BACKEND_STORAGE_DIR = EnvStr(None)
SGLANG_HICACHE_NIXL_BACKEND_STORAGE_DIR = EnvStr(None)
# Max fraction of cache (by token count) that can be pinned; 0 = disable pinning.

2
python/sglang/srt/managers/cache_controller.py
View File

@@ -576,6 +576,8 @@ class HiCacheController:
model_name: Optional[str] = None,
storage_backend_extra_config: Optional[dict] = None,
):
if storage_backend_extra_config is None:
storage_backend_extra_config = {}
if is_dp_attention_enabled():
self.tp_rank = get_attention_tp_rank()

8
python/sglang/srt/managers/scheduler_output_processor_mixin.py
View File

@@ -412,6 +412,12 @@ class SchedulerOutputProcessorMixin:
req.check_finished(new_accepted_len)
if (
self.server_args.disaggregation_decode_enable_offload_kvcache
and not req.finished()
):
self.decode_offload_manager.offload_kv_cache(req)
if req.finished():
# delete feature to save memory
if req.multimodal_inputs is not None:
@@ -425,7 +431,7 @@ class SchedulerOutputProcessorMixin:
if self.server_args.disaggregation_decode_enable_offload_kvcache:
# Asynchronously offload KV cache; release_kv_cache will be called after Device->Host transfer completes
if not self.decode_offload_manager.offload_kv_cache(req):
release_kv_cache(req, self.tree_cache)
self.decode_offload_manager.finalize_release_on_finish(req)
else:
release_kv_cache(req, self.tree_cache)