[Simulator][Compatibility] Adapt to latest KV cache pool interfaces (#40418)

This commit is contained in:
Ruiyan Ma
2026-09-20 18:00:06 +08:00
committed by GitHub
parent 0024efa0de
commit efa7be2091
3 changed files with 23 additions and 4 deletions
@@ -176,6 +176,11 @@ class C_KVCacheConfiguratorHook(BaseHook):
)
except RuntimeError:
return 0
except ValueError as error:
# New runtimes reject deliberately undersized SWA probes.
if "cannot hold even one request" not in str(error):
raise
return 0
return config.max_total_num_tokens
lower, upper = 0, 1
@@ -249,6 +254,20 @@ class C_KVCacheConfiguratorHook(BaseHook):
if hasattr(token_pool, name):
setattr(token_pool, name, value)
# Some pools derive their physical row widths while the dimensions
# above are compacted. Restore that derived geometry before host
# pools inspect the device-pool rows for HiCache sizing.
if all(
hasattr(token_pool, name)
for name in ("head_num", "head_dim", "row_dim")
):
token_pool.row_dim = token_pool.head_num * token_pool.head_dim
if all(
hasattr(token_pool, name)
for name in ("head_num", "v_head_dim", "v_row_dim")
):
token_pool.v_row_dim = token_pool.head_num * token_pool.v_head_dim
if (
hasattr(token_pool, "kv_cache_dim")
and token_pool.kv_cache_dim == 2
@@ -13,9 +13,9 @@ from test_simulation_sglang_serving import (
def test_in_process_runner_reports_each_cache_tier(tmp_path):
runner = make_sglang_runner(tmp_path)
benchmark_config = BenchmarkConfig(request_rate=10, ignore_request_timestamp=False)
cached_ds = make_fixed_dataset(1000, 8)
evict_l1_ds = make_fixed_dataset(2000, 10)
evict_l2_ds = make_fixed_dataset(3000, 20)
cached_ds = make_fixed_dataset(1000, 3)
evict_l1_ds = make_fixed_dataset(2000, 5)
evict_l2_ds = make_fixed_dataset(3000, 10)
try:
metrics = runner.benchmark(benchmark_config, dataset=cached_ds)
@@ -78,7 +78,7 @@ def make_sglang_runner(tmp_path: Path):
hicache_ratio=2,
hicache_storage_backend="file",
hicache_storage_prefetch_policy="wait_complete",
max_total_tokens=10 * 1024,
max_total_tokens=4 * 1024,
page_size=256,
skip_tokenizer_init=True,
)