Skip to content

Commit 926d483

Browse files
committed
rename cache_type to prefix_sharing_algorithm
1 parent c774c0d commit 926d483

File tree

3 files changed

+9
-5
lines changed

3 files changed

+9
-5
lines changed

app_tests/integration_tests/llm/shortfin/conftest.py

Lines changed: 5 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -83,7 +83,11 @@ def model_test_dir(request, tmp_path_factory):
8383
"prefill_batch_sizes": batch_sizes,
8484
"decode_batch_sizes": batch_sizes,
8585
"transformer_block_count": 26,
86-
"paged_kv_cache": {"block_seq_stride": 16, "device_block_count": 256},
86+
"paged_kv_cache": {
87+
"block_seq_stride": 16,
88+
"device_block_count": 256,
89+
"prefix_sharing_algorithm": "none",
90+
},
8791
}
8892
logger.info(f"Saving edited config to: {edited_config_path}\n")
8993
logger.info(f"Config: {json.dumps(config, indent=2)}")

shortfin/python/shortfin_apps/llm/components/config_struct.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -86,7 +86,7 @@ class PagedKVCacheParams:
8686
# Size of the cache on each device.
8787
device_block_count: int
8888

89-
cache_type: str = "base" # currently supporting base and trie
89+
prefix_sharing_algorithm: str = "none" # currently supporting none and trie
9090

9191

9292
@dataclass_json(undefined=Undefined.RAISE)

shortfin/python/shortfin_apps/llm/components/service.py

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -68,19 +68,19 @@ def __init__(
6868
page_pool = PagePool(
6969
devices=self.main_fiber.devices_dict.values(), config=page_pool_config
7070
)
71-
if model_params.paged_kv_cache.cache_type == "trie":
71+
if model_params.paged_kv_cache.prefix_sharing_algorithm == "trie":
7272
self.page_cache = TriePagedAttentionCache(
7373
page_pool=page_pool,
7474
tokens_per_page=model_params.paged_kv_cache.block_seq_stride,
7575
)
76-
elif model_params.paged_kv_cache.cache_type == "base":
76+
elif model_params.paged_kv_cache.prefix_sharing_algorithm == "none":
7777
self.page_cache = BasePagedAttentionCache(
7878
page_pool=page_pool,
7979
tokens_per_page=model_params.paged_kv_cache.block_seq_stride,
8080
)
8181
else:
8282
raise ValueError(
83-
f"Unknown model_params.paged_kv_cache.cache_type {model_params.paged_kv_cache.cache_type}. Currently only supporting 'trie' and 'base'."
83+
f"Unknown model_params.paged_kv_cache.prefix_sharing_algorithm {model_params.paged_kv_cache.prefix_sharing_algorithm}. Currently only supporting 'trie' and 'none'."
8484
)
8585

8686
self.program_isolation = PROG_ISOLATIONS[program_isolation]

0 commit comments

Comments
 (0)