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

[AMD][DSV4] Enable hicache on deepseek-v4 fp8 unified attn (#37778)

Browse Source
This commit is contained in:
Thomas Wang
2026-09-17 11:29:23 -07:00
committed by GitHub
parent 1f60ddef5d
commit 6c73368c32
9 changed files with 246 additions and 24 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/mem_cache/deepseek_v4_memory_pool.py
+33 -21
View File
@@ -939,32 +939,21 @@ class DeepSeekV4TokenToKVPool(BaseSWAKVPool):
item_lens.append(row_bytes)
return data_ptrs, data_lens, item_lens
def unified_region_buffers(self, ratio: int) -> Tuple[List[torch.Tensor], int]:
def _unified_page_views(
self, buffers: List[torch.Tensor], ratio: int
) -> Tuple[List[torch.Tensor], int]:
# HiCache expects byte rows containing whole pages;
# the unified pool stores individual token rows after its SWA region.
assert self._unified_kv, "unified_region_buffers requires unified_kv layout"
assert ratio in (4, 128), f"unsupported compression ratio: {ratio}"
if self._unified_kv_fp8:
# item_bytes below prices kv_buffer alone, so the rope pool would never
# be offloaded and a fetched page would carry stale rope -- wrong output,
# no crash.
# TODO(danli103): give rope its own host pool, the way C4_INDEXER
# already parallels C4.
raise NotImplementedError(
"HiCache offload is not supported with "
"SGLANG_DSV4_UNIFIED_KV_FP8=1 (the host pool assumes a single "
"unified pool; the rope pool would never be offloaded)."
)
# Bf16 kv layout: [rows, 1024B]
# Fp8 kv layout: [rows, 512B] fp8 nope, [rows, 128B] bf16 rope
swa_pages = self.unified_kv_pool.swa_pages
head_dim = self.unified_kv_pool.head_dim
rows_per_page = self.page_size // ratio
stage_ratios = self.compression_ratios[self._stage_start : self._stage_end]
local_layer_ids = [i for i, r in enumerate(stage_ratios) if r == ratio]
views: List[torch.Tensor] = []
for local_layer_id in local_layer_ids:
buf = self.unified_kv_pool.kv_buffer[local_layer_id]
buf = buffers[local_layer_id]
compress_rows = buf.shape[0] - swa_pages
assert compress_rows % rows_per_page == 0, (
f"compressed rows {compress_rows} not a multiple of "
@@ -973,16 +962,39 @@ class DeepSeekV4TokenToKVPool(BaseSWAKVPool):
num_pages = compress_rows // rows_per_page
page_view = (
buf.narrow(0, swa_pages, compress_rows)
.reshape(num_pages, rows_per_page * head_dim)
.reshape(num_pages, rows_per_page * buf.shape[1])
.view(torch.uint8)
)
views.append(page_view)
item_bytes = (
rows_per_page * head_dim * self.unified_kv_pool.kv_buffer[0].element_size()
)
item_bytes = rows_per_page * buffers[0].shape[1] * buffers[0].element_size()
return views, item_bytes
def unified_region_buffers(self, ratio: int) -> Tuple[List[torch.Tensor], int]:
"""
Main compressed region of one stage: bf16 latents, or fp8 nope.
"""
assert self._unified_kv, "unified_region_buffers requires unified_kv layout"
assert ratio in (4, 128), f"unsupported compression ratio: {ratio}"
return self._unified_page_views(self.unified_kv_pool.kv_buffer, ratio)
def unified_rope_region_buffers(
self, ratio: int
) -> Optional[Tuple[List[torch.Tensor], int]]:
"""
The bf16 rope half of an fp8 two-pool row, or None when there isn't one.
A row index addresses both pools, so this mirrors exactly the rows
``unified_region_buffers`` does and only the row width differs. It needs
its own host pool: offloading the nope half alone leaves whatever rope the
row held before, which is wrong output rather than a crash.
"""
if not self._unified_kv_fp8:
return None
assert self._unified_kv, "unified_rope_region_buffers requires unified_kv"
assert ratio in (4, 128), f"unsupported compression ratio: {ratio}"
return self._unified_page_views(self.unified_kv_pool.kv_buffer_rope, ratio)
def get_state_buf_infos(self) -> Tuple[List[int], List[int], List[int]]:
data_ptrs: List[int] = []
data_lens: List[int] = []
python/sglang/srt/mem_cache/hicache_storage.py
+4
View File
@@ -72,6 +72,10 @@ class PoolName(str, Enum):
# so it needs a second pool alongside DEEPSEEK_V4_C4_INDEXER.
DEEPSEEK_V4_C4_INDEXER_SCALE = "deepseek_v4_c4_indexer_scale"
DEEPSEEK_V4_C128 = "deepseek_v4_c128"
# fp8 unified_kv splits a row across a packed fp8 nope pool and a parallel
# bf16 rope pool, so each compressed region mirrors to two host pools.
DEEPSEEK_V4_C4_ROPE = "deepseek_v4_c4_rope"
DEEPSEEK_V4_C128_ROPE = "deepseek_v4_c128_rope"
DEEPSEEK_V4_C4_STATE = "deepseek_v4_c4_state"
DEEPSEEK_V4_C4_INDEXER_STATE = "deepseek_v4_c4_indexer_state"
DEEPSEEK_V4_C128_STATE = "deepseek_v4_c128_state"
python/sglang/srt/mem_cache/hybrid_cache/hybrid_pool_assembler.py
+83
View File
@@ -578,6 +578,58 @@ def _dsv4_indexer_regions(kvcache: Any, page_size: int) -> list[_IndexerRegion]:
]
def _dsv4_rope_sibling(
kvcache: Any, ratio: int
) -> Optional[tuple[PoolName, list, int]]:
"""
``(name, device_buffers, item_bytes)`` for the bf16 rope half of an fp8
two-pool unified_kv row; None for every other layout, which keeps the whole
row in one buffer.
"""
if not getattr(kvcache, "_unified_kv", False):
return None
region = kvcache.unified_rope_region_buffers(ratio)
if region is None:
return None
buffers, item_bytes = region
name = (
PoolName.DEEPSEEK_V4_C4_ROPE if ratio == 4 else PoolName.DEEPSEEK_V4_C128_ROPE
)
return name, buffers, item_bytes
def _build_dsv4_rope_entry(
kvcache: Any,
ratio: int,
*,
device_pool: Any,
layer_mapping: dict[int, int],
num_host_pages: int,
slot_page_size: int,
transfer_layer_num: int,
) -> Optional[PoolEntry]:
sibling = _dsv4_rope_sibling(kvcache, ratio)
if sibling is None:
return None
name, device_buffers, item_bytes = sibling
return build_pool_entry(
name=name,
host_pool=DeepSeekV4PagedHostPool(
pool_name=str(name),
device_buffers=device_buffers,
item_bytes=item_bytes,
num_host_pages=num_host_pages,
slot_page_size=slot_page_size,
layout=get_memory().hicache_mem_layout,
allocator_type=_get_allocator_type(),
page_aligned_only=True,
),
device_pool=device_pool,
layer_mapping=layer_mapping,
transfer_layer_num=transfer_layer_num,
)
def build_deepseek_v4_hicache_stack(
*,
params: CacheInitParams,
@@ -697,6 +749,7 @@ def build_deepseek_v4_hicache_stack(
slot_page_size=page_size,
layout=get_memory().hicache_mem_layout,
allocator_type=_get_allocator_type(),
page_aligned_only=is_unified_kv,
)
entries.append(
build_pool_entry(
@@ -727,6 +780,19 @@ def build_deepseek_v4_hicache_stack(
)
)
# Build c4 rope buffer when using unified fp8 kv
c4_rope_entry = _build_dsv4_rope_entry(
kvcache,
4,
device_pool=kvcache.c4_kv_pool,
layer_mapping=c4_layer_mapping,
num_host_pages=num_host_pages,
slot_page_size=page_size,
transfer_layer_num=transfer_layer_num,
)
if c4_rope_entry is not None:
entries.append(c4_rope_entry)
if not is_unified_kv:
c4_state_host_pool = DeepSeekV4StateHostPool(
pool_name=str(PoolName.DEEPSEEK_V4_C4_STATE),
@@ -789,6 +855,7 @@ def build_deepseek_v4_hicache_stack(
slot_page_size=c128_slot_page_size,
layout=get_memory().hicache_mem_layout,
allocator_type=_get_allocator_type(),
page_aligned_only=is_unified_kv,
)
# C128 state pool is intentionally not registered with hicache.
# page_size=256 % 128 == 0, so state pool is not consumed on load.
@@ -817,6 +884,18 @@ def build_deepseek_v4_hicache_stack(
),
]
)
# Build c128 rope buffer when using unified fp8 kv
c128_rope_entry = _build_dsv4_rope_entry(
kvcache,
128,
device_pool=kvcache.c128_kv_pool,
layer_mapping=c128_layer_mapping,
num_host_pages=c128_num_host_pages,
slot_page_size=c128_slot_page_size,
transfer_layer_num=transfer_layer_num,
)
if c128_rope_entry is not None:
entries.append(c128_rope_entry)
host_pool_group = HostPoolGroup(entries)
cache_controller = HybridCacheController(
@@ -1416,8 +1495,11 @@ class _DeepSeekV4Strategy(StackStrategy):
)
# NPU drives C128 as an independent tree component, so adding a KV-derived
# sidecar would duplicate transfers. Add that sidecar only on GPU.
# The *_ROPE entries only resolve under unified fp8 kv; entry_map filters
# them out everywhere else.
_sidecar_srcs = [
(PoolName.DEEPSEEK_V4_C4, PoolName.KV),
(PoolName.DEEPSEEK_V4_C4_ROPE, PoolName.KV),
(PoolName.DEEPSEEK_V4_C4_INDEXER, PoolName.KV),
(PoolName.DEEPSEEK_V4_C4_INDEXER_SCALE, PoolName.KV),
(PoolName.DEEPSEEK_V4_C4_STATE, PoolName.SWA),
@@ -1426,6 +1508,7 @@ class _DeepSeekV4Strategy(StackStrategy):
]
if ComponentType.C128 not in cache.components:
_sidecar_srcs.append((PoolName.DEEPSEEK_V4_C128, PoolName.KV))
_sidecar_srcs.append((PoolName.DEEPSEEK_V4_C128_ROPE, PoolName.KV))
sidecars = [
SidecarPoolSpec(
pool_name=name,
python/sglang/srt/mem_cache/memory_pool_host.py
+4 -3
View File
@@ -194,9 +194,10 @@ class DeepSeekV4PagedHostPool(HiSparseHostPoolMixin, HostKVCache):
self.pool_name = pool_name
self.layer_num = len(device_buffers)
self.item_bytes = item_bytes
# A page row of the FP4 indexer buffers is a grouped slot layout rather
# than a flat token array, so the token-granular copy used for fused
# DSv4 C4 rows does not apply and only whole pages may move.
# The token-granular copy below addresses the fused DSv4 C4 row. Neither
# the FP4 indexer, whose page rows group their slots, nor unified_kv,
# whose page row is a flat run of tokens, has that layout, so those may
# only move whole pages.
self.page_aligned_only = page_aligned_only
self.num_host_pages = num_host_pages
self.slot_page_size = slot_page_size
python/sglang/srt/mem_cache/storage/mooncake_store/mooncake_store.py
+2
View File
@@ -815,9 +815,11 @@ class MooncakeStore(HiCacheStorage, MooncakeBaseStore):
PoolName.INDEXER,
PoolName.DRAFT_INDEXER,
PoolName.DEEPSEEK_V4_C4,
PoolName.DEEPSEEK_V4_C4_ROPE,
PoolName.DEEPSEEK_V4_C4_INDEXER,
PoolName.DEEPSEEK_V4_C4_INDEXER_SCALE,
PoolName.DEEPSEEK_V4_C128,
PoolName.DEEPSEEK_V4_C128_ROPE,
PoolName.DEEPSEEK_V4_C4_STATE,
PoolName.DEEPSEEK_V4_C4_INDEXER_STATE,
PoolName.DEEPSEEK_V4_C128_STATE,