[HiCache] Add MLA host-dedup primitives (#36800)

Co-authored-by: Zhangheng <hzh0425@apache.org>

python/sglang/srt/environ.py

@@ -720,6 +720,8 @@ class Envs:
     # ===================================================================
     # Per-call cudaHostRegister limit in GB.
     SGLANG_HICACHE_HOST_REGISTER_CHUNK_GB = EnvInt(256)
+    # Base token count for each MLA/DSA dedup broadcast chunk.
+    SGLANG_MLA_DEDUP_CHUNK_TOKENS = EnvInt(2048)
     SGLANG_HICACHE_HF3FS_CONFIG_PATH = EnvStr(None)
     SGLANG_HICACHE_DECODE_OFFLOAD_STRIDE = EnvInt(None)
     SGLANG_HICACHE_FILE_BACKEND_STORAGE_DIR = EnvStr(None)

python/sglang/srt/mem_cache/mla_host_dedup.py

@@ -0,0 +1,305 @@
+"""Deduplicate MLA/DSA host cache across attention-TP ranks."""
+
+from __future__ import annotations
+
+import logging
+import math
+from dataclasses import dataclass
+from typing import List, Optional
+
+import torch
+
+from sglang.srt.environ import envs
+from sglang.srt.layers.dp_attention import is_dp_attention_enabled
+from sglang.srt.mem_cache.memory_pool import (
+    DSATokenToKVPool,
+    MLATokenToKVPool,
+    MLATokenToKVPoolFP4,
+)
+from sglang.srt.runtime_context import get_parallel
+from sglang.srt.utils import is_cuda
+
+logger = logging.getLogger(__name__)
+
+
+# These backends tolerate buffer-less host pools on non-source ranks.
+_DEDUP_COMPATIBLE_STORAGE = frozenset({None, "", "file"})
+
+
+def storage_supports_host_dedup(storage_backend: Optional[str]) -> bool:
+    """Whether MLA/DSA host-memory dedup can engage with this storage backend."""
+    return storage_backend in _DEDUP_COMPATIBLE_STORAGE
+
+
+def mla_dedup_rank_and_size() -> tuple[int, int]:
+    """Attn-TP rank/size when DP attention is enabled, model-TP otherwise."""
+    parallel = get_parallel()
+    if is_dp_attention_enabled():
+        return parallel.attn_tp_rank, parallel.attn_tp_size
+    return parallel.tp_rank, parallel.tp_size
+
+
+def mla_host_dedup_eligible(kv_cache, storage_backend: Optional[str]) -> bool:
+    """Rank-independent gate. CUDA only; FP4 excluded (its per-rank scale
+    buffer is not covered by the broadcast)."""
+    return (
+        isinstance(kv_cache, MLATokenToKVPool)
+        and not isinstance(kv_cache, MLATokenToKVPoolFP4)
+        and is_cuda()
+        and storage_supports_host_dedup(storage_backend)
+    )
+
+
+class MLAHostDedupBroadcaster:
+    """Layerwise MLA/DSA broadcast over a dedicated NCCL group."""
+
+    def __init__(
+        self,
+        device_pool: MLATokenToKVPool,
+        group: torch.distributed.ProcessGroup,
+        src_global_rank: int,
+    ):
+        self.device_pool = device_pool
+        self.group = group
+        self.src_global_rank = src_global_rank
+        self.is_src = mla_dedup_rank_and_size()[0] == 0
+        self.layer_num = device_pool.layer_num
+        self.device = device_pool.device
+        self.chunk_tokens = envs.SGLANG_MLA_DEDUP_CHUNK_TOKENS.get()
+        if self.chunk_tokens <= 0:
+            raise ValueError(
+                "SGLANG_MLA_DEDUP_CHUNK_TOKENS must be positive, "
+                f"got {self.chunk_tokens}."
+            )
+        self.kv_staging = torch.empty(
+            self.layer_num * self.chunk_tokens * device_pool.kv_cache_dim,
+            dtype=device_pool.kv_buffer[0].dtype,
+            device=self.device,
+        )
+        self.idx_bufs = None
+        self.idx_elem = None
+        self.idx_staging = None
+        if isinstance(device_pool, DSATokenToKVPool):
+            self.idx_bufs = device_pool.index_k_with_scale_buffer
+            self.idx_elem = math.prod(self.idx_bufs[0].shape[1:]) or 1
+            self.idx_staging = torch.empty(
+                self.layer_num * self.chunk_tokens * self.idx_elem,
+                dtype=self.idx_bufs[0].dtype,
+                device=self.device,
+            )
+        logger.info(
+            "MLA host-dedup broadcast chunk configured: base_tokens=%d, "
+            "effective_layer_tokens=%d",
+            self.chunk_tokens,
+            self.layer_num * self.chunk_tokens,
+        )
+
+    @classmethod
+    def build(
+        cls,
+        device_pool,
+        tp_group: torch.distributed.ProcessGroup,
+        attn_tp_group: Optional[torch.distributed.ProcessGroup],
+    ) -> MLAHostDedupBroadcaster:
+        """Build and initialize the NCCL group before host-pool allocation."""
+        from sglang.srt.distributed.parallel_state import create_custom_parallel_group
+
+        base_group = tp_group
+        if is_dp_attention_enabled() and attn_tp_group is not None:
+            base_group = attn_tp_group
+        group_ranks = torch.distributed.get_process_group_ranks(base_group)
+        group = create_custom_parallel_group(
+            group_ranks=list(group_ranks), backend="nccl"
+        )
+        broadcaster = cls(device_pool, group, src_global_rank=group_ranks[0])
+        broadcaster._warmup_group()
+        return broadcaster
+
+    def _warmup_group(self) -> None:
+        """Initialize the NCCL communicator before serving."""
+        warmup = self.kv_staging[:1]
+        if self.is_src:
+            warmup.zero_()
+        torch.distributed.broadcast(warmup, src=self.src_global_rank, group=self.group)
+        torch.cuda.synchronize(self.device)
+        logger.info("MLA host-dedup NCCL broadcast group warmup completed")
+
+    def prepare_broadcast(
+        self, device_indices: torch.Tensor, load_stream
+    ) -> tuple[torch.Tensor, Optional[torch.Tensor]]:
+        """Prepare reusable KV/indexer indices for one layerwise load."""
+        indices = device_indices
+        if not indices.is_cuda:
+            indices = indices.to(self.device)
+        if indices.is_cuda:
+            indices.record_stream(load_stream)
+
+        page_idx = None
+        if self.idx_bufs is not None:
+            page_size = self.device_pool.page_size
+            if page_size > 1:
+                if indices.numel() % page_size != 0:
+                    raise ValueError(
+                        "DSA dedup broadcast expects page-aligned device indices: "
+                        f"got {indices.numel()} indices for page_size={page_size}."
+                    )
+                # Preserve logical page order across rank-local allocations.
+                page_idx = indices[::page_size] // page_size
+            else:
+                page_idx = indices
+            if page_idx.is_cuda:
+                page_idx.record_stream(load_stream)
+        return indices, page_idx
+
+    def broadcast_loaded_layer(
+        self,
+        layer_id: int,
+        prepared: tuple[torch.Tensor, Optional[torch.Tensor]],
+    ) -> None:
+        """Broadcast one loaded KV layer and its optional DSA indexer layer."""
+        indices, page_idx = prepared
+        self._bcast_layer(
+            self.device_pool.kv_buffer,
+            self.kv_staging,
+            indices,
+            self.device_pool.kv_cache_dim,
+            layer_id,
+        )
+        if self.idx_bufs is not None:
+            assert page_idx is not None
+            self._bcast_layer(
+                self.idx_bufs,
+                self.idx_staging,
+                page_idx,
+                self.idx_elem,
+                layer_id,
+            )
+
+    def _bcast_layer(
+        self,
+        buf_list,
+        staging,
+        target,
+        elem,
+        layer_id: int,
+    ) -> None:
+        """Broadcast one layer in chunks using the shared staging buffer."""
+        n = target.shape[0]
+        rows_per_chunk = staging.numel() // elem
+        assert rows_per_chunk > 0
+        layer_buf = buf_list[layer_id]
+        row_shape = layer_buf.shape[1:]
+
+        for start in range(0, n, rows_per_chunk):
+            cur = min(rows_per_chunk, n - start)
+            idx = target[start : start + cur]
+            chunk = staging[: cur * elem]
+            chunk_rows = chunk.view(cur, *row_shape)
+            if self.is_src:
+                torch.index_select(layer_buf, 0, idx, out=chunk_rows)
+            torch.distributed.broadcast(
+                chunk, src=self.src_global_rank, group=self.group
+            )
+            if not self.is_src:
+                layer_buf.index_copy_(0, idx, chunk_rows)
+
+    def destroy(self) -> None:
+        if self.group is None:
+            return
+        try:
+            torch.distributed.destroy_process_group(self.group)
+        except Exception:
+            pass
+        self.group = None
+
+
+@dataclass
+class MLAHostDedupContext:
+    """All state owned by the optional MLA host-dedup path."""
+
+    broadcaster: MLAHostDedupBroadcaster
+    prefetch_hits_sync_groups: Optional[List[torch.distributed.ProcessGroup]]
+    prefetch_completion_sync_groups: Optional[List[torch.distributed.ProcessGroup]]
+    producer_stream: Optional[object] = None
+    last_write_finish_event: Optional[object] = None
+
+    @property
+    def is_src(self) -> bool:
+        return self.broadcaster.is_src
+
+    @property
+    def is_dummy_rank(self) -> bool:
+        return not self.is_src
+
+    def destroy(self) -> None:
+        self.broadcaster.destroy()
+        groups = (self.prefetch_hits_sync_groups or []) + (
+            self.prefetch_completion_sync_groups or []
+        )
+        for group in groups:
+            try:
+                torch.distributed.destroy_process_group(group)
+            except Exception:
+                pass
+        self.prefetch_hits_sync_groups = None
+        self.prefetch_completion_sync_groups = None
+
+
+def maybe_create_mla_host_dedup_context(
+    kv_cache,
+    tp_group: torch.distributed.ProcessGroup,
+    attn_cp_group: Optional[torch.distributed.ProcessGroup],
+    attn_tp_group: Optional[torch.distributed.ProcessGroup],
+    storage_backend: Optional[str],
+    enabled: bool = False,
+) -> Optional[MLAHostDedupContext]:
+    """Create dedup state before host allocation, or preserve the original path."""
+    if not enabled:
+        return None
+    if not mla_host_dedup_eligible(kv_cache, storage_backend):
+        return None
+    if mla_dedup_rank_and_size()[1] <= 1:
+        return None
+
+    broadcaster = MLAHostDedupBroadcaster.build(kv_cache, tp_group, attn_tp_group)
+    prefetch_hits_sync_groups = None
+    prefetch_completion_sync_groups = None
+    if storage_backend is not None:
+        prefetch_hits_sync_groups = _prebuild_prefetch_sync_groups(
+            tp_group, attn_cp_group, attn_tp_group
+        )
+        prefetch_completion_sync_groups = _prebuild_prefetch_sync_groups(
+            tp_group, attn_cp_group, attn_tp_group
+        )
+    return MLAHostDedupContext(
+        broadcaster,
+        prefetch_hits_sync_groups,
+        prefetch_completion_sync_groups,
+    )
+
+
+def _prebuild_prefetch_sync_groups(
+    tp_group: torch.distributed.ProcessGroup,
+    attn_cp_group: Optional[torch.distributed.ProcessGroup],
+    attn_tp_group: Optional[torch.distributed.ProcessGroup],
+) -> List[torch.distributed.ProcessGroup]:
+    """Prebuild one set of HiCache storage synchronization groups."""
+    from sglang.srt.distributed.parallel_state import create_custom_parallel_group
+
+    groups: List[torch.distributed.ProcessGroup] = []
+    seen_rank_sets = set()
+    if attn_cp_group is not None or attn_tp_group is not None:
+        base_groups = [attn_cp_group, attn_tp_group]
+    else:
+        base_groups = [tp_group]
+    for group in base_groups:
+        if group is None or torch.distributed.get_world_size(group=group) == 1:
+            continue
+        ranks = tuple(torch.distributed.get_process_group_ranks(group))
+        if ranks in seen_rank_sets:
+            continue
+        seen_rank_sets.add(ranks)
+        groups.append(
+            create_custom_parallel_group(group_ranks=list(ranks), backend="gloo")
+        )
+    return groups

python/sglang/srt/mem_cache/pool_host/dsa.py

@@ -59,7 +59,9 @@ class DSAIndexerPoolHost(HostKVCache):
         pin_memory: bool = True,
         device: str = "cpu",
         allocator_type: str = "default",
+        is_dummy: bool = False,
     ):
