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

python/sglang/srt/managers/scheduler.py

@@ -1127,6 +1127,15 @@ class Scheduler(
         self.init_all_cuda_graphs()
 
         model_runner = self.tp_worker.model_runner
+        if model_runner.token_to_kv_pool.post_capture_active:
+            kv_cache_builder.prepare_hicache_staging(
+                tp_worker=self.tp_worker,
+                draft_plan=(
+                    self.draft_worker.hicache_draft_plan
+                    if self.draft_worker is not None
+                    else None
+                ),
+            )
         device_module = torch.get_device_module(model_runner.device)
         self.schedule_stream = None if use_mlx() else device_module.Stream(priority=0)
         # Match run_batch / _pp_launch_batch so warmup allocations stay reusable.

python/sglang/srt/mem_cache/kv_cache_builder.py

@@ -38,10 +38,13 @@ from sglang.srt.configs.model_config import ModelImpl, is_deepseek_dsa
 from sglang.srt.environ import envs
 from sglang.srt.hardware_backend.mlx.runtime import use_mlx
 from sglang.srt.managers.mm_schedule import init_mm_embedding_cache
+from sglang.srt.mem_cache.base_swa_memory_pool import BaseSWAKVPool
 from sglang.srt.mem_cache.cache_init_params import CacheInitParams
 from sglang.srt.mem_cache.deepseek_v4_memory_pool import DeepSeekV4TokenToKVPool
 from sglang.srt.mem_cache.hicache_auto_size import auto_size_hicache
-from sglang.srt.mem_cache.memory_pool import MHATokenToKVPool
+from sglang.srt.mem_cache.memory_pool import HybridLinearKVPool, MHATokenToKVPool
+from sglang.srt.mem_cache.pool_host.base import _WRITE_BACK_STAGING_PAGE_CHUNK
+from sglang.srt.mem_cache.pool_host.mha import prepare_mha_write_back_staging
 from sglang.srt.mem_cache.registry import TreeCacheBuildContext, create_tree_cache
 from sglang.srt.mem_cache.swa_memory_pool import SWAKVPool
 from sglang.srt.mem_cache.unified_radix_cache import UnifiedRadixCache
@@ -53,6 +56,7 @@ from sglang.srt.runtime_context import (
     get_parallel,
     get_schedule,
 )
+from sglang.srt.speculative.base_spec_worker import HiCacheDraftMode
 from sglang.srt.utils import is_hip
 
 if TYPE_CHECKING:
@@ -67,6 +71,54 @@ if TYPE_CHECKING:
     from sglang.srt.speculative.spec_info import SpeculativeAlgorithm
 
 
+def prepare_hicache_staging(
+    *, tp_worker: BaseTpWorker, draft_plan: Optional[HiCacheDraftPlan] = None
+) -> None:
+    """Materialize MHA transfer buffers before the final KV budget is measured."""
+    memory = get_memory()
+    page_size = get_schedule().page_size
+    if memory.hicache_mem_layout != "page_first" or not (
+        memory.enable_hierarchical_cache
+        or get_disagg().disaggregation_decode_retraction_backup == "host_pool"
+    ):
+        return
+
+    def prepare(pool, packed_drafts=(), *, sidecar=False):
+        if isinstance(pool, SWAKVPool):
+            prepare(pool.full_kv_pool)
+            prepare(pool.swa_kv_pool, tuple(p.swa_kv_pool for p in packed_drafts))
+        elif isinstance(pool, HybridLinearKVPool):
+            prepare(pool.full_kv_pool, packed_drafts, sidecar=sidecar)
+        elif isinstance(pool, MHATokenToKVPool):
+            # Ratio-based host pools can only shrink with post-capture KV sizing.
+            # Sidecars instead inherit their target host pool's capacity.
+            page_capacity = _WRITE_BACK_STAGING_PAGE_CHUNK
+            if memory.hicache_size <= 0 and not sidecar:
+                page_capacity = int(pool.size * memory.hicache_ratio) // page_size + 1
+            staging = prepare_mha_write_back_staging(
+                pool,
+                layer_num=pool.layer_num + len(packed_drafts),
+                page_size=page_size,
+                page_capacity=page_capacity,
+            )
+            if staging is not None:
+                logger.info(
+                    "HiCache staging prepared before KV sizing: %.1f MiB, %d layers",
+                    sum(buffer.nbytes for buffer in staging) / (1 << 20),
+                    pool.layer_num + len(packed_drafts),
+                )
+
+    runner = tp_worker.model_runner
+    prepare(runner.token_to_kv_pool, runner.mtp_draft_device_pools)
+    if draft_plan is not None and draft_plan.mode == HiCacheDraftMode.SIDECAR:
+        for pool in draft_plan.device_pools:
+            # SWA sidecars follow only the draft's SWA component.
+            prepare(
+                pool.swa_kv_pool if isinstance(pool, BaseSWAKVPool) else pool,
+                sidecar=True,
+            )
+
+
 def get_draft_kv_pool(
     *,
     draft_worker: BaseTpWorker,
@@ -95,8 +147,6 @@ def maybe_register_hicache_draft(
     tree_cache,
     draft_plan: HiCacheDraftPlan,
 ) -> None:
-    from sglang.srt.speculative.base_spec_worker import HiCacheDraftMode
-
     if draft_plan.mode != HiCacheDraftMode.SIDECAR:
         return
 

python/sglang/srt/mem_cache/memory_pool.py

@@ -1958,6 +1958,8 @@ class KVCache(abc.ABC):
 
 
 class MHATokenToKVPool(KVCache):
+    hicache_write_back_staging: Optional[Tuple[torch.Tensor, torch.Tensor]] = None
+
     def __init__(
         self,
         size: int,

python/sglang/srt/mem_cache/pool_host/mha.py

@@ -71,6 +71,54 @@ if _is_npu:
 logger = logging.getLogger(__name__)
 
 
+def prepare_mha_write_back_staging(
+    device_pool: MHATokenToKVPool,
+    *,
+    layer_num: int,
+    page_size: int,
+    page_capacity: int = _WRITE_BACK_STAGING_PAGE_CHUNK,
+    retain: bool = True,
+) -> tuple[torch.Tensor, torch.Tensor] | None:
+    if not (_is_cuda or _is_hip) or layer_num == 0:
+        return None
+    head_num = device_pool.head_num
+    if device_pool.head_dim == device_pool.v_head_dim:
+        head_num = device_pool.row_dim // device_pool.head_dim
+    if not all(
+        can_use_write_back_jit_kernel(
+            element_size=head_num * dim * device_pool.store_dtype.itemsize
+        )
+        for dim in (device_pool.head_dim, device_pool.v_head_dim)
+    ):
+        return None
+    token_capacity = min(page_capacity, _WRITE_BACK_STAGING_PAGE_CHUNK) * page_size
+    shapes = tuple(
+        (token_capacity, layer_num, head_num, dim)
+        for dim in (device_pool.head_dim, device_pool.v_head_dim)
+    )
+    staging = device_pool.hicache_write_back_staging
+    if staging is None:
+        staging = (
+            torch.empty(
+                shapes[0], dtype=device_pool.store_dtype, device=device_pool.device
+            ),
+            torch.empty(
+                shapes[1], dtype=device_pool.store_dtype, device=device_pool.device
+            ),
+        )
+        if retain:
+            device_pool.hicache_write_back_staging = staging
+    else:
+        for buffer, shape in zip(staging, shapes):
+            if (
+                buffer.dtype != device_pool.store_dtype
+                or tuple(buffer.shape[1:]) != shape[1:]
+                or buffer.shape[0] < shape[0]
+            ):
+                raise ValueError("HiCache staging geometry changed after preparation")
+    return staging
+
+
 class MHATokenToKVPoolHost(HostKVCache):
     device_pool: MHATokenToKVPool | None = None
     mtp_draft_device_pools: tuple[MHATokenToKVPool, ...] = ()
@@ -219,32 +267,24 @@ class MHATokenToKVPoolHost(HostKVCache):
         self.staging_k_buffer = None
         self.staging_v_buffer = None
         self.can_use_write_back_jit = False
-        if self.layout != "page_first" or (_is_npu or _is_xpu or _is_mps):
+        if self.layout != "page_first":
             return
-
-        # The staged write-back JIT kernel builds with hipcc and has a ROCm
-        # path, so enable it on HIP too (consistent with the CUDA path).
-        self.can_use_write_back_jit = (
-            _is_cuda or _is_hip
-        ) and can_use_write_back_jit_kernel(
-            element_size=self.element_dim * self.dtype.itemsize,
+        page_capacity = min(self.page_num, _WRITE_BACK_STAGING_PAGE_CHUNK)
+        staging = prepare_mha_write_back_staging(
+            self.device_pool,
+            layer_num=self.layer_num,
+            page_size=self.page_size,
+            page_capacity=page_capacity,
+            retain=False,
         )
-        if not self.can_use_write_back_jit:
+        if staging is None:
             return
-
-        self.staging_page_capacity = min(self.page_num, _WRITE_BACK_STAGING_PAGE_CHUNK)
-        self.staging_token_capacity = self.staging_page_capacity * self.page_size
-        self.staging_k_buffer = torch.empty(
-            (
-                self.staging_token_capacity,
-                self.layer_num,
-                self.head_num,
-                self.head_dim,
-            ),
-            dtype=self.dtype,
-            device=self.device_pool.device,
+        self.can_use_write_back_jit = True
+        self.staging_page_capacity = page_capacity
+        self.staging_token_capacity = page_capacity * self.page_size
+        self.staging_k_buffer, self.staging_v_buffer = (
+            buffer[: self.staging_token_capacity] for buffer in staging
         )
-        self.staging_v_buffer = torch.empty_like(self.staging_k_buffer)
 
     @property
     def k_buffer(self):
@@ -1070,51 +1110,6 @@ class AsymmetricMHATokenToKVPoolHost(MHATokenToKVPoolHost):
     kernels derive copy sizes from each call's first tensor.
     """
 
-    def _init_write_back_staging_buffers(self):
-        self.staging_page_capacity = 0
-        self.staging_token_capacity = 0
-        self.staging_k_buffer = None
-        self.staging_v_buffer = None
-        self.can_use_write_back_jit = False
-        if self.layout != "page_first" or (_is_npu or _is_xpu or _is_mps):
-            return
-
-        # K and V have different element sizes. Use the single-buffer staged
-        # kernel for each side, which specializes to its native stride.
-        can_use_staged_jit = (_is_cuda or _is_hip) and all(
-            can_use_write_back_jit_kernel(element_size=element_size)
-            for element_size in (
-                self._k_token_stride_size(),
-                self._v_token_stride_size(),
-            )
-        )
-        if not can_use_staged_jit:
-            return
-
-        self.can_use_write_back_jit = True
-        self.staging_page_capacity = min(self.page_num, _WRITE_BACK_STAGING_PAGE_CHUNK)
-        self.staging_token_capacity = self.staging_page_capacity * self.page_size
-        self.staging_k_buffer = torch.empty(
-            (
-                self.staging_token_capacity,
-                self.layer_num,
-                self.head_num,
-                self.head_dim,
-            ),
-            dtype=self.dtype,
-            device=self.device_pool.device,
-        )
-        self.staging_v_buffer = torch.empty(
-            (
-                self.staging_token_capacity,
-                self.layer_num,
-                self.head_num,
-                self.v_head_dim,
-            ),
-            dtype=self.dtype,
-            device=self.device_pool.device,
-        )
-
     def get_size_per_token(self):
         self.head_num = self.device_pool.head_num
         self.head_dim = self.device_pool.head_dim

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,),
             )