diff --git a/docs/references/environment_variables.md b/docs/references/environment_variables.md index 17f115f9a..5702f196f 100644 --- a/docs/references/environment_variables.md +++ b/docs/references/environment_variables.md @@ -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 diff --git a/python/sglang/srt/disaggregation/decode_kvcache_offload_manager.py b/python/sglang/srt/disaggregation/decode_kvcache_offload_manager.py index 6abd49144..7e969c1ac 100644 --- a/python/sglang/srt/disaggregation/decode_kvcache_offload_manager.py +++ b/python/sglang/srt/disaggregation/decode_kvcache_offload_manager.py @@ -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) diff --git a/python/sglang/srt/disaggregation/kv_events.py b/python/sglang/srt/disaggregation/kv_events.py index d256a7e9f..992798c79 100644 --- a/python/sglang/srt/disaggregation/kv_events.py +++ b/python/sglang/srt/disaggregation/kv_events.py @@ -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] diff --git a/python/sglang/srt/environ.py b/python/sglang/srt/environ.py index 90a73c00c..5399fd2a4 100644 --- a/python/sglang/srt/environ.py +++ b/python/sglang/srt/environ.py @@ -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. diff --git a/python/sglang/srt/managers/cache_controller.py b/python/sglang/srt/managers/cache_controller.py index a0e5bcc1b..ff3376ca9 100644 --- a/python/sglang/srt/managers/cache_controller.py +++ b/python/sglang/srt/managers/cache_controller.py @@ -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() diff --git a/python/sglang/srt/managers/scheduler_output_processor_mixin.py b/python/sglang/srt/managers/scheduler_output_processor_mixin.py index c56a92217..53c4e31c0 100644 --- a/python/sglang/srt/managers/scheduler_output_processor_mixin.py +++ b/python/sglang/srt/managers/scheduler_output_processor_mixin.py @@ -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) diff --git a/test/registered/disaggregation/test_disaggregation_decode_offload.py b/test/registered/disaggregation/test_disaggregation_decode_offload.py new file mode 100644 index 000000000..7088ee5a0 --- /dev/null +++ b/test/registered/disaggregation/test_disaggregation_decode_offload.py @@ -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() diff --git a/test/registered/disaggregation/test_specv2_kvcache_offloading.py b/test/registered/disaggregation/test_specv2_kvcache_offloading.py index aa635c9e1..64afa7b4a 100644 --- a/test/registered/disaggregation/test_specv2_kvcache_offloading.py +++ b/test/registered/disaggregation/test_specv2_kvcache_offloading.py @@ -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)