+        self._is_dummy = is_dummy
         self.device_pool = device_pool
         self.page_size = anchor_host.page_size
         self.layout = layout
@@ -92,6 +94,20 @@ class DSAIndexerPoolHost(HostKVCache):
             self.indexer_size_per_token * self.layer_num * self.indexer_dtype.itemsize
         )
 
+        self.can_use_jit = False
+        self.can_use_write_back_jit = False
+        if is_dummy:
+            self.index_k_with_scale_buffer = None
+            self.index_k_device_ptrs = None
+            logger.info(
+                "DSAIndexerPoolHost dummy mode: allocator-only, size=%d tokens, "
+                "skipping indexer buffer allocation",
+                self.size,
+            )
+            self.lock = threading.RLock()
+            self.clear()
+            return
+
         buf_elem_size = self.page_num * self.layer_num * self.indexer_page_stride_size
         requested_bytes = buf_elem_size * self.indexer_dtype.itemsize
         available_bytes = host_memory_budget_bytes()
@@ -120,8 +136,6 @@ class DSAIndexerPoolHost(HostKVCache):
                 layout,
             )
         self.init_kv_buffer()
-        self.can_use_jit = False
-        self.can_use_write_back_jit = False
         self._init_write_back_staging_buffers()
         self.lock = threading.RLock()
         self.clear()
