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

[HiCache][LoRA] Isolate storage pages by extra key (#38577)

Browse Source 
Co-authored-by: Shuwen Wang <47200617+alphabetc1@users.noreply.github.com>
This commit is contained in:
Yanbin Jiang
2026-09-12 18:18:14 +08:00
committed by GitHub
co-authored by Shuwen Wang
parent b9cb96496d
commit 0b415fa573
21 changed files with 516 additions and 132 deletions
Download Patch File Download Diff File Expand all files Collapse all files
test/registered/e2e/hicache/test_hicache_storage_lora.py
+156
View File
@@ -0,0 +1,156 @@
"""Storage round trips preserve LoRA and salt isolation."""
import json
import os
import random
import shutil
import tempfile
import unittest
from typing import Dict, Optional
import requests
from sglang.benchmark.utils import get_tokenizer
from sglang.test.ci.ci_register import register_cuda_ci
from sglang.test.test_utils import (
DEFAULT_MODEL_NAME_FOR_TEST,
DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
DEFAULT_URL_FOR_TEST,
CustomTestCase,
popen_launch_server,
terminate_and_kill_process_tree,
)
register_cuda_ci(est_time=300, stage="base-b", runner_config="1-gpu-large")
LORA_NAME = "sql"
LORA_PATH = "philschmid/code-llama-3-1-8b-text-to-sql-lora"
PAGE_SIZE = 64
PROMPT_TOKENS = 768
class TestHiCacheStorageLoRAIsolation(CustomTestCase):
@classmethod
def setUpClass(cls):
cls.temp_dir = tempfile.mkdtemp()
cls.model = DEFAULT_MODEL_NAME_FOR_TEST
cls.base_url = DEFAULT_URL_FOR_TEST
cls.tokenizer = get_tokenizer(cls.model)
extra_config = {"hicache_storage_pass_prefix_keys": True}
other_args = [
"--enable-hierarchical-cache",
"--mem-fraction-static",
"0.6",
"--hicache-ratio",
"1.2",
"--page-size",
str(PAGE_SIZE),
"--enable-cache-report",
"--hicache-storage-prefetch-policy",
"wait_complete",
"--hicache-storage-backend",
"file",
"--hicache-storage-backend-extra-config",
json.dumps(extra_config),
"--enable-lora",
"--lora-paths",
f"{LORA_NAME}={LORA_PATH}",
"--max-loras-per-batch",
"2",
# Triton keeps radix caching enabled under deterministic inference.
"--enable-deterministic-inference",
"--attention-backend",
"triton",
]
env = {**os.environ, "SGLANG_HICACHE_FILE_BACKEND_STORAGE_DIR": cls.temp_dir}
cls.process = popen_launch_server(
cls.model,
cls.base_url,
timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
other_args=other_args,
env=env,
)
@classmethod
def tearDownClass(cls):
if getattr(cls, "process", None):
terminate_and_kill_process_tree(cls.process)
shutil.rmtree(cls.temp_dir, ignore_errors=True)
def send_request(
self,
prompt: str,
lora_path: Optional[str],
max_tokens: int = 32,
cache_salt: Optional[str] = None,
) -> Dict:
payload = {
"text": prompt,
"sampling_params": {
"temperature": 0.0,
"max_new_tokens": max_tokens,
"ignore_eos": True,
},
}
if lora_path is not None:
payload["lora_path"] = lora_path
if cache_salt is not None:
payload["cache_salt"] = cache_salt
response = requests.post(f"{self.base_url}/generate", json=payload, timeout=120)
self.assertEqual(response.status_code, 200, response.text)
return response.json()
@staticmethod
def cached_tokens(response_json: Dict) -> int:
return int(response_json.get("meta_info", {}).get("cached_tokens", 0))
def flush_device_cache(self):
# A short unrelated request first so the pages of interest get offloaded.
self.send_request(self.gen_prompt(1), lora_path=None, max_tokens=150)
res = requests.post(
f"{self.base_url}/flush_cache", params={"timeout": 30}, timeout=40
)
res.raise_for_status()
def gen_prompt(self, token_num: int) -> str:
vocab = list(self.tokenizer.get_vocab().values())
return self.tokenizer.decode(random.choices(vocab, k=token_num))
def test_adapter_pages_are_isolated_in_storage(self):
prompt = self.gen_prompt(PROMPT_TOKENS)
hit_floor = PROMPT_TOKENS - 2 * PAGE_SIZE
# Cold pass with the adapter populates host and storage.
lora_first = self.send_request(prompt, lora_path=LORA_NAME)
self.flush_device_cache()
# Read adapter pages before any base request stores the same prompt.
lora_again = self.send_request(prompt, lora_path=LORA_NAME)
self.assertGreater(
self.cached_tokens(lora_again),
hit_floor,
"the adapter's pages were not served from storage after the flush",
)
self.assertEqual(lora_first["text"], lora_again["text"])
self.flush_device_cache()
# Base and salted pages must miss existing namespaces, then round-trip.
for cache_salt in (None, "tenant-a"):
with self.subTest(cache_salt=cache_salt):
first = self.send_request(prompt, lora_path=None, cache_salt=cache_salt)
self.assertLess(self.cached_tokens(first), PAGE_SIZE)
self.flush_device_cache()
again = self.send_request(prompt, lora_path=None, cache_salt=cache_salt)
self.assertGreater(self.cached_tokens(again), hit_floor)
self.assertEqual(first["text"], again["text"])
self.flush_device_cache()
lora_third = self.send_request(prompt, lora_path=LORA_NAME)
self.assertGreater(self.cached_tokens(lora_third), hit_floor)
self.assertEqual(lora_first["text"], lora_third["text"])
if __name__ == "__main__":
unittest.main()
test/registered/unit/disaggregation/test_decode_hicache_tree_core.py
+8 -3
View File
@@ -41,7 +41,7 @@ class TestDecodeHiCacheTreeCore(CustomTestCase):
rid="req-0",
origin_input_ids=[0, 1, 2, 3, 4, 5, 6, 7],
extra_key="model",
cache_salt=None,
cache_salt="tenant-a",
)
result = SimpleNamespace(
device_indices=torch.tensor([10, 11]),
@@ -56,7 +56,12 @@ class TestDecodeHiCacheTreeCore(CustomTestCase):
self.assertEqual(prefix_match.l3_storage_hit_length, 2)
tree_cache.query_storage_hit_length.assert_called_once_with(
22, [4, 5, 6, 7], "h2", ["h0", "h1"]
22,
[4, 5, 6, 7],
"h2",
["h0", "h1"],
extra_key="model",
cache_salt="tenant-a",
)
DecodeHiCachePreallocMixin._start_hicache_prefetch(harness, req, prefix_match)
@@ -69,7 +74,7 @@ class TestDecodeHiCacheTreeCore(CustomTestCase):
"h2",
["h0", "h1"],
extra_key="model",
cache_salt=None,
cache_salt="tenant-a",
)
def test_stale_prefetch_anchor_degrades_to_l2(self):
test/registered/unit/disaggregation/test_specv2_kvcache_offloading.py
+24
View File
@@ -26,6 +26,8 @@ from sglang.srt.managers.scheduler_components.batch_result_processor import (
)
from sglang.srt.mem_cache.allocator import BaseTokenToKVPoolAllocator
from sglang.srt.mem_cache.base_prefix_cache import BasePrefixCache
from sglang.srt.mem_cache.radix_cache import RadixKey
from sglang.srt.mem_cache.utils import get_hash_str, get_storage_hash_str
from sglang.srt.runtime_context import get_context
from sglang.test.ci.ci_register import register_cpu_ci
from sglang.test.test_utils import CustomTestCase
@@ -44,6 +46,8 @@ def _make_mock_req(
"""Create a mock Req with the KV cache state needed for testing."""
req = MagicMock()
req.rid = rid
req.extra_key = None # base traffic: storage hashes chain from tokens alone
req.cache_salt = None
req.origin_input_ids = list(range(origin_len))
req.kv = ReqKvInfo(
req_pool_idx=req_pool_idx,
@@ -121,6 +125,26 @@ class _FinishedEvent:
class TestReleaseFinishedReq(unittest.TestCase):
"""Tests for _release_finished_req overallocation cleanup."""
def test_decode_offload_hash_chain_matches_prefill(self):
"""Decode pages must keep the prefill namespace across offload chunks."""
manager, _ = _make_manager(pool_size=8, page_size=2)
manager.cache_controller = MagicMock(get_hash_str=get_hash_str)
tokens = [1, 2, 3, 4, 5, 6]
for extra_key, cache_salt in [
(None, None),
("lora-a", None),
(None, "tenant-a"),
("lora-a", "tenant-a"),
]:
with self.subTest(extra_key=extra_key, cache_salt=cache_salt):
namespace = dict(extra_key=extra_key, cache_salt=cache_salt)
prefix = manager._compute_prefix_hash(tokens[:4], **namespace)
tail = manager._compute_prefix_hash(tokens[4:], prefix[-1], **namespace)
self.assertEqual(
prefix + tail,
get_storage_hash_str(RadixKey(tokens, **namespace), page_size=2),
)
def test_no_overallocation(self):
"""Without spec v2, kv_committed == kv_allocated; no extra free."""
manager, freed = _make_manager(pool_size=32)
test/registered/unit/mem_cache/test_mem_cache_utils.py
+54 -6
View File
@@ -56,10 +56,11 @@ def _legacy_page_hashes(key, page_size, prior_hash=None):
class _HashKey:
def __init__(self, token_ids, is_bigram=False, cache_salt=None):
def __init__(self, token_ids, is_bigram=False, cache_salt=None, extra_key=None):
self.token_ids = token_ids
self.is_bigram = is_bigram
self.cache_salt = cache_salt
self.extra_key = extra_key
def __len__(self):
if self.is_bigram:
@@ -75,8 +76,13 @@ class _HashKey:
self.token_ids[start : stop + 1],
is_bigram=True,
cache_salt=self.cache_salt,
extra_key=self.extra_key,
)
return _HashKey(self.token_ids[start:stop], cache_salt=self.cache_salt)
return _HashKey(
self.token_ids[start:stop],
cache_salt=self.cache_salt,
extra_key=self.extra_key,
)
if self.is_bigram:
return (self.token_ids[index], self.token_ids[index + 1])
return self.token_ids[index]
@@ -267,6 +273,52 @@ class TestGetHashStr(unittest.TestCase):
)
class TestStorageHashNamespace(unittest.TestCase):
def test_node_hashes_isolate_namespaces_and_continue_the_chain(self):
root = SimpleNamespace(parent=None, key=_HashKey(array("q")), hash_value=None)
tokens = array("q", range(1, 129))
def child(extra_key=None, cache_salt=None):
return SimpleNamespace(
parent=root,
key=_HashKey(tokens, extra_key=extra_key, cache_salt=cache_salt),
hash_value=None,
)
plain = compute_node_hash_values(child(), page_size=64)
self.assertEqual(plain, get_hash_str(tokens, None, page_size=64))
# Also guard ambiguous concatenations: ("a", "bc") vs ("ab", "c").
namespaced = [
compute_node_hash_values(child(*namespace), page_size=64)
for namespace in [
("lora-a", None),
("lora-b", None),
(None, "tenant-a"),
("lora-a", "tenant-a"),
("a", "bc"),
("ab", "c"),
]
]
for i in range(len(plain)):
page_hashes = {plain[i], *(hashes[i] for hashes in namespaced)}
self.assertEqual(len(page_hashes), 1 + len(namespaced))
# Continue the parent chain without re-seeding.
parent = child("lora-a", "tenant-a")
parent.hash_value = namespaced[3]
grand = SimpleNamespace(
parent=parent,
key=_HashKey(
array("q", range(200, 264)), extra_key="lora-a", cache_salt="tenant-a"
),
hash_value=None,
)
self.assertEqual(
compute_node_hash_values(grand, page_size=64),
get_hash_str(array("q", range(200, 264)), namespaced[3][-1], page_size=64),
)
class TestHashStrToInt64(unittest.TestCase):
def test_zero_hash(self):
result = hash_str_to_int64("0" * 64)
@@ -333,10 +385,6 @@ class TestComputeNodeHashValues(unittest.TestCase):
compute_node_event_hash_values(self._make_node(key), page_size=8),
_legacy_page_hashes(key, page_size=8, prior_hash=seed),
)
self.assertEqual(
compute_node_hash_values(self._make_node(key), page_size=8),
_legacy_page_hashes(key, page_size=8),
)
other = _HashKey(array("q", range(1, 17)), cache_salt="tenant-b")
self.assertNotEqual(
test/registered/unit/mem_cache/test_radix_cache_unit.py
+26
View File
@@ -796,6 +796,32 @@ class TestRadixCache(CustomTestCase):
]
self.assertNotEqual(unsalted_hashes, stored[0].block_hashes)
def test_extra_key_does_not_move_published_block_hashes(self):
"""Adding extra_key preserves event hashes and split-parent links."""
for cache_salt in (None, "tenant-a"):
published = []
for extra_key in (None, "lora-a"):
cache = RadixCache.create_simulated(
page_size=2, enable_kv_cache_events=True
)
namespace = dict(extra_key=extra_key, cache_salt=cache_salt)
for tokens in ([1, 2, 3, 4, 5, 6], [1, 2, 7, 8]):
cache.insert(
InsertParams(
key=RadixKey(array("q", tokens), **namespace),
value=torch.tensor(tokens, dtype=torch.int64),
)
)
published.append(
[
(event.parent_block_hash, tuple(event.block_hashes))
for event in cache.take_events()
if isinstance(event, BlockStored)
]
)
self.assertEqual(published[0], published[1])
self.assertIsNotNone(published[1][-1][0])
def test_cache_salt_event_hashes_are_preserved_across_node_split(self):
cache = RadixCache.create_simulated(page_size=2, enable_kv_cache_events=True)
original = RadixKey(array("q", [1, 2, 3, 4]), cache_salt="tenant-a")
test/registered/unit/mem_cache/test_rust_tree_core_integration.py
+14 -10
View File
@@ -52,7 +52,7 @@ from sglang.srt.mem_cache.unified_cache.cache_action import (
SWARebuild,
)
from sglang.srt.mem_cache.unified_cache.component_type import ComponentType
from sglang.srt.mem_cache.utils import hash_str_to_int64
from sglang.srt.mem_cache.utils import get_storage_hash_str, hash_str_to_int64
from sglang.srt.runtime_context import get_context
@@ -1058,21 +1058,25 @@ def test_storage_backup_spec_round_trips_the_backuped_node():
core = _tree_core(page_size=2)
core.set_hicache_enabled()
core.enable_storage = True
_insert(core, [1, 2], [10, 11])
_insert(core, [1, 2, 7, 8], [10, 11, 12, 13])
parent = core.match_prefix(MatchPrefixParams(key=_key([1, 2]))).best_match_node
child = core.match_prefix(MatchPrefixParams(key=_key([1, 2, 7, 8]))).best_match_node
key = RadixKey(
array("q", [1, 2, 7, 8]), extra_key="adapter-a", cache_salt="tenant-a"
)
for length in (2, 4):
_pump_insert(
core,
InsertParams(key=key[:length], value=torch.arange(10, 10 + length)),
)
parent = core.match_prefix(MatchPrefixParams(key=key[:2])).best_match_node
child = core.match_prefix(MatchPrefixParams(key=key)).best_match_node
core.commit_backup(parent, torch.tensor([100, 101], dtype=torch.int64), {})
core.commit_backup(child, torch.tensor([102, 103], dtype=torch.int64), {})
spec = core.build_storage_backup_spec(child, pass_prefix_keys=True)
assert spec.host_value.tolist() == [102, 103]
assert spec.token_ids == array("q", [7, 8])
parent_hashes = mem_cache.get_hash_str(array("q", [1, 2]), None, 2)
assert spec.prefix_keys == parent_hashes
assert spec.hash_value == mem_cache.get_hash_str(
array("q", [7, 8]), parent_hashes[-1], 2
)
hashes = get_storage_hash_str(key, page_size=2)
assert spec.prefix_keys == hashes[:1]
assert spec.hash_value == hashes[1:]
assert spec.comp_xfers == {}
test/registered/unit/mem_cache/test_unified_radix_cache_unittest.py
+10 -2
View File
@@ -3912,10 +3912,18 @@ class UnifiedRadixCacheSuite:
storage_dir, seq, extra_key=extra_key, cache_salt=cache_salt
)
# A root anchor has no namespace of its own. The fetched span must use
# the request namespace supplied to prefetch_from_storage.
# A root anchor has no namespace; probe and prefetch must use the request's.
cons, _, _ = build_fixture(self.cfg)
self._init_buffer_hicache(cons, storage_dir)
self.assertEqual(
cons.query_storage_hit_length(
cons.root_node_handle(),
array("q", seq),
extra_key=extra_key,
cache_salt=cache_salt,
),
len(seq),
)
root_req = "salted-root-prefetch"
cons.prefetch_from_storage(
root_req,