Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
Original file line number Diff line number Diff line change
Expand Up @@ -1007,9 +1007,11 @@ def build_full_draft_pools(
controller = tree_cache.cache_controller
host_pool_group = controller.mem_pool_host

# Note(kpham-sgl): DCP x DSpark draft KV is replicated and spans the virtual
# loc space, so match the target host's logical_size instead of physical size.
draft_host_pool = _build_mha_mla_host_pool(
pool=pool,
host_to_device_ratio=host_pool_group.size / pool.size,
host_to_device_ratio=host_pool_group.logical_size / pool.size,

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

can add a comment here to explain why use logical_size

page_size=controller.page_size,
layout=server_args.hicache_mem_layout,
allocator_type=_get_allocator_type(server_args),
Expand Down
2 changes: 1 addition & 1 deletion python/sglang/srt/mem_cache/kv_cache_builder.py
Original file line number Diff line number Diff line change
Expand Up @@ -117,7 +117,7 @@ def _register_legacy_hicache_draft(
# so that host indices stay 1-to-1 between target and draft KV caches.
primary_host_pool = tree_cache.cache_controller.mem_pool_host
host_pool_kwargs = dict(
host_to_device_ratio=primary_host_pool.size / pool.size,
host_to_device_ratio=primary_host_pool.logical_size / pool.size,
host_size=0,
page_size=page_size,
layout=server_args.hicache_mem_layout,
Expand Down
8 changes: 4 additions & 4 deletions python/sglang/srt/server_args.py
Original file line number Diff line number Diff line change
Expand Up @@ -7386,11 +7386,11 @@ def _resolve_hicache_dcp_compatibility(self):
"backup and the storage keys must become dcp_rank-aware "
"first. Run HiCache+DCP with L1/L2 only."
)
if self.speculative_algorithm is not None:
if self.speculative_algorithm not in (None, "DSPARK"):
raise NotImplementedError(
"HiCache with --dcp-size > 1 does not support speculative "
"decoding yet (the draft-model host pool has no DCP index "
"translation)."
"HiCache with --dcp-size > 1 only supports DSPARK speculative "
"decoding; other draft-model host pools have no DCP index "
"translation."
)
if self.enable_lmcache:
raise NotImplementedError(
Expand Down
34 changes: 34 additions & 0 deletions test/registered/unit/mem_cache/test_hybrid_pool_assembler.py
Original file line number Diff line number Diff line change
Expand Up @@ -2,6 +2,7 @@

import unittest
from types import SimpleNamespace
from unittest.mock import patch

from sglang.srt.mem_cache.hybrid_cache.hybrid_pool_assembler import (
_split_hicache_size,
Expand Down Expand Up @@ -54,6 +55,39 @@ def test_full_builder_unwraps_empty_hybrid_linear_pool(self):
self.assertEqual(specs, [])
self.assertEqual(entries, [])

def test_full_builder_sizes_sidecar_for_anchor_logical_space(self):
draft_kv_pool = SimpleNamespace(layer_num=1, size=800)
draft_host_pool = SimpleNamespace(layer_num=1)
tree_cache = SimpleNamespace(
cache_controller=SimpleNamespace(
mem_pool_host=SimpleNamespace(size=100, logical_size=800),
page_size=512,
)
)
server_args = SimpleNamespace(hicache_mem_layout="page_first")

with (
patch(
"sglang.srt.mem_cache.hybrid_cache.hybrid_pool_assembler."
"_build_mha_mla_host_pool",
return_value=draft_host_pool,
) as build_host_pool,
patch(
"sglang.srt.mem_cache.hybrid_cache.hybrid_pool_assembler."
"_get_allocator_type",
return_value="default",
),
):
specs, entries = build_full_draft_pools(
draft_kv_pool=draft_kv_pool,
tree_cache=tree_cache,
server_args=server_args,
)

self.assertEqual(build_host_pool.call_args.kwargs["host_to_device_ratio"], 1.0)
self.assertEqual(len(specs), 1)
self.assertIs(entries[0].host_pool, draft_host_pool)


if __name__ == "__main__":
unittest.main()
Loading