diff --git a/python/sglang/srt/mem_cache/deepseek_v4_memory_pool.py b/python/sglang/srt/mem_cache/deepseek_v4_memory_pool.py index deaeb52a9..a43748b11 100644 --- a/python/sglang/srt/mem_cache/deepseek_v4_memory_pool.py +++ b/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" - 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] = [] diff --git a/python/sglang/srt/mem_cache/hicache_storage.py b/python/sglang/srt/mem_cache/hicache_storage.py index 4603cf219..2b41258fd 100644 --- a/python/sglang/srt/mem_cache/hicache_storage.py +++ b/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" diff --git a/python/sglang/srt/mem_cache/hybrid_cache/hybrid_pool_assembler.py b/python/sglang/srt/mem_cache/hybrid_cache/hybrid_pool_assembler.py index c10f11320..5acc1cd69 100644 --- a/python/sglang/srt/mem_cache/hybrid_cache/hybrid_pool_assembler.py +++ b/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, diff --git a/python/sglang/srt/mem_cache/memory_pool_host.py b/python/sglang/srt/mem_cache/memory_pool_host.py index c701dc16d..9cae60cfc 100644 --- a/python/sglang/srt/mem_cache/memory_pool_host.py +++ b/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 - # 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 diff --git a/python/sglang/srt/mem_cache/storage/mooncake_store/mooncake_store.py b/python/sglang/srt/mem_cache/storage/mooncake_store/mooncake_store.py index 2146298d0..e49c88e9a 100644 --- a/python/sglang/srt/mem_cache/storage/mooncake_store/mooncake_store.py +++ b/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, diff --git a/rust/sglang-radix-tree/src/python_bindings.rs b/rust/sglang-radix-tree/src/python_bindings.rs index 666538fa6..7a871b0ee 100644 --- a/rust/sglang-radix-tree/src/python_bindings.rs +++ b/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 { "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), diff --git a/rust/sglang-radix-tree/src/unified_tree_core.rs b/rust/sglang-radix-tree/src/unified_tree_core.rs index 02f7db8dd..2dadbcc4e 100644 --- a/rust/sglang-radix-tree/src/unified_tree_core.rs +++ b/rust/sglang-radix-tree/src/unified_tree_core.rs @@ -348,6 +348,8 @@ pub enum PoolName { DeepseekV4C2, DeepseekV4C2Indexer, DeepseekV4C2IndexerScale, + DeepseekV4C4Rope, + DeepseekV4C128Rope, DeepseekV4C4State, DeepseekV4C4IndexerState, DeepseekV4C128State, diff --git a/test/registered/unit/mem_cache/test_dsv4_hicache_l2.py b/test/registered/unit/mem_cache/test_dsv4_hicache_l2.py index 9046a6941..bff743331 100644 --- a/test/registered/unit/mem_cache/test_dsv4_hicache_l2.py +++ b/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, diff --git a/test/registered/unit/mem_cache/test_dsv4_unified_fp8_pool.py b/test/registered/unit/mem_cache/test_dsv4_unified_fp8_pool.py index f9ddfb49f..19d51cacb 100644 --- a/test/registered/unit/mem_cache/test_dsv4_unified_fp8_pool.py +++ b/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()