diff --git a/python/sglang/srt/managers/scheduler.py b/python/sglang/srt/managers/scheduler.py index f633e64b9..c89a6314e 100644 --- a/python/sglang/srt/managers/scheduler.py +++ b/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. diff --git a/python/sglang/srt/mem_cache/kv_cache_builder.py b/python/sglang/srt/mem_cache/kv_cache_builder.py index 0007d143f..f2484bb53 100644 --- a/python/sglang/srt/mem_cache/kv_cache_builder.py +++ b/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 diff --git a/python/sglang/srt/mem_cache/memory_pool.py b/python/sglang/srt/mem_cache/memory_pool.py index 0ae0ec0af..1ffc80733 100644 --- a/python/sglang/srt/mem_cache/memory_pool.py +++ b/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, diff --git a/python/sglang/srt/mem_cache/pool_host/mha.py b/python/sglang/srt/mem_cache/pool_host/mha.py index 371e4eec0..5ac3e1bdb 100644 --- a/python/sglang/srt/mem_cache/pool_host/mha.py +++ b/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 diff --git a/test/registered/mem_cache/test_post_capture_kv_sizing.py b/test/registered/mem_cache/test_post_capture_kv_sizing.py index 0233d8b93..bbe344509 100644 --- a/test/registered/mem_cache/test_post_capture_kv_sizing.py +++ b/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() diff --git a/test/registered/unit/mem_cache/test_asymmetric_mha_pool_host_unit.py b/test/registered/unit/mem_cache/test_asymmetric_mha_pool_host_unit.py index 5313ce66e..d6fb6de05 100644 --- a/test/registered/unit/mem_cache/test_asymmetric_mha_pool_host_unit.py +++ b/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): diff --git a/test/registered/unit/model_executor/model_runner_components/test_startup_weight_load.py b/test/registered/unit/model_executor/model_runner_components/test_startup_weight_load.py index b10cf53ae..1816cd510 100644 --- a/test/registered/unit/model_executor/model_runner_components/test_startup_weight_load.py +++ b/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,), )