Fix LMCache unit test and init bug (#14005)

Signed-off-by: DongDongJu <commisori28@gmail.com>
This commit is contained in:
Dongjoo Seo
2025-11-30 07:57:32 -08:00
committed by GitHub
parent dc7bdc7329
commit c15c864b6f

View File

@@ -20,6 +20,7 @@ except ImportError as e:
"LMCache is not installed. Please install it by running `pip install lmcache`"
) from e
if TYPE_CHECKING:
from sglang.srt.configs.model_config import ModelConfig
from sglang.srt.managers.schedule_batch import Req
@@ -211,7 +212,18 @@ class LMCRadixCache(RadixCache):
if not is_insert:
return
kv_committed_len = req.pop_committed_kv_cache()
from sglang.srt.server_args import get_global_server_args
global_server_args = get_global_server_args()
topk = global_server_args.speculative_eagle_topk
enable_kv_committed_len = topk is None or topk == 1
if enable_kv_committed_len:
kv_committed_len = req.kv_committed_len
else:
kv_committed_len = len(req.origin_input_ids) + max(
len(req.output_ids) - 1, 0
)
token_ids = (req.origin_input_ids + req.output_ids)[:kv_committed_len]
kv_indices = self.req_to_token_pool.req_to_token[
req.req_pool_idx, :kv_committed_len
@@ -254,27 +266,3 @@ class LMCRadixCache(RadixCache):
)
except Exception: # pragma: no cover
pass
if __name__ == "__main__":
from sglang.srt.mem_cache.cache_init_params import CacheInitParams
params = CacheInitParams(
req_to_token_pool=None,
token_to_kv_pool_allocator=None,
page_size=1,
disable=False,
enable_kv_cache_events=False,
)
cache = LMCRadixCache(
params=params,
model_config=None,
tp_size=1,
rank=0,
tp_group=None,
)
cache.insert(RadixKey([1, 2, 3]), torch.tensor([10, 11, 12], dtype=torch.int64))
cache.insert(
RadixKey([1, 2, 3, 4]), torch.tensor([10, 11, 12, 13], dtype=torch.int64)
)
cache.pretty_print()