From 8a0863c72840ce1ff4fd2e5e2e7b8dd20782c8d8 Mon Sep 17 00:00:00 2001 From: HZY Date: Tue, 8 Sep 2026 17:40:21 +0200 Subject: [PATCH] [HiCache] Add MLA host-dedup primitives (#36800) Co-authored-by: Zhangheng --- python/sglang/srt/environ.py | 2 + python/sglang/srt/mem_cache/mla_host_dedup.py | 305 ++++++++++++++++++ python/sglang/srt/mem_cache/pool_host/dsa.py | 27 +- python/sglang/srt/mem_cache/pool_host/mla.py | 99 +++++- .../test_mla_host_dedup_primitives.py | 258 +++++++++++++++ 5 files changed, 688 insertions(+), 3 deletions(-) create mode 100644 python/sglang/srt/mem_cache/mla_host_dedup.py create mode 100644 test/registered/unit/mem_cache/test_mla_host_dedup_primitives.py diff --git a/python/sglang/srt/environ.py b/python/sglang/srt/environ.py index 0b9b87903..576772665 100644 --- a/python/sglang/srt/environ.py +++ b/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) diff --git a/python/sglang/srt/mem_cache/mla_host_dedup.py b/python/sglang/srt/mem_cache/mla_host_dedup.py new file mode 100644 index 000000000..08145441c --- /dev/null +++ b/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 diff --git a/python/sglang/srt/mem_cache/pool_host/dsa.py b/python/sglang/srt/mem_cache/pool_host/dsa.py index df76d5e18..7ef693fb0 100644 --- a/python/sglang/srt/mem_cache/pool_host/dsa.py +++ b/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( diff --git a/python/sglang/srt/mem_cache/pool_host/mla.py b/python/sglang/srt/mem_cache/pool_host/mla.py index 41c651c3b..f3762c064 100644 --- a/python/sglang/srt/mem_cache/pool_host/mla.py +++ b/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): diff --git a/test/registered/unit/mem_cache/test_mla_host_dedup_primitives.py b/test/registered/unit/mem_cache/test_mla_host_dedup_primitives.py new file mode 100644 index 000000000..57c6905d8 --- /dev/null +++ b/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()