@@ -230,6 +244,9 @@ class DSAIndexerPoolHost(HostKVCache):
     ):
         if not is_draft and not self._is_device_layer_owned(device_pool, layer_id):
             return
+        assert not getattr(self, "_is_dummy", False), (
+            "load on a dummy (non-src DSA) host pool"
+        )
         # MTP draft layers do not participate in CP layer sharding.
         host_layer_id = layer_id if is_draft else self._host_layer_index(layer_id)
         device_layer_id = 0 if is_draft else layer_id
@@ -292,6 +309,9 @@ class DSAIndexerPoolHost(HostKVCache):
         *,
         is_draft: bool = False,
     ):
+        assert not getattr(self, "_is_dummy", False), (
+            "backup on a dummy (non-src DSA) host pool"
+        )
         # MTP draft layers do not participate in CP layer sharding.
         host_layer_id = layer_id if is_draft else self._host_layer_index(layer_id)
         device_layer_id = 0 if is_draft else layer_id
@@ -336,6 +356,9 @@ class DSAIndexerPoolHost(HostKVCache):
     def backup_from_device_all_layer(
         self, device_pool, host_indices, device_indices, io_backend
     ):
+        assert not getattr(self, "_is_dummy", False), (
+            "backup on a dummy (non-src DSA) host pool"
+        )
         if self._is_device_layer_sharded(device_pool):
             for layer_id in self._owned_device_layer_ids(device_pool):
                 self._backup_from_device_per_layer(

python/sglang/srt/mem_cache/pool_host/mla.py

@@ -1,6 +1,7 @@
 from __future__ import annotations
 
 import logging
+import threading
 from typing import Optional, Sequence
 
 import torch
@@ -22,8 +23,12 @@ from sglang.srt.mem_cache.memory_pool import MLATokenToKVPool
 from sglang.srt.mem_cache.pool_host.base import (
     _WRITE_BACK_STAGING_PAGE_CHUNK,
     HostKVCache,
+    sync_fixed_hicache_size,
+)
+from sglang.srt.mem_cache.pool_host.common import (
+    ALLOC_MEMORY_FUNCS,
+    get_allocator_from_storage,
 )
-from sglang.srt.mem_cache.pool_host.common import ALLOC_MEMORY_FUNCS
 from sglang.srt.mem_cache.pool_host.hisparse import HiSparseHostPoolMixin
 from sglang.srt.utils import is_cuda, is_hip, is_mps, is_npu, is_xpu
 
@@ -68,9 +73,28 @@ class MLATokenToKVPoolHost(HiSparseHostPoolMixin, HostKVCache):
         dcp_rank: int = 0,
         *,
         pool_label: str = "kv",
+        is_dummy: bool = False,
     ):
         self.override_kv_cache_dim = override_kv_cache_dim
         self.mtp_draft_device_pools = tuple(mtp_draft_device_pools)
