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

fix(hicache/umbp): support DeepSeek-V4 hybrid HostPoolGroup (multi-po… (#30762)

Browse Source 
Co-authored-by: Zhangheng <hzh0425@apache.org>
This commit is contained in:
AMD-yanfeiwang
2026-08-13 15:40:07 +08:00
committed by GitHub
co-authored by Zhangheng
parent b7f87a2513
commit a34f81251f
6 changed files with 706 additions and 30 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/mem_cache/memory_pool_host.py
+4
View File
@@ -671,6 +671,10 @@ class LogicalHostPool:
def clear(self):
self.free_slots = torch.arange(self.size, dtype=torch.int64)
def destroy(self) -> None:
"""Logical anchors own no backing buffers or registrations to release."""
return None
def available_size(self):
return len(self.free_slots)
python/sglang/srt/mem_cache/storage/umbp/umbp_store.py
+279 -17
View File
@@ -18,6 +18,10 @@ from sglang.srt.mem_cache.hicache_storage import (
HiCacheStorage,
HiCacheStorageConfig,
HiCacheStorageExtraInfo,
PoolHitPolicy,
PoolName,
PoolTransfer,
PoolTransferResult,
)
from sglang.srt.mem_cache.memory_pool_host import HostKVCache
@@ -573,6 +577,15 @@ class UMBPStore(HiCacheStorage):
# .default_dram_page_size (2 MiB by default). The
# partial-tail safety net in PoolClient handles any
# size mismatch.
#
# Logical-anchor host pools (the DeepSeek-V4 HiCache HostPoolGroup
# whose KV anchor is a LogicalHostPool that owns only page indices
# and no physical KV tensor) return None from get_page_buffer_meta()
# by design — the real per-page byte sizes live in the v2 side pools
# (SWA / compressed KV / indexer / state), which each carry their own
# dimensions. There is no single page size that fits all of them, so
# we leave dram_page_size at 0 and let the mori master use its
# default with the PoolClient partial-tail safety net.
page_byte_size = None
if "dram_page_size" in extra:
page_byte_size = int(extra["dram_page_size"])
@@ -584,14 +597,17 @@ class UMBPStore(HiCacheStorage):
# would over-count by the indexer buffer that is never put to UMBP).
dummy = torch.zeros(mem_pool_host.page_size, dtype=torch.int64)
if self.is_mla_backend:
_, esz = mem_pool_host.get_page_buffer_meta(dummy)
meta = mem_pool_host.get_page_buffer_meta(dummy)
elif storage_config is not None and getattr(
storage_config, "should_split_heads", False
):
sf = storage_config.tp_lcm_size // storage_config.tp_size
_, esz = mem_pool_host.get_split_heads_page_buffer_meta(dummy, sf)
meta = mem_pool_host.get_split_heads_page_buffer_meta(dummy, sf)
else:
_, esz = mem_pool_host.get_page_buffer_meta(dummy)
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.
esz = meta[1] if meta else None
page_byte_size = int(esz[0]) if esz else 0
if (
@@ -787,6 +803,12 @@ 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.
self.registered_pools: dict = {}
self._kv_anchor_is_logical = False
self.client = UMBPClient(cfg)
if mem_pool_host is not None:
self.register_mem_pool_host(mem_pool_host)
@@ -867,23 +889,44 @@ 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().
self._kv_anchor_is_logical = self.mem_pool_host.kv_buffer is None
self._zero_copy_registered = False
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.
self._zero_copy_registered = False
if self._register_host_buffer_for_zero_copy(mem_pool_host):
self._zero_copy_registered = True
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).
"""
if self.client is None:
return
return False
try:
is_distributed = bool(self.client.is_distributed())
except Exception:
is_distributed = False
if not is_distributed:
return
return False
if not hasattr(self.client, "register_memory"):
return
return False
if getattr(self, "_disable_zero_copy_register", False):
logger.info(
"UMBPStore: skipping host KV buffer RDMA registration because "
@@ -891,9 +934,20 @@ class UMBPStore(HiCacheStorage):
"Falling back to the staging-buffer transfer path; per-transfer "
"size is capped by distributed.staging_buffer_size."
)
return
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:
return False
try:
kv_buffer = mem_pool_host.kv_buffer
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
@@ -901,7 +955,7 @@ class UMBPStore(HiCacheStorage):
# 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(mem_pool_host, "allocator", None)
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)
@@ -917,20 +971,37 @@ class UMBPStore(HiCacheStorage):
"distributed.staging_buffer_size.",
exc,
)
return
return False
if ok:
self._zero_copy_registered = True
logger.info(
"UMBPStore: registered host KV buffer for RDMA zero-copy "
"(ptr=0x%x, size=%d MB)",
host_ptr,
host_size // (1024 * 1024),
)
else:
logger.warning(
"UMBPStore: register_memory returned false; staying on staging "
"buffer fallback path."
)
return True
logger.warning(
"UMBPStore: register_memory returned false; staying on staging "
"buffer fallback path."
)
return False
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
self._register_host_buffer_for_zero_copy(host_pool)
# ------------------------------------------------------------------
# Key suffix generation — mirrors MooncakeStore
@@ -1007,6 +1078,11 @@ class UMBPStore(HiCacheStorage):
host_indices: torch.Tensor,
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)
# Normalize sizes to list of per-key sizes
@@ -1076,6 +1152,11 @@ class UMBPStore(HiCacheStorage):
page_count = len(host_indices) // self.mem_pool_host.page_size
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)
if isinstance(buffer_sizes, int):
@@ -1138,6 +1219,187 @@ class UMBPStore(HiCacheStorage):
hit_count = self.client.batch_exists_consecutive(query_keys)
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).
# ------------------------------------------------------------------
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.
"""
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}"]
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",
]
else:
suffixes = [f"_{self.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
def batch_exists_v2(
self,
keys: List[str],
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)
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(
keys[:final_pages], transfer
)
exists = list(self.client.batch_exists(component_keys))
if len(exists) != len(component_keys):
logger.error(
"UMBP v2 batch_exists result-size mismatch for pool %s: "
"expected=%d actual=%d; treating the storage prefix as a miss",
transfer.name,
len(component_keys),
len(exists),
)
final_pages = 0
break
# Collapse per-object results into per-page presence.
page_exists = [
all(exists[i * key_multiplier : (i + 1) * key_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
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):
if all(
page_exists[i]
for i in range(max(0, prefix_len - trailing), prefix_len)
):
boundary = prefix_len
break
if boundary:
hit_count[transfer.name] = boundary
final_pages = min(final_pages, boundary)
return PoolTransferResult(final_pages, hit_count)
def _batch_io_v2(self, transfers: List[PoolTransfer], is_set: bool) -> dict:
"""Unified per-pool zero-copy I/O. Returns {pool_name: per-page bools}."""
results: dict = {}
for transfer in transfers:
host_pool = self.registered_pools.get(transfer.name)
if host_pool is None:
raise ValueError(f"Unregistered UMBP hybrid pool: {transfer.name}")
keys = transfer.keys or []
host_indices = transfer.host_indices
page_size = getattr(host_pool, "page_size", 1) or 1
if not keys or host_indices is None:
results[transfer.name] = [False] * len(keys)
continue
assert len(keys) == len(host_indices) // page_size
key_strs, key_multiplier = self._get_hybrid_page_component_keys(
keys, transfer
)
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)
)
]
if len(io_results) != len(key_strs):
logger.error(
"UMBP v2 %s result-size mismatch for pool %s: "
"expected=%d actual=%d; treating every page as failed",
"set" if is_set else "get",
transfer.name,
len(key_strs),
len(io_results),
)
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])
for i in range(len(keys))
]
return results
def batch_get_v2(
self,
transfers: List[PoolTransfer],
extra_info: Optional[HiCacheStorageExtraInfo] = None,
) -> dict:
return self._batch_io_v2(transfers, is_set=False)
def batch_set_v2(
self,
transfers: List[PoolTransfer],
extra_info: Optional[HiCacheStorageExtraInfo] = None,
) -> dict:
return self._batch_io_v2(transfers, is_set=True)
# ------------------------------------------------------------------
# Legacy ABC interface (required by HiCacheStorage)
# ------------------------------------------------------------------