[HiCache] MLA/DSA host dedup with rotating layer owners, chunked D2H backup issue and idle-sync skip (opt-in)

python/sglang/srt/arg_groups/field_order.py

@@ -520,5 +520,6 @@
     "enable_decoder_swa_bounded_replay",
     "enable_response_store",
     "disaggregation_decode_host_receive_threshold",
+    "enable_mla_hicache_host_dedup",
 )
 # fmt: on

python/sglang/srt/arg_groups/fields/memory.py

@@ -142,6 +142,14 @@ class Memory(msgspec.Struct):
             ],
         ),
     ] = "page_first"
+    enable_mla_hicache_host_dedup: A[
+        bool,
+        "Deduplicate the replicated MLA/DSA KV in the HiCache host tier across "
+        "attention-TP ranks: each rank backs up and loads back a round-robin share "
+        "of the layers, and every loaded layer is broadcast from its owner. Needs "
+        "--hicache-io-backend direct and --hicache-mem-layout page_first_direct, "
+        "without a storage backend.",
+    ] = False
     hicache_storage_backend: A[
         Optional[str],
         Arg(

python/sglang/srt/environ.py

@@ -769,9 +769,15 @@ class Envs:
     SGLANG_HICACHE_TMA_TRANSFER = EnvBool(True)
     # Base token count for each MLA/DSA dedup broadcast chunk.
     SGLANG_MLA_DEDUP_CHUNK_TOKENS = EnvInt(2048)
+    # Direct-I/O D2H backup: issue each all-layer copy batch in batches of this
+    # many pages, so each batch's DMA runs while the next is issued (0 = one batch).
+    SGLANG_HICACHE_D2H_ISSUE_CHUNK_PAGES = EnvInt(0)
     SGLANG_HICACHE_HF3FS_CONFIG_PATH = EnvStr(None)
     SGLANG_HICACHE_DECODE_OFFLOAD_STRIDE = EnvInt(None)
     SGLANG_HICACHE_SKIP_HOST_DUPLICATE_RECLAIM = EnvBool(False)
+    # Skip the per-step HiCache ready-count CPU all-reduce while no D<->H
+    # transfer ack is outstanding (aggregated, L2-only, single PP stage).
+    SGLANG_ENABLE_HICACHE_IDLE_SYNC_SKIP = EnvBool(False)
     SGLANG_HICACHE_FILE_BACKEND_STORAGE_DIR = EnvStr(None)
     # File-backend LRU eviction (opt-in; sizes accept SI/IEC suffixes, "0" disables).
     SGLANG_HICACHE_FILE_BACKEND_MAX_SIZE = EnvStr(None)

python/sglang/srt/managers/cache_controller.py

@@ -35,6 +35,7 @@
 if TYPE_CHECKING:
     from sglang.srt.mem_cache.allocator import BaseTokenToKVPoolAllocator
     from sglang.srt.mem_cache.memory_pool_host import LogicalHostPool
+    from sglang.srt.mem_cache.mla_host_dedup import MLAHostDedupContext
     from sglang.srt.mem_cache.pool_host import HostKVCache
 
 from sglang.srt.layers.dp_attention import (
@@ -81,13 +82,19 @@ def finish_event(self):
 
 
 class LayerDoneCounter:
-    def __init__(self, num_layers: int):
+    def __init__(
+        self,
+        num_layers: int,
+        on_layer_ready: Optional[Callable[[int, int], None]] = None,
+    ):
         self.num_layers = num_layers
         # extra producer and consumer counters for overlap mode
         self.num_counters = 3
         self.events = [LayerLoadingEvent(num_layers) for _ in range(self.num_counters)]
         self.producer_index = -1
         self.consumer_index = -1
+        # Called on the forward stream after each wait, with (consumer, layer).
+        self.on_layer_ready = on_layer_ready
 
     def update_producer(self):
         self.producer_index = (self.producer_index + 1) % self.num_counters
@@ -103,6 +110,8 @@ def wait_until(self, threshold: int):
         if self.consumer_index < 0:
             return
         self.events[self.consumer_index].wait(threshold)
+        if self.on_layer_ready is not None:
+            self.on_layer_ready(self.consumer_index, threshold)
 
     def reset(self):
         self.producer_index = -1
@@ -311,6 +320,7 @@ def __init__(
         storage_backend_extra_config: Optional[dict] = None,
         enable_storage_metrics: bool = False,
         host_memory_mode: str = "cache",
+        mla_dedup_context: Optional[MLAHostDedupContext] = None,
     ):
         self.tp_group = tp_group
         self.host_memory_mode = host_memory_mode
@@ -356,7 +366,17 @@ def __init__(
 
         self.device = self.mem_pool_device.device
         self.transfer_layer_id_max = self.mem_pool_device.layer_num
-        self.layer_done_counter = LayerDoneCounter(self.transfer_layer_id_max)
+        # --enable-mla-hicache-host-dedup: each rank loads its own layers, and the
+        # forward broadcasts every loaded layer from its owner as it waits for it.
+        self.mla_dedup = mla_dedup_context
+        self.layer_done_counter = LayerDoneCounter(
+            self.transfer_layer_id_max,
+            on_layer_ready=(
+                None
+                if mla_dedup_context is None
+                else mla_dedup_context.broadcast_ready_layers
+            ),
+        )
         self.mem_pool_device.register_layer_transfer_counter(self.layer_done_counter)
 
         if write_policy not in [
@@ -543,6 +563,11 @@ def attach_storage_backend(
         """
         if self.enable_storage:
             raise RuntimeError("Storage backend already attached.")
+        if self.mla_dedup is not None:
+            raise RuntimeError(
+                "Cannot attach a storage backend with MLA host dedup: each rank's "
+                "host pool holds only its own layers."
+            )
 
         # Defensive: a previous partial detach may have flipped `enable_storage` but
         # left background threads alive. Attaching on top of them is unsafe.
@@ -989,6 +1014,10 @@ def start_loading(self) -> int:
             on_layer_done=producer_event.complete,
             transfer_layer_id_max=self.transfer_layer_id_max,
         )
+        if self.mla_dedup is not None:
+            self.mla_dedup.start_load(
+                producer_id, op.device_indices, completion.finish_event
+            )
 
         self.mem_pool_device_allocator.set_hicache_transfer_done_event(
             (id(self), "load"), completion.finish_event

python/sglang/srt/mem_cache/hybrid_cache/hybrid_cache_controller.py

@@ -45,6 +45,7 @@
 
 if TYPE_CHECKING:
     from sglang.srt.mem_cache.allocator import BaseTokenToKVPoolAllocator
+    from sglang.srt.mem_cache.mla_host_dedup import MLAHostDedupContext
 
 from sglang.srt.mem_cache.utils import get_storage_hash_str
 from sglang.srt.utils import broadcast_pyobj
@@ -209,6 +210,7 @@ def __init__(
         transfer_layer_id_max: Optional[int] = None,
         enable_storage_metrics: bool = False,
         host_memory_mode: str = "cache",
+        mla_dedup_context: Optional[MLAHostDedupContext] = None,
     ):
         startup_storage_backend = storage_backend
         self.extra_host_mem_release_queues: dict[PoolName, Queue[torch.Tensor]] = {}
@@ -235,6 +237,7 @@ def __init__(
             storage_backend_extra_config=storage_backend_extra_config,
             enable_storage_metrics=enable_storage_metrics,
             host_memory_mode=host_memory_mode,
+            mla_dedup_context=mla_dedup_context,
         )
         # Hybrid transfer IDs span every component pool, including holes for
         # uncached layers that the anchor pool alone cannot describe.
@@ -243,7 +246,10 @@ def __init__(
             and transfer_layer_id_max != self.transfer_layer_id_max
         ):
             self.transfer_layer_id_max = transfer_layer_id_max
-            self.layer_done_counter = LayerDoneCounter(self.transfer_layer_id_max)
+            self.layer_done_counter = LayerDoneCounter(
+                self.transfer_layer_id_max,
+                on_layer_ready=self.layer_done_counter.on_layer_ready,
+            )
 
         self.storage_host_pool = mem_pool_host.anchor_entry.host_pool
         if startup_storage_backend is not None:

python/sglang/srt/mem_cache/hybrid_cache/hybrid_pool_assembler.py

@@ -24,6 +24,10 @@
     DeepSeekV4StateHostPool,
     LogicalHostPool,
 )
+from sglang.srt.mem_cache.mla_host_dedup import (
+    MLAHostDedupLayerOwners,
+    maybe_create_hicache_mla_host_dedup,
+)
 from sglang.srt.mem_cache.pool_host import HostPoolGroup, PoolEntry
 from sglang.srt.mem_cache.pool_host.common import get_allocator_type
 from sglang.srt.mem_cache.pool_host.host_pool_decl import HostPoolDecl
@@ -142,6 +146,7 @@ def build_kv_host_pool(
     host_size: Optional[float] = None,
     mtp_draft_device_pools: tuple[Any, ...] = (),
     pool_label: str = "kv",
+    dedup_owners: Optional[MLAHostDedupLayerOwners] = None,
 ):
     if use_mla is None:
         use_mla = is_mla_pool(kv_pool)
@@ -160,6 +165,8 @@ def build_kv_host_pool(
         kwargs["override_kv_cache_dim"] = override_kv_cache_dim
     if mtp_draft_device_pools:
         kwargs["mtp_draft_device_pools"] = mtp_draft_device_pools
+    if dedup_owners is not None:
+        kwargs["dedup_owners"] = dedup_owners
     parallel = get_parallel()
     if parallel.dcp_enabled:
         assert use_mla, (
@@ -237,6 +244,7 @@ def build_kv_only_group(
     override_kv_cache_dim: Optional[int] = None,
     host_size: Optional[float] = None,
     mtp_draft_device_pools: tuple[Any, ...] = (),
+    dedup_owners: Optional[MLAHostDedupLayerOwners] = None,
 ) -> HostPoolGroup:
     """Anchor-only host pool group for a flat MHA/MLA device pool."""
     transfer_layer_id_max = len(full_layer_mapping)
@@ -247,6 +255,7 @@ def build_kv_only_group(
         override_kv_cache_dim=override_kv_cache_dim,
         host_size=host_size,
         mtp_draft_device_pools=mtp_draft_device_pools,
+        dedup_owners=dedup_owners,
     )
     if mtp_draft_device_pools:
         full_layer_mapping = with_packed_draft_layer_mapping(
@@ -420,13 +429,17 @@ def build_kv_only_stack(
     enable_storage_metrics: bool = False,
 ) -> tuple[HostPoolGroup, HybridCacheController]:
     transfer_layer_id_max = len(full_layer_mapping)
+    mla_dedup = maybe_create_hicache_mla_host_dedup(
+        kv_pool, params, enabled=get_memory().enable_mla_hicache_host_dedup
+    )
     host_pool_group = build_kv_only_group(
         page_size=params.page_size,
         kv_pool=kv_pool,
         full_layer_mapping=full_layer_mapping,
         use_mla=use_mla,
         override_kv_cache_dim=override_kv_cache_dim,
         mtp_draft_device_pools=params.mtp_draft_device_pools,
+        dedup_owners=None if mla_dedup is None else mla_dedup.owners,
     )
     cache_controller = HybridCacheController(
         params.token_to_kv_pool_allocator,
@@ -446,6 +459,7 @@ def build_kv_only_stack(
         transfer_layer_id_max=transfer_layer_id_max,
         enable_storage_metrics=enable_storage_metrics,
         host_memory_mode=get_memory().hicache_host_memory_mode,
+        mla_dedup_context=mla_dedup,
     )
     return host_pool_group, cache_controller
 
@@ -1387,6 +1401,7 @@ def build_hybrid_mamba_swa_stack(
 def build_host_pool_group(
     *,
     config: HostPoolGroupConfig,
+    dedup_owners: Optional[MLAHostDedupLayerOwners] = None,
 ) -> HostPoolGroup:
     """Allocate host pools and transfer entries from prepared configs."""
     root = config.pools[0]
@@ -1395,6 +1410,7 @@ def build_host_pool_group(
             kv_pool=root.decl.device_pool,
             page_size=config.transfer_page_size,
             mtp_draft_device_pools=root.packed_draft_device_pools,
+            dedup_owners=dedup_owners,
         )
     }
     entries = []
@@ -2158,7 +2174,13 @@ def build(
             transfer_page_size=params.page_size,
             packed_draft_device_pools=params.mtp_draft_device_pools,
         )
-        host_pool_group = build_host_pool_group(config=config)
+        mla_dedup = maybe_create_hicache_mla_host_dedup(
+            kvcache, params, enabled=get_memory().enable_mla_hicache_host_dedup
+        )
+        host_pool_group = build_host_pool_group(
+            config=config,
+            dedup_owners=None if mla_dedup is None else mla_dedup.owners,
+        )
         cache_controller = HybridCacheController(
             params.token_to_kv_pool_allocator,
             host_pool_group,
@@ -2178,6 +2200,7 @@ def build(
             transfer_layer_id_max=kvcache.layer_num,
             enable_storage_metrics=enable_storage_metrics,
             host_memory_mode=get_memory().hicache_host_memory_mode,
+            mla_dedup_context=mla_dedup,
         )
         return StackBuildResult(
             host_pool_group=host_pool_group,