HiCache: Add @rank_consensus to various functions (#37425)

Co-authored-by: Zhangheng <hzh0425@apache.org>
2026-09-14 10:44:21 +08:00
parent ca8ecc6a6f
commit bf9773e1da
24 changed files with 77 additions and 9 deletions
@@ -34,7 +34,10 @@ class TestHiCacheSpecFileStorage(HiCacheSpecStorageMixin, CustomTestCase):

    @classmethod
    def _get_spec_server_env(cls):
-        return {"SGLANG_HICACHE_FILE_BACKEND_STORAGE_DIR": cls.temp_dir}
+        return {
+            "SGLANG_ENABLE_RANK_CONSENSUS_CHECKER": "1",
+            "SGLANG_HICACHE_FILE_BACKEND_STORAGE_DIR": cls.temp_dir,
+        }

    @classmethod
    def _count_file_storage_pages(cls):
@@ -156,6 +156,7 @@ class TestHiCacheSpecMooncakeStorage(
    @classmethod
    def _get_spec_server_env(cls):
        return {
+            "SGLANG_ENABLE_RANK_CONSENSUS_CHECKER": "1",
            "MOONCAKE_MASTER": f"127.0.0.1:{cls.mooncake_master_port}",
            "MOONCAKE_PROTOCOL": "tcp",
            "MC_MS_AUTO_DISC": "0",
@@ -44,6 +44,7 @@ class TestHiCache(CustomTestCase, MMLUMixin):
                "--hicache-storage-backend",
                "file",
            ],
+            env={"SGLANG_ENABLE_RANK_CONSENSUS_CHECKER": "1"},
        )

    @classmethod
@@ -103,6 +103,7 @@ class HiCacheStorageBaseMixin:

        additional_server_args, env_vars = cls._get_additional_server_args_and_env()
        env_vars["SGLANG_ENABLE_DETERMINISTIC_INFERENCE"] = "1"
+        env_vars["SGLANG_ENABLE_RANK_CONSENSUS_CHECKER"] = "1"
        server_args = cls._get_base_server_args()
        if additional_server_args:
            server_args.update(additional_server_args)
@@ -63,6 +63,7 @@ class TestHiCacheStorageRuntimeAttachDetach(CustomTestCase):
            "SGLANG_HICACHE_FILE_BACKEND_STORAGE_DIR": cls.temp_dir,
            # Make runs less flaky for CI/dev.
            "SGLANG_ENABLE_DETERMINISTIC_INFERENCE": "1",
+            "SGLANG_ENABLE_RANK_CONSENSUS_CHECKER": "1",
            **cls.extra_env,
        }

