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,
rust/sglang-radix-tree/src/python_bindings.rs
+4
View File
@@ -247,6 +247,8 @@ fn pool_name_str(name: PoolName) -> &'static str {
PoolName::DeepseekV4C2 => "deepseek_v4_c2",
PoolName::DeepseekV4C2Indexer => "deepseek_v4_c2_indexer",
PoolName::DeepseekV4C2IndexerScale => "deepseek_v4_c2_indexer_scale",
PoolName::DeepseekV4C4Rope => "deepseek_v4_c4_rope",
PoolName::DeepseekV4C128Rope => "deepseek_v4_c128_rope",
PoolName::DeepseekV4C4State => "deepseek_v4_c4_state",
PoolName::DeepseekV4C4IndexerState => "deepseek_v4_c4_indexer_state",
PoolName::DeepseekV4C128State => "deepseek_v4_c128_state",
@@ -273,6 +275,8 @@ fn parse_pool_name(name: &str) -> PyResult<PoolName> {
"deepseek_v4_c2" => Ok(PoolName::DeepseekV4C2),
"deepseek_v4_c2_indexer" => Ok(PoolName::DeepseekV4C2Indexer),
"deepseek_v4_c2_indexer_scale" => Ok(PoolName::DeepseekV4C2IndexerScale),
"deepseek_v4_c4_rope" => Ok(PoolName::DeepseekV4C4Rope),
"deepseek_v4_c128_rope" => Ok(PoolName::DeepseekV4C128Rope),
"deepseek_v4_c4_state" => Ok(PoolName::DeepseekV4C4State),
"deepseek_v4_c4_indexer_state" => Ok(PoolName::DeepseekV4C4IndexerState),
"deepseek_v4_c128_state" => Ok(PoolName::DeepseekV4C128State),
rust/sglang-radix-tree/src/unified_tree_core.rs
+2
View File
@@ -348,6 +348,8 @@ pub enum PoolName {
DeepseekV4C2,
DeepseekV4C2Indexer,
DeepseekV4C2IndexerScale,
DeepseekV4C4Rope,
DeepseekV4C128Rope,
DeepseekV4C4State,
DeepseekV4C4IndexerState,
DeepseekV4C128State,
test/registered/unit/mem_cache/test_dsv4_hicache_l2.py
+1
View File
@@ -242,6 +242,7 @@ class TestDSV4PoolAssembly(CustomTestCase):
c4_kv_pool=SimpleNamespace(kernel_page_size=32),
c4_indexer_kv_pool=indexer_pool,
unified_region_buffers=lambda ratio: (c4_buffer, 256),
unified_rope_region_buffers=lambda ratio: None,
)
params = SimpleNamespace(
page_size=128,
test/registered/unit/mem_cache/test_dsv4_unified_fp8_pool.py
+113
View File
@@ -6,6 +6,7 @@ import torch
from sglang.srt.mem_cache.deepseek_v4_memory_pool import (
DSV4_FP8_NOPE_ROW_BYTES,
DSV4_FP8_QUANT_TILE,
DeepSeekV4TokenToKVPool,
DeepSeekV4UnifiedKVPool,
dsv4_unified_row_bytes,
)
@@ -124,5 +125,117 @@ class TestDSV4UnifiedFp8PoolAllocation(CustomTestCase):
self._pool(fp8=False).get_unified_kv_rope(0)
class _StubTokenToKVPool:
"""
Only the attributes the region math reads, so it can be exercised without standing
up a whole DeepSeekV4TokenToKVPool.
"""
_unified_page_views = DeepSeekV4TokenToKVPool._unified_page_views
unified_region_buffers = DeepSeekV4TokenToKVPool.unified_region_buffers
unified_rope_region_buffers = DeepSeekV4TokenToKVPool.unified_rope_region_buffers
def __init__(self, unified_kv_pool, page_size, stage_ratios, fp8):
self.unified_kv_pool = unified_kv_pool
self.page_size = page_size
self.compression_ratios = list(stage_ratios)
self._stage_start = 0
self._stage_end = len(stage_ratios)
self._unified_kv = True
self._unified_kv_fp8 = fp8
class TestDSV4UnifiedRegionBuffers(CustomTestCase):
"""
HiCache mirrors the compressed region of every unified_kv layer. Under fp8
that region lives in two pools, and offloading only the nope half refills a
fetched page with stale rope (wrong output, no crash), so the pairing is what
these tests pin.
"""
# Two c4 layers to cover the per-ratio layer filter, one c128 layer. Blocks
# must be even: a c4 layer holds `num_blocks * 32` compressed rows plus one
# padding page of 64, which only divides into whole pages when it is.
STAGE_RATIOS = [4, 4, 128]
NUM_SLOTS = 3
NUM_BLOCKS = 6
PAGE_SIZE = 256
SWA_RING = 8
NUM_PAGES = 4
def _stub(self, fp8):
pool = DeepSeekV4UnifiedKVPool(
stage_ratios=self.STAGE_RATIOS,
num_slots=self.NUM_SLOTS,
num_blocks=self.NUM_BLOCKS,
page_size=self.PAGE_SIZE,
qk_nope_head_dim=NOPE_DIM,
qk_rope_head_dim=ROPE_DIM,
device="cpu",
memory_saver_adapter=_StubMemorySaver(),
custom_mem_pool=None,
swa_ring_size=self.SWA_RING,
fp8=fp8,
)
return pool, _StubTokenToKVPool(pool, self.PAGE_SIZE, self.STAGE_RATIOS, fp8)
def test_bf16_keeps_the_whole_row_in_one_region(self):
"""
The bf16 arm must not grow a second host pool: its row is undivided,
and a stray rope region would mirror a buffer that does not exist.
"""
_, stub = self._stub(fp8=False)
for ratio, num_layers in ((4, 2), (128, 1)):
views, item_bytes = stub.unified_region_buffers(ratio)
rows_per_page = self.PAGE_SIZE // ratio
self.assertEqual(len(views), num_layers)
self.assertEqual(item_bytes, rows_per_page * (NOPE_DIM + ROPE_DIM) * 2)
self.assertIsNone(stub.unified_rope_region_buffers(ratio))
def test_fp8_pairs_a_rope_region_with_the_nope_one(self):
_, stub = self._stub(fp8=True)
for ratio, num_layers in ((4, 2), (128, 1)):
rows_per_page = self.PAGE_SIZE // ratio
for (views, item_bytes), row_bytes in (
(stub.unified_region_buffers(ratio), DSV4_FP8_NOPE_ROW_BYTES),
(stub.unified_rope_region_buffers(ratio), ROPE_DIM * 2),
):
self.assertEqual(len(views), num_layers)
self.assertEqual(item_bytes, rows_per_page * row_bytes)
for view in views:
self.assertEqual(view.dtype, torch.uint8)
self.assertEqual(list(view.shape), [self.NUM_PAGES, item_bytes])
def test_regions_start_past_the_swa_ring(self):
"""
The host pool takes each view's data_ptr as the device base, so a view
that still covered the ring would offload ring rows as page 0.
"""
pool, stub = self._stub(fp8=True)
swa_pages = pool.swa_pages
for (views, _), device_buffers in (
(stub.unified_region_buffers(4), pool.kv_buffer),
(stub.unified_rope_region_buffers(4), pool.kv_buffer_rope),
):
for view, buf in zip(views, device_buffers):
row_bytes = buf.shape[1] * buf.element_size()
self.assertEqual(
view.data_ptr(), buf.data_ptr() + swa_pages * row_bytes
)
def test_fp8_regions_cover_the_same_bytes_as_bf16(self):
"""
Both halves together must mirror the whole row; the 0.625x is the same
saving the row-width test pins.
"""
_, bf16 = self._stub(fp8=False)
_, fp8 = self._stub(fp8=True)
for ratio in (4, 128):
_, whole = bf16.unified_region_buffers(ratio)
_, nope = fp8.unified_region_buffers(ratio)
_, rope = fp8.unified_rope_region_buffers(ratio)
self.assertAlmostEqual((nope + rope) / whole, 0.625)
if __name__ == "__main__":
unittest.main()