From 6c91590e1bab0a0275feeb6d32ef92f501625cac Mon Sep 17 00:00:00 2001 From: shuwenn <47200617+alphabetc1@users.noreply.github.com> Date: Sat, 21 Mar 2026 09:38:44 +0800 Subject: [PATCH] [HiCache] refactor: hicache normalization flow and compatibility checks (#19669) --- .../sglang/srt/managers/cache_controller.py | 4 + python/sglang/srt/mem_cache/hiradix_cache.py | 7 - python/sglang/srt/server_args.py | 125 ++++++++++++------ .../unit/server_args/test_server_args.py | 94 +++++++++++++ 4 files changed, 184 insertions(+), 46 deletions(-) diff --git a/python/sglang/srt/managers/cache_controller.py b/python/sglang/srt/managers/cache_controller.py index 0fedbcf3d..e4f8d0b70 100644 --- a/python/sglang/srt/managers/cache_controller.py +++ b/python/sglang/srt/managers/cache_controller.py @@ -728,6 +728,10 @@ class HiCacheController: return host_indices, device_indices.index_select(0, idx) elif self.mem_pool_host.layout == "page_first_direct": return host_indices, device_indices.cpu() + else: + raise ValueError( + f"Unsupported layout {self.mem_pool_host.layout!r} for io backend 'direct'" + ) elif self.io_backend == "kernel_ascend": return host_indices, device_indices.cpu() else: diff --git a/python/sglang/srt/mem_cache/hiradix_cache.py b/python/sglang/srt/mem_cache/hiradix_cache.py index e13d43c37..2d3e579bb 100644 --- a/python/sglang/srt/mem_cache/hiradix_cache.py +++ b/python/sglang/srt/mem_cache/hiradix_cache.py @@ -58,13 +58,6 @@ class HiRadixCache(RadixCache): def __init__(self, params: CacheInitParams, server_args: ServerArgs): self._enable_metrics_flag = params.enable_metrics - if server_args.hicache_io_backend == "direct": - # FIXME: move this logic into server_args parsing - if server_args.hicache_mem_layout == "page_first": - server_args.hicache_mem_layout = "page_first_direct" - logger.warning( - "Page first layout is not supported with direct IO backend, switching to page first direct layout" - ) if not server_args.disable_hicache_numa_detect: bind_to_closest_numa_node_cuda() diff --git a/python/sglang/srt/server_args.py b/python/sglang/srt/server_args.py index dcf5c5683..11e2fcd6b 100644 --- a/python/sglang/srt/server_args.py +++ b/python/sglang/srt/server_args.py @@ -2786,56 +2786,103 @@ class ServerArgs: ) def _handle_hicache(self): + """Normalize hicache-related knobs into a valid runtime configuration. + + Resolution order: + 1) Layout <-> I/O compatibility for direct conflicts. + 2) Storage <-> layout compatibility (may rewrite layout). + 3) I/O <-> decode-attention compatibility (may rewrite I/O or decode backend). + 4) Re-run step (1) if step (3) changed I/O backend. + """ + # Skip all normalization when neither hicache nor decode-offload path is active. + if not ( + self.enable_hierarchical_cache + or self.disaggregation_decode_enable_offload_kvcache + ): + return + + # Step 1: Initial layout-io compatibility normalization. + self._resolve_layout_io_compatibility() + + # Step 2: Storage-layout normalization without changing io backend. + self._resolve_storage_layout_compatibility() + + # Step 3: IO-decode backend compatibility (may change io backend). + io_changed = self._resolve_io_decode_attention_compatibility() + + # Step 4: Re-normalize layout after io backend changes. + if io_changed: + self._resolve_layout_io_compatibility() + + def _resolve_layout_io_compatibility(self): if ( self.hicache_mem_layout == "page_first_direct" and self.hicache_io_backend == "kernel" ): self.hicache_io_backend = "direct" logger.warning( - "Kernel io backend does not support page first direct layout" + "Kernel io backend does not support page first direct layout, switching to direct io backend" ) if ( - self.enable_hierarchical_cache - or self.disaggregation_decode_enable_offload_kvcache - ) and self.hicache_io_backend == "kernel": - # fix for the compatibility issue with FlashAttention3 decoding and HiCache kernel backend - # Only override when the *effective* decode backend would be FA3. - # Otherwise, respect the user's chosen attention backend (e.g., aiter on ROCm). - effective_decode_backend = ( - self.decode_attention_backend - if self.decode_attention_backend is not None - else self.attention_backend + self.hicache_mem_layout == "page_first" + and self.hicache_io_backend == "direct" + ): + self.hicache_mem_layout = "page_first_direct" + logger.warning( + "Page first layout is not supported with direct IO backend, switching to page first direct layout" ) - if effective_decode_backend == "fa3": - if self.decode_attention_backend is None: - # If decode backend wasn't explicitly set, pick a safe default that works with HiCache kernel IO. - if not self.use_mla_backend(): - self.decode_attention_backend = ( - "flashinfer" if is_flashinfer_available() else "triton" - ) - else: - self.decode_attention_backend = ( - "flashinfer" if is_sm100_supported() else "triton" - ) - else: - # If user explicitly requested FA3 decode, fall back to direct IO. - self.hicache_io_backend = "direct" - logger.warning( - "FlashAttention3 decode backend is not compatible with hierarchical cache. " - "Setting hicache_io_backend to vanilla I/O, which may lead to suboptimal performance with small page sizes." - ) - if self.hicache_storage_backend == "mooncake": - if self.hicache_mem_layout == "layer_first": - if self.hicache_io_backend == "direct": - self.hicache_mem_layout = "page_first_direct" - elif self.hicache_io_backend == "kernel": - self.hicache_mem_layout = "page_first" - logger.warning( - f"Mooncake storage backend does not support layer_first layout, " - f"switching to {self.hicache_mem_layout} layout for {self.hicache_io_backend} io backend" - ) + def _resolve_storage_layout_compatibility(self): + if ( + self.hicache_storage_backend != "mooncake" + or self.hicache_mem_layout != "layer_first" + ): + return + + if self.hicache_io_backend == "direct": + new_layout = "page_first_direct" + elif self.hicache_io_backend == "kernel": + new_layout = "page_first" + else: + # Keep current behavior for unknown backends (e.g., kernel_ascend). + new_layout = self.hicache_mem_layout + + self.hicache_mem_layout = new_layout + logger.warning( + f"Mooncake storage backend does not support layer_first layout, " + f"switching to {new_layout} layout for {self.hicache_io_backend} io backend" + ) + + def _resolve_io_decode_attention_compatibility(self) -> bool: + if self.hicache_io_backend != "kernel": + return False + + # Only patch settings when the effective decode backend is FA3. + effective_decode_backend = ( + self.decode_attention_backend or self.attention_backend + ) + if effective_decode_backend != "fa3": + return False + + if self.decode_attention_backend is not None: + self.hicache_io_backend = "direct" + logger.warning( + "FlashAttention3 decode backend is not compatible with hierarchical cache. " + "Setting hicache_io_backend to vanilla I/O, which may lead to suboptimal performance with small page sizes." + ) + return True + + # If decode backend is implicit, pick a safe backend without changing io backend. + if not self.use_mla_backend(): + self.decode_attention_backend = ( + "flashinfer" if is_flashinfer_available() else "triton" + ) + else: + self.decode_attention_backend = ( + "flashinfer" if is_sm100_supported() else "triton" + ) + return False def _handle_speculative_decoding(self): if ( diff --git a/test/registered/unit/server_args/test_server_args.py b/test/registered/unit/server_args/test_server_args.py index 41fd5c04c..0b4cdda3f 100644 --- a/test/registered/unit/server_args/test_server_args.py +++ b/test/registered/unit/server_args/test_server_args.py @@ -332,5 +332,99 @@ class TestSSLArgs(unittest.TestCase): self.assertTrue(server_args.enable_ssl_refresh) +class TestHiCacheArgs(unittest.TestCase): + def _make_args(self, **overrides) -> ServerArgs: + args = ServerArgs(model_path="dummy") + for key, value in overrides.items(): + setattr(args, key, value) + return args + + def _assert_hicache_fields( + self, + args: ServerArgs, + *, + expected_io_backend: str, + expected_mem_layout: str, + expected_decode_backend: str | None = None, + ): + self.assertEqual(args.hicache_io_backend, expected_io_backend) + self.assertEqual(args.hicache_mem_layout, expected_mem_layout) + if expected_decode_backend is not None: + self.assertEqual(args.decode_attention_backend, expected_decode_backend) + + def test_hicache_io_backend_and_mem_layout_compatibility(self): + cases = [ + { + "name": "kernel_with_page_first_direct", + "overrides": { + "enable_hierarchical_cache": True, + "hicache_io_backend": "kernel", + "hicache_mem_layout": "page_first_direct", + }, + "expected_io_backend": "direct", + "expected_mem_layout": "page_first_direct", + }, + { + "name": "direct_with_page_first", + "overrides": { + "enable_hierarchical_cache": True, + "hicache_io_backend": "direct", + "hicache_mem_layout": "page_first", + }, + "expected_io_backend": "direct", + "expected_mem_layout": "page_first_direct", + }, + { + "name": "mooncake_with_layer_first", + "overrides": { + "enable_hierarchical_cache": True, + "hicache_storage_backend": "mooncake", + "hicache_io_backend": "direct", + "hicache_mem_layout": "layer_first", + }, + "expected_io_backend": "direct", + "expected_mem_layout": "page_first_direct", + }, + { + "name": "fa3_kernel_with_explicit_decode_backend", + "overrides": { + "enable_hierarchical_cache": True, + "hicache_io_backend": "kernel", + "hicache_mem_layout": "page_first", + "attention_backend": "triton", + "decode_attention_backend": "fa3", + }, + "expected_io_backend": "direct", + "expected_mem_layout": "page_first_direct", + }, + ] + + for case in cases: + with self.subTest(case=case["name"]): + args = self._make_args(**case["overrides"]) + args._handle_hicache() + self._assert_hicache_fields( + args, + expected_io_backend=case["expected_io_backend"], + expected_mem_layout=case["expected_mem_layout"], + ) + + @patch.object(ServerArgs, "use_mla_backend", return_value=False) + @patch("sglang.srt.server_args.is_flashinfer_available", return_value=False) + def test_decode_attention_backend_with_implicit_fa3( + self, _mock_flashinfer, _mock_use_mla_backend + ): + args = self._make_args( + enable_hierarchical_cache=True, + hicache_io_backend="kernel", + attention_backend="fa3", + decode_attention_backend=None, + ) + + args._handle_hicache() + + self.assertEqual(args.decode_attention_backend, "triton") + + if __name__ == "__main__": unittest.main()