@@ -125,6 +125,7 @@ class TestHiCacheStorageUMBPBackend(CustomTestCase):
        env.update(
            {
                "SGLANG_ENABLE_DETERMINISTIC_INFERENCE": "1",
+                "SGLANG_ENABLE_RANK_CONSENSUS_CHECKER": "1",
                "SGLANG_ENABLE_UNIFIED_RADIX_TREE": "1",
                "SGLANG_DSV4_FP4_EXPERTS": "0",
                "SGLANG_HACK_FLASHMLA_BACKEND": "unified_kv_triton",
@@ -32,6 +32,7 @@ class HiCacheBaseServer(CustomTestCase):

    model_name = DEFAULT_MODEL_NAME_FOR_TEST
    hicache_args = []
+    server_env: dict = {}

    @classmethod
    def setUpClass(cls):
@@ -47,6 +48,7 @@ class HiCacheBaseServer(CustomTestCase):
            cls.base_url,
            timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
            other_args=cls.hicache_args,
+            env=cls.server_env,
        )

    @classmethod
@@ -58,6 +60,7 @@ class TestHiCacheStandard(HiCacheBaseServer, MMLUMixin):
    """Standard HiCache configuration tests"""

    model_name = DEFAULT_MODEL_NAME_FOR_TEST
+    server_env = {"SGLANG_ENABLE_RANK_CONSENSUS_CHECKER": "1"}
    hicache_args = [
        "--enable-hierarchical-cache",
        "--mem-fraction-static",
@@ -74,6 +77,7 @@ class TestHiCacheMLA(HiCacheBaseServer, MMLUMixin, MGSMEnMixin):
    """HiCache with MLA model tests"""

    model_name = DEFAULT_MLA_MODEL_NAME_FOR_TEST
+    server_env = {"SGLANG_ENABLE_RANK_CONSENSUS_CHECKER": "1"}
    hicache_args = [
        "--trust-remote-code",
        "--enable-hierarchical-cache",
@@ -92,6 +96,7 @@ class TestHiCacheEagle(HiCacheBaseServer, MMLUMixin):

    model_name = DEFAULT_TARGET_MODEL_EAGLE3
    needs_tokenizer = True
+    server_env = {"SGLANG_ENABLE_RANK_CONSENSUS_CHECKER": "1"}
    hicache_args = [
        "--enable-hierarchical-cache",
        "--hicache-ratio",
@@ -123,6 +128,7 @@ class TestHiCachePage(HiCacheBaseServer, MMLUMixin):
    """HiCache with custom page size tests"""

    model_name = DEFAULT_MODEL_NAME_FOR_TEST
+    server_env = {"SGLANG_ENABLE_RANK_CONSENSUS_CHECKER": "1"}
    hicache_args = [
        "--enable-hierarchical-cache",
        "--page-size",
@@ -57,7 +57,11 @@ class TestPPWithHiCache(unittest.TestCase):
            if value is not True:
                final_server_args.append(str(value))

-        env_vars = {**os.environ, **cls._mooncake_env()}
+        env_vars = {
+            **os.environ,
+            **cls._mooncake_env(),
+            "SGLANG_ENABLE_RANK_CONSENSUS_CHECKER": "1",
+        }

        try:
            cls.process = popen_launch_server(
@@ -35,6 +35,7 @@ class TestQwen35WithHiCache(CustomTestCase):
        cls.base_url = DEFAULT_URL_FOR_TEST
        cls.storage_dir = tempfile.mkdtemp(prefix="qwen35-hicache-")
        env = {
+            "SGLANG_ENABLE_RANK_CONSENSUS_CHECKER": "1",
            "SGLANG_HICACHE_FILE_BACKEND_STORAGE_DIR": cls.storage_dir,
        }
        cls.process = popen_launch_server(
@@ -110,7 +110,10 @@ class TestUnifiedRadixTreeInt8MambaCheckpointE2E(TestInt8MambaCheckpointE2E):
                cls.base_url,
                timeout=cls.timeout,
                other_args=cls.other_args,
-                env={"SGLANG_ENABLE_UNIFIED_RADIX_TREE": "1"},
+                env={
+                    "SGLANG_ENABLE_RANK_CONSENSUS_CHECKER": "1",
+                    "SGLANG_ENABLE_UNIFIED_RADIX_TREE": "1",
+                },
            )

    @classmethod
@@ -33,6 +33,7 @@ class TestMamba2ExtraBufferKL(KLDivergenceMixin, DefaultServerBase):
    """NemotronH (Mamba2) + extra_buffer: cache-hit logprobs match cold recompute."""

    model = "nvidia/NVIDIA-Nemotron-Nano-9B-v2"
+    server_env = {"SGLANG_ENABLE_RANK_CONSENSUS_CHECKER": "1"}

    # Decode-seeded reuse is the regression trigger (the graphed decode
    # track-save); the broken path fails at KL ~1.5, so 0.005 discriminates
@@ -40,6 +40,7 @@ class TestRadixCacheFCFS(CustomTestCase):
                "--schedule-policy",
                "fcfs",
            ],
+            env={"SGLANG_ENABLE_RANK_CONSENSUS_CHECKER": "1"},
        )

    @classmethod
@@ -68,6 +69,7 @@ class TestRadixCacheLPM(TestRadixCacheFCFS):
                "--schedule-policy",
                "lpm",
            ],
+            env={"SGLANG_ENABLE_RANK_CONSENSUS_CHECKER": "1"},
        )


@@ -89,6 +91,7 @@ class TestRadixCacheNonOverlapLPM(TestRadixCacheFCFS):
                "--schedule-policy",
                "lpm",
            ],
+            env={"SGLANG_ENABLE_RANK_CONSENSUS_CHECKER": "1"},
        )


@@ -31,6 +31,7 @@ class TestRadixCacheHit(CustomTestCase):
            cls.model,
            cls.base_url,
            timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
+            env={"SGLANG_ENABLE_RANK_CONSENSUS_CHECKER": "1"},
        )

    @classmethod
@@ -14,6 +14,7 @@ class TestSWARadixCacheKL(KLDivergenceMixin, DefaultServerBase):
    model = MODEL
    kl_div_thres = 0.02  # it was 0.002
    kl_div_decode_max_new_tokens = 2048
+    server_env = {"SGLANG_ENABLE_RANK_CONSENSUS_CHECKER": "1"}
    other_args = [
        "--tp-size",
        "1",
@@ -99,7 +99,10 @@ class TestUnifiedQwen3HiCachePP(UnifiedRadixTreeTestMixin, CustomTestCase):
                "--hicache-mem-layout",
                cls.hicache_mem_layout,
            ],
-            env={"SGLANG_ENABLE_UNIFIED_RADIX_TREE": "1"},
+            env={
+                "SGLANG_ENABLE_RANK_CONSENSUS_CHECKER": "1",
+                "SGLANG_ENABLE_UNIFIED_RADIX_TREE": "1",
+            },
        )
        cls.input_ids = get_input_ids(cls.model, num_samples=18)

@@ -150,6 +153,7 @@ class TestUnifiedQwen3HiCachePPL3(AccuracyTwoPassMixin, CustomTestCase):
                "file",
            ],
            env={
+                "SGLANG_ENABLE_RANK_CONSENSUS_CHECKER": "1",
                "SGLANG_ENABLE_UNIFIED_RADIX_TREE": "1",
                "SGLANG_HICACHE_FILE_BACKEND_STORAGE_DIR": cls.hicache_dir,
            },
@@ -63,7 +63,10 @@ class TestUnifiedQwen3HiCacheCP(UnifiedRadixTreeTestMixin, CustomTestCase):
                "--hicache-mem-layout",
                cls.hicache_mem_layout,
            ],
-            env={"SGLANG_ENABLE_UNIFIED_RADIX_TREE": "1"},
+            env={
+                "SGLANG_ENABLE_RANK_CONSENSUS_CHECKER": "1",
+                "SGLANG_ENABLE_UNIFIED_RADIX_TREE": "1",
+            },
        )
        cls.input_ids = get_input_ids(cls.model, num_samples=18)

@@ -92,7 +92,10 @@ class TestUnifiedKimiLinearDcpHiCache(UnifiedRadixTreeTestMixin, CustomTestCase)
                str(MAX_MAMBA_CACHE_SIZE),
                "--enable-metrics",
            ],
-            env={"SGLANG_ENABLE_UNIFIED_RADIX_TREE": "1"},
+            env={
+                "SGLANG_ENABLE_RANK_CONSENSUS_CHECKER": "1",
+                "SGLANG_ENABLE_UNIFIED_RADIX_TREE": "1",
+            },
        )
        cls.input_ids = get_input_ids(cls.model, num_samples=18, trust_remote_code=True)

@@ -70,6 +70,7 @@ class TestGLM5UnifiedRadixCacheL3Accuracy(AccuracyTwoPassMixin, CustomTestCase):
            ],
            env={
                "SGLANG_HICACHE_FILE_BACKEND_STORAGE_DIR": cls.hicache_dir,
+                "SGLANG_ENABLE_RANK_CONSENSUS_CHECKER": "1",
                "SGLANG_ENABLE_UNIFIED_RADIX_TREE": "1",
            },
        )
@@ -63,7 +63,10 @@ class TestUnifiedMambaRadixCache(UnifiedRadixTreeTestMixin, CustomTestCase):
                "--mamba-max-states-per-path",
                "3",
            ],
-            env={"SGLANG_ENABLE_UNIFIED_RADIX_TREE": "1"},
+            env={
+                "SGLANG_ENABLE_RANK_CONSENSUS_CHECKER": "1",
+                "SGLANG_ENABLE_UNIFIED_RADIX_TREE": "1",
+            },
        )
        cls.input_ids = get_input_ids(cls.model, num_samples=18)

@@ -122,7 +125,10 @@ class TestUnifiedMambaHiCache(UnifiedRadixTreeTestMixin, CustomTestCase):
                "4",
                "--weight-loader-prefetch-checkpoints",
            ],
-            env={"SGLANG_ENABLE_UNIFIED_RADIX_TREE": "1"},
+            env={
+                "SGLANG_ENABLE_RANK_CONSENSUS_CHECKER": "1",
+                "SGLANG_ENABLE_UNIFIED_RADIX_TREE": "1",
+            },
        )
        cls.input_ids = get_input_ids(cls.model, num_samples=18)

@@ -190,6 +196,7 @@ class TestUnifiedMambaHiCacheL3(AccuracyTwoPassMixin, CustomTestCase):
                "4",
            ],
            env={
+                "SGLANG_ENABLE_RANK_CONSENSUS_CHECKER": "1",
                "SGLANG_ENABLE_UNIFIED_RADIX_TREE": "1",
                "SGLANG_HICACHE_FILE_BACKEND_STORAGE_DIR": cls.hicache_dir,
            },
@@ -90,6 +90,7 @@ class TestUnifiedMiMoHiCacheLoadBackKL(CustomTestCase):
                "page_first",
            ],
            env={
+                "SGLANG_ENABLE_RANK_CONSENSUS_CHECKER": "1",
                "SGLANG_ENABLE_UNIFIED_RADIX_TREE": "1",
                "SGLANG_USE_CUDA_IPC_TRANSPORT": "1",
            },
@@ -37,7 +37,10 @@ class TestUnifiedSWARadixCache(UnifiedRadixTreeTestMixin, CustomTestCase):
                "0.7",
                "--cuda-graph-backend-prefill=disabled",
            ],
-            env={"SGLANG_ENABLE_UNIFIED_RADIX_TREE": "1"},
+            env={
+                "SGLANG_ENABLE_RANK_CONSENSUS_CHECKER": "1",
+                "SGLANG_ENABLE_UNIFIED_RADIX_TREE": "1",
+            },
        )
        cls.input_ids = get_input_ids(cls.model, num_samples=18)