Preallocate HiCache MHA staging before post-capture KV sizing (#40256)

test/registered/mem_cache/test_post_capture_kv_sizing.py

@@ -47,6 +47,13 @@ class TestPostCaptureKVSizing(CustomTestCase):
             timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
             env={**os.environ, "SGLANG_ENABLE_POST_CAPTURE_KV_SIZING": "1"},
             return_stdout_stderr=(cls.stdout, cls.stderr),
+            other_args=[
+                "--enable-hierarchical-cache",
+                "--hicache-mem-layout",
+                "page_first",
+                "--hicache-size",
+                "1",
+            ],
         )
 
     @classmethod
@@ -78,6 +85,10 @@ class TestPostCaptureKVSizing(CustomTestCase):
             "or the resize path did not run.",
         )
         self.assertGreater(float(m.group(1)), 0)
+        logs = self._server_logs()
+        staging = logs.find("HiCache staging prepared before KV sizing:")
+        self.assertGreaterEqual(staging, 0, "HiCache staging was not prepared")
+        self.assertLess(staging, m.start())
 
     def test_server_info_pool_sized(self):
         info = requests.get(f"{self.base_url}/server_info").json()

test/registered/unit/mem_cache/test_asymmetric_mha_pool_host_unit.py

@@ -84,7 +84,14 @@ class TestAsymmetricMHATokenToKVPoolHost(CustomTestCase):
         host = _make_host("page_first")
         host.page_num = 4
         host.v_head_dim = 8
-        host.device_pool = SimpleNamespace(device="cuda")
+        host.device_pool = SimpleNamespace(
+            device="cuda",
+            head_num=host.head_num,
+            head_dim=host.head_dim,
+            v_head_dim=host.v_head_dim,
+            store_dtype=host.dtype,
+            hicache_write_back_staging=None,
+        )
         cpu_empty = torch.empty
 
         def _cpu_empty(shape, *, dtype, device):

test/registered/unit/model_executor/model_runner_components/test_startup_weight_load.py

@@ -761,6 +761,7 @@ class TestStartupWeightLoadSchedulerRouting(CustomTestCase):
         worker = _SchedulerWorker(trace, post_capture_active=True)
         draft_worker = (
             SimpleNamespace(
+                hicache_draft_plan=None,
                 prewarm_sampling=lambda: trace.append("draft_prewarm"),
                 _draft_model_runners=lambda: (worker.model_runner,),
             )