+        self._is_dummy = is_dummy
+
+        if is_dummy:
+            self._init_dummy(
+                device_pool,
+                host_to_device_ratio,
+                host_size,
+                page_size,
+                layout,
+                pin_memory,
+                device,
+                allocator_type,
+                dcp_size,
+                dcp_rank,
+                pool_label,
+            )
+            return
+
         super().__init__(
             device_pool,
             host_to_device_ratio,
@@ -114,9 +138,73 @@ class MLATokenToKVPoolHost(HiSparseHostPoolMixin, HostKVCache):
             ]
         self._init_write_back_staging_buffers()
 
+    def _init_dummy(
+        self,
+        device_pool: MLATokenToKVPool,
+        host_to_device_ratio: float,
+        host_size: int,
+        page_size: int,
+        layout: str,
+        pin_memory: bool,
+        device: str,
+        allocator_type: str,
+        dcp_size: int,
+        dcp_rank: int,
+        pool_label: str,
+    ) -> None:
+        self.device_pool = device_pool
+        self.pool_label = pool_label
+        self.dcp_size = dcp_size
+        self.dcp_rank = dcp_rank
+        assert page_size % dcp_size == 0, (
+            f"HiCache host pool page_size ({page_size}) must be a multiple of "
+            f"dcp_size ({dcp_size})."
+        )
+        self.page_size = page_size // dcp_size
+        self.layout = layout
+        self.pin_memory = pin_memory
+        self.device = device
+        self.allocator = get_allocator_from_storage(allocator_type)
+
+        self.dtype = device_pool.store_dtype
+        self.size_per_token = self.get_size_per_token()
+        if host_size > 0:
+            self.size = sync_fixed_hicache_size(
+                int(host_size * 1e9 // self.size_per_token), host_size
+            )
+        else:
+            self.size = int(device_pool.size * host_to_device_ratio)
+        self.page_num = self.size // self.page_size + 1
+        self.size = self.page_num * self.page_size
+        self.start_layer = device_pool.start_layer
+        self.end_layer = device_pool.end_layer
+
+        self.token_stride_size = self.kv_cache_dim * self.dtype.itemsize
+        self.layout_dim = self.token_stride_size * self.layer_num
+        self.can_use_jit = False
+        self.can_use_write_back_jit = False
+        self.staging_page_capacity = 0
+        self.staging_token_capacity = 0
+        self.staging_buffer = None
+        self.kv_buffer = None
+        self.data_refs = None
+        self.data_ptrs = None
+
+        logger.info(
+            "MLATokenToKVPoolHost dummy mode: allocator-only, size=%d tokens, "
+            "saving %.2f GB host memory",
+            self.size,
+            self.size * self.size_per_token / 1e9,
+        )
+
+        self.lock = threading.RLock()
+        self.clear()
+
     def get_contiguous_buf_infos(self):
         """Return (data_ptrs, data_lens, item_lens) in the same format as device pool,
         for registering host memory with the disaggregation transfer engine."""
+        if self._is_dummy:
+            return [], [], []
         data_ptrs = [int(self.data_ptrs[i].item()) for i in range(self.layer_num)]
         data_lens = [self.kv_buffer[i].nbytes for i in range(self.layer_num)]
         item_lens = [self.token_stride_size * self.page_size] * self.layer_num
@@ -257,6 +345,9 @@ class MLATokenToKVPoolHost(HiSparseHostPoolMixin, HostKVCache):
     ):
         if not is_draft and not self._is_device_layer_owned(device_pool, layer_id):
             return
