Hoist hisparse and decode-offload setup out of init_cache_with_memory_pool (#25605)
This commit is contained in:
@@ -453,6 +453,29 @@ class Scheduler(
|
|||||||
# Init cache and memory pool
|
# Init cache and memory pool
|
||||||
self.init_cache_with_memory_pool()
|
self.init_cache_with_memory_pool()
|
||||||
|
|
||||||
|
if self.enable_hisparse:
|
||||||
|
# Coordinator was created inside ModelRunner.initialize() before CUDA graph capture
|
||||||
|
self.hisparse_coordinator = self.tp_worker.model_runner.hisparse_coordinator
|
||||||
|
self.hisparse_coordinator.set_decode_producer_stream(self.forward_stream)
|
||||||
|
|
||||||
|
if (
|
||||||
|
self.server_args.disaggregation_mode == "decode"
|
||||||
|
and self.server_args.disaggregation_decode_enable_offload_kvcache
|
||||||
|
):
|
||||||
|
self.decode_offload_manager = DecodeKVCacheOffloadManager(
|
||||||
|
req_to_token_pool=self.req_to_token_pool,
|
||||||
|
token_to_kv_pool_allocator=self.token_to_kv_pool_allocator,
|
||||||
|
tp_group=(
|
||||||
|
self.attn_tp_cpu_group
|
||||||
|
if self.server_args.enable_dp_attention
|
||||||
|
else self.tp_cpu_group
|
||||||
|
),
|
||||||
|
tree_cache=self.tree_cache,
|
||||||
|
server_args=self.server_args,
|
||||||
|
)
|
||||||
|
else:
|
||||||
|
self.decode_offload_manager = None
|
||||||
|
|
||||||
# Register draft KV pool (when spec + HiCache co-enabled).
|
# Register draft KV pool (when spec + HiCache co-enabled).
|
||||||
kv_cache_builder.maybe_register_hicache_draft(
|
kv_cache_builder.maybe_register_hicache_draft(
|
||||||
tree_cache=self.tree_cache,
|
tree_cache=self.tree_cache,
|
||||||
@@ -985,25 +1008,6 @@ class Scheduler(
|
|||||||
):
|
):
|
||||||
self.tree_cache = StreamingSession(self.tree_cache)
|
self.tree_cache = StreamingSession(self.tree_cache)
|
||||||
|
|
||||||
if self.enable_hisparse:
|
|
||||||
# Coordinator was created inside ModelRunner.initialize() before CUDA graph capture
|
|
||||||
self.hisparse_coordinator = self.tp_worker.model_runner.hisparse_coordinator
|
|
||||||
self.hisparse_coordinator.set_decode_producer_stream(self.forward_stream)
|
|
||||||
|
|
||||||
if (
|
|
||||||
server_args.disaggregation_mode == "decode"
|
|
||||||
and server_args.disaggregation_decode_enable_offload_kvcache
|
|
||||||
):
|
|
||||||
self.decode_offload_manager = DecodeKVCacheOffloadManager(
|
|
||||||
req_to_token_pool=self.req_to_token_pool,
|
|
||||||
token_to_kv_pool_allocator=self.token_to_kv_pool_allocator,
|
|
||||||
tp_group=params.tp_cache_group,
|
|
||||||
tree_cache=self.tree_cache,
|
|
||||||
server_args=self.server_args,
|
|
||||||
)
|
|
||||||
else:
|
|
||||||
self.decode_offload_manager = None
|
|
||||||
|
|
||||||
embedding_cache_size = envs.SGLANG_VLM_CACHE_SIZE_MB.get()
|
embedding_cache_size = envs.SGLANG_VLM_CACHE_SIZE_MB.get()
|
||||||
init_mm_embedding_cache(embedding_cache_size * 1024 * 1024)
|
init_mm_embedding_cache(embedding_cache_size * 1024 * 1024)
|
||||||
|
|
||||||
|
|||||||
Reference in New Issue
Block a user