chenchenghao/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 11 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

2
docs/references/environment_variables.md
View File

@@ -164,6 +164,8 @@ SGLang supports various environment variables that can be used to configure its
| `SGLANG_WAIT_WEIGHTS_READY_TIMEOUT` | Timeout period for waiting on weights | `120` |
| `SGLANG_DISABLE_OUTLINES_DISK_CACHE` | Disable Outlines disk cache | `true` |
| `SGLANG_USE_CUSTOM_TRITON_KERNEL_CACHE` | Use SGLang's custom Triton kernel cache implementation for lower overheads (automatically enabled on CUDA) | `false` |
| `SGLANG_HICACHE_DECODE_OFFLOAD_STRIDE` | Decode-side incremental KV cache offload stride. Rounded down to a multiple of `--page-size` (min is `--page-size`). If unset/invalid/<=0, it falls back to `--page-size`. | Not set (uses `--page-size`) |
## Function Calling / Tool Use

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)

169
test/registered/disaggregation/test_disaggregation_decode_offload.py Normal file
View File

@@ -0,0 +1,169 @@
import os
import shutil
import unittest
from types import SimpleNamespace
from sglang.srt.utils import kill_process_tree
from sglang.test.ci.ci_register import register_cuda_ci
from sglang.test.run_eval import run_eval
from sglang.test.server_fixtures.disaggregation_fixture import (
PDDisaggregationServerBase,
)
from sglang.test.test_utils import (
DEFAULT_MODEL_NAME_FOR_TEST,
DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
popen_launch_pd_server,
)
# Registering the test for CUDA CI with appropriate parameters
# Increasing estimated time since we run evaluation twice
register_cuda_ci(est_time=600, suite="stage-b-test-large-2-gpu")
class TestDisaggregationDecodeOffload(PDDisaggregationServerBase):
"""
Test class for verifying KV cache offloading on the decode side in a
prefill-decode disaggregation setup.
"""
@classmethod
def setUpClass(cls):
# Set environment variable to make offloading more frequent for testing purposes
cls.old_stride = os.environ.get("SGLANG_HICACHE_DECODE_OFFLOAD_STRIDE")
cls.hicache_dir = "/tmp/hicache_test"
os.environ["SGLANG_HICACHE_FILE_BACKEND_STORAGE_DIR"] = cls.hicache_dir
os.environ["SGLANG_HICACHE_DECODE_OFFLOAD_STRIDE"] = "16"
# Ensure a clean cache directory
if os.path.exists(cls.hicache_dir):
shutil.rmtree(cls.hicache_dir)
os.makedirs(cls.hicache_dir, exist_ok=True)
super().setUpClass()
cls.model = DEFAULT_MODEL_NAME_FOR_TEST
# Non-blocking start of prefill and decode servers
cls.start_prefill()
cls.start_decode()
# Wait for both servers to be ready before proceeding
cls.wait_server_ready(cls.prefill_url + "/health")
cls.wait_server_ready(cls.decode_url + "/health")
cls.launch_lb()
@classmethod
def tearDownClass(cls):
# Restore the original environment variable state
super().tearDownClass()
if cls.old_stride is not None:
os.environ["SGLANG_HICACHE_DECODE_OFFLOAD_STRIDE"] = cls.old_stride
else:
os.environ.pop("SGLANG_HICACHE_DECODE_OFFLOAD_STRIDE", None)
os.environ.pop("SGLANG_HICACHE_FILE_BACKEND_STORAGE_DIR", None)
# Clean up the cache directory
if os.path.exists(cls.hicache_dir):
shutil.rmtree(cls.hicache_dir)
@classmethod
def start_prefill(cls):
prefill_args = [
"--trust-remote-code",
"--disaggregation-mode",
"prefill",
"--tp",
"1",
"--page-size",
"16",
"--enable-hierarchical-cache",
"--hicache-storage-backend",
"file",
"--hicache-ratio",
"2",
]
prefill_args += cls.transfer_backend + cls.rdma_devices
cls.process_prefill = popen_launch_pd_server(
cls.model,
cls.prefill_url,
timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
other_args=prefill_args,
)
@classmethod
def start_decode(cls):
decode_args = [
"--trust-remote-code",
"--disaggregation-mode",
"decode",
"--tp",
"1",
"--base-gpu-id",
"1",
"--disaggregation-decode-enable-offload-kvcache",
"--num-reserved-decode-tokens",
"128",
"--hicache-ratio",
"2",
"--page-size",
"16",
"--hicache-storage-backend",
"file",
]
decode_args += cls.transfer_backend + cls.rdma_devices
cls.process_decode = popen_launch_pd_server(
cls.model,
cls.decode_url,
timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
other_args=decode_args,
)
def test_mmlu_double_eval(self):
"""
Run two rounds of MMLU evaluation:
1. First round: Decode node offloads KV cache back to disk (HiCache).
2. Restart All Nodes to clear memory cache.
3. Second round: Prefill node loads KV cache from disk (HiCache).
Verify that both rounds produce consistent scores.
"""
args = SimpleNamespace(
base_url=f"http://{self.base_host}:{self.lb_port}",
model=self.model,
eval_name="mmlu",
num_examples=64,
num_threads=32,
)
metrics1 = run_eval(args)
# Ensure all offloads are committed to disk
import time
time.sleep(10)
kill_process_tree(self.process_prefill.pid)
kill_process_tree(self.process_decode.pid)
kill_process_tree(self.process_lb.pid)
self.process_prefill.wait()
self.process_decode.wait()
self.process_lb.wait()
self.start_prefill()
self.start_decode()
self.launch_lb()
self.wait_server_ready(self.prefill_url + "/health")
self.wait_server_ready(self.decode_url + "/health")
metrics2 = run_eval(args)
# Assert score is above a minimum threshold for both rounds
self.assertGreater(metrics1["score"], 0.65)
self.assertGreater(metrics2["score"], 0.65)
# Score should be consistent: round 2 should be >= round 1, or at least within a 0.05 margin if slightly lower
self.assertGreaterEqual(metrics2["score"], metrics1["score"] - 0.05)
if __name__ == "__main__":
unittest.main()

5
test/registered/disaggregation/test_specv2_kvcache_offloading.py
View File

@@ -25,9 +25,11 @@ def _make_mock_req(
kv_committed_len: int,
kv_allocated_len: int,
prefix_indices_len: int = 0,
rid: int = 0,
):
"""Create a mock Req with the KV cache state needed for testing."""
req = MagicMock()
req.rid = rid
req.req_pool_idx = req_pool_idx
req.kv_committed_len = kv_committed_len
req.kv_allocated_len = kv_allocated_len
@@ -74,6 +76,7 @@ def _make_manager(pool_size: int, page_size: int = 1):
manager.token_to_kv_pool_allocator = allocator
manager.page_size = page_size
manager.tree_cache = tree_cache
manager.offloaded_state = {}
return manager, freed_indices
@@ -169,7 +172,7 @@ class TestReleaseFinishedReq(unittest.TestCase):
prefix_indices_len=5,
)
manager._release_finished_req(req, prefill_offloaded_len=0)
manager._release_finished_req(req, start_offset=0)
self.assertEqual(manager.tree_cache.protected_size_, 5)