+        assert not getattr(self, "_is_dummy", False), (
+            "load on a dummy (non-src MLA) host pool"
+        )
         host_indices = self.maybe_dcp_kernel_indices(host_indices)
         device_indices = self.maybe_dcp_kernel_indices(device_indices)
         # MTP draft layers do not participate in CP layer sharding.
@@ -353,6 +444,9 @@ class MLATokenToKVPoolHost(HiSparseHostPoolMixin, HostKVCache):
         *,
         is_draft: bool = False,
     ):
+        assert not getattr(self, "_is_dummy", False), (
+            "backup on a dummy (non-src MLA) host pool"
+        )
         # Indices arrive already translated by backup_from_device_all_layer.
         # MTP draft layers do not participate in CP layer sharding.
         host_layer_id = layer_id if is_draft else self._host_layer_index(layer_id)
@@ -421,6 +515,9 @@ class MLATokenToKVPoolHost(HiSparseHostPoolMixin, HostKVCache):
     def backup_from_device_all_layer(
         self, device_pool, host_indices, device_indices, io_backend
     ):
+        assert not getattr(self, "_is_dummy", False), (
+            "backup on a dummy (non-src MLA) host pool"
+        )
         host_indices = self.maybe_dcp_kernel_indices(host_indices)
         device_indices = self.maybe_dcp_kernel_indices(device_indices)
         if self._is_device_layer_sharded(device_pool):

