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

[Unified Cache][6/N]: Add UMBP external linker (#37578)

Browse Source 
Co-authored-by: Zhangheng <hzh0425@apache.org>
This commit is contained in:
Niko Ma
2026-09-05 09:21:43 +08:00
committed by GitHub
co-authored by Zhangheng
parent bc727bc4ee
commit f1f2380d2b
9 changed files with 1584 additions and 207 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/mem_cache/registry.py
+6
View File
@@ -204,6 +204,12 @@ def _create_unified_radix_cache(
)
linker_cls = MooncakeDirectLinker
elif backend == "mori":
from sglang.srt.mem_cache.storage.umbp.umbp_direct_linker import (
UMBPDirectLinker,
)
linker_cls = UMBPDirectLinker
else:
raise ValueError(
f"Unknown unified cache external linker backend: {backend!r}"
python/sglang/srt/mem_cache/storage/umbp/umbp_direct_linker.py
+1050
View File
File diff suppressed because it is too large Load Diff
python/sglang/srt/mem_cache/storage/umbp/umbp_host_allocator.py
+27 -13
View File
@@ -45,6 +45,9 @@ class UMBPHostTensorAllocator(HostTensorAllocator):
)
self._numa_node = _int_env("SGLANG_HICACHE_HOST_NUMA_NODE", -1)
self._prefault = _bool_env("SGLANG_HICACHE_HOST_PREFAULT", True)
# Standalone mode needs fd-shareable backing; allocation precedes
# config parsing.
self._standalone_process = bool(os.getenv("UMBP_STANDALONE_ADDRESS"))
self._handles: Dict[int, Any] = {}
def allocate(
@@ -62,11 +65,18 @@ class UMBPHostTensorAllocator(HostTensorAllocator):
element_size = torch.empty((), dtype=dtype).element_size()
nbytes = math.prod(int(dim) for dim in dims) * element_size
requested_backing = (
self._mod.UMBPHostBufferBacking.AnonymousHugetlb
if self._use_hugepage
else self._mod.UMBPHostBufferBacking.Anonymous
)
if self._standalone_process:
requested_backing = (
self._mod.UMBPHostBufferBacking.AnonymousShmHugetlb
if self._use_hugepage
else self._mod.UMBPHostBufferBacking.AnonymousShm
)
else:
requested_backing = (
self._mod.UMBPHostBufferBacking.AnonymousHugetlb
if self._use_hugepage
else self._mod.UMBPHostBufferBacking.Anonymous
)
handle = self._allocator.alloc(
nbytes,
@@ -101,15 +111,19 @@ class UMBPHostTensorAllocator(HostTensorAllocator):
handle.mapped_size,
self._numa_node,
)
if (
self._use_hugepage
and handle.actual_backing == self._mod.UMBPHostBufferBacking.Anonymous
):
demoted = handle.actual_backing == (
self._mod.UMBPHostBufferBacking.AnonymousShm
if self._standalone_process
else self._mod.UMBPHostBufferBacking.Anonymous
)
if self._use_hugepage and demoted:
logger.warning(
"UMBPHostTensorAllocator: requested AnonymousHugetlb backing "
"but kernel demoted to Anonymous (4 KiB pages). Check "
"vm.nr_hugepages and HugePages_Free in /proc/meminfo. "
"Performance and AINIC MR-size benefits will not apply."
"UMBPHostTensorAllocator: requested %s backing but kernel "
"demoted to %s (4 KiB pages). Check vm.nr_hugepages and "
"HugePages_Free in /proc/meminfo. Performance and AINIC "
"MR-size benefits will not apply.",
requested_backing,
handle.actual_backing,
)
return tensor.view(dims)
python/sglang/srt/mem_cache/storage/umbp/umbp_store.py
+358 -192
View File
@@ -38,6 +38,8 @@ def _import_umbp_client():
UMBPIoBackend = getattr(umbp_mod, "UMBPIoBackend", None)
UMBPDurabilityMode = getattr(umbp_mod, "UMBPDurabilityMode", None)
UMBPDistributedConfig = getattr(umbp_mod, "UMBPDistributedConfig", None)
UMBPStandaloneProcessConfig = getattr(umbp_mod, "UMBPStandaloneProcessConfig", None)
UMBPDeploymentMode = getattr(umbp_mod, "UMBPDeploymentMode", None)
return (
UMBPClient,
@@ -46,6 +48,8 @@ def _import_umbp_client():
UMBPIoBackend,
UMBPDurabilityMode,
UMBPDistributedConfig,
UMBPStandaloneProcessConfig,
UMBPDeploymentMode,
)
@@ -134,6 +138,10 @@ def _select_rank_config_value(
# knobs outside this list go through the "spdk_passthrough" escape hatch.
_COMMON_EXTRA_KEYS = frozenset(
{
"node_address",
"node_id",
"node_tags",
"tags",
"dram_capacity_bytes",
"ssd_enabled",
"ssd_storage_dir",
@@ -164,6 +172,8 @@ _COMMON_EXTRA_KEYS = frozenset(
"kv_events_subscriber",
"kv_events_endpoint",
"kv_events_topic",
"disable_zero_copy_register",
"extra_backend_tag",
}
)
@@ -175,24 +185,25 @@ _STANDALONE_ONLY_EXTRA_KEYS = frozenset(
"eviction_policy",
"eviction_candidate_window",
"auto_promote_on_read",
"standalone_address",
"standalone_auto_start",
"standalone_startup_timeout_ms",
}
)
_DISTRIBUTED_ONLY_EXTRA_KEYS = frozenset(
{
"master_address",
"node_address",
"node_id",
"auto_heartbeat",
"io_engine_host",
"io_engine_port",
"staging_buffer_size",
"ranged_scratch_size",
"ssd_staging_buffer_size",
"ssd_staging_buffer_slots",
"peer_service_port",
"cache_remote_fetches",
"dram_page_size",
"disable_zero_copy_register",
}
)
@@ -238,7 +249,13 @@ class UMBPStore(HiCacheStorage):
self,
storage_config: HiCacheStorageConfig = None,
mem_pool_host: HostKVCache = None,
*,
per_rank_keyspace: bool = False,
):
# per_rank_keyspace: the direct linker already organises keys by its own
# cache-group rank, so HiCache's shared-SSD leader/follower deduplication
# must not be layered on top. Default False preserves the existing
# HiCache L3 behaviour.
(
UMBPClient,
UMBPConfig,
@@ -246,6 +263,8 @@ class UMBPStore(HiCacheStorage):
UMBPIoBackend,
UMBPDurabilityMode,
UMBPDistributedConfig,
UMBPStandaloneProcessConfig,
UMBPDeploymentMode,
) = _import_umbp_client()
if storage_config is not None:
@@ -260,6 +279,7 @@ class UMBPStore(HiCacheStorage):
self.pp_rank = 0
self.pp_size = 1
self.tp_size = 1
self._umbp_deployment_mode_enum = UMBPDeploymentMode
cfg = UMBPConfig.from_environment()
# UMBPStore owns role selection explicitly. Do not inherit LOCAL_RANK /
@@ -268,6 +288,12 @@ class UMBPStore(HiCacheStorage):
# and skip writes.
cfg.role = UMBPRole.Standalone
extra = getattr(storage_config, "extra_config", None) or {}
prefix_parts = []
if extra.get("extra_backend_tag") is not None:
prefix_parts.append(str(extra["extra_backend_tag"]))
if storage_config is not None and storage_config.model_name:
prefix_parts.append("-".join(storage_config.model_name.split("/")))
self.config_prefix = "_".join(prefix_parts) if prefix_parts else None
explicit_tenant_id = (
os.getenv("UMBP_SPDK_PROXY_TENANT_ID") is not None
or "spdk_proxy_tenant_id" in extra
@@ -461,6 +487,58 @@ class UMBPStore(HiCacheStorage):
master_address = extra.get(
"master_address", _optional_env_str("UMBP_MASTER_ADDRESS")
)
standalone_extra_address = extra.get("standalone_address")
standalone_env_address = _optional_env_str("UMBP_STANDALONE_ADDRESS")
standalone_address = standalone_extra_address or standalone_env_address
# Verify the client did not silently fall back to local mode.
self._standalone_process_expected = bool(standalone_address)
if master_address and standalone_address:
raise ValueError(
"master_address and standalone_address are mutually exclusive "
"(distributed vs. standalone-process mode)."
)
if (
mem_pool_host is not None
and standalone_extra_address
and not standalone_env_address
):
raise ValueError(
"standalone_address in hicache-storage-backend-extra-config is "
"not supported when a host KV pool is present. The host memory "
"pool allocator chooses Anonymous vs. AnonymousShm before "
"extra_config is parsed, so set UMBP_STANDALONE_ADDRESS in the "
"process environment instead."
)
# Both remote modes use the same worker identity.
def _resolve_node_address() -> str:
node_address = extra.get(
"node_address", _optional_env_str("UMBP_NODE_ADDRESS")
)
if node_address is None:
return _default_node_address()
return _select_rank_config_value(
node_address, unique_rank, "node_address", str
)
def _resolve_node_id(node_address: str) -> str:
node_id = extra.get("node_id", _optional_env_str("UMBP_NODE_ID"))
if node_id is None:
return (
f"{node_address}:dp{dp_rank_hint if dp_rank_hint is not None else 0}"
f":pp{self.pp_rank}:tp{self.local_rank}"
)
return _select_rank_config_value(node_id, unique_rank, "node_id", str)
def _resolve_node_tags() -> List[str]:
raw_tags = extra.get("node_tags", extra.get("tags"))
if raw_tags is None:
raw_tags = _optional_env_str("UMBP_NODE_TAGS")
if raw_tags is None:
return []
if isinstance(raw_tags, str):
return [tag.strip() for tag in raw_tags.split(",") if tag.strip()]
return [str(tag) for tag in raw_tags]
_warn_extra_config_scope(extra, distributed_enabled=bool(master_address))
if master_address and UMBPDistributedConfig is not None:
@@ -470,33 +548,11 @@ class UMBPStore(HiCacheStorage):
if "ssd_copy_worker_threads" not in extra:
cfg.copy_pipeline.worker_threads = 1
node_address = extra.get(
"node_address", _optional_env_str("UMBP_NODE_ADDRESS")
)
if node_address is None:
node_address = _default_node_address()
else:
node_address = _select_rank_config_value(
node_address,
unique_rank,
"node_address",
str,
)
node_address = _resolve_node_address()
dist_cfg.master_config.node_address = node_address
node_id = extra.get("node_id", _optional_env_str("UMBP_NODE_ID"))
if node_id is None:
dist_cfg.master_config.node_id = (
f"{node_address}:dp{dp_rank_hint if dp_rank_hint is not None else 0}"
f":pp{self.pp_rank}:tp{self.local_rank}"
)
else:
dist_cfg.master_config.node_id = _select_rank_config_value(
node_id,
unique_rank,
"node_id",
str,
)
dist_cfg.master_config.node_id = _resolve_node_id(node_address)
if hasattr(dist_cfg.master_config, "tags"):
dist_cfg.master_config.tags = _resolve_node_tags()
if "auto_heartbeat" in extra:
dist_cfg.master_config.auto_heartbeat = _strict_bool(
@@ -531,6 +587,10 @@ class UMBPStore(HiCacheStorage):
if "staging_buffer_size" in extra:
dist_cfg.staging_buffer_size = int(extra["staging_buffer_size"])
if "ranged_scratch_size" in extra and hasattr(
dist_cfg, "ranged_scratch_size"
):
dist_cfg.ranged_scratch_size = int(extra["ranged_scratch_size"])
if "ssd_staging_buffer_size" in extra and hasattr(
dist_cfg, "ssd_staging_buffer_size"
@@ -604,8 +664,8 @@ class UMBPStore(HiCacheStorage):
meta = mem_pool_host.get_split_heads_page_buffer_meta(dummy, sf)
else:
meta = mem_pool_host.get_page_buffer_meta(dummy)
# meta is None for a logical-anchor group (see note above);
# esz is the per-page element-size list otherwise.
# A hybrid logical anchor returns None here by design; leave
# dram_page_size at 0 and let the per-pool v2 sizes handle it.
esz = meta[1] if meta else None
page_byte_size = int(esz[0]) if esz else 0
@@ -647,6 +707,52 @@ class UMBPStore(HiCacheStorage):
dist_cfg.io_engine.port,
dist_cfg.peer_service_port,
)
elif standalone_address:
if UMBPStandaloneProcessConfig is None:
raise RuntimeError(
"Installed mori does not expose UMBPStandaloneProcessConfig"
)
standalone_cfg = UMBPStandaloneProcessConfig()
standalone_cfg.address = str(standalone_address)
auto_start = extra.get(
"standalone_auto_start",
_optional_env_str("UMBP_STANDALONE_AUTO_START"),
)
if auto_start is not None:
standalone_cfg.auto_start = _strict_bool(
auto_start, "standalone_auto_start"
)
startup_timeout_ms = extra.get(
"standalone_startup_timeout_ms",
_optional_env_int("UMBP_STANDALONE_STARTUP_TIMEOUT_MS"),
)
if startup_timeout_ms is not None:
standalone_cfg.startup_timeout_ms = int(startup_timeout_ms)
if standalone_cfg.startup_timeout_ms <= 0:
raise ValueError("standalone_startup_timeout_ms must be > 0")
if all(
hasattr(standalone_cfg, field)
for field in ("worker_node_address", "worker_node_id", "tags")
):
worker_node_address = _resolve_node_address()
standalone_cfg.worker_node_address = worker_node_address
standalone_cfg.worker_node_id = _resolve_node_id(worker_node_address)
standalone_cfg.tags = _resolve_node_tags()
else:
logger.warning(
"UMBPStore standalone-process mode: installed mori does not "
"expose worker identity on UMBPStandaloneProcessConfig; a "
"distributed-backed standalone server cannot build per-worker "
"external-KV identities."
)
cfg.standalone_process = standalone_cfg
logger.info(
"UMBPStore standalone-process mode: address=%s, auto_start=%s, "
"startup_timeout_ms=%s",
standalone_cfg.address,
standalone_cfg.auto_start,
standalone_cfg.startup_timeout_ms,
)
self.storage_config = storage_config
@@ -657,8 +763,25 @@ class UMBPStore(HiCacheStorage):
self.is_mla_follower = False
tp_size = self.tp_size
use_spdk = cfg.ssd.ssd_backend in ("spdk", "spdk_proxy")
distributed_enabled = cfg.distributed is not None
if not distributed_enabled and self.is_mla_backend and tp_size > 1:
remote_process_enabled = (
cfg.distributed is not None or cfg.standalone_process is not None
)
# Shared SSD exists to deduplicate MLA KV, which TP replicates: one
# rank owns the bytes and the others read them back. A caller that
# already keys per rank has nothing to deduplicate -- and would be
# broken by the scheme, because a follower would be sent looking for
# keys the leader never wrote under the follower's own suffix.
#
# It is also the reason embedded mode could not run: followers are the
# one role whose client reports no ranged multi-buffer I/O, which
# page-granular objects require. Standalone never hit this, not by
# design but because remote_process_enabled short-circuits it there.
if (
not remote_process_enabled
and not per_rank_keyspace
and self.is_mla_backend
and tp_size > 1
):
cfg.ssd.enabled = True
if self.local_rank == 0:
# Leader: copy every DRAM write to shared SSD.
@@ -802,11 +925,10 @@ class UMBPStore(HiCacheStorage):
safe_cap = int(cfg.ssd.capacity_bytes * 0.95)
cfg.ssd.spdk_proxy_tenant_quota_bytes = max(1, safe_cap // dp_size_hint)
# Initialize registration state before the optional constructor-time
# register_mem_pool_host() call below. In particular, do not overwrite
# the logical-anchor flag after that call has detected a LogicalHostPool.
# Initialize before the optional constructor-time pool registration.
self.registered_pools: dict = {}
self._kv_anchor_is_logical = False
self._registered_regions: set = set()
self.client = UMBPClient(cfg)
if mem_pool_host is not None:
@@ -888,45 +1010,99 @@ class UMBPStore(HiCacheStorage):
"page_head",
], "UMBP store only supports page_first, page_first_direct, or page_head layout"
# Hybrid logical anchors (e.g. DeepSeek-V4's KV anchor LogicalHostPool)
# own only allocation indices and hold no physical KV tensor. Compute
# this once and reuse: there is nothing to register for RDMA here, v1
# I/O no-ops on it, and the real per-pool buffers are registered through
# register_mem_host_pool_v2().
# A logical anchor owns indices; side pools carry the data.
self._kv_anchor_is_logical = self.mem_pool_host.kv_buffer is None
self._zero_copy_registered = False
# Side-pool registration needs the mode even for a logical anchor.
self._is_standalone_process = False
if self.client is not None:
deployment_mode = None
mode_enum = self._umbp_deployment_mode_enum
try:
deployment_mode = self.client.get_deployment_mode()
if mode_enum is not None:
self._is_standalone_process = (
deployment_mode == mode_enum.StandaloneProcess
)
except Exception as exc:
if self._standalone_process_expected:
raise RuntimeError(
"UMBPStore expected standalone-process mode from "
"UMBP_STANDALONE_ADDRESS, but get_deployment_mode() failed."
) from exc
if self._standalone_process_expected:
if mode_enum is None:
raise RuntimeError(
"UMBPStore expected standalone-process mode, but "
"UMBPDeploymentMode is not exposed by mori.umbp."
)
if deployment_mode != mode_enum.StandaloneProcess:
raise RuntimeError(
"UMBPStore expected standalone-process mode, but the "
f"UMBP client reported deployment_mode={deployment_mode!r}."
)
if self._kv_anchor_is_logical:
return
# In distributed mode, pre-register the entire host KV buffer with the
# underlying RDMA IOEngine so PoolClient can take the zero-copy path
# for batch_get_into_ptr / batch_put_from_ptr (skips the staging
# buffer memcpy + lock and removes the per-call `staging_buffer_size`
# cap). Standalone returns true as no-op by IUMBPClient contract;
# we still gate on is_distributed() below to avoid a pointless call.
if self._register_host_buffer_for_zero_copy(mem_pool_host):
self._zero_copy_registered = True
self._zero_copy_registered = self._register_host_buffer_for_zero_copy(
mem_pool_host
)
@staticmethod
def _pool_physical_buffers(host_pool: HostKVCache) -> List[Any]:
"""Return every non-empty physical tensor exposed by a host pool."""
getter = getattr(host_pool, "get_hybrid_pool_buffer", None)
buffers = getter() if getter is not None else None
if not buffers:
buffers = [getattr(host_pool, "kv_buffer", None)]
flat: List[Any] = []
for buffer in buffers:
if buffer is None:
continue
for tensor in buffer if isinstance(buffer, (list, tuple)) else [buffer]:
# Empty views may share storage; register_memory rejects zero bytes.
if tensor is not None and tensor.numel() > 0:
flat.append(tensor)
return flat
@staticmethod
def _buffer_extent(buffer, allocator) -> tuple:
"""Registerable (base pointer, size) of the allocation behind a tensor."""
storage = buffer.untyped_storage()
base = int(storage.data_ptr())
size = int(storage.nbytes())
# Hugepage-backed mmaps are rounded up to the hugepage boundary, and
# ibv_reg_mr on AINIC / ROCm needs whole hugepages covered.
mapped_size_fn = getattr(allocator, "mapped_size_for", None)
mapped_size = (
mapped_size_fn(base)
if mapped_size_fn is not None
else getattr(allocator, "mapped_size", 0)
)
return base, max(size, int(mapped_size or 0))
def _register_host_buffer_for_zero_copy(self, host_pool: HostKVCache) -> bool:
"""Register a host pool's KV buffer with the RDMA IOEngine for zero-copy.
Shared by the single-pool path (register_mem_pool_host) and the
multi-pool path (register_mem_host_pool_v2). Returns True when the
buffer was successfully registered, False on any skip/failure (the
caller then transparently falls back to the staging-buffer path).
"""
"""Register host buffers; standalone failures are fatal without fallback."""
if self.client is None:
return False
is_standalone_process = getattr(self, "_is_standalone_process", False)
try:
is_distributed = bool(self.client.is_distributed())
except Exception:
is_distributed = False
if not is_distributed:
if not (is_distributed or is_standalone_process):
return False
if not hasattr(self.client, "register_memory"):
return False
if getattr(self, "_disable_zero_copy_register", False):
if is_standalone_process:
raise RuntimeError(
"disable_zero_copy_register is not supported in UMBP "
"standalone-process mode: there is no staging-buffer "
"fallback path."
)
logger.info(
"UMBPStore: skipping host KV buffer RDMA registration because "
"disable_zero_copy_register=true (UMBP_DISABLE_ZERO_COPY_REGISTER). "
@@ -934,69 +1110,62 @@ class UMBPStore(HiCacheStorage):
"size is capped by distributed.staging_buffer_size."
)
return False
# NOTE(layer_first): this only handles the page_first layout, where a
# host pool exposes a single contiguous `kv_buffer` that we can register
# for RDMA in one shot. If UMBP later supports a layer_first layout, or
# side pools that expose multiple buffers via get_hybrid_pool_buffer()
# (e.g. DSAIndexerPoolHost, whose buffer lives in
# index_k_with_scale_buffer rather than kv_buffer), this branch must be
# extended to register every per-layer / per-buffer region. Otherwise
# such pools bypass zero-copy and silently fall back to the slower
# staging-buffer path.
kv_buffer = getattr(host_pool, "kv_buffer", None)
if kv_buffer is None:
buffers = self._pool_physical_buffers(host_pool)
if not buffers:
if is_standalone_process:
raise RuntimeError(
f"UMBPStore: {type(host_pool).__name__} exposes no host buffer "
"to register; standalone-process mode has no fallback path."
)
return False
try:
host_ptr = int(kv_buffer.data_ptr())
host_size = int(kv_buffer.numel() * kv_buffer.element_size())
# When the buffer is backed by hugepages the mmap region is
# rounded up to the hugepage boundary. RDMA ibv_reg_mr on
# some NICs (AINIC / ROCm) requires the registered region to
# cover complete hugepages, so use the full mapped_size
# instead of the logical tensor size.
allocator = getattr(host_pool, "allocator", None)
mapped_size_fn = getattr(allocator, "mapped_size_for", None)
if mapped_size_fn is not None:
mapped_size = mapped_size_fn(host_ptr)
else:
mapped_size = getattr(allocator, "mapped_size", 0)
if mapped_size > host_size:
host_size = mapped_size
ok = bool(self.client.register_memory(host_ptr, host_size))
except Exception as exc:
logger.warning(
"UMBPStore: register_memory failed (%s); falling back to staging "
"buffer path. Per-transfer size will be capped by "
"distributed.staging_buffer_size.",
exc,
)
return False
if ok:
mode = "standalone-process" if is_standalone_process else "distributed"
allocator = getattr(host_pool, "allocator", None)
# Already registered storage counts as covered.
covered = 0
for buffer in buffers:
try:
host_ptr, host_size = self._buffer_extent(buffer, allocator)
if host_ptr in self._registered_regions:
covered += 1
continue
ok = bool(self.client.register_memory(host_ptr, host_size))
except Exception as exc:
if is_standalone_process:
raise RuntimeError(
"UMBPStore: register_memory failed in standalone-process "
f"mode and cannot fall back: {exc}"
) from exc
logger.warning(
"UMBPStore: register_memory failed (%s); falling back to staging "
"buffer path. Per-transfer size will be capped by "
"distributed.staging_buffer_size.",
exc,
)
return False
if not ok:
if is_standalone_process:
raise RuntimeError(
"UMBPStore: register_memory returned false in "
"standalone-process mode; no fallback path exists."
)
logger.warning(
"UMBPStore: register_memory returned false; staying on staging "
"buffer fallback path."
)
return False
self._registered_regions.add(host_ptr)
covered += 1
logger.info(
"UMBPStore: registered host KV buffer for RDMA zero-copy "
"(ptr=0x%x, size=%d MB)",
"UMBPStore: registered host buffer for zero-copy "
"(ptr=0x%x, size=%d MB, mode=%s)",
host_ptr,
host_size // (1024 * 1024),
mode,
)
return True
logger.warning(
"UMBPStore: register_memory returned false; staying on staging "
"buffer fallback path."
)
return False
return covered == len(buffers)
def register_mem_host_pool_v2(self, host_pool: HostKVCache, host_pool_name):
"""Register an additional hybrid side pool (DeepSeek-V4 HostPoolGroup).
The controller calls this once per PoolEntry in the group, including the
KV anchor. The KV anchor is logical (no physical tensor) so we skip it;
its allocation-index role is unrelated to storage I/O. Every other pool
(SWA / compressed KV / indexer / state) carries a real page_first KV
buffer that must be (a) resolvable by name at v2 I/O time and (b)
registered with the RDMA IOEngine for zero-copy transfers.
"""
# KV anchor is either already registered via register_mem_pool_host()
# (non-hybrid single pool) or purely logical (hybrid group). Skip it.
if host_pool_name == PoolName.KV:
return
self.registered_pools[host_pool_name] = host_pool
@@ -1078,8 +1247,6 @@ class UMBPStore(HiCacheStorage):
extra_info: Optional[HiCacheStorageExtraInfo] = None,
) -> List[bool]:
if self._kv_anchor_is_logical:
# DeepSeek-V4's KV anchor is logical only; the physical KV data is
# carried by the v2 side pools, so there is nothing to read here.
return [True] * len(keys)
key_strs, buffer_ptrs, buffer_sizes = self._batch_preprocess(keys, host_indices)
@@ -1152,8 +1319,6 @@ class UMBPStore(HiCacheStorage):
return [True] * page_count
if self._kv_anchor_is_logical:
# DeepSeek-V4's KV anchor is logical only; the physical KV data is
# written by the v2 side pools, so there is nothing to write here.
return [True] * len(keys)
key_strs, buffer_ptrs, buffer_sizes = self._batch_preprocess(keys, host_indices)
@@ -1219,47 +1384,66 @@ class UMBPStore(HiCacheStorage):
return hit_count // key_multiplier
# ------------------------------------------------------------------
# Multi-pool v2 interface (DeepSeek-V4 hybrid HiCache HostPoolGroup)
#
# The DeepSeek-V4 HiCache stack splits KV state across several page_first
# side pools (SWA / compressed KV / indexer / state), coordinated by a
# logical KV anchor that owns only page indices. The controller registers
# each real pool through register_mem_host_pool_v2() and drives storage
# via these _v2 methods, one PoolTransfer per pool. This mirrors the proven
# MooncakeStore / HiCacheHF3FS design, specialized for UMBP's page_first,
# single-object-per-page layout (each page -> exactly one storage object).
# Multi-pool v2 interface
# ------------------------------------------------------------------
def _get_hybrid_page_component_keys(self, page_keys, transfer: PoolTransfer):
"""Map per-page logical keys to per-object storage keys for a side pool.
For UMBP every registered side pool is page_first and stores one object
per page (MLA: a single K object; MHA: a K and a V object), so the
component-key count is an exact multiple of the page count. The pool
name is embedded in the suffix so pages that share a hash across pools
never collide.
"""
def _get_hybrid_page_component_keys(
self, page_keys, transfer: PoolTransfer, *, rank_suffix: Optional[str] = None
):
"""Expand logical page keys for one registered hybrid side pool."""
pool_name = transfer.name
host_pool = self.registered_pools.get(pool_name)
if host_pool is None:
raise ValueError(f"Unregistered UMBP hybrid pool: {pool_name}")
if self.is_mla_backend:
# Single compressed object per page.
suffixes = [f"_{self.mla_suffix}_{pool_name}"]
mla_suffix = self.mla_suffix if rank_suffix is None else rank_suffix
mha_suffix = (
getattr(self, "mha_suffix", mla_suffix)
if rank_suffix is None
else rank_suffix
)
components = getattr(host_pool, "components", None)
if pool_name == PoolName.MAMBA:
conv_num = len(getattr(host_pool, "conv_buffer", None) or [])
suffixes = [f"_{mha_suffix}_conv_{i}" for i in range(conv_num)]
if getattr(host_pool, "temporal_state_elem_size", 1) > 0:
suffixes = [f"_{mha_suffix}_temporal"] + suffixes
elif components is not None and len(components) == 1:
suffixes = [f"_{mla_suffix}_{pool_name}"]
elif components is not None and len(components) == 2:
# Packed DevicePoolEntry K/V components share one stored object.
suffixes = (
[f"_{mha_suffix}_{pool_name}"]
if host_pool.packed
else [
f"_{mha_suffix}_{pool_name}_k",
f"_{mha_suffix}_{pool_name}_v",
]
)
elif components is not None:
raise ValueError(
f"Unsupported UMBP component count for pool {pool_name}: "
f"{len(components)}"
)
elif self.is_mla_backend:
suffixes = [f"_{mla_suffix}_{pool_name}"]
elif getattr(host_pool, "v_buffer", None) is not None:
# Ordinary MHA side pool mirrors a K/V pool.
suffixes = [
f"_{self.mha_suffix}_{pool_name}_k",
f"_{self.mha_suffix}_{pool_name}_v",
f"_{mha_suffix}_{pool_name}_k",
f"_{mha_suffix}_{pool_name}_v",
]
else:
suffixes = [f"_{self.mha_suffix}_{pool_name}"]
suffixes = [f"_{mha_suffix}_{pool_name}"]
key_multiplier = len(suffixes)
component_keys = [
f"{page_key}{suffix}" for page_key in page_keys for suffix in suffixes
]
return component_keys, key_multiplier
if self.config_prefix:
component_keys = [
f"{self.config_prefix}_{component_key}"
for component_key in component_keys
]
return component_keys, len(suffixes)
def batch_exists_v2(
self,
@@ -1267,20 +1451,18 @@ class UMBPStore(HiCacheStorage):
pool_transfers: Optional[List[PoolTransfer]] = None,
extra_info: Optional[HiCacheStorageExtraInfo] = None,
) -> PoolTransferResult:
if self._kv_anchor_is_logical:
# Logical KV anchor: no physical KV object exists in UMBP, so the
# usable prefix is bounded entirely by the required side pools.
kv_pages = len(keys)
else:
kv_pages = self.batch_exists(keys, extra_info)
kv_pages = (
len(keys)
if self._kv_anchor_is_logical
else self.batch_exists(keys, extra_info)
)
hit_count: dict = {PoolName.KV: kv_pages} if kv_pages else {}
final_pages = kv_pages
for transfer in pool_transfers or []:
if final_pages == 0:
break
component_keys, key_multiplier = self._get_hybrid_page_component_keys(
component_keys, multiplier = self._get_hybrid_page_component_keys(
keys[:final_pages], transfer
)
exists = list(self.client.batch_exists(component_keys))
@@ -1294,18 +1476,15 @@ class UMBPStore(HiCacheStorage):
)
final_pages = 0
break
# Collapse per-object results into per-page presence.
page_exists = [
all(exists[i * key_multiplier : (i + 1) * key_multiplier])
all(exists[i * multiplier : (i + 1) * multiplier])
for i in range(final_pages)
]
boundary = 0
if transfer.hit_policy == PoolHitPolicy.ALL_PAGES:
try:
boundary = page_exists.index(False)
except ValueError:
boundary = final_pages
boundary = (
page_exists.index(False) if False in page_exists else final_pages
)
elif transfer.hit_policy == PoolHitPolicy.TRAILING_PAGES:
trailing = max(1, len(transfer.keys) if transfer.keys else 1)
for prefix_len in range(final_pages, 0, -1):
@@ -1334,38 +1513,26 @@ class UMBPStore(HiCacheStorage):
if not keys or host_indices is None:
results[transfer.name] = [False] * len(keys)
continue
assert len(keys) == len(host_indices) // page_size
if len(keys) != len(host_indices) // page_size:
raise ValueError(
f"UMBP v2 pool {transfer.name} has {len(keys)} keys for "
f"{len(host_indices)} indices with page_size={page_size}."
)
key_strs, key_multiplier = self._get_hybrid_page_component_keys(
keys, transfer
key_strs, multiplier = self._get_hybrid_page_component_keys(keys, transfer)
ptrs, sizes = host_pool.get_page_buffer_meta(host_indices)
if not len(key_strs) == len(ptrs) == len(sizes):
raise ValueError(
f"UMBP v2 buffer-meta mismatch for pool {transfer.name}: "
f"keys={len(key_strs)} ptrs={len(ptrs)} sizes={len(sizes)}"
)
operation = (
self.client.batch_put_from_ptr
if is_set
else self.client.batch_get_into_ptr
)
ptr_list, element_size_list = host_pool.get_page_buffer_meta(host_indices)
# page_first side pools emit exactly one (ptr, size) per component
# key; assert the invariant so any future layout change is caught
# loudly instead of silently corrupting the key<->buffer zip.
assert len(key_strs) == len(ptr_list) == len(element_size_list), (
f"UMBP v2 buffer-meta mismatch for pool {transfer.name}: "
f"keys={len(key_strs)} ptrs={len(ptr_list)} sizes={len(element_size_list)}"
)
if is_set:
# UMBP performs its own key-level deduplication, so skip the
# extra batch_exists round-trip and put directly (mirrors
# batch_set_v1).
io_results = [
bool(r)
for r in self.client.batch_put_from_ptr(
key_strs, list(ptr_list), list(element_size_list)
)
]
else:
io_results = [
bool(r)
for r in self.client.batch_get_into_ptr(
key_strs, list(ptr_list), list(element_size_list)
)
]
io_results = [bool(value) for value in operation(key_strs, ptrs, sizes)]
if len(io_results) != len(key_strs):
logger.error(
"UMBP v2 %s result-size mismatch for pool %s: "
@@ -1378,9 +1545,8 @@ class UMBPStore(HiCacheStorage):
results[transfer.name] = [False] * len(keys)
continue
# Collapse per-object results back to per-page results.
results[transfer.name] = [
all(io_results[i * key_multiplier : (i + 1) * key_multiplier])
all(io_results[i * multiplier : (i + 1) * multiplier])
for i in range(len(keys))
]
return results
python/sglang/srt/server_args.py
+1 -1
View File
@@ -2842,7 +2842,7 @@ class ServerArgs:
str,
Arg(
help="Storage backend for --enable-unified-cache-external-linker.",
choices=["mooncake"],
choices=["mooncake", "mori"],
),
NS("memory"),
] = "mooncake"