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

[Unified Cache][AMD] Support DeepSeek-V4 unified KV in direct external linkers (#38269)

Browse Source 
Co-authored-by: amd-danli103 <danli103@amd.com>
Co-authored-by: Duyi-Wang <duyi.wang@amd.com>
Co-authored-by: TianDi101 <tiandi920722@gmail.com>
This commit is contained in:
Niko Ma
2026-09-11 01:49:29 -07:00
committed by GitHub
co-authored by amd-danli103 Duyi-Wang TianDi101
parent 0bae67648a
commit 822e73ccdd
7 changed files with 572 additions and 56 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/mem_cache/allocator/swa.py
+6 -4
View File
@@ -431,8 +431,9 @@ class SWATokenToKVPoolAllocator(BaseTokenToKVPoolAllocator):
def free_swa(self, free_index: torch.Tensor):
"""Release the SWA peers of an arbitrary slot set and clear their mapping.
Synchronizes at page_size > 1; kv-row segments go through free_swa_segment()."""
if free_index.numel() == 0:
No-op for a per-request ring, which owns no paged SWA peers. Otherwise
synchronizes at page_size > 1; kv-row segments use free_swa_segment()."""
if self._swa_req_ring or free_index.numel() == 0:
return
if self.page_size == 1:
@@ -455,8 +456,9 @@ class SWATokenToKVPoolAllocator(BaseTokenToKVPoolAllocator):
def free_swa_segment(self, free_index: torch.Tensor, *, start_pos: int):
"""free_swa() for a kv-row segment; same start-alignment contract as
free_segment(), and fixed-shape at every page size."""
if free_index.numel() == 0:
free_segment(), and fixed-shape at every page size. No-op for a
per-request ring, as in free_swa()."""
if self._swa_req_ring or free_index.numel() == 0:
return
self._free_swa_pages(free_index, start_pos=start_pos)
python/sglang/srt/mem_cache/hybrid_cache/linker_pool_assembler.py
+46 -42
View File
@@ -239,34 +239,34 @@ def _build_deepseek_v4_device_pool_group(
) -> DevicePoolGroup:
from sglang.srt.mem_cache.deepseek_v4_memory_pool import HiSparseC4DevicePool
from sglang.srt.mem_cache.hybrid_cache.hybrid_pool_assembler import (
_dsv4_compressed_region_buffers,
_dsv4_indexer_regions,
_resolve_deepseek_v4_layer_mappings,
)
mappings = _resolve_deepseek_v4_layer_mappings(kvcache)
if getattr(kvcache, "_unified_kv", False) or isinstance(
kvcache.c4_kv_pool, HiSparseC4DevicePool
):
raise ValueError(
"The direct external linker does not support unified-KV or HiSparse."
)
if kvcache.swa_page_size != page_size:
raise ValueError(
"DeepSeek V4 SWA page size must match the tree page size: "
f"{kvcache.swa_page_size} != {page_size}."
)
if isinstance(kvcache.c4_kv_pool, HiSparseC4DevicePool):
raise ValueError("The direct external linker does not support HiSparse.")
entries = [
DevicePoolEntry(
name=PoolName.SWA,
indices_from_pool=PoolName.SWA,
device_pool=kvcache.swa_kv_pool,
components=[kvcache.swa_kv_pool.kv_buffer],
layer_mapping=mappings.swa,
page_size=page_size,
rows_are_pages=True,
mappings = _resolve_deepseek_v4_layer_mappings(kvcache)
is_unified_kv = getattr(kvcache, "_unified_kv", False)
entries = []
if not is_unified_kv:
if kvcache.swa_page_size != page_size:
raise ValueError(
"DeepSeek V4 SWA page size must match the tree page size: "
f"{kvcache.swa_page_size} != {page_size}."
)
entries.append(
DevicePoolEntry(
name=PoolName.SWA,
indices_from_pool=PoolName.SWA,
device_pool=kvcache.swa_kv_pool,
components=[kvcache.swa_kv_pool.kv_buffer],
layer_mapping=mappings.swa,
page_size=page_size,
rows_are_pages=True,
)
)
]
def add(name, source, pool, buffers, layer_mapping):
if layer_mapping:
@@ -282,11 +282,14 @@ def _build_deepseek_v4_device_pool_group(
)
)
c4_buffers, _ = _dsv4_compressed_region_buffers(kvcache, 4)
c128_buffers, _ = _dsv4_compressed_region_buffers(kvcache, 128)
add(
PoolName.DEEPSEEK_V4_C4,
PoolName.KV,
kvcache.c4_kv_pool,
kvcache.c4_kv_pool.kv_buffer,
c4_buffers,
mappings.c4,
)
for region in _dsv4_indexer_regions(kvcache, page_size):
@@ -301,29 +304,30 @@ def _build_deepseek_v4_device_pool_group(
PoolName.DEEPSEEK_V4_C128,
PoolName.KV,
kvcache.c128_kv_pool,
kvcache.c128_kv_pool.kv_buffer,
c128_buffers,
mappings.c128,
)
add(
PoolName.DEEPSEEK_V4_C4_STATE,
PoolName.SWA,
kvcache.compress_state_pools,
_deepseek_v4_state_views(
if not is_unified_kv:
add(
PoolName.DEEPSEEK_V4_C4_STATE,
PoolName.SWA,
kvcache.compress_state_pools,
mappings.c4_state_global_layers,
),
mappings.c4_state,
)
add(
PoolName.DEEPSEEK_V4_C4_INDEXER_STATE,
PoolName.SWA,
kvcache.indexer_compress_state_pools,
_deepseek_v4_state_views(
_deepseek_v4_state_views(
kvcache.compress_state_pools,
mappings.c4_state_global_layers,
),
mappings.c4_state,
)
add(
PoolName.DEEPSEEK_V4_C4_INDEXER_STATE,
PoolName.SWA,
kvcache.indexer_compress_state_pools,
mappings.c4_state_global_layers,
),
mappings.c4_state,
)
_deepseek_v4_state_views(
kvcache.indexer_compress_state_pools,
mappings.c4_state_global_layers,
),
mappings.c4_state,
)
return DevicePoolGroup(
entries,
mappings.transfer_layer_num,
python/sglang/srt/mem_cache/unified_cache/components/swa_component.py
+5 -6
View File
@@ -5,6 +5,7 @@ from typing import TYPE_CHECKING, Callable, Optional, Sequence
import torch
from sglang.srt.environ import envs
from sglang.srt.mem_cache.allocator.swa import is_swa_req_ring
from sglang.srt.mem_cache.base_prefix_cache import (
DecLockRefParams,
EvictParams,
@@ -308,11 +309,9 @@ class SWAComponent(TreeComponent):
ct = self.component_type
state = {"len": float("inf")}
# unified_kv never caches the SWA ring (per-request, not content-stable),
# so SWA bookkeeping must not gate the match here.
swa_device_only_hicache = (
not self.tree_core.has_swa_host_pool and self.tree_core.enable_hicache
)
# A per-request SWA ring is not stored in tree nodes, so its bookkeeping
# must not gate prefix matching.
swa_req_ring = is_swa_req_ring(self.cache.token_to_kv_pool_allocator)
def validator(node: UnifiedTreeNode) -> bool:
cd = node.component_data[ct]
@@ -320,7 +319,7 @@ class SWAComponent(TreeComponent):
# — load_back will restore SWA from host before use.
if cd.value is None and (match_device_only or cd.host_value is None):
state["len"] = 0
if swa_device_only_hicache and (node.backuped or not node.evicted):
if swa_req_ring and (node.backuped or not node.evicted):
return True
return False
state["len"] += len(node.key)
python/sglang/srt/mem_cache/unified_cache/unified_cache_linker.py
+28 -2
View File
@@ -25,6 +25,7 @@ from typing import TYPE_CHECKING, NamedTuple
import torch
from sglang.srt.mem_cache.allocator.swa import is_swa_req_ring
from sglang.srt.mem_cache.base_prefix_cache import (
DecLockRefParams,
InsertParams,
@@ -160,6 +161,15 @@ class UnifiedCacheLinkerWrapper:
self.cache = cache
self.cache_linker = cache_linker
swa = cache.components.get(ComponentType.SWA)
self._skip_swa = swa is not None and is_swa_req_ring(
cache.token_to_kv_pool_allocator
)
self._components = tuple(
component
for component in cache._components_tuple
if not (self._skip_swa and component is swa)
)
# rid -> what match found, consumed by the next init_load_back.
self.hit_markers: dict[str, ExternalCacheHitMarker] = {}
# Loads in flight, each pinning its inserted endpoint until DMA completes.
@@ -192,7 +202,7 @@ class UnifiedCacheLinkerWrapper:
return result
lookup_transfers = []
for component in cache._components_tuple:
for component in self._components:
transfer = component.build_external_linker_transfer(
LinkerTransferPhase.LOOKUP, None, tail_hashes
)
@@ -290,7 +300,7 @@ class UnifiedCacheLinkerWrapper:
# Build per-component linker transfers.
component_transfers: list[tuple[TreeComponent, PoolTransfer]] = []
for component in cache._components_tuple:
for component in self._components:
transfer = component.build_external_linker_transfer(
LinkerTransferPhase.LOAD, None, tail_hashes
)
@@ -313,6 +323,20 @@ class UnifiedCacheLinkerWrapper:
prefix_len,
)
# Components omitted from the linker do not run their PREPARE hook.
# Keep a non-restorable SWA range as tombstones instead of rebuilding
# it from an uninitialized FULL-to-SWA mapping during cache.insert().
if self._skip_swa:
if req.kv is None:
from sglang.srt.managers.schedule_batch import ReqKvInfo
req.kv = ReqKvInfo(
kv_allocated_len=prefix_len,
swa_evicted_seqlen=prefix_len,
)
else:
req.kv.swa_evicted_seqlen = max(req.kv.swa_evicted_seqlen, prefix_len)
# Insert the newly loaded tail into the tree.
prefix_indices = torch.cat(
[req.prefix_indices.to(torch.int64), full_transfer.device_indices]
@@ -484,6 +508,8 @@ class UnifiedCacheLinkerWrapper:
node_id
)
if transfers is not None:
if self._skip_swa:
transfers = [t for t in transfers if t.name != PoolName.SWA]
self._offload_node(node_id, transfers)
def _offload_node(self, node_id: NodeId, transfers: list[PoolTransfer]) -> None: