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

python/sglang/srt/mem_cache/deepseek_v4_memory_pool.py

@@ -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"
+        # Bf16 kv layout: [rows, 1024B]
-        assert ratio in (4, 128), f"unsupported compression ratio: {ratio}"
+        # Fp8 kv layout:  [rows, 512B] fp8 nope, [rows, 128B] bf16 rope
-        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)."
-            )
-
         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 = (
+        item_bytes = rows_per_page * buffers[0].shape[1] * buffers[0].element_size()
-            rows_per_page * head_dim * self.unified_kv_pool.kv_buffer[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

@@ -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

@@ -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

@@ -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
+        # The token-granular copy below addresses the fused DSv4 C4 row. Neither
-        # than a flat token array, so the token-granular copy used for fused
+        # the FP4 indexer, whose page rows group their slots, nor unified_kv,
-        # DSv4 C4 rows does not apply and only whole pages may move.
+        # 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

@@ -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

@@ -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

@@ -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

@@ -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

@@ -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()