test/registered/unit/mem_cache/test_mla_host_dedup_primitives.py

@@ -0,0 +1,258 @@
+import unittest
+from types import SimpleNamespace
+from unittest import mock
+
+import torch
+
+from sglang.srt.environ import envs
+from sglang.srt.mem_cache.mla_host_dedup import (
+    MLAHostDedupBroadcaster,
+    MLAHostDedupContext,
+    maybe_create_mla_host_dedup_context,
+)
+from sglang.srt.mem_cache.pool_host.dsa import DSAIndexerPoolHost
+from sglang.srt.mem_cache.pool_host.mla import MLATokenToKVPoolHost
+from sglang.test.ci.ci_register import register_cpu_ci
+
+register_cpu_ci(est_time=1, suite="base-a-test-cpu")
+
+
+def _device_pool_stub(*, layer_num: int, **fields) -> SimpleNamespace:
+    return SimpleNamespace(
+        layer_num=layer_num,
+        layer_shard_enabled=False,
+        **fields,
+    )
+
+
+class _FakeStream:
+    pass
+
+
+class TestMLAHostDedupPrimitives(unittest.TestCase):
+    def test_disabled_flag_is_a_noop(self):
+        with mock.patch(
+            "sglang.srt.mem_cache.mla_host_dedup.mla_host_dedup_eligible"
+        ) as eligible:
+            context = maybe_create_mla_host_dedup_context(
+                object(), object(), None, None, None, enabled=False
+            )
+
+        self.assertIsNone(context)
+        eligible.assert_not_called()
+
+    def test_dummy_host_pools_keep_allocator_metadata_only(self):
+        mla_device_pool = _device_pool_stub(
+            layer_num=2,
+            store_dtype=torch.float16,
+            kv_lora_rank=4,
+            qk_rope_head_dim=2,
+            size=8,
+            start_layer=0,
+            end_layer=2,
+        )
+        mla_host = MLATokenToKVPoolHost(
+            mla_device_pool,
+            host_to_device_ratio=2,
+            host_size=0,
+            page_size=2,
+            layout="page_first",
+            pin_memory=False,
+            is_dummy=True,
+        )
+
+        self.assertTrue(mla_host._is_dummy)
+        self.assertIsNone(mla_host.kv_buffer)
+        self.assertIsNone(mla_host.data_ptrs)
+        self.assertEqual(mla_host.get_contiguous_buf_infos(), ([], [], []))
+        slots = mla_host.alloc(2)
+        self.assertEqual(slots.tolist(), [0, 1])
+        with self.assertRaisesRegex(AssertionError, "load on a dummy"):
+            mla_host.load_to_device_per_layer(
+                mla_device_pool, slots, slots, layer_id=0, io_backend="kernel"
+            )
+
+        dsa_device_pool = _device_pool_stub(
+            layer_num=2,
+            store_dtype=torch.float16,
+            size=8,
+            start_layer=0,
+            end_layer=2,
+            index_head_dim=8,
+            quant_block_size=4,
+        )
+        indexer_host = DSAIndexerPoolHost(
+            dsa_device_pool,
+            mla_host,
+            layout="page_first",
+            pin_memory=False,
+            is_dummy=True,
+        )
+
+        self.assertTrue(indexer_host._is_dummy)
+        self.assertIsNone(indexer_host.index_k_with_scale_buffer)
+        self.assertIsNone(indexer_host.index_k_device_ptrs)
+        self.assertEqual(indexer_host.size, mla_host.size)
+        with self.assertRaisesRegex(AssertionError, "load on a dummy"):
+            indexer_host.load_to_device_per_layer(
+                dsa_device_pool, slots, slots, layer_id=0, io_backend="kernel"
+            )
+
+    def test_layer_broadcast_reuses_full_staging_capacity(self):
+        broadcaster = MLAHostDedupBroadcaster.__new__(MLAHostDedupBroadcaster)
+        broadcaster.is_src = True
+        broadcaster.src_global_rank = 0
+        broadcaster.group = object()
+
+        layer_buffers = [
+            torch.arange(24, dtype=torch.float32).reshape(6, 1, 4),
+            torch.arange(24, 48, dtype=torch.float32).reshape(6, 1, 4),
+        ]
+        target = torch.tensor([0, 2, 5], dtype=torch.int64)
+        staging = torch.empty(2 * 3 * 4, dtype=torch.float32)
+
+        with mock.patch.object(torch.distributed, "broadcast") as broadcast:
+            broadcaster._bcast_layer(layer_buffers, staging, target, 4, layer_id=1)
+
+        broadcast.assert_called_once()
+        expected = layer_buffers[1].index_select(0, target)
+        torch.testing.assert_close(
+            staging[: expected.numel()].view_as(expected), expected
+        )
+
+        broadcaster.is_src = False
+        received = [torch.zeros_like(layer) for layer in layer_buffers]
+        with mock.patch.object(torch.distributed, "broadcast"):
+            broadcaster._bcast_layer(received, staging, target, 4, layer_id=1)
+        torch.testing.assert_close(received[1].index_select(0, target), expected)
+
+    def test_chunk_tokens_uses_environment(self):
+        device_pool = _device_pool_stub(
+            layer_num=2,
+            device=torch.device("cpu"),
+            kv_cache_dim=4,
+            kv_buffer=[torch.empty(3, 1, 4), torch.empty(3, 1, 4)],
+        )
+
+        with (
+            envs.SGLANG_MLA_DEDUP_CHUNK_TOKENS.override(7),
+            mock.patch(
+                "sglang.srt.mem_cache.mla_host_dedup.mla_dedup_rank_and_size",
+                return_value=(0, 2),
+            ),
+        ):
+            broadcaster = MLAHostDedupBroadcaster(
+                device_pool, group=object(), src_global_rank=0
+            )
+
+        self.assertEqual(broadcaster.chunk_tokens, 7)
+        self.assertEqual(broadcaster.kv_staging.numel(), 2 * 7 * 4)
+
+    def test_chunk_tokens_must_be_positive(self):
+        device_pool = _device_pool_stub(
+            layer_num=2,
+            device=torch.device("cpu"),
+            kv_cache_dim=4,
+            kv_buffer=[torch.empty(3, 1, 4), torch.empty(3, 1, 4)],
+        )
+
+        with (
+            envs.SGLANG_MLA_DEDUP_CHUNK_TOKENS.override(0),
+            mock.patch(
+                "sglang.srt.mem_cache.mla_host_dedup.mla_dedup_rank_and_size",
+                return_value=(0, 2),
+            ),
+            self.assertRaisesRegex(ValueError, "must be positive"),
+        ):
+            MLAHostDedupBroadcaster(device_pool, group=object(), src_global_rank=0)
+
+    def test_build_eagerly_warms_dedicated_nccl_group(self):
+        tp_group = object()
+        dedicated_group = object()
+        device_pool = _device_pool_stub(
+            layer_num=2,
+            device=torch.device("cpu"),
+            kv_cache_dim=4,
+            kv_buffer=[torch.empty(3, 1, 4)],
+        )
+
+        with (
+            mock.patch(
+                "sglang.srt.mem_cache.mla_host_dedup.is_dp_attention_enabled",
+                return_value=False,
+            ),
+            mock.patch(
+                "sglang.srt.mem_cache.mla_host_dedup.mla_dedup_rank_and_size",
+                return_value=(0, 2),
+            ),
+            mock.patch.object(
+                torch.distributed,
+                "get_process_group_ranks",
+                return_value=[4, 5],
+            ),
+            mock.patch(
+                "sglang.srt.distributed.parallel_state.create_custom_parallel_group",
+                return_value=dedicated_group,
+            ) as create_group,
+            mock.patch.object(torch.distributed, "broadcast") as broadcast,
+            mock.patch.object(torch.cuda, "synchronize") as synchronize,
+        ):
+            broadcaster = MLAHostDedupBroadcaster.build(
+                device_pool, tp_group, attn_tp_group=None
+            )
+
+        create_group.assert_called_once_with(group_ranks=[4, 5], backend="nccl")
+        broadcast.assert_called_once()
+        self.assertEqual(broadcast.call_args.args[0].numel(), 1)
+        self.assertIs(broadcast.call_args.kwargs["group"], dedicated_group)
+        self.assertEqual(broadcast.call_args.kwargs["src"], 4)
+        synchronize.assert_called_once_with(device_pool.device)
+        self.assertIs(broadcaster.group, dedicated_group)
+
+    def test_indexer_pages_preserve_logical_order(self):
+        broadcaster = MLAHostDedupBroadcaster.__new__(MLAHostDedupBroadcaster)
+        broadcaster.device = torch.device("cpu")
+        broadcaster.device_pool = SimpleNamespace(page_size=4)
+        broadcaster.idx_bufs = [object()]
+
+        device_indices = torch.tensor([8, 9, 10, 11, 0, 1, 2, 3])
+        prepared_indices, page_indices = broadcaster.prepare_broadcast(
+            device_indices, _FakeStream()
+        )
+
+        self.assertIs(prepared_indices, device_indices)
+        torch.testing.assert_close(page_indices, torch.tensor([2, 0]))
+
+    def test_indexer_rejects_partial_pages(self):
+        broadcaster = MLAHostDedupBroadcaster.__new__(MLAHostDedupBroadcaster)
+        broadcaster.device = torch.device("cpu")
+        broadcaster.device_pool = SimpleNamespace(page_size=4)
+        broadcaster.idx_bufs = [object()]
+
+        with self.assertRaisesRegex(ValueError, "page-aligned device indices"):
+            broadcaster.prepare_broadcast(torch.arange(7), _FakeStream())
+
+    def test_context_destroys_all_owned_process_groups(self):
+        broadcaster = mock.Mock()
+        hit_group = object()
+        completion_group = object()
+        context = MLAHostDedupContext(
+            broadcaster=broadcaster,
+            prefetch_hits_sync_groups=[hit_group],
+            prefetch_completion_sync_groups=[completion_group],
+        )
+
+        with mock.patch.object(torch.distributed, "destroy_process_group") as destroy:
+            context.destroy()
+
+        broadcaster.destroy.assert_called_once()
+        self.assertEqual(
+            destroy.call_args_list,
+            [mock.call(hit_group), mock.call(completion_group)],
+        )
+        self.assertIsNone(context.prefetch_hits_sync_groups)
+        self.assertIsNone(context.prefetch_completion_sync_groups)
+
+
+if __name__ == "__main__":
+    unittest.main()