minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

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

Browse Source 
Co-authored-by: Zhangheng <hzh0425@apache.org>
This commit is contained in:
HZY
2026-09-08 23:40:21 +08:00
committed by GitHub
co-authored by Zhangheng
parent 482e9f257b
commit 8a0863c728
5 changed files with 688 additions and 3 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/environ.py
+2
View File
@@ -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
+305
View File
@@ -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
+25 -2
View File
@@ -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
+98 -1
View File
